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