Metadata-Version: 2.4
Name: upet
Version: 0.3.1
Summary: A universal interatomic potential for advanced materials modeling
Author: Filippo Bigi, Matthias Kellner, Paolo Pegolo, Davide Tisi, Guillaume Fraux, Sergey Pozdnyakov, Philip Loche
Author-email: Arslan Mazitov <arslan.mazitov@epfl.ch>, Michele Ceriotti <michele.ceriotti@epfl.ch>
Maintainer-email: Arslan Mazitov <arslan.mazitov@epfl.ch>, Filippo Bigi <filippo.bigi@epfl.ch>
License: BSD-3-Clause
Project-URL: repository, https://github.com/lab-cosmo/upet
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.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: metatrain<2026.5,>=2026.4
Requires-Dist: metatomic-ase<0.2,>=0.1.4
Requires-Dist: nvalchemi-toolkit-ops<0.5.0,>=0.4.0
Requires-Dist: huggingface_hub
Requires-Dist: hf_xet
Requires-Dist: packaging
Requires-Dist: platformdirs
Requires-Dist: tqdm
Provides-Extra: nvalchemi
Requires-Dist: nvalchemi-toolkit==0.2.0; extra == "nvalchemi"
Dynamic: license-file

<div align="center" width="600">
  <picture>
    <source srcset="https://github.com/lab-cosmo/upet/raw/refs/heads/main/docs/static/images/upet-logo-with-text-dark.svg" media="(prefers-color-scheme: dark)">
    <img src="https://github.com/lab-cosmo/upet/raw/refs/heads/main/docs/static/images/upet-logo-with-text.svg" alt="Figure">
  </picture>
</div>

> [!NOTE]
> The PET-MAD-1.6 models trained for 102 elements at the r2SCAN level of theory are
> now available! On top of the MAD-1.5 training data, they were additionally trained 
> on catalytic surfaces, and therefore have better accuracy for surface reactions 
> and adsorption energies. They also come in a new **M** size. Check out the
> [documentation](https://lab-cosmo.github.io/upet/latest/models.html)
> and updated [preprint](https://arxiv.org/abs/2603.02089) for more details.
> Try them out and let us know what you think!
```py
from upet.ase import UPETCalculator
calculator = UPETCalculator(model="pet-mad-s", version="1.6.0", device="cuda")
```

> [!NOTE]
> A new experimental integration of UPET with the
> [NVIDIA ALCHEMI Toolkit](https://developer.nvidia.com/alchemi) is now available!
> It allows for GPU-native batched inference, relaxations, and MD simulations with
> a torch-compiled version of UPET. Check out the
> [documentation](https://lab-cosmo.github.io/upet/latest/usage/nvalchemi.html)
> and try it yourself!
```py
from upet.nvalchemi import UPETWrapper
model = UPETWrapper.from_checkpoint(model="pet-mad-s", version="1.6.0", device="cuda")
```


# UPET: Universal Models for Advanced Atomistic Simulations

**UPET** is a family of universal interatomic potentials for advanced materials
modeling across the periodic table. These models are based on the **Point Edge
Transformer (PET)** architecture trained on various popular atomistic datasets,
and they are capable of predicting energies and forces in complex atomistic
workflows. The package also ships **PET-MAD-DOS**, a universal model for
predicting the electronic density of states (DOS) of materials and molecules,
as well as their Fermi levels and bandgaps.

📖 **Full documentation:** <https://lab-cosmo.github.io/upet/latest/>

## Installation

Install UPET from PyPI:

```bash
pip install upet
```

Or directly from the GitHub repository:

```bash
pip install upet@git+https://github.com/lab-cosmo/upet.git
```

See the [installation guide](https://lab-cosmo.github.io/upet/latest/installation.html)
for additional methods (specific versions, `uv`, etc.).

## Quick start

Run a single-point evaluation with the ASE-compatible `UPETCalculator`:

```python
from upet.ase import UPETCalculator
from ase.build import bulk

atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
calculator = UPETCalculator(model="pet-mad-s", version="1.6.0", device="cpu")
atoms.calc = calculator

energy = atoms.get_potential_energy()
forces = atoms.get_forces()
```

For DOS calculations, you can use the `PETMADDOSCalculator`:

```python
from upet.ase.dos import PETMADDOSCalculator
from ase.build import bulk

atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
calculator = PETMADDOSCalculator(version="latest", device="cpu")
results = calculator.calculate(atoms)
```
where `results` is a dictionary whose keys include `dos_raw`, `dos_denoised`,
`fermi_level` and `bandgap`. Each key corresponds to its output quantity, e.g.
`dos_denoised` is the denoised DOS obtained by applying a denoising algorithm to the
raw predicted DOS. The DOS has units of states/eV and is projected on an energy grid
with intervals of 0.05 eV. The bandgap and Fermi level have units of eV.


The first call downloads the checkpoint from the
[HuggingFace repository](https://huggingface.co/lab-cosmo/upet) and caches it
locally, so subsequent calls are fast. If `version` is omitted, the latest
available version is used.

To list available models, sizes, and versions:

```python
from upet import list_upet

list_upet(model="pet-mad", size="s")
list_upet(model="pet-mad")
list_upet()
```

## Pre-trained models

UPET ships several pre-trained model families (PET-MAD, PET-OAM, PET-OMat,
PET-OMATPES, PET-SPICE) at multiple sizes. PET-MAD-1.6 (XS, S, M) is the
recommended choice for molecular dynamics, while PET-OAM is the one for
materials discovery. The full table with supported elements, training sets,
and recommended use cases is in the
[models documentation](https://lab-cosmo.github.io/upet/latest/models.html).
All checkpoints are available on the
[HuggingFace repository](https://huggingface.co/lab-cosmo/upet).

## Going further

The documentation covers the complete feature surface:

- [ASE workflows](https://lab-cosmo.github.io/upet/latest/usage/ase.html):
  non-conservative forces, uncertainty quantification, rotational averaging,
  dispersion corrections, and PET-MAD-DOS for DOS / Fermi levels / bandgaps.
- [Batched evaluation with metatrain](https://lab-cosmo.github.io/upet/latest/usage/metatrain.html).
- [LAMMPS](https://lab-cosmo.github.io/upet/latest/usage/lammps.html) (including
  KOKKOS GPU support),
  [NVIDIA ALCHEMI Toolkit](https://lab-cosmo.github.io/upet/latest/usage/nvalchemi.html)
  (NVIDIA's GPU-native batched inference and MD toolkit),
  [i-PI](https://lab-cosmo.github.io/upet/latest/usage/ipi.html),
  [TorchSim](https://lab-cosmo.github.io/upet/latest/usage/torchsim.html), and
  [GROMACS](https://lab-cosmo.github.io/upet/latest/usage/gromacs.html) interfaces.
- [Fine-tuning](https://lab-cosmo.github.io/upet/latest/fine-tuning.html) and
  [example gallery](https://lab-cosmo.github.io/upet/latest/generated_examples/index.html).
- [FAQ and known issues](https://lab-cosmo.github.io/upet/latest/faq.html).

More worked examples for **ASE, i-PI, and LAMMPS** are also available in the
[Atomistic Cookbook](https://atomistic-cookbook.org/examples/pet-mad/pet-mad.html).

## Citing UPET Models

If you found our models useful, please cite the corresponding articles. The
full list with copy-pasteable BibTeX is also available in the
[documentation](https://lab-cosmo.github.io/upet/latest/cite.html).

PET-MAD-1.6 or PET-MAD-1.5:
```bibtex
@misc{PET-MAD-1.6-2026,
      title={High-quality, high-information datasets for universal atomistic machine learning},
      author={Cesare Malosso and Filippo Bigi and Paolo Pegolo and Joseph W. Abbott and Philip Loche and Mariana Rossi and Tiago J. Goncalves and Sandip De and Michele Ceriotti and Arslan Mazitov},
      year={2026},
      eprint={2603.02089},
      archivePrefix={arXiv},
      primaryClass={cond-mat.mtrl-sci},
      url={https://arxiv.org/abs/2603.02089},
}
```

Current UPET architecture, PET-OAM, PET-OMat, PET-OMAD, PET-OMATPES, PET-SPICE:
```bibtex
@misc{pushing-unconstrained-2026,
      title={Pushing the limits of unconstrained machine-learned interatomic potentials},
      author={Filippo Bigi and Paolo Pegolo and Arslan Mazitov and Michele Ceriotti},
      year={2026},
      eprint={2601.16195},
      archivePrefix={arXiv},
      primaryClass={physics.chem-ph},
      url={https://arxiv.org/abs/2601.16195},
}
```

PET-MAD:
```bibtex
@misc{PET-MAD-2025,
      title={PET-MAD as a lightweight universal interatomic potential for advanced materials modeling},
      author={Mazitov, Arslan and Bigi, Filippo and Kellner, Matthias and Pegolo, Paolo and Tisi, Davide and Fraux, Guillaume and Pozdnyakov, Sergey and Loche, Philip and Ceriotti, Michele},
      journal={Nature Communications},
      volume={16},
      number={1},
      pages={10653},
      year={2025},
      url={https://doi.org/10.1038/s41467-025-65662-7},
}
```

PET-MAD-DOS:
```bibtex
@misc{PET-MAD-DOS-2025,
      title={A universal machine learning model for the electronic density of states},
      author={Wei Bin How and Pol Febrer and Sanggyu Chong and Arslan Mazitov and Filippo Bigi and Matthias Kellner and Sergey Pozdnyakov and Michele Ceriotti},
      year={2025},
      eprint={2508.17418},
      archivePrefix={arXiv},
      primaryClass={physics.chem-ph},
      url={https://arxiv.org/abs/2508.17418},
}
```

For a general citation for the PET architecture, you can use
```bibtex
@misc{PET-ECSE-2023,
  title = {Smooth, Exact Rotational Symmetrization for Deep Learning on Point Clouds},
  journal = {Advances in {{Neural Information Processing Systems}}},
  author = {Pozdnyakov, Sergey and Ceriotti, Michele},
  year = 2023,
  volume = {36},
  pages = {79469--79501},
}
```

## Maintainers

This project is [maintained](https://github.com/lab-cosmo/.github/blob/main/Maintainers.md) by [@abmazitov](https://github.com/abmazitov), who will reply to issues and pull requests opened on this repository as soon as possible. You can mention him directly if you have not received an answer after a couple of days.
