Skip to content
/ graphite Public

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

License

Notifications You must be signed in to change notification settings

LLNL/graphite

Repository files navigation

graphite

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

Implemented or replicated works

Gallery

Spectroscopy-guided amorphous material generation

Unconditional generation of amorphous carbons via diffusion model. Color is meant to help give you a sense of depth.

Conditional generation of amorphous carbons based on a given XANES spectrum.

Generation of multi-element, amorphous carbon nitrides.

Score dynamics

Generative rollouts of molecular dynamics over picosecond timesteps via conditional diffusion model.

Atomic structure denoiser

Simple and effective atomic denoiser for structure characterization.

Installation

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

Installation dependencies:

  • pytorch>=2.0.1
  • torch_geometric
  • torch-scatter
  • torch-cluster

Reasons: many model implementations in this repo are written mostly based on PyTorch, but some operations such as scatter and graph pooling are from PyTorch Geometric (PyG), which offers very optimized CUDA implementations. Additionally, most models in this repo treat atomic/molecular data as graphs following the PyG data format (torch_geometric.data.Data).

Also, this repo has scattering (via Torch.scatter_reduce) and clustering (e.g., radius_graph and knn_graph) codes such that torch-scatter and torch-cluster are not strictly required. Still, it is recommended to install torch-scatter and torch-cluster if you favor CUDA-optimized compute speed.

Lastly, for development purposes, you may want to install packages such as ase, MDAnalysis, rdkit, lightning, etc.

Example installation process:

conda create -n graphite
conda activate graphite
conda install pytorch pytorch-cuda=12.1 -c pytorch -c nvidia
pip install torch_geometric
pip install torch_scatter torch_cluster -f https://data.pyg.org/whl/torch-2.3.0+cu121.html

# For development depending on your use case
pip install ase jupyterlab ipywidgets seaborn lightning tensorboard

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

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

To uninstall:

pip uninstall graphite

How to use

graphite is intended for research and prototyping. It is a general collection of simple codes containing helper functions, custom graph convolutions, model templates, and so on. Full-scale production codes for specific 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).
    • The demo notebooks are not always up-to-date. We will try to update the notebooks after every major change to the source code.
  • The media folder contains media files.

Release

LLNL-CODE-836648

About

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

Resources

License

Code of conduct

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published