Your First Training Run#

This tutorial trains a MACE potential on a small dataset from start to finish, explaining each step. It requires the gnn extra (pip install -e ".[gnn]").

1. Prepare the data#

xnn consumes structures as plain Python dictionaries. Required keys are pos (an (N, 3) array of positions) and atomic_numbers (an (N,) array); training targets and periodicity are optional:

import numpy as np

structures = []
for frame in my_trajectory:            # however you obtain your data
    structures.append(
        {
            "pos": frame.positions,             # (N, 3), required
            "atomic_numbers": frame.numbers,    # (N,),   required
            "cell": frame.cell,                 # (3, 3), optional (periodic)
            "pbc": frame.pbc,                   # (3,),   optional
            "energy": frame.energy,             # scalar, training target
            "forces": frame.forces,             # (N, 3), training target
            # "stress": frame.stress,           # (3, 3), optional target
        }
    )

If your data lives in a file format ASE can read (extxyz, CIF, VASP, …), skip the dictionaries entirely and load the file directly (requires the ase extra):

from xnn.common.data import AtomicDataset

dataset = AtomicDataset.from_file("my_trajectory.extxyz", cutoff=4.0)

Energy / forces / stress targets stored in the file are picked up automatically; see Data Pipeline.

2. Build the dataset#

AtomicDataset turns the structure dictionaries into AtomicGraph objects, building the (PBC-aware) neighbor list at the given cutoff:

from xnn.common.data import AtomicDataset

cutoff = 4.0
train_set = AtomicDataset(structures[:800], cutoff)
val_set = AtomicDataset(structures[800:], cutoff)

Molecular and periodic systems go through the same class; periodicity is handled entirely inside the neighbor list and edge vectors, and models never see the difference.

3. Configure the model#

The Config dataclass brings everything together: model, data, and optimizer settings. Model-specific options go into cfg.model.extra:

from xnn.common.config import Config

cfg = Config()
cfg.model.name = "mace"
cfg.model.cutoff = 4.0
cfg.model.extra = {
    "species": [18],              # atomic numbers in the dataset (Argon here)
    "max_ell": 3,                 # spherical-harmonic order of the edges
    "num_channels": 32,           # feature channels
    "num_interactions": 2,        # message-passing layers (T)
    "correlation": 3,             # order of the symmetric contraction
    "avg_num_neighbors": 20.0,    # normalization; compute from your data
}

cfg.data.batch_size = 8
cfg.optim.lr = 1e-3
cfg.optim.epochs = 100
cfg.optim.energy_weight = 1.0
cfg.optim.force_weight = 10.0    # > 0 enables force training
cfg.device = "auto"
cfg.output_dir = "runs/my_first_run"

Alternatively, load the same settings from a YAML file (see Configuration), or start from the templates in configs/.

4. Train#

Trainer builds the model from the config, wraps it in ForceStressOutput (force and stress heads are switched on by nonzero loss weights), sets up the Adam optimizer, scheduler, and data loaders, and runs the training loop:

from xnn.common.train import Trainer

trainer = Trainer(cfg, train_set, val_set)
trainer.fit()

Progress is printed per epoch; best.pt (lowest validation loss) and last.pt are written to cfg.output_dir. A checkpoint is a dictionary {"model": state_dict, "cfg": Config}. A held-out test set is optional: pass it as a third dataset (Trainer(cfg, train_set, val_set, test_set)) or set cfg.data.test_fraction to carve one out of the training data; it is evaluated once after the last epoch (see Training).

5. Predict#

Load the checkpoint and evaluate on new structures:

import torch
from xnn.common.data import AtomicDataset
from xnn.common.models import build_model, ForceStressOutput

ckpt = torch.load("runs/my_first_run/best.pt", weights_only=False)
model = build_model(ckpt["cfg"].model)
model.load_state_dict(ckpt["model"])
model = ForceStressOutput(model).eval()

graph = AtomicDataset(test_structures, cutoff)[0]
out = model(graph)
print(out["energy"], out["forces"].shape)

6. Where to go from here#

  • Drive the same run from a config file and the CLI: Train a Model

  • Attach the model to ASE for molecular dynamics: Run Molecular Dynamics with ASE

  • Export it for production LAMMPS runs: Export to LAMMPS and TorchScript

  • The example notebooks in examples/gnn/ repeat this workflow on real Argon MD data and validate every step against the reference MACE, NequIP, and Allegro implementations.