Metadata-Version: 2.4
Name: equitrain
Version: 1.1.0
Summary: A training framework for machine learning interatomic potentials (MLIPs)
Author: Philipp Benner, Ahmad Lutfi, Carlos Madariaga, Jose M. Pizarro, Janosh Riebesell
Maintainer-email: Philipp Benner <philipp.benner@bam.de>
License-Expression: MIT
Keywords: MLIP,machine learning,materials
Classifier: Programming Language :: Python :: 3
Classifier: Operating System :: OS Independent
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: matscipy
Requires-Dist: beautifulsoup4
Requires-Dist: pqdm
Requires-Dist: ase
Requires-Dist: h5py
Requires-Dist: pymatgen
Requires-Dist: tqdm
Provides-Extra: dev
Requires-Dist: mypy==1.0.1; extra == "dev"
Requires-Dist: pre-commit; extra == "dev"
Requires-Dist: ruff==0.15.11; extra == "dev"
Requires-Dist: pytest; extra == "dev"
Provides-Extra: docu
Requires-Dist: mkdocs; extra == "docu"
Requires-Dist: mkdocs-material; extra == "docu"
Requires-Dist: pymdown-extensions<12.2; extra == "docu"
Requires-Dist: mkdocs-click; extra == "docu"
Requires-Dist: mkdocstrings[python]; extra == "docu"
Requires-Dist: torch>=2.0.0; extra == "docu"
Provides-Extra: torch
Requires-Dist: torch>=2.0.0; extra == "torch"
Requires-Dist: torch_geometric>=2.4.0; extra == "torch"
Requires-Dist: torch_ema>=0.3.0; extra == "torch"
Requires-Dist: accelerate>=0.28.0; extra == "torch"
Provides-Extra: jax
Requires-Dist: jax>=0.4; extra == "jax"
Requires-Dist: jaxlib>=0.4; extra == "jax"
Requires-Dist: flax; extra == "jax"
Requires-Dist: jraph; extra == "jax"
Provides-Extra: ani
Requires-Dist: torch>=2.0.0; extra == "ani"
Requires-Dist: torch_geometric>=2.4.0; extra == "ani"
Requires-Dist: torch_ema>=0.3.0; extra == "ani"
Requires-Dist: accelerate>=0.28.0; extra == "ani"
Requires-Dist: torchani>=2.2; extra == "ani"
Requires-Dist: ase>=3.22.0; extra == "ani"
Provides-Extra: orb
Requires-Dist: torch>=2.0.0; extra == "orb"
Requires-Dist: torch_geometric>=2.4.0; extra == "orb"
Requires-Dist: torch_ema>=0.3.0; extra == "orb"
Requires-Dist: accelerate>=0.28.0; extra == "orb"
Requires-Dist: orb-models>=3.0; extra == "orb"
Requires-Dist: cuml-cu11; sys_platform == "linux" and extra == "orb"
Provides-Extra: m3gnet
Requires-Dist: torch>=2.0.0; extra == "m3gnet"
Requires-Dist: torch_geometric>=2.4.0; extra == "m3gnet"
Requires-Dist: torch_ema>=0.3.0; extra == "m3gnet"
Requires-Dist: accelerate>=0.28.0; extra == "m3gnet"
Requires-Dist: matgl>=4.0.0; extra == "m3gnet"
Provides-Extra: mace
Requires-Dist: torch>=2.0.0; extra == "mace"
Requires-Dist: torch_geometric>=2.4.0; extra == "mace"
Requires-Dist: torch_ema>=0.3.0; extra == "mace"
Requires-Dist: accelerate>=0.28.0; extra == "mace"
Requires-Dist: mace-torch>=0.3.6; extra == "mace"
Provides-Extra: mace-jax
Requires-Dist: mace-jax>=0.2.0; extra == "mace-jax"
Requires-Dist: jax>=0.4; extra == "mace-jax"
Requires-Dist: jaxlib>=0.4; extra == "mace-jax"
Requires-Dist: flax; extra == "mace-jax"
Requires-Dist: jraph; extra == "mace-jax"
Dynamic: license-file

# Equitrain: Training and Fine-Tuning Machine Learning Interatomic Potentials

Equitrain is a Python toolkit for preprocessing atomistic datasets, training
machine-learning interatomic potentials (MLIPs), fine-tuning existing
checkpoints, and running evaluation or prediction through one CLI/API.

## Features

- Unified Torch and JAX training entry points.
- Model wrappers for MACE, SevenNet, ORB, ANI, and M3GNet.
- Native HDF5 preprocessing for large atomistic datasets.
- Torch reaction-relative losses for barrier and reaction energies.
- Fine-tuning adapters for Delta/L<sup>2</sup>-SP, Freeze, and LoRA workflows.
- ASE calculator helpers for batched prediction and relaxation.

## Supported Models

| Wrapper | Backends | Upstream / Companion Project | Notes |
| --- | --- | --- | --- |
| `mace` | Torch, JAX | [`mace-model`](https://github.com/bamescience/mace-model) | Companion repository for MACE model definitions, conversion, and foundation-model export. |
| `sevennet` | Torch | [`MDIL-SNU/SevenNet`](https://github.com/MDIL-SNU/SevenNet) | Torch SevenNet checkpoints and models. |
| `orb` | Torch | [`orbital-materials/orb-models`](https://github.com/orbital-materials/orb-models) | Torch ORB force-field models. |
| `ani` | Torch, JAX | [`aiqm/torchani`](https://github.com/aiqm/torchani) | Torch uses TorchANI; JAX uses a JAX-native bundle. |
| `m3gnet` | Torch, JAX | [`materialsvirtuallab/matgl`](https://github.com/materialsvirtuallab/matgl) | Torch uses MatGL; JAX uses a JAX-native bundle. |

For MACE, use [`mace-model`](https://github.com/bamescience/mace-model) for
model construction/conversion and `equitrain` for preprocessing, training,
fine-tuning, checkpointing, evaluation, and prediction.

## Documentation

Full documentation is published at
[https://bamescience.github.io/equitrain/](https://bamescience.github.io/equitrain/):

- [Installation](https://bamescience.github.io/equitrain/installation/)
- [Quickstart](https://bamescience.github.io/equitrain/quickstart/)
- [Data and Preprocessing](https://bamescience.github.io/equitrain/data/)
- [CLI](https://bamescience.github.io/equitrain/cli/)
- [Training Options](https://bamescience.github.io/equitrain/training-options/)
- [Python API](https://bamescience.github.io/equitrain/python-api/)
- [API Reference](https://bamescience.github.io/equitrain/api-reference/)
- [Model Wrappers](https://bamescience.github.io/equitrain/model-wrappers/)
- [JAX Bundles](https://bamescience.github.io/equitrain/jax-bundles/)
- [Fine-Tuning](https://bamescience.github.io/equitrain/fine-tuning/)
- [Phonon Fine-Tuning Paper](https://bamescience.github.io/equitrain/phonon-finetuning-paper/)
- [Calculators](https://bamescience.github.io/equitrain/calculators/)
- [Reaction-Relative Losses](https://bamescience.github.io/equitrain/reaction-relative-losses/)
- [Resources](https://bamescience.github.io/equitrain/resources/)

The documentation source is in `docs/`. Build or serve it locally with:

```bash
pip install -e '.[docu]'
mkdocs serve
```

## Optional EquiLoad HDF5 Loader

Equitrain's portable PyTorch/PyG loader remains the default.
[EquiLoad](https://github.com/BAMeScience/EquiLoad) is a separately installed native
extension. This integration is tested with the `v0.1.0` release. Build it against
the exact Torch installation in the Equitrain environment, because PyTorch's C++
ABI is not stable between Torch releases:

```bash
git clone --branch v0.1.0 https://github.com/BAMeScience/EquiLoad.git equiload
cd equiload
bash scripts/install-python.sh --venv /path/to/equitrain/.venv
cd ..
equitrain --backend torch --torch-loader equiload ...
```

Use the tagged release when building EquiLoad. Earlier development checkouts
also reported version `0.1.0` but lacked the global sample-index fix required
for correct multi-file loading. Rebuild those older installations from the tag.

Do not install a wheel built in another environment. EquiLoad is CPU-side; Equitrain
moves its returned batches to the active accelerator device. It currently requires
Torch 2.6 through 2.14; the standard Python loader remains available for other
supported Equitrain Torch versions.

The EquiLoad loader is currently single-process and does not support weighted resampling or
relative reaction losses. It returns the same nested PyG batch structure as the standard
loader and moves batches to Equitrain's active accelerator device. Without `equiload`
installed, the default `--torch-loader python` requires no Rust tooling.

## Installation

```bash
pip install equitrain
```

Until the package is fully available on PyPI, install from a local clone:

```bash
git clone https://github.com/BAMeScience/equitrain.git
cd equitrain
python3.10 -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
pip install uv
uv pip install -e '.[dev,docu]'
```

Install model/runtime extras as needed:

```bash
pip install 'equitrain[torch,mace]'
pip install 'equitrain[jax,mace-jax]'
pip install 'equitrain[torch,ani]'
```

## Minimal Workflow

Preprocess data:

```bash
equitrain-preprocess \
    --train-file data-train.xyz \
    --valid-file data-valid.xyz \
    --compute-statistics \
    --atomic-energies average \
    --output-dir data \
    --r-max 4.5
```

Train a Torch/MACE model:

```bash
equitrain -v \
    --train-file data/train.h5 \
    --valid-file data/valid.h5 \
    --output-dir runs/mace \
    --model path/to/mace.model \
    --model-wrapper mace \
    --epochs 10 \
    --tqdm
```

Evaluate and predict:

```bash
equitrain-evaluate -v \
    --test-file data/test.h5 \
    --model path/to/mace.model \
    --model-wrapper mace \
    --output-dir evaluation_mace

equitrain-predict \
    --predict-file data/valid.h5 \
    --model path/to/mace.model \
    --model-wrapper mace \
    --output-dir predictions_mace
```

See the [Quickstart](https://bamescience.github.io/equitrain/quickstart/) for the full workflow, including JAX
bundles and fine-tuned checkpoint export.

## Fine-Tuning Note

Equitrain's Delta adapter is a residual-parameter implementation of
L<sup>2</sup>-SP ("Starting Point") regularization from Li, Grandvalet, and
Davoine, 2018,
[*Explicit Inductive Bias for Transfer Learning with Convolutional Networks*](https://proceedings.mlr.press/v80/li18a.html).
It parameterizes fine-tuning as `theta = theta_0 + delta`, so weight decay on
trainable deltas regularizes `||delta||_2^2`.

Delta combined with `freeze_layers` is targeted L<sup>2</sup>-SP
(L<sup>2</sup>-TSP): the L<sup>2</sup>-SP penalty applies only to selected
trainable delta layers while frozen layers remain exactly at their pre-trained
starting values. See [Fine-Tuning](https://bamescience.github.io/equitrain/fine-tuning/).

## Resources

Example data-preparation scripts are in `resources/data`, training scripts are
in `resources/training`, and initial model examples are in `resources/models`.
