Allegro, block by block: reproducing the original implementation with xnn#

This notebook checks every piece needed to reproduce the original Allegro model (mir-group/allegro, v0.3.0, the e3nn-era reference implementation of the paper) using the xnn re-implementation (xnn.gnn.models.allegro). For each architectural block we

  1. state the defining equation(s) from the paper,

  2. show the corresponding xnn building block,

  3. run it on a small toy system, and

  4. compare it numerically against the original allegro package (ground truth).

The paper (provided alongside this notebook): Musaelian, Batzner et al., Learning local equivariant representations for large-scale atomistic dynamics, Nat. Commun. 14, 579 (2023). Equation numbers below refer to it. Its Methods pin the nequip/e3nn 0.4.4 stack we compare against.

Ground truth. Where a block has no learnable weights (radial basis shape, spherical harmonics, the Wigner-3j tensor products) the two agree to machine precision out of the box. Where a block has weights, we transplant them from allegro into xnn and check the outputs match to ~\(10^{-16}\). The notebook ends by transplanting an entire Allegro model and showing energy and forces are bit-for-bit identical.

0. Setup: float64 for exact comparison#

# 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)
from e3nn import o3
import xnn, allegro, nequip, e3nn

print("xnn   :", xnn.__version__)
print("allegro:", allegro.__version__, "(original mir-group/allegro -- ground truth)")
print("nequip :", nequip.__version__, "| e3nn:", e3nn.__version__,
      "| torch:", torch.__version__)

xnn   : 0.1.0
allegro: 0.3.0 (original mir-group/allegro -- ground truth)
nequip : 0.6.2 | e3nn: 0.4.4 | torch: 2.5.1+cu121

A toy system#

from xnn.common.data import structure_to_graph

SPECIES = [1, 6, 8]          # H, C, O
CUTOFF  = 5.0                # r_max
rng = np.random.default_rng(1)
pos = rng.uniform(0, 4, (7, 3))
Z   = np.array(([1, 6, 8] * 7)[:7])
graph = structure_to_graph({"pos": pos, "atomic_numbers": Z}, CUTOFF)
print("atoms:", graph.num_nodes, "| edges:", graph.num_edges)

def report(name, diff, tol=1e-12):
    tag = "OK " if diff <= tol else "!! "
    print(f"{tag}{name:<48} max|xnn - allegro| = {diff:.2e}")

atoms: 7 | edges: 42

The Allegro architecture in one picture#

Allegro is strictly local: no message passing. The energy is a sum of per-species-scaled pair energies (paper eqs 5–6),

\[E = \sum_i \Big[\sigma_{Z_i}\,\varepsilon_i + \mu_{Z_i}\Big],\qquad \varepsilon_i = \tfrac{1}{\sqrt{\lambda}}\sum_{j\in\mathcal N(i)} E_{ij},\]

and each ordered pair \(ij\) carries two latents that interact at every layer (paper Fig. 1): an invariant scalar latent \(x^{ij,L}\) and an equivariant tensor latent \(V^{ij,L}_{n,\ell,p}\):

\[\begin{split} \begin{align} & x^{ij,0} = \mathrm{MLP}_{\rm two-body}\big(Z_i, Z_j, B(r_{ij})\big) && \text{two-body embedding (eq 7)}\\[3pt] & V^{ij,0}_{n,\ell,p} = w^{ij,0}_{n,\ell,p}\,\vec Y^{ij}_{\ell,p},\quad w = \mathrm{MLP}^{0}_{\rm embed}(x^{ij,0}) && \text{initial tensors (eqs 9–10)}\\[3pt] & V^{ij,L} = V^{ij,L-1}\otimes \underbrace{\sum_{k\in\mathcal N(i)} w^{ik,L}_{n,\ell_2,p_2}\vec Y^{ik}_{\ell_2,p_2}}_{\text{embedded environment}} && \text{weighted-env tensor product (eqs 11–14)}\\[3pt] & x^{ij,L} = \mathrm{MLP}^L_{\rm latent}\big(x^{ij,L-1}\,\|\, \text{scalar outputs of the TP}\big) && \text{latent update + resnet (eq 15)}\\[3pt] & V^{ij,L}_{n,\ell,p} = \textstyle\sum_{n'} w_{nn'} V^{ij,L}_{n',(\dots)\to(\ell,p)} && \text{channel mixing linear (eq 16)}\\[3pt] & E_{ij} = \mathrm{MLP}_{\rm output}(x^{ij,L_{\max}}) && \text{pair energy (eq 17)} \end{align} \end{split}\]

We reproduce each block in turn.

Block 1: Radial basis (trainable “normalized sinc” Bessel × polynomial cutoff)#

The reference code’s AllegroBesselBasis is \(B_n(r) = \sin(b_n r)\,/\,(\pi r / r_c)\cdot u(r)\) with trainable frequencies \(b_n = n\pi/r_c\) and the degree-\(p\) polynomial envelope \(u\) (eq 8; “Bessel basis functions with a polynomial envelope”). The xnn featurizer reproduces it with BesselRBF(..., trainable=True, prefactor=r_c/π), the same family as NequIP/MACE, differing only in prefactor and frequency parameterization (transplant: freqs = bessel_weights · r_c).

from xnn.common.models import build_model
from xnn.common.config import from_dict

LMAX, NL, NF, NRBF, AVG = 2, 2, 8, 8, 8.0
TB, LAT, EE = [16, 32], [32], [16]

cfg = from_dict({"model": {"name": "allegro", "cutoff": CUTOFF, "n_features": NF,
    "n_interactions": NL, "n_rbf": NRBF,
    "extra": {"species": SPECIES, "l_max": LMAX, "avg_num_neighbors": AVG,
              "two_body_latent": TB, "latent": LAT, "edge_eng": EE}}})
xal = build_model(cfg.model)

from nequip.model import model_from_config
n_model = model_from_config(dict(
    model_builders=["allegro.model.Allegro"],
    r_max=CUTOFF, num_layers=NL, l_max=LMAX, parity="o3_full",
    num_tensor_features=NF, num_bessels_per_basis=NRBF, PolynomialCutoff_p=6.0,
    avg_num_neighbors=AVG, chemical_symbols=["H", "C", "O"],
    two_body_latent_mlp_latent_dimensions=TB, latent_mlp_latent_dimensions=LAT,
    env_embed_mlp_latent_dimensions=[], edge_eng_mlp_latent_dimensions=EE,
), initialize=True)
seq = n_model.model          # the Allegro sequential (no force wrapper here)
print("original blocks:", [n for n, _ in seq.named_children()])

# transplant the trainable Bessel frequencies and compare the basis on real edges
with torch.no_grad():
    xal.edge_feat.rbf.freqs.copy_(seq.radial_basis.bessel_weights * CUTOFF)
edge = xal.edge_feat(graph)
r = edge["edge_length"]
ref = torch.sin(r[:, None] * seq.radial_basis.bessel_weights) / (np.pi * r[:, None] / CUTOFF)
p = 6.0; x = r[:, None] / CUTOFF
env = (1 - (p+1)*(p+2)/2*x**p + p*(p+2)*x**(p+1) - p*(p+1)/2*x**(p+2)) * (x < 1)
report("Bessel basis x cutoff  B_n(r) u(r)", (edge["edge_radial"] - ref*env).abs().max().item())

original blocks: ['radial_basis', 'typeembed', 'spharm', 'allegro', 'edge_eng', 'edge_eng_sum', 'total_energy_sum']
OK Bessel basis x cutoff  B_n(r) u(r)               max|xnn - allegro| = 8.60e-16

Block 2: Two-body scalar embedding · eq 7#

The pair’s initial invariant latent embeds the chemistry of the pair and the distance. The reference code implements eq 7 as ProductTypeEmbedding: learned per-type embeddings of centre and neighbour, concatenated and multiplied elementwise with a linear projection of the radial basis. In xnn this is the type_embeddings parameter plus basis_embed. Both weights transplant.

def copy_fcn(fcn, mod):
    '''upstream ScalarMLPFunction -> e3nn FullyConnectedNet (same math, verified)'''
    sd = dict(mod.named_parameters())
    with torch.no_grad():
        for i in range(len(fcn.hs) - 1):
            getattr(fcn, f"layer{i}").weight.copy_(sd[f"_forward._weight_{i}"])

with torch.no_grad():
    xal.type_embeddings.copy_(seq.typeembed.type_embeddings)
copy_fcn(xal.basis_embed, seq.typeembed.basis_mlp)

# xnn two-body embedding on the toy graph (centre = edge_index[1] in xnn)
types = xal.z_to_index[graph.atomic_numbers].clamp(min=0)
ei = graph.edge_index
x_embed = torch.cat((xal.type_embeddings[0][types[ei[1]]],
                     xal.type_embeddings[1][types[ei[0]]]), dim=-1) \
          * xal.basis_embed(edge["edge_radial"])

# original, driven with the same per-edge inputs (its centre is row 0)
from nequip.data import AtomicDataDict
data = {AtomicDataDict.ATOM_TYPE_KEY: types.unsqueeze(-1),
        AtomicDataDict.EDGE_INDEX_KEY: ei.flip(0),
        AtomicDataDict.EDGE_EMBEDDING_KEY: edge["edge_radial"],
        AtomicDataDict.POSITIONS_KEY: graph.pos,
        AtomicDataDict.EDGE_CELL_SHIFT_KEY: graph.cell_shifts.to(graph.pos.dtype)}
n_embed = seq.typeembed(dict(data))[AtomicDataDict.EDGE_EMBEDDING_KEY]
report("two-body product embedding  x^(ij,0)", (x_embed - n_embed).abs().max().item())

OK two-body product embedding  x^(ij,0)             max|xnn - allegro| = 8.88e-16

Block 3: Spherical harmonics & initial tensor features · eqs 9–10#

The angular basis is \(\vec Y^{ij}_{\ell,p} = Y^m_\ell(\hat r_{ij})\) with \(\vec r_{ij} = \vec r_j - \vec r_i\) (paper notation), the opposite orientation to the xnn/MACE edge vector, so the model flips internally, exactly as the xnn NequIP does. The initial equivariant features are the spherical harmonics weighted into \(n_{\rm tensor}\) channels by weights generated from the scalar latent (eqs 9–10): _ChannelWeighter vs the original MakeWeightedChannels.

from allegro.nn._strided import MakeWeightedChannels
from xnn.gnn.models.allegro import _ChannelWeighter

vec = graph.edge_vectors()
ir_sh = o3.Irreps.spherical_harmonics(LMAX)
ref_sh = o3.spherical_harmonics(ir_sh, -vec, normalize=True, normalization="component")
_, xsh, _ = xal.edge_feat.embed(-vec)
report("spherical harmonics  Y(r_j - r_i)", (xsh - ref_sh).abs().max().item())

m_wc = MakeWeightedChannels(irreps_in=ir_sh, multiplicity_out=NF)
x_wc = _ChannelWeighter(ir_sh, NF)
w = torch.randn(graph.num_edges, m_wc.weight_numel)
report("weighted channels  V^(ij,0) = w Y", (x_wc(xsh, w) - m_wc(xsh, w)).abs().max().item(), tol=0.0)

OK spherical harmonics  Y(r_j - r_i)                max|xnn - allegro| = 0.00e+00
OK weighted channels  V^(ij,0) = w Y                max|xnn - allegro| = 0.00e+00

Block 4: The weighted-environment tensor product · eqs 11–14#

The core of Allegro. The embedded environment of atom \(i\) is the weighted sum \(\sum_{k\in\mathcal N(i)} w^{ik,L}_{n,\ell_2,p_2}\vec Y^{ik}\) (eq 14, with learned weights from the scalar latents, the “density trick”, eqs 12–13; the reference code subtracts the edge’s own term and normalizes by \(1/\sqrt{\lambda-1}\)). The pair tensors are then contracted with it through a per-channel (“uuu”), weightless tensor product whose Wigner-3j blocks carry \(\sqrt{2\ell_{\rm out}+1}\): xnn’s _Contracter vs the original strided Contracter, bit-for-bit.

from allegro.nn._strided import Contracter
from xnn.gnn.models.allegro import _Contracter

arg_irs = [ir for _, ir in ir_sh]
env_irs = [ir for _, ir in ir_sh]
# allowed o3_full irreps, pruned to those reachable from arg x env (as the model does)
out_irs = [ir for ir in (o3.Irrep(l, p) for l in range(LMAX + 1) for p in (1, -1))
           if any(ir in a * e for a in arg_irs for e in env_irs)]
instr, full_out = [], []
for ir_out in out_irs:
    for i1, ir1 in enumerate(arg_irs):
        for i2, ir2 in enumerate(env_irs):
            if ir_out in ir1 * ir2:
                instr.append((i1, i2, ir_out)); full_out.append(ir_out)

m_tp = Contracter(
    irreps_in1=o3.Irreps([(NF, ir) for ir in arg_irs]),
    irreps_in2=o3.Irreps([(NF, ir) for ir in env_irs]),
    irreps_out=o3.Irreps([(NF, ir) for ir in full_out]),
    instructions=[(i1, i2, k) for k, (i1, i2, _) in enumerate(instr)],
    connection_mode="uuu", shared_weights=False, has_weight=False)
x_tp = _Contracter(arg_irs, env_irs, instr)

E = graph.num_edges
t1 = torch.randn(E, NF, sum(ir.dim for ir in arg_irs))
t2 = torch.randn(E, NF, sum(ir.dim for ir in env_irs))
report("weightless uuu tensor product (eqs 11-13)",
       (x_tp(t1, t2) - m_tp(t1, t2)).abs().max().item())

OK weightless uuu tensor product (eqs 11-13)        max|xnn - allegro| = 1.07e-14

Block 5: Scalar latent MLPs & the cumulative-softmax resnet · eq 15#

The scalar outputs of each TP are concatenated with the previous latent and compressed by the latent MLP (eq 15). All Allegro MLPs are variance-preserving scalar MLPs; e3nn.nn.FullyConnectedNet is numerically identical to the original ScalarMLPFunction given the same weights. The residual update uses normalized cumulative-softmax coefficients (Supplementary Note 2), reproduced in xnn (_resnet_params).

from allegro.nn._fc import ScalarMLPFunction
from e3nn.nn import FullyConnectedNet
import torch.nn.functional as Fn

up = ScalarMLPFunction(mlp_input_dimension=24, mlp_latent_dimensions=[32, 64],
                       mlp_output_dimension=16, mlp_nonlinearity="silu")
fc = FullyConnectedNet([24, 32, 64, 16], Fn.silu)
copy_fcn(fc, up)
t = torch.randn(9, 24)
report("variance-preserving scalar MLP (eq 15)", (fc(t) - up(t)).abs().max().item())

# resnet coefficients: zeros -> exp() -> all-equal cumulative softmax
params = torch.zeros(NL + 1)
coeff = (params - params.max()).exp(); cumsum = coeff.cumsum(0) + 1e-12
print("layer resnet coefficients (old, new) per layer:",
      [(float((cumsum[i-1]/cumsum[i]).sqrt()), float((coeff[i]/cumsum[i]).sqrt()))
       for i in range(1, NL + 1)])

OK variance-preserving scalar MLP (eq 15)           max|xnn - allegro| = 1.33e-15
layer resnet coefficients (old, new) per layer: [(0.7071067811867243, 0.7071067811863707), (0.8164965809277941, 0.5773502691895296)]

Block 6: Channel-mixing linear · eq 16#

The outputs of all TP paths with the same irrep are linearly mixed back into num_tensor_features channels with a \(1/\sqrt{\text{mul}\cdot n_{\rm paths}}\) normalization (eq 16). xnn’s _StridedLinear keeps the original’s flat weight layout, so the weight vector transplants directly.

from allegro.nn._strided import Linear as StridedLinearRef
from xnn.gnn.models.allegro import _StridedLinear

m_lin = StridedLinearRef(o3.Irreps([(NF, ir) for ir in full_out]),
                         o3.Irreps([(NF, ir) for ir in out_irs]),
                         shared_weights=True, internal_weights=True)
x_lin = _StridedLinear(full_out, out_irs, NF)
with torch.no_grad():
    x_lin.w.copy_(m_lin.w)
t = torch.randn(E, NF, sum(ir.dim for ir in full_out))
report("channel-mixing linear (eq 16)",
       (x_lin(t) - m_lin(t.reshape(E, -1))).abs().max().item())

OK channel-mixing linear (eq 16)                    max|xnn - allegro| = 1.33e-15

Block 7: Pair energy, edgewise sum, per-species scale/shift · eqs 5–6, 17#

The final scalar latent is read out to the pair energy \(E_{ij}\) by the output MLP (eq 17); pair energies are summed onto their centre atom with the \(1/\sqrt{\lambda}\) normalization, then the per-species scale/shift of eq 5 (atom_scale / atom_ref in xnn, PerSpeciesRescale upstream) and autograd forces \(\vec F = -\nabla E\). With a zero scale the energy is exactly the sum of shifts and forces vanish.

from xnn.common.models import ForceStressOutput

cfg0 = from_dict({"model": {"name": "allegro", "cutoff": CUTOFF, "n_features": NF,
    "n_interactions": 1, "extra": {"species": SPECIES, "l_max": 1,
    "two_body_latent": TB, "latent": LAT, "edge_eng": EE,
    "atomic_energies": [0.5, -1.3, -2.1], "atomic_scales": 0.0}}})
m0 = ForceStressOutput(build_model(cfg0.model))
o0 = m0(graph)
E0_expected = sum({1: 0.5, 6: -1.3, 8: -2.1}[int(z)] for z in Z)
report("sigma=0 energy == sum of shifts", abs(float(o0["energy"]) - E0_expected))
report("sigma=0 forces == 0", o0["forces"].abs().max().item())

OK sigma=0 energy == sum of shifts                  max|xnn - allegro| = 0.00e+00
OK sigma=0 forces == 0                              max|xnn - allegro| = 0.00e+00

Capstone: transplant a whole Allegro model and compare energy & forces#

We rebuild the original with the per-species scale/shift and force output, transplant every weight (Bessel frequencies, type/basis embedding, all latent / env-embed MLPs, the channel-mixing linears, resnet coefficients, the output MLP, and the scale/shift), and run both on the toy molecule, the original through its own data pipeline.

from nequip.data import AtomicData
from nequip.data.transforms import TypeMapper

E0 = [0.5, -1.3, -2.1]; SIG = [1.7, 0.9, 1.1]
n_full = model_from_config(dict(
    model_builders=["allegro.model.Allegro", "PerSpeciesRescale", "ForceOutput"],
    r_max=CUTOFF, num_layers=NL, l_max=LMAX, parity="o3_full",
    num_tensor_features=NF, num_bessels_per_basis=NRBF, PolynomialCutoff_p=6.0,
    avg_num_neighbors=AVG, chemical_symbols=["H", "C", "O"],
    two_body_latent_mlp_latent_dimensions=TB, latent_mlp_latent_dimensions=LAT,
    env_embed_mlp_latent_dimensions=[], edge_eng_mlp_latent_dimensions=EE,
    per_species_rescale_shifts=E0, per_species_rescale_scales=SIG,
), initialize=True)
seq = n_full.model.func
al = seq.allegro

cfgF = from_dict({"model": {"name": "allegro", "cutoff": CUTOFF, "n_features": NF,
    "n_interactions": NL, "n_rbf": NRBF,
    "extra": {"species": SPECIES, "l_max": LMAX, "avg_num_neighbors": AVG,
              "two_body_latent": TB, "latent": LAT, "edge_eng": EE,
              "atomic_energies": E0, "atomic_scales": SIG}}})
xfull = build_model(cfgF.model)

def transplant_full(x, seq, n_layers):
    al = seq.allegro
    with torch.no_grad():
        x.edge_feat.rbf.freqs.copy_(seq.radial_basis.bessel_weights * float(al.r_max))
        x.type_embeddings.copy_(seq.typeembed.type_embeddings)
        copy_fcn(x.basis_embed, seq.typeembed.basis_mlp)
        for i in range(n_layers):
            copy_fcn(x.latents[i], al.latents[i])
            copy_fcn(x.env_embed_mlps[i], al.env_embed_mlps[i])
            x.linears[i].w.copy_(al.linears[i].w)
        copy_fcn(x.final_latent, al.final_latent)
        copy_fcn(x.edge_eng, seq.edge_eng._module)
        x._resnet_params.copy_(al._latent_resnet_coefficients_params)

transplant_full(xfull, seq, NL)

ox = ForceStressOutput(xfull)(structure_to_graph({"pos": pos, "atomic_numbers": Z}, CUTOFF))
tm = TypeMapper(chemical_symbols=["H", "C", "O"])
dd = AtomicData.to_AtomicDataDict(tm(AtomicData.from_points(
    pos=torch.tensor(pos), r_max=CUTOFF, atomic_numbers=torch.tensor(Z))))
dd[AtomicDataDict.POSITIONS_KEY].requires_grad_(True)
om = n_full(dd)
E_m = float(om[AtomicDataDict.TOTAL_ENERGY_KEY].sum())
F_m = -torch.autograd.grad(om[AtomicDataDict.TOTAL_ENERGY_KEY].sum(),
                           dd[AtomicDataDict.POSITIONS_KEY])[0]

print(f"total energy   xnn    = {float(ox['energy']):.10f} eV")
print(f"total energy   allegro = {E_m:.10f} eV")
report("FULL MODEL total energy", abs(float(ox["energy"]) - E_m))
report("FULL MODEL per-atom forces",
       (ox["forces"].detach() - F_m).abs().max().item())

total energy   xnn    = -8.7580160536 eV
total energy   allegro = -8.7580160536 eV
OK FULL MODEL total energy                          max|xnn - allegro| = 1.78e-15
OK FULL MODEL per-atom forces                       max|xnn - allegro| = 6.22e-15

Summary#

Block

Paper eq

Weights?

agreement

1. Trainable “normalized sinc” Bessel × cutoff

8

transplanted (\(b_n\))

machine precision

2. Two-body product embedding

7

transplanted

machine precision

3. Spherical harmonics + weighted channels

9–10

transplanted

bit-identical (0)

4. Weighted-env “uuu” tensor product

11–14

none (w3j)

machine precision

5. Scalar latent MLPs + cumulative-softmax resnet

15

transplanted

machine precision

6. Channel-mixing linear

16

transplanted (flat layout)

machine precision

7. Pair energy, \(1/\sqrt\lambda\) sums, \(\sigma_Z/\mu_Z\)

5–6, 17

n/a

exact

Full model

5–17

all transplanted

energy 0, forces ~1e-15

xnn.gnn.models.allegro is a faithful reproduction of the original Allegro, depending only on e3nn (no allegro/nequip/opt_einsum_fx), sharing the xnn equivariant-GNN abstractions with NequIP and MACE, and TorchScript / LAMMPS-deployable (the pair_allegro path). The companion notebooks train and test it on realistic Argon MD data.