MEGNet utilities

This module provides some convenience utilities used in the unlockNN backend. For most end users, the only function of interest here is:

unlocknn.megnet_utils.default_megnet_config(nfeat_bond: int = 100, r_cutoff: float = 5.0, gaussian_width: float = 0.5) → dict

Get sensible defaults for MEGNetModel configuration.

These arguments are taken from the MEGNet README file.

Examples

Create a MEGNetModel using these defaults:

>>> model = MEGNetModel(**default_megnet_config())

Backend functions

unlocknn.megnet_utils.scale_targets(meg_model: MEGNetModel, targets: List[Union[float, ndarray]], num_atoms: List[int]) → List[Union[float, ndarray]]

Scale target values using given MEGNetModel’s scaler.

Parameters

  • meg_model – The MEGNetModel whose scaler to use.

  • targets – A list of target values.

  • num_atoms – A list of the number of atoms in each structure corresponding to the target values.

Returns

The scaled target values.

unlocknn.megnet_utils.create_megnet_input(meg_model: MEGNetModel, inputs: List[Union[Structure, Dict[str, Union[ndarray, List[Union[int, float]]]]]], targets: Optional[List[Union[float, ndarray]]] = None, batch_size: int = 128, scrub_failed_structs: bool = False, shuffle: bool = True) → Tuple[Union[GraphBatchDistanceConvert, GraphBatchGenerator], List[Dict[str, Union[ndarray, List[Union[int, float]]]]]]

Create generator for model inputs.

Parameters

  • meg_model – The MEGNetModel whose graph converter to use.

  • inputs – The input, either as graphs or structures.

  • targets – The input targets, if any.

  • batch_size – The batch size for the generator.

  • scrub_failed_structures – Whether to discard structures that could not be converted to graphs.

  • shuffle – Whether the generator should shuffle the order of the structure/target pairs.

Returns

The input generator.

A list of the model input graphs.