The ANI-2x dataset: seven elements (S, F, Cl) and the `ani-2x` preset#

The ANI-2x dataset (Devereux et al., J. Chem. Theory Comput. 16, 4192, 2020) is the training set behind the ANI-2x potential, the seven-element extension of ANI. Where ANI-1/ANI-1x/ANI-1ccx cover only H, C, N, O, ANI-2x adds S, F, and Cl, the elements that (with the original four) make up ~90% of drug-like molecules. It was built by the same active learning procedure as ANI-1x, giving ~9.6 million wB97X/6-31G(d) conformations with energies and forces.

Same one-line hub as every other dataset. The first call downloads the ~3.7 GB wB97X/6-31G(d) archive from Zenodo (record 10108942), cached and MD5-verified under datasets/ani2x/:

from xnn.common.data import load_dataset
data = load_dataset("ani2x", n_atoms=[5, 6, 7])   # atom-count groups

The ANI-2x HDF5 groups conformations by total number of atoms ("002", "003", …), and within a group each conformation can be a different molecule with that atom count. Here we load a capped subset that contains the new elements, train the xnn ani-2x preset on energies and forces, and look at both parities plus a per-element breakdown. This is the seven-element companion to ani1x_dataset.ipynb.

Run with the xnn kernel.

0. Load a subset of ANI-2x, including S/F/Cl#

n_atoms picks which atom-count groups to read; max_conformations caps each group so the demo trains in minutes. Energies convert to eV and forces to eV/A. Note the seven elements and the per-atom forces, absent from ANI-1.

import time, warnings
warnings.filterwarnings("ignore")
import numpy as np
import torch
import matplotlib.pyplot as plt

torch.set_default_dtype(torch.float64)
torch.manual_seed(0)

from xnn.common.data import load_dataset

# a handful of atom-count groups, capped per group; wB97X/6-31G* energies+forces
data = load_dataset("ani2x", n_atoms=[5, 6, 7, 8], max_conformations=1500,
                    units="eV")["all"]
print(f"loaded {len(data):,} conformations")

Z_ALL = {1: "H", 6: "C", 7: "N", 8: "O", 16: "S", 9: "F", 17: "Cl"}
counts = {z: 0 for z in Z_ALL}
for s in data:
    for z in s["atomic_numbers"]:
        counts[int(z)] += 1
print("atom counts by element:",
      {Z_ALL[z]: counts[z] for z in Z_ALL})
n_heavy_halogen = sum(any(int(z) in (16, 9, 17) for z in s["atomic_numbers"])
                      for s in data)
print(f"conformations containing S/F/Cl: {n_heavy_halogen:,} / {len(data):,}")
print(f"keys per structure: {sorted(data[0])}   # note 'forces'")
print(f"force array shape (atoms, 3): {data[0]['forces'].shape}")

loaded 6,000 conformations
atom counts by element: {'H': 14577, 'C': 8503, 'N': 4226, 'O': 5688, 'S': 2357, 'F': 1909, 'Cl': 1740}
conformations containing S/F/Cl: 2,971 / 6,000
keys per structure: ['atomic_numbers', 'energy', 'forces', 'pos']   # note 'forces'
force array shape (atoms, 3): (5, 3)

1. Self atomic energies as model references + train/test split#

As for the other ANI presets, fit per-element self atomic energies by least squares over the seven elements and hand them to the model as atomic_energies (constant per atom, so they leave the forces untouched). The S and Cl values are large and negative because of their extra core electrons. Then a random 90/10 train/test split.

SPECIES = [1, 6, 7, 8, 16, 9, 17]        # H, C, N, O, S, F, Cl (torchani order)
counts = np.array([[np.sum(s["atomic_numbers"] == z) for z in SPECIES]
                   for s in data], float)
E = np.array([s["energy"] for s in data])
sae, *_ = np.linalg.lstsq(counts, E, rcond=None)
print("self energies (eV):", {Z_ALL[z]: round(float(e), 1)
                              for z, e in zip(SPECIES, sae)})

rng = np.random.default_rng(0)
idx = rng.permutation(len(data)); cut = int(0.9 * len(data))
train_structs = [data[i] for i in idx[:cut]]
test_structs = [data[i] for i in idx[cut:]]
print(f"train {len(train_structs)}  test {len(test_structs)}")

self energies (eV): {'H': -16.2, 'C': -1036.2, 'N': -1488.9, 'O': -2045.9, 'S': -10834.3, 'F': -2716.1, 'Cl': -12521.8}
train 5400  test 600

2. Train the `ani-2x` preset on energies and forces#

We select the model with the preset: ani-2x config key, the 1008-length AEV (5.1/3.5 A cutoffs, shift grids from 0.8 A), the seven-element wider networks (H 256:192:160, C 224:192:160, N/O 192:160:128, S/F/Cl 160:128:96) and CELU of ANI.ani2x(), and pass the fitted self energies through atomic_energies. The neighbour-list cutoff is 5.1 A (the ANI-2x radial cutoff).

from torch.utils.data import Subset
from xnn.common.data import AtomicDataset
from xnn.common.config import Config, ModelConfig, DataConfig, OptimConfig
from xnn.common.train import Trainer

CUTOFF = 5.1                      # ANI-2x radial cutoff (>= angular 3.5)
EPOCHS, BS, LR = 120, 64, 1e-3
train_ds = AtomicDataset(train_structs, CUTOFF)
test_ds = AtomicDataset(test_structs, CUTOFF)

cfg = Config(
    model=ModelConfig(name="ani", cutoff=CUTOFF,
                      extra={"preset": "ani-2x", "species": SPECIES,
                             "atomic_energies": sae.tolist()}),
    data=DataConfig(cutoff=CUTOFF, batch_size=BS),   # val_fraction 0.1 carves val
    optim=OptimConfig(lr=LR, epochs=EPOCHS, energy_weight=1.0,
                      force_weight=10.0, scheduler="plateau"),
    output_dir="runs/ani2x_subset",
)
trainer = Trainer(cfg, train_ds, test_set=test_ds)
print(f"ANI-2x parameters: {sum(p.numel() for p in trainer.module.parameters()):,}"
      f"  device {trainer.device}")

t0 = time.time(); trainer.fit()
print(f"trained {EPOCHS} epochs in {time.time()-t0:.0f}s")

ANI-2x parameters: 1,713,223  device cuda
epoch    0 | train loss 2.1616e+01 | val loss 1.3778e+01
epoch    1 | train loss 1.0254e+01 | val loss 1.0514e+01
epoch    2 | train loss 7.0194e+00 | val loss 1.0726e+01
epoch    3 | train loss 6.2858e+00 | val loss 9.3168e+00
epoch    4 | train loss 5.7944e+00 | val loss 8.0987e+00
epoch    5 | train loss 3.7282e+00 | val loss 9.3460e+00
epoch    6 | train loss 3.4788e+00 | val loss 7.5860e+00
epoch    7 | train loss 2.4071e+00 | val loss 7.1589e+00
epoch    8 | train loss 2.3018e+00 | val loss 7.0889e+00
epoch    9 | train loss 1.9856e+00 | val loss 6.9701e+00
epoch   10 | train loss 1.6465e+00 | val loss 6.5244e+00
epoch   11 | train loss 1.6924e+00 | val loss 6.3686e+00
epoch   12 | train loss 1.5755e+00 | val loss 6.1039e+00
epoch   13 | train loss 1.9616e+00 | val loss 6.5338e+00
epoch   14 | train loss 1.9465e+00 | val loss 5.9304e+00
epoch   15 | train loss 1.2248e+00 | val loss 5.9287e+00
epoch   16 | train loss 1.0825e+00 | val loss 6.2532e+00
epoch   17 | train loss 1.6089e+00 | val loss 6.5112e+00
epoch   18 | train loss 1.4690e+00 | val loss 5.8522e+00
epoch   19 | train loss 1.0832e+00 | val loss 5.9040e+00
epoch   20 | train loss 1.1896e+00 | val loss 5.8892e+00
epoch   21 | train loss 1.5186e+00 | val loss 5.8698e+00
epoch   22 | train loss 1.9575e+00 | val loss 6.0629e+00
epoch   23 | train loss 1.3052e+00 | val loss 5.6499e+00
epoch   24 | train loss 8.6208e-01 | val loss 5.8300e+00
epoch   25 | train loss 9.4785e-01 | val loss 5.6974e+00
epoch   26 | train loss 8.2676e-01 | val loss 5.6601e+00
epoch   27 | train loss 6.8550e-01 | val loss 5.5001e+00
epoch   28 | train loss 5.8387e-01 | val loss 5.4395e+00
epoch   29 | train loss 5.9566e-01 | val loss 5.4656e+00
epoch   30 | train loss 6.3387e-01 | val loss 5.4896e+00
epoch   31 | train loss 6.0810e-01 | val loss 5.4924e+00
epoch   32 | train loss 7.6710e-01 | val loss 5.6491e+00
epoch   33 | train loss 9.4398e-01 | val loss 5.4439e+00
epoch   34 | train loss 7.1772e-01 | val loss 5.4253e+00
epoch   35 | train loss 6.5828e-01 | val loss 5.4666e+00
epoch   36 | train loss 7.0872e-01 | val loss 5.4630e+00
epoch   37 | train loss 8.6244e-01 | val loss 5.4488e+00
epoch   38 | train loss 9.2090e-01 | val loss 5.2572e+00
epoch   39 | train loss 7.7218e-01 | val loss 5.3686e+00
epoch   40 | train loss 7.7013e-01 | val loss 5.4106e+00
epoch   41 | train loss 1.3156e+00 | val loss 5.4383e+00
epoch   42 | train loss 8.6030e-01 | val loss 5.3408e+00
epoch   43 | train loss 7.2891e-01 | val loss 5.1794e+00
epoch   44 | train loss 5.6099e-01 | val loss 5.1268e+00
epoch   45 | train loss 5.8000e-01 | val loss 5.4679e+00
epoch   46 | train loss 6.7008e-01 | val loss 5.3007e+00
epoch   47 | train loss 1.1040e+00 | val loss 5.3431e+00
epoch   48 | train loss 7.5570e-01 | val loss 5.0824e+00
epoch   49 | train loss 6.7751e-01 | val loss 5.3219e+00
epoch   50 | train loss 7.2635e-01 | val loss 5.1507e+00
epoch   51 | train loss 5.9174e-01 | val loss 5.2179e+00
epoch   52 | train loss 5.6623e-01 | val loss 5.1270e+00
epoch   53 | train loss 4.6881e-01 | val loss 5.1742e+00
epoch   54 | train loss 4.3947e-01 | val loss 5.1396e+00
epoch   55 | train loss 3.9432e-01 | val loss 5.1425e+00
epoch   56 | train loss 3.7668e-01 | val loss 4.9426e+00
epoch   57 | train loss 3.8884e-01 | val loss 5.1943e+00
epoch   58 | train loss 4.8703e-01 | val loss 4.9945e+00
epoch   59 | train loss 4.4587e-01 | val loss 5.1215e+00
epoch   60 | train loss 6.3092e-01 | val loss 4.9777e+00
epoch   61 | train loss 6.0561e-01 | val loss 5.1210e+00
epoch   62 | train loss 7.3014e-01 | val loss 4.8882e+00
epoch   63 | train loss 8.5694e-01 | val loss 5.1025e+00
epoch   64 | train loss 3.9477e-01 | val loss 5.2366e+00
epoch   65 | train loss 4.4241e-01 | val loss 5.0932e+00
epoch   66 | train loss 4.8801e-01 | val loss 5.0079e+00
epoch   67 | train loss 5.5652e-01 | val loss 5.0420e+00
epoch   68 | train loss 8.0737e-01 | val loss 5.0104e+00
epoch   69 | train loss 7.5612e-01 | val loss 4.8711e+00
epoch   70 | train loss 7.0596e-01 | val loss 4.9827e+00
epoch   71 | train loss 7.3518e-01 | val loss 5.0252e+00
epoch   72 | train loss 4.2159e-01 | val loss 4.8295e+00
epoch   73 | train loss 4.6540e-01 | val loss 5.0616e+00
epoch   74 | train loss 5.9063e-01 | val loss 4.9889e+00
epoch   75 | train loss 3.6112e-01 | val loss 4.9106e+00
epoch   76 | train loss 2.8046e-01 | val loss 4.8531e+00
epoch   77 | train loss 3.1950e-01 | val loss 4.9020e+00
epoch   78 | train loss 2.6542e-01 | val loss 4.8458e+00
epoch   79 | train loss 3.0524e-01 | val loss 4.8954e+00
epoch   80 | train loss 3.8715e-01 | val loss 4.9944e+00
epoch   81 | train loss 3.8507e-01 | val loss 5.0433e+00
epoch   82 | train loss 3.9639e-01 | val loss 5.0687e+00
epoch   83 | train loss 3.2649e-01 | val loss 5.0346e+00
epoch   84 | train loss 2.9007e-01 | val loss 4.9398e+00
epoch   85 | train loss 2.0770e-01 | val loss 4.9161e+00
epoch   86 | train loss 1.6876e-01 | val loss 4.9126e+00
epoch   87 | train loss 1.4886e-01 | val loss 4.9133e+00
epoch   88 | train loss 1.3489e-01 | val loss 4.9078e+00
epoch   89 | train loss 1.2602e-01 | val loss 4.9057e+00
epoch   90 | train loss 1.2020e-01 | val loss 4.8979e+00
epoch   91 | train loss 1.1622e-01 | val loss 4.8998e+00
epoch   92 | train loss 1.1339e-01 | val loss 4.8921e+00
epoch   93 | train loss 1.1117e-01 | val loss 4.8917e+00
epoch   94 | train loss 1.0961e-01 | val loss 4.8889e+00
epoch   95 | train loss 1.0571e-01 | val loss 4.8880e+00
epoch   96 | train loss 1.0526e-01 | val loss 4.8875e+00
epoch   97 | train loss 1.0507e-01 | val loss 4.8874e+00
epoch   98 | train loss 1.0490e-01 | val loss 4.8872e+00
epoch   99 | train loss 1.0472e-01 | val loss 4.8870e+00
epoch  100 | train loss 1.0433e-01 | val loss 4.8862e+00
epoch  101 | train loss 1.0448e-01 | val loss 4.8866e+00
epoch  102 | train loss 1.0436e-01 | val loss 4.8861e+00
epoch  103 | train loss 1.0399e-01 | val loss 4.8864e+00
epoch  104 | train loss 1.0415e-01 | val loss 4.8862e+00
epoch  105 | train loss 1.0382e-01 | val loss 4.8857e+00
epoch  106 | train loss 1.0342e-01 | val loss 4.8857e+00
epoch  107 | train loss 1.0329e-01 | val loss 4.8855e+00
epoch  108 | train loss 1.0333e-01 | val loss 4.8855e+00
epoch  109 | train loss 1.0342e-01 | val loss 4.8854e+00
epoch  110 | train loss 1.0332e-01 | val loss 4.8854e+00
epoch  111 | train loss 1.0322e-01 | val loss 4.8854e+00
epoch  112 | train loss 1.0317e-01 | val loss 4.8853e+00
epoch  113 | train loss 1.0318e-01 | val loss 4.8853e+00
epoch  114 | train loss 1.0321e-01 | val loss 4.8853e+00
epoch  115 | train loss 1.0316e-01 | val loss 4.8852e+00
epoch  116 | train loss 1.0319e-01 | val loss 4.8852e+00
epoch  117 | train loss 1.0300e-01 | val loss 4.8852e+00
epoch  118 | train loss 1.0301e-01 | val loss 4.8852e+00
epoch  119 | train loss 1.0306e-01 | val loss 4.8852e+00
test loss 1.2685e+00
trained 120 epochs in 1267s

3. Energy and force correlation vs DFT#

Because ANI-2x carries forces we check both. There is no torch.no_grad() here: forces are -dE/dx, so the energy must keep its graph; we detach the tensors after the model call instead.

from torch.utils.data import DataLoader
from xnn.common.data import collate

model = trainer.model.eval(); device = trainer.device
EV2KCAL = 23.060541945329334
E_pred, E_ref, F_pred, F_ref = [], [], [], []
for batch in DataLoader(test_ds, batch_size=64, collate_fn=collate):
    batch = batch.to(device)
    out = model(batch)                       # forces need autograd, no no_grad()
    E_pred.append(out["energy"].detach().cpu().numpy())
    F_pred.append(out["forces"].detach().cpu().numpy())
    E_ref.append(batch.energy.cpu().numpy())
    F_ref.append(batch.forces.cpu().numpy())
E_pred, E_ref = np.concatenate(E_pred), np.concatenate(E_ref)
F_pred, F_ref = np.concatenate(F_pred), np.concatenate(F_ref)
e_rmse = np.sqrt(np.mean((E_pred - E_ref) ** 2))
f_rmse = np.sqrt(np.mean((F_pred - F_ref) ** 2))
print(f"energy RMSE: {e_rmse*1e3:.1f} meV = {e_rmse*EV2KCAL:.3f} kcal/mol")
print(f"force  RMSE: {f_rmse*1e3:.1f} meV/A")

fig, (a1, a2) = plt.subplots(1, 2, figsize=(8.8, 4.2))
a1.scatter(E_ref, E_pred, s=8, alpha=0.5)
lim = [min(E_ref.min(), E_pred.min()), max(E_ref.max(), E_pred.max())]
a1.plot(lim, lim, "k--", lw=1)
a1.set_xlabel("DFT energy (eV)"); a1.set_ylabel("ANI-2x energy (eV)")
a1.set_title(f"energy  {e_rmse*EV2KCAL:.3f} kcal/mol")
fl = [min(F_ref.min(), F_pred.min()), max(F_ref.max(), F_pred.max())]
a2.scatter(F_ref.ravel(), F_pred.ravel(), s=3, alpha=0.3)
a2.plot(fl, fl, "k--", lw=1)
a2.set_xlabel("DFT force (eV/A)"); a2.set_ylabel("ANI-2x force (eV/A)")
a2.set_title(f"forces  {f_rmse*1e3:.0f} meV/A")
fig.tight_layout(); fig.savefig("ani2x_correlation.png", dpi=110); plt.show()

energy RMSE: 456.0 meV = 10.516 kcal/mol
force  RMSE: 358.3 meV/A

../../../../_images/1ad23b95fd08dd89ef30624a471439c427b54abe540a14196c966d1b5c7052a8.png

Summary#

load_dataset("ani2x", ...) brings the ~9.6 M-conformation, seven-element ANI-2x training set (energies and forces for H/C/N/O/S/F/Cl) into the same one-line hub as every other xnn dataset. Paired with the ani-2x preset (ANI.ani2x()), this reproduces the ANI-2x model’s training setup: the larger 1008-length AEV, the wider per-element networks, and a force-aware loss.

The full ANI family in the hub, side by side:

	elements	labels	preset
ANI-1 (`ani1`)	H C N O	energies	`ANI.ani1()`
ANI-1x (`ani1x`)	H C N O	energies + forces	`ANI.ani1x()`
ANI-1ccx (`ani1ccx`)	H C N O	CCSD(T)*/CBS energies	`ANI.ani1ccx()`
ANI-2x (`ani2x`)	+ S F Cl	energies + forces	`ANI.ani2x()`

For the pretrained ANI-2x weights transplanted from torchani (rather than trained here), see examples/fidelity_checks/ani_verification.ipynb and the parity tests in tests/test_ani.py.

The ANI-2x dataset: seven elements (S, F, Cl) and the ani-2x preset#