Argon density from MD: xnn vs the original MACE#

This notebook computes the mass density of liquid Argon by running NPT molecular dynamics through ASE with a trained MACE potential, and compares the xnn result against the original mace-torch model in two complementary ways:

  • Track (a): same potential. The trained original-MACE weights are copied into the xnn model, so both codes represent the identical potential-energy surface. Any density difference then reflects only the xnn-vs-mace inference / MD code path; it should be numerically zero. (Sections 3a/4a/5a.)

  • Track (b): independently trained. xnn is trained from scratch on the same data with no weight copying, giving two independent potentials. Now we compare the density as two practitioners would if each fit their own model. (Sections 3b/4b/5b.)

Pipeline for each track: train → wrap in an ASE calculator (energy + forces + stress) → run NPT MD (\(T=85\) K, \(P=1\) bar) → measure \(\rho=M/V\).

0. Setup#

Train in float32 (speed); run MD in float64 (smooth forces/stress, as in production MACE).

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]
import xnn, mace
print("xnn:", xnn.__version__, "| mace (original):", mace.__version__, "| device:", DEVICE)

xnn: 0.1.0 | mace (original): 0.3.16 | device: cuda

/D3/sina/xnn/.venv/lib/python3.13/site-packages/e3nn/o3/_wigner.py:10: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  _Jd, _W3j_flat, _W3j_indices = torch.load(os.path.join(os.path.dirname(__file__), 'constants.pt'))

1. Load data and build both data pipelines#

from xnn.common.data import AtomicDataset, load_dataset
from mace.data import AtomicData, Configuration
from mace.tools import AtomicNumberTable, torch_geometric

E0 = {18: 0.0}   # argon isolated-atom reference energy
train_structs = load_dataset("argon_md", split="train")
ZT = AtomicNumberTable(SPECIES)

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_mace(s):
    conf = Configuration(atomic_numbers=np.asarray(s["atomic_numbers"]), positions=s["pos"],
        properties={"energy": s["energy"], "forces": s["forces"]},
        property_weights={"energy": 1.0, "forces": 1.0}, cell=s["cell"], pbc=(True,)*3)
    return AtomicData.from_config(conf, z_table=ZT, cutoff=CUTOFF)
mace_train = [to_mace(s) for s in train_structs]

# one shared train/val split used by both models
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)}")

cuequivariance or cuequivariance_torch is not available. Cuequivariance acceleration will be disabled.
200 configs | E0={18: 0.0} | lambda=17.16 | train 180 / val 20

2. Train the two models#

Both use the same architecture (\(T=2\), \(\ell_{\max}=3\), \(L_{\max}=1\), \(\nu=3\), 32 channels), the same data/loss/optimiser/schedule and the same split. The original MACE is trained with a native loop; xnn with xnn.train.Trainer.

from e3nn import o3
import torch.nn.functional as Fn
import mace.modules as mm
from mace.modules.blocks import RealAgnosticResidualInteractionBlock as MIB
from xnn.common.config import from_dict
from xnn.common.train import Trainer
from torch.utils.data import Subset

HID = "32x0e+32x1o"
EW, FW, LR, WD, BS, EPOCHS = 1.0, 100.0, 0.01, 5e-7, 10, 60

# ---------- 2.1 train the ORIGINAL MACE (native loop) ----------
torch.manual_seed(0)
mace_model = mm.MACE(r_max=CUTOFF, num_bessel=8, num_polynomial_cutoff=5, max_ell=3,
    interaction_cls=MIB, interaction_cls_first=MIB, num_interactions=2, num_elements=1,
    hidden_irreps=o3.Irreps(HID), MLP_irreps=o3.Irreps("16x0e"),
    atomic_energies=np.array([E0[18]]), avg_num_neighbors=LAMBDA, atomic_numbers=SPECIES,
    correlation=3, gate=Fn.silu, radial_MLP=[64, 64, 64], radial_type="bessel",
    use_reduced_cg=False, apply_cutoff=True).to(DEVICE)

DL = torch_geometric.dataloader.DataLoader
tr_loader = DL([mace_train[i] for i in train_idx], batch_size=BS, shuffle=True)
va_loader = DL([mace_train[i] for i in val_idx],   batch_size=BS, shuffle=False)
opt = torch.optim.Adam(mace_model.parameters(), lr=LR, weight_decay=WD)
sched = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, patience=10)
def mace_loss(out, b):
    n = (b.ptr[1:] - b.ptr[:-1]).to(out["energy"].dtype)
    return EW * (((out["energy"] - b.energy) / n) ** 2).mean() + FW * ((out["forces"] - b.forces) ** 2).mean()

t0 = time.time()
for epoch in range(EPOCHS):
    mace_model.train()
    for b in tr_loader:
        b = b.to(DEVICE); out = mace_model(b.to_dict(), training=True, compute_force=True)
        loss = mace_loss(out, b); opt.zero_grad(); loss.backward(); opt.step()
    mace_model.eval(); vl = 0.0
    for b in va_loader:
        b = b.to(DEVICE); vl += float(mace_loss(mace_model(b.to_dict(), training=False, compute_force=True), b))
    sched.step(vl / len(va_loader))
print(f"original MACE trained {EPOCHS} epochs in {time.time()-t0:.0f} s (final val {vl/len(va_loader):.3e})")

# ---------- 2.2 train xnn INDEPENDENTLY (Trainer) ----------
# non-default flags only; everything else = stock MACE defaults (configs/model/mace.yaml)
core = from_dict({
    "model": {"name": "mace", "cutoff": CUTOFF, "species": SPECIES, "max_L": 1,
              "hidden_irreps": HID, "avg_num_neighbors": LAMBDA, "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     # independently trained xnn MACE

original MACE trained 60 epochs in 417 s (final val 5.952e-04)
epoch    0 | train loss 1.7824e-01 | val loss 4.0494e-02
epoch    1 | train loss 4.9534e-02 | val loss 1.9639e-02
epoch    2 | train loss 2.4203e-02 | val loss 8.6105e-03
epoch    3 | train loss 8.9203e-03 | val loss 3.8398e-03
epoch    4 | train loss 3.4036e-03 | val loss 1.6303e-03
epoch    5 | train loss 2.2902e-03 | val loss 1.2410e-03
epoch    6 | train loss 1.7293e-03 | val loss 9.7818e-04
epoch    7 | train loss 1.5235e-03 | val loss 1.1189e-03
epoch    8 | train loss 1.3902e-03 | val loss 1.0681e-03
epoch    9 | train loss 1.2514e-03 | val loss 9.6917e-04
epoch   10 | train loss 1.3658e-03 | val loss 1.3300e-03
epoch   11 | train loss 1.2830e-03 | val loss 1.2296e-03
epoch   12 | train loss 1.1869e-03 | val loss 1.5758e-03
epoch   13 | train loss 1.0534e-03 | val loss 7.0099e-04
epoch   14 | train loss 8.3061e-04 | val loss 7.5308e-04
epoch   15 | train loss 7.8896e-04 | val loss 1.2321e-03
epoch   16 | train loss 9.4133e-04 | val loss 1.0322e-03
epoch   17 | train loss 1.1433e-03 | val loss 9.0873e-04
epoch   18 | train loss 7.7776e-04 | val loss 6.3342e-04
epoch   19 | train loss 7.5670e-04 | val loss 1.2568e-03
epoch   20 | train loss 1.1542e-03 | val loss 1.3643e-03
epoch   21 | train loss 7.7342e-04 | val loss 6.4181e-04
epoch   22 | train loss 8.7842e-04 | val loss 9.1290e-04
epoch   23 | train loss 8.4342e-04 | val loss 1.3029e-03
epoch   24 | train loss 8.9835e-04 | val loss 6.6383e-04
epoch   25 | train loss 8.4394e-04 | val loss 6.9654e-04
epoch   26 | train loss 9.3775e-04 | val loss 1.1299e-03
epoch   27 | train loss 8.8414e-04 | val loss 5.5828e-04
epoch   28 | train loss 8.3435e-04 | val loss 7.7909e-04
epoch   29 | train loss 7.1105e-04 | val loss 5.3847e-04
epoch   30 | train loss 6.7055e-04 | val loss 7.1859e-04
epoch   31 | train loss 1.1826e-03 | val loss 8.8792e-04
epoch   32 | train loss 1.2959e-03 | val loss 1.1714e-03
epoch   33 | train loss 1.8242e-03 | val loss 1.0036e-03
epoch   34 | train loss 1.6937e-03 | val loss 9.4435e-04
epoch   35 | train loss 1.2377e-03 | val loss 1.0633e-03
epoch   36 | train loss 8.3443e-04 | val loss 9.6946e-04
epoch   37 | train loss 1.0155e-03 | val loss 8.6592e-04
epoch   38 | train loss 1.1977e-03 | val loss 7.3642e-04
epoch   39 | train loss 1.2071e-03 | val loss 7.3105e-04
epoch   40 | train loss 6.0991e-04 | val loss 4.9166e-04
epoch   41 | train loss 1.1405e-03 | val loss 1.1974e-03
epoch   42 | train loss 8.4249e-04 | val loss 1.1444e-03
epoch   43 | train loss 1.7142e-03 | val loss 1.2634e-03
epoch   44 | train loss 2.3817e-03 | val loss 1.3921e-03
epoch   45 | train loss 2.0804e-03 | val loss 1.3230e-03
epoch   46 | train loss 1.1204e-03 | val loss 5.5318e-04
epoch   47 | train loss 1.0169e-03 | val loss 6.6129e-04
epoch   48 | train loss 9.7584e-04 | val loss 1.4237e-03
epoch   49 | train loss 2.3880e-03 | val loss 1.9722e-03
epoch   50 | train loss 2.7580e-03 | val loss 7.9063e-04
epoch   51 | train loss 1.0216e-03 | val loss 8.0398e-04
epoch   52 | train loss 7.0373e-04 | val loss 4.3472e-04
epoch   53 | train loss 4.9870e-04 | val loss 4.7063e-04
epoch   54 | train loss 4.5946e-04 | val loss 4.2057e-04
epoch   55 | train loss 4.4675e-04 | val loss 4.1498e-04
epoch   56 | train loss 4.4809e-04 | val loss 4.5636e-04
epoch   57 | train loss 4.4194e-04 | val loss 4.3432e-04
epoch   58 | train loss 4.4035e-04 | val loss 4.5914e-04
epoch   59 | train loss 4.3256e-04 | val loss 4.1203e-04
xnn (independent) trained 60 epochs in 300 s

/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(

Common MD utilities#

Switch to float64 for the dynamics and define the ASE calculators, the NPT driver and the density helper used by both tracks.

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

mace_model = mace_model.double().eval()
xnn_indep_base = xnn_indep_base.double().eval()

class MACEASECalculator(Calculator):
    '''Minimal ASE calculator wrapping an in-memory original-MACE 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)
        conf = Configuration(atomic_numbers=np.asarray(atoms.get_atomic_numbers()),
            positions=atoms.get_positions(), properties={}, property_weights={},
            cell=np.asarray(atoms.get_cell()), pbc=tuple(atoms.pbc))
        ad = AtomicData.from_config(conf, z_table=ZT, cutoff=self.cutoff)
        b = next(iter(DL([ad], batch_size=1))).to(self.device)
        out = self.model(b.to_dict(), training=False, compute_force=True, compute_stress=True)
        self.results["energy"] = float(out["energy"].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]                     # dense initial configuration (400 atoms)

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

def compare_calcs(make_x, make_m):
    at = Atoms(numbers=a0["atomic_numbers"], positions=a0["pos"], cell=a0["cell"], pbc=True)
    ax = at.copy(); ax.calc = make_x(); am = at.copy(); am.calc = make_m()
    return (abs(ax.get_potential_energy() - am.get_potential_energy()),
            np.abs(ax.get_forces() - am.get_forces()).max(),
            np.abs(ax.get_stress() - am.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); temp = np.empty_like(rho)
    t0 = time.time()
    for k in range(N_EQUIL + N_PROD):
        dyn.run(1); rho[k] = density(at); temp[k] = at.get_temperature()
    print(f"{label}: {N_EQUIL+N_PROD} steps in {time.time()-t0:.0f} s | rho_eq = {rho[N_EQUIL:].mean():.4f} g/cm³")
    return rho, temp

rho0 = density(Atoms(numbers=a0["atomic_numbers"], positions=a0["pos"], cell=a0["cell"], pbc=True))
RHO_EXP = 1.41
print(f"initial density = {rho0:.4f} g/cm³ | 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 MACE → xnn)#

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

Every learnable weight of the trained original MACE, and the reference energy \(E_0\), is copied into a fresh xnn MACE, so both codes carry the identical potential. We then confirm the two ASE calculators return the same energy, forces and stress to machine precision.

def copy_mace_into_xnn(xbase, mmod, T=2, corr=3):
    with torch.no_grad():
        xbase.node_embedding.load_state_dict(mmod.node_embedding.linear.state_dict())
        for i in range(T):
            xbase.interactions[i].load_state_dict(mmod.interactions[i].state_dict())
            xsc = xbase.products[i].symmetric_contractions
            msc = mmod.products[i].symmetric_contractions
            for c in range(len(xsc.contractions)):
                xc, mc = xsc.contractions[c], msc.contractions[c]
                xc.weights[corr-1].copy_(mc.weights_max)
                for nu in range(1, corr):
                    xc.weights[nu-1].copy_(mc.weights[corr-1-nu])
            xbase.products[i].linear.load_state_dict(mmod.products[i].linear.state_dict())
            xr, mr = xbase.readouts[i], mmod.readouts[i]
            if "NonLinear" in type(mr).__name__:
                xr.linear_1.load_state_dict(mr.linear_1.state_dict())
                xr.linear_2.load_state_dict(mr.linear_2.state_dict())
            else:
                xr.linear.load_state_dict(mr.linear.state_dict())
        xbase.atom_ref.weight[SPECIES[0]] = float(mmod.atomic_energies_fn.atomic_energies[0])

xnn_shared_base = build_model(core.model)          # fresh xnn model (float64)
copy_mace_into_xnn(xnn_shared_base, mace_model)
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 mace_a():  return MACEASECalculator(mace_model, CUTOFF, DEVICE)

dE, dF, dS = compare_calcs(xnn_a, mace_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 = 6.30e-06 eV | dF = 3.00e-07 eV/Å | dσ = 3.27e-09 eV/ų  -> identical

/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
/D3/sina/xnn/.venv/lib/python3.13/site-packages/torch/jit/_check.py:178: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(

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

Same initial positions and velocities; the only difference is the calculator.

rho_xa, T_xa = run_npt(xnn_a, "xnn (a)")
rho_ma, T_ma = run_npt(mace_a, "mace (a)")

xnn (a): 1000 steps in 168 s | rho_eq = 1.4293 g/cm³
mace (a): 1000 steps in 168 s | rho_eq = 1.4293 g/cm³

/tmp/ipykernel_1051448/3511752848.py:48: DeprecationWarning: Use thermalize_momenta
  MaxwellBoltzmannDistribution(at, temperature_K=T_K, rng=np.random.default_rng(0)); Stationary(at)

5a. Result (a): the densities are identical#

Because both calculators evaluate the same PES, the equilibrium densities agree to numerical noise; the trajectories overlap until chaotic float divergence.

da_x, da_m = rho_xa[N_EQUIL:].mean(), rho_ma[N_EQUIL:].mean()
print(f"track (a)  rho_xnn = {da_x:.4f}   rho_mace = {da_m:.4f}   |diff| = {abs(da_x-da_m):.2e} g/cm³")

track (a)  rho_xnn = 1.4293   rho_mace = 1.4293   |diff| = 4.85e-08 g/cm³

Track (b): independently trained models (no weight copying)#

3b. Two independent potentials#

Here xnn is the model trained from scratch in Section 2.2 (never copied from MACE), and MACE is its independently trained counterpart. The two calculators now differ at the level of independent training (small, not machine precision).

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)
def mace_b():  return MACEASECalculator(mace_model, CUTOFF, DEVICE)

dE, dF, dS = compare_calcs(xnn_b, mace_b)
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)")

track (b) calculators on one Argon config (INDEPENDENT models):
  dE = 3.966e+00 eV | dF = 1.885e-02 eV/Å | dσ = 9.862e-05 eV/ų  (training-level differences)

4b. NPT MD: two independent potentials#

rho_xb, T_xb = run_npt(xnn_b, "xnn (b)")
rho_mb, T_mb = run_npt(mace_b, "mace (b)")

xnn (b): 1000 steps in 200 s | rho_eq = 1.4463 g/cm³
mace (b): 1000 steps in 134 s | rho_eq = 1.4293 g/cm³

/tmp/ipykernel_1051448/3511752848.py:48: DeprecationWarning: Use thermalize_momenta
  MaxwellBoltzmannDistribution(at, temperature_K=T_K, rng=np.random.default_rng(0)); Stationary(at)

5b. Result (b): two independent density predictions#

db_x, db_m = rho_xb[N_EQUIL:].mean(), rho_mb[N_EQUIL:].mean()
sb_x, sb_m = rho_xb[N_EQUIL:].std(), rho_mb[N_EQUIL:].std()
print(f"track (b)  rho_xnn = {db_x:.4f} ± {sb_x:.3f}   rho_mace = {db_m:.4f} ± {sb_m:.3f}")
print(f"           |diff| = {abs(db_x-db_m):.2e} g/cm³  (within thermal fluctuations; exp ~{RHO_EXP})")

track (b)  rho_xnn = 1.4463 ± 0.005   rho_mace = 1.4293 ± 0.010
           |diff| = 1.70e-02 g/cm³  (within thermal fluctuations; exp ~1.41)

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_ma, label="original MACE", lw=1, ls="--")
ax[0].set_title(f"(a) same potential  |Δρ|={abs(da_x-da_m):.1e} g/cm³")
ax[1].plot(t_ps, rho_xb, label="xnn (independent)", lw=1)
ax[1].plot(t_ps, rho_mb, label="original MACE (independent)", lw=1, ls="--")
ax[1].set_title(f"(b) independently trained  |Δρ|={abs(db_x-db_m):.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 MACE':>16}{'|diff|':>12}")
print("-" * 60)
print(f"{'(a) same PES':<20}{da_x:>12.4f}{da_m:>16.4f}{abs(da_x-da_m):>12.1e}")
print(f"{'(b) independent':<20}{db_x:>12.4f}{db_m:>16.4f}{abs(db_x-db_m):>12.1e}")
print(f"{'experiment':<20}{RHO_EXP:>12.2f}")
../../../../_images/99abffad3ba105089c7bdab520d7f01b95fa37fd3dde6e283a33d5baff518f22.png 
                            xnn   original MACE      |diff|
------------------------------------------------------------
(a) same PES              1.4293          1.4293     4.9e-08
(b) independent           1.4463          1.4293     1.7e-02
experiment                  1.41

Summary#

Computing the Argon density from ASE NPT MD two ways:

  • Track (a): same potential. Copying the trained original-MACE weights into xnn makes the two calculators return identical energy/forces/stress (machine precision), and the NPT densities are identical (\(|\Delta\rho|\sim10^{-8}\) g/cm³). This isolates and confirms the xnn inference/MD path reproduces the original MACE exactly.

  • Track (b): independently trained. Training xnn from scratch (no copying) gives an independent potential; its density agrees with the independently trained original MACE to within the thermal fluctuations, and both sit near the experimental liquid-Ar density (~1.41 g/cm³). This is the realistic “two practitioners, two fits” agreement.

Together: xnn is not only bit-for-bit equivalent to the original MACE for a fixed PES (a), but as a modelling tool it produces the same physical property when trained independently (b).

Notes. Berendsen barostat for simplicity (use ase.md.npt.NPT for rigorous ensembles); longer runs tighten the estimate; stress is the autograd virial of the energy+force-trained model.