tencent-ailab/mdm

MDM

MDM: Molecular Diffusion Model for 3D Molecule Generation

[Paper]

📢 News

• If you are interested in generating molecules based on the protein pocket structure, please refer to our new paper 'PMDM' which is recently accepted by Nature Communications !!

Dependencies

• RDKit
• PyTorch
• Scipy
• torch-scatter
• torch-geometric

You can also clone our environment file.

# Clone the environment
conda env create -f MDM.yml
# Activate the environment
conda activate MDM

Data preparation

QM9 dataset

Download the dataset and split it.

cd ./qm9/data/prepare/
python ./qm9/data/prepare/download.py

You can also download the data from https://drive.google.com/file/d/1JZ_Z5bjS0RsX_BRWtrplMN9vZpL78-T7/view?usp=drive_link and put it under data/QM9/qm9/raw.

Geom dataset

1. Download the file at https://dataverse.harvard.edu/file.xhtml?fileId=4360331&version=2.0 (Warning: 50gb): wget https://dataverse.harvard.edu/api/access/datafile/4360331

2. Untar it and move it to data/geom/ tar -xzvf 4360331

3. pip install msgpack

4. python3 build_geom_dataset.py --conformations 1

Training

QM9 dataset

python train.py --config './configs/qm9_full_epoch.yml'

Geom dataset

python train.py --config './configs/geom_full.yml'

Sampling and evaluation

python test_eval.py --ckpt <checkpoint> --sampling_type generalized --w_global_pos 1 -- w_global_node 1 --w_local_pos 4 --w_local_node 5

Conditional training and sampling

Train a conditional MDM for desired properties

python train_qm9_condition.py --config './configs/qm9_full_epoch.yml' --context {property name} --config_name {config_name}

The property name includes homo | lumo | alpha | gap | mu | Cv. For example, you could set --context alpha to train MDM conditioned on alpha. MDM also supports multiple properties conditioned generation. For example, you could set --context alpha gap to train MDM conditioned on alpha and gap.

Sampling

python eval_qm9_condition_quality.py --ckpt {saved_chekpoint} --num_samples {num_samples} 

You should use the saved checkpoint of train_qm9_condition.py as {saved_checkpoint}

Sampling for evaluation

Train a specific classifier

If you would like to train the classifier by yourself

cd qm9/property_prediction

python main_qm9_prop.py --num_workers 2 --lr 5e-4 --property alpha --exp_name exp_class_alpha --model_name egnn

Sampling and calculating the MAE loss

python eval_qm9_condition.py --ckpt {saved_chekpoint} --classifiers_path {saved_cls_checkpoint}

You should use the saved checkpoint of train_qm9_condition.py as {saved_checkpoint} and the saved checkpoint of main_qm9_prop.py as {saved_cls_checkpoint}

Citation

@article{huang2022mdm,
  title={MDM: Molecular Diffusion Model for 3D Molecule Generation},
  author={Huang, Lei and Zhang, Hengtong and Xu, Tingyang and Wong, Ka-Chun},
  journal={arXiv preprint arXiv:2209.05710},
  year={2022}
}

mdm's Issues

Training Time

Dear Authors,

How long does your method need to train?

Pretrained checkpoints

Hi, thanks for this wonderful work! Would you like to release pre-trained checkpoints for QM9 and GeomDrug?

AttributeError

Hello, what could be the possible reason for this?
AttributeError: 'EGNNSparse' object has no attribute '__collect__'

Code for evaluating conditional generation on QM9

Dear author, I tried to run the code for evaluating conditional generation on QM9（eval_conditional_qm9.py）, but I found the code didn't work. Could you update the code please~
Best regards!

error when running 'test_eval.py'

Hi, thanks for your code, it's great : )
I got an error when running the test and eval command in README.md: python test_eval.py --ckpt <checkpoint> --sampling_type generalized --w_global_pos 1 -- w_global_node 1 --w_local_pos 4 --w_local_node 5.
the error is as follows:
Traceback (most recent call last):
File "/data/0shared/xxx/xxx/MDM/eval.py", line 69, in retrieve_qm9_smiles
with open(file_name, 'rb') as f:
FileNotFoundError: [Errno 2] No such file or directory: '../qm9/temp/qm9_smiles.pickle'

Hope you'll help me with it, thanks a lot!!!

Code Error

Dear Authors,

When I run your code, there exists an error:

KeyError: 'positions is not a file in the archive'

Missing commands for conditional generation tasks

Hello, thank you for sharing your code. I encountered some issues while running it and would appreciate your assistance. I noticed that the provided code includes commands for running unconditional generation tasks but lacks commands for conditional generation. Could you please provide the corresponding commands for conditional generation to help us run the code more effectively? Thank you.