MACE¶

class mlip.models.mace.models.Mace(config: ~mlip.models.mace.config.MaceConfig, 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 MACE model flax module. It is derived from the MLIPNetwork class.

References

Ilyes Batatia, Dávid Péter Kovács, Gregor N. C. Simm, Christoph Ortner, and Gábor Csányi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields, 2023. URL: https://arxiv.org/abs/2206.07697.

config¶

Hyperparameters / configuration for the MACE model, see MaceConfig.

Type:: mlip.models.mace.config.MaceConfig

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.mace.config.MaceConfig(*, num_layers: Annotated[int, Gt(gt=0)] = 2, num_channels: Annotated[int, Gt(gt=0)] = 128, l_max: Annotated[int, Ge(ge=0)] = 3, node_symmetry: Annotated[int, FieldInfo(annotation=NoneType, required=True, metadata=[Gt(gt=0)])] | None = None, correlation: Annotated[int, Gt(gt=0)] = 3, readout_irreps: tuple[Annotated[str, AfterValidator(func=_check_irreps)], ...] = ('16x0e', '0e'), num_readout_heads: Annotated[int, Gt(gt=0)] = 1, include_pseudotensors: bool = False, num_bessel: Annotated[int, Gt(gt=0)] = 8, activation: Activation = Activation.SILU, radial_envelope: RadialEnvelope = RadialEnvelope.POLYNOMIAL, symmetric_tensor_product_basis: bool = False, atomic_energies: str | dict[int, float] | None = None, avg_num_neighbors: float | None = None, avg_r_min: float | None = None, num_species: int | None = None)¶

The configuration / hyperparameters of the MACE model.

num_layers¶

Number of MACE layers. Default is 2.

Type:: int

num_channels¶

The number of channels. Default is 128.

Type:: int

l_max¶

Highest degree of spherical harmonics used for the directional encoding of edge vectors, and during the convolution block. Default is 3, it is recommended to keep it at 3.

Type:: int

node_symmetry¶

Highest degree of node features kept after the node-wise power expansion of features, also called Atomic Cluster Expansion (ACE). The default behaviour is to assign l_max, although high values of node_symmetry may have a significant impact on runtime. It should be less or equal to l_max.

Type:: int | None

correlation¶

Maximum correlation order, by default it is 3.

Type:: int

readout_irreps¶

Irreps for the readout block, passed as a tuple of irreps string representations for each of the layers in the readout block. Currently, this MACE model only supports two layers, and it defaults to ("16x0e", "0e").

Type:: tuple[str, …]

num_readout_heads¶

Number of readout heads. The default is 1. For fine-tuning, additional heads must be added.

Type:: int

include_pseudotensors¶

If False (default), only parities p = (-1)**l will be kept. If True, all parities will be kept, e.g., "1e" pseudo-vectors returned by the cross product on R3.

Type:: bool

num_bessel¶

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

Type:: int

activation¶

The activation function used in the non-linear readout block. The options are "silu", "elu", "relu", "tanh", "sigmoid", and "swish". The default is "silu".

Type:: mlip.models.options.Activation

radial_envelope¶

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

Type:: mlip.models.options.RadialEnvelope

symmetric_tensor_product_basis¶

Whether to use a symmetric tensor product basis (default is False).

Type:: bool

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

avg_r_min¶

The mean minimum neighbour distance in Angstrom. If None (default), use the value from the dataset info.

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