OneScience-Group/SurfDock

1---2license: mit3language:4- en5tags:6- OneScience7- SCNet8- DCU9- protein-ligand-docking10- virtual-screening11- diffusion-model12frameworks:13- PyTorch14---15 16<p align="center">17  <strong><span style="font-size: 30px;">SurfDock</span></strong>18</p>19 20# Model Introduction21 22SurfDock is a surface-informed diffusion generative model for protein–ligand complex prediction and structure-based virtual screening. The model incorporates protein surface geometry and chemical information into the diffusion-based docking process to generate and screen protein–small-molecule binding conformations.23 24Paper:25 26> **SurfDock is a surface-informed diffusion generative model for reliable and accurate protein–ligand complex prediction**  27> Duanhua Cao, Mingan Chen, Rui Zhang, et al.  28> *Nature Methods*, 2024  29> DOI: https://doi.org/10.1038/s41592-024-02516-y30 31# Model Description32 33SurfDock is a surface-informed diffusion generative model for protein–ligand complex prediction and structure-based virtual screening. The model first preprocesses the target protein structure and computes geometric and physicochemical information on the protein surface. In parallel, it uses ESM to extract protein sequence representations, providing both structural and sequence information for subsequent ligand conformation generation.34 35During inference, SurfDock uses a diffusion generative model to sample candidate ligand conformations within the protein binding site, and then evaluates and ranks the generated poses using a pose confidence model. For virtual screening tasks, a screening score model can be further applied to rescore candidate protein–ligand conformations, producing results that are better suited for downstream ranking and screening.36 37By integrating protein surface information, protein language model representations, and diffusion-based conformation generation, SurfDock can be used for protein–ligand docking, candidate binding pose generation, pose scoring, and structure-based virtual screening.38 39# Use Cases40 41| Use Case | Description |42| --- | --- |43| Protein–ligand docking | Predict the binding conformation of a ligand within a protein binding site |44| Structure-based virtual screening | Perform batch docking and scoring for small-molecule libraries |45| Ligand conformation generation | Sample candidate binding poses using a diffusion model |46| Protein surface modeling | Use protein surface geometry and physicochemical information to assist docking |47 48 49# Usage Instructions50 51## 1. OneCode Usage52 53You can use the OneCode online environment for intelligent one-click AI4S programming:54 55[Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)56 57## 2. Manual Installation and Usage58 59### Hardware Requirements60 61- SurfDock involves computationally intensive diffusion sampling, ESM representation extraction, and graph neural network computation. GPU/DCU acceleration is recommended.62 63### Environment Setup64 65#### DCU Environment66 67```bash68# Activate DTK and CONDA first69conda create -n onescience311 python=3.11 -y70conda activate onescience31171 72pip install onescience[bio] \73  -i http://mirrors.onescience.ai:3141/pypi/simple/ \74  --trusted-host mirrors.onescience.ai75```76 77### Environment Notes78 79- During actual execution, if dependencies are missing or version incompatibilities occur, refer to the dependency versions declared in `environment.yaml` in the repository root and install or adjust the corresponding dependencies as needed.80- SurfDock's protein surface processing pipeline depends on PyMesh. If the current Python version is not fully compatible with the upstream PyMesh package, compatibility adjustments should be made according to the actual call relationships.81 82 83### Weights and Data Preparation84 85#### SurfDock Model Weights86 87The current repository already includes the main model weights required for inference:88 89```text90weight/91├── docking/92│   ├── best_ema_inference_epoch_model.pt93│   └── model_parameters.yml94├── posepredict/95│   ├── best_model.pt96│   └── model_parameters.yml97└── screen/98    ├── best_model.pt99    └── model_parameters.yml100```101 102Therefore, after downloading the complete repository, you generally do not need to download the main SurfDock model weights separately.103 104Check:105 106```bash107ls -lh weight/docking/108ls -lh weight/posepredict/109ls -lh weight/screen/110```111 112#### ESM Model113 114SurfDock uses ESM to extract protein sequence representations.115 116Official installation method:117 118```bash119git clone https://github.com/facebookresearch/esm model/esm120cd model/esm121pip install -e .122cd ../..123```124 125The official inference script uses:126 127```text128esm2_t33_650M_UR50D129```130 131and runs:132 133```bash134python model/esm/scripts/extract.py \135  "esm2_t33_650M_UR50D" \136  input.fasta \137  output_dir \138  --repr_layers 33 \139  --include "per_tok" \140  --truncation_seq_length 4096141```142 143to extract protein residue-level embeddings.144 145- If the current environment cannot access the internet, it is recommended to cache the corresponding ESM weights in advance. `esm2_t33_650M_UR50D` requires the following two files:146 147```text148https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t33_650M_UR50D.pt149https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t33_650M_UR50D-contact-regression.pt150```151 152After downloading, it is recommended to place them in the current user's Torch Hub checkpoints cache directory:153 154```text155~/.cache/torch/hub/checkpoints/156```157 158The final file paths should be:159 160```text161~/.cache/torch/hub/checkpoints/esm2_t33_650M_UR50D.pt162~/.cache/torch/hub/checkpoints/esm2_t33_650M_UR50D-contact-regression.pt163```164 165#### PDBBind Data166 167If you need to retrain SurfDock, you must obtain PDBBind separately. According to the official instructions, the preprocessed data previously provided by EquiBind is no longer publicly distributed due to PDBBind licensing restrictions, so users need to obtain and process the data themselves through the official PDBBind channels.168 169After processing, place the data in:170 171```text172model/data/PDBBind_processed/173```174`model/data/splits/` contains data split information and is not equivalent to the complete PDBBind dataset.175 176## 3. Quick Start177 178### Download the Model Package179 180```bash181hf download \182  OneScience-Group/SurfDock \183  --local-dir ./SurfDock184 185cd SurfDock186```187- SurfDock uses ESM to extract protein sequence representations, so the ESM model must be downloaded separately. See the Weights and Data Preparation section for details.188 189# Example Data190 191The current repository already provides docking and screening examples:192 193```text194model/data/eval_sample_dirs/195model/data/Screen_sample_dirs/196```197 198When using your own data, organize the protein and ligand inputs according to the example directory structure, and modify the following paths in the corresponding bash scripts:199 200```text201data_dir202surface_out_dir203out_csv_file204Screen_lib_path205docking_out_dir206```207 208as needed.209 210# Inference Examples211 212The following commands are assumed to be executed from the SurfDock repository root directory.213 214## Protein–Ligand Docking Example215 216Run:217 218```bash219cd scripts/bash_scripts/test_scripts220bash eval_samples.sh221```222 223The script mainly performs the following steps automatically:224 225```text2261. Preprocess the protein structure2272. Compute the protein surface2283. Build the inference input CSV2294. Extract ESM embeddings2305. Run SurfDock diffusion sampling2316. Save docking results232```233 234Before running, it is recommended to check:235 236```bash237vim scripts/bash_scripts/test_scripts/eval_samples.sh238```239 240Pay particular attention to:241 242```text243gpu_string244data_dir245surface_out_dir246out_csv_file247esmbedding_dir248docking_out_dir249```250 251Users also need to adjust the `CUDA_VISIBLE_DEVICES` and `accelerate launch` settings in the script according to the actual adaptation method.252 253## Virtual Screening Example254 255Run:256 257```bash258cd scripts/bash_scripts/test_scripts259bash screen_pipeline.sh260```261 262Before running, check:263 264```bash265vim scripts/bash_scripts/test_scripts/screen_pipeline.sh266```267 268Modify the following as needed:269 270```text271gpu_string272data_dir273surface_out_dir274out_csv_file275esmbedding_dir276Screen_lib_path277docking_out_dir278```279 280Here, `Screen_lib_path` specifies the small-molecule library to be screened, for example the official sample:281 282```text283model/data/Screen_sample_dirs/test_samples/1a0q/1a0q_ligand_for_Screen.sdf284```285 286The screening workflow is mainly:287 288```text289Protein preprocessing290        ↓291Protein surface computation292        ↓293ESM embedding294        ↓295SurfDock generates candidate conformations296        ↓297Screen model rescoring298        ↓299Output screening results300```301 302## Skip Completed Protein Preprocessing303 304This parameter can be modified in the following two scripts:305 306```text307scripts/bash_scripts/test_scripts/eval_samples.sh308scripts/bash_scripts/test_scripts/screen_pipeline.sh309```310 311The scripts use:312 313```bash314target_have_processed=true315```316 317to control whether target preprocessing is skipped. When set to `true`, the scripts skip target protein preprocessing steps such as OpenBabel/reduce and proceed directly to subsequent surface computation, CSV construction, ESM embedding, and inference.318 319To reprocess the target protein:320 321```bash322target_have_processed=false323```324 325When set to `false`, the scripts rerun the target protein preprocessing steps.326 327## Generate ESM Embeddings Separately328 329First, build the FASTA file:330 331```bash332python model/datasets/esm_embedding_preparation.py \333  --out_file ./protein.fasta \334  --protein_ligand_csv ./input.csv335```336 337Extract ESM representations:338 339```bash340python model/esm/scripts/extract.py \341  "esm2_t33_650M_UR50D" \342  ./protein.fasta \343  ./esm_embedding_output \344  --repr_layers 33 \345  --include "per_tok" \346  --truncation_seq_length 4096347```348 349Extract pocket embeddings:350 351```bash352python model/datasets/get_pocket_embedding.py \353  --protein_pocket_csv ./input.csv \354  --embeddings_dir ./esm_embedding_output \355  --pocket_emb_save_dir ./esm_embedding_pocket_output356```357 358Merge them into the `.pt` file required for SurfDock inference:359 360```bash361python model/datasets/esm_pocket_embeddings_to_pt.py \362  --esm_embeddings_path ./esm_embedding_pocket_output \363  --output_path ./esm2_pocket_embeddings.pt364```365 366# Training Instructions367 368## Retraining SurfDock369 370Retraining requires preparing the PDBBind data first and completing protein surface and ESM embedding preprocessing.371 372The training-related scripts are located in:373 374```text375scripts/bash_scripts/train_SurfDock_docking_module/376```377 378First, prepare the ESM embedding files required for training by following the "Generate ESM Embeddings Separately" section. The training script in the current repository is located at:379 380```text381scripts/bash_scripts/train_SurfDock_docking_module/train_SurfDock.sh382```383 384After completing the ESM embeddings, check the data, model, and output paths in `train_SurfDock.sh`, then run:385 386```bash387cd scripts/bash_scripts/train_SurfDock_docking_module388bash train_SurfDock.sh389```390 391## Retraining SurfScore392 393The SurfScore training script is located at:394 395```text396scripts/bash_scripts/train_SurfScore/train_SurfScore.sh397```398 399Before running, check the PDBBind data, cache, ESM embedding, output directory, and GPU parameters in the script, then run:400 401```bash402cd scripts/bash_scripts/train_SurfScore403bash train_SurfScore.sh404```405 406# Output Description407 408## Docking Output409 410The docking results from `eval_samples.sh` are saved in the directory specified by `docking_out_dir` in the script.411 412The main outputs include:413 414```text415Generated protein–ligand conformations416Confidence/scoring results417Runtime logs418Intermediate CSV419ESM embeddings420Protein surface files421```422 423## Screening Output424 425`screen_pipeline.sh` generates the screening results in `docking_out_dir` and uses:426 427```text428weight/screen/best_model.pt429```430 431to rescore the docking poses.432 433The final results can be used to rank and screen candidate small molecules.434 435# OneScience Official Information436 437| Platform | OneScience Main Repository | Skills Repository |438| --- | --- | --- |439| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |440| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |441 442 443# Citation and License444 445- The official SurfDock source code repository is released under the **MIT License**, which permits use, modification, distribution, sublicensing, and commercial use. The original copyright notice and MIT License text should be retained when copying or distributing the software.446- PDBBind data is governed by its own license and terms of use. SurfDock's MIT License does not automatically apply to PDBBind data.447- This repository is a **DCU-adapted version** of SurfDock, with adjustments to parts of the runtime environment, dependency configuration, and execution workflow. The use of the repository code, model weights, and related data remains subject to the licenses and terms of use of their respective original projects.448