Metadata-Version: 2.4
Name: equitrain
Version: 1.0.0
Summary: A training framework for machine learning interatomic potentials (MLIPs)
Author: Janosh Riebesell
Author-email: Carlos Madariaga <carlos.madariaga@bam.de>, "Jose M. Pizarro" <jose.pizarro-blanco@bam.de>, Philipp Benner <philipp.benner@bam.de>
License: MIT License
        
        Copyright (c) 2024 BAMeScience
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
        
Keywords: MLIP,machine learning,materials
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
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; 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
```

## 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`.
