87 lines
2.9 KiB
Markdown
87 lines
2.9 KiB
Markdown
|
Inverse Molecular Design with Multi-Conditional Diffusion Guidance
|
||
|
|
================================================================
|
||
|
|
|
||
|
|
Paper: https://arxiv.org/abs/2401.13858
|
||
|
|
|
||
|
|
This is the code for MCD: a Multi-Conditional Diffusion Model for inverse small molecule and polymer designs and generations. The denoising model architecture in `mcd/models` looks like:
|
||
|
|
|
||
|
|
<div style="display: flex;" markdown="1">
|
||
|
|
<img src="asset/reverse.png" style="width: 50%;" alt="Description of the first image">
|
||
|
|
<img src="asset/arch.png" style="width: 50%;" alt="Description of the second image">
|
||
|
|
</div>
|
||
|
|
|
||
|
|
|
||
|
|
## Requirements
|
||
|
|
All dependencies are specified in the `requirements.txt` file.
|
||
|
|
|
||
|
|
This code was developed and tested with Python 3.9.16, PyTorch 2.0.0, and PyG 2.3.0, pytorch-lightning 2.0.1.
|
||
|
|
|
||
|
|
For molecular generation evaluation, we should first install rdkit:
|
||
|
|
|
||
|
|
Then `fcd_torch`: `pip install fcd_torch` (https://github.com/insilicomedicine/fcd_torch).
|
||
|
|
|
||
|
|
And `mini_moses` package: `pip install git+https://github.com/igor-krawczuk/mini-moses` (https://github.com/igor-krawczuk/mini-moses),
|
||
|
|
|
||
|
|
## Usage
|
||
|
|
|
||
|
|
We could train the model on an A6000 GPU card. Here is an example to running the code for polymer graphs:
|
||
|
|
|
||
|
|
```
|
||
|
|
python main.py --config-name=config.yaml \
|
||
|
|
model.ensure_connected=True \
|
||
|
|
dataset.task_name='O2-N2-CO2' \
|
||
|
|
dataset.guidance_target='O2-N2-CO2'
|
||
|
|
```
|
||
|
|
All default configurations can be found in `configs/config.yaml`. In this example, we set `model.ensure_connected=True` to ensure that all generated components are retained during graph-to-molecule conversion (see paper Section 3.2).
|
||
|
|
|
||
|
|
Other examples for small molecule generation:
|
||
|
|
|
||
|
|
```
|
||
|
|
python main.py --config-name=config.yaml \
|
||
|
|
dataset.task_name='bace_b' \
|
||
|
|
dataset.guidance_target='Class'
|
||
|
|
|
||
|
|
python main.py --config-name=config.yaml \
|
||
|
|
dataset.task_name='bbbp_b' \
|
||
|
|
dataset.guidance_target='p_np'
|
||
|
|
|
||
|
|
python main.py --config-name=config.yaml \
|
||
|
|
dataset.task_name='hiv_b' \
|
||
|
|
dataset.guidance_target='HIV_active'
|
||
|
|
```
|
||
|
|
|
||
|
|
We could generate polymer graphs by conditioning on single gas permeability.
|
||
|
|
|
||
|
|
```
|
||
|
|
|
||
|
|
python main.py --config-name=config.yaml \
|
||
|
|
dataset.task_name='O2' \
|
||
|
|
dataset.guidance_target='O2'
|
||
|
|
|
||
|
|
python main.py --config-name=config.yaml \
|
||
|
|
dataset.task_name='N2' \
|
||
|
|
dataset.guidance_target='N2'
|
||
|
|
|
||
|
|
python main.py --config-name=config.yaml \
|
||
|
|
dataset.task_name='CO2' \
|
||
|
|
dataset.guidance_target='CO2'
|
||
|
|
```
|
||
|
|
|
||
|
|
Feel free to test the code on your own dataset!
|
||
|
|
|
||
|
|
## Citation
|
||
|
|
|
||
|
|
If you find this repository useful, please cite our paper:
|
||
|
|
|
||
|
|
```
|
||
|
|
@misc{liu2024inverse,
|
||
|
|
title={Inverse Molecular Design with Multi-Conditional Diffusion Guidance},
|
||
|
|
author={Gang Liu and Jiaxin Xu and Tengfei Luo and Meng Jiang},
|
||
|
|
year={2024},
|
||
|
|
eprint={2401.13858},
|
||
|
|
archivePrefix={arXiv},
|
||
|
|
primaryClass={cs.LG}
|
||
|
|
}
|
||
|
|
```
|
||
|
|
|