Directional Message Passing Neural Network (DimeNet and DimeNet++)

Reference implementation of the DimeNet model proposed in the paper:

Directional Message Passing for Molecular Graphs
by Johannes Gasteiger, Janek Groß, Stephan Günnemann
Published at ICLR 2020.

As well as DimeNet++, its significantly faster successor:

Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules
by Johannes Gasteiger, Shankari Giri, Johannes T. Margraf, Stephan Günnemann
Published at the ML for Molecules workshop, NeurIPS 2020.

Note that the author's name has changed from Johannes Klicpera to Johannes Gasteiger.

If you are interested in energy and force predictions, we now recommend our new model:
GemNet.

The new DimeNet++ model is both 8x faster and 10% more accurate, so we recommend using this model instead of the original.

There are some slight differences between this repository and the original (TF1) DimeNet model, such as slightly different training and initialization in TF2. This implementation uses orthogonal Glorot initialization in the output layer for the targets alpha, R2, U0, U, H, G, and Cv and zero initialization for Mu, HOMO, LUMO, and ZPVE. The paper only used zero initialization for the output layer.

The following table gives an overview of all MAEs:

git clone git@github.com:waqarahmed89/dimenet.git
cd dimenet
#
conda create -y -n dimenet python=3.8
conda activate dimenet
conda install -y ipykernel
ipython kernel install --user --name=dimenet
python -m pip install --upgrade pip setuptools
pip install --upgrade tensorflow==2.8 tensorflow_addons numpy==1.20 scipy sympy tqdm jupyterlab
conda install -c anaconda idna
conda install -c anaconda pyyaml

This repository contains a notebook for training the model (train.ipynb) and for generating predictions on the test set with a trained model (predict.ipynb). It also contains a script for training the model on a cluster with Sacred and SEML (train_seml.py). For faster experimentation we also offer two sets of pretrained models, which you can find in the pretrained folder.

# To run training in background
jupyter nbconvert --to notebook --execute train.ipynb 

Unfortunately there are a few issues/bugs in the code (and paper) that we can't fix without retraining the models. So far, these are:

• The second distance used for calculating the angles is switched (DimeNet).
• The envelope function is implicitly divided by the distance. This is accounted for in the radial bessel basis layer but leads to an incorrect spherical basis layer (DimeNet and DimeNet++).
• DimeNet was evaluated on MD17's Benzene17 dataset, but compared to sGDML on Benzene18, which gives sGDML an unfair advantage.
• In TensorFlow AddOns <0.12 there is a bug when checkpointing. The earlier versions require explicitly passing the _optimizer variable of the MovingAverage optimizer. This is only relevant if you actually load checkpoints from disk and continue training (DimeNet and DimeNet++).
• The radial basis functions in the interaction block actually use d_kj and not d_ji. The best way to fix this is by just using d_ji instead of d_kj in the SBF and leaving the RBF unchanged (DimeNet and DimeNet++).