NequIP

NequIP is an open-source code for building E(3)-equivariant interatomic potentials.

PLEASE NOTE: the NequIP code is under active development and is still in beta versions 0.x.x. In general changes to the patch version (the third number) indicate backward compatible beta releases, but please be aware that file formats and APIs may change. Bug reports are also welcomed in the GitHub issues!

Installation

NequIP requires:

Python >= 3.6
PyTorch >= 1.8

To install:

Install PyTorch Geometric, following their installation instructions and making sure to install with the correct version of CUDA. (Please note that torch_geometric>=1.7.1) is required.)
Install our fork of pytorch_ema for using an Exponential Moving Average on the weights:

$ pip install "git+https://github.com/Linux-cpp-lisp/pytorch_ema@context_manager#egg=torch_ema"

We use Weights&Biases to keep track of experiments. This is not a strict requirement — you can use our package without it — but it may make your life easier. If you want to use it, create an account here and install the Python package:

pip install wandb

Install NequIP

git clone https://github.com/mir-group/nequip.git
cd nequip
pip install .

Installation Issues

We recommend running the tests using pytest:

pip install pytest
pytest ./tests/

While the tests are somewhat compute intensive, we've known them to hang on certain systems that have GPUs. If this happens to you, please report it along with information on your software environment in the Issues!

Usage

! PLEASE NOTE: the first few training epochs/calls to a NequIP model can be painfully slow. This is expected behaviour as the profile-guided optimization of TorchScript models takes a number of calls to warm up before optimizing the model. This occurs regardless of whether the entire model is compiled because many core components from e3nn are compiled and optimized through TorchScript.

Basic network training

To train a network, you run nequip-train with a YAML config file that describes your data set, model hyperparameters, and training options.

$ nequip-train configs/example.yaml

A number of example configuration files are provided:

configs/minimal.yaml: A minimal example of training a toy model on force data.
configs/minimal_eng.yaml: The same, but for a toy model that predicts and trains on only energy labels.
configs/example.yaml: Training a more realistic model on forces and energies.
configs/full.yaml: A complete configuration file containing all available options along with documenting comments.

Training runs can be restarted using nequip-restart; training that starts fresh or restarts depending on the existance of the working directory can be launched using nequip-requeue. All nequip-* commands accept the --help option to show their call signatures and options.

Evaluating trained models (and their error)

The nequip-evaluate command can be used to evaluate a trained model on a specified dataset, optionally computing error metrics or writing the results to an XYZ file for further processing.

The simplest command is:

$ nequip-evaluate --train-dir /path/to/training/session/

which will evaluate the original training error metrics over any part of the original dataset not used in the training or validation sets.

For more details on this command, please run nequip-evaluate --help.

Deploying models

The nequip-deploy command is used to deploy the result of a training session into a model that can be stored and used for inference. It compiles a NequIP model trained in Python to TorchScript. The result is an optimized model file that has no dependency on the nequip Python library, or even on Python itself:

nequip-deploy build path/to/training/session/ path/to/deployed.pth

For more details on this command, please run nequip-deploy --help.

Using models in Python

Both deployed and undeployed models can be used in Python code; for details, see the end of the Developer's tutorial mentioned again below.

An ASE calculator is also provided in nequip.dynamics.

LAMMPS Integration

NequIP is integrated with the popular Molecular Dynamics code LAMMPS which allows for MD simulations over large time- and length-scales and gives users access to the full suite of LAMMPS features.

The interface is implemented as pair_style nequip. Using it requires two simple steps:

Deploy a trained NequIP model, as discussed above.

nequip-deploy build path/to/training/session/ path/to/deployed.pth

The result is an optimized model file that has no Python dependency and can be used by standalone C++ programs such as LAMMPS.

Change the LAMMPS input file to the nequip pair_style and point it to the deployed NequIP model:

pair_style	nequip
pair_coeff	* * deployed.pth

For installation instructions, please see the pair_nequip repository.

Developer's tutorial

A more in-depth introduction to the internals of NequIP can be found in the tutorial notebook. This notebook discusses theoretical background as well as the Python interfaces that can be used to train and call models.

Please note that for most common usecases, including customized models, the nequip-* commands should be prefered for training models.

References & citing

The theory behind NequIP is described in our preprint (1). NequIP's backend builds on e3nn, a general framework for building E(3)-equivariant neural networks (2). If you use this repository in your work, please consider citing NequIP (1) and e3nn (3):

Authors

NequIP is being developed by:

Simon Batzner
Albert Musaelian
Lixin Sun
Anders Johansson
Mario Geiger
Tess Smidt

under the guidance of Boris Kozinsky at Harvard.

Contact & questions

If you have questions, please don't hesitate to reach out at batzner[at]g[dot]harvard[dot]edu.

If you find a bug or have a proposal for a feature, please post it in the Issues. If you have a question, topic, or issue that isn't obviously one of those, try our GitHub Disucssions.

Xiangyan93 / nequip