MACE

class mlip.models.mace.config.MaceConfig(*, add_atomic_energies: bool = True, num_layers: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 2, num_channels: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 128, l_max: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)] = 3, node_symmetry: ~typing.Annotated[int, FieldInfo(annotation=NoneType, required=True, metadata=[Ge(ge=0)])] | None = None, correlation: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 3, readout_irreps: tuple[~typing.Annotated[str, ~pydantic.functional_validators.AfterValidator(func=~mlip.typing.fields._check_irreps)], ...] = ('16x0e', '0e'), num_readout_heads: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 1, include_pseudotensors: bool = False, num_rbf: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 8, activation: ~mlip.models.options.Activation = Activation.SILU, radial_envelope: ~mlip.models.options.RadialEnvelope = RadialEnvelope.POLYNOMIAL, radial_mlp_hidden: list[~typing.Annotated[int, FieldInfo(annotation=NoneType, required=True, metadata=[Gt(gt=0)])]] = <factory>, radial_mlp_activation: ~mlip.models.options.Activation = Activation.SILU, avg_num_neighbors: float | None = None, avg_r_min: float | None = None, num_species: int | None = None, gate_nodes: bool = False, residual_connection_first_layer: bool = False, soft_normalization: float | None = None, predict_partial_charges: bool = False, use_coulomb_term: bool = False, use_total_charge_embedding: bool = False, embed_activation: ~mlip.models.options.Activation = Activation.SILU, deterministic_scatter_ops: bool = False, symmetric_contraction_backend: ~typing.Literal['e3j', 'e3nn', 'e3nn_symmetric', 'gaunt_tp'] = 'e3j', use_gaunt_tp_message_passing: bool = False)

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_rbf

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

radial_mlp_hidden

Sizes of the radial MLP hidden layers. Default is [64, 64, 64].

Type:

list[int]

radial_mlp_activation

Activation function for the radial MLP. Default is "silu".

Type:

mlip.models.options.Activation

add_atomic_energies

Whether to add atomic energies to the final energies. Default is True.

Type:

bool

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

gate_nodes

Whether to use a gating for the self-interaction. Default is False. See our white paper for a description of this option that is not present in the original MACE architecture.

Type:

bool

residual_connection_first_layer

Include a skip connection on the first layer, the default is false.

Type:

bool

soft_normalization

Node features will be regularized so that their norm stay below this parameter’s value (soft saturation). The default is None.

Type:

float | None

predict_partial_charges

Whether the model will be trained to predict charges.

Type:

bool

use_coulomb_term

Whether to use the Coulomb term in the model for long range interactions. Default is False.

Type:

bool

use_total_charge_embedding

Whether to use the total charge embedding. Default is False.

Type:

bool

embed_activation

Activation function for the embedding block. Default is “silu”.

Type:

mlip.models.options.Activation

deterministic_scatter_ops

Whether to use deterministic scatter operations in the forward pass, ensuring deterministic energy outputs. Setting to True makes prediction slower. Default is False.

Type:

bool

symmetric_contraction_backend

Which backend to use for the symmetric contraction, if wanting to revert to e3nn. Default is e3j.

Type:

Literal[‘e3j’, ‘e3nn’, ‘e3nn_symmetric’, ‘gaunt_tp’]

use_gaunt_tp_message_passing

Whether to use Gaunt tensor products as a replacement for Clebsch-Gordan tensor products in the message passing block. Default is False.

Type:

bool

class mlip.models.mace.network.Mace(config: MaceConfig, dataset_info: DatasetInfo, parent: Module | Scope | _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__(graph: Graph) → Graph

Calculate node-wise energy contributions and update the input Graph.

Parameters:

graph – Input Graph object containing edge vectors, node species, senders, receivers, and node mask.

Returns:

The updated Graph with per-node energy contributions stored in the “energy” node feature.

class mlip.models.mace.layer.MaceLayer(use_residuals: bool, last_layer: bool, num_channels: int, source_irreps: ~e3nn_jax._src.irreps.Irreps, interaction_irreps: ~e3nn_jax._src.irreps.Irreps, node_irreps: ~e3nn_jax._src.irreps.Irreps, activation: ~mlip.models.options.Activation | ~typing.Callable[[~jax.Array], ~jax.Array], num_species: int, l_max: int, num_rbf: int, radial_mlp_hidden: list[int], radial_mlp_activation: ~mlip.models.options.Activation | ~typing.Callable[[~jax.Array], ~jax.Array], avg_num_neighbors: float, correlation: int, soft_normalization: float | None, output_irreps: ~e3nn_jax._src.irreps.Irreps, readout_mlp_irreps: ~e3nn_jax._src.irreps.Irreps, num_readout_heads: int = 1, gate_nodes: bool = False, deterministic_scatter_ops: bool = False, symmetric_contraction_backend: ~typing.Literal['e3j', 'e3nn', 'e3nn_symmetric', 'gaunt_tp'] = 'e3j', use_gaunt_tp_message_passing: bool = False, parent: ~flax.linen.module.Module | ~flax.core.scope.Scope | ~flax.linen.module._Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None)

A single MACE layer.

This module assumes that the input graph has been transformed by mlip.models.blocks.MaceEmbeddingBlock upstream.

__call__(graph: Graph) → Graph

Apply the MaceLayer on a featured graph.

Parameters:

graph – an input graph which should have the following features assigned: * edges.features['spherical_embedding'] * edges.features['radial_embedding']

Returns:

• nodes.features['latent']

• nodes.features['outputs']

Return type:

a graph with the following features updated

class mlip.models.mace.blocks.MaceEmbeddingBlock(num_species: int, num_charges: int | None, num_channels: int, l_max: int, r_max: float, num_rbf: int, radial_envelope: RadialEnvelope, avg_r_min: float | None = None, activation_fn: Callable = <PjitFunction of <function silu>>, parent: Module | Scope | _Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None)

Initial embedding block for MACE.

This module will assign the following graph features:

• nodes.features['embedding']: initial embedding of atomic numbers

• edges.features['radial_embedding']: RBF-encoding of interatomic distances.

• edges.features['spherical_embedding']: harmonic embedding of unit edge vectors.

__call__(graph: Graph) → Graph

Call self as a function.