GemNet: Universal Directional Graph Neural Networks for Molecules
Reference implementation in PyTorch of the geometric message passing neural network (GemNet). You can find its original TensorFlow 2 implementation in another repository. GemNet is a model for predicting the overall energy and the forces acting on the atoms of a molecule. It was proposed in the paper:
GemNet: Universal Directional Graph Neural Networks for Molecules
by Johannes Gasteiger, Florian Becker, Stephan Günnemann
Published at NeurIPS 2021
and further analyzed in
How Robust are Modern Graph Neural Network Potentials in Long and Hot Molecular Dynamics Simulations?
by Sina Stocker*, Johannes Gasteiger*, Florian Becker, Stephan Günnemann and Johannes T. Margraf
2022
*Both authors contributed equally to this research. Note that the author's name has changed from Johannes Klicpera to Johannes Gasteiger.
Run the code
Adjust config.yaml (or config_seml.yaml) to your needs.
This repository contains notebooks for training the model (train.ipynb) and for generating predictions on a molecule loaded from ASE (predict.ipynb). It also contains a script for training the model on a cluster with Sacred and SEML (train_seml.py). Further, a notebook is provided to show how GemNet can be used for MD simulations (ase_example.ipynb).
Compute scaling factors
You can either use the precomputed scaling_factors (in scaling_factors.json) or compute them yourself by running fit_scaling.py. Scaling factors are used to ensure a consistent scale of activations at initialization. They are the same for all GemNet variants.
Contact
Please contact [email protected] if you have any questions.
Cite
Please cite our papers if you use the model or this code in your own work:
@inproceedings{gasteiger_gemnet_2021,
title = {GemNet: Universal Directional Graph Neural Networks for Molecules},
author = {Gasteiger, Johannes and Becker, Florian and G{\"u}nnemann, Stephan},
booktitle={Conference on Neural Information Processing Systems (NeurIPS)},
year = {2021}
}
@article{stocker_gnn_2022,
title = {How Robust are Modern Graph Neural Network Potentials in Long and Hot Molecular Dynamics Simulations?},
author = {Stocker, Sina and Gasteiger, Johannes and Becker, Florian and G{\"u}nnemann, Stephan and Margraf, Johannes T.},
year = {2022}
}