xnn.gnn.models.cace#

CACE (Cheng 2024): Cartesian Atomic Cluster Expansion potential.

A faithful, self-contained re-implementation of the original BingqingCheng/cace reference code on the xnn abstractions: it subclasses GNNPotential (species bookkeeping, per-species energy shift atom_ref), reuses the shared BesselRBF / PolynomialCutoff radial pieces and the CartesianAngularBasis, and adds the genuinely CACE-specific blocks – the tensor-product edge-type encoding, the trainable radial channel coupling, the Cartesian symmetrizer, and the two message-passing mechanisms. Given the same weights it reproduces the original package to machine precision (tests/test_cace.py), needing neither spherical harmonics nor e3nn.

Architecture (paper eqs 1-15, npj Comput Mater 10, 157, 2024):

  • every element is embedded into a low-dimensional vector theta (length n_atom_basis, typically 1-4); an edge type is the flattened tensor product theta_i (x) theta_j – c = n_atom_basis^2 channels (eq 1);

  • the edge basis chi = T_c(s_i, s_j) R_n(r_ji) L_l(r_hat_ji) combines the edge type with a (trainable) Bessel radial basis times a polynomial cutoff and the Cartesian angular monomials x^lx y^ly z^lz (eq 2);

  • summing edges onto nodes gives the atom-centered A basis (eq 6), whose raw radial channels are mixed per (l, c) by a learned linear map (eq 5, the “radial channel coupling”);

  • products of A entries whose angular indices pair up with shared factors are summed with multinomial prefactors into the polynomially independent invariant B features of body order nu (eqs 7-10 and fig 1i);

  • optional message passing (eqs 11-14): per-edge messages m1 = F(r_ji) A_j (learned exponential-decay filter, Ar) and m2 = H(B_j) chi (recursive edge embedding, Bchi) are aggregated, normalized by 1/sqrt(avg_num_neighbors) and combined with a per-node memory term (M) into the next A;

  • the B features of all layers are concatenated and read out by the sum of a linear layer and an MLP (eq 15), plus the per-species reference energy atom_ref.

Upstream conventions preserved: normalized edge vectors (receiver - sender, eps = 1e-9); trainable Bessel basis with the MACE sqrt(2/cutoff) prefactor; torch.rand initialization of the radial-coupling and Ar filter parameters; the shared radial transform applied both to the initial A basis and to the aggregated Bchi messages; readout = linear + [32, 16] SiLU MLP.

Unlike the other xnn GNNs this model has no TorchScript/LAMMPS export path – upstream CACE has none either (MD runs through the ASE calculator, which uses the regular eager forward).

Classes

CACE(species[, cutoff, n_atom_basis, n_rbf, ...])

Faithful CACE (Cheng 2024): Cartesian atomic cluster expansion.
Show Source