Argon density from MD: `xnn` vs the original Allegro#

Computes the mass density of liquid Argon from NPT molecular dynamics through ASE with a trained Allegro potential, in the same two complementary tracks as the MACE/NequIP companions:

Track (a): same potential. The trained original-Allegro weights are copied into xnn; any density difference reflects only the inference/MD code path (should be ~zero).
Track (b): independently trained. xnn is trained from scratch on the same data; densities are compared as two practitioners would.

Pipeline per track: train → ASE calculator (energy + forces + stress) → NPT MD (\(T=85\) K, \(P=1\) bar) → \(\rho = M/V\).

0. Setup: train `float32`, run MD `float64`#

# 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 time
import numpy as np
import torch
import matplotlib.pyplot as plt
import ase.io
import ase.units as u

torch.set_default_dtype(torch.float32)
torch.manual_seed(0)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
CUTOFF, SPECIES = 6.0, [18]
DATA = "../../../datasets/argon_md"
import xnn, allegro
print("xnn:", xnn.__version__, "| allegro (original):", allegro.__version__,
      "| device:", DEVICE)

xnn: 0.1.0 | allegro (original): 0.3.0 | device: cuda

1. Load data and build both data pipelines#

from xnn.common.data import AtomicDataset, load_dataset
from nequip.data import AtomicData
from nequip.data.dataloader import DataLoader as NequipDataLoader
from nequip.data.transforms import TypeMapper

E0 = {18: 0.0}   # argon isolated-atom reference energy
train_structs = load_dataset("argon_md", split="train")
ds_all = AtomicDataset(train_structs, CUTOFF)
train_structs = [s for i, s in enumerate(train_structs) if ds_all[i].num_edges > 0]

TM = TypeMapper(chemical_symbols=["Ar"])
xnn_train = AtomicDataset(train_structs, CUTOFF)
LAMBDA = float(sum(xnn_train[i].num_edges for i in range(len(xnn_train))) /
               sum(xnn_train[i].num_nodes for i in range(len(xnn_train))))

def to_upstream(s):
    d = AtomicData.from_points(
        pos=torch.tensor(s["pos"], dtype=torch.get_default_dtype()), r_max=CUTOFF,
        atomic_numbers=torch.tensor(s["atomic_numbers"]),
        cell=torch.tensor(s["cell"], dtype=torch.get_default_dtype()),
        pbc=torch.tensor([True]*3),
        total_energy=torch.tensor([s["energy"]], dtype=torch.get_default_dtype()),
        forces=torch.tensor(s["forces"], dtype=torch.get_default_dtype()))
    return TM(d)
al_train = [to_upstream(s) for s in train_structs]

g = torch.Generator().manual_seed(0)
perm = torch.randperm(len(train_structs), generator=g).tolist()
n_val = max(1, int(0.1 * len(train_structs)))
val_idx, train_idx = perm[:n_val], perm[n_val:]
print(f"{len(train_structs)} configs | E0={E0} | lambda={LAMBDA:.2f} | "
      f"train {len(train_idx)} / val {len(val_idx)}")

193 configs | E0={18: 0.0} | lambda=17.78 | train 174 / val 19

2. Train the two models#

Same architecture (2 layers, \(\ell_{\max}=2\) o3_full, 32 tensor channels), same data/loss/optimiser/schedule/split. Original with a native loop; xnn with xnn.train.Trainer.

from nequip.model import model_from_config
from xnn.common.config import from_dict
from xnn.common.train import Trainer
from torch.utils.data import Subset

NL, LMAX, NF = 2, 2, 32
TB, LAT, EE = [32, 64, 128], [128], [32]
EW, FW, LR, WD, BS, EPOCHS = 1.0, 100.0, 0.01, 5e-7, 10, 60

UPSTREAM_HP = dict(
    r_max=CUTOFF, num_layers=NL, l_max=LMAX, parity="o3_full",
    num_tensor_features=NF, num_bessels_per_basis=8, PolynomialCutoff_p=6.0,
    avg_num_neighbors=LAMBDA, chemical_symbols=["Ar"],
    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[18]], per_species_rescale_scales=[1.0])

# ---------- 2.1 train the ORIGINAL Allegro (native loop) ----------
torch.manual_seed(0)
al_model = model_from_config(dict(
    model_builders=["allegro.model.Allegro", "PerSpeciesRescale", "ForceOutput"],
    **UPSTREAM_HP), initialize=True).to(DEVICE)

tr_loader = NequipDataLoader([al_train[i] for i in train_idx], batch_size=BS, shuffle=True)
va_loader = NequipDataLoader([al_train[i] for i in val_idx],   batch_size=BS, shuffle=False)
opt = torch.optim.Adam(al_model.parameters(), lr=LR, weight_decay=WD)
sched = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, patience=10)
def al_loss(out, b):
    n = (b.ptr[1:] - b.ptr[:-1]).to(out["total_energy"].dtype)
    return EW * (((out["total_energy"].squeeze(-1) - b.total_energy.squeeze(-1)) / n) ** 2).mean() \
         + FW * ((out["forces"] - b.forces) ** 2).mean()

t0 = time.time()
for epoch in range(EPOCHS):
    al_model.train()
    for b in tr_loader:
        b = b.to(DEVICE)
        loss = al_loss(al_model(AtomicData.to_AtomicDataDict(b)), b)
        opt.zero_grad(); loss.backward(); opt.step()
    al_model.eval(); vl = 0.0
    for b in va_loader:
        b = b.to(DEVICE)
        vl += float(al_loss(al_model(AtomicData.to_AtomicDataDict(b)), b))
    sched.step(vl / len(va_loader))
print(f"original Allegro trained {EPOCHS} epochs in {time.time()-t0:.0f} s "
      f"(final val {vl/len(va_loader):.3e})")

# ---------- 2.2 train xnn INDEPENDENTLY (Trainer) ----------
core = from_dict({
    "model": {"name": "allegro", "cutoff": CUTOFF, "n_features": NF,
              "n_interactions": NL, "species": SPECIES, "l_max": LMAX,
              "avg_num_neighbors": LAMBDA, "two_body_latent": TB, "latent": LAT,
              "edge_eng": EE, "atomic_energies": [E0[18]]},
    "data": {"batch_size": BS},
    "optim": {"lr": LR, "weight_decay": WD, "epochs": EPOCHS, "energy_weight": EW,
              "force_weight": FW, "scheduler": "plateau"},
    "device": DEVICE, "seed": 0, "output_dir": "runs/argon_md_indep",
})
t0 = time.time()
trainer = Trainer(core, Subset(xnn_train, train_idx), Subset(xnn_train, val_idx))
trainer.fit()
print(f"xnn (independent) trained {EPOCHS} epochs in {time.time()-t0:.0f} s")
xnn_indep_base = trainer.model.model

original Allegro trained 60 epochs in 193 s (final val 4.061e-04)
epoch    0 | train loss 9.3339e-01 | val loss 1.7936e-01
epoch    1 | train loss 7.4252e-02 | val loss 1.2085e-02
epoch    2 | train loss 1.1073e-02 | val loss 2.6839e-03
epoch    3 | train loss 4.4131e-03 | val loss 1.6396e-03
epoch    4 | train loss 2.8103e-03 | val loss 1.3049e-03
epoch    5 | train loss 2.0568e-03 | val loss 1.1053e-03
epoch    6 | train loss 1.7918e-03 | val loss 9.4939e-04
epoch    7 | train loss 1.3174e-03 | val loss 8.5242e-04
epoch    8 | train loss 1.0816e-03 | val loss 7.6886e-04
epoch    9 | train loss 9.3924e-04 | val loss 7.2583e-04
epoch   10 | train loss 8.5024e-04 | val loss 7.0451e-04
epoch   11 | train loss 7.8161e-04 | val loss 6.7192e-04
epoch   12 | train loss 7.4336e-04 | val loss 6.5199e-04
epoch   13 | train loss 6.9064e-04 | val loss 6.1449e-04
epoch   14 | train loss 6.5856e-04 | val loss 6.3887e-04
epoch   15 | train loss 6.9874e-04 | val loss 5.7859e-04
epoch   16 | train loss 6.5901e-04 | val loss 5.5251e-04
epoch   17 | train loss 5.9613e-04 | val loss 5.2888e-04
epoch   18 | train loss 6.0417e-04 | val loss 5.3255e-04
epoch   19 | train loss 5.5718e-04 | val loss 5.0468e-04
epoch   20 | train loss 5.2535e-04 | val loss 4.8798e-04
epoch   21 | train loss 5.1400e-04 | val loss 4.8761e-04
epoch   22 | train loss 5.2655e-04 | val loss 4.9559e-04
epoch   23 | train loss 5.0366e-04 | val loss 4.8256e-04
epoch   24 | train loss 4.7429e-04 | val loss 4.5582e-04
epoch   25 | train loss 4.6373e-04 | val loss 4.4049e-04
epoch   26 | train loss 4.4726e-04 | val loss 4.6189e-04
epoch   27 | train loss 4.4362e-04 | val loss 4.5016e-04
epoch   28 | train loss 4.3721e-04 | val loss 4.2293e-04
epoch   29 | train loss 4.1618e-04 | val loss 4.2283e-04
epoch   30 | train loss 4.2706e-04 | val loss 4.6601e-04
epoch   31 | train loss 4.2393e-04 | val loss 3.9880e-04
epoch   32 | train loss 4.1589e-04 | val loss 4.0192e-04
epoch   33 | train loss 3.9876e-04 | val loss 3.9462e-04
epoch   34 | train loss 4.1270e-04 | val loss 3.8881e-04
epoch   35 | train loss 4.1359e-04 | val loss 3.9523e-04
epoch   36 | train loss 3.9687e-04 | val loss 3.6866e-04
epoch   37 | train loss 3.7622e-04 | val loss 3.9155e-04
epoch   38 | train loss 3.8085e-04 | val loss 3.7154e-04
epoch   39 | train loss 3.8232e-04 | val loss 3.7094e-04
epoch   40 | train loss 3.6297e-04 | val loss 3.8782e-04
epoch   41 | train loss 3.7186e-04 | val loss 3.8430e-04
epoch   42 | train loss 3.6196e-04 | val loss 3.7639e-04
epoch   43 | train loss 3.6581e-04 | val loss 3.7024e-04
epoch   44 | train loss 3.9327e-04 | val loss 3.7947e-04
epoch   45 | train loss 3.7136e-04 | val loss 3.6896e-04
epoch   46 | train loss 3.5533e-04 | val loss 3.3163e-04
epoch   47 | train loss 3.6300e-04 | val loss 3.3912e-04
epoch   48 | train loss 3.5166e-04 | val loss 3.5246e-04
epoch   49 | train loss 3.4953e-04 | val loss 3.9673e-04
epoch   50 | train loss 3.8315e-04 | val loss 3.5802e-04
epoch   51 | train loss 3.5842e-04 | val loss 3.6046e-04
epoch   52 | train loss 3.3230e-04 | val loss 3.3324e-04
epoch   53 | train loss 3.5310e-04 | val loss 3.4285e-04
epoch   54 | train loss 3.4336e-04 | val loss 3.5681e-04
epoch   55 | train loss 3.3566e-04 | val loss 3.3169e-04
epoch   56 | train loss 3.2688e-04 | val loss 3.3325e-04
epoch   57 | train loss 3.3887e-04 | val loss 3.3419e-04
epoch   58 | train loss 3.1849e-04 | val loss 3.3383e-04
epoch   59 | train loss 3.1927e-04 | val loss 3.3496e-04
xnn (independent) trained 60 epochs in 523 s

Common MD utilities#

Switch to float64; rebuild the trained original with StressForceOutput (adds the barostat’s stress via the strain trick, same convention as the xnn ForceStressOutput) and load the trained weights.

torch.set_default_dtype(torch.float64)
from xnn.common.models import build_model, ForceStressOutput
from ase import Atoms
from ase.calculators.calculator import Calculator, all_changes
from ase.md.nptberendsen import NPTBerendsen
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution, Stationary
from xnn.common.deploy import XNNCalculator

al_md_model = model_from_config(dict(
    model_builders=["allegro.model.Allegro", "PerSpeciesRescale", "StressForceOutput"],
    **UPSTREAM_HP), initialize=True)
# keep the float64 rebuild's exact Wigner-3j constants (the float32 model's
# _big_w3j buffers are float32-rounded); load only the learned state
sd = {k: v.double() for k, v in al_model.state_dict().items() if "_big_w3j" not in k}
al_md_model.load_state_dict(sd, strict=False)   # only misses a dummy buffer
al_md_model = al_md_model.to(DEVICE).eval()
xnn_indep_base = xnn_indep_base.double().eval()

class AllegroASECalculator(Calculator):
    '''Minimal ASE calculator wrapping an in-memory original-Allegro model.'''
    implemented_properties = ["energy", "forces", "stress"]
    def __init__(self, model, cutoff, device="cuda", **kw):
        super().__init__(**kw); self.model, self.cutoff, self.device = model, cutoff, device
    def calculate(self, atoms=None, properties=("energy",), system_changes=all_changes):
        super().calculate(atoms, properties, system_changes)
        d = TM(AtomicData.from_ase(atoms.copy(), r_max=self.cutoff))
        b = next(iter(NequipDataLoader([d], batch_size=1))).to(self.device)
        out = self.model(AtomicData.to_AtomicDataDict(b))
        self.results["energy"] = float(out["total_energy"].sum().detach())
        self.results["forces"] = out["forces"].detach().cpu().numpy()
        s = out["stress"][0].detach().cpu().numpy()
        self.results["stress"] = np.array([s[0,0], s[1,1], s[2,2], s[1,2], s[0,2], s[0,1]])

T_K, P_BAR, DT = 85.0, 1.0, 5 * u.fs
N_EQUIL, N_PROD = 300, 700
AMU_A3_TO_G_CM3 = 1.6605390666
a0 = train_structs[0]

def density(atoms):
    return atoms.get_masses().sum() / atoms.get_volume() * AMU_A3_TO_G_CM3

def compare_calcs(make_x, make_a):
    at = Atoms(numbers=a0["atomic_numbers"], positions=a0["pos"], cell=a0["cell"], pbc=True)
    ax = at.copy(); ax.calc = make_x(); aa = at.copy(); aa.calc = make_a()
    return (abs(ax.get_potential_energy() - aa.get_potential_energy()),
            np.abs(ax.get_forces() - aa.get_forces()).max(),
            np.abs(ax.get_stress() - aa.get_stress()).max())

def run_npt(make_calc, label):
    at = Atoms(numbers=a0["atomic_numbers"], positions=a0["pos"], cell=a0["cell"], pbc=True)
    at.calc = make_calc()
    MaxwellBoltzmannDistribution(at, temperature_K=T_K, rng=np.random.default_rng(0)); Stationary(at)
    dyn = NPTBerendsen(at, timestep=DT, temperature_K=T_K, pressure_au=P_BAR * u.bar,
                       taut=100 * u.fs, taup=1000 * u.fs, compressibility_au=2e-4 / u.bar)
    rho = np.empty(N_EQUIL + N_PROD)
    t0 = time.time()
    for k in range(N_EQUIL + N_PROD):
        dyn.run(1); rho[k] = density(at)
    print(f"{label}: {N_EQUIL+N_PROD} steps in {time.time()-t0:.0f} s | "
          f"rho_eq = {rho[N_EQUIL:].mean():.4f} g/cm³")
    return rho

RHO_EXP = 1.41
print(f"initial density = {density(Atoms(numbers=a0['atomic_numbers'], positions=a0['pos'], cell=a0['cell'], pbc=True)):.4f} g/cm³ "
      f"| target T={T_K} K, P={P_BAR} bar | exp ~{RHO_EXP} g/cm³")

initial density = 1.7910 g/cm³ | target T=85.0 K, P=1.0 bar | exp ~1.41 g/cm³

Track (a): same potential (weights copied Allegro → xnn)#

3a. Copy the trained original-Allegro weights into `xnn`#

def copy_fcn(fcn, mod):
    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}"].to(getattr(fcn, f"layer{i}").weight.dtype))

def copy_allegro_into_xnn(x, al_gm, n_layers):
    seq = al_gm.model.func
    al = seq.allegro
    with torch.no_grad():
        x.edge_feat.rbf.freqs.copy_(seq.radial_basis.bessel_weights.double() * float(al.r_max))
        x.type_embeddings.copy_(seq.typeembed.type_embeddings.double())
        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.double())
        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.double())
        psr = seq.per_species_rescale
        for k, z in enumerate(SPECIES):
            x.atom_ref.weight[z] = float(psr.shifts[k])
            x.atom_scale[z] = float(psr.scales[k])

xnn_shared_base = build_model(core.model)          # fresh float64 model
copy_allegro_into_xnn(xnn_shared_base, al_md_model, NL)
xnn_shared = ForceStressOutput(xnn_shared_base, compute_forces=True,
                                compute_stress=True).to(DEVICE).double().eval()

def xnn_a():    return XNNCalculator(xnn_shared, cutoff=CUTOFF, device=DEVICE)
def allegro_a(): return AllegroASECalculator(al_md_model, CUTOFF, DEVICE)

dE, dF, dS = compare_calcs(xnn_a, allegro_a)
print("track (a) calculators on one Argon config (SAME potential):")
print(f"  dE = {dE:.2e} eV | dF = {dF:.2e} eV/Å | dσ = {dS:.2e} eV/Å³  -> identical")

track (a) calculators on one Argon config (SAME potential):
  dE = 1.28e-06 eV | dF = 6.60e-08 eV/Å | dσ = 3.08e-10 eV/Å³  -> identical

4a/5a. NPT MD: same potential through both codes#

rho_xa = run_npt(xnn_a, "xnn (a)")
rho_aa = run_npt(allegro_a, "allegro (a)")
da_x, da_a = rho_xa[N_EQUIL:].mean(), rho_aa[N_EQUIL:].mean()
print(f"track (a)  rho_xnn = {da_x:.4f}   rho_allegro = {da_a:.4f}   "
      f"|diff| = {abs(da_x-da_a):.2e} g/cm³")

xnn (a): 1000 steps in 154 s | rho_eq = 1.4133 g/cm³
allegro (a): 1000 steps in 102 s | rho_eq = 1.4133 g/cm³
track (a)  rho_xnn = 1.4133   rho_allegro = 1.4133   |diff| = 4.82e-08 g/cm³

/tmp/ipykernel_1044677/1079695994.py:52: DeprecationWarning: Use thermalize_momenta
  MaxwellBoltzmannDistribution(at, temperature_K=T_K, rng=np.random.default_rng(0)); Stationary(at)

Track (b): independently trained models#

3b/4b/5b. Two independent potentials → two densities#

xnn_indep = ForceStressOutput(xnn_indep_base, compute_forces=True,
                               compute_stress=True).to(DEVICE).double().eval()
def xnn_b():    return XNNCalculator(xnn_indep, cutoff=CUTOFF, device=DEVICE)

dE, dF, dS = compare_calcs(xnn_b, allegro_a)
print("track (b) calculators on one Argon config (INDEPENDENT models):")
print(f"  dE = {dE:.3e} eV | dF = {dF:.3e} eV/Å | dσ = {dS:.3e} eV/Å³ "
      "(training-level differences)")

rho_xb = run_npt(xnn_b, "xnn (b)")
rho_ab = rho_aa                                  # same trained original potential
db_x, db_a = rho_xb[N_EQUIL:].mean(), rho_ab[N_EQUIL:].mean()
print(f"track (b)  rho_xnn = {db_x:.4f} ± {rho_xb[N_EQUIL:].std():.3f}   "
      f"rho_allegro = {db_a:.4f} ± {rho_ab[N_EQUIL:].std():.3f}")
print(f"           |diff| = {abs(db_x-db_a):.2e} g/cm³ (within thermal fluctuations; exp ~{RHO_EXP})")

track (b) calculators on one Argon config (INDEPENDENT models):
  dE = 4.380e-02 eV | dF = 3.025e-02 eV/Å | dσ = 8.178e-05 eV/Å³ (training-level differences)
xnn (b): 1000 steps in 141 s | rho_eq = 1.4180 g/cm³
track (b)  rho_xnn = 1.4180 ± 0.012   rho_allegro = 1.4133 ± 0.011
           |diff| = 4.76e-03 g/cm³ (within thermal fluctuations; exp ~1.41)

/tmp/ipykernel_1044677/1079695994.py:52: DeprecationWarning: Use thermalize_momenta
  MaxwellBoltzmannDistribution(at, temperature_K=T_K, rng=np.random.default_rng(0)); Stationary(at)

6. Overview: both tracks#

t_ps = np.arange(N_EQUIL + N_PROD) * (DT / u.fs) / 1000.0
xc = N_EQUIL * (DT / u.fs) / 1000.0
fig, ax = plt.subplots(1, 2, figsize=(12, 4.4), sharey=True)
ax[0].plot(t_ps, rho_xa, label="xnn", lw=1)
ax[0].plot(t_ps, rho_aa, label="original Allegro", lw=1, ls="--")
ax[0].set_title(f"(a) same potential  |Δρ|={abs(da_x-da_a):.1e} g/cm³")
ax[1].plot(t_ps, rho_xb, label="xnn (independent)", lw=1)
ax[1].plot(t_ps, rho_ab, label="original Allegro (independent)", lw=1, ls="--")
ax[1].set_title(f"(b) independently trained  |Δρ|={abs(db_x-db_a):.1e} g/cm³")
for a in ax:
    a.axvline(xc, color="gray", ls=":", lw=1)
    a.axhline(RHO_EXP, color="k", ls="-.", lw=1, label=f"exp ≈ {RHO_EXP}")
    a.set_xlabel("time [ps]"); a.legend(fontsize=8)
ax[0].set_ylabel("density [g/cm³]")
plt.tight_layout(); plt.savefig("argon_density_md.png", dpi=120); plt.show()

print(f"{'':<20}{'xnn':>12}{'original Allegro':>18}{'|diff|':>12}")
print("-" * 62)
print(f"{'(a) same PES':<20}{da_x:>12.4f}{da_a:>18.4f}{abs(da_x-da_a):>12.1e}")
print(f"{'(b) independent':<20}{db_x:>12.4f}{db_a:>18.4f}{abs(db_x-db_a):>12.1e}")
print(f"{'experiment':<20}{RHO_EXP:>12.2f}")

../../../../_images/705cdc00336639e1abe78f536cd931e8a8e2a718841f32fe8fd8ac293544965c.png

                            xnn  original Allegro      |diff|
--------------------------------------------------------------
(a) same PES              1.4133            1.4133     4.8e-08
(b) independent           1.4180            1.4133     4.8e-03
experiment                  1.41

Summary#

Track (a): copying the trained original-Allegro weights into xnn gives identical energy/forces/stress and identical NPT densities; the xnn inference/MD path reproduces the original Allegro exactly.
Track (b): an independently trained xnn Allegro lands within the thermal fluctuations of the original, both near the experimental liquid-Ar density (~1.41 g/cm³).

Same conclusions, same pipeline, third model family: Allegro joins MACE and NequIP as a faithful, dependency-light member of the xnn equivariant-GNN family (shared EquivariantGNN base, featurizers, ForceStressOutput, ASE/LAMMPS deploy).

Argon density from MD: xnn vs the original Allegro#

0. Setup: train float32, run MD float64#