MACE, block by block: reproducing the original implementation with xnn#
This notebook checks every piece needed to reproduce the original MACE model
(ACEsuit/mace) using the xnn
re-implementation (xnn.gnn.models.mace). For each architectural block we
state the defining equation(s) from the MACE papers,
show the corresponding
xnnbuilding block,run it on a small toy system, and
compare it numerically against the original
mace-torchblock (our ground truth).
The three papers used here (provided alongside this notebook):
Batatia et al., The design space of E(3)-equivariant atom-centred interatomic potentials: NeurIPS-era preprint arXiv:2205.06643 (2022) and the journal version Nat. Mach. Intell. 7, 56 (2025). This is the Multi-ACE framework paper: it defines the one-particle basis, the atomic (\(A\)) basis, the product basis, the symmetrised (\(B\)) basis and the message/update equations in full generality. Equation numbers below (e.g. Multi-ACE eq 13–21) refer to the Nat. Mach. Intell. version.
Kovács et al., MACE-OFF: J. Am. Chem. Soc. 147, 17598 (2025). Section 2.1 gives the cleanest closed-form statement of the MACE architecture actually used in practice (eqs 1–8). We anchor each block to these.
Ground truth. We import the original
mace-torchpackage and compare thexnnblocks against it. Where a block has no learnable weights (radial basis, cutoff, spherical harmonics, the Clebsch–Gordan \(U\) tensors) the two agree to machine precision out of the box. Where a block has weights, we transplant the weights frommace-torchintoxnnand check that the outputs then match to ~\(10^{-16}\). The notebook ends by transplanting an entire MACE model and showing the total energy and per-atom forces are bit-for-bit identical.
0. Setup#
We work in float64 throughout (MACE’s default), which is what makes
exact numerical comparison meaningful.
# silence the expected warnings
import logging
import warnings
# 1. Silence the cuEquivariance library warning log
logging.getLogger("cuequivariance").setLevel(logging.ERROR)
# 2. Silence the TorchScript UserWarning
warnings.filterwarnings(
"ignore",
category=UserWarning,
message="The TorchScript type system doesn't support",
)
# 3. Silence the torch.load FutureWarning from e3nn
warnings.filterwarnings(
"ignore",
category=FutureWarning,
message="You are using `torch.load` with `weights_only=False`",
)
import numpy as np
import torch
torch.set_default_dtype(torch.float64) # MACE default; required for exact comparison
torch.manual_seed(0)
from e3nn import o3
import xnn, mace, e3nn
print("xnn :", xnn.__version__)
print("mace :", mace.__version__, "(original ACEsuit/mace -- ground truth)")
print("e3nn :", e3nn.__version__)
print("torch:", torch.__version__, "| CUDA:", torch.cuda.is_available())
xnn : 0.1.0
mace : 0.3.16 (original ACEsuit/mace -- ground truth)
e3nn : 0.4.4
torch: 2.5.1+cu121 | CUDA: True
A toy system#
A handful of atoms of three elements (H, C, O). Everything downstream is checked on this single small graph so each block stays inspectable.
from xnn.common.data import structure_to_graph
SPECIES = [1, 6, 8] # H, C, O -> element channels in MACE
CUTOFF = 5.0 # r_max
rng = np.random.default_rng(1)
pos = rng.uniform(0, 4, (7, 3))
Z = np.array(([1, 6, 8] * 7)[:7])
graph = structure_to_graph({"pos": pos, "atomic_numbers": Z}, CUTOFF)
print("atoms :", graph.num_nodes)
print("edges :", graph.num_edges, " (neighbour list within r_max =", CUTOFF, ")")
print("Z :", Z.tolist())
def report(name, diff, tol=1e-12):
tag = "OK " if diff <= tol else "!! "
print(f"{tag}{name:<46} max|xnn - mace| = {diff:.2e}")
atoms : 7
edges : 42 (neighbour list within r_max = 5.0 )
Z : [1, 6, 8, 1, 6, 8, 1]
The MACE architecture in one picture#
For each atom \(i\) MACE builds a site energy and sums them (\(E=\sum_i E_i\)). One layer \(t\) does (MACE-OFF eqs 1–8):
We now reproduce each underbraced piece in turn.
Block 1: Chemical (element) embedding · MACE-OFF eq 1 / Multi-ACE eq 13#
The initial node feature is a learnable embedding of the atomic number into \(k\) channels (MACE-OFF eq 1):
In the Multi-ACE language this is the \(T_{kc}(\theta_i,\theta_j)\) map applied to
the one-hot element attribute \(\theta\) (eq 13). In xnn the one-hot is produced
by _GNNBase.node_attr() and the linear map \(W\) is MACE.node_embedding
(an o3.Linear); the original MACE calls it LinearNodeEmbeddingBlock. Same
parameters, so transplanting \(W\) makes them identical.
from mace.modules.blocks import LinearNodeEmbeddingBlock
from xnn.common.models import build_model
from xnn.common.config import from_dict
# build an xnn MACE just to grab its node embedding + one-hot machinery
cfg = from_dict(
{
"model": {
"name": "mace",
"cutoff": CUTOFF,
"n_features": 8,
"n_interactions": 2,
"extra": {
"species": SPECIES,
"max_ell": 2,
"max_L": 1,
"num_channels": 8,
"correlation": 3,
"hidden_irreps": "8x0e+8x1o",
},
}
}
)
xnn_mace = build_model(cfg.model)
node_attrs = xnn_mace.node_attr(graph.atomic_numbers) # one-hot theta (N, 3)
node_attr_irreps = xnn_mace.node_attr_irreps # 3x0e
feat_irreps = o3.Irreps("8x0e")
mace_embed = LinearNodeEmbeddingBlock(node_attr_irreps, feat_irreps)
xnn_mace.node_embedding.load_state_dict(mace_embed.linear.state_dict()) # transplant W
h0_xnn = xnn_mace.node_embedding(node_attrs)
h0_mace = mace_embed(node_attrs)
report("element embedding h^(0)", (h0_xnn - h0_mace).abs().max().item())
print("one-hot theta (first 3 atoms):\n", node_attrs[:3].int().numpy())
cuequivariance or cuequivariance_torch is not available. Cuequivariance acceleration will be disabled.
OK element embedding h^(0) max|xnn - mace| = 0.00e+00
one-hot theta (first 3 atoms):
[[1 0 0]
[0 1 0]
[0 0 1]]
Block 2: Radial basis (Bessel × polynomial cutoff, then a radial MLP) · Multi-ACE eq 9 / MACE-OFF radial block#
MACE’s learnable radial function is (Multi-ACE eq 9 / preprint eq 27)
where \(R_n\) are Bessel basis functions and \(f_{\rm cut}\) is a smooth polynomial cutoff (MACE-OFF “radial embedding block”). The Bessel functions are
and the polynomial envelope of degree \(p\) is the C²-smooth Klicpera form.
xnn provides BesselRBF, PolynomialCutoff (the \(R_n f_{\rm cut}\) product is the
SphericalHarmonicEdgeEmbedding.edge_radial); the per-path MLP is the
FullyConnectedNet conv_tp_weights inside each interaction. We compare the
basis and the cutoff directly against mace.modules.radial.
from xnn.gnn.featurizers.radial import BesselRBF
from xnn.gnn.featurizers.cutoff import PolynomialCutoff as XPoly
from mace.modules.radial import BesselBasis, PolynomialCutoff as MPoly
from mace.modules.blocks import RadialEmbeddingBlock
r = torch.linspace(0.2, 4.9, 60)
NRBF, P = 8, 5
xb = BesselRBF(NRBF, CUTOFF)(r)
mb = BesselBasis(CUTOFF, NRBF, trainable=False)(r.unsqueeze(-1))
report("Bessel basis R_n(r)", (xb - mb).abs().max().item())
xc = XPoly(CUTOFF, p=P)(r)
mc = MPoly(CUTOFF, p=P)(r)
report("polynomial cutoff f_cut(r)", (xc - mc).abs().max().item())
# full radial embedding R_n(r) * f_cut(r)
xr = BesselRBF(NRBF, CUTOFF)(r) * XPoly(CUTOFF, p=P)(r)[:, None]
mre = RadialEmbeddingBlock(CUTOFF, NRBF, P, radial_type="bessel")
mr, _ = mre(r.unsqueeze(-1), None, None, None)
report("radial embedding R_n * f_cut", (xr - mr).abs().max().item())
OK Bessel basis R_n(r) max|xnn - mace| = 5.97e-16
OK polynomial cutoff f_cut(r) max|xnn - mace| = 0.00e+00
OK radial embedding R_n * f_cut max|xnn - mace| = 4.72e-16
Block 3: Spherical harmonics of the edge direction · \(Y^{m}_{l}(\hat r_{ij})\) in eq 2#
The angular part of the one-particle basis is the real spherical harmonics of the
unit edge vector, up to degree \(\ell_{\max}\) (max_ell). Both codes call the
same e3nn routine with the same normalisation
(normalize=True, normalization="component"), so they agree to machine precision.
In xnn this lives in SphericalHarmonicEdgeEmbedding (edge_sh).
vec = graph.edge_vectors()
ir_sh = o3.Irreps.spherical_harmonics(2) # 1x0e+1x1o+1x2e (max_ell=2)
edge = xnn_mace.edge_feat(graph) # xnn edge featurizer
xsh = edge["edge_sh"]
msh = o3.spherical_harmonics(ir_sh, vec, normalize=True, normalization="component")
report("spherical harmonics Y_l^m(r_hat)", (xsh - msh).abs().max().item())
print("edge_sh irreps:", ir_sh, " shape:", tuple(xsh.shape))
OK spherical harmonics Y_l^m(r_hat) max|xnn - mace| = 0.00e+00
edge_sh irreps: 1x0e+1x1o+1x2e shape: (42, 9)
Blocks 4 & 5: One-particle basis \(\phi\) and atomic basis \(A\) · MACE-OFF eqs 2–3 / Multi-ACE eqs 14, 19#
The one-particle (edge) basis couples the neighbour feature, the radial function and the spherical harmonics through Clebsch–Gordan coefficients (MACE-OFF eq 2):
and the permutation-invariant atomic basis sums it over neighbours (MACE-OFF eq 3 / Multi-ACE eq 14):
In xnn this is exactly the body of RealAgnostic(Residual)InteractionBlock:
the weighted o3.TensorProduct (conv_tp, with per-edge weights from the radial
MLP) followed by a scatter-sum over edge_index and a channel-mixing o3.Linear.
We build the original and xnn interaction blocks with identical irreps, copy
the state dict across, and confirm the message \(A\) (and the residual
self-connection \(sc\)) match bit-for-bit.
from mace.modules.blocks import RealAgnosticResidualInteractionBlock as MIB
from xnn.gnn.models.mace import RealAgnosticResidualInteractionBlock as XIB
na = o3.Irreps("3x0e") # node attrs (one-hot, 3 elements)
nf = o3.Irreps("8x0e") # node feats (k=8 scalar channels)
ea = o3.Irreps.spherical_harmonics(2) # edge attrs (SH)
ef = o3.Irreps("8x0e") # edge radial feats (n_bessel=8)
ti = o3.Irreps("8x0e+8x1o+8x2e") # interaction (target) irreps
hi = o3.Irreps("8x0e+8x1o") # hidden irreps
kw = dict(node_attrs_irreps=na, node_feats_irreps=nf, edge_attrs_irreps=ea,
edge_feats_irreps=ef, target_irreps=ti, hidden_irreps=hi,
avg_num_neighbors=8.0, radial_MLP=[16, 16])
m_int, x_int = MIB(**kw), XIB(**kw)
x_int.load_state_dict(m_int.state_dict()) # transplant ALL interaction weights
N, E = graph.num_nodes, graph.num_edges
nfeat = torch.randn(N, nf.dim)
edge_attrs = torch.randn(E, ea.dim); edge_feats = torch.randn(E, ef.dim)
ei = graph.edge_index
xA, xsc = x_int(node_attrs, nfeat, edge_attrs, edge_feats, ei)
mA, msc = m_int(node_attrs, nfeat, edge_attrs, edge_feats, ei, cutoff=None)
report("atomic basis / message A", (xA - mA).abs().max().item())
report("residual self-connection sc", (xsc - msc).abs().max().item())
OK atomic basis / message A max|xnn - mace| = 0.00e+00
OK residual self-connection sc max|xnn - mace| = 0.00e+00
Block 6: Product basis and the generalized Clebsch–Gordan \(U\) tensors · MACE-OFF eq 4 / Multi-ACE eqs 15–16, 20, 42#
Higher body order comes from the product basis: the \(\nu\)-fold product of the atomic basis with itself (MACE-OFF/Multi-ACE):
symmetrised into the equivariant \(B\) basis through the generalized Clebsch–Gordan coefficients \(\mathcal C^{LM}_{\eta,\mathbf v}\) (Multi-ACE eq 20), which are themselves products of ordinary CG coefficients (Multi-ACE eq 42):
The coupling coefficients are precomputed as the \(U\) tensors. xnn’s
U_matrix_real is adapted from mace-torch; here we confirm it is bit-identical
to the original mace.modules.symmetric_contraction.U_matrix_real for correlation
orders \(\nu=1,2,3\) (body orders 2, 3, 4; MACE uses \(\nu=3\)).
from xnn.gnn.models.mace import U_matrix_real as xU
from mace.modules.symmetric_contraction import U_matrix_real as mU
coupling = o3.Irreps("1x0e+1x1o+1x2e+1x3o")
for nu in (1, 2, 3):
u_x = xU(coupling, "0e", nu, dtype=torch.float64)[-1]
u_m = mU(coupling, o3.Irreps("0e"), nu, dtype=torch.float64, use_cueq_cg=False)[-1]
report(f"U tensor (generalized CG), nu={nu} shape={tuple(u_x.shape)}",
(u_x - u_m).abs().max().item(), tol=0.0)
OK U tensor (generalized CG), nu=1 shape=(16, 1) max|xnn - mace| = 0.00e+00
OK U tensor (generalized CG), nu=2 shape=(16, 16, 4) max|xnn - mace| = 0.00e+00
OK U tensor (generalized CG), nu=3 shape=(16, 16, 16, 23) max|xnn - mace| = 0.00e+00
Block 7: Symmetric contraction → message \(m\) · MACE-OFF eqs 4–5 / Multi-ACE eq 21#
The learned, per-element contraction over the \(U\) basis turns the atomic basis \(A\) into the message (MACE-OFF eq 5 / Multi-ACE eq 21):
This is the genuinely MACE-specific operation. xnn’s SymmetricContraction
mirrors the original mace.modules.symmetric_contraction.SymmetricContraction.
The two store the per-correlation weights in a different order
(mace keeps the top order as weights_max, the rest descending; xnn keeps an
ascending weights[ν-1] list), so we map the weights and then check the contracted
output matches to ~\(10^{-16}\).
from xnn.gnn.models.mace import SymmetricContraction as XSC
from mace.modules.symmetric_contraction import SymmetricContraction as MSC
irreps_in = o3.Irreps("8x0e+8x1o+8x2e+8x3o")
irreps_out = o3.Irreps("8x0e+8x1o")
CORR, NEL = 3, 3
xsc = XSC(irreps_in, irreps_out, correlation=CORR, num_elements=NEL)
msc = MSC(irreps_in, irreps_out, correlation=CORR, num_elements=NEL,
irrep_normalization="component", path_normalization="element",
use_reduced_cg=False)
def transplant_sc(xsc, msc, corr):
with torch.no_grad():
for c in range(len(xsc.contractions)):
xc, mc = xsc.contractions[c], msc.contractions[c]
xc.weights[corr - 1].copy_(mc.weights_max) # top order
for nu in range(1, corr):
xc.weights[nu - 1].copy_(mc.weights[corr - 1 - nu]) # lower orders
transplant_sc(xsc, msc, CORR)
B = 5
nfeat_dim = irreps_in.count((0, 1))
x = torch.randn(B, nfeat_dim, o3.Irreps([ir.ir for ir in irreps_in]).dim)
y = torch.zeros(B, NEL); y[torch.arange(B), torch.randint(0, NEL, (B,))] = 1.0 # element one-hot
report("symmetric contraction m", (xsc(x, y) - msc(x, y)).abs().max().item())
OK symmetric contraction m max|xnn - mace| = 8.88e-16
Block 8: Equivariant product block + update with self-connection · MACE-OFF eq 6 / Multi-ACE eqs 22, 30–31#
The full higher-order block raises the body order via the symmetric contraction, mixes channels with a linear map, and adds the element-dependent self-connection (residual update, MACE-OFF eq 6 / preprint eqs 30–31):
xnn’s _EquivariantProductBasis matches the original
EquivariantProductBasisBlock. We transplant the contraction weights (same mapping
as Block 7) and the output o3.Linear, then confirm the block output matches.
from mace.modules.blocks import EquivariantProductBasisBlock as MPB
from xnn.gnn.models.mace import _EquivariantProductBasis as XPB
mpb = MPB(node_feats_irreps=irreps_in, target_irreps=irreps_out,
correlation=CORR, num_elements=NEL, use_sc=True)
xpb = XPB(irreps_in, irreps_out, correlation=CORR, num_elements=NEL, use_sc=True)
transplant_sc(xpb.symmetric_contractions, mpb.symmetric_contractions, CORR)
xpb.linear.load_state_dict(mpb.linear.state_dict())
nfeats = torch.randn(B, nfeat_dim, o3.Irreps([ir.ir for ir in irreps_in]).dim)
sc = torch.randn(B, irreps_out.dim)
report("equivariant product + update", (xpb(nfeats, sc, y) - mpb(nfeats, sc, y)).abs().max().item())
OK equivariant product + update max|xnn - mace| = 4.44e-16
Block 9: Readouts · MACE-OFF eqs 7–8#
The site energy is a sum of read-outs of the per-layer features; the read-out is linear for the early layers and a small MLP (gated) for the last layer (MACE-OFF eq 8):
xnn provides _LinearReadout and _NonLinearReadout, matching the original
LinearReadoutBlock / NonLinearReadoutBlock (same linear_1 → gate → linear_2).
from mace.modules.blocks import LinearReadoutBlock, NonLinearReadoutBlock
from xnn.gnn.models.mace import _LinearReadout, _NonLinearReadout
import torch.nn.functional as Fn
# linear readout (early layers)
mlin = LinearReadoutBlock(hi)
xlin = _LinearReadout(hi); xlin.linear.load_state_dict(mlin.linear.state_dict())
hin = torch.randn(N, hi.dim)
report("linear readout", (xlin(hin) - mlin(hin)).abs().max().item())
# non-linear (gated) readout (final layer)
hidden_scalars = o3.Irreps("8x0e")
mnl = NonLinearReadoutBlock(hidden_scalars, o3.Irreps("16x0e"), Fn.silu)
xnl = _NonLinearReadout(hidden_scalars, o3.Irreps("16x0e"), Fn.silu)
xnl.linear_1.load_state_dict(mnl.linear_1.state_dict())
xnl.linear_2.load_state_dict(mnl.linear_2.state_dict())
hs = torch.randn(N, hidden_scalars.dim)
report("non-linear (gated) readout", (xnl(hs) - mnl(hs)).abs().max().item())
OK linear readout max|xnn - mace| = 0.00e+00
OK non-linear (gated) readout max|xnn - mace| = 0.00e+00
Block 10: Site-energy sum, reference energy \(E_0\), and conservative forces · MACE-OFF eq 7 & §2.1#
Per-atom energies are summed to the total energy, with a per-element reference
energy \(E_{0,z}\) added (the “isolated-atom” baseline; MACE-OFF eq 6 keeps the
isolated-atom limit exact). In xnn, \(E_0\) is the atom_ref embedding and the
sum is aggregate_energy. Finally,
are obtained by autograd: xnn’s ForceStressOutput wraps any model and
differentiates through it (the original computes the same conservative forces).
We demonstrate the property MACE relies on: that with zero interactions
(\(T=0\)) the energy is exactly the sum of \(E_0\) and forces vanish.
from xnn.common.models import ForceStressOutput
cfg0 = from_dict({"model": {"name": "mace", "cutoff": CUTOFF, "n_features": 8,
"n_interactions": 0, "extra": {"species": SPECIES, "max_ell": 2, "max_L": 0,
"num_channels": 8, "atomic_energies": [0.5, -1.3, -2.1]}}})
m0 = ForceStressOutput(build_model(cfg0.model))
o0 = m0(graph)
E0_expected = sum({1: 0.5, 6: -1.3, 8: -2.1}[int(z)] for z in Z)
report("T=0 energy == sum of E0", abs(float(o0["energy"]) - E0_expected))
report("T=0 forces == 0", o0["forces"].abs().max().item())
OK T=0 energy == sum of E0 max|xnn - mace| = 0.00e+00
OK T=0 forces == 0 max|xnn - mace| = 0.00e+00
Capstone: transplant a whole MACE model and compare energy & forces#
Blocks 1–10 are the complete set of pieces. As the final check we build a full
xnn MACE and the original mace.modules.MACE with identical hyper-parameters
(\(T=2\) layers, \(\ell_{\max}=2\), \(L_{\max}=1\), correlation \(\nu=3\), 8 channels),
transplant every weight (embedding, both interactions, both product blocks,
both readouts, and \(E_0\)), and run both on the toy molecule.
Total energy is a scalar and per-atom forces are physical quantities (both are
independent of the internal edge ordering), so they must agree exactly if the two
models are the same function. We use the original code’s own data pipeline for the
mace side.
import mace.modules as mm
from mace.tools import AtomicNumberTable, torch_geometric
from mace.data import AtomicData, Configuration
T, NCH, MAXL, MAX_ELL, NRBF = 2, 8, 1, 2, 8
HID, MLP, RMLP, AVG = "8x0e+8x1o", "16x0e", [16, 16], 2.0
E0 = np.array([0.5, -1.3, -2.1])
cfg = from_dict({"model": {"name": "mace", "cutoff": CUTOFF, "n_features": NCH,
"n_interactions": T, "extra": {"species": SPECIES, "max_ell": MAX_ELL,
"correlation": 3, "num_channels": NCH, "max_L": MAXL, "hidden_irreps": HID,
"MLP_irreps": MLP, "radial_MLP": RMLP, "num_bessel": NRBF,
"avg_num_neighbors": AVG, "atomic_energies": E0.tolist()}}})
xfull = build_model(cfg.model)
xmodel = ForceStressOutput(xfull, compute_forces=True).double()
mmodel = mm.MACE(
r_max=CUTOFF, num_bessel=NRBF, num_polynomial_cutoff=5, max_ell=MAX_ELL,
interaction_cls=MIB, interaction_cls_first=MIB, num_interactions=T,
num_elements=len(SPECIES), hidden_irreps=o3.Irreps(HID), MLP_irreps=o3.Irreps(MLP),
atomic_energies=E0, avg_num_neighbors=AVG, atomic_numbers=SPECIES, correlation=3,
gate=torch.nn.functional.silu, radial_MLP=RMLP, radial_type="bessel",
use_reduced_cg=False, apply_cutoff=True).double()
# ---- transplant every weight: mace -> xnn ----
with torch.no_grad():
xfull.node_embedding.load_state_dict(mmodel.node_embedding.linear.state_dict())
for i in range(T):
xfull.interactions[i].load_state_dict(mmodel.interactions[i].state_dict())
transplant_sc(xfull.products[i].symmetric_contractions,
mmodel.products[i].symmetric_contractions, 3)
xfull.products[i].linear.load_state_dict(mmodel.products[i].linear.state_dict())
xr, mr = xfull.readouts[i], mmodel.readouts[i]
if "NonLinear" in type(mr).__name__:
xr.linear_1.load_state_dict(mr.linear_1.state_dict())
xr.linear_2.load_state_dict(mr.linear_2.state_dict())
else:
xr.linear.load_state_dict(mr.linear.state_dict())
for z, e in zip(SPECIES, E0):
xfull.atom_ref.weight[z] = float(e)
# ---- xnn prediction ----
gx = structure_to_graph({"pos": pos, "atomic_numbers": Z}, CUTOFF)
ox = xmodel(gx); E_x = float(ox["energy"]); F_x = ox["forces"].detach().numpy()
# ---- original mace prediction (its own data pipeline) ----
zt = AtomicNumberTable(SPECIES)
conf = Configuration(atomic_numbers=Z, positions=pos, properties={}, property_weights={})
ad = AtomicData.from_config(conf, z_table=zt, cutoff=CUTOFF)
batch = next(iter(torch_geometric.dataloader.DataLoader([ad], batch_size=1)))
om = mmodel(batch.to_dict(), compute_force=True)
E_m = float(om["energy"]); F_m = om["forces"].detach().numpy()
print(f"total energy xnn = {E_x:.10f} eV")
print(f"total energy mace = {E_m:.10f} eV")
report("FULL MODEL total energy", abs(E_x - E_m))
report("FULL MODEL per-atom forces", np.abs(F_x - F_m).max())
total energy xnn = -5.6090262073 eV
total energy mace = -5.6090262073 eV
OK FULL MODEL total energy max|xnn - mace| = 0.00e+00
OK FULL MODEL per-atom forces max|xnn - mace| = 2.43e-16
Summary#
Every block needed to reproduce the original MACE was checked against
mace-torch on the toy system:
Block |
MACE eq |
Weights? |
agreement |
|---|---|---|---|
1. Element embedding |
MACE-OFF 1 / Multi-ACE 13 |
transplanted |
machine precision |
2. Bessel basis · cutoff · radial MLP |
Multi-ACE 9 |
none |
machine precision |
3. Spherical harmonics |
eq 2 (\(Y\)) |
none (same e3nn call) |
machine precision |
4–5. One-particle \(\phi\) + atomic basis \(A\) |
MACE-OFF 2–3 / Multi-ACE 14,19 |
transplanted |
bit-identical |
6. Generalized CG \(U\) tensors |
Multi-ACE 16,20,42 |
none |
bit-identical (0) |
7. Symmetric contraction → message |
MACE-OFF 5 / Multi-ACE 21 |
transplanted |
~1e-16 |
8. Product block + self-connection |
MACE-OFF 6 / Multi-ACE 22,30–31 |
transplanted |
~1e-16 |
9. Linear / gated readouts |
MACE-OFF 7–8 |
transplanted |
machine precision |
10. \(E_0\) + autograd forces |
MACE-OFF 7, §2.1 |
n/a |
exact (\(T=0\)) |
Full model |
eqs 1–8 |
all transplanted |
energy 0, forces ~1e-16 |
xnn.gnn.models.mace is therefore a faithful reproduction of the original MACE
architecture, depending only on e3nn (no mace-torch / cuequivariance).
The companion notebook mace_argon_train_test.ipynb trains and tests this MACE
on a realistic Argon MD dataset.