xnn.gnn.models.mace.MACE#

class xnn.gnn.models.mace.MACE(species, cutoff=4.0, max_ell=3, max_L=0, num_channels=32, n_rbf=8, num_interactions=2, correlation=3, MLP_irreps='16x0e', radial_MLP=None, interaction='RealAgnosticResidualInteractionBlock', interaction_first='RealAgnosticResidualInteractionBlock', gate='silu', avg_num_neighbors=1.0, hidden_irreps=None, num_cutoff_basis=5, radial_type='bessel', distance_transform='None', pair_repulsion=False, atomic_energies=None, scale=1.0, shift=0.0)[source]#

Bases: EquivariantGNN

Faithful MACE with a flexible number of interaction layers (T = 0..N).

Subclasses EquivariantGNN, inheriting species bookkeeping, the per-element reference energy atom_ref, and the SphericalHarmonicEdgeEmbedding edge featurizer, and adds the MACE-specific interaction blocks and the learned symmetric-contraction product basis. The number of message-passing layers T = num_interactions is fully flexible (T = 0 gives a pure atom_ref/pair-repulsion baseline).

Parameters:

species (list of int) – Atomic numbers of the elements the model supports.
cutoff (float, optional) – Radial cutoff distance r_max in angstrom, by default 4.0.
max_ell (int, optional) – Maximum degree l of the spherical-harmonic edge attributes, by default 3.
max_L (int, optional) – Maximum output irrep order L of the hidden node features, by default 0 (invariant features only).
num_channels (int, optional) – Number of feature channels (multiplicity) in the hidden irreps, by default 32.
n_rbf (int, optional) – Number of radial basis functions, by default 8.
num_interactions (int, optional) – Number of message-passing layers T, by default 2. Must be >= 0.
correlation (int or list of int, optional) – Correlation order (body order minus one) of the symmetric contraction, by default 3. An int is broadcast across all interactions.
MLP_irreps (str, optional) – Hidden irreps of the final non-linear readout MLP, by default "16x0e".
radial_MLP (list of int or None, optional) – Hidden layer widths of the radial MLP producing tensor-product weights. None defaults to [64, 64, 64].
interaction (str, optional) – Interaction-block class name for layers after the first, by default "RealAgnosticResidualInteractionBlock".
interaction_first (str, optional) – Interaction-block class name for the first layer, by default "RealAgnosticResidualInteractionBlock".
gate (str or None, optional) – Name of the scalar gate activation for the non-linear readout, by default "silu".
avg_num_neighbors (float, optional) – Average neighbour count used to normalise messages, by default 1.0.
hidden_irreps (str or None, optional) – Explicit hidden irreps. None derives them from num_channels and max_L.
num_cutoff_basis (int, optional) – Polynomial cutoff degree p for the edge envelope (and ZBL), by default 5.
radial_type (str, optional) – Radial basis type ("bessel" or "gaussian"), by default "bessel".
distance_transform (str, optional) – Chemistry-aware warp of the distance fed to the radial basis (the cutoff envelope always sees the raw distance): "None" (default), "Agnesi" or "Soft" (see xnn.gnn.featurizers.radial; "Agnesi" is what the MACE-MP-0b and later foundation models use).
pair_repulsion (bool, optional) – If True, add a _ZBLPairRepulsion short-range term, by default False.
atomic_energies (torch.Tensor or None, optional) – Per-element reference energies (E0s) used to initialise atom_ref.
scale (float, optional) – Affine rescaling of the per-atom interaction energy (readouts plus pair repulsion), E_i = E0_i + scale * E_int,i + shift – the upstream ScaleShiftMACE convention. The defaults (1, 0) recover the plain MACE energy expression, so one class covers both upstream variants.
shift (float, optional) – Affine rescaling of the per-atom interaction energy (readouts plus pair repulsion), E_i = E0_i + scale * E_int,i + shift – the upstream ScaleShiftMACE convention. The defaults (1, 0) recover the plain MACE energy expression, so one class covers both upstream variants.

Raises:

ValueError – If num_interactions is negative.
NotImplementedError – If distance_transform is not one of the supported options.

node_features_energy(atomic_numbers, edge_index, edge_vec)[source]#

TorchScript-compatible core: tensors in, features + energies out.

The single implementation reused by node_energy() (the deploy entry point) and forward(), so it must avoid the AtomicGraph dataclass and any Python-only constructs. Accumulates the interaction energy – the optional ZBL pair-repulsion term plus one readout per round of interaction + product basis – maps it through scale_shift (identity unless constructed with scale/shift, the ScaleShiftMACE convention) and adds the per-element reference energy. The radial basis sees the (optionally distance_transform-warped) distance while the cutoff envelope always sees the raw one. The invariant (l = 0) channels of every layer’s node features are collected alongside (what LatentEwald consumes).

Parameters:

atomic_numbers (torch.Tensor) – Per-atom atomic numbers, shape (N,).
edge_index (torch.Tensor) – Edge index of shape (2, E); row 0 is the source (neighbour) and row 1 the destination (centre) node of each edge.
edge_vec (torch.Tensor) – Edge displacement vectors, shape (E, 3) (already accounting for any periodic cell shifts).

Returns:

The concatenated invariant node features (N, node_feature_dim) and the per-atom energy (N,).

Return type:

tuple of torch.Tensor

node_energy(atomic_numbers, edge_index, edge_vec)[source]#

Per-atom energy, shape (N,) (thin wrapper over node_features_energy(); the deploy wrappers call this).

Parameters:

atomic_numbers (Tensor)
edge_index (Tensor)
edge_vec (Tensor)

Return type:

Tensor

forward(data)[source]#

Predict per-node and total energy for an atomic graph.

Thin wrapper over node_energy() (the scriptable tensor core): builds the edge displacement vectors from the graph and aggregates the per-node energies to a per-structure total.

Parameters:: data (xnn.common.data.AtomicGraph) – The input atomic graph (atomic numbers, edge index, positions, …).
Returns:: "node_energy" (per-node energies) and "energy" (per-structure total energy).
Return type:: dict of str to torch.Tensor

classmethod from_config(cfg)[source]#

Construct a MACE from a core model config.

Reads the MACE-specific hyper-parameters from cfg.extra (falling back to defaults) and the shared fields from cfg directly. Only the xnn canonical key names are read here; upstream MACE-CLI spellings (r_max, num_radial_basis, atomic_numbers, E0s, …) are translated to these names at config-load time by xnn.common.config.translate. Values copied from an upstream yaml are coerced: species accepts a "[1, 6, 8]" string, radial_MLP a "[64, 64, 64]" string, and atomic_energies (MACE E0s) a list aligned with species, a {Z: E0} dict, or the string form of either.

Alternatively, extra["foundation"] names (or points to) a pretrained MACE foundation checkpoint: the model is then built by from_foundation() (optionally with extra["head"] and extra["dtype"]) and every architecture key is taken from the checkpoint instead of the config.

Parameters:: cfg (xnn.common.config.schema.ModelConfig) – The core model config, whose extra dict carries the MACE architecture options.
Returns:: The instantiated model.
Return type:: MACE

classmethod from_foundation(source, head=None, dtype=None)[source]#

Load a pretrained MACE foundation model into an xnn MACE.

Downloads (and caches) the requested checkpoint if needed, unpickles it with the mace-torch package, and converts it weight-for-weight into this implementation – covering the ScaleShiftMACE energy expression, the Agnesi distance transform, ZBL pair repulsion, the density-normalized interaction generation, and multi-head checkpoints (sliced to one head). See xnn.gnn.models.mace_foundation for the alias registry and the conversion details.

Parameters:

source (str or Path or torch.nn.Module) – A registered alias (e.g. "mace-mp-0-medium", "mace-off23-small"; see FOUNDATION_MODELS), a checkpoint URL or local path, or an already-loaded mace-torch model instance.
head (str, optional) – Which head of a multi-head checkpoint to keep. Defaults to the checkpoint’s only head; required (with the options listed in the error) when there are several.
dtype (torch.dtype or str, optional) – Final dtype of the converted model; None keeps the checkpoint’s (float64 for most foundation models).

Returns:

The converted model, ready for evaluation or fine-tuning.

Return type:

MACE