SurfDock / README.md
anzhi2710gmailcom's picture
Upload folder using huggingface_hub
10f2621 verified
|
Raw
History Blame Contribute Delete
12.9 kB
---
license: mit
language:
- en
tags:
- OneScience
- SCNet
- DCU
- protein-ligand-docking
- virtual-screening
- diffusion-model
frameworks:
- PyTorch
---
<p align="center">
<strong><span style="font-size: 30px;">SurfDock</span></strong>
</p>
# 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.