File size: 12,949 Bytes
10f2621 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 | ---
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.
|