README.md

3D scenario Neural Atoms

Models for which Neural Atoms is currently implemented:

• SchNet [arXiv] [code]
• DimeNet++ [arXiv] [code]
• PaiNN [arXiv] [code]

Currently supported datasets:

• OE62 [arXiv] [download]

Installation

To setup a conda environment with the required dependencies, please follow the OCP installation instructions. They should work identically in this repository. We further recommend installing the jupyter package to access our example training and evaluation notebooks, as well as the seml package [github] to run and manage (especially longer) experiments from the CLI. To reproduce the long-range binning analyses from the Ewald message passing paper, please install the simple-dftd3 package [installation instructions] including the Python API.

Commands

For SchNet model.

CUDA_VISIBLE_DEVICES=0 python train_and_evaluate.py --cfg=schnet_oe62_na.yml

For PaiNN model.

CUDA_VISIBLE_DEVICES=0 python train_and_evaluate.py --cfg=painn_oe62_na.yml

For DimeNet++ model.

CUDA_VISIBLE_DEVICES=0 python train_and_evaluate.py --cfg=dpp_oe62_na.yml

One can modify the configs in 3D_Molecule/configs_oe62.

The NeuralAtom block is in 3D_Molecule/ocpmodels/models/neural_atom_block.py, while the attention components is in 3D_Molecule/ocpmodels/models/na_pooling.py.

To change the number of NeuralAtoms, please refers to lines 65 and 66 in neural_atom_block.py.

Acknowledge

This repo is based on the Ewald Message Passing, we thanks the author for publishing their codes and data.