Awesome

SE(3) Equivariant Graph Neural Networks with Complete Local Frames

Reference implementation in PyTorch of the equivariant graph neural network (ClofNet). You can find the paper here.

Run the code

Build environment

for newtonian system experiments

conda create -n clof python=3.7 -y
conda activate clof
conda install -y -c pytorch pytorch=1.7.0 torchvision torchaudio cudatoolkit=10.2 -y

for conformation generation task

conda install -y -c rdkit rdkit==2020.03.2.0
conda install -y scikit-learn pandas decorator ipython networkx tqdm matplotlib
conda install -y -c conda-forge easydict
pip install pyyaml wandb
pip install torch-scatter==2.0.6 -f https://pytorch-geometric.com/whl/torch-1.7.0+cu102.html
pip install torch-sparse==0.6.8 -f https://pytorch-geometric.com/whl/torch-1.7.0+cu102.html
pip install torch-cluster==1.5.9 -f https://pytorch-geometric.com/whl/torch-1.7.0+cu102.html
pip install torch-spline-conv==1.2.0 -f https://pytorch-geometric.com/whl/torch-1.7.0+cu102.html
pip install torch-geometric==1.6.3

Newtonian many-body system

This task is inspired by (Kipf et al., 2018; Fuchs et al., 2020; Satorras et al., 2021b), where a 5-body charged system is controlled by the electrostatic force field. Note that the force direction between any two particles is always along the radial direction in the original setting. To validate the effectiveness of ClofNet on modeling arbitrary force directions, we also impose two external force fields into the original system, a gravity field and a Lorentz-like dynamical force field, which can provide more complex and dynamical force directions.
The original source code for generating trajectories comes from Kipf et al., 2018 and is modified by EGNN. We further extend the version of EGNN to three new settings, as described in Section 7.1 of our paper. We sincerely thank the solid contribution of these two works.

Create Many-body dataset

cd newtonian/dataset
bash script.sh

Run experiments

for the ES(5) setting, run

python -u main_newtonian.py --max_training_samples 3000 --norm_diff True --LR_decay True --lr 0.01 --outf saved/newtonian \
--data_mode small --decay 0.9 --epochs 400 --exp_name clof_vel_small_5body --model clof_vel --n_layers 4 --data_root <root_of_data>

for the ES(20) setting, run

python -u main_newtonian.py --max_training_samples 3000 --norm_diff True --LR_decay True --lr 0.01 --outf saved/newtonian \
--data_mode small_20body --decay 0.9 --epochs 600 --exp_name clof_vel_small_20body --model clof_vel --n_layers 4 --data_root <root_of_data>

for the G+ES(20) setting, run

python -u main_newtonian.py --max_training_samples 3000 --norm_diff True --LR_decay True --lr 0.01 --outf saved/newtonian \
--data_mode static_20body --decay 0.9 --epochs 200 --exp_name clof_vel_static_20body --model clof_vel --n_layers 4 --data_root <root_of_data>

for the L+ES(20) setting, run

python -u main_newtonian.py --max_training_samples 3000 --norm_diff True --LR_decay True --decay 0.9 --lr 0.01 --outf saved/newtonian \
--data_mode dynamic_20body --epochs 600 --exp_name clof_vel_dynamic_20body --model clof_vel --n_layers 4 --data_root <root_of_data>

Conformation Generation

Equilibrium conformation generation targets on predicting stable 3D structures from 2D molecular graphs. Following ConfGF, we evaluate the proposed ClofNet on the GEOM-QM9 and GEOM-Drugs datasets (Axelrod & Gomez-Bombarelli, 2020) as well as the ISO17 dataset (Sch¨utt et al., 2017). For the score-based generation framework, we build our algorithm based on the public codebase of ConfGF. We sincerely thank their solid contribution for this field.

Dataset

Offical Dataset: The offical raw GEOM dataset is avaiable [here].
Preprocessed dataset: We use the preprocessed datasets (GEOM, ISO17) published by ConfGF([google drive folder]).

Train

cd confgen
python -u script/train.py --config_path ./config/qm9_clofnet.yml

Generation

python -u script/gen.py --config_path ./config/qm9_clofnet.yml --generator EquiGF --eval_epoch [epoch] --start 0 --end 1

Evaluation

python -u script/get_task1_results.py --input /root/to/generation --core 10 --threshold 0.5
python -u script/get_task1_results.py --input /casp/v-hezha1/workspace/EquiNODE/code_publish/ClofNet/confgen/generation/clofnet4qm9/EquiGF_s0e1epoch398min_sig0.000repeat2.pkl --core 10 --threshold 0.5

Cite

Please cite our paper if you use the model or this code in your own work:

@inproceedings{weitao_clofnet_2021,
  title     = {{SE(3)} Equivariant Graph Neural Networks with Complete Local Frames},
  author =  {Weitao Du and
            He Zhang and
            Yuanqi Du and
            Qi Meng and
            Wei Chen and
            Nanning Zheng and
            Bin Shao and
            Tie{-}Yan Liu},
  booktitle={International Conference on Machine Learning, {ICML} 2022, 17-23 July
               2022, Baltimore, Maryland, {USA}},
  year = {2021}
}