graphite

A repository for implementing graph network models based on atomic structures.

Implemented or replicated works

• ALIGNN
  • Atomistic Line Graph Neural Network for improved materials property predictions
• ALIGNN-d (our work, see demo here)
  • Efficient and interpretable graph network representation for angle-dependent properties applied to optical spectroscopy
• Gated GCN
  • Benchmarking Graph Neural Networks
• NequIP (code implementation adopted from e3nn)
  • E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials
• Atomic Structure denoiser (our work, see demo here)
  • Score-based denoising for atomic structure identification
• MeshGraphNets
  • Learning Mesh-Based Simulation with Graph Networks

Installation

The installation time is typically within 10 minutes on a normal local machine.

Dependencies:

• PyTorch (pytorch>=1.8.1)
• PyTorch-Geometric (pyg>=2.0.1): for implementing graph network operations.
• [Optional] Atomic Simulation Environment (ase): for reading/writing atomic structures.
• [Optional] Euclidean neural networks (e3nn>=0.4.4): dependency for the NequIP models.

An example for the installation process:

conda create -n graphite
conda activate graphite
conda install pytorch pytorch-cuda=11.7 -c pytorch -c nvidia
conda install pyg -c pyg

# Optional install dependency, but required for some of the model implementations
pip install ase e3nn

# Other useful packages for development (optional)
pip install jupyterlab seaborn lightning tensorboard MDAnalysis

Then, to install graphite, clone this repo and run:

pip install -e /path/to/the/repo

The -e option signifies an editable install, which is well suited for development; this allows you to edit the source code without having to re-install.

To uninstall:

pip uninstall graphite

How to use

graphite is intended to be a general collection of codes (e.g., helper functions, custom graph convolutions, and template graph models) for research purposes. Production codes for certain applications and deployments should be hosted elsewhere.

• The src folder contains the source code.
• The notebooks folder contains Jupyter notebooks that demonstrate running or training models.
  • Some demos require additional packages (e.g., PyTorch Lightning for automated training). Please see Installation and the instructions in the demos.

Release

LLNL-CODE-836648