The ANI-1x dataset in one line: forces, active learning, and the ani-1x preset

The ANI-1x dataset (Smith et al., J. Chem. Phys. 148, 241733, 2018; released in Sci. Data 7, 134, 2020) is the training set behind the ANI-1x potential. It differs from the ANI-1 dataset in the two ways that matter here:

• Active learning, not brute force. Instead of densely sampling every molecule (ANI-1’s ~20 M conformations from Normal-Mode Sampling), ANI-1x iteratively adds only the conformations where an ensemble of models disagrees: ~5 M conformations that are fewer but more diverse and more transferable.

• Forces. ANI-1x ships wB97X energies and forces (plus higher-level CCSD(T)/CBS energies), where ANI-1 is energy-only.

Same one-line hub as every other dataset:

from xnn.common.data import load_dataset
splits = load_dataset("ani1x", split="train")   # ani1x-release.h5 -> xnn dicts

The first call downloads the one 5.6 GB HDF5 file from figshare (cached and MD5-verified under datasets/ani1x/); afterwards it is instant and offline. Here we load a small subset, train the xnn ani-1x preset on energies and forces, and look at both parities. This is the companion to ani1_dataset.ipynb, which does the energy-only ANI-1 version; read them side by side to see the difference between the two models.

Run with the xnn kernel.

0. Load a subset of ANI-1x: note the forces

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 few hundred molecule groups, capped so the demo trains in minutes. wb97x_dz
# is the level of theory the ANI-1x model was fit to; energies -> eV, forces ->
# eV/A. Per-conformation NaN entries (un-computed properties) are dropped for us.
data = load_dataset("ani1x", max_molecules=200, max_conformations=40,
                    level="wb97x_dz", units="eV")["all"]
print(f"loaded {len(data):,} conformations")
sizes = [len(s['atomic_numbers']) for s in data]
elems = sorted(set(int(z) for s in data for z in s['atomic_numbers']))
print(f"elements {elems}  atoms/mol {min(sizes)}-{max(sizes)}")
print(f"keys per structure: {sorted(data[0])}")   # note 'forces' -- absent in ANI-1
print(f"force array shape (atoms, 3): {data[0]['forces'].shape}")

loaded 5,027 conformations
elements [1, 6, 7, 8]  atoms/mol 15-44
keys per structure: ['atomic_numbers', 'energy', 'forces', 'pos']
force array shape (atoms, 3): (20, 3)

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

Fit per-element self atomic energies by least squares (the ANI EnergyShifter idea). Rather than subtracting them from the data by hand, we hand them to the model as atomic_energies: per-element reference energies the network learns the residual on top of (they are constant per atom, so they leave the forces untouched). Then a random 90/10 train/test split.

SPECIES = [1, 6, 7, 8]
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: round(float(e), 2) 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): {1: -16.35, 6: -1036.15, 7: -1489.82, 8: -2046.74}
train 4524  test 503

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

The one-line difference from the ANI-1 demo: ANI-1x has forces, so force_weight > 0. We select the model with the preset: ani-1x config key (the 384-length AEV, per-element network widths, and CELU of ANI.ani1x()) and pass the fitted self energies through the atomic_energies key.

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.2                      # ANI-1x radial cutoff (>= angular 3.5)
EPOCHS, BS, LR = 120, 64, 1e-3
train_ds = AtomicDataset(train_structs, CUTOFF)
test_ds = AtomicDataset(test_structs, CUTOFF)
n_val = max(1, len(train_ds) // 10)
val_idx = list(range(n_val)); tr_idx = list(range(n_val, len(train_ds)))

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

hist = {"train": [], "val": []}
trainer._log = lambda ep, tr, va: (hist["train"].append(tr.get("loss")),
                                   hist["val"].append(va.get("loss")))
t0 = time.time(); trainer.fit()
print(f"trained {EPOCHS} epochs in {time.time()-t0:.1f}s")

ANI-1x parameters: 326,660  device cuda
test loss 4.6513e-01
trained 120 epochs in 2497.8s

3. Energy and force correlation vs DFT

Because ANI-1x carries forces, we can check both. Note 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-1x 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-1x force (eV/A)")
a2.set_title(f"forces  {f_rmse*1e3:.0f} meV/A")
fig.tight_layout(); fig.savefig("ani1x_correlation.png", dpi=110); plt.show()

energy RMSE: 950.7 meV = 21.924 kcal/mol
force  RMSE: 214.3 meV/A

Summary

load_dataset("ani1x", ...) brings the ~5 M-conformation, active-learning ANI-1x training set (energies and forces) into the same one-line hub as every other xnn dataset. Paired with the ani-1x preset (ANI.ani1x()), this reproduces the ANI-1x model’s training setup end-to-end: the leaner 384-length AEV, per-element networks, and a force-aware loss.

Side by side with ani1_dataset.ipynb:

	ANI-1 (ani1)	ANI-1x (ani1x)
sampling	dense Normal-Mode (~20 M)	active learning (~5 M)
labels	energies	energies + forces
AEV / preset	768-length, ANI.ani1()	384-length, ANI.ani1x()
loss here	energy only	energy + force

The dataset and the preset are the two halves of one model: ani-1 + load_dataset("ani1") reproduces ANI-1, and ani-1x + load_dataset("ani1x") reproduces ANI-1x. For the pretrained ANI-1x weights transplanted from torchani (rather than trained here), see examples/fidelity_checks/ani_verification.ipynb.