TorchNEP

A pure PyTorch implementation of the NEP4 (Neuroevolution Potential) training framework.

Many ready-to-use examples (the training inputs and trained models of the TorchNEP paper) are available at TorchNEP_models.

✨ Features

• 🔄 GPUMD-compatible — output nep.txt files load directly into GPUMD for MD simulation
• 🎯 Two-stage training — Stage 1: force-focused; Stage 2: energy-focused
• 🖥️ Multi-GPU training — distributed data parallel (DDP) on one node or multiple nodes
• ⚡ Fast on NVIDIA and AMD — torch.compile support with automatic backend selection, tuned on both CUDA and ROCm GPUs
• 💾 Memory-friendly — the dataset stays in host memory and batches are streamed to the GPU, so memory scales with batch size, not dataset size
• 🔧 Fine-tuning — load any nep.txt or checkpoint.pt to fine-tune; optionally slim the model to only the element types present in the new dataset
• 🛡️ ZBL — universal ZBL repulsive potential with optional typewise cutoffs
• 📈 Plots — loss curves, parity plots and error breakdowns straight from the output files (torchnep.plot)

_{Excellent single-GPU training speed (left) and near-linear multi-node parallel scaling (right).}

📦 Installation

TorchNEP needs only torch >= 2.0 and numpy, but neither is installed automatically — install the PyTorch build that matches your CUDA/CPU setup first (see the official guide; numpy comes with it).

Then install TorchNEP with:

pip install torchnep -U

or install the latest development version directly from GitHub:

pip install git+https://github.com/mushroomfire/torchnep.git

or install from source code:

git clone https://github.com/mushroomfire/torchnep.git
cd torchnep
pip install .

Optional extras: torchnep[ase] (ASE calculator), torchnep[plot] (figures: matplotlib, polars), torchnep[all] (everything).

📄 Training data (extended-XYZ)

TorchNEP reads extended-XYZ files. The parser is strict — the rules below are enforced, and violations raise on load.

Comment line tags

• Lattice="ax ay az bx by bz cx cy cz" — mandatory. Nine floats in Å giving the three lattice vectors as rows. Every frame is treated as fully periodic, so pbc=... is ignored. For isolated clusters/molecules or a non-periodic direction, use a vacuum box wider than the NEP cutoff.
• energy=<value>— optional, eV. System energy.
• virial="vxx vxy vxz vyx vyy vyz vzx vzy vzz" — optional, eV. Must have exactly 9 components. Positive values denote compressed states, negative denote stretched states (GPUMD convention).
• stress="sxx sxy sxz syx syy syz szx szy szz" — optional, eV/ų. Must have exactly 9 components. Positive = stretched, negative = compressed — opposite sign to virial. If both virial and stress are present, virial wins.

Per-atom columns

The Properties=... schema declares column layout. TorchNEP reads only three fields and silently ignores everything else (e.g. Z:I:1):

• species:S:1 — chemical symbol (case-sensitive; must match the type list in nep.in).
• pos:R:3 — Cartesian position in Å.
• force:R:3 or forces:R:3 — reference force in eV/Å (optional).

⚙️ Training Parameters

Model architecture (GPUMD-compatible)

Cutoff rules, checked on load: the angular cutoff is ≥ 3 Å and ≤ the radial cutoff for every species, the radial cutoff ≤ 100 Å, and the ZBL outer cutoff (zbl, every zbl.in row) is between 1 and 3 Å. The ZBL cutoff is therefore always below the smallest angular cutoff, so the ZBL term can safely be evaluated on the angular neighbor list.

Training hyperparameters

Runtime arguments

Everything that is not about hyperparameter values lives on the Python function (train_nep / train_nep_sharded):

📊 Output Files

🚀 Launch training

Single GPU / CPU / MPS — train_nep

# run_train.py
from torchnep import train_nep
train_nep("nep.in", "train.xyz", output_dir="output")

# with a validation set (either form):
train_nep("nep.in", "train.xyz", output_dir="output", valid_file="valid.xyz")
train_nep("nep.in", "train.xyz", output_dir="output", valid_ratio=0.1)

# export the exact valid_ratio split as GPUMD-ready files — train the same
# partition in GPUMD (or anything else) and compare loss curves directly:
from torchnep import export_valid_split
export_valid_split("train.xyz", valid_ratio=0.1, run_seed=42,
                   output_dir="split")   # writes split/train.xyz + split/test.xyz
# coverage-aware split (same strategy as valid_strategy="stratified"):
export_valid_split("train.xyz", valid_ratio=0.1, run_seed=42,
                   output_dir="split", strategy="stratified")

python run_train.py

Multi-GPU, single node — train_nep_sharded

Each rank loads only 1/N of the structures, so total GPU memory for the data store scales as 1/N.

# run_train.py
from torchnep import train_nep_sharded
train_nep_sharded("nep.in", "train.xyz", output_dir="output")

torchrun --standalone --nproc_per_node=4 run_train.py    # 4 GPUs on this node

Multi-GPU, multi-node (SLURM) — train_nep_sharded

For M nodes × N GPUs each, the key SLURM directives are:

#SBATCH --nodes=2                  # M nodes
#SBATCH --ntasks-per-node=1        # 1 srun task per node; torchrun fans out to all GPUs
#SBATCH --gpus-per-node=4          # N GPUs per node
#SBATCH --cpus-per-task=16         # CPU cores per node

MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)
MASTER_PORT=$((20000 + SLURM_JOB_ID % 40000))

srun --nodes=$SLURM_NNODES --ntasks-per-node=1 bash -c "
  torchrun \
    --nnodes=$SLURM_NNODES \
    --nproc_per_node=\$SLURM_GPUS_ON_NODE \
    --node_rank=\$SLURM_NODEID \
    --rdzv_id=$SLURM_JOB_ID \
    --rdzv_backend=c10d \
    --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
    run_train.py
"

🔁 Restart and Resume

Two ways to resume:

# 1) automatic: looks for checkpoint.pt in output_dir (restart=True default)
train_nep("nep.in", "train.xyz", output_dir="output")

# 2) explicit: continue from a specific checkpoint
train_nep("nep.in", "train.xyz", output_dir="output",
          resume_from="output/checkpoint_stage1.pt")

What you can safely change on restart

🔧 Fine-Tuning

Fine-tuning starts from a pre-trained model's weights instead of random initialisation. The architecture (nep.in parameters) must match the source model, but the new dataset's element types may be a subset of the original.

Basic fine-tuning

train_nep(
    "nep.in",
    "new_data.xyz",
    output_dir="finetune_output",
    finetune_from="pretrained/nep.txt",   # or a "pretrained/checkpoint.pt"
    slim_types=True,
)

finetune_from accepts:

• nep.txt — GPUMD text format (works with models trained by GPUMD or TorchNEP)
• checkpoint.pt — full checkpoint (weights are extracted automatically)

If the new dataset contains fewer element types than the original model, slim_types=True removes the unused types before training begins, shrinking the model and speeding up training.

Standalone model slimming

from torchnep.model import NEPModel, slim_model
from torchnep.data import parse_nep_in

config = parse_nep_in("nep.in")
model = NEPModel(config)
model.load_weights_from_nep_txt("nep.txt")

slimmed = slim_model(model, ["Cr", "Ni"])
slimmed.save_nep_txt("nep_slim.txt", max_NN_radial, max_NN_angular)

🔮 Prediction

Single-structure prediction

from torchnep.nep import NEPCalculator
import numpy as np

calc = NEPCalculator("nep.txt")
result = calc.compute(
    species=["Cr", "Cr", "Ni"],
    positions=np.array([[0,0,0],[1.5,0,0],[3,0,0]]),
    cell=np.eye(3) * 6.0,
)
print(result["energy"])         # (N,) per-atom energy (eV); sum for total
print(result["forces"])         # (N, 3) forces (eV/Å)
print(result["virial"])         # (N, 9) per-atom virial (eV)

# Split the NEP (neural-network) part from the ZBL repulsive part:
result = calc.compute(..., return_components=True)
print(result["energy_nep"], result["energy_zbl"])   # sum == result["energy"]

ASE calculator

If ASE is installed, any ASE workflow (relaxation, MD, EOS, …) can drive a trained model:

from ase.io import read
from torchnep.ase_calculator import NEP

atoms = read("POSCAR")
atoms.calc = NEP("nep.txt", dtype='float32', device='cuda')
print(atoms.get_potential_energy())   # eV
print(atoms.get_forces())             # (N, 3) eV/Å
print(atoms.get_stress())             # Voigt 6-vector eV/ų (periodic cells)

# NEP / ZBL / total breakdown of energy, forces, and stress:
parts = atoms.calc.get_components()
print(parts["nep"]["energy"], parts["zbl"]["energy"], parts["total"]["energy"])

Full-dataset prediction

Runs streamed, batched inference on an entire .xyz file and writes GPUMD-compatible output files. The file is indexed once, then processed chunk by chunk (~chunk_atoms atoms each: read → neighbor lists → device-sized batches → rows appended to the outputs), so host memory is bounded by the chunk and device memory by the batch — a dataset of any size finishes on any machine, only the wall time differs. A progress bar (tqdm if installed) shows frames done.

from torchnep import predict_dataset

predict_dataset(
    "nep.txt",
    "test.xyz",
    output_dir="results",
    output_descriptor=0,   # 0=off, 1=per-frame mean, 2=per-atom (matches GPUMD)
    batch_size=None,       # auto from free GPU memory (OOM-halving retry); or an int
    chunk_atoms=None,      # atoms per streamed chunk, default 200000 (env TORCHNEP_PREDICT_CHUNK_ATOMS)
)
# writes energy_train.out, force_train.out, virial_train.out,
# stress_train.out, and (when output_descriptor != 0) descriptor.out

Multi-GPU / multi-node prediction

predict_dataset_sharded has the same arguments and writes the same files, but is launched with one process per GPU (torchrun --nproc_per_node=8 script.py, or srun under Slurm).

from torchnep import predict_dataset_sharded
predict_dataset_sharded("nep.txt", "huge.xyz", output_dir="results")

📈 Plotting

torchnep.plot.NEPPlotter draws figures straight from the files a run writes (loss.out, energy/force/virial/stress_train.out, *_test.out; predict_dataset writes the same files). Needs matplotlib: pip install torchnep[plot].

from torchnep.plot import NEPPlotter

p = NEPPlotter()
p.dashboard("run", out="dashboard.png")                 # loss curves + E/F/stress parity, train and valid
p.loss("run", out="loss.png")                            # the loss panel alone
p.parity("run", out="parity.png")                        # E/F/stress parity plots (virial=True for the virial)
p.parity("run", kind="density", margins=True,            # hexbin density + error distributions
         out="parity_density.png")
p.parity("pred", shift_energy="element", xyz="test.xyz", # data of another DFT reference: remove
         out="pred.png")                                 # the per-element energy offset first

NEPPlotter(font, fontsize, dpi, colors, cmaps, panel_labels, label_format, frame, rc) sets the style for every figure.

_{p.dashboard("run")}

_{p.parity("run", kind="density", margins=True)}

🗂️ Source layout

The torchnep/ package is organised as follows:

📚 Citation

If you use TorchNEP in your research, please cite the following paper:

@misc{wu2026torchne,
      title={TorchNEP: Ultra-Efficient and Accurate Training of Neuroevolution Potentials}, 
      author={Yong-Chao Wu and Xiaoya Chang and Tero Mäkinen and Amin Esfandiarpour and Jian-Li Shao and Tapio Ala-Nissila and Zheyong Fan and Mikko Alava},
      year={2026},
      eprint={2606.19557},
      archivePrefix={arXiv},
      primaryClass={physics.comp-ph},
      url={https://arxiv.org/abs/2606.19557}, 
}