Orchestr.AI

Unified Platform for Machine-Learning Force Fields for Quantum Dots 🚀

Orchestr.AI is a unified, modular, and engine‑agnostic framework for training state‑of‑the‑art machine learning force fields (MLFFs) for quantum dots (QDs).
It integrates multiple ML engines under a single interface:

✅ SchNet ✅ PaiNN ✅ SO3net ✅ FieldSchNet ✅ NequIP ✅ Allegro ✅ MACE

1. Installation

For the installation of the Orchestr.AI platform and all the required packages, we recommend to create a conda environment using Python 3.12. Orchestr.AI supports three installation modes depending on your system and preference:

Mode	Tool	Recommended For
Conda	conda	Simple users
Micromamba	micromamba	Faster HPC setup
Micromamba + UV	micromamba + uv	Fastest recommended setup
Legacy	Conda manual	Compatibility

Clone the repository

git clone https://github.com/nlesc-nano/Orchestr.AI.git
cd Orchestr.AI

Choose installation method

Show installation options

Add PRINT_VERSIONS=0 before the command to skip version checks and speed up installation.

Option A — Conda

PRINT_VERSIONS=0 bash scripts/setup_envs.sh
source scripts/orchestr_ai_env.sh

Option B — Micromamba

PRINT_VERSIONS=0 bash scripts/setup_envs_micromamba.sh
source scripts/orchestr_ai_env.sh

Option C — Micromamba + UV

PRINT_VERSIONS=0 bash scripts/setup_envs_micromamba_uv.sh
source scripts/orchestr_ai_env.sh

Option D — Legacy (manual setup)

conda env create -f environment.yaml
conda activate orchestr_ai_legacy
pip install mace-torch==0.3.14
pip install -e .
source scripts/orchestr_ai_single_env.sh

Installation notes

Example of internal dispatch

Running in env prefix: orchestr_ai-core
[Orchestr.AI] Dispatch: engine 'nequip' → env 'orchestr_ai-nequip'
Running in env prefix: orchestr_ai-nequip

What this setup does

• Creates three environments:
  • core → SchNet, PaiNN, SO3net, FieldSchNet
  • nequip → NequIP, Allegro
  • mace → MACE
• Installs orchestr_ai package in all environments
• Enables automatic environment dispatch

Custom environment names

Modify these scripts:

• setup_envs.sh
• setup_envs_micromamba.sh
• setup_envs_micromamba_uv.sh

Environment variables for custom names:

• ORCHESTRAI_CORE_CONDA_ENV
• ORCHESTRAI_NEQUIP_CONDA_ENV
• ORCHESTRAI_MACE_CONDA_ENV

Example:

export ORCHESTRAI_CORE_CONDA_ENV=mycore
export ORCHESTRAI_NEQUIP_CONDA_ENV=mynequip
export ORCHESTRAI_MACE_CONDA_ENV=mymace

Running on SLURM (HPC)

Click to expand SLURM usage

Multi-environment mode (recommended)

sbatch run_training.sh input.yaml --engine nequip --train-after-generate

Flow:

core → dispatch → engine-specific env → training

Summary

Mode	Dispatch	Recommended
Multi-env	Yes	Yes
Single-env	No	Legacy

---

Getting started

The current version of the platform is developed for being run in a cluster. Thus, in this repository one can find the necessary code, a bash script example for submitting jobs in a slurm queue system and an input file example.

This platform is currently being subject of several changes. Thus, on the meanwhile, descriptions of the files will be included here so they can be used.

2. Preprocessing tools

An input file example for the preprocessing of the data can be found in config_files/preprocessing/preprocess_config.yaml. The initial data for being processed should be placed in a consistent way to the paths indicated in the input file. This preprocessing tool is used for preparing the xyz files in the useful formats after DFT calculations with CP2K.

By default, the preprocessing code assumes that the input file is preprocess_config.yaml. If that is the case, it can be run as:

python -m orchestr_ai.preprocessing.generate_mlff_dataset

However, if an user wants to specify a different custom configuration file for the preprocessing, the code can be run as:

python -m orchestr_ai.preprocessing.generate_mlff_dataset --config my_experiment.yaml

3. Training Guide

Orchestr.AI supports two ways to train:

A) Using Unified YAML (Recommended)

(The file is typically named input.yaml inside config_files/training/.)

This YAML contains:

• shared hyperparameters
• dataset config
• the name of the engine (platform:)
• training/evaluation settings

You control workflow using these two main flags:

✔ --only-generate

Only generate engine-specific YAML + converted dataset.
Training will not run.

✔ --train-after-generate

Generate engine YAML → then start training automatically.

Priority Rule

If both flags are given:

--only-generate takes priority ⇒ NO training begins.

3.1 CLI Usage

The main entry:

python -m orchestr_ai.training

Arguments:

--config                Path to unified YAML file (input.yaml)
--engine                Override engine name (optional)
--input                 Override input XYZ file
--only-generate         Only produce engine YAML
--train-after-generate  Produce YAML + immediately train
--benchmark             Run cross-engine benchmarks
--post-process          Summaries from benchmark results

3.2 Recommended Workflow: Unified YAML

3.2.1 Generate engine YAML and immediately train

python -m orchestr_ai.training --config input.yaml --engine <engine_name> --train-after-generate
python -m orchestr_ai.training --config input.yaml --engine schnet --train-after-generate

3.2.2 Only generate engine YAML (no training)

python -m orchestr_ai.training --config input.yaml --engine <engine_name> --only-generate
python -m orchestr_ai.training --config input.yaml --engine schnet --only-generate

3.3 Running on SLURM

The repo includes running_files/run_training.sh.

3.3.1 Generate + Train

sbatch run_training.sh input.yaml --engine <engine_name> --train-after-generate
sbatch run_training.sh input.yaml --engine schnet --train-after-generate

3.3.2 Only Generate

sbatch run_training.sh input.yaml --engine <engine_name> --only-generate
sbatch run_training.sh input.yaml --engine schnet --only-generate

Same applies for any engine:

sbatch run_training.sh input.yaml --engine nequip --train-after-generate
sbatch run_training.sh input.yaml --engine allegro --train-after-generate

3.4 Using Engine-Specific YAML Files

(Example: config_files/training/schnet.yaml, config_files/training/nequip.yaml)

You do not need to use --only-generate or --train-after-generate.

Train directly

python -m orchestr_ai.training --config schnet.yaml --engine schnet

or SLURM:

sbatch run_training.sh schnet.yaml --engine schnet

3.5 Override dataset path (optional)

But you can specify input data inside input.yaml.

python -m orchestr_ai.training --config schnet.yaml --engine schnet --input data/new.xyz

3.6 Summary Table

Task	Recommended Command
Generate engine YAML only	python -m orchestr_ai.training --config input.yaml --engine <engine_name> --only-generate
Generate + Train	python -m orchestr_ai.training --config input.yaml --engine <engine_name> --train-after-generate
Train using engine-specific YAML	python -m orchestr_ai.training --config <engine_name>.yaml --engine <engine_name>
SLURM – Generate only	sbatch run_training.sh input.yaml --engine <engine_name> --only-generate
SLURM – Generate + Train	sbatch run_training.sh input.yaml --engine <engine_name> --train-after-generate

---

Inference code

After the training has finished, an user can run the inference code that generates the MLFF:

python -m orchestr_ai.training.inference

By default, it will look for a input file called input.yaml. Thus, if an user wants to specify another input file, one can do the following:

python -m orchestr_ai.training.inference --config input_file.yaml

After inference, if an user wants to use fine-tuning, that option is also available in the following way:

python -m orchestr_ai.training.fine_tuning

If an input file different from the default one was used, the procedure is the following:

python -m orchestr_ai.training.fine_tuning --config input_file.yaml

Postprocessing

More details will be added in future versions, but the postprocessing code is run as:

python -m orchestr_ai.postprocessing

The postprocessing part of the code, requires also to install the following packages: plotly, kneed.

CLI Mode - Extract Training Metrics from TensorBoard Event Files

This script, analysis/extract_metrics.py, extracts scalar training metrics from TensorBoard event files and saves them to a CSV file.

• -p/--path: Path to the TensorBoard event file. (Required).
• -o/--output_file: Provides the path to the CSV file containing the training metrics.

• Prerequisites Required Python Packages:
  • tensorboard You can install these using pip:

  pip install tensorboard

Command-Line Usage:

To run the script use the following command:

python analysis/extract_metrics.py -p <event_file_path> [-o <output_file_name>]

CLI Mode - Plotting Training Metrics for SchNet and Nequip

The analysis/plot.py script allows you to visualize training progress for your models. It accepts several command-line options to control its behavior. Here’s what each option means:

• --platform: Specifies the model platform. Use either schnet or nequip.
• --file: Provides the path to the CSV file containing the training metrics.
• --cols: Sets the number of columns for the subplot grid (default is 2).
• --out: Defines the output file name for the saved plot. The name should end with .png.

Plotting SchNet Results

To plot the results for SchNet, use the following command:

python analysis/plot.py --platform schnet --file "path/to/schnet_metrics.csv" --cols 2 --out schnet_plot.png

Replace "path/to/schnet_metrics.csv" with the actual path to your SchNet metrics CSV file.

Plotting Nequip Results

To plot the results for Nequip, use the following command:

python analysis/plot.py --platform nequip --file "path/to/nequip_metrics.csv" --cols 2 --out nequip_plot.png

Replace "path/to/nequip_metrics.csv" with the actual path to your Nequip metrics CSV file.

These commands will generate plots for the respective platforms and save them as PNG files in the current working directory.

GUI Mode: Interactive Metrics Extraction and Plotting with Streamlit

The analysis/app.py script offers a Streamlit GUI to extract metrics from TensorBoard event files and visualize SchNet/NequIP training progress with static (Matplotlib, saveable) or interactive (Plotly, display-only) plots.

Prerequisites: - streamlit, plotly

pip install streamlit plotly

Launching the GUI:

streamlit run analysis/app.py