This is the official code repository of our ICML paper "An End-to-End Framework for Molecular Conformation Generation via Bilevel Programming" (2021).
You can follow the instructions here. We adopt an environment the same as our another previous project.
# Create conda environment
conda create --name ConfVAE python=3.7
# Activate the environment
conda activate ConfVAE
# Install packages
conda install pytorch torchvision cudatoolkit=10.2 -c pytorch
conda install rdkit==2020.03.3 -c rdkit
conda install tqdm networkx scipy scikit-learn h5py tensorboard -c conda-forge
pip install torchdiffeq==0.0.1
# Install PyTorch Geometric
conda install pytorch-geometric -c rusty1s -c conda-forgeThe official datasets are available here.
The dataset file is a pickled Python list consisting of rdkit.Chem.rdchem.Mol objects. Each conformation is stored individually as a Mol object. For example, if a dataset contains 3 molecules, where the first molecule has 4 conformations, the second one and the third one have 5 and 6 conformations respectively, then the pickled Python list will contain 4+5+6 Mol objects in total.
The output format is identical to the input format.
Example: training a model for QM9 molecules.
python train_vae.py \
--train_dataset ./data/qm9/train_QM9.pkl \
--val_dataset ./data/qm9/val_QM9.pklMore training options can be found in train_vae.py.
Example: generating conformations for each molecule in the QM9 test-split, with twice the number of test set for each molecule.
python eval_vae.py \
--ckpt ./logs/VAE_QM9 \
--dataset ./data/iclr/qm9/test_QM9.pkl \
--num_samples -2More generation options can be found in eval_vae.py.
Please consider citing our work if you find it helpful.
@inproceedings{
xu2021end,
title={An End-to-End Framework for Molecular Conformation Generation via Bilevel Programming},
author={Xu, Minkai and Wang, Wujie and Luo, Shitong and Shi, Chence and Bengio, Yoshua and Gomez-Bombarelli, Rafael and Tang, Jian},
booktitle={International Conference on Machine Learning},
year={2021}
}
If you have any question, please contact me at minkai.xu@umontreal.ca or xuminkai@mila.quebec.
Please also check our another concurrent work on molecular conformation generation, which has also been accepted in ICML'2021 (Long Talk): Learning Gradient Fields for Molecular Conformation Generation. [Code]
