Skip to content
aboys-cbPublic

About

An automated toolkit for training neuroevolution potential (NEP), integrating tools like GPUMD, VASP, and NEP for streamlined workflows including perturbation, active learning, single-point energy calculations, and potential training.

Resources

Stars

51 stars

Watchers

1 watching

Forks

Repository files navigation

English | 简体中文

NepTrain

NepTrain is a command-line tool for the complete lifecycle of neuroevolution potential (NEP) models. It can run training, molecular dynamics, structure selection, and labeling as standalone tasks, or compose the same steps into a resumable active-learning workflow.

train → explore → select → label → evaluate → update

Standalone commands and automated workflows use the same scientific adapters and execution targets. The workflow layer owns planning, state transitions, and acceptance; it does not duplicate the scientific calculation logic.

Runtime messages, CLI help, notifications, and generated project/report text use English. README files and documentation remain available in English and Chinese. User-provided names, paths, and external program messages are preserved.

What NepTrain provides

  • GPUMD or TorchNEP training with a canonical nep.txt result.
  • GPUMD or LAMMPS sampling, including LAMMPS DynSpin for spin MD.
  • VASP, ABACUS, MACE, DeepMD/DPA, or TACE labeling.
  • Element-set-aware farthest-point sampling with deterministic provenance.
  • Local process, local Slurm, and SSH + Slurm execution targets.
  • Content-addressed task bundles, atomic publication, hash validation, and resumable workflow state.
  • Automatic training and new-label evaluation parity plots (diagonal, RMSE, error distributions), training loss curves, and selection PCA coverage plots with selected structures highlighted. Each plot includes a JSON provenance report.

Installation

NepTrain requires Python 3.10 or later:

pip install NepTrain

This README, the documentation, and repository examples track the current source. For these examples, install from the repository root and use the same revision on the controller and workers:

python -m pip install -e .
# After installing a suitable PyTorch runtime:
python -m pip install -e '.[torchnep]'

The published package may not yet include the latest repository features.

Install only the optional runtime you need:

# TorchNEP training
pip install torch
pip install 'NepTrain[torchnep]'

# MACE teacher labeling
pip install 'NepTrain[mace]'

# DeepMD / DPA teacher labeling
pip install 'NepTrain[deepmd]'

# TACE teacher labeling (TACE is currently installed from its source repository)
pip install \
  'TACE @ git+https://github.com/xvzemin/tace.git@4b977dcc13ee87d8ba6cceba3ffb7abe43c087c8'

# Optional cuEquivariance acceleration on Ampere-or-newer CUDA 12 GPUs
pip install \
  'TACE[cueq12] @ git+https://github.com/xvzemin/tace.git@4b977dcc13ee87d8ba6cceba3ffb7abe43c087c8'
export TACE_USE_CUE=1

# SOAP descriptors when manual selection has no NEP model
pip install 'NepTrain[soap]'

LAMMPS, GPUMD, VASP, ABACUS, and first-principles resource files are supplied by the user or computing platform.

First successful run

Use the deterministic toy workflow to check the NepTrain installation without submitting DFT or Slurm jobs:

neptrain smoke --profile ordinary --workflow

For a real project, validate every configured execution target before submission:

neptrain doctor --project project.yaml

doctor collects missing required data, training inputs, route structures, and resource errors before reporting the summary. Optional tests and disabled automatic convergence produce warnings. It also reads the project backends, stage targets, setup scripts, and target environment. It checks the actual GPUMD/TorchNEP, LAMMPS/GPUMD, VASP/ABACUS, or MACE/DeepMD/TACE runtime required by each target.

Standalone steps

Train one model:

neptrain train train.xyz \
  --backend torchnep \
  --config nep.in \
  --device cuda \
  -o nep.txt

Run a structure × temperature MD batch:

neptrain md structures/ \
  --backend lammps \
  --model nep.txt \
  --temperature 300 500 700 \
  --steps 100000 \
  --max-concurrent 12 \
  -o trajectories.xyz

Select representative structures:

neptrain select trajectories.xyz \
  --base train.xyz \
  --nep nep.txt \
  --max-selected 64 \
  --out selected.xyz \
  --report selected.selection.json

Label structures with VASP:

neptrain label selected.xyz \
  --backend vasp \
  --input-file INCAR \
  --resources /shared/potpaw_PBE \
  --potcar-manifest vasp-resources.json \
  --structures-per-job 1 \
  --max-concurrent 20 \
  -o labeled.xyz

Submitted standalone tasks use one control surface:

neptrain task status runs/label-...
neptrain task logs runs/label-...
neptrain task wait runs/label-...
neptrain task retry runs/label-...
neptrain task cancel runs/label-...

The final output is published only after every shard passes identity, order, and artifact validation. Existing output is not overwritten unless --force is given.

Automated workflow

Create a schema-v8 project:

neptrain workflow init \
  --profile slurm \
  --ensemble npt \
  --dft-backend vasp \
  --directory fe-project
cd fe-project

After filling in structures, training data, templates, resource manifests, and execution targets:

neptrain doctor --project project.yaml
neptrain workflow run project.yaml --prepare-only
neptrain workflow run workflow

New projects always use the current active-learning flow. Without workflow.convergence, automatic accuracy convergence is disabled; add system-specific thresholds using the commented example generated in project.yaml.

Each sampling generation starts by training and activating a model on its incoming dataset. evaluate checks that sampling model against newly labeled structures before update merges them for the next generation. Even without new labels, a subsequent sampling generation still starts with training. Optional test datasets are diagnostic only and do not decide convergence or block progression when unavailable. Completion requires new-label accuracy, production coverage, and the configured passing streak, followed by final training.

Inspect and control it with:

neptrain workflow status workflow --jobs
neptrain workflow status workflow --details
neptrain workflow resume workflow
neptrain workflow restart workflow --generation 3 --from label --dry-run
neptrain workflow stop workflow
neptrain workflow extend workflow 15

Workflow control commands default to readable summaries; use --json for scripts. extend workflow 15 sets the total sampling budget to 15; it does not add 15 generations. Follow the printed next action to start or resume.

Before a stage restart, NepTrain refreshes unconfirmed job states. If a cluster fault leaves a job stuck in cleanup, workflow restart ... --force abandons that attempt and uses a new task directory; late results are ignored. Stop the controller first. Old jobs are not cancelled by this flag and may still consume resources. Add --dry-run to preview without changes; see the recovery guide.

Workflow status shows the current stage, dataset growth, training progress, and MD progress as a ps range with readable-trajectory counts. It shows the latest three new-label evaluations, convergence thresholds and coverage gaps, plus plot paths. --details adds all evaluations and optional test diagnostics; --jobs groups batches by generation, stage, attempt, and target. Job observations are cached and have their own timestamps; reading workflow status does not query the scheduler. --json retains individual records for scripts.

NepTrain accepts only schema_version: 8. Unknown fields and legacy project formats fail explicitly instead of being migrated silently.

Built-in LAMMPS and GPUMD sampling adapts trajectory dump spacing to the run length. Custom LAMMPS templates can use {{ dump_interval }}. Slurm analysis targets can also declare memory_ladder: [4G, 8G, 16G]; selection OOM advances to the next tier with a bounded, traceable retry.

Optional Feishu progress notifications can be configured directly in project.yaml:

notifications:
  feishu:
    webhook: https://open.feishu.cn/open-apis/bot/v2/hook/REPLACE
    secret: REPLACE
    timeout_seconds: 5

neptrain doctor --project project.yaml sends a real connectivity probe. Notifications cover successful startup/resume, user pause/stop, stage restarts, generation progress, failures, and one final completion report. Normal resume is not a stall. Messages identify the workflow and whether existing jobs are retained. Sustained monitoring failures trigger one warning after five minutes and one recovery message.

Delivery uses a persistent outbox with bounded background retries. Notification failures never change computation or convergence. workflow status shows pending deliveries, retries, historical failures, and the last successful delivery. Existing delivered records remain valid; enabling notifications midway through a run does not replay all historical generations. See the notification documentation for shutdown and recovery behavior.

Labeling backends

Backend Runtime Main boundary
VASP User-provided VASP and pinned POTCAR manifest Ordinary labels; DeltaSpin-enabled VASP 6 also supports spin/mforce
ABACUS User-provided ABACUS and pinned UPF/ORB manifest Ordinary or DeltaSpin spin/mforce labels
MACE NepTrain[mace] Ordinary structures; no mforce
DeepMD / DPA NepTrain[deepmd] DPA-3 and supported DPA-4 formats; no mforce
TACE Official TACE source install Checkpoints supported by tace-eval; spin requires real noncollinear magnetic-force output

Teacher-model labeling still enters through neptrain label. The hidden model-worker command is an internal runner protocol, not a second user-facing CLI.

Documentation and tutorials

The website documentation is built from one Chinese source tree with complete English translation catalogs. Both languages are checked with strict Sphinx warnings in CI.

Scientific-data boundaries

  • Structure identity uses the versioned structure-id.v3 contract.
  • Spin datasets use canonical spin and mforce extxyz arrays.
  • VASP and ABACUS resources are pinned by relative path and SHA256 manifest.
  • Model, dataset, task, result, and publication identities are content addressed.
  • Workflow ledger entries are scientific commit points; controller state is execution intent only.

Support

VASP 6 DeltaSpin spin labeling: configuration and INCAR.

Spin DFT labeling can select elements in the project, for example labeling.constrained_elements: [Fe]. Per-atom spin_active masks take precedence. See the constraint rules.

About

An automated toolkit for training neuroevolution potential (NEP), integrating tools like GPUMD, VASP, and NEP for streamlined workflows including perturbation, active learning, single-point energy calculations, and potential training.

Resources

Stars

51 stars

Watchers

1 watching

Forks

Releases

Packages

Used by

Contributors

Languages