TorchNEP

Train NEP machine-learned interatomic potentials in PyTorch — the models run directly in GPUMD.

Documentation · Quick start · Examples · Citation

TorchNEP is a from-scratch implementation of NEP4, the neuroevolution potential architecture.

GPUMD-compatible — nep.txt files load directly into GPUMD for molecular dynamics
Two-stage training — a force-focused stage, then an energy-focused stage
Multi-GPU and multi-node — data-parallel training with near-linear scaling
Fast on NVIDIA and AMD — torch.compile with automatic backend selection, tuned on CUDA and ROCm
Memory-friendly — the dataset stays in host memory; GPU memory scales with the batch, not the dataset
Fine-tuning, ZBL, plots — start from any nep.txt, add short-range repulsion, plot every run

Training speed and scaling

Installation

Install the PyTorch build for your hardware first, then:

pip install torchnep -U

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

Quick start

from torchnep import train_nep

train_nep("nep.in", "train.xyz", output_dir="output", valid_ratio=0.1)

The best model is written to output/nep_best.txt. The documentation covers nep.in, the training data format, multi-GPU training, restart and fine-tuning, prediction, the ASE calculator and plotting.

Citation

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

Metadata

Release files for torchnep 1.0.6a1

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for torchnep 1.0.6a1
File	Size	Uploaded
torchnep-1.0.6a1.tar.gz	174.2 kB	Sep 17, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for torchnep 1.0.6a1
File	Interpreter	ABI	Platform	Reset
torchnep-1.0.6a1-py3-none-any.whl 183.0 kB Sep 17, 2026	Python 3	none	any	Details

Total release size: 357.3 kB

Release files / torchnep-1.0.6a1.tar.gz

Download URL	torchnep-1.0.6a1.tar.gz
Size	174.2 kB
Tags	Source
SHA-256 checksum How to use checksums	`373ff27ed2ff13db92bc41c7b15b2d32e6af3eeb50ac48b914e8efdbb293e1ad`
BLAKE2b-256 checksum How to use checksums	`6c47f5541e93e69a102c59876af7d97ee95efac17051b7f8f4e92156c321c7e0`
Upload date	Sep 17, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	Yes
Uploaded via	`twine/7.0.0 CPython/3.13.14`

Provenance

Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.

Identity

Publishing platform: GitHub Actions
Publishing repository: github.com/mushroomfire/torchnep
Publishing commit: github.com/mushroomfire/torchnep/tree/c58eb8ce2081408201751b4d11b51db38a55c28e

Workflow

Publishing configuration: github.com/mushroomfire/torchnep/blob/c58eb8ce2081408201751b4d11b51db38a55c28e/.github/workflows/publish.yml
Publishing logs: github.com/mushroomfire/torchnep/actions/runs/35266300750/attempts/1

Artifact

Subject: torchnep-1.0.6a1.tar.gz
SHA-256 checksum: 373ff27ed2ff13db92bc41c7b15b2d32e6af3eeb50ac48b914e8efdbb293e1ad Verifying attestations

Release files / torchnep-1.0.6a1-py3-none-any.whl

Download URL	torchnep-1.0.6a1-py3-none-any.whl
Size	183.0 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`41855adfbedeef9d43e433c8016d3e0a58c34c4c85307c7610643bd064e70bad`
BLAKE2b-256 checksum How to use checksums	`4e39027bdbe47b96005249225935303736fb9d0d6ec4a2e4051795025af2886d`
Upload date	Sep 17, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	Yes
Uploaded via	`twine/7.0.0 CPython/3.13.14`

Provenance

Identity

Publishing platform: GitHub Actions
Publishing repository: github.com/mushroomfire/torchnep
Publishing commit: github.com/mushroomfire/torchnep/tree/c58eb8ce2081408201751b4d11b51db38a55c28e

Workflow

Publishing configuration: github.com/mushroomfire/torchnep/blob/c58eb8ce2081408201751b4d11b51db38a55c28e/.github/workflows/publish.yml
Publishing logs: github.com/mushroomfire/torchnep/actions/runs/35266300750/attempts/1

Artifact

Subject: torchnep-1.0.6a1-py3-none-any.whl
SHA-256 checksum: 41855adfbedeef9d43e433c8016d3e0a58c34c4c85307c7610643bd064e70bad Verifying attestations