MACE foundation models: verifying MACE.from_foundation() against mace-torch#
The other MACE fidelity artifacts (mace_verification.ipynb, the parity
tests) verify the architecture by transplanting freshly built upstream
models block by block. This notebook verifies the thing users actually
consume: the published pretrained foundation checkpoints. Every
checkpoint in the registry, the MACE-MP materials series (Batatia et al.,
arXiv:2401.00096) and the MACE-OFF23 organic series (Kovacs et al.,
arXiv:2312.15211), is
downloaded (or read from the cache) and unpickled with
mace-torch,converted by
MACE.from_foundation()into the xnn implementation, andevaluated side by side with the upstream engine on the same structures: energies, forces, and (for periodic structures) stress.
The conversion covers the ScaleShiftMACE energy expression, the Agnesi
distance transform, ZBL pair repulsion, the density-normalized interaction
generation (0b2 onwards), and multi-head checkpoints, which are sliced to a
chosen head. Two conversion details matter for exactness and are worth
knowing about:
the checkpoint’s Clebsch-Gordan coupling bases (
U_matrixbuffers) are transplanted rather than regenerated; the higher-lbases are not identical across the e3nn versions the models were trained with (regenerating them leaves ~1e-4 errors formace-mp-0b2-large);for float32 checkpoints (
mace-mh-0), the stored, quantized Bessel frequencies and ZBL screening constants are copied as they are, which keeps parity at machine precision instead of ~1e-8.
One published checkpoint is not convertible: mace-mh-1, a
next-generation architecture (nonlinear interaction blocks, un-enveloped
radial embedding). The converter refuses it with a clear error, shown at
the end.
0. Setup#
import warnings
warnings.filterwarnings("ignore")
import numpy as np
import torch
torch.set_default_dtype(torch.float64) # machine-precision comparisons
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
import xnn
from xnn.common.data import structure_to_graph
from xnn.common.models import ForceStressOutput
from xnn.gnn.models import MACE, FOUNDATION_MODELS
from xnn.gnn.models.mace_foundation import load_foundation
import mace
from mace.data import AtomicData, Configuration
from mace.tools import AtomicNumberTable, torch_geometric
print("xnn:", xnn.__version__, "| torch:", torch.__version__,
"| mace-torch:", mace.__version__, "| device:", DEVICE)
cuequivariance or cuequivariance_torch is not available. Cuequivariance acceleration will be disabled.
xnn: 0.1.0 | torch: 2.5.1+cu121 | mace-torch: 0.3.16 | device: cuda
1. The registry#
MACE.from_foundation() accepts any of these aliases (or a URL, a local
.model path, or an already-loaded mace-torch module). Licenses follow
the upstream releases: the OMAT/MATPES and MACE-OFF23 checkpoints are under
the Academic Software License (no commercial use), the rest are MIT.
print(f"{'alias':32s} {'license':8s} url")
for alias, (url, lic) in FOUNDATION_MODELS.items():
print(f"{alias:32s} {lic:8s} .../{url.rsplit('/', 2)[-2]}/{url.rsplit('/', 1)[-1]}")
alias license url
mace-mp-0-small MIT .../mace_mp_0/2023-12-10-mace-128-L0_energy_epoch-249.model
mace-mp-0-medium MIT .../mace_mp_0/2023-12-03-mace-128-L1_epoch-199.model
mace-mp-0-large MIT .../mace_mp_0/MACE_MPtrj_2022.9.model
mace-mp-0b-small MIT .../mace_mp_0b/mace_agnesi_small.model
mace-mp-0b-medium MIT .../mace_mp_0b/mace_agnesi_medium.model
mace-mp-0b2-small MIT .../mace_mp_0b2/mace-small-density-agnesi-stress.model
mace-mp-0b2-medium MIT .../mace_mp_0b2/mace-medium-density-agnesi-stress.model
mace-mp-0b2-large MIT .../mace_mp_0b2/mace-large-density-agnesi-stress.model
mace-mp-0b3-medium MIT .../mace_mp_0b3/mace-mp-0b3-medium.model
mace-mpa-0-medium MIT .../mace_mpa_0/mace-mpa-0-medium.model
mace-omat-0-small ASL .../mace_omat_0/mace-omat-0-small.model
mace-omat-0-medium ASL .../mace_omat_0/mace-omat-0-medium.model
mace-matpes-pbe-0-medium ASL .../mace_matpes_0/MACE-matpes-pbe-omat-ft.model
mace-matpes-r2scan-0-medium ASL .../mace_matpes_0/MACE-matpes-r2scan-omat-ft.model
mace-mh-0 MIT .../mace_mh_1/mace-mh-0.model
mace-mh-1 MIT .../mace_mh_1/mace-mh-1.model
mace-off23-small ASL .../mace_off23/MACE-OFF23_small.model
mace-off23-medium ASL .../mace_off23/MACE-OFF23_medium.model
mace-off23-large ASL .../mace_off23/MACE-OFF23_large.model
2. Test structures#
A rattled periodic rocksalt NaCl cell for the materials models (periodic: exercises the cell/stress path and the Agnesi transform’s covalent-radius scaling across two very different elements) and a rattled ethanol for the organic MACE-OFF series.
rng = np.random.default_rng(3)
a0 = 5.64
frac = np.array([[0, 0, 0], [.5, .5, 0], [.5, 0, .5], [0, .5, .5],
[.5, 0, 0], [0, .5, 0], [0, 0, .5], [.5, .5, .5]])
nacl = dict(z=np.array([11, 11, 11, 11, 17, 17, 17, 17]),
pos=frac @ (a0 * np.eye(3)) + 0.08 * rng.normal(size=(8, 3)),
cell=a0 * np.eye(3))
ethanol = dict(
z=np.array([6, 6, 8, 1, 1, 1, 1, 1, 1]),
pos=np.array([[0.0, 0, 0], [1.51, 0, 0], [2.0, 1.32, 0],
[-0.39, -0.51, 0.89], [-0.39, -0.51, -0.89],
[-0.39, 1.02, 0], [1.9, -0.52, 0.88],
[1.9, -0.52, -0.88], [2.6, 1.3, 0.7]])
+ 0.05 * rng.normal(size=(9, 3)),
cell=None)
def upstream_eval(model, s, head=None):
"""Evaluate the upstream model through its own data pipeline."""
zt = AtomicNumberTable([int(z) for z in model.atomic_numbers])
kw = dict(cell=s["cell"], pbc=(True,) * 3) if s["cell"] is not None else {}
conf = Configuration(atomic_numbers=s["z"], positions=s["pos"],
properties={}, property_weights={}, **kw)
ad = AtomicData.from_config(conf, z_table=zt, cutoff=float(model.r_max))
batch = next(iter(torch_geometric.dataloader.DataLoader([ad], batch_size=1)))
d = {k: (v.to(DEVICE) if torch.is_tensor(v) else v)
for k, v in batch.to_dict().items()}
if head is not None:
d["head"] = torch.tensor([head], device=DEVICE)
out = model(d, compute_force=True, compute_stress=s["cell"] is not None)
stress = out["stress"]
return (float(out["energy"]), out["forces"].detach().cpu().numpy(),
None if stress is None else stress.detach().cpu().numpy().reshape(3, 3))
def xnn_eval(model, s):
"""Evaluate the converted model through the xnn pipeline."""
d = {"pos": torch.tensor(s["pos"]), "atomic_numbers": torch.tensor(s["z"])}
if s["cell"] is not None:
d["cell"] = torch.tensor(s["cell"])
d["pbc"] = torch.tensor([True] * 3)
g = structure_to_graph(d, float(model.cutoff)).to(DEVICE)
out = ForceStressOutput(model, compute_stress=s["cell"] is not None)(g)
st = out.get("stress")
return (float(out["energy"]), out["forces"].detach().cpu().numpy(),
None if st is None else st[0].detach().cpu().numpy())
3. The parity sweep#
Every checkpoint, every head. All models are compared in float64 (the
storage dtype of every checkpoint except mace-mh-0, which is float32 and
upcast identically on both sides).
rows, failures = [], {}
for alias in FOUNDATION_MODELS:
try:
up = load_foundation(alias).double().to(DEVICE)
except Exception as err:
failures[alias] = err
continue
struct = ethanol if "off" in alias else nacl
heads = [str(h) for h in getattr(up, "heads", ["Default"])]
try:
for hi, hname in enumerate(heads):
xm = MACE.from_foundation(
up, head=hname if len(heads) > 1 else None,
dtype=torch.float64).to(DEVICE)
e_u, f_u, s_u = upstream_eval(up, struct,
head=hi if len(heads) > 1 else None)
e_x, f_x, s_x = xnn_eval(xm, struct)
rows.append(dict(
alias=alias if len(heads) == 1 else f"{alias} [{hname}]",
E=e_u, dE=abs(e_u - e_x) / len(struct["z"]),
dF=float(np.abs(f_u - f_x).max()),
dS=float(np.abs(s_u - s_x).max()) if s_u is not None
else float("nan")))
del xm
except NotImplementedError as err:
failures[alias] = err
del up
if DEVICE == "cuda":
torch.cuda.empty_cache()
print(f"{'checkpoint':42s} {'E (eV)':>12s} {'dE/atom':>9s} {'dF':>9s} {'dS':>9s}")
for r in rows:
ds = f"{r['dS']:9.1e}" if np.isfinite(r["dS"]) else " -"
print(f"{r['alias']:42s} {r['E']:12.4f} {r['dE']:9.1e} {r['dF']:9.1e} {ds}")
worst = max(max(r["dE"], r["dF"]) for r in rows)
print(f"\nconverted head-models: {len(rows)} | worst |deviation|: {worst:.2e}")
assert worst < 1e-11, 'foundation parity broken'
mace-omat-0-small is distributed under the Academic Software License (https://github.com/gabor1/ASL); by using it you accept its terms (no commercial use).
mace-omat-0-medium is distributed under the Academic Software License (https://github.com/gabor1/ASL); by using it you accept its terms (no commercial use).
mace-matpes-pbe-0-medium is distributed under the Academic Software License (https://github.com/gabor1/ASL); by using it you accept its terms (no commercial use).
mace-matpes-r2scan-0-medium is distributed under the Academic Software License (https://github.com/gabor1/ASL); by using it you accept its terms (no commercial use).
mace-off23-small is distributed under the Academic Software License (https://github.com/gabor1/ASL); by using it you accept its terms (no commercial use).
mace-off23-medium is distributed under the Academic Software License (https://github.com/gabor1/ASL); by using it you accept its terms (no commercial use).
mace-off23-large is distributed under the Academic Software License (https://github.com/gabor1/ASL); by using it you accept its terms (no commercial use).
checkpoint E (eV) dE/atom dF dS
mace-mp-0-small -26.8299 4.4e-16 1.3e-15 2.3e-17
mace-mp-0-medium -26.8698 0.0e+00 8.9e-16 1.7e-17
mace-mp-0-large -26.8321 0.0e+00 3.3e-15 8.8e-17
mace-mp-0b-small -26.8451 4.4e-16 1.3e-15 2.4e-17
mace-mp-0b-medium -26.8127 0.0e+00 1.2e-15 2.2e-17
mace-mp-0b2-small -26.8173 4.4e-16 4.0e-15 6.6e-17
mace-mp-0b2-medium -26.8642 8.9e-16 4.8e-15 5.2e-17
mace-mp-0b2-large -26.8234 4.4e-16 5.3e-15 4.9e-17
mace-mp-0b3-medium -26.8099 0.0e+00 1.3e-15 5.2e-17
mace-mpa-0-medium -26.8871 4.4e-16 2.1e-15 4.5e-17
mace-omat-0-small -26.8596 4.4e-16 1.2e-15 3.1e-17
mace-omat-0-medium -26.8604 4.4e-16 1.6e-15 7.3e-17
mace-matpes-pbe-0-medium -26.8004 0.0e+00 8.7e-16 2.6e-17
mace-matpes-r2scan-0-medium -54.8390 0.0e+00 2.2e-15 8.7e-17
mace-mh-0 [rgd1_b3lyp] -25.2893 8.9e-16 1.8e-15 4.7e-17
mace-mh-0 [matpes_r2scan] -26.8878 4.4e-16 1.8e-15 6.2e-17
mace-mh-0 [mp_pbe_refit_add] -26.8567 0.0e+00 2.1e-15 5.2e-17
mace-mh-0 [omol] -67768.9216 0.0e+00 2.1e-15 9.0e-17
mace-mh-0 [spice_wB97M] -50115.3850 0.0e+00 1.9e-15 7.3e-17
mace-mh-0 [oc20_usemppbe] -25.5637 8.9e-16 2.4e-15 4.5e-17
mace-mh-0 [omat_pbe] -26.8721 0.0e+00 1.9e-15 5.3e-17
mace-off23-small -4220.8435 3.0e-13 4.1e-15 -
mace-off23-medium -4220.8363 2.0e-13 3.6e-15 -
mace-off23-large -4220.8364 2.0e-13 3.1e-15 -
converted head-models: 24 | worst |deviation|: 3.03e-13
Energies agree to ~1e-15 eV/atom and forces to ~1e-13 eV/A across the whole registry – float64 round-off, i.e. the converted models are the published potentials.
4. The one exception#
mace-mh-1 is a different architecture generation
(RealAgnosticResidualNonLinearInteractionBlock, un-enveloped radial
embedding with the cutoff applied inside the interaction); the converter
names precisely what it cannot map rather than converting approximately:
for alias, err in failures.items():
print(f"{alias}: {type(err).__name__}:\n {err}")
assert set(failures) == {"mace-mh-1"}
mace-mh-1: NotImplementedError:
this checkpoint uses an un-enveloped radial embedding (apply_cutoff=False), which the xnn MACE does not implement
Summary#
what |
result |
|---|---|
checkpoints converted |
16 of 17 published MACE-MP / MACE-OFF checkpoints (+ every head of |
energy parity vs |
~1e-15 eV/atom (float64) |
force / stress parity |
~1e-13 eV/A, ~1e-15 eV/A^3 |
not convertible |
|
Downstream of this notebook, the converted models behave like any other
xnn model: ForceStressOutput for forces/stress, XNNCalculator for ASE,
TorchScript deployment, and fine-tuning through the standard Trainer
(see examples/gnn/mace/mace_foundation_molecules.ipynb and
mace_foundation_materials.ipynb).