Quickstart#

The fastest way to see xnn in action is the bundled quickstart script, which builds toy data, trains a small SchNet for a few epochs, and predicts energies and forces with the trained model:

python examples/quickstart.py

It runs on CPU in a few seconds (and picks up a GPU automatically when one is present). The rest of this page walks through exactly what the script does.

1. Build data as plain dictionaries#

xnn consumes structures as plain dictionaries with keys pos and atomic_numbers (and optionally cell, pbc, energy, forces, stress). The quickstart generates a toy set of small random H/C/O structures with a smooth synthetic target, just to have a learnable signal:

import numpy as np

def toy_structures(n=24, periodic=False):
    """Random small structures with a smooth synthetic energy/forces target."""
    rng = np.random.default_rng(0)
    out = []
    for _ in range(n):
        natoms = rng.integers(3, 6)
        pos = rng.uniform(0, 4, size=(natoms, 3))
        z = rng.choice([1, 6, 8], size=natoms)
        # toy target: pairwise gaussian well
        d = np.linalg.norm(pos[:, None] - pos[None], axis=-1)
        e = float(-np.exp(-((d - 1.5) ** 2)).sum())
        s = {"pos": pos, "atomic_numbers": z, "energy": e,
             "forces": rng.normal(0, 0.1, size=(natoms, 3))}
        if periodic:
            s["cell"] = np.eye(3) * 6.0
            s["pbc"] = [True, True, True]
        out.append(s)
    return out

AtomicDataset converts each dictionary into an AtomicGraph, the single data object every xnn model consumes. Real data loads just as easily: any ASE-readable file format (extxyz, CIF, VASP, …) with AtomicDataset.from_file("trajectory.extxyz", cutoff), and standard benchmark datasets download in one line with load_dataset() (e.g. load_dataset("rmd17", molecule="aspirin", cutoff=5.0)); see Data Pipeline.

2. Configure and train#

Everything in xnn funnels through one Config dataclass, one dataset class, and one trainer. The quickstart trains a small SchNet, but any registered model name works here:

from xnn.common.config import Config
from xnn.common.data import AtomicDataset
from xnn.common.train import Trainer

cfg = Config()
cfg.model.name = "schnet"        # schnet|hdnnp|ani|physnet|nequip|mace|allegro|cace|bamboo
cfg.model.cutoff = 5.0
cfg.model.n_features = 32
cfg.model.n_interactions = 2
cfg.data.batch_size = 4          # >1 == batch training; set 1 to disable
cfg.optim.epochs = 3
cfg.optim.force_weight = 1.0
cfg.device = "auto"              # picks cuda if present, else cpu
cfg.output_dir = "runs/quickstart"

dataset = AtomicDataset(toy_structures(periodic=False), cfg.data.cutoff)
trainer = Trainer(cfg, dataset)
trainer.fit()

Training prints a per-epoch summary and writes best.pt and last.pt checkpoints to cfg.output_dir (default runs/exp).

3. Predict energies and forces#

The trainer wraps the model in ForceStressOutput, which adds conservative forces (and optionally stress) by automatic differentiation of the predicted energy. trainer.module is that trained, wrapped model, ready for inference:

model = trainer.module.to("cpu").eval()
g = dataset[0]
out = model(g)                   # {"energy", "node_energy", "forces", ...}
print("energy:", round(float(out["energy"].detach()), 4),
      "| forces shape:", tuple(out["forces"].shape))

Note

Forces are computed by autograd, so do not wrap inference in torch.no_grad(). Also predict with the trained module (as above): a freshly built SchNet predicts exactly its energy shift, because its readout head starts zero-initialized (the DTNN convention).

Run end to end, the script prints something like:

registered models: ['allegro', 'ani', 'bamboo', 'cace', 'hdnnp', 'mace', 'nequip', 'opls', 'physnet', 'reaxff', 'schnet']
training on cuda ...
epoch    0 | train loss 2.0641e+00 | val loss 9.3133e-01
epoch    1 | train loss 1.8768e+00 | val loss 8.7541e-01
epoch    2 | train loss 1.7326e+00 | val loss 7.9425e-01
energy: -0.5449 | forces shape: (5, 3)

The pieces also work on their own#

Each layer of xnn is independently importable: data, featurizers, and models compose but do not require each other:

# Data on its own
from xnn.common.data import AtomicDataset, build_neighbor_list
ds = AtomicDataset(structures, cutoff=5.0)
graph = ds[0]

# Featurizers on their own (AtomicGraph -> model inputs)
from xnn.dnn.featurizers import AEV, RadialSymmetryFunctions
from xnn.gnn.featurizers import SphericalHarmonicEdgeEmbedding
descriptor = AEV(species=[1, 6, 8])(graph)              # (N, D) invariant AEV
edges = SphericalHarmonicEdgeEmbedding(l_max=2)(graph)  # equivariant edges

# Models on their own
from xnn.common.models import build_model, ForceStressOutput, available_models
model = ForceStressOutput(build_model(cfg.model))

From the command line#

The same workflow is available through the xnn command:

xnn train --config configs/train.yaml --set optim.epochs=50
xnn export --config configs/train.yaml --ckpt runs/exp/best.pt --to lammps

Next steps#