graph-pes is a framework built to accelerate the development of machine-learned potential energy surface (PES) models that act on graph representations of atomic structures.

Links: Google Colab Quickstart - Documentation - PyPI

Statement of need

graph-pes is a toolkit for building, training, and deploying machine-learned potential energy surfaces (PES) models that act on graph representations of atomic structures.

As a researcher who wants to train and use existing MLIPs, you can use the graph-pes-train command to train many different architectures from scratch on your own data, or fine-tune several existing foundation models. Once trained, you can use our drivers to run optimisations, single point energy calculations, and molecular dynamics simulations with your model with a variety of existing tools (LAMMPS, ASE, and torch-sim).

As a researcher wanting to work on MLIP methodology, graph-pes makes implementing new architectures easy, allows you to experiment with various different training strategies, and provides a clean, well-documented API for building things yourself.

Features

• Experiment with new model architectures by inheriting from our GraphPESModel base class.
• Train your own or existing model architectures (e.g., SchNet, NequIP, PaiNN, MACE, TensorNet, OrB etc.).
• Use and fine-tune foundation models via a unified interface: MACE-MP0, MACE-OFF, MatterSim, GO-MACE, Egret and Orb v2/3.
• Easily configure distributed training, learning rate scheduling, weights and biases logging, and other features using our graph-pes-train command line interface.
• Use our data-loading pipeline within your own training loop.
• Run molecular dynamics simulations with any GraphPESModel using torch-sim, LAMMPS or ASE

Quickstart

pip install -q graph-pes
wget https://tinyurl.com/graph-pes-minimal-config -O config.yaml
graph-pes-train config.yaml

Alternatively, for a 0-install quickstart experience, please see this Google Colab, which you can also find in our documentation.

Contributing

Contributions are welcome! If you find any issues or have suggestions for new features, please open an issue or submit a pull request on the GitHub repository. We use uv to manage dependencies and run commands. Install it here, and sync the dependencies using uv sync --all-extras.

Once you have made your changes, you can:

• run tests locally: uv run pytest tests/
• build the documentation: uv run sphinx-build docs/source docs/build

Please see our CONTRIBUTING.md file for more details.

Citing graph-pes

We kindly ask that you cite graph-pes in your work if it has been useful to you. A manuscript is currently in preparation - in the meantime, please cite the Zenodo DOI found in the CITATION.cff file.