| | --- |
| | license: cc-by-nc-sa-2.0 |
| | --- |
| | <div align='center'> |
| | <h1> MolCRAFT Series for Drug Design </h1> |
| |
|
| | [](https://github.com/AlgoMole/MolCRAFT/tree/master) |
| | [](https://MolCRAFT-GenSI.github.io/) |
| | [](https://drive.google.com/drive/folders/16KiwfMGUIk4a6mNU20GnUd0ah-mjNlhC?usp=share_link) |
| |
|
| | </div> |
| |
|
| |
|
| | Welcome to the official repository for the MolCRAFT series of projects! This series focuses on developing and improving deep learning models for **structure-based drug design (SBDD)** and **molecule optimization (SBMO)**. Our goal is to create molecules with high binding affinity and plausible 3D conformations. |
| |
|
| | This repository contains the source code for the following projects: |
| |
|
| | * [**MolCRAFT**: Structure-Based Drug Design in Continuous Parameter Space](https://arxiv.org/abs/2404.12141) (ICML'24) |
| | * [**MolJO**: Empower Structure-Based Molecule Optimization with Gradient Guided Bayesian Flow Networks](https://arxiv.org/abs/2411.13280) (ICML'25) |
| | * [**MolPilot**: Piloting Structure-Based Drug Design via Modality-Specific Optimal Schedule](https://arxiv.org/abs/2505.07286) (ICML'25) |
| |
|
| | ## 📜 Overview |
| |
|
| | The MolCRAFT series addresses critical challenges in generative models for SBDD, including modeling molecular geometries, handling hybrid continuous-discrete spaces, and optimizing molecules against protein targets. Each project introduces novel methodologies and achieves **state-of-the-art** performance on relevant benchmarks. |
| |
|
| | ## 🧭 Navigation |
| |
|
| | | Folder | TL, DR | Description | |
| | | --------------------------- | --------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | |
| | | [MolCRAFT](https://huggingface.co/GenSI/MolCRAFT) | Unified Space for Molecule Generation | MolCRAFT is the first SBDD generative model based on Bayesian Flow Network (BFN) operating in the unified continuous parameter space for different modalities, with variance reduction sampling strategy to generate high-quality samples with more than 10x speedup. |
| | | [MolJO](https://huggingface.co/GenSI/MolJO) | Gradient-Guided Molecule Optimization | MolJO is a gradient-based Structure-Based Molecule Optimization (SBMO) framework derived within BFN. It employs joint guidance across continuous coordinates and discrete atom types, alongside a backward correction strategy for effective optimization. |
| | | [MolPilot](https://huggingface.co/GenSI/MolPilot) | Optimal Scheduling | MolPilot enhances SBDD by introducing a VLB-Optimal Scheduling (VOS) strategy for the twisted multimodal probability paths, significantly improving molecular geometries and interaction modeling, achieving 95.9% PB-Valid rate. | |
| |
|
| |
|
| | --- |
| |
|
| | ## 🚀 Projects |
| |
|
| | ### MolCRAFT (Let's Craft the Molecules) |
| |
|
| | <!-- <p align="center"><img src="asset/molcraft.gif" width="60%"></p> --> |
| |
|
| | * **Description**: MolCRAFT is the first SBDD model that employs BFN and operates in a **continuous parameter space**. It introduces a novel noise-reduced sampling strategy to generate molecules with superior binding affinity and more stable 3D structures. MolCRAFT has demonstrated its ability to accurately model interatomic interactions, achieving reference-level Vina Scores. |
| | * **Key Contributions**: |
| | * Operates in continuous parameter space for SBDD within BFN framework. |
| | * Novel variance reduction sampling strategy that improves both sample quality and efficiency. |
| | * Achieves state-of-the-art binding affinity and structural stability. |
| |
|
| | --- |
| |
|
| | ### MolJO (Molecule Joint Optimization) |
| |
|
| | <!--  --> |
| |
|
| | * **Description**: MolJO is a gradient-based SBMO framework that leverages a continuous and differentiable space derived through Bayesian inference. It facilitates **joint guidance signals across different modalities** (continuous coordinates and discrete atom types) while preserving SE(3)-equivariance. MolJO introduces a novel backward correction strategy for an effective trade-off between exploration and exploitation. |
| | * **Key Contributions**: |
| | * Gradient-based SBMO framework with joint guidance across different modalities. |
| | * Backward correction strategy for optimized exploration-exploitation. |
| | * State-of-the-art performance in practical optimization tasks, including multi-objective and constrained optimization for R-group redesign, scaffold hopping, etc. |
| |
|
| | --- |
| |
|
| | ### MolPilot (How to Pilot the Aircraft) |
| |
|
| | <!--  --> |
| |
|
| | * **Description**: MolPilot addresses challenges in geometric structure modeling by focusing on the **twisted probability path of multi-modalities** (continuous 3D positions and discrete 2D topologies). It proposes a VLB-Optimal Scheduling (VOS) strategy, optimizing the Variational Lower Bound as a path integral for SBDD. MolPilot significantly enhances molecular geometries and interaction modeling. |
| | * **Key Contributions**: |
| | * Addresses multi-modality challenges in SBDD. |
| | * Introduces VLB-Optimal Scheduling (VOS) strategy, generally applicable to a wide range of frameworks including diffusions. |
| | * Achieves 95.9% PoseBusters passing rate on CrossDock with significantly improved molecular geometries. |
| |
|
| | --- |
| | # MolCRAFT |
| | Official implementation of ICML 2024 ["MolCRAFT: Structure-Based Drug Design in Continuous Parameter Space"](https://arxiv.org/abs/2404.12141). |
| |
|
| | 🎉 Our demo is now available [here](http://61.241.63.126:8000). Welcome to have a try! |
| |
|
| | <!--  --> |
| |
|
| | ## Environment |
| | It is highly recommended to install via docker if a Linux server with NVIDIA GPU is available. |
| |
|
| | Otherwise, you might check [README for env](docker/README.md) for further details of docker or conda setup. |
| |
|
| | ### Prerequisite |
| | A docker with `nvidia-container-runtime` enabled on your Linux system is required. |
| |
|
| | > [!TIP] |
| | > - This repo provides an easy-to-use script to install docker and nvidia-container-runtime, in `./docker` run `sudo ./setup_docker_for_host.sh` to set up your host machine. |
| | > - For details, please refer to the [install guide](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html). |
| | |
| | |
| | ### Install via Docker |
| | We highly recommend you to set up the environment via docker, since all you need to do is a simple `make` command. |
| | ```bash |
| | cd ./docker |
| | make |
| | ``` |
| | |
| | |
| | ----- |
| | ## Data |
| | We use the same data as [TargetDiff](https://github.com/guanjq/targetdiff/tree/main?tab=readme-ov-file#data). Data used for training / evaluating the model should be put in the `data` folder by default, and accessible in the [data](https://drive.google.com/drive/folders/1j21cc7-97TedKh_El5E34yI8o5ckI7eK?usp=share_link) Google Drive folder. |
| | |
| | To train the model from scratch, download the lmdb file and split file into data folder: |
| | * `crossdocked_v1.1_rmsd1.0_pocket10_processed_final.lmdb` |
| | * `crossdocked_pocket10_pose_split.pt` |
| |
|
| | To evaluate the model on the test set, download _and_ unzip the `test_set.zip` into data folder. It includes the original PDB files that will be used in Vina Docking. |
| |
|
| |
|
| | --- |
| | ## Training |
| | Run `make -f scripts.mk` (without the need for data preparation), or alternatively (with data folder correctly configured), |
| | ```bash |
| | python train_bfn.py --exp_name ${EXP_NAME} --revision ${REVISION} |
| | ``` |
| |
|
| | where the default values should be set the same as: |
| | ```bash |
| | python train_bfn.py --sigma1_coord 0.03 --beta1 1.5 --lr 5e-4 --time_emb_dim 1 --epochs 15 --max_grad_norm Q --destination_prediction True --use_discrete_t True --num_samples 10 --sampling_strategy end_back_pmf |
| | ``` |
| |
|
| | ### Testing |
| | For quick evaluation of the official checkpoint, refer to `make evaluate` in `scripts.mk`: |
| | ```bash |
| | python train_bfn.py --test_only --no_wandb --ckpt_path ./checkpoints/${CKPT_NAME} |
| | ``` |
| |
|
| | ### Debugging |
| | For quick debugging training process, run `make debug -f scripts.mk`: |
| | ```bash |
| | python train_bfn.py --no_wandb --debug --epochs 1 |
| | ``` |
| |
|
| | ## Sampling |
| | We provide the pretrained MolCRAFT checkpoint [here](https://drive.google.com/file/d/1TcUQM7Lw1klH2wOVBu20cTsvBTcC1WKu/view?usp=share_link). |
| |
|
| |
|
| | ### Sampling for pockets in the testset |
| | Run `make evaluate -f scripts.mk`, or alternatively, |
| | ```bash |
| | python train_bfn.py --config_file configs/default.yaml --exp_name ${EXP_NAME} --revision ${REVISION} --test_only --num_samples ${NUM_MOLS_PER_POCKET} --sample_steps 100 |
| | ``` |
| |
|
| | The output molecules `vina_docked.pt` for all 100 test pockets will be saved in `./logs/${USER}_bfn_sbdd/${EXP_NAME}/${REVISION}/test_outputs/${TIMESTAMP}` folders. |
| |
|
| | ### Sampling from pdb file |
| | To sample from a whole protein pdb file, we need the corresponding reference ligand to clip the protein pocket (a 10A region around the reference position). |
| |
|
| | Below is an example that stores the generated 10 molecules under `output` folder. The configurations are managed in the ``call()`` function of ``sample_for_pocket.py``. |
| |
|
| | ```bash |
| | python sample_for_pocket.py ${PDB_PATH} ${SDF_PATH} |
| | ``` |
| |
|
| | ## Evaluation |
| | ### Evaluating molecules |
| | For binding affinity (Vina Score / Min / Dock) and molecular properties (QED, SA), it is calculated upon sampling. |
| |
|
| | For PoseCheck (strain energy, clashes) and other conformational results (bond length, bond angle, torsion angle, RMSD), please refer to `test` folder. |
| |
|
| | ### Evaluating meta files |
| | We provide samples for all SBDD baselines in the [sample](https://drive.google.com/drive/folders/1A3Mthm9ksbfUnMCe5T2noGsiEV1RfChH?usp=sharing) Google Drive folder. |
| |
|
| | You may download the `all_samples.tar.gz` and then `tar xzvf all_samples.tar.gz`, which extracts all the pt files into `samples` folder for evaluation. |
| |
|
| | <!-- ## Demo |
| | ### Host our web app demo locally |
| |
|
| | With ``gradio`` and ``gradio_molecule3d`` installed, you can simply run ``python app.py`` to open the demo locally. Port mapping has been set in Makefile if you are using docker. You should also forward this port if you run the docker in an ssh server. We will share a permanent demo link later. |
| |
|
| | Great thanks to @duerrsimon for his kind support in resolving rendering issues! --> |
| |
|
| | ## Citation |
| |
|
| | ``` |
| | @article{qu2024molcraft, |
| | title={MolCRAFT: Structure-Based Drug Design in Continuous Parameter Space}, |
| | author={Qu, Yanru and Qiu, Keyue and Song, Yuxuan and Gong, Jingjing and Han, Jiawei and Zheng, Mingyue and Zhou, Hao and Ma, Wei-Ying}, |
| | journal={ICML 2024}, |
| | year={2024} |
| | } |
| | |
| | @article{song2024unified, |
| | title={Unified Generative Modeling of 3D Molecules via Bayesian Flow Networks}, |
| | author={Song, Yuxuan and Gong, Jingjing and Qu, Yanru and Zhou, Hao and Zheng, Mingyue and Liu, Jingjing and Ma, Wei-Ying}, |
| | journal={ICLR 2024}, |
| | year={2024} |
| | } |
| | ``` |
