NequIP

class mlip.models.nequip.models.Nequip(config: ~mlip.models.nequip.config.NequipConfig, dataset_info: ~mlip.data.dataset_info.DatasetInfo, parent: ~flax.linen.module.Module | ~flax.core.scope.Scope | ~flax.linen.module._Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None)

The NequIP model flax module. It is derived from the MLIPNetwork class.

References

• Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13(1), May 2022. ISSN: 2041-1723. URL: https://dx.doi.org/10.1038/s41467-022-29939-5.

config

Hyperparameters / configuration for the NequIP model, see NequipConfig.

Type:

mlip.models.nequip.config.NequipConfig

dataset_info

Hyperparameters dictated by the dataset (e.g., cutoff radius or average number of neighbors).

Type:

mlip.data.dataset_info.DatasetInfo

__call__(edge_vectors: Array, node_species: Array, senders: Array, receivers: Array) → Array

Compute node-wise energy summands. This function must be overridden by the implementation of MLIPNetwork.

class mlip.models.nequip.config.NequipConfig(*, num_layers: Annotated[int, Gt(gt=0)] = 2, node_irreps: Annotated[str, AfterValidator(func=_check_irreps)] = '128x0e + 128x0o + 64x1o + 64x1e + 4x2e + 4x2o', l_max: Annotated[int, Ge(ge=0)] = 3, num_bessel: Annotated[int, Gt(gt=0)] = 8, radial_net_nonlinearity: Activation = Activation.SWISH, radial_net_n_hidden: Annotated[int, Gt(gt=0)] = 64, radial_net_n_layers: Annotated[int, Gt(gt=0)] = 2, radial_envelope: RadialEnvelope = RadialEnvelope.POLYNOMIAL, scalar_mlp_std: Annotated[float, Gt(gt=0)] = 4.0, atomic_energies: str | dict[int, float] | None = None, avg_num_neighbors: float | None = None, num_species: Annotated[int, FieldInfo(annotation=NoneType, required=True, metadata=[Gt(gt=0)])] | None = None)

The configuration / hyperparameters of the NequIP model.

num_layers

Number of NequIP layers. Default is 2.

Type:

int

node_irreps

The O3 representation space of node features, with number of channels that may depend on the degree l. Default is "128x0e + 128x0o + 64x1o + 64x1e + 4x2e + 4x2o".

Type:

str

l_max

Maximal degree of spherical harmonics used for the angular encoding of edge vectors. Default is 3.

Type:

int

num_bessel

The number of Bessel basis functions to use (default is 8).

Type:

int

radial_net_nonlinearity

Activation function for radial MLP. Default is raw_swish.

Type:

mlip.models.options.Activation

radial_net_n_hidden

Number of hidden features in radial MLP. Default is 64.

Type:

int

radial_net_n_layers

Number of layers in radial MLP. Default is 2.

Type:

int

radial_envelope

The radial envelope function, by default it is "polynomial_envelope". The only other option is "soft_envelope".

Type:

mlip.models.options.RadialEnvelope

scalar_mlp_std

Standard deviation of weight init. of radial MLP. Default is 4.0.

Type:

float

atomic_energies

How to treat the atomic energies. If set to None (default) or the string "average", then the average atomic energies stored in the dataset info are used. It can also be set to the string "zero" which means not to use any atomic energies in the model. Lastly, one can also pass an atomic energies dictionary via this parameter different from the one in the dataset info, that is used.

Type:

str | dict[int, float] | None

avg_num_neighbors

The mean number of neighbors for atoms. If None (default), use the value from the dataset info. It is used to rescale messages by this value.

Type:

float | None

num_species

The number of elements (atomic species descriptors) allowed. If None (default), infer the value from the atomic energies map in the dataset info.

Type:

int | None