github.com/Yangxinsix/curator @main

Active learning for machine-learned interatomic potentials.

Train, explore, select, label, evaluate, and deploy from one modular toolkit.

Why CURATOR · Installation · Quick start · CLI · Configuration · Citation

CURATOR is a config-driven framework for building robust machine-learned interatomic potentials (MLIPs). It brings equivariant neural networks, model fine-tuning and knowledge distillation, atomistic simulation, uncertainty-aware batch selection, first-principles labeling, evaluation, and deployment into a single workflow.

Reference data  →  Train  →  Simulate  →  Select  →  Label  ↺

_{Evaluate throughout the loop · Deploy when the potential is ready}

Every stage can run independently, while MyQueue can connect the stages into autonomous, iteration-aware HPC workflows.

Why CURATOR

All stages use the same YAML configuration system without requiring the full workflow to run as one monolith.

Installation

Install PyTorch for the CPU or CUDA environment you intend to use, following the official PyTorch instructions, then install CURATOR.

Stable release

python -m pip install --upgrade pip
python -m pip install curator-torch

Official wheels include CURATOR's native C++ neighbor list, which is the default backend for data loading. Installing from source instead requires a C++17 compiler and Python development headers.

Latest development version

git clone https://github.com/Yangxinsix/curator.git
cd curator
python -m pip install -e .

For the current development branch, Python 3.10 or newer is recommended.

Optional extras

[!NOTE] GPU packages are platform-specific. Confirm that the PyTorch, CUDA, and cuEquivariance builds are compatible before installing acceleration extras.

Quick start

1. Evaluate the included LiFePO₄ model

The repository includes a small dataset and checkpoint, so you can verify an installation without training first:

curator-evaluate \
  --data example/LiFePO4.traj \
  --model example/best_model.ckpt \
  --device cpu \
  --out runs/evaluate

Metrics and plots are written to runs/evaluate/LiFePO4/:

runs/evaluate/LiFePO4/
├── metrics.json
├── parity_energy.png
├── parity_forces_xyz.png
├── hist_energy_error.png
├── hist_force_error_norm.png
└── bar_force_mae_by_element.png

Add --save-data to also write results.npz, or --no-plot for metrics-only evaluation.

2. Train the included example

cd example/train
curator-train cfg=config.yaml

Before a production run, review data.datapath, device, batch size, precision, and trainer.max_epochs in the config. Training writes the resolved config, log, checkpoints, and deployable model into run_path.

3. Run the active-learning stages

Each stage consumes a YAML config and hands an artifact to the next stage:

curator-train    cfg=train.yaml
curator-simulate cfg=simulate.yaml
curator-select   cfg=select.yaml
curator-label    cfg=label.yaml

dataset ──▶ train ──▶ model ──▶ simulate ──▶ pool ──▶ select ──▶ indices
   ▲                                                                  │
   └──────────────────────────── label ◀──────────────────────────────┘

Ready-to-edit examples live in example/, while reusable defaults and component groups live in curator/configs/.

Command-line tools

Installing CURATOR provides the following commands:

Two useful post-training commands are:

# Export a TorchScript model
curator-deploy model.ckpt --target_path compiled_model.pt

# Export a LAMMPS ML-IAP model
curator-deploy model.ckpt --mliap \
  --element-types Fe Li O P \
  --target_path mliap_model.pt

Run any command with --help for its current options. Hydra-driven commands also accept direct overrides such as device=cpu, run_path=runs/train, or trainer.max_epochs=10.

Configuration and outputs

CURATOR composes package defaults with a user YAML supplied through cfg=<path>. The resolved configuration is saved beside the run artifacts, so every experiment remains inspectable and reproducible.

The most important component groups are:

curator/configs/
├── model/representation/   # PaiNN, NequIP, MACE, Allegro, eSEN
├── data/                   # single- and multi-domain datasets
├── finetune/               # full, head-only, LoRA
├── simulator/              # engines, callbacks, uncertainty
├── task/                   # objectives, distillation, optimizers, schedulers
├── trainer/                # Lightning runtime, logging, callbacks
└── annotator/              # VASP and GPAW labeling

Interfaces and documentation

• Documentation
• Training tutorial
• Fine-tuning and knowledge distillation
• LAMMPS interface guide
• LAMMPS ML-IAP documentation
• OpenMM interface notes
• Example configurations and notebooks

Found a bug or have an idea for a new model, simulator, selector, or annotator? Please open an issue.

Citation

If CURATOR contributes to your research, please cite the CURATOR preprint:

@article{yang2024curator,
  title   = {CURATOR: Building Robust Machine Learning Potentials for Atomistic
             Simulations Autonomously with Batch Active Learning},
  author  = {Yang, Xin and Petersen, Martin Hoffmann and Sechi, Renata and others},
  journal = {ChemRxiv},
  year    = {2024},
  doi     = {10.26434/chemrxiv-2024-p5t3l}
}