graph-pes

graph-pes is a set of tools designed to accelerate the development of machine-learned potential energy surfaces (ML-PESs) that act on graph representations of atomic structures.

A 3-in-1 toolset:

graph_pes: a Python library of components for training ML-PESs

This includes:

• a data pipeline for turning ase.Atoms objects into AtomicGraph objects

• a set of operations for common tasks on AtomicGraph objects, including edge-trimming, neighbour indexing, summations, and batching,

• a base class for all model of the potential energy surface (GraphPESModel) that automatically infers force and stress predictions if these are not provided by the model’s implementation

• reference implementations of popular models, including NequIP, PaiNN, MACE and TensorNet

graph_pes is written in vanilla PyTorch that is 100% compatible with TorchScript compilation for use within LAMMPS

graph-pes-train: a command line tool for training ML-PESs

• get up and running quickly with sensible defaults to train new models from scratch

• fine-tune existing models on new datasets

• easily configure advanced features such as distributed training, learning rate scheduling, stochastic weight averaging, and logging to Weights & Biases

pair_style graph_pes: a LAMMPS pair style for GPU-accelerated MD

• use this to drive GPU-accelerated molecular dynamics (MD) simulations with any model that inherits from GraphPESModel (i.e. both ones we’ve implemented and also your own)

• we’ve included helper scripts to automate the LAMMPS build process for you

Quick-start

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

Alternatively, open any of these notebooks to get started. Install graph-pes to follow along locally, or run the code in the cloud using Google Colab.

Train a model

Analyse a model

Fine-tune a model

Run LAMMPS MD

Implement your own model

Implement your own training loop

Installation

We recommend installing graph-pes in a new environment, e.g. using conda:

conda create -n graph-pes python=3.10 -y
conda activate graph-pes

Install graph-pes from PyPI using pip (installs all dependencies):

pip install graph-pes

Please see the GitHub repository for the source code.