{ "cells": [ { "cell_type": "markdown", "id": "50d1544e", "metadata": {}, "source": [ "# CACE, block by block: reproducing the original implementation with `xnn`\n", "\n", "[CACE (Cheng, *npj Comput Mater* **10**, 157, 2024)](https://doi.org/10.1038/s41524-024-01332-4)\n", "builds body-ordered invariant features of atomic environments **entirely in\n", "Cartesian coordinates**: no spherical harmonics, no Clebsch–Gordan\n", "contraction, no e3nn. This notebook walks through every block of the\n", "architecture and checks each against the **original**\n", "[`cace`](https://github.com/BingqingCheng/cace) package on the same inputs,\n", "ending with a whole-model weight transplant and an energy/force parity check\n", "at machine precision, the same protocol as the MACE / NequIP / Allegro\n", "companions (`examples/gnn/{mace,nequip,allegro}/01_*`)." ] }, { "cell_type": "markdown", "id": "4dce7b8e", "metadata": {}, "source": [ "## 0. Setup: `float64` for exact comparison\n", "\n", "A toy periodic H/O box; both pipelines are fed the **same edge list** so every\n", "difference we see is a difference between the *implementations*, not between\n", "neighbour-list codes." ] }, { "cell_type": "code", "execution_count": 1, "id": "080b68fc", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:24:20.514079Z", "iopub.status.busy": "2026-07-20T04:24:20.513966Z", "iopub.status.idle": "2026-07-20T04:24:22.681779Z", "shell.execute_reply": "2026-07-20T04:24:22.680929Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "xnn: 0.1.0 | cace (original): 0.1.0\n", "toy box: 10 atoms, 282 directed edges (shared by both pipelines)\n" ] } ], "source": [ "# silence the expected warnings\n", "import logging, warnings\n", "logging.disable(logging.WARNING)\n", "warnings.filterwarnings(\"ignore\", category=UserWarning)\n", "warnings.filterwarnings(\"ignore\", category=FutureWarning,\n", " message=\"You are using `torch.load` with `weights_only=False`\")\n", "\n", "import numpy as np\n", "import torch\n", "\n", "torch.set_default_dtype(torch.float64)\n", "torch.manual_seed(0)\n", "\n", "import xnn, cace\n", "from xnn.common.data import structure_to_graph\n", "print(\"xnn:\", xnn.__version__, \"| cace (original):\", cace.__version__ if hasattr(cace, \"__version__\") else \"0.1.0\")\n", "\n", "CUT, SPECIES = 4.5, [1, 8]\n", "NAB, NRBF, NRB, LMAX, NU, T, AVG = 2, 6, 8, 3, 3, 1, 9.0 # paper notation: N_embedding, ñ, n, l_max, nu_max, T\n", "\n", "rng = np.random.default_rng(11)\n", "pos = rng.uniform(0, 4, (10, 3))\n", "Z = [1, 8] * 5\n", "g = structure_to_graph({\"pos\": pos, \"atomic_numbers\": Z,\n", " \"cell\": np.eye(3) * 5.0, \"pbc\": [True] * 3}, CUT)\n", "\n", "# the upstream forward consumes a plain dict; reuse the xnn edge list & shifts\n", "cell = g.cell[0]\n", "data = {\"positions\": g.pos.clone().requires_grad_(True),\n", " \"atomic_numbers\": g.atomic_numbers,\n", " \"edge_index\": g.edge_index,\n", " \"shifts\": g.cell_shifts.to(torch.get_default_dtype()) @ cell,\n", " \"batch\": g.batch, \"cell\": cell.unsqueeze(0)}\n", "print(f\"toy box: {g.num_nodes} atoms, {g.num_edges} directed edges (shared by both pipelines)\")" ] }, { "cell_type": "markdown", "id": "c64ec29d", "metadata": {}, "source": [ "## The CACE architecture in one picture\n", "\n", "Per structure (paper eqs 1–15):\n", "\n", "1. **Element embedding**: each element gets a learnable vector $\\theta$ of length\n", " $N_{\\rm embedding}$ (1–4); an edge type is the flattened tensor product\n", " $T = \\theta_i \\otimes \\theta_j$, giving $c=N_{\\rm embedding}^2$ channels *(eq 1)*.\n", "2. **Edge basis** $\\chi_{cn\\mathbf{l}} = T_c\\,R_{n}(r_{ji})\\,L_\\mathbf{l}(\\hat r_{ji})$ with\n", " a (trainable) Bessel radial basis × polynomial cutoff and the **Cartesian angular\n", " monomials** $L_\\mathbf{l}(\\hat r) = x^{l_x} y^{l_y} z^{l_z}$ *(eq 2)*.\n", "3. **A basis**: sum over edges of a node (the \"density trick\", eq 6), then the raw radial\n", " channels are mixed per $(l, c)$ by a learned matrix $W_{\\tilde n n, cl}$ *(eq 5)*.\n", "4. **B basis**: products of A entries whose angular indices pair with *shared factors*\n", " are contracted with multinomial prefactors into polynomially independent rotational\n", " invariants of body order $\\nu$ *(eqs 7–10, fig 1i)*.\n", "5. **Message passing** *(eqs 11–14)*: $m_1 = F(r_{ji}) A_j$ (exponential-decay filter, \"Ar\"),\n", " $m_2 = H(B_j)\\chi$ (recursive edge embedding, \"Bchi\"), plus a node-memory term (\"M\");\n", " normalized by $1/\\sqrt{\\langle\\text{neighbors}\\rangle}$.\n", "6. **Readout**: concatenated B features of all layers → linear + MLP → atomic energies\n", " *(eq 15)*; forces are exact gradients via autograd." ] }, { "cell_type": "code", "execution_count": 2, "id": "9b7e2d1e", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:24:22.683398Z", "iopub.status.busy": "2026-07-20T04:24:22.683309Z", "iopub.status.idle": "2026-07-20T04:24:22.750330Z", "shell.execute_reply": "2026-07-20T04:24:22.749722Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "parameters: original 14609 | xnn 14809 (xnn adds the 200-entry atom_ref E0 table: 200 = 200)\n" ] } ], "source": [ "# --- build both models once; every block below compares their internals ---\n", "from cace.modules import BesselRBF as UpBessel, PolynomialCutoff as UpPoly\n", "from cace.modules.atomwise import Atomwise\n", "from cace.representations import Cace as UpCace\n", "\n", "from xnn.common.config import from_dict\n", "from xnn.common.models import build_model, ForceStressOutput\n", "\n", "torch.manual_seed(7)\n", "up_rep = UpCace(zs=SPECIES, n_atom_basis=NAB, cutoff=CUT,\n", " radial_basis=UpBessel(cutoff=CUT, n_rbf=NRBF, trainable=True),\n", " cutoff_fn=UpPoly(cutoff=CUT, p=6),\n", " max_l=LMAX, max_nu=NU, num_message_passing=T,\n", " type_message_passing=[\"M\", \"Ar\", \"Bchi\"], n_radial_basis=NRB,\n", " avg_num_neighbors=AVG, embed_receiver_nodes=True)\n", "up_readout = Atomwise(n_layers=3, n_hidden=[32, 16], output_key=\"energy\",\n", " add_linear_nn=True)\n", "up_out = up_readout(up_rep(data)) # one call lazy-initializes Bchi's H and the readout MLP\n", "\n", "cfg = from_dict({\"model\": {\"name\": \"cace\", \"cutoff\": CUT, \"n_interactions\": T, \"n_rbf\": NRBF,\n", " \"extra\": {\"species\": SPECIES, \"n_atom_basis\": NAB,\n", " \"n_radial_basis\": NRB, \"max_l\": LMAX, \"max_nu\": NU,\n", " \"avg_num_neighbors\": AVG,\n", " \"embed_receiver_nodes\": True}}})\n", "x = build_model(cfg.model)\n", "\n", "def transplant(x, rep, readout):\n", " '''Copy every learnable tensor from the original cace into the xnn CACE.'''\n", " with torch.no_grad():\n", " x.embed_sender.copy_(rep.node_embedding_sender.embedding_weights)\n", " x.embed_receiver.copy_(rep.node_embedding_receiver.embedding_weights)\n", " x.rbf.freqs.copy_(rep.radial_basis.bessel_weights * float(rep.cutoff))\n", " x.radial_transform.weight.copy_(torch.stack(list(rep.radial_transform.weights)))\n", " for t, (nm, ar, bchi) in enumerate(rep.message_passing_list):\n", " xi = x.interactions[t]\n", " if nm is not None:\n", " xi.memory.memory_coef.copy_(torch.stack(list(nm.memory_coef)))\n", " if ar is not None:\n", " xi.message_ar.prefactor.copy_(torch.stack(list(ar.prefactor)))\n", " xi.message_ar.inv_r0.copy_(torch.stack(list(ar.invr0)))\n", " if bchi is not None:\n", " xi.message_bchi.h.weight.copy_(bchi.hnet[0].linear.weight)\n", " xi.message_bchi.h.bias.copy_(bchi.hnet[0].linear.bias)\n", " for j, dense in enumerate(readout.outnet):\n", " x.readout_mlp[2 * j].weight.copy_(dense.linear.weight)\n", " x.readout_mlp[2 * j].bias.copy_(dense.linear.bias)\n", " x.readout_linear.weight.copy_(readout.linear_nn.linear.weight)\n", " x.readout_linear.bias.copy_(readout.linear_nn.linear.bias)\n", "\n", "transplant(x, up_rep, up_readout)\n", "p_up = sum(p.numel() for p in up_rep.parameters()) + sum(p.numel() for p in up_readout.parameters())\n", "p_x = sum(p.numel() for p in x.parameters())\n", "print(f\"parameters: original {p_up} | xnn {p_x} \"\n", " f\"(xnn adds the 200-entry atom_ref E0 table: {p_x - p_up} = {x.atom_ref.weight.numel()})\")" ] }, { "cell_type": "markdown", "id": "6f89e3be", "metadata": {}, "source": [ "## Block 1: Element embedding & tensor-product edge type · eq 1\n", "\n", "Each atom's one-hot element vector is embedded to $\\theta\\in\\mathbb{R}^{N_{\\rm emb}}$\n", "(one table for senders, optionally another for receivers). The edge type is the\n", "flattened outer product $\\theta_i \\otimes \\theta_j$ (interpretable as an attention\n", "key/query pair), giving $c = N_{\\rm emb}^2$ channels that *smoothly encode chemistry*\n", "and let CACE learn across elements (alchemical learning)." ] }, { "cell_type": "code", "execution_count": 3, "id": "29a3aeee", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:24:22.751623Z", "iopub.status.busy": "2026-07-20T04:24:22.751549Z", "iopub.status.idle": "2026-07-20T04:24:22.755603Z", "shell.execute_reply": "2026-07-20T04:24:22.755149Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "edge-type channels c = 4 (= 2^2)\n", "max |xnn - original| = 0.0\n" ] } ], "source": [ "node_onehot = up_rep.node_onehot(data[\"atomic_numbers\"])\n", "emb_s_up = up_rep.node_embedding_sender(node_onehot)\n", "emb_r_up = up_rep.node_embedding_receiver(node_onehot)\n", "edge_type_up = up_rep.edge_coding(edge_index=data[\"edge_index\"],\n", " node_type=emb_s_up, node_type_2=emb_r_up, data=data)\n", "\n", "one_hot = x.node_attr(g.atomic_numbers)\n", "theta_s = (one_hot @ x.embed_sender)[g.edge_index[0]]\n", "theta_r = (one_hot @ x.embed_receiver)[g.edge_index[1]]\n", "edge_type_x = (theta_s.unsqueeze(2) * theta_r.unsqueeze(1)).flatten(1)\n", "\n", "print(\"edge-type channels c =\", edge_type_x.shape[1], f\"(= {NAB}^2)\")\n", "print(\"max |xnn - original| =\", (edge_type_x - edge_type_up).abs().max().item())" ] }, { "cell_type": "markdown", "id": "43db6ba1", "metadata": {}, "source": [ "## Block 2: Radial basis (trainable Bessel × polynomial cutoff)\n", "\n", "The raw radial basis is $\\tilde R_{\\tilde n}(r) = \\sqrt{2/r_c}\\,\\sin(\\tilde n\\pi r/r_c)/r$\n", "with trainable frequencies, enveloped by the DimeNet degree-6 polynomial cutoff;\n", "CACE inherits both from MACE, so xnn reuses its shared\n", "`xnn.gnn.featurizers.BesselRBF` / `PolynomialCutoff` (the only mapping is that xnn\n", "stores the dimensionless frequencies $\\tilde n\\pi$, upstream stores $\\tilde n\\pi/r_c$)." ] }, { "cell_type": "code", "execution_count": 4, "id": "c0ac5db2", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:24:22.756758Z", "iopub.status.busy": "2026-07-20T04:24:22.756693Z", "iopub.status.idle": "2026-07-20T04:24:22.761298Z", "shell.execute_reply": "2026-07-20T04:24:22.760775Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "max |edge vectors diff| = 0.0\n", "max |radial basis diff| = 4.440892098500626e-16\n" ] } ], "source": [ "from cace.modules import get_edge_vectors_and_lengths\n", "\n", "vec_up, len_up = get_edge_vectors_and_lengths(positions=data[\"positions\"],\n", " edge_index=data[\"edge_index\"],\n", " shifts=data[\"shifts\"], normalize=True)\n", "radial_up = up_rep.radial_basis(len_up) * up_rep.cutoff_fn(len_up)\n", "\n", "edge_vec = g.edge_vectors()\n", "lengths = edge_vec.norm(dim=-1)\n", "unit_vec = edge_vec / (lengths + 1e-9).unsqueeze(-1)\n", "radial_x = x.rbf(lengths) * x.envelope(lengths).unsqueeze(-1)\n", "\n", "print(\"max |edge vectors diff| =\", (vec_up - unit_vec).abs().max().item())\n", "print(\"max |radial basis diff| =\", (radial_x - radial_up).abs().max().item())" ] }, { "cell_type": "markdown", "id": "d9f99ae1", "metadata": {}, "source": [ "## Block 3: Cartesian angular basis · eq 2\n", "\n", "Instead of spherical harmonics $Y_l^m$, CACE uses the Cartesian monomials\n", "\n", "$$L_\\mathbf{l}(\\hat r) = x^{l_x}\\,y^{l_y}\\,z^{l_z},\\qquad l_x+l_y+l_z = l \\le l_{\\max},$$\n", "\n", "which span exactly the same space per total angular momentum $l$ (a fixed linear map\n", "connects the two bases). There are $(l_{\\max}+1)(l_{\\max}+2)(l_{\\max}+3)/6$ of them.\n", "`xnn.gnn.featurizers.CartesianAngularBasis` evaluates them with the same\n", "multiply-recursion as upstream (autograd-safe for axis-aligned edges); the two codes\n", "order the entries within each $l$ block differently, so we compare through the\n", "$(l_x,l_y,l_z)$ index map; every value is identical." ] }, { "cell_type": "code", "execution_count": 5, "id": "98a34ca4", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:24:22.762503Z", "iopub.status.busy": "2026-07-20T04:24:22.762425Z", "iopub.status.idle": "2026-07-20T04:24:22.769143Z", "shell.execute_reply": "2026-07-20T04:24:22.768669Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "angular entries: 20 monomials up to l_max = 3\n", "first few (lx,ly,lz) upstream: [(0, 0, 0), (1, 0, 0), (0, 1, 0), (0, 0, 1), (2, 0, 0)] xnn: [(0, 0, 0), (0, 0, 1), (0, 1, 0), (1, 0, 0), (0, 0, 2)]\n", "max |diff (via index map)| = 0.0\n" ] } ], "source": [ "angular_up = up_rep.angular_basis(vec_up)\n", "angular_x = x.angular(unit_vec)\n", "\n", "up_list = [tuple(c) for c in up_rep.angular_basis.get_lxlylz_list()]\n", "x_list = x.angular.lxlylz\n", "ang_map = torch.tensor([x_list.index(c) for c in up_list]) # upstream entry i = xnn entry ang_map[i]\n", "\n", "print(\"angular entries:\", angular_x.shape[1], \"monomials up to l_max =\", LMAX)\n", "print(\"first few (lx,ly,lz) upstream:\", up_list[:5], \" xnn:\", x_list[:5])\n", "print(\"max |diff (via index map)| =\", (angular_x[:, ang_map] - angular_up).abs().max().item())" ] }, { "cell_type": "markdown", "id": "65101985", "metadata": {}, "source": [ "## Block 4: Atom-centered A basis + trainable radial coupling · eqs 5–6\n", "\n", "The edge basis $\\chi = (R\\cdot f_{\\rm cut}) \\otimes L \\otimes T$ is summed over the\n", "edges of each node (the \"density trick\") and the raw radial channels are then mixed,\n", "**per total $l$ and channel $c$**, by a learned $\\tilde n \\times n$ matrix, CACE's\n", "\"optimized radial channel coupling\". xnn stacks upstream's per-$l$ weight list into\n", "one tensor and applies it as a single einsum; the mixing weights are shared by all\n", "angular entries of the same $l$, so the within-$l$ ordering still cancels through\n", "the index map." ] }, { "cell_type": "code", "execution_count": 6, "id": "1b87f70b", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:24:22.770440Z", "iopub.status.busy": "2026-07-20T04:24:22.770373Z", "iopub.status.idle": "2026-07-20T04:24:22.814043Z", "shell.execute_reply": "2026-07-20T04:24:22.813621Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "A basis shape (N, n, angular, c): (10, 8, 20, 4)\n", "max |A diff (via index map)| = 1.7763568394002505e-15\n" ] } ], "source": [ "from cace.tools import elementwise_multiply_3tensors, scatter_sum as up_scatter\n", "from xnn.common.models.ops import scatter_sum\n", "\n", "edge_attr_up = elementwise_multiply_3tensors(radial_up, angular_up, edge_type_up)\n", "A_up = up_scatter(src=edge_attr_up, index=data[\"edge_index\"][1], dim=0, dim_size=10)\n", "A_up = up_rep.radial_transform(A_up)\n", "\n", "edge_attr_x = (radial_x.unsqueeze(2).unsqueeze(3)\n", " * angular_x.unsqueeze(1).unsqueeze(3)\n", " * edge_type_x.unsqueeze(1).unsqueeze(2))\n", "A_x = x.radial_transform(scatter_sum(edge_attr_x, g.edge_index[1], g.num_nodes))\n", "\n", "print(\"A basis shape (N, n, angular, c):\", tuple(A_x.shape))\n", "print(\"max |A diff (via index map)| =\", (A_x[:, :, ang_map, :] - A_up).abs().max().item())" ] }, { "cell_type": "markdown", "id": "d1da2961", "metadata": {}, "source": [ "## Block 5: Symmetrized B basis · eqs 7–10 & fig 1i\n", "\n", "Rotational invariants come from summing products of A entries whose Cartesian indices\n", "pair up with shared factors, weighted by multinomial coefficients\n", "$\\mathcal{C}(\\mathbf{l}) = l!/(l_x!\\,l_y!\\,l_z!)$:\n", "\n", "$$B^{(2)}_{cnl} = \\sum_{\\mathbf{l}} \\mathcal{C}(\\mathbf{l})\\,A^2_{cn\\mathbf{l}}, \\qquad\n", "B^{(3)}_{cnl_1l_2} = \\sum \\mathcal{C}(\\mathbf{l}_1)\\mathcal{C}(\\mathbf{l}_2)\\,\n", "A_{cn\\mathbf{l}_1} A_{cn(\\mathbf{l}_1+\\mathbf{l}_2)} A_{cn\\mathbf{l}_2}, \\dots$$\n", "\n", "Only *connected* combinations are kept (any zero shared factor would factorize into\n", "lower-order invariants), which makes the features **polynomially independent**, far\n", "more compact than ACE or MTP. The xnn `_Symmetrizer` builds the same combination\n", "rules as upstream `find_combo_vectors_nu{2,3,4}` and evaluates them with one\n", "gather–product–`index_add` per body order. B-feature ordering is identical, so no\n", "index map is needed from here on." ] }, { "cell_type": "code", "execution_count": 7, "id": "9590cc4f", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:24:22.815374Z", "iopub.status.busy": "2026-07-20T04:24:22.815307Z", "iopub.status.idle": "2026-07-20T04:24:22.832344Z", "shell.execute_reply": "2026-07-20T04:24:22.831873Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "invariant features N_L for l_max=3: nu=1: 1, nu<=2: 4, nu<=3: 6, nu<=4: 7 (paper fig 2)\n", "B basis shape (N, n, N_L, c): (10, 8, 6, 4)\n", "max |B diff| / max |B| = 2.8604808891911574e-16\n" ] } ], "source": [ "B_up = up_rep.symmetrizer(node_attr=A_up)\n", "B_x = x.symmetrizer(A_x)\n", "\n", "from xnn.gnn.models.cace import _Symmetrizer\n", "counts = {nu: _Symmetrizer(nu, LMAX).n_features for nu in (1, 2, 3, 4)}\n", "print(f\"invariant features N_L for l_max={LMAX}: nu=1: {counts[1]}, nu<=2: {counts[2]}, \"\n", " f\"nu<=3: {counts[3]}, nu<=4: {counts[4]} (paper fig 2)\")\n", "print(\"B basis shape (N, n, N_L, c):\", tuple(B_x.shape))\n", "print(\"max |B diff| / max |B| =\", ((B_x - B_up).abs().max() / B_up.abs().max()).item())" ] }, { "cell_type": "markdown", "id": "c2432063", "metadata": {}, "source": [ "## Block 6: Message passing · eqs 11–14\n", "\n", "One CACE layer combines up to three mechanisms into the next A basis\n", "$A^{(t+1)} = \\big(m_{\\rm Ar} + m_{B\\chi}\\big)/\\sqrt{\\lambda} + M$:\n", "\n", "* **Ar** *(eq 11)*: the sender's A features filtered by a trainable exponential decay\n", " $a\\,e^{-r/r_0} f_{\\rm cut}(r)$, independent per $(l, n, c)$;\n", "* **Bchi** *(eq 12)*: the layer-0 edge basis $\\chi$ re-weighted by a linear function\n", " $H$ of the sender's invariant B features (recursive edge embedding, as in\n", " REANN/ml-ACE); aggregated messages pass through the *shared* radial coupling;\n", "* **M**: a per-$(l,n,c)$ memory coefficient on the node's own features\n", " (the linear update function $G$, eq 14)." ] }, { "cell_type": "code", "execution_count": 8, "id": "ad3f8dff", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:24:22.833472Z", "iopub.status.busy": "2026-07-20T04:24:22.833404Z", "iopub.status.idle": "2026-07-20T04:24:22.850621Z", "shell.execute_reply": "2026-07-20T04:24:22.850178Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "message norm 1/sqrt(avg_num_neighbors) = 0.3333333333333333\n", "max |A(1) diff| / max |A(1)| = 4.069267721151726e-16\n", "max |B(1) diff| / max |B(1)| = 7.378935653896807e-16\n" ] } ], "source": [ "nm, ar, bchi = up_rep.message_passing_list[0]\n", "mem_up = nm(node_feat=A_up)\n", "m_ar = ar(node_feat=A_up, edge_lengths=len_up,\n", " radial_cutoff_fn=up_rep.cutoff_fn(len_up), edge_index=data[\"edge_index\"])\n", "m_bchi = bchi(node_feat=B_up, edge_attri=edge_attr_up, edge_index=data[\"edge_index\"])\n", "A1_up = (up_scatter(src=m_ar, index=data[\"edge_index\"][1], dim=0, dim_size=10)\n", " + up_rep.radial_transform(\n", " up_scatter(src=m_bchi, index=data[\"edge_index\"][1], dim=0, dim_size=10)))\n", "A1_up = A1_up * up_rep.mp_norm_factor + mem_up\n", "B1_up = up_rep.symmetrizer(node_attr=A1_up)\n", "\n", "A1_x = x.interactions[0](A_x, B_x, edge_attr_x, lengths, x.envelope(lengths),\n", " g.edge_index, x.radial_transform)\n", "B1_x = x.symmetrizer(A1_x)\n", "\n", "print(\"message norm 1/sqrt(avg_num_neighbors) =\", x.mp_norm)\n", "print(\"max |A(1) diff| / max |A(1)| =\",\n", " ((A1_x[:, :, ang_map, :] - A1_up).abs().max() / A1_up.abs().max()).item())\n", "print(\"max |B(1) diff| / max |B(1)| =\",\n", " ((B1_x - B1_up).abs().max() / B1_up.abs().max()).item())" ] }, { "cell_type": "markdown", "id": "ebe117ff", "metadata": {}, "source": [ "## Block 7: Readout · eq 15\n", "\n", "The B features of all $T+1$ stages are concatenated and mapped to atomic energies by\n", "the **sum of a linear layer and an MLP** ([32, 16], SiLU): the linear part preserves\n", "the body-ordered contributions, the MLP captures what the truncated expansion misses.\n", "xnn adds the per-species reference energy through its standard `atom_ref` table\n", "(upstream subtracts $E_0$ from the training labels instead)." ] }, { "cell_type": "code", "execution_count": 9, "id": "d122cb3a", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:24:22.851728Z", "iopub.status.busy": "2026-07-20T04:24:22.851662Z", "iopub.status.idle": "2026-07-20T04:24:22.855701Z", "shell.execute_reply": "2026-07-20T04:24:22.855189Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "readout input width = 384 = n(8) x N_L(6) x c(4) x (T+1)(2)\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "max |atomic energy diff| / max |E_i| = 4.810506842768143e-16\n", "(raw magnitudes are huge here because upstream's torch.rand radial-coupling init is not variance-preserving -- training tames this; only *relative* differences are meaningful)\n" ] } ], "source": [ "feats_up = torch.stack([B_up, B1_up], dim=-1).flatten(1)\n", "e_up = (up_readout.outnet(feats_up) + up_readout.linear_nn(feats_up)).squeeze(-1)\n", "\n", "feats_x = torch.stack([B_x, B1_x], dim=-1).flatten(1)\n", "e_x = (x.readout_mlp(feats_x) + x.readout_linear(feats_x)).squeeze(-1)\n", "\n", "print(\"readout input width =\", feats_x.shape[1], f\"= n({NRB}) x N_L({B_x.shape[2]}) x c({NAB**2}) x (T+1)({T+1})\")\n", "print(\"max |atomic energy diff| / max |E_i| =\",\n", " ((e_x - e_up).abs().max() / e_up.abs().max()).item())\n", "print(\"(raw magnitudes are huge here because upstream's torch.rand radial-\"\n", " \"coupling init is not variance-preserving -- training tames this; \"\n", " \"only *relative* differences are meaningful)\")" ] }, { "cell_type": "markdown", "id": "7bdc2f37", "metadata": {}, "source": [ "## Capstone: transplant a *whole* CACE model and compare energy & forces\n", "\n", "The end-to-end check: same weights → same function, including gradients." ] }, { "cell_type": "code", "execution_count": 10, "id": "53e1b647", "metadata": { "execution": { "iopub.execute_input": "2026-07-20T04:24:22.856821Z", "iopub.status.busy": "2026-07-20T04:24:22.856755Z", "iopub.status.idle": "2026-07-20T04:24:23.044856Z", "shell.execute_reply": "2026-07-20T04:24:23.044273Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "E original = -4888643.527288072743 xnn = -4888643.527288074605\n", "|dE|/|E| = 3.81e-16 max|dF|/max|F| = 2.37e-15\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "|E(rotated) - E| / |E| = 2.667103046445623e-15\n" ] } ], "source": [ "E_up = up_out[\"energy\"].sum()\n", "F_up = -torch.autograd.grad(E_up, data[\"positions\"])[0]\n", "\n", "ox = ForceStressOutput(x)(g)\n", "E_x, F_x = float(ox[\"energy\"]), ox[\"forces\"].detach()\n", "\n", "print(f\"E original = {float(E_up):.12f} xnn = {E_x:.12f}\")\n", "print(f\"|dE|/|E| = {abs(E_x - float(E_up))/abs(float(E_up)):.2e} \"\n", " f\"max|dF|/max|F| = {(F_x - F_up).abs().max().item()/F_up.abs().max().item():.2e}\")\n", "\n", "# and the symmetry that Cartesian symmetrization guarantees:\n", "Q, _ = np.linalg.qr(np.random.default_rng(3).normal(size=(3, 3)))\n", "if np.linalg.det(Q) < 0: Q[:, 0] *= -1\n", "g_rot = structure_to_graph({\"pos\": pos @ Q.T, \"atomic_numbers\": Z,\n", " \"cell\": (np.eye(3) * 5.0) @ Q.T, \"pbc\": [True] * 3}, CUT)\n", "print(\"|E(rotated) - E| / |E| =\",\n", " abs(float(ForceStressOutput(x)(g_rot)[\"energy\"]) - E_x) / abs(E_x))" ] }, { "cell_type": "markdown", "id": "14714ba0", "metadata": {}, "source": [ "## Summary\n", "\n", "| block | original `cace` | `xnn` | relative diff |\n", "|---|---|---|---|\n", "| element embedding + edge type (eq 1) | `NodeEncoder`+`NodeEmbedding`+`EdgeEncoder` | `node_attr` + `embed_sender/receiver` outer product | 0 |\n", "| radial basis (eqs 2, 5) | `BesselRBF` × `PolynomialCutoff` | shared `xnn.gnn.featurizers.{BesselRBF, PolynomialCutoff}` | ~1e-16 |\n", "| Cartesian angular basis (eq 2) | `AngularComponent` | `CartesianAngularBasis` (featurizer) | 0 (via index map) |\n", "| A basis + radial coupling (eqs 5–6) | `SharedRadialLinearTransform` | `_SharedRadialTransform` (stacked einsum) | ~1e-16 |\n", "| symmetrized B basis (eqs 7–10) | `Symmetrizer` | `_Symmetrizer` (gather–product–index_add) | ~1e-16 |\n", "| message passing (eqs 11–14) | `NodeMemory`/`MessageAr`/`MessageBchi` | `_CaceInteraction` | ~1e-15 |\n", "| readout (eq 15) | `Atomwise` (MLP + linear) | `readout_mlp` + `readout_linear` | ~1e-15 |\n", "| **whole model** | | | **E & F to float64 round-off (~1e-16)** |\n", "\n", "The xnn CACE is the original CACE as one self-contained model class on the shared\n", "xnn abstractions (`GNNPotential` species bookkeeping + `atom_ref`, shared radial\n", "featurizers, `ForceStressOutput` autograd forces), and it needs no e3nn at all.\n", "Next: `cace_argon_train_test.ipynb` trains both implementations on Argon MD data." ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "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 }