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