smsharma / eqnn-jax

Annotated implementations of equivariant graph neural networks in Jax: EGNN, SEGNN, NequIP.

$E(3)$ Equivariant Graph Neural Networks in Jax

Implementation of $E(3)$ equivariant graph neural networks in Jax.

Models

The following equivariant models are implemented:

• EGNN (Satorras et al 2021)
• SEGNN (Brandstetter et al 2021)
• NequIP (Batzner et al 2021)

Additionally, the following non-equivariant models are implemented:

• Graph Network (Battaglia et al 2018)
• PointNet++ (Qi et al 2017)
• DiffPool (Ying et al 2018)
• Set Transformer (Lee et al 2019)

Requirements and tests

To install requirements:

pip install -r requirements.txt

To run tests (testing equivariance and periodic boundary conditions):

cd tests
pytest .

Basic usage and examples

See notebooks/examples.ipynb for example usage of GNN, SEGNN, NequIP, and EGNN.

Cosmological benchmark

The cosmological benchmarking dataset, available in TFRecord format, can be downloaded from Zenodo under the DOI 10.5281/zenodo.11479419. To download the dataset into benchmarks/galaxies/quijote_records, run:

bash benchmarks/galaxies/download_tfrecords.sh

To run the graph-level task:

python benchmarks/galaxies/train_cosmology.py

To run the node-level task:

python benchmarks/galaxies/train_velocities.py

Attribution

See CITATION.cff for citation information. The implementation of SEGNN was partially inspired by segnn-jax.