TorchNEP

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

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) across on one node or multiple nodes
• Fine-tuning — load any nep.txt or checkpoint.pt to do fine-tuning; optionally slim the model to only the element types present in the new dataset
• ZBL — Universal ZBL repulsive potential with optional typewise cutoffs

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 from source code:

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

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)

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)

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 batched GPU inference on an entire .xyz file and writes GPUMD-compatible output files.

from torchnep import predict_dataset

predict_dataset(
    "nep.txt",
    "test.xyz",
    output_dir="results",
    dtype="float32",       # float32 or float64
    batch_size=500,
    output_descriptor=0,   # 0=off, 1=per-frame mean, 2=per-atom (matches GPUMD)
)
# writes energy_train.out, force_train.out, virial_train.out,
# stress_train.out, and (when output_descriptor != 0) descriptor.out

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}, 
}