ChemFlow: Navigating Chemical Space with Latent Flows

This repo implements the paper 🔗: Navigating Chemical Space with Latent Flows by Guanghao Wei*, Yining Huang*, Chenru Duan, Yue Song, and Yuanqi Du.

Flows can uncover meaningful structures of latent spaces learned by generative models! We propose a unifying framework to characterize latent structures by flows/diffusions for optimization and traversal.

Live Demo

Try our live demo here!

Quick Start

Install all dependencies with conda env create -f environment.yml.
- (Optional) Install AutoDock-GPU for docking binding affinity. See Notes on Compiling AutoDock-GPU.
- (Recommended) mv .env.defaults .env and specify PROJECT_PATH in .env. This is later used to run the experiments in the project root directory.
Download data and put it in the data directory.
Train the VAE model by running python experiments/train_vae.py.
- (Optional) Download our pre-trained VAE model checkpoint, see Download Data & Model Checkpoints.
For supervised learning
1. Prepare the data by running python experiments/prepare_random_data.py.
2. Train the supervised surrogate predictor by running bash experiments/supervised/train_prop_predictor.sh.
3. Train the energy network with supervised semantic guidance by running bash experiments/supervised/train_wavepde_prop.sh.
For unsupervised learning
1. Train the energy network with unsupervised diversity guidance by running python experiments/train_wavepde.py.
2. Compute the pearson correlation coefficient by running python experiments/unsupervised/corr.py. Refer to notebooks/experiments/unsupervised/corr.ipynb for more details.
3. Modify experiments/utils/traversal_step.py in place with the best correlation coefficient index.
To reproduce the experiment results from the paper, run the following commands:
- bash experiments/optimization/optimization.sh for similarity constrained optimization.
- bash experiments/optimization/uc_optim.sh for unconstrained optimization.
- python experiments/optimization/optimization_multi.py for multi-objective optimization.
- bash experiments/success_rate/success_rate.sh for molecule manipulation tasks.

Additional Arguments for the scripts

We used lightning(doc) and tap(doc) to parse the arguments. Following is an example command to pass in arguments configured by lightning:

python experiments/supervised/train_prop_predictor.py \
    -e 50 \
    --model.optimizer sgd \
    --data.n 11000 \
    --data.batch_size 100 \
    --data.binding_affinity true \
    --data.prop 1err

Download Data & Model Checkpoints

We extracted 4,253,577 molecules from the three commonly used datasets for drug discovery including MOSES, ZINC250K(download), and ChEMBL.

The processed dataset and VAE model checkpoints are available at Google Drive.
- Data processing notebooks refers to notebooks/datasets.ipynb.

Notes on Compiling AutoDock-GPU

The conda version of AutoDock-GPU is not compatible with RTX 3080 & 3090. So don't use environment.yml to install AutoDock-GPU. Make sure to follow this issue to compile the source code. A good reference for the SM code can be found here.

Some commands might be useful:

export GPU_INCLUDE_PATH=/usr/local/cuda/include
export GPU_LIBRARY_PATH=/usr/local/cuda/lib64

make DEVICE=CUDA NUMWI=128 TARGETS=86

To test if the compilation is successful, run the following command:

obabel -:"CCN(CCCCl)OC1=CC2=C(Cl)C1C3=C2CCCO3" -O demo.pdbqt -p 7.4 --partialcharge gasteiger --gen3d
autodock_gpu_128wi -M data/raw/1err/1err.maps.fld -L demo.pdbqt -s 0 -N demo

Cite Us

@misc{wei2024navigating,
      title={Navigating Chemical Space with Latent Flows}, 
      author={Guanghao Wei and Yining Huang and Chenru Duan and Yue Song and Yuanqi Du},
      year={2024},
      eprint={2405.03987},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}