Metadata-Version: 2.4
Name: metatrain
Version: 2026.3
Summary: Training and evaluating machine learning models for atomistic systems.
Author: metatrain developers
License-Expression: BSD-3-Clause
Project-URL: homepage, https://metatensor.github.io/metatrain
Project-URL: documentation, https://metatensor.github.io/metatrain
Project-URL: repository, https://github.com/metatensor/metatrain
Project-URL: issues, https://github.com/metatensor/metatrain/issues
Project-URL: changelog, https://metatensor.github.io/metatrain/latest/dev-docs/changelog.html
Keywords: machine learning,molecular modeling
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: POSIX
Classifier: Operating System :: MacOS :: MacOS X
Classifier: Operating System :: Microsoft :: Windows
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
Classifier: Topic :: Scientific/Engineering :: Chemistry
Classifier: Topic :: Scientific/Engineering :: Physics
Classifier: Topic :: Software Development :: Libraries
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: ase
Requires-Dist: huggingface_hub
Requires-Dist: numpy
Requires-Dist: metatensor-learn<0.6,>=0.5.0
Requires-Dist: metatensor-operations<0.6,>=0.5.0
Requires-Dist: metatensor-torch<0.11,>=0.10.0
Requires-Dist: metatomic-torch<0.2,>=0.1.15
Requires-Dist: metatomic-ase<0.2,>=0.1.1
Requires-Dist: jsonschema
Requires-Dist: pydantic>=2.12
Requires-Dist: typing_extensions
Requires-Dist: omegaconf>=2.3.0
Requires-Dist: python-hostlist
Requires-Dist: tqdm
Requires-Dist: vesin>=0.5.2
Provides-Extra: soap-bpnn
Requires-Dist: torch-spex<0.2,>=0.1; extra == "soap-bpnn"
Requires-Dist: sphericart-torch==1.0.8; platform_system == "Windows" and extra == "soap-bpnn"
Requires-Dist: wigners; extra == "soap-bpnn"
Provides-Extra: pet
Provides-Extra: gap
Requires-Dist: featomic-torch<0.8,>=0.7.4; extra == "gap"
Requires-Dist: skmatter; extra == "gap"
Requires-Dist: scipy; extra == "gap"
Provides-Extra: phace
Requires-Dist: physical_basis; extra == "phace"
Requires-Dist: wigners; extra == "phace"
Requires-Dist: opt-einsum; extra == "phace"
Provides-Extra: llpr
Provides-Extra: classifier
Provides-Extra: mace
Requires-Dist: mace-torch<0.3.15,>=0.3.14; extra == "mace"
Requires-Dist: e3nn; extra == "mace"
Provides-Extra: flashmd
Provides-Extra: flashmd-symplectic
Provides-Extra: dpa3
Requires-Dist: deepmd-kit[torch]>=3.1.0; extra == "dpa3"
Requires-Dist: torch>=2.7; extra == "dpa3"
Dynamic: license-file

<h1 align="center">
    <img src="https://raw.githubusercontent.com/metatensor/metatrain/refs/heads/main/docs/src/logo/metatrain-horizontal-dark.svg" alt="Metatensor logo" width="600"/>
</h1>

<h4 align="center">

[![tests status](https://img.shields.io/github/checks-status/metatensor/metatrain/main)](https://github.com/metatensor/metatrain/actions?query=branch%3Amain)
[![documentation](https://img.shields.io/badge/📚_documentation-latest-sucess)](https://metatensor.github.io/metatrain)
[![coverage](https://codecov.io/gh/metatensor/metatrain/branch/main/graph/badge.svg)](https://codecov.io/gh/metatensor/metatrain)
</h4>

<!-- marker-introduction -->

`metatrain` is a command line interface (CLI) to **train** and **evaluate** atomistic
models of various architectures. It features a common `yaml` option inputs to configure
training and evaluation. Trained models are exported as standalone files that can be
used directly in various molecular dynamics (MD) engines (e.g. `ASE`, `LAMMPS`, `i-PI`,
`TorchSim`, `ESPResSo`,...) using the [metatomic](https://docs.metatensor.org/metatomic)
interface.

The idea behind `metatrain` is to have a general training hub that provides a
homogeneous environment and user interface, transforming every ML architecture into an
end-to-end model that can be connected to MD engines. Any custom architecture compatible
with [TorchScript](https://pytorch.org/docs/stable/jit.html) can be integrated into
`metatrain`, gaining automatic access to a training and evaluation interface, as well as
compatibility with various MD engines.

<!-- marker-architectures -->

# List of Implemented Architectures

Currently `metatrain` supports the following architectures for building an atomistic
model:

| Name                                          | Description                                                                                                                          |
|-----------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------|
| [PET][arch-pet]                               | Point Edge Transformer (PET), interatomic machine learning potential                                                                 |
| [SOAP-BPNN][arch-soap_bpnn]                   | A Behler-Parrinello neural network with SOAP features                                                                                |
| [MACE][arch-mace]                             | A higher order equivariant message passing neural network.                                                                           |
| [PhACE][arch-phace]                           | SO(3)-equivariant message-passing model with physical radial functions and fast tensor products.                                     |
| [GAP][arch-gap]                               | Sparse Gaussian Approximation Potential (GAP) using Smooth Overlap of Atomic Positions (SOAP).                                       |
| [FlashMD][arch-flashmd]                       | An architecture for the direct prediction of molecular dynamics                                                                      |
| [DPA3][arch-dpa3]                             | An invariant graph neural network based on line graph series representations                                                         |
| [Symplectic FlashMD][arch-flashmd_symplectic] | A symplectic variant of FlashMD for the direct prediction of molecular dynamics.                                                     |

<!-- marker-arch-links -->

<!-- links for the different architectures. To be replaced if we are building the docs locally or
on a PR, since the docs use this README file directly.-->
[arch-dpa3]: https://docs.metatensor.org/metatrain/latest/architectures/generated/dpa3.html
[arch-flashmd]: https://docs.metatensor.org/metatrain/latest/architectures/generated/flashmd.html
[arch-gap]: https://docs.metatensor.org/metatrain/latest/architectures/generated/gap.html
[arch-mace]: https://docs.metatensor.org/metatrain/latest/architectures/generated/mace.html
[arch-pet]: https://docs.metatensor.org/metatrain/latest/architectures/generated/pet.html
[arch-phace]: https://docs.metatensor.org/metatrain/latest/architectures/generated/phace.html
[arch-soap_bpnn]: https://docs.metatensor.org/metatrain/latest/architectures/generated/soap_bpnn.html
[arch-flashmd_symplectic]: https://docs.metatensor.org/metatrain/latest/architectures/generated/flashmd_symplectic.html

<!-- marker-documentation -->

# Documentation

For details, tutorials, and examples, please visit our
[documentation](https://metatensor.github.io/metatrain/latest/).

<!-- marker-installation -->

# Installation

Install `metatrain` with pip:

```bash
pip install metatrain
```

Install specific models by specifying the model name. For example, to install the SOAP-BPNN model:

```bash
pip install metatrain[soap-bpnn]
```

We also offer a conda installation:

```bash
conda install -c conda-forge metatrain
```

> ⚠️ The conda installation does not install model-specific dependencies and will only
> work for architectures without optional dependencies such as PET.

After installation, you can use mtt from the command line to train your models!

<!-- marker-quickstart -->

# Quickstart

To train a model, use the following command:

```bash
mtt train options.yaml
```

Where options.yaml is a configuration file specifying training options. For example, the
following configuration trains a *SOAP-BPNN* model on the QM9 dataset:

```yaml
# architecture used to train the model
architecture:
  name: soap_bpnn
  training:
    num_epochs: 5  # a very short training run

# Mandatory section defining the parameters for system and target data of the training set
training_set:
  systems: "qm9_reduced_100.xyz"  # file where the positions are stored
  targets:
    energy:
      key: "U0"      # name of the target value
      unit: "eV"     # unit of the target value

test_set: 0.1        # 10% of the training_set are randomly split for test
validation_set: 0.1  # 10% of the training_set are randomly split for validation
```

<!-- marker-shell -->

# Shell Completion

`metatrain` comes with completion definitions for its commands for bash and zsh. You
must manually configure your shell to enable completion support.

To make the completions available, source the definitions in your shell’s startup file
(e.g., `~/.bash_profile`, `~/.zshrc`, or `~/.profile`):

```bash
source $(mtt --shell-completion)
```

<!-- marker-issues -->

# Having problems or ideas?

Having a problem with metatrain? Please let us know by submitting an issue.

Submit new features or bug fixes through a pull request.

<!-- marker-contributing -->

# Contributors

Thanks goes to all people who make metatrain possible:

[![Contributors](https://contrib.rocks/image?repo=metatensor/metatrain)](https://github.com/metatensor/metatrain/graphs/contributors)

The overall metatrain project is [maintained](https://github.com/lab-cosmo/.github/blob/main/Maintainers.md) by @frostedoyster, @pfebrer, and @PicoCentauri who will reply to issues and pull requests opened on this repository as soon as possible. You can mention them directly if you did not receive an answer after a couple of days. Additionally, different architectures are maintained by separate maintainers, you can find their names in the corresponding [documentation](https://docs.metatensor.org/metatrain/latest/architectures/)

<!-- marker-cite -->

# Citing metatrain

If you found metatrain useful for your work, please cite the corresponding article:

F. Bigi, J.W. Abbott, P. Loche et. al.<br>
Metatensor and metatomic: foundational libraries for interoperable atomistic machine learning, (2026).<br>
[https://doi.org/10.1063/5.0304911](https://doi.org/10.1063/5.0304911)

```bibtex
@article{bigi_metatensor_2026,
  title = {Metatensor and Metatomic: {{Foundational}} Libraries for Interoperable Atomistic Machine Learning},
  shorttitle = {Metatensor and Metatomic},
  author = {Bigi, Filippo and Abbott, Joseph W. and Loche, Philip and Mazitov, Arslan and Tisi, Davide and Langer, Marcel F. and Goscinski, Alexander and Pegolo, Paolo and Chong, Sanggyu and Goswami, Rohit and Febrer, Pol and Chorna, Sofiia and Kellner, Matthias and Ceriotti, Michele and Fraux, Guillaume},
  year = 2026,
  month = feb,
  journal = {J. Chem. Phys.},
  volume = {164},
  number = {6},
  pages = {064113},
  issn = {0021-9606},
  doi = {10.1063/5.0304911},
}
```

<!-- marker-end --> <!-- Update docs/src/cite.rst if you add more sections before this marker -->
