xnn.gnn.models.mace#

MACE (Batatia et al. 2022): higher body-order equivariant message passing.

A faithful, self-contained MACE built on the xnn equivariant-GNN abstractions: it subclasses EquivariantGNN (species bookkeeping, per-element reference energy atom_ref, and the SphericalHarmonicEdgeEmbedding edge featurizer) and adds the genuinely MACE-specific pieces – the real RealAgnostic(Residual)InteractionBlock and a learned symmetric contraction over Clebsch-Gordan paths (correlation order).

The CG U basis (U_matrix_real()) is bit-identical to mace-torch and the symmetric contraction reproduces it to ~1e-16 given the same weights (see tests/test_gnn.py). Only e3nn is required – no mace-torch, cuequivariance or opt_einsum_fx.

The model is TorchScript-deployable: the tensor-only MACE.node_energy() core compiles under torch.jit.script (used by the LAMMPS/TorchScript exporters in xnn.common.deploy) and reproduces the eager model to machine precision (see tests/test_mace.py).

Difference from upstream MACE: num_interactions (the number of message-passing layers T) is fully flexible – T = 0 (a pure atom_ref/pair-repulsion baseline) through any T = N – rather than being fixed to 2. All architecture options are read from ModelConfig.extra (see MACE.from_config()); upstream MACE-CLI spellings (r_max, atomic_numbers, E0s, …) are translated to the xnn names at config-load time by the key-translation registry in xnn.common.config.translate.

The CG coupling and the symmetric contraction are independent implementations (plain torch.einsum, built on e3nn’s o3.wigner_3j; no codegen/cueq deps), verified bit-identical against ACEsuit/mace (MIT licence). All conventions – CG normalization, coupling-path ordering, parameter and buffer names – follow upstream so trained mace-torch weights transplant directly.

Functions

U_matrix_real(irreps_in, irreps_out, correlation)

Symmetric coupling basis of correlation copies of irreps_in -> irreps_out.

Classes

`MACE`(species[, cutoff, max_ell, max_L, ...])	Faithful MACE with a flexible number of interaction layers (T = 0..N).
`RealAgnosticDensityInteractionBlock`(...)	Non-residual interaction with learned density normalization.
`RealAgnosticDensityResidualInteractionBlock`(...)	Residual interaction with learned density normalization.
`RealAgnosticInteractionBlock`(...)	Non-residual interaction: the skip connection is applied to the message.
`RealAgnosticResidualInteractionBlock`(...)	Residual interaction: self-connection computed from the input features.
`SymmetricContraction`(irreps_in, irreps_out, ...)	Per-element symmetric contraction over all output irreps (the MACE product basis).