--- license: mit language: - en tags: - OneScience - SCNet - DCU - protein-ligand-docking - virtual-screening - diffusion-model frameworks: - PyTorch ---
SurfDock
# Model Introduction SurfDock 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. Paper: > **SurfDock is a surface-informed diffusion generative model for reliable and accurate protein–ligand complex prediction** > Duanhua Cao, Mingan Chen, Rui Zhang, et al. > *Nature Methods*, 2024 > DOI: https://doi.org/10.1038/s41592-024-02516-y # Model Description SurfDock 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. During 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. By 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. # Use Cases | Use Case | Description | | --- | --- | | Protein–ligand docking | Predict the binding conformation of a ligand within a protein binding site | | Structure-based virtual screening | Perform batch docking and scoring for small-molecule libraries | | Ligand conformation generation | Sample candidate binding poses using a diffusion model | | Protein surface modeling | Use protein surface geometry and physicochemical information to assist docking | # Usage Instructions ## 1. OneCode Usage You can use the OneCode online environment for intelligent one-click AI4S programming: [Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home) ## 2. Manual Installation and Usage ### Hardware Requirements - SurfDock involves computationally intensive diffusion sampling, ESM representation extraction, and graph neural network computation. GPU/DCU acceleration is recommended. ### Environment Setup #### DCU Environment ```bash # Activate DTK and CONDA first conda create -n onescience311 python=3.11 -y conda activate onescience311 pip install onescience[bio] \ -i http://mirrors.onescience.ai:3141/pypi/simple/ \ --trusted-host mirrors.onescience.ai ``` ### Environment Notes - 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. - 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. ### Weights and Data Preparation #### SurfDock Model Weights The current repository already includes the main model weights required for inference: ```text weight/ ├── docking/ │ ├── best_ema_inference_epoch_model.pt │ └── model_parameters.yml ├── posepredict/ │ ├── best_model.pt │ └── model_parameters.yml └── screen/ ├── best_model.pt └── model_parameters.yml ``` Therefore, after downloading the complete repository, you generally do not need to download the main SurfDock model weights separately. Check: ```bash ls -lh weight/docking/ ls -lh weight/posepredict/ ls -lh weight/screen/ ``` #### ESM Model SurfDock uses ESM to extract protein sequence representations. Official installation method: ```bash git clone https://github.com/facebookresearch/esm model/esm cd model/esm pip install -e . cd ../.. ``` The official inference script uses: ```text esm2_t33_650M_UR50D ``` and runs: ```bash python model/esm/scripts/extract.py \ "esm2_t33_650M_UR50D" \ input.fasta \ output_dir \ --repr_layers 33 \ --include "per_tok" \ --truncation_seq_length 4096 ``` to extract protein residue-level embeddings. - 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: ```text https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t33_650M_UR50D.pt https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t33_650M_UR50D-contact-regression.pt ``` After downloading, it is recommended to place them in the current user's Torch Hub checkpoints cache directory: ```text ~/.cache/torch/hub/checkpoints/ ``` The final file paths should be: ```text ~/.cache/torch/hub/checkpoints/esm2_t33_650M_UR50D.pt ~/.cache/torch/hub/checkpoints/esm2_t33_650M_UR50D-contact-regression.pt ``` #### PDBBind Data If 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. After processing, place the data in: ```text model/data/PDBBind_processed/ ``` `model/data/splits/` contains data split information and is not equivalent to the complete PDBBind dataset. ## 3. Quick Start ### Download the Model Package ```bash hf download \ OneScience-Group/SurfDock \ --local-dir ./SurfDock cd SurfDock ``` - 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. # Example Data The current repository already provides docking and screening examples: ```text model/data/eval_sample_dirs/ model/data/Screen_sample_dirs/ ``` When 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: ```text data_dir surface_out_dir out_csv_file Screen_lib_path docking_out_dir ``` as needed. # Inference Examples The following commands are assumed to be executed from the SurfDock repository root directory. ## Protein–Ligand Docking Example Run: ```bash cd scripts/bash_scripts/test_scripts bash eval_samples.sh ``` The script mainly performs the following steps automatically: ```text 1. Preprocess the protein structure 2. Compute the protein surface 3. Build the inference input CSV 4. Extract ESM embeddings 5. Run SurfDock diffusion sampling 6. Save docking results ``` Before running, it is recommended to check: ```bash vim scripts/bash_scripts/test_scripts/eval_samples.sh ``` Pay particular attention to: ```text gpu_string data_dir surface_out_dir out_csv_file esmbedding_dir docking_out_dir ``` Users also need to adjust the `CUDA_VISIBLE_DEVICES` and `accelerate launch` settings in the script according to the actual adaptation method. ## Virtual Screening Example Run: ```bash cd scripts/bash_scripts/test_scripts bash screen_pipeline.sh ``` Before running, check: ```bash vim scripts/bash_scripts/test_scripts/screen_pipeline.sh ``` Modify the following as needed: ```text gpu_string data_dir surface_out_dir out_csv_file esmbedding_dir Screen_lib_path docking_out_dir ``` Here, `Screen_lib_path` specifies the small-molecule library to be screened, for example the official sample: ```text model/data/Screen_sample_dirs/test_samples/1a0q/1a0q_ligand_for_Screen.sdf ``` The screening workflow is mainly: ```text Protein preprocessing ↓ Protein surface computation ↓ ESM embedding ↓ SurfDock generates candidate conformations ↓ Screen model rescoring ↓ Output screening results ``` ## Skip Completed Protein Preprocessing This parameter can be modified in the following two scripts: ```text scripts/bash_scripts/test_scripts/eval_samples.sh scripts/bash_scripts/test_scripts/screen_pipeline.sh ``` The scripts use: ```bash target_have_processed=true ``` to 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. To reprocess the target protein: ```bash target_have_processed=false ``` When set to `false`, the scripts rerun the target protein preprocessing steps. ## Generate ESM Embeddings Separately First, build the FASTA file: ```bash python model/datasets/esm_embedding_preparation.py \ --out_file ./protein.fasta \ --protein_ligand_csv ./input.csv ``` Extract ESM representations: ```bash python model/esm/scripts/extract.py \ "esm2_t33_650M_UR50D" \ ./protein.fasta \ ./esm_embedding_output \ --repr_layers 33 \ --include "per_tok" \ --truncation_seq_length 4096 ``` Extract pocket embeddings: ```bash python model/datasets/get_pocket_embedding.py \ --protein_pocket_csv ./input.csv \ --embeddings_dir ./esm_embedding_output \ --pocket_emb_save_dir ./esm_embedding_pocket_output ``` Merge them into the `.pt` file required for SurfDock inference: ```bash python model/datasets/esm_pocket_embeddings_to_pt.py \ --esm_embeddings_path ./esm_embedding_pocket_output \ --output_path ./esm2_pocket_embeddings.pt ``` # Training Instructions ## Retraining SurfDock Retraining requires preparing the PDBBind data first and completing protein surface and ESM embedding preprocessing. The training-related scripts are located in: ```text scripts/bash_scripts/train_SurfDock_docking_module/ ``` First, 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: ```text scripts/bash_scripts/train_SurfDock_docking_module/train_SurfDock.sh ``` After completing the ESM embeddings, check the data, model, and output paths in `train_SurfDock.sh`, then run: ```bash cd scripts/bash_scripts/train_SurfDock_docking_module bash train_SurfDock.sh ``` ## Retraining SurfScore The SurfScore training script is located at: ```text scripts/bash_scripts/train_SurfScore/train_SurfScore.sh ``` Before running, check the PDBBind data, cache, ESM embedding, output directory, and GPU parameters in the script, then run: ```bash cd scripts/bash_scripts/train_SurfScore bash train_SurfScore.sh ``` # Output Description ## Docking Output The docking results from `eval_samples.sh` are saved in the directory specified by `docking_out_dir` in the script. The main outputs include: ```text Generated protein–ligand conformations Confidence/scoring results Runtime logs Intermediate CSV ESM embeddings Protein surface files ``` ## Screening Output `screen_pipeline.sh` generates the screening results in `docking_out_dir` and uses: ```text weight/screen/best_model.pt ``` to rescore the docking poses. The final results can be used to rank and screen candidate small molecules. # OneScience Official Information | Platform | OneScience Main Repository | Skills Repository | | --- | --- | --- | | Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills | | GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills | # Citation and License - 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. - PDBBind data is governed by its own license and terms of use. SurfDock's MIT License does not automatically apply to PDBBind data. - 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.