Orchestr.AI: A Unified Workflow for Training and Benchmarking of Machine-Learned Interatomic Potentials
Zain Ul Abideen, Mario Fernández-Pendás, Muhammad Usman, Camilo Rodríguez Quintero, Masuma Suleymanova, Ignacio Arganda-Carreras, Ivan Infante*
* Corresponding author
Input files for the rMD17 and 3BPA platform tests
This repository contains the input files for the Orchestr.AI platform tests reported in the paper. To test the platform, Orchestr.AI was run end to end on two public benchmark datasets, rMD17 and 3BPA, training five of the engines it supports (SchNet, PaiNN, NequIP, Allegro and MACE) on the same splits.
Everything needed to rerun the tests is here: split definitions, engine templates, data-preparation scripts and SLURM workflows. Trained models, logs and result tables are not stored; they are produced by running the workflow, and the results are reported in the paper.
| rMD17 | 3BPA | |
|---|---|---|
| Systems | aspirin (21 atoms), benzene (12), ethanol (9), paracetamol (20), uracil (12) | 3-(benzyloxy)pyridin-2-amine (27 atoms) |
| Reference data | Revised MD17, Christensen & von Lilienfeld (2020) | Kovács et al. (2021) |
| Splits | 5 official rMD17 splits per molecule | 4 splits: trained at 300 K, tested at 300 K, 600 K, 1200 K and on dihedral scans |
| Train / val / test | 950 / 50 / 1000 | 450 / 50 / 1669 – 7047 |
| Engines | SchNet, PaiNN, NequIP, Allegro, MACE | SchNet, PaiNN, NequIP, Allegro, MACE |
| Training runs | 5 molecules × 5 engines × 5 splits = 125 | 5 engines × 4 splits = 20 |
| Raw data in this repo | split index files only (download the .npz files, see below) |
✓ included in 3BPA/ |
| Folder | rMD17_benchmark/ |
3BPA_benchmark/ |
Both platform tests follow the same three stages. Each stage is one command, run from inside the benchmark folder.
run_pipeline.py (run_pipeline_3BPA.py for 3BPA) drives stage 1 and can also run its steps separately:
| Option | What runs |
|---|---|
--all (default) |
build splits → generate YAMLs → validate YAMLs |
--data_only |
build splits only |
--yaml_gen |
generate and validate YAMLs from existing splits |
input-files/
├── rMD17_benchmark/
│ ├── rmd17/splits/ # official rMD17 split indices (index_train_01..05.csv, index_test_01..05.csv)
│ ├── All_Engine_Yaml/ # engine templates: schnet, painn, nequip, allegro, mace
│ ├── run_pipeline.py # stage 1 driver
│ ├── test_new.py # builds the splits from rmd17_<molecule>.npz
│ ├── make_engine_yamls.py # writes one YAML per molecule × engine × split
│ ├── validate_engine_yamls.py
│ ├── submit_data_gen.sh # SLURM: stage 1
│ ├── submit_test.sh # submits all training jobs (stage 2)
│ ├── run_training.sh # SLURM: one training job
│ ├── stage_from_yaml.py # copies a job's data files to node-local scratch
│ ├── submit_collect_result.sh # SLURM: stage 3
│ ├── collect_results.py
│ └── split_checksums.md5 # MD5 of every split file used in the paper
├── 3BPA_benchmark/
│ ├── 3BPA/ # raw data: train_300K.xyz, test_300K/600K/1200K/dih.xyz
│ ├── All_Engine_Yaml/
│ ├── run_pipeline_3BPA.py
│ ├── test_3BPA.py
│ ├── make_engine_yamls_3BPA.py
│ ├── validate_engine_yamls_3BPA.py
│ ├── submit_data_gen.sh
│ ├── submit_test.sh
│ ├── run_training.sh
│ ├── stage_from_yaml.py
│ ├── submit_collect_result.sh
│ ├── collect_results_3bpa_updated.py
│ └── split_checksums.md5
├── docs/images/ # figures used in this README
└── README.md
| Script | Stage | Role |
|---|---|---|
submit_data_gen.sh |
1 | SLURM job that runs run_pipeline*.py --all |
test_new.py / test_3BPA.py |
1 | Build the train / val / test files, the combined all.xyz, the SchNetPack all.db and split_file.npz |
make_engine_yamls*.py |
1 | Fill each template with the split's file paths, the seed and the element list |
validate_engine_yamls*.py |
1 | Check that every file referenced by every generated YAML exists |
All_Engine_Yaml/*.yaml |
1 | Engine templates holding all hyperparameters |
submit_test.sh |
2 | Loop over molecule × engine × split and submit run_training.sh from each run folder |
run_training.sh |
2 | SLURM job: stage data to scratch, train with Orchestr.AI, copy results back |
stage_from_yaml.py |
2 | List the data files a YAML references and rewrite their paths for the scratch copy |
submit_collect_result.sh |
3 | SLURM job that runs the collect script |
collect_results.py / collect_results_3bpa_updated.py |
3 | Parse every job's .out file and write the benchmark tables |
The five molecules come from the Revised MD17 dataset (100 000 frames per molecule, CC0). The .npz files are not stored here because of their size (aspirin and paracetamol exceed 100 MB). Download the dataset and copy the five files into rMD17_benchmark/:
rMD17_benchmark/rmd17_aspirin.npz
rMD17_benchmark/rmd17_benzene.npz
rMD17_benchmark/rmd17_ethanol.npz
rMD17_benchmark/rmd17_paracetamol.npz
rMD17_benchmark/rmd17_uracil.npz
The five official rMD17 splits (rmd17/splits/) are used for every molecule. Energies and forces are converted from kcal/mol and kcal/mol/Å to eV and eV/Å.
| Split | Train indices | Test indices | Train / val / test |
|---|---|---|---|
| split01 – split05 | index_train_XX.csv (1000, shuffled → 950 train + 50 val) |
index_test_XX.csv (1000) |
950 / 50 / 1000 |
The 3BPA data of Kovács et al. (2021) (CC BY 4.0) is included in 3BPA_benchmark/3BPA/.
The 500 frames of train_300K.xyz (MD at 300 K) are shuffled once into 450 training and 50 validation frames. All four splits share this training and validation set; each split tests on a different condition: MD at 300 K, 600 K and 1200 K, and dihedral scans.
| Split | Test set | Test condition | Test frames | Train / val / test |
|---|---|---|---|---|
| split01 | test_300K.xyz |
MD at 300 K (same temperature as training) | 1669 | 450 / 50 / 1669 |
| split02 | test_600K.xyz |
MD at 600 K | 2138 | 450 / 50 / 2138 |
| split03 | test_1200K.xyz |
MD at 1200 K | 2139 | 450 / 50 / 2139 |
| split04 | test_dih.xyz |
dihedral scans | 7047 | 450 / 50 / 7047 |
| File | Used by |
|---|---|
*_train*.xyz, *_val50.xyz, *_test*.xyz |
NequIP, Allegro, MACE |
*_all.db (SchNetPack database of train + val + test, in this order) + split_file.npz |
SchNet, PaiNN |
*_all.xyz |
source of *_all.db |
split_checksums.md5 holds the MD5 checksum of every train / val / test file used in the paper. After stage 1, run this inside the benchmark folder:
md5sum -c split_checksums.md5Every line should end in OK.
Install Orchestr.AI by following the installation instructions in the Orchestr.AI repository. The Micromamba + UV option is the one used for these tests. Orchestr.AI creates three environments (core, NequIP, MACE); the scripts activate the core environment, and Orchestr.AI dispatches each engine to its own environment.
Then adjust the scripts to your cluster. These lines appear in run_training.sh, submit_data_gen.sh and submit_collect_result.sh in both benchmark folders:
| What | Lines in the scripts | Change to |
|---|---|---|
| SLURM account, partition, qos | #SBATCH --account=<your_account>, --partition=<your_partition>, --qos=<your_qos> |
your values; delete any line your cluster does not use |
| GPU | #SBATCH --gres=gpu:a100:1 (one NVIDIA A100 per job in our runs) |
the GPU request for your cluster |
| micromamba | eval "$($HOME/.local/bin/micromamba shell hook -s bash)" |
the location of your micromamba |
| Environment names | ORCHESTRAI_CORE_CONDA_ENV="orchestr_ai-core-micromamba-uv2" (and ..._NEQUIP_..., ..._MACE_...) |
the names of your Orchestr.AI environments (the setup script's default is orchestr_ai-core-micromamba-uv, orchestr_ai-nequip-micromamba-uv, orchestr_ai-mace-micromamba-uv) |
git clone https://github.com/nlesc-nano/input-files.git
cd input-files/rMD17_benchmark # or: cd input-files/3BPA_benchmark
# rMD17 only: copy the five rmd17_<molecule>.npz files into this folder first
# 1) Data preparation: splits, per-run YAMLs, path check
sbatch submit_data_gen.sh
md5sum -c split_checksums.md5 # optional: confirm the splits match the paper
# 2) Training: submits one SLURM job per molecule × engine × split
bash submit_test.sh
# 3) Results: once all training jobs have finished
sbatch submit_collect_result.shsubmit_test.sh is run with bash, not sbatch: it calls sbatch run_training.sh itself for every run folder. Each training job does the following:
Running the workflow creates the following (shown for rMD17; 3BPA uses output_Dataset_3BPA/3bpa/ and benchmark_results_3BPA/):
rMD17_benchmark/
├── output_Dataset/
│ ├── pipeline_<timestamp>.log
│ └── rmd17_aspirin/
│ ├── split01/
│ │ ├── rmd17_aspirin_split01_train950.xyz
│ │ ├── rmd17_aspirin_split01_val50.xyz
│ │ ├── rmd17_aspirin_split01_test1000.xyz
│ │ ├── rmd17_aspirin_split01_all.xyz
│ │ ├── rmd17_aspirin_split01_all.db
│ │ └── split_file.npz
│ ├── split02/ … split05/
│ └── mace/ # likewise schnet/, painn/, nequip/, allegro/
│ └── split01/
│ ├── mace_split01.yaml # generated in stage 1
│ ├── rmd17_aspirin_mace_split01-<jobid>.out
│ └── <jobid>/standardized/ # model, checkpoints and logs of the job
└── benchmark_results/
└── <timestamp>/
├── benchmark_raw_results.csv / .txt # one row per run
└── benchmark_summary.csv / .txt # mean over splits
- Orchestr.AI: the platform tested here
If you use these input files, please cite:
@article{abideen2026orchestrai,
title = {Orchestr.AI: A Unified Workflow for Training and Benchmarking of Machine-Learned Interatomic Potentials},
author = {Zain Ul Abideen and Mario Fernández-Pendás and Muhammad Usman and Camilo Rodríguez Quintero and Masuma Suleymanova and Ignacio Arganda-Carreras and Ivan Infante},
journal = {},
year = {2026}
}and the original datasets:
@article{christensen2020rmd17,
title = {On the role of gradients for machine learning of molecular energies and forces},
author = {Christensen, Anders S. and von Lilienfeld, O. Anatole},
journal = {Machine Learning: Science and Technology},
volume = {1},
number = {4},
pages = {045018},
year = {2020},
doi = {10.1088/2632-2153/abba6f}
}
@article{kovacs2021ace,
title = {Linear Atomic Cluster Expansion Force Fields for Organic Molecules: Beyond RMSE},
author = {Kov{\'a}cs, D{\'a}vid P{\'e}ter and van der Oord, Cas and Kucera, Jiri and Allen, Alice E. A. and Cole, Daniel J. and Ortner, Christoph and Cs{\'a}nyi, G{\'a}bor},
journal = {Journal of Chemical Theory and Computation},
volume = {17},
number = {12},
pages = {7696--7711},
year = {2021},
doi = {10.1021/acs.jctc.1c00647}
}The 3BPA data is redistributed under CC BY 4.0 (Kovács et al., 2021) and the rMD17 split files under CC0. For the scripts and templates, please contact the corresponding author (ivan.infante@bcmaterials.net) for terms of use.