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

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

# 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:

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:

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:

git clone https://github.com/facebookresearch/esm model/esm
cd model/esm
pip install -e .
cd ../..

The official inference script uses:

esm2_t33_650M_UR50D

and runs:

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:
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:

~/.cache/torch/hub/checkpoints/

The final file paths should be:

~/.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:

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

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:

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:

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:

cd scripts/bash_scripts/test_scripts
bash eval_samples.sh

The script mainly performs the following steps automatically:

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:

vim scripts/bash_scripts/test_scripts/eval_samples.sh

Pay particular attention to:

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:

cd scripts/bash_scripts/test_scripts
bash screen_pipeline.sh

Before running, check:

vim scripts/bash_scripts/test_scripts/screen_pipeline.sh

Modify the following as needed:

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:

model/data/Screen_sample_dirs/test_samples/1a0q/1a0q_ligand_for_Screen.sdf

The screening workflow is mainly:

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:

scripts/bash_scripts/test_scripts/eval_samples.sh
scripts/bash_scripts/test_scripts/screen_pipeline.sh

The scripts use:

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:

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:

python model/datasets/esm_embedding_preparation.py \
  --out_file ./protein.fasta \
  --protein_ligand_csv ./input.csv

Extract ESM representations:

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:

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:

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:

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:

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:

cd scripts/bash_scripts/train_SurfDock_docking_module
bash train_SurfDock.sh

Retraining SurfScore

The SurfScore training script is located at:

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:

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:

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:

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

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.
Downloads last month
-
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support