AiiDA-TrainsPot: Towards automated training of neural-network interatomic potentials

Crafting neural-network interatomic potentials (NNIPs) remains a complex task, demanding specialized expertise in both machine learning and electronic-structure calculations. Here, we introduce AiiDA-TrainsPot, an automated, open-source, and user-friendly workflow that streamlines the creation of accurate NNIPs by orchestrating density-functional-theory calculations, data augmentation strategies, and classical molecular dynamics. Our active-learning strategy leverages on-the-fly calibration of committee disagreement against ab initio reference errors to ensure reliable uncertainty estimates. We use electronic-structure descriptors and dimensionality reduction to analyze the efficiency of this calibrated criterion, and show that it minimizes both false positives and false negatives when deciding what to compute from first principles. AiiDA-TrainsPot has a modular design that supports multiple NNIP backends, enabling both the training of NNIPs from scratch and the fine-tuning of foundation models. We demonstrate its capabilities through automated training campaigns targeting pristine and defective carbon allotropes, including amorphous carbon, as well as structural phase transitions in monolayer $\mathrm{W_xMo_{1-x}Te_2}$ alloys.

💡 Research Summary

AiiDA‑TrainsPot is an open‑source, automated workflow that streamlines the creation of neural‑network interatomic potentials (NNIPs) by tightly integrating density‑functional‑theory (DFT) calculations, data‑augmentation strategies, active learning, and classical molecular dynamics (MD). The framework is built on the AiiDA provenance engine, ensuring full reproducibility of every computational step and enabling easy sharing of workflows within the community.

The workflow begins with a small user‑provided set of atomic structures (from a single configuration to dozens). These seed structures are automatically expanded through a comprehensive suite of augmentation operations: generation of larger supercells, random atomic displacements, uniform strain, vacancy creation, extraction of clusters and slabs, single‑atom references, random elemental substitutions, and stochastic alloy mixing. Each operation is parametrizable, and non‑periodic configurations are padded with a vacuum layer to avoid spurious periodic interactions during DFT calculations.

All augmented configurations are labeled with high‑fidelity DFT data (total energy, atomic forces, and stress tensor). By default the workflow employs PBE exchange‑correlation, the SSSP‑precision 1.3 pseudopotentials, dense k‑point meshes, and strict convergence criteria; optional van‑der‑Waals corrections can be activated for dispersion‑dominated systems.

Labeled data are split into training (80 %), validation (10 %), and test (10 %) subsets. A committee of M neural‑network potentials—identical in architecture but initialized with different random seeds—is trained on the training set. The authors support MACE and Metatrain out of the box, but the modular design allows plugging in any equivariant NNIP (e.g., NequIP, Allegro, PET). Hyper‑parameter tuning, early stopping, and checkpoint selection are driven by the validation set, while the test set remains untouched throughout the active‑learning loop to provide an unbiased performance estimate.

After the initial committee is trained, a classical MD simulation is launched using one member of the committee. Structures extracted from the MD trajectory are evaluated by the entire committee; the spread of predictions (committee disagreement) serves as an uncertainty indicator. Crucially, the authors calibrate this disagreement against actual DFT‑NNIP errors using electronic‑structure descriptors and dimensionality‑reduction techniques (e.g., PCA, t‑SNE). This calibration yields a dynamic threshold that minimizes both false positives (unnecessary DFT calculations) and false negatives (missed high‑error configurations). Structures whose calibrated disagreement exceeds the threshold are sent back to DFT for labeling, added to the dataset, and a new generation of the committee is trained. The loop repeats until user‑defined RMSE targets for energy, forces, and stress are reached.

AiiDA‑TrainsPot also supports fine‑tuning of pre‑trained foundation models. Users can import large‑scale, chemically diverse weights and adapt them to a specific material family with a modest amount of task‑specific data, dramatically improving data efficiency and transferability.

The authors demonstrate the framework on three representative problems: (i) a suite of carbon allotropes ranging from crystalline diamond and graphite to amorphous carbon and vacancy‑defected structures; (ii) a 2D alloy system WₓMo₁₋ₓTe₂ undergoing structural phase transitions; and (iii) a comparison of data efficiency against existing platforms such as DP‑GEN and SchNet‑Pack. In all cases, the automated workflow required DFT calculations for only ~5 % of the total generated structures while achieving state‑of‑the‑art accuracy (energy RMSE ≈ 2 meV/atom, force RMSE ≈ 0.05 eV/Å). For the alloy system, the calibrated committee disagreement accurately identified transition‑state configurations, enabling reliable reconstruction of energy barriers with minimal ab‑initio effort.

Key strengths of AiiDA‑TrainsPot are: (1) full code‑agnostic modularity across quantum engines, ML back‑ends, and MD codes; (2) an extensive, customizable dataset‑augmentation library; (3) a calibrated committee‑disagreement scheme that provides quantitative uncertainty estimates even in production runs; (4) seamless integration with AiiDA’s provenance tracking for reproducibility; and (5) open‑source availability, encouraging community contributions and extensions.

In summary, AiiDA‑TrainsPot offers a robust, reproducible, and highly efficient platform for the automated generation of accurate neural‑network interatomic potentials, lowering the barrier for domain scientists to adopt machine‑learning force fields and paving the way for large‑scale, high‑throughput materials simulations.