ANI-1, block by block: reproducing the original implementation with `xnn`#

ANI-1 (Smith, Isayev & Roitberg, Chem. Sci. 8, 3192, 2017) is a High-Dimensional Neural Network Potential whose per-atom descriptor is the Atomic Environment Vector (AEV): element-resolved radial and angular symmetry functions (paper eqns 2–4). One neural network per element maps the AEV to an atomic energy; the atomic energies (plus per-element self energies) sum to the total energy.

The reference implementation is aiqm/torchani (the PyTorch ANI the paper’s authors maintain). This notebook shows, block by block, that xnn reproduces torchani element-for-element:

block	what we check
cutoff	`xnn` `CosineCutoff` vs `torchani.aev.cutoff_cosine`
AEV (ANI-1x, 384)	`AEV.ani1x()` vs `torchani.AEVComputer` with the ANI-1x grid
AEV (ANI-1, 768)	`AEV.ani1()` vs `torchani.AEVComputer` with the paper’s grid
element nets + E, F	transplant torchani’s pretrained ANI-1x weights → energies & forces
ensemble	8-model ANI-1x ensemble mean

Everything runs in float64. Run this notebook with the xnn-ani kernel (torch + torchani==2.2.4 + xnn); see pyproject.toml’s [ani] extra.

Conventions note. torchani follows NeuroChem, not the paper’s eqn 3 literally: it multiplies the radial term by 0.25 and scales cos(θ) by 0.95 inside acos (to keep the gradient finite at ±1). xnn bakes both in by default (radial_prefactor=0.25, angular_cos_factor=0.95), so the AEVs match exactly. Set them to 1.0 for the literal Behler-Parrinello form.

0. Setup#

import warnings; warnings.filterwarnings("ignore")
import math
import numpy as np
import torch
import torchani

torch.set_default_dtype(torch.float64)

from xnn.dnn.models.ani import ANI
from xnn.dnn.featurizers import AEV
from xnn.dnn.featurizers.aev import _even_shifts, _angle_shifts
from xnn.common.featurizers import CosineCutoff
from xnn.common.data import AtomicGraph
from xnn.common.data.neighborlist import build_neighbor_list
from xnn.common.models.outputs import ForceStressOutput

print("torch", torch.__version__, "| torchani", torchani.__version__)

# H, C, N, O -> torchani's internal species index
IDX = {1: 0, 6: 1, 7: 2, 8: 3}


def xnn_graph(Z, pos, cutoff, requires_grad=False):
    "Build a single-structure AtomicGraph from atomic numbers + positions."
    Z = torch.as_tensor(Z, dtype=torch.long)
    pos = torch.as_tensor(pos, dtype=torch.float64)
    if requires_grad:
        pos = pos.clone().requires_grad_(True)
    ei, cs = build_neighbor_list(pos, cutoff)
    return AtomicGraph(pos=pos, atomic_numbers=Z, edge_index=ei, cell_shifts=cs,
                       batch=torch.zeros(len(Z), dtype=torch.long),
                       n_atoms=torch.tensor([len(Z)]))


def torchani_aev(Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ, ns=4):
    "A torchani AEVComputer built from explicit constants."
    t = lambda x: torch.tensor(x, dtype=torch.float64)
    return torchani.AEVComputer(Rcr, Rca, t(EtaR), t(ShfR), t(EtaA), t(Zeta),
                                t(ShfA), t(ShfZ), ns)

torch 2.5.1+cpu | torchani 2.2.4

A handful of H/C/N/O molecules#

Small, distinct organic geometries (water, ammonia, methanol, formamide) with a little random distortion so no symmetry hides a bug. We compare xnn and torchani on every one of them.

rng = np.random.default_rng(0)

_MOLS = {
    "water":     ([8, 1, 1],
                  [[0, 0, 0], [0.76, 0.59, 0], [-0.76, 0.59, 0]]),
    "ammonia":   ([7, 1, 1, 1],
                  [[0, 0, 0.12], [0, 0.94, -0.27], [0.81, -0.47, -0.27],
                   [-0.81, -0.47, -0.27]]),
    "methanol":  ([6, 8, 1, 1, 1, 1],
                  [[-0.05, 0.66, 0], [-0.05, -0.75, 0], [1.0, 1.0, 0],
                   [-0.56, 1.05, 0.88], [-0.56, 1.05, -0.88], [0.85, -1.09, 0]]),
    "formamide": ([6, 7, 8, 1, 1, 1],
                  [[0, 0.10, 0], [1.27, 0.55, 0], [-0.98, 0.82, 0],
                   [-0.20, -0.95, 0], [1.42, 1.55, 0], [2.05, -0.09, 0]]),
}
MOLS = {name: (np.array(Z, dtype=np.int64),
               np.array(pos, dtype=np.float64) + 0.03 * rng.standard_normal((len(Z), 3)))
        for name, (Z, pos) in _MOLS.items()}
print("molecules:", ", ".join(MOLS))

molecules: water, ammonia, methanol, formamide

1. Cutoff function (paper eqn 2)#

\(f_C(r) = \tfrac12\cos(\pi r / R_C) + \tfrac12\) for \(r \le R_C\), else 0.

r = torch.linspace(0.05, 5.2, 200, dtype=torch.float64)
xn = CosineCutoff(5.2)(r)
ta = torchani.aev.cutoff_cosine(r, 5.2) * (r < 5.2)
print(f"max|Δ| cutoff = {(xn - ta).abs().max():.2e}")

max|Δ| cutoff = 2.22e-16

2. AEV: the ANI-1x grid (384 elements)#

AEV.ani1x() uses the exact constants torchani ships for ANI-1x (radial cutoff 5.2 Å, 16 radial shifts; angular cutoff 3.5 Å, 4 radial × 8 angular shifts). We build a torchani.AEVComputer from the same constants and compare the full 384-vector per atom, and separately the radial (0:64) and angular (64:384) blocks.

aev_x = AEV.ani1x()
tani = torchani_aev(5.2, 3.5, [16.0], _even_shifts(5.2, 16),
                    [8.0], [32.0], _even_shifts(3.5, 4), _angle_shifts(8))

print(f"{'molecule':<10} {'max|Δ| full':>14} {'max|Δ| radial':>16} {'max|Δ| angular':>16}")
worst = 0.0
for name, (Z, pos) in MOLS.items():
    g = xnn_graph(Z, pos, aev_x.cutoff)
    x = aev_x(g)
    sp = torch.tensor([[IDX[z] for z in Z]])
    _, t = tani((sp, torch.as_tensor(pos[None])))
    t = t[0]
    d_full = (x - t).abs().max().item()
    d_rad = (x[:, :64] - t[:, :64]).abs().max().item()
    d_ang = (x[:, 64:] - t[:, 64:]).abs().max().item()
    worst = max(worst, d_full)
    print(f"{name:<10} {d_full:>14.2e} {d_rad:>16.2e} {d_ang:>16.2e}")
print(f"\nworst over all molecules: {worst:.2e}  (AEV length {aev_x.output_dim})")
assert worst < 1e-10

molecule      max|Δ| full    max|Δ| radial   max|Δ| angular
water            2.78e-17         2.78e-17         0.00e+00
ammonia          1.11e-16         5.55e-17         1.11e-16
methanol         2.50e-16         8.33e-17         2.50e-16
formamide        2.78e-16         2.78e-17         2.78e-16

worst over all molecules: 2.78e-16  (AEV length 384)

3. AEV: the original ANI-1 grid (768 elements)#

AEV.ani1() is the paper’s parameterisation: radial cutoff 4.6 Å with 32 shifts, angular cutoff 3.1 Å with 8 × 8 shifts → a 768-vector for H, C, N, O (paper §3.4). Same check, against a torchani computer built with these constants.

aev_1 = AEV.ani1()
tani1 = torchani_aev(4.6, 3.1, [16.0], _even_shifts(4.6, 32),
                     [8.0], [8.0], _even_shifts(3.1, 8), _angle_shifts(8))

worst = 0.0
for name, (Z, pos) in MOLS.items():
    g = xnn_graph(Z, pos, aev_1.cutoff)
    x = aev_1(g)
    sp = torch.tensor([[IDX[z] for z in Z]])
    _, t = tani1((sp, torch.as_tensor(pos[None])))
    d = (x - t[0]).abs().max().item()
    worst = max(worst, d)
    print(f"{name:<10} max|Δ| = {d:.2e}")
print(f"\nworst: {worst:.2e}  (AEV length {aev_1.output_dim})")
assert worst < 1e-10

water      max|Δ| = 3.33e-16
ammonia    max|Δ| = 8.88e-16
methanol   max|Δ| = 8.88e-16
formamide  max|Δ| = 6.66e-16

worst: 8.88e-16  (AEV length 768)

4. Element networks → energies and forces (pretrained weights)#

Now the whole model. We load torchani’s pretrained ANI-1x (member 0 of its 8-model ensemble), transplant its per-element Linear weights into an xnn ANI.ani1x(), and compare total energies and (autograd) forces. The per-element architectures (H 160:128:96, C 144:112:96, N/O 128:112:96), the CELU(0.1) activation and the per-element self energies all line up, so the transplant is one-to-one.

model = torchani.models.ANI1x(periodic_table_index=False).double()
member0 = model.neural_networks[0]
zsym = {1: "H", 6: "C", 7: "N", 8: "O"}


def transplant(member, xa):
    "Copy a torchani ANIModel's per-element Linear weights into an xnn ANI."
    nets = dict(member.named_children())
    for z in [1, 6, 7, 8]:
        src = [l for l in nets[zsym[z]] if isinstance(l, torch.nn.Linear)]
        dst = [l for l in xa.element_nets.nets[str(z)] if isinstance(l, torch.nn.Linear)]
        for s, d in zip(src, dst):
            d.weight.data = s.weight.data.clone()
            d.bias.data = s.bias.data.clone()


xa = ANI.ani1x()          # self energies default to torchani's (Hartree)
transplant(member0, xa)
wrapped = ForceStressOutput(xa)

print(f"{'molecule':<10} {'E xnn (Ha)':>16} {'E torchani':>16} {'|ΔE|':>10} {'max|ΔF|':>10}")
wE = wF = 0.0
for name, (Z, pos) in MOLS.items():
    g = xnn_graph(Z, pos, xa.cutoff, requires_grad=True)
    out = wrapped(g)
    E_x = out["energy"].item()
    F_x = out["forces"].detach().numpy()

    coords = torch.as_tensor(pos[None]).clone().requires_grad_(True)
    sp = torch.tensor([[IDX[z] for z in Z]])
    e = model.energy_shifter(member0(model.aev_computer((sp, coords)))).energies
    F_t = -torch.autograd.grad(e.sum(), coords)[0][0].numpy()
    E_t = e.item()

    dE, dF = abs(E_x - E_t), np.abs(F_x - F_t).max()
    wE, wF = max(wE, dE), max(wF, dF)
    print(f"{name:<10} {E_x:>16.8f} {E_t:>16.8f} {dE:>10.2e} {dF:>10.2e}")
print(f"\nworst |ΔE| = {wE:.2e} Ha   worst max|ΔF| = {wF:.2e} Ha/Å")
assert wE < 1e-6 and wF < 1e-6

/D3/sina/xnn/.venv-ani/lib/python3.13/site-packages/torchani/resources/
molecule        E xnn (Ha)       E torchani       |ΔE|    max|ΔF|
water          -76.38793938     -76.38793938   9.62e-10   2.03e-08
ammonia        -56.52444686     -56.52444686   1.48e-09   3.75e-09
methanol      -115.67180357    -115.67180357   7.47e-10   1.45e-08
formamide     -169.81418123    -169.81418123   2.51e-10   2.20e-08

worst |ΔE| = 1.48e-09 Ha   worst max|ΔF| = 2.20e-08 Ha/Å

5. The full 8-model ANI-1x ensemble#

The released ANI-1x is the mean of 8 networks. We transplant all 8 into 8 xnn models and average, then compare to torchani’s built-in ensemble output (which also adds the self energies).

members = [ANI.ani1x() for _ in range(len(model.neural_networks))]
for xa_k, m_k in zip(members, model.neural_networks):
    transplant(m_k, xa_k)

wrapped_members = [ForceStressOutput(m) for m in members]

print(f"{'molecule':<10} {'E xnn mean':>16} {'E torchani':>16} {'|ΔE|':>10}")
worst = 0.0
for name, (Z, pos) in MOLS.items():
    g = xnn_graph(Z, pos, members[0].cutoff)
    E_x = float(np.mean([w(g)["energy"].item() for w in wrapped_members]))

    coords = torch.as_tensor(pos[None])
    sp = torch.tensor([[IDX[z] for z in Z]])
    E_t = model((sp, coords)).energies.item()
    dE = abs(E_x - E_t)
    worst = max(worst, dE)
    print(f"{name:<10} {E_x:>16.8f} {E_t:>16.8f} {dE:>10.2e}")
print(f"\nworst |ΔE| = {worst:.2e} Ha")
assert worst < 1e-6

molecule        E xnn mean       E torchani       |ΔE|
water          -76.38818940     -76.38818940   3.61e-11
ammonia        -56.52408408     -56.52408408   1.96e-09
methanol      -115.67195817    -115.67195817   9.51e-10
formamide     -169.81364721    -169.81364721   7.59e-10

worst |ΔE| = 1.96e-09 Ha

6. ANI-1ccx: the same architecture, coupled-cluster weights#

The ANI-1ccx potential (Smith et al., Nat. Commun. 10, 2903, 2019) is the ANI-1x architecture transfer-learned to CCSD(T)*/CBS coupled-cluster data; in xnn, ANI.ani1ccx() simply reuses ANI.ani1x() with the coupled-cluster self energies. The very same transplant therefore moves torchani’s pretrained ANI-1ccx across one-to-one: first a single member, then the full 8-network ensemble the potential is released as (paper SI S1.2.4).

model_ccx = torchani.models.ANI1ccx(periodic_table_index=False).double()
xccx = ANI.ani1ccx()      # ANI-1x architecture + CCSD(T)*/CBS self energies
transplant(model_ccx.neural_networks[0], xccx)
wrapped_ccx = ForceStressOutput(xccx)

print(f"{'molecule':<10} {'E xnn (Ha)':>16} {'E torchani':>16} {'|\u0394E|':>10} {'max|\u0394F|':>10}")
wE = wF = 0.0
for name, (Z, pos) in MOLS.items():
    g = xnn_graph(Z, pos, xccx.cutoff, requires_grad=True)
    out = wrapped_ccx(g)
    E_x, F_x = out["energy"].item(), out["forces"].detach().numpy()

    coords = torch.as_tensor(pos[None]).clone().requires_grad_(True)
    sp = torch.tensor([[IDX[z] for z in Z]])
    e = model_ccx.energy_shifter(
        model_ccx.neural_networks[0](model_ccx.aev_computer((sp, coords)))).energies
    F_t = -torch.autograd.grad(e.sum(), coords)[0][0].numpy()

    dE, dF = abs(E_x - e.item()), np.abs(F_x - F_t).max()
    wE, wF = max(wE, dE), max(wF, dF)
    print(f"{name:<10} {E_x:>16.8f} {e.item():>16.8f} {dE:>10.2e} {dF:>10.2e}")
print(f"\nworst |\u0394E| = {wE:.2e} Ha   worst max|\u0394F| = {wF:.2e} Ha/\u00c5")
assert wE < 1e-6 and wF < 1e-6

# the released ANI-1ccx, like ANI-1x, is the MEAN of 8 networks (paper SI S1.2.4)
members_ccx = [ANI.ani1ccx() for _ in range(len(model_ccx.neural_networks))]
for xk, mk in zip(members_ccx, model_ccx.neural_networks):
    transplant(mk, xk)
wrapped_members_ccx = [ForceStressOutput(m) for m in members_ccx]
worst = 0.0
for name, (Z, pos) in MOLS.items():
    g = xnn_graph(Z, pos, xccx.cutoff)
    E_x = float(np.mean([w(g)["energy"].item() for w in wrapped_members_ccx]))
    coords = torch.as_tensor(pos[None])
    sp = torch.tensor([[IDX[z] for z in Z]])
    E_t = model_ccx((sp, coords)).energies.item()
    worst = max(worst, abs(E_x - E_t))
print(f"\n8-model ANI-1ccx ensemble: worst |\u0394E| = {worst:.2e} Ha")
assert worst < 1e-6

/D3/sina/xnn/.venv-ani/lib/python3.13/site-packages/torchani/resources/
molecule        E xnn (Ha)       E torchani       |ΔE|    max|ΔF|
water          -76.38344647     -76.38344647   5.19e-10   1.01e-08
ammonia        -56.50648961     -56.50648961   1.04e-09   1.53e-08
methanol      -115.61587813    -115.61587813   5.72e-10   2.03e-08
formamide     -169.71677096    -169.71677096   1.91e-10   1.87e-08

worst |ΔE| = 1.04e-09 Ha   worst max|ΔF| = 2.03e-08 Ha/Å

8-model ANI-1ccx ensemble: worst |ΔE| = 1.35e-09 Ha

7. ANI-2x: seven elements (adds S, F, Cl)#

ANI-2x (Devereux et al., J. Chem. Theory Comput. 16, 4192, 2020) extends ANI to seven elements. Its AEV is larger (1008 elements: radial cutoff 5.1 A, angular cutoff 3.5 A, shift grids starting at 0.8 A) and the per-element networks are wider. ANI.ani2x() builds it over the torchani element order (H, C, N, O, S, F, Cl); transplanting torchani’s pretrained ANI-2x weights reproduces its energies and forces on molecules containing the new elements.

model2x = torchani.models.ANI2x(periodic_table_index=False).double()
x2x = ANI.ani2x()                                  # 7 elements, 1008-length AEV
print("AEV length:", x2x.featurizer.output_dim, " cutoff:", x2x.cutoff)

# seven-element index and symbol maps (torchani's order for ANI-2x)
IDX2 = {z: i for i, z in enumerate(x2x.species)}
zsym2 = {1: "H", 6: "C", 7: "N", 8: "O", 16: "S", 9: "F", 17: "Cl"}


def transplant2x(member, xa):
    "Copy a torchani ANI-2x member's per-element Linear weights into an xnn ANI."
    nets = dict(member.named_children())
    for z in x2x.species:
        src = [l for l in nets[zsym2[z]] if isinstance(l, torch.nn.Linear)]
        dst = [l for l in xa.element_nets.nets[str(z)] if isinstance(l, torch.nn.Linear)]
        for s, d in zip(src, dst):
            d.weight.data = s.weight.data.clone()
            d.bias.data = s.bias.data.clone()


transplant2x(model2x.neural_networks[0], x2x)
wrapped2x = ForceStressOutput(x2x)

# small molecules that exercise S, F, and Cl
rng2 = np.random.default_rng(2)
MOLS2X = {
    "H2S":            ([16, 1, 1],
                       [[0, 0, 0], [0.96, 0.94, 0], [-0.96, 0.94, 0]]),
    "CH3F":           ([6, 9, 1, 1, 1],
                       [[0, 0, 0], [0, 0, 1.38], [1.03, 0, -0.36],
                        [-0.51, 0.89, -0.36], [-0.51, -0.89, -0.36]]),
    "CH3Cl":          ([6, 17, 1, 1, 1],
                       [[0, 0, 0], [0, 0, 1.78], [1.03, 0, -0.36],
                        [-0.51, 0.89, -0.36], [-0.51, -0.89, -0.36]]),
    "methanethiol":   ([6, 16, 1, 1, 1, 1],
                       [[0, 0, 0], [1.42, 0.72, 0], [-0.55, 0.36, 0.89],
                        [-0.55, 0.36, -0.89], [-0.10, -1.09, 0], [1.28, 2.00, 0]]),
}
MOLS2X = {name: (np.array(Z, dtype=np.int64),
                 np.array(pos, dtype=np.float64)
                 + 0.03 * rng2.standard_normal((len(Z), 3)))
          for name, (Z, pos) in MOLS2X.items()}

print(f"{'molecule':<14} {'E xnn (Ha)':>16} {'E torchani':>16} {'|dE|':>10} {'max|dF|':>10}")
wE = wF = 0.0
for name, (Z, pos) in MOLS2X.items():
    g = xnn_graph(Z, pos, x2x.cutoff, requires_grad=True)
    out = wrapped2x(g)
    E_x, F_x = out["energy"].item(), out["forces"].detach().numpy()

    coords = torch.as_tensor(pos[None]).clone().requires_grad_(True)
    sp = torch.tensor([[IDX2[int(z)] for z in Z]])
    e = model2x.energy_shifter(
        model2x.neural_networks[0](model2x.aev_computer((sp, coords)))).energies
    F_t = -torch.autograd.grad(e.sum(), coords)[0][0].numpy()

    dE, dF = abs(E_x - e.item()), np.abs(F_x - F_t).max()
    wE, wF = max(wE, dE), max(wF, dF)
    print(f"{name:<14} {E_x:>16.8f} {e.item():>16.8f} {dE:>10.2e} {dF:>10.2e}")
print(f"\nworst |dE| = {wE:.2e} Ha   worst max|dF| = {wF:.2e} Ha/A")
assert wE < 1e-5 and wF < 1e-5

# and the full 8-model ANI-2x ensemble (mean energy)
members2x = [ANI.ani2x() for _ in model2x.neural_networks]
for xk, mk in zip(members2x, model2x.neural_networks):
    transplant2x(mk, xk)
wrapped_members2x = [ForceStressOutput(m) for m in members2x]
worst = 0.0
for name, (Z, pos) in MOLS2X.items():
    g = xnn_graph(Z, pos, x2x.cutoff)
    E_x = float(np.mean([w(g)["energy"].item() for w in wrapped_members2x]))
    coords = torch.as_tensor(pos[None])
    sp = torch.tensor([[IDX2[int(z)] for z in Z]])
    E_t = model2x((sp, coords)).energies.item()
    worst = max(worst, abs(E_x - E_t))
print(f"8-model ANI-2x ensemble: worst |dE| = {worst:.2e} Ha")
assert worst < 1e-5

/D3/sina/xnn/.venv-ani/lib/python3.13/site-packages/torchani/resources/
AEV length: 1008  cutoff: 5.1
molecule            E xnn (Ha)       E torchani       |dE|    max|dF|
H2S               -399.35810752    -399.35810752   3.17e-10   4.09e-09
CH3F              -139.69268288    -139.69268288   2.07e-09   4.11e-09
CH3Cl             -500.06826174    -500.06826174   9.95e-10   1.30e-09
methanethiol      -438.61078843    -438.61078843   4.19e-10   3.81e-09

worst |dE| = 2.07e-09 Ha   worst max|dF| = 4.11e-09 Ha/A
8-model ANI-2x ensemble: worst |dE| = 8.59e-10 Ha

Summary#

Block by block (cutoff, the ANI-1x (384) and ANI-1 (768) AEVs, the per-element networks with pretrained weights, the 8-model ensemble, the ANI-1ccx transplant, and the seven-element ANI-2x transplant), xnn reproduces torchani to numerical precision (float64 round-off), for both energies and forces. The xnn ANI is a faithful, from-scratch re-implementation of the ANI method; the same AEV/ANI code trains from scratch on the ANI datasets (see examples/dnn/ani/).

ANI-1, block by block: reproducing the original implementation with xnn#