{ "cells": [ { "cell_type": "markdown", "id": "87e3e46b", "metadata": {}, "source": [ "# MACE, block by block: reproducing the original implementation with `xnn`\n", "\n", "This notebook checks **every piece needed to reproduce the original MACE model**\n", "([ACEsuit/mace](https://github.com/ACEsuit/mace)) using the `xnn`\n", "re-implementation (`xnn.gnn.models.mace`). For each architectural block we\n", "\n", "1. state the **defining equation(s)** from the MACE papers,\n", "2. show the corresponding `xnn` building block,\n", "3. run it on a small **toy system**, and\n", "4. compare it **numerically against the original `mace-torch`** block (our ground truth).\n", "\n", "The three papers used here (provided alongside this notebook):\n", "\n", "* **Batatia et al., *The design space of E(3)-equivariant atom-centred interatomic\n", " potentials***: NeurIPS-era preprint arXiv:2205.06643 (2022) and the journal\n", " version *Nat. Mach. Intell.* **7**, 56 (2025). This is the *Multi-ACE* framework\n", " paper: it defines the one-particle basis, the atomic ($A$) basis, the product\n", " basis, the symmetrised ($B$) basis and the message/update equations in full\n", " generality. Equation numbers below (e.g. *Multi-ACE eq 13–21*) refer to the\n", " Nat. Mach. Intell. version.\n", "* **Kovács et al., *MACE-OFF***: *J. Am. Chem. Soc.* **147**, 17598 (2025).\n", " Section 2.1 gives the cleanest closed-form statement of the MACE architecture\n", " actually used in practice (eqs 1–8). We anchor each block to these.\n", "\n", "> **Ground truth.** We import the *original* `mace-torch` package and compare the\n", "> `xnn` blocks against it. Where a block has no learnable weights (radial basis,\n", "> cutoff, spherical harmonics, the Clebsch–Gordan $U$ tensors) the two agree to\n", "> machine precision out of the box. Where a block has weights, we **transplant the\n", "> weights** from `mace-torch` into `xnn` and check that the outputs then match to\n", "> ~$10^{-16}$. The notebook ends by transplanting an *entire* MACE model and\n", "> showing the total energy and per-atom forces are bit-for-bit identical.\n" ] }, { "cell_type": "markdown", "id": "b5016b62", "metadata": {}, "source": [ "## 0. Setup\n", "\n", "We work in `float64` throughout (MACE's default), which is what makes\n", "exact numerical comparison meaningful." ] }, { "cell_type": "code", "execution_count": 1, "id": "b0ed4396", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:18:13.989303Z", "iopub.status.busy": "2026-07-20T04:18:13.989194Z", "iopub.status.idle": "2026-07-20T04:18:20.750365Z", "shell.execute_reply": "2026-07-20T04:18:20.749450Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "xnn : 0.1.0\n", "mace : 0.3.16 (original ACEsuit/mace -- ground truth)\n", "e3nn : 0.4.4\n", "torch: 2.5.1+cu121 | CUDA: True\n" ] } ], "source": [ "# silence the expected warnings\n", "import logging\n", "import warnings\n", "\n", "# 1. Silence the cuEquivariance library warning log\n", "logging.getLogger(\"cuequivariance\").setLevel(logging.ERROR)\n", "\n", "# 2. Silence the TorchScript UserWarning\n", "warnings.filterwarnings(\n", " \"ignore\",\n", " category=UserWarning,\n", " message=\"The TorchScript type system doesn't support\",\n", ")\n", "\n", "# 3. Silence the torch.load FutureWarning from e3nn\n", "warnings.filterwarnings(\n", " \"ignore\",\n", " category=FutureWarning,\n", " message=\"You are using `torch.load` with `weights_only=False`\",\n", ")\n", "\n", "import numpy as np\n", "import torch\n", "\n", "torch.set_default_dtype(torch.float64) # MACE default; required for exact comparison\n", "torch.manual_seed(0)\n", "\n", "from e3nn import o3\n", "\n", "import xnn, mace, e3nn\n", "\n", "print(\"xnn :\", xnn.__version__)\n", "print(\"mace :\", mace.__version__, \"(original ACEsuit/mace -- ground truth)\")\n", "print(\"e3nn :\", e3nn.__version__)\n", "print(\"torch:\", torch.__version__, \"| CUDA:\", torch.cuda.is_available())" ] }, { "cell_type": "markdown", "id": "af91090f", "metadata": {}, "source": [ "### A toy system\n", "\n", "A handful of atoms of three elements (H, C, O). Everything downstream is checked\n", "on this single small graph so each block stays inspectable." ] }, { "cell_type": "code", "execution_count": 2, "id": "f3f899ee", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:18:20.752439Z", "iopub.status.busy": "2026-07-20T04:18:20.752157Z", "iopub.status.idle": "2026-07-20T04:18:20.795049Z", "shell.execute_reply": "2026-07-20T04:18:20.794345Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "atoms : 7\n", "edges : 42 (neighbour list within r_max = 5.0 )\n", "Z : [1, 6, 8, 1, 6, 8, 1]\n" ] } ], "source": [ "from xnn.common.data import structure_to_graph\n", "\n", "SPECIES = [1, 6, 8] # H, C, O -> element channels in MACE\n", "CUTOFF = 5.0 # r_max\n", "rng = np.random.default_rng(1)\n", "pos = rng.uniform(0, 4, (7, 3))\n", "Z = np.array(([1, 6, 8] * 7)[:7])\n", "\n", "graph = structure_to_graph({\"pos\": pos, \"atomic_numbers\": Z}, CUTOFF)\n", "print(\"atoms :\", graph.num_nodes)\n", "print(\"edges :\", graph.num_edges, \" (neighbour list within r_max =\", CUTOFF, \")\")\n", "print(\"Z :\", Z.tolist())\n", "\n", "def report(name, diff, tol=1e-12):\n", " tag = \"OK \" if diff <= tol else \"!! \"\n", " print(f\"{tag}{name:<46} max|xnn - mace| = {diff:.2e}\")" ] }, { "cell_type": "markdown", "id": "2306c48e", "metadata": {}, "source": [ "## The MACE architecture in one picture\n", "\n", "For each atom $i$ MACE builds a site energy and sums them\n", "($E=\\sum_i E_i$). One layer $t$ does (MACE-OFF eqs 1–8):\n", "\n", "$$\n", "\\begin{align}\n", "& h^{(0)}_{i,k00}=\\sum_z W_{kz}\\,\\delta_{z z_i} & & {\\text{element embedding}}\n", "\\\\[5pt]\n", "\n", "& \\phi^{(t)}_{ij,k\\mathbf{\\eta}_1l_3m_3} = \\sum_{l_1l_2m_1m_2} C^{l_3m_3}_{\\mathbf{\\eta}_1,l_1m_1l_2m_2}\\,R^{(t)}(r_{ij})_{k\\mathbf{\\eta}_1,l_1l_2l_3}\\,Y^{m_1}_{l_1}(\\hat r_{ij})\\,h^{(t)}_{j,kl_2m_2} & & {\\text{one-particle basis}}\n", "\\\\[5pt]\n", "\n", "& A^{(t)}_{i,kl_3m_3}= \\sum_{\\tilde{k}\\mathbf{\\eta}_1}W^{(t)}_{k\\tilde{k}\\mathbf{\\eta}_1l_3} \\sum_{j\\in\\mathcal N(i)}\\phi^{(t)}_{ij,k\\mathbf{\\eta}_1l_3m_3} & & {\\text{atomic basis}}\n", "\\\\[5pt]\n", "\n", "& \\mathbf B^{(t), \\mathbf{\\nu}}_{i,\\eta} = \\sum_{\\mathbf{lm}} \\mathcal C^{LM}_{\\mathbf{\\eta}_\\nu,\\mathbf{lm}}\\prod_{\\xi=1}^{\\mathbf{\\nu}} A^{(t)}_{i} & & {\\text{product / symmetrised basis}}\n", "\\\\[5pt]\n", "\n", "& m^{(t)}_{i}= \\sum_{\\eta} W\\,\\mathbf B^{(t),\\mathbf{v}}_{i,\\mathbf{\\eta}_\\mathbf{v}kLM} & & {\\text{message}}\n", "\\\\[5pt]\n", "\n", "& h^{(t+1)}_i = W\\,m^{(t)}_i + W_{z_i}\\,h^{(t)}_i & & {\\text{update + self-connection}}\n", "\\\\[5pt]\n", "\n", "& E_i=\\sum_{t} \\mathcal R^{(t)}(h^{(t)}_i) & & {\\text{readouts}}\n", "\\\\[5pt]\n", "\n", "& F = -\\,\\partial E/\\partial r & & {\\text{forces}}\n", "\\end{align}\n", "$$\n", "\n", "We now reproduce each underbraced piece in turn." ] }, { "cell_type": "markdown", "id": "eef8780c", "metadata": {}, "source": [ "## Block 1: Chemical (element) embedding · MACE-OFF eq 1 / Multi-ACE eq 13\n", "\n", "The initial node feature is a learnable embedding of the atomic number into $k$\n", "channels (MACE-OFF eq 1):\n", "\n", "$$h^{(0)}_{i,k00}=\\sum_z W_{kz}\\,\\delta_{z z_i}.$$\n", "\n", "In the Multi-ACE language this is the $T_{kc}(\\theta_i,\\theta_j)$ map applied to\n", "the one-hot element attribute $\\theta$ (eq 13). In `xnn` the one-hot is produced\n", "by `_GNNBase.node_attr()` and the linear map $W$ is `MACE.node_embedding`\n", "(an `o3.Linear`); the original MACE calls it `LinearNodeEmbeddingBlock`. Same\n", "parameters, so transplanting $W$ makes them identical." ] }, { "cell_type": "code", "execution_count": 3, "id": "8730430e", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:18:20.796706Z", "iopub.status.busy": "2026-07-20T04:18:20.796558Z", "iopub.status.idle": "2026-07-20T04:18:26.764297Z", "shell.execute_reply": "2026-07-20T04:18:26.763045Z" } }, "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": [ "OK element embedding h^(0) max|xnn - mace| = 0.00e+00\n", "one-hot theta (first 3 atoms):\n", " [[1 0 0]\n", " [0 1 0]\n", " [0 0 1]]\n" ] } ], "source": [ "from mace.modules.blocks import LinearNodeEmbeddingBlock\n", "from xnn.common.models import build_model\n", "from xnn.common.config import from_dict\n", "\n", "# build an xnn MACE just to grab its node embedding + one-hot machinery\n", "cfg = from_dict(\n", " {\n", " \"model\": {\n", " \"name\": \"mace\",\n", " \"cutoff\": CUTOFF,\n", " \"n_features\": 8,\n", " \"n_interactions\": 2,\n", " \"extra\": {\n", " \"species\": SPECIES,\n", " \"max_ell\": 2,\n", " \"max_L\": 1,\n", " \"num_channels\": 8,\n", " \"correlation\": 3,\n", " \"hidden_irreps\": \"8x0e+8x1o\",\n", " },\n", " }\n", " }\n", ")\n", "xnn_mace = build_model(cfg.model)\n", "\n", "node_attrs = xnn_mace.node_attr(graph.atomic_numbers) # one-hot theta (N, 3)\n", "node_attr_irreps = xnn_mace.node_attr_irreps # 3x0e\n", "feat_irreps = o3.Irreps(\"8x0e\")\n", "\n", "mace_embed = LinearNodeEmbeddingBlock(node_attr_irreps, feat_irreps)\n", "xnn_mace.node_embedding.load_state_dict(mace_embed.linear.state_dict()) # transplant W\n", "\n", "h0_xnn = xnn_mace.node_embedding(node_attrs)\n", "h0_mace = mace_embed(node_attrs)\n", "report(\"element embedding h^(0)\", (h0_xnn - h0_mace).abs().max().item())\n", "print(\"one-hot theta (first 3 atoms):\\n\", node_attrs[:3].int().numpy())" ] }, { "cell_type": "markdown", "id": "0f1270d9", "metadata": {}, "source": [ "## Block 2: Radial basis (Bessel × polynomial cutoff, then a radial MLP) · Multi-ACE eq 9 / MACE-OFF radial block\n", "\n", "MACE's learnable radial function is (Multi-ACE eq 9 / preprint eq 27)\n", "\n", "$$R^{(t)}_{k l_1 l_2 L}(r_{ij}) = \\mathrm{MLP}\\big(\\,R_n(r_{ij})\\,f_{\\rm cut}(r_{ij})\\,\\big),$$\n", "\n", "where $R_n$ are **Bessel** basis functions and $f_{\\rm cut}$ is a smooth\n", "**polynomial cutoff** (MACE-OFF \"radial embedding block\"). The Bessel functions are\n", "\n", "$$R_n(r)=\\sqrt{\\tfrac{2}{r_{\\rm cut}}}\\,\\frac{\\sin(n\\pi r/r_{\\rm cut})}{r},\\qquad n=1\\dots N,$$\n", "\n", "and the polynomial envelope of degree $p$ is the C²-smooth Klicpera form.\n", "\n", "`xnn` provides `BesselRBF`, `PolynomialCutoff` (the $R_n f_{\\rm cut}$ product is the\n", "`SphericalHarmonicEdgeEmbedding.edge_radial`); the per-path MLP is the\n", "`FullyConnectedNet` `conv_tp_weights` inside each interaction. We compare the\n", "basis and the cutoff directly against `mace.modules.radial`." ] }, { "cell_type": "code", "execution_count": 4, "id": "0c48669e", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:18:26.766205Z", "iopub.status.busy": "2026-07-20T04:18:26.765973Z", "iopub.status.idle": "2026-07-20T04:18:26.773691Z", "shell.execute_reply": "2026-07-20T04:18:26.773058Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "OK Bessel basis R_n(r) max|xnn - mace| = 5.97e-16\n", "OK polynomial cutoff f_cut(r) max|xnn - mace| = 0.00e+00\n", "OK radial embedding R_n * f_cut max|xnn - mace| = 4.72e-16\n" ] } ], "source": [ "from xnn.gnn.featurizers.radial import BesselRBF\n", "from xnn.gnn.featurizers.cutoff import PolynomialCutoff as XPoly\n", "from mace.modules.radial import BesselBasis, PolynomialCutoff as MPoly\n", "from mace.modules.blocks import RadialEmbeddingBlock\n", "\n", "r = torch.linspace(0.2, 4.9, 60)\n", "NRBF, P = 8, 5\n", "\n", "xb = BesselRBF(NRBF, CUTOFF)(r)\n", "mb = BesselBasis(CUTOFF, NRBF, trainable=False)(r.unsqueeze(-1))\n", "report(\"Bessel basis R_n(r)\", (xb - mb).abs().max().item())\n", "\n", "xc = XPoly(CUTOFF, p=P)(r)\n", "mc = MPoly(CUTOFF, p=P)(r)\n", "report(\"polynomial cutoff f_cut(r)\", (xc - mc).abs().max().item())\n", "\n", "# full radial embedding R_n(r) * f_cut(r)\n", "xr = BesselRBF(NRBF, CUTOFF)(r) * XPoly(CUTOFF, p=P)(r)[:, None]\n", "mre = RadialEmbeddingBlock(CUTOFF, NRBF, P, radial_type=\"bessel\")\n", "mr, _ = mre(r.unsqueeze(-1), None, None, None)\n", "report(\"radial embedding R_n * f_cut\", (xr - mr).abs().max().item())" ] }, { "cell_type": "markdown", "id": "1c1ffc7c", "metadata": {}, "source": [ "## Block 3: Spherical harmonics of the edge direction · $Y^{m}_{l}(\\hat r_{ij})$ in eq 2\n", "\n", "The angular part of the one-particle basis is the real spherical harmonics of the\n", "unit edge vector, up to degree $\\ell_{\\max}$ (`max_ell`). Both codes call the\n", "**same** `e3nn` routine with the same normalisation\n", "(`normalize=True, normalization=\"component\"`), so they agree to machine precision.\n", "In `xnn` this lives in `SphericalHarmonicEdgeEmbedding` (`edge_sh`)." ] }, { "cell_type": "code", "execution_count": 5, "id": "b083673e", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:18:26.775285Z", "iopub.status.busy": "2026-07-20T04:18:26.775160Z", "iopub.status.idle": "2026-07-20T04:18:26.979599Z", "shell.execute_reply": "2026-07-20T04:18:26.978325Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "OK spherical harmonics Y_l^m(r_hat) max|xnn - mace| = 0.00e+00\n", "edge_sh irreps: 1x0e+1x1o+1x2e shape: (42, 9)\n" ] } ], "source": [ "vec = graph.edge_vectors()\n", "ir_sh = o3.Irreps.spherical_harmonics(2) # 1x0e+1x1o+1x2e (max_ell=2)\n", "\n", "edge = xnn_mace.edge_feat(graph) # xnn edge featurizer\n", "xsh = edge[\"edge_sh\"]\n", "msh = o3.spherical_harmonics(ir_sh, vec, normalize=True, normalization=\"component\")\n", "report(\"spherical harmonics Y_l^m(r_hat)\", (xsh - msh).abs().max().item())\n", "print(\"edge_sh irreps:\", ir_sh, \" shape:\", tuple(xsh.shape))" ] }, { "cell_type": "markdown", "id": "2ef0d9d5", "metadata": {}, "source": [ "## Blocks 4 & 5: One-particle basis $\\phi$ and atomic basis $A$ · MACE-OFF eqs 2–3 / Multi-ACE eqs 14, 19\n", "\n", "The one-particle (edge) basis couples the neighbour feature, the radial function\n", "and the spherical harmonics through Clebsch–Gordan coefficients (MACE-OFF eq 2):\n", "\n", "$$\\phi^{(t)}_{ij,k l_3 m_3}= \\sum_{l_1 m_1, l_2 m_2}\n", " C^{l_3 m_3}_{l_1 m_1 l_2 m_2}\\,R^{(t)}_{k l_1 l_2 l_3}(r_{ij})\\,Y^{m_1}_{l_1}(\\hat r_{ij})\\,h^{(t)}_{j,k l_2 m_2},$$\n", "\n", "and the permutation-invariant **atomic basis** sums it over neighbours\n", "(MACE-OFF eq 3 / Multi-ACE eq 14):\n", "\n", "$$A^{(t)}_{i,k l_3 m_3}= \\sum_{k'} W_{k k'}\\sum_{j\\in\\mathcal N(i)}\\phi^{(t)}_{ij,k' l_3 m_3}.$$\n", "\n", "In `xnn` this is exactly the body of `RealAgnostic(Residual)InteractionBlock`:\n", "the weighted `o3.TensorProduct` (`conv_tp`, with per-edge weights from the radial\n", "MLP) followed by a scatter-sum over `edge_index` and a channel-mixing `o3.Linear`.\n", "We build the original and `xnn` interaction blocks with identical irreps, copy\n", "the state dict across, and confirm the message $A$ (and the residual\n", "self-connection $sc$) match **bit-for-bit**." ] }, { "cell_type": "code", "execution_count": 6, "id": "bd526de3", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:18:26.982242Z", "iopub.status.busy": "2026-07-20T04:18:26.981354Z", "iopub.status.idle": "2026-07-20T04:18:28.248227Z", "shell.execute_reply": "2026-07-20T04:18:28.246249Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "OK atomic basis / message A max|xnn - mace| = 0.00e+00\n", "OK residual self-connection sc max|xnn - mace| = 0.00e+00\n" ] } ], "source": [ "from mace.modules.blocks import RealAgnosticResidualInteractionBlock as MIB\n", "from xnn.gnn.models.mace import RealAgnosticResidualInteractionBlock as XIB\n", "\n", "na = o3.Irreps(\"3x0e\") # node attrs (one-hot, 3 elements)\n", "nf = o3.Irreps(\"8x0e\") # node feats (k=8 scalar channels)\n", "ea = o3.Irreps.spherical_harmonics(2) # edge attrs (SH)\n", "ef = o3.Irreps(\"8x0e\") # edge radial feats (n_bessel=8)\n", "ti = o3.Irreps(\"8x0e+8x1o+8x2e\") # interaction (target) irreps\n", "hi = o3.Irreps(\"8x0e+8x1o\") # hidden irreps\n", "kw = dict(node_attrs_irreps=na, node_feats_irreps=nf, edge_attrs_irreps=ea,\n", " edge_feats_irreps=ef, target_irreps=ti, hidden_irreps=hi,\n", " avg_num_neighbors=8.0, radial_MLP=[16, 16])\n", "\n", "m_int, x_int = MIB(**kw), XIB(**kw)\n", "x_int.load_state_dict(m_int.state_dict()) # transplant ALL interaction weights\n", "\n", "N, E = graph.num_nodes, graph.num_edges\n", "nfeat = torch.randn(N, nf.dim)\n", "edge_attrs = torch.randn(E, ea.dim); edge_feats = torch.randn(E, ef.dim)\n", "ei = graph.edge_index\n", "\n", "xA, xsc = x_int(node_attrs, nfeat, edge_attrs, edge_feats, ei)\n", "mA, msc = m_int(node_attrs, nfeat, edge_attrs, edge_feats, ei, cutoff=None)\n", "report(\"atomic basis / message A\", (xA - mA).abs().max().item())\n", "report(\"residual self-connection sc\", (xsc - msc).abs().max().item())" ] }, { "cell_type": "markdown", "id": "1f36fef0", "metadata": {}, "source": [ "## Block 6: Product basis and the generalized Clebsch–Gordan $U$ tensors · MACE-OFF eq 4 / Multi-ACE eqs 15–16, 20, 42\n", "\n", "Higher body order comes from the **product basis**: the $\\nu$-fold product of the\n", "atomic basis with itself (MACE-OFF/Multi-ACE):\n", "\n", "$$\\mathbf A^{(t)}_{i,k\\mathbf v}=\\prod_{\\xi=1}^{\\nu} A^{(t)}_{i,k v_\\xi},\\qquad\n", " \\mathbf v=(v_1,\\dots,v_\\nu),$$\n", "\n", "symmetrised into the equivariant $B$ basis through the **generalized\n", "Clebsch–Gordan coefficients** $\\mathcal C^{LM}_{\\eta,\\mathbf v}$ (Multi-ACE eq 20),\n", "which are themselves products of ordinary CG coefficients (Multi-ACE eq 42):\n", "\n", "$$\\mathbf B^{(t)}_{i,k\\eta,LM}=\\sum_{\\mathbf v}\\mathcal C^{LM}_{\\eta,\\mathbf v}\\,\\mathbf A^{(t)}_{i,k\\mathbf v}.$$\n", "\n", "The coupling coefficients are precomputed as the $U$ tensors. `xnn`'s\n", "`U_matrix_real` is adapted from `mace-torch`; here we confirm it is **bit-identical**\n", "to the original `mace.modules.symmetric_contraction.U_matrix_real` for correlation\n", "orders $\\nu=1,2,3$ (body orders 2, 3, 4; MACE uses $\\nu=3$)." ] }, { "cell_type": "code", "execution_count": 7, "id": "ba308fdc", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:18:28.250764Z", "iopub.status.busy": "2026-07-20T04:18:28.250100Z", "iopub.status.idle": "2026-07-20T04:19:44.841737Z", "shell.execute_reply": "2026-07-20T04:19:44.836303Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "OK U tensor (generalized CG), nu=1 shape=(16, 1) max|xnn - mace| = 0.00e+00\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "OK U tensor (generalized CG), nu=2 shape=(16, 16, 4) max|xnn - mace| = 0.00e+00\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "OK U tensor (generalized CG), nu=3 shape=(16, 16, 16, 23) max|xnn - mace| = 0.00e+00\n" ] } ], "source": [ "from xnn.gnn.models.mace import U_matrix_real as xU\n", "from mace.modules.symmetric_contraction import U_matrix_real as mU\n", "\n", "coupling = o3.Irreps(\"1x0e+1x1o+1x2e+1x3o\")\n", "for nu in (1, 2, 3):\n", " u_x = xU(coupling, \"0e\", nu, dtype=torch.float64)[-1]\n", " u_m = mU(coupling, o3.Irreps(\"0e\"), nu, dtype=torch.float64, use_cueq_cg=False)[-1]\n", " report(f\"U tensor (generalized CG), nu={nu} shape={tuple(u_x.shape)}\",\n", " (u_x - u_m).abs().max().item(), tol=0.0)" ] }, { "cell_type": "markdown", "id": "3fc41752", "metadata": {}, "source": [ "## Block 7: Symmetric contraction → message $m$ · MACE-OFF eqs 4–5 / Multi-ACE eq 21\n", "\n", "The learned, per-element contraction over the $U$ basis turns the atomic basis $A$\n", "into the message (MACE-OFF eq 5 / Multi-ACE eq 21):\n", "\n", "$$m^{(t)}_{i,kLM}=\\sum_{\\nu}\\sum_{\\eta_\\nu} W^{(t)}_{z_i,k\\eta_\\nu L}\\,\\mathbf B^{(t)}_{i,k\\eta_\\nu LM}.$$\n", "\n", "This is the genuinely MACE-specific operation. `xnn`'s `SymmetricContraction`\n", "mirrors the original `mace.modules.symmetric_contraction.SymmetricContraction`.\n", "The two store the per-correlation weights in a different order\n", "(`mace` keeps the top order as `weights_max`, the rest descending; `xnn` keeps an\n", "ascending `weights[ν-1]` list), so we map the weights and then check the contracted\n", "output matches to ~$10^{-16}$." ] }, { "cell_type": "code", "execution_count": 8, "id": "637dc648", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:19:44.844167Z", "iopub.status.busy": "2026-07-20T04:19:44.844000Z", "iopub.status.idle": "2026-07-20T04:22:17.832118Z", "shell.execute_reply": "2026-07-20T04:22:17.831371Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "OK symmetric contraction m max|xnn - mace| = 8.88e-16\n" ] } ], "source": [ "from xnn.gnn.models.mace import SymmetricContraction as XSC\n", "from mace.modules.symmetric_contraction import SymmetricContraction as MSC\n", "\n", "irreps_in = o3.Irreps(\"8x0e+8x1o+8x2e+8x3o\")\n", "irreps_out = o3.Irreps(\"8x0e+8x1o\")\n", "CORR, NEL = 3, 3\n", "xsc = XSC(irreps_in, irreps_out, correlation=CORR, num_elements=NEL)\n", "msc = MSC(irreps_in, irreps_out, correlation=CORR, num_elements=NEL,\n", " irrep_normalization=\"component\", path_normalization=\"element\",\n", " use_reduced_cg=False)\n", "\n", "def transplant_sc(xsc, msc, corr):\n", " with torch.no_grad():\n", " for c in range(len(xsc.contractions)):\n", " xc, mc = xsc.contractions[c], msc.contractions[c]\n", " xc.weights[corr - 1].copy_(mc.weights_max) # top order\n", " for nu in range(1, corr):\n", " xc.weights[nu - 1].copy_(mc.weights[corr - 1 - nu]) # lower orders\n", "transplant_sc(xsc, msc, CORR)\n", "\n", "B = 5\n", "nfeat_dim = irreps_in.count((0, 1))\n", "x = torch.randn(B, nfeat_dim, o3.Irreps([ir.ir for ir in irreps_in]).dim)\n", "y = torch.zeros(B, NEL); y[torch.arange(B), torch.randint(0, NEL, (B,))] = 1.0 # element one-hot\n", "report(\"symmetric contraction m\", (xsc(x, y) - msc(x, y)).abs().max().item())" ] }, { "cell_type": "markdown", "id": "3d68ac68", "metadata": {}, "source": [ "## Block 8: Equivariant product block + update with self-connection · MACE-OFF eq 6 / Multi-ACE eqs 22, 30–31\n", "\n", "The full higher-order block raises the body order via the symmetric contraction,\n", "mixes channels with a linear map, and adds the element-dependent **self-connection**\n", "(residual update, MACE-OFF eq 6 / preprint eqs 30–31):\n", "\n", "$$h^{(t+1)}_{i,kLM}=\\sum_{\\tilde k} W^{(t)}_{k\\tilde kL}\\,m^{(t)}_{i,\\tilde kLM}\n", " +\\sum_{\\tilde k} W^{(t)}_{z_i,k\\tilde kL}\\,h^{(t)}_{i,\\tilde kLM}.$$\n", "\n", "`xnn`'s `_EquivariantProductBasis` matches the original\n", "`EquivariantProductBasisBlock`. We transplant the contraction weights (same mapping\n", "as Block 7) and the output `o3.Linear`, then confirm the block output matches." ] }, { "cell_type": "code", "execution_count": 9, "id": "08deff03", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:22:17.834604Z", "iopub.status.busy": "2026-07-20T04:22:17.833919Z", "iopub.status.idle": "2026-07-20T04:23:46.300975Z", "shell.execute_reply": "2026-07-20T04:23:46.298130Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "OK equivariant product + update max|xnn - mace| = 4.44e-16\n" ] } ], "source": [ "from mace.modules.blocks import EquivariantProductBasisBlock as MPB\n", "from xnn.gnn.models.mace import _EquivariantProductBasis as XPB\n", "\n", "mpb = MPB(node_feats_irreps=irreps_in, target_irreps=irreps_out,\n", " correlation=CORR, num_elements=NEL, use_sc=True)\n", "xpb = XPB(irreps_in, irreps_out, correlation=CORR, num_elements=NEL, use_sc=True)\n", "transplant_sc(xpb.symmetric_contractions, mpb.symmetric_contractions, CORR)\n", "xpb.linear.load_state_dict(mpb.linear.state_dict())\n", "\n", "nfeats = torch.randn(B, nfeat_dim, o3.Irreps([ir.ir for ir in irreps_in]).dim)\n", "sc = torch.randn(B, irreps_out.dim)\n", "report(\"equivariant product + update\", (xpb(nfeats, sc, y) - mpb(nfeats, sc, y)).abs().max().item())" ] }, { "cell_type": "markdown", "id": "940e9a64", "metadata": {}, "source": [ "## Block 9: Readouts · MACE-OFF eqs 7–8\n", "\n", "The site energy is a sum of read-outs of the per-layer features; the read-out is\n", "**linear** for the early layers and a small **MLP** (gated) for the last layer\n", "(MACE-OFF eq 8):\n", "\n", "$$\\mathcal R^{(t)}(h^{(t)}_i)=\\begin{cases}\\sum_k W^{(t)}_k\\,h^{(t)}_{i,k00}&t xnn ----\n", "with torch.no_grad():\n", " xfull.node_embedding.load_state_dict(mmodel.node_embedding.linear.state_dict())\n", " for i in range(T):\n", " xfull.interactions[i].load_state_dict(mmodel.interactions[i].state_dict())\n", " transplant_sc(xfull.products[i].symmetric_contractions,\n", " mmodel.products[i].symmetric_contractions, 3)\n", " xfull.products[i].linear.load_state_dict(mmodel.products[i].linear.state_dict())\n", " xr, mr = xfull.readouts[i], mmodel.readouts[i]\n", " if \"NonLinear\" in type(mr).__name__:\n", " xr.linear_1.load_state_dict(mr.linear_1.state_dict())\n", " xr.linear_2.load_state_dict(mr.linear_2.state_dict())\n", " else:\n", " xr.linear.load_state_dict(mr.linear.state_dict())\n", " for z, e in zip(SPECIES, E0):\n", " xfull.atom_ref.weight[z] = float(e)\n", "\n", "# ---- xnn prediction ----\n", "gx = structure_to_graph({\"pos\": pos, \"atomic_numbers\": Z}, CUTOFF)\n", "ox = xmodel(gx); E_x = float(ox[\"energy\"]); F_x = ox[\"forces\"].detach().numpy()\n", "\n", "# ---- original mace prediction (its own data pipeline) ----\n", "zt = AtomicNumberTable(SPECIES)\n", "conf = Configuration(atomic_numbers=Z, positions=pos, properties={}, property_weights={})\n", "ad = AtomicData.from_config(conf, z_table=zt, cutoff=CUTOFF)\n", "batch = next(iter(torch_geometric.dataloader.DataLoader([ad], batch_size=1)))\n", "om = mmodel(batch.to_dict(), compute_force=True)\n", "E_m = float(om[\"energy\"]); F_m = om[\"forces\"].detach().numpy()\n", "\n", "print(f\"total energy xnn = {E_x:.10f} eV\")\n", "print(f\"total energy mace = {E_m:.10f} eV\")\n", "report(\"FULL MODEL total energy\", abs(E_x - E_m))\n", "report(\"FULL MODEL per-atom forces\", np.abs(F_x - F_m).max())" ] }, { "cell_type": "markdown", "id": "f017b109", "metadata": {}, "source": [ "## Summary\n", "\n", "Every block needed to reproduce the original MACE was checked against\n", "`mace-torch` on the toy system:\n", "\n", "| Block | MACE eq | Weights? | agreement |\n", "|---|---|---|---|\n", "| 1. Element embedding | MACE-OFF 1 / Multi-ACE 13 | transplanted | machine precision |\n", "| 2. Bessel basis · cutoff · radial MLP | Multi-ACE 9 | none | machine precision |\n", "| 3. Spherical harmonics | eq 2 ($Y$) | none (same e3nn call) | machine precision |\n", "| 4–5. One-particle $\\phi$ + atomic basis $A$ | MACE-OFF 2–3 / Multi-ACE 14,19 | transplanted | **bit-identical** |\n", "| 6. Generalized CG $U$ tensors | Multi-ACE 16,20,42 | none | **bit-identical (0)** |\n", "| 7. Symmetric contraction → message | MACE-OFF 5 / Multi-ACE 21 | transplanted | ~1e-16 |\n", "| 8. Product block + self-connection | MACE-OFF 6 / Multi-ACE 22,30–31 | transplanted | ~1e-16 |\n", "| 9. Linear / gated readouts | MACE-OFF 7–8 | transplanted | machine precision |\n", "| 10. $E_0$ + autograd forces | MACE-OFF 7, §2.1 | n/a | exact ($T=0$) |\n", "| **Full model** | eqs 1–8 | **all transplanted** | **energy 0, forces ~1e-16** |\n", "\n", "`xnn.gnn.models.mace` is therefore a **faithful reproduction of the original MACE\n", "architecture**, depending only on `e3nn` (no `mace-torch` / `cuequivariance`).\n", "The companion notebook `mace_argon_train_test.ipynb` trains and tests this MACE\n", "on a realistic Argon MD dataset." ] } ], "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 }