vish317 / physics-aware-multiplex-gnn

Code for our Nature Scientific Reports paper "A universal framework for accurate and efficient geometric deep learning of molecular systems"

PAMNet: A Universal Framework for Accurate and Efficient Geometric Deep Learning of Molecular Systems

Official implementation of PAMNet (Physics-aware Multiplex Graph Neural Network) in our paper A universal framework for accurate and efficient geometric deep learning of molecular systems accepted by Nature Scientific Reports (doi: 10.1038/s41598-023-46382-8).

PAMNet is an improved version of MXMNet and outperforms state-of-the-art baselines regarding both accuracy and efficiency in diverse tasks including small molecule property prediction, RNA 3D structure prediction, and protein-ligand binding affinity prediction.

This implementation is also applicable to:

1. Our preprint Efficient and Accurate Physics-aware Multiplex Graph Neural Networks for 3D Small Molecules and Macromolecule Complexes.
2. Our paper Physics-aware Graph Neural Network for Accurate RNA 3D Structure Prediction on Machine Learning for Structural Biology Workshop at NeurIPS 2022.

If you have any questions, feel free to open an issue or reach out to: [email protected].

Overall Architecture

Setup

Environment

• Python : 3.7.4
• CUDA : 10.1

Optional: Install Open Babel 3.1.1 for binding affinity prediction on PDBbind:

1. Download source file
2. conda install filename

The other dependencies can be installed with:

pip install -r requirements.txt

Datasets

QM9 for small molecule property prediction:

The training script (main_qm9.py) will automatically download the QM9 dataset and preprocess it.

PDBbind for protein-ligand binding affinity prediction:

1. Download PDBbind_dataset.tar.gz from dropbox
2. Unzip the downloaded file under ./data/PDBbind. There will be two subfolders (core-set and refined-set) after the unzip
3. Run python preprocess_pdbbind.py to preprocess the dataset to construct graphs

RNA-Puzzles for RNA 3D structure prediction:

1. Download classics_train_val.tar from Stanford Digital Repository
2. Unzip the downloaded file under ./data/RNA-Puzzles. There will be one subfolder classics_train_val containing example_train and example_valafter the unzip
3. Run python preprocess_rna_puzzles.py to preprocess the dataset to construct graphs

How to Run

Arguments

  --gpu             GPU number
  --seed            random seed
  --dataset         dataset to be used
  --epochs          number of epochs to train
  --lr              initial learning rate
  --wd              weight decay value
  --n_layer         number of hidden layers
  --dim             size of input hidden units
  --batch_size      batch size
  --cutoff_l        distance cutoff used in the local layer
  --cutoff_g        distance cutoff used in the global layer
  --model           model to be used on QM9
  --target          index of target (0~11) for prediction on QM9

Example command for training and evaluation

Small molecule property prediction on QM9:

python -u main_qm9.py --dataset 'QM9' --model 'PAMNet' --target=7 --epochs=900 --batch_size=32 --dim=128 --n_layer=6 --lr=1e-4

Protein-ligand binding affinity prediction on PDBbind:

python -u main_pdbbind.py --dataset 'PDBbind' --epochs=170 --batch_size=32 --dim=128 --n_layer=3 --lr=1e-3

RNA 3D structure prediction on RNA-Puzzles:

python -u main_rna_puzzles.py --dataset 'RNA-Puzzles' --epochs=15 --batch_size=8  --dim=16 --n_layer=1 --lr=1e-4

Citation

If you find our model and code helpful in your work, please consider citing us:

@article{zhang2023universal,
  title={A Universal Framework for Accurate and Efficient Geometric Deep Learning of Molecular Systems},
  author={Zhang, Shuo and Liu, Yang and Xie, Lei},
  journal={Scientific Reports},
  volume={13},
  number={1},
  pages={19171},
  year={2023},
  publisher={Nature Publishing Group UK London}
}

@article{zhang2022physics,
  title={Physics-aware graph neural network for accurate RNA 3D structure prediction},
  author={Zhang, Shuo and Liu, Yang and Xie, Lei},
  journal={arXiv preprint arXiv:2210.16392},
  year={2022}
}