OneScience-Group/DiffDock

<p align="center"> <strong> <span style="font-size: 30px;">DiffDock</span> </strong> </p>

Model Overview

DiffDock is a diffusion model for protein-ligand molecular docking proposed by Corso et al. It supports single-complex molecular docking, batch molecular docking, evaluation, and related workflows.

Model Description

DiffDock formulates docking as a generative modeling problem. Given a protein receptor and a small-molecule ligand, it jointly models the ligand's translational, rotational, and torsional degrees of freedom through diffusion processes on SO(3)/SE(3), and generates candidate binding conformations.

Weights and datasets are not available at the moment. They will be uploaded to Hugging Face soon, and command-line downloads will be supported later.

Use Cases

Usage

1. Using OneCode

You can try intelligent one-click AI4S programming through the OneCode online environment:

Try intelligent one-click AI4S programming

2. Manual Installation and Usage

Hardware Requirements

• —Running on a GPU or DCU is recommended.
• —CPU can be used for connectivity checks, but it is relatively slow.
• —DCU users need to install DTK in advance. DTK 25.04.2 or later is recommended, or the OneScience-recommended version that matches the current cluster.

Software Requirements

Common DiffDock dependencies include PyTorch, PyTorch Geometric, RDKit, OpenBabel, e3nn, torch-scatter, torch-cluster, NumPy, SciPy, tqdm, PyYAML, and others. GNINA energy minimization during dataset evaluation requires installing gnina separately.

Environment Checks

• —NVIDIA GPU:

nvidia-smi

• —Hygon DCU:

hy-smi

Quick Start

1. Install the Runtime Environment

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

If the following code cannot find required libraries at runtime, activate CUDA as shown below.

source ${ROCM_PATH}/cuda/env.sh
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/python3.11/site-packages/fastpt/torch/lib:$LD_LIBRARY_PATH"

2. Download the Model Package and Database

# If the model ID on the Hugging Face page uses different capitalization, use the actual published name.
hf download OneScience-Group/diffdock --local-dir ./diffdock
cd diffdock

Training Weights and Datasets

Weights and datasets are not available at the moment. They will be uploaded to Hugging Face soon, and command-line downloads will be supported later.

3. Check Files in the Package

The current package does not include pretrained weights. The command below should only show source code, configuration files, and example input files:

find . -maxdepth 3 -type f

If sampling or evaluation is required, first obtain a checkpoint through training, or place external DiffDock score/confidence weights in a local directory. Make sure the weight directory contains model_parameters.yml.

4. Score Model Training

Before training, update the data paths in configs/training.yml to local paths. Common fields include:

• —data.pdbbind_dir: processed PDBBind data directory.
• —data.moad_dir: processed MOAD data directory.
• —data.split_train: list of training-set complexes.
• —data.split_val: list of validation-set complexes.
• —runtime.log_dir: training output directory.

Start training:

cd scripts
bash train.sh

After training succeeds, the output directory usually contains:

• —model_parameters.yml: model architecture and training parameters.
• —best_model.pt: weights with the best validation loss.
• —best_ema_model.pt: EMA weights.
• —best_inference_epoch_model.pt or best_ema_inference_epoch_model.pt: weights saved when inference validation is enabled.
• —last_model.pt: training state from the final epoch.

5. Molecular Docking Sampling

Sampling requires a trained score model directory, for example:

outputs/train/diffdock_cg_example/
├── model_parameters.yml
└── best_model.pt

Update the key fields in configs/sampling.yml to the actual paths:

• —model.model_dir: score model directory.
• —model.ckpt: score checkpoint file name.
• —confidence.confidence_model_dir: confidence model directory; set it to null when rerank is not used.
• —input.protein_path: protein PDB path. You can use data/6o5u_protein_processed.pdb.
• —input.ligand_description: SMILES string or ligand SDF/MOL2 path. You can use data/6o5u_ligand.sdf.
• —runtime.out_dir: sampling output directory.

Start single-complex sampling:

cd scripts
bash infer.sh

6. CSV Batch Sampling

Set input.protein_ligand_csv in configs/sampling.yml to the CSV path, and set the single-complex fields to null as needed. The CSV is recommended to contain the following columns:

complex_name,protein_path,ligand_description,protein_sequence

Here, ligand_description can be either a SMILES string or an SDF/MOL2 file path.

7. Dataset Evaluation

Before evaluation, update the dataset, model, and output directories in configs/evaluate.yml to actual paths:

python -m scripts.evaluate --config configs/evaluate.yml

If gnina.gnina_minimize=true is enabled, first make sure the following command can be run directly in the current environment:

gnina --help

8. Confidence Model Training

The confidence training entry point is provided by the installed OneScience package. This Hugging Face package does not duplicate that module. After completing score model training, run the following in the OneScience environment:

python -m onescience.confidence.diffdock.confidence_train \
    --original_model_dir outputs/train/diffdock_cg_example \
    --data_dir /path/to/PDBBind_processed \
    --split_train /path/to/splits/timesplit_no_lig_overlap_train \
    --split_val /path/to/splits/timesplit_no_lig_overlap_val

Data Preparation

DiffDock training usually requires preprocessed PDBBind or MOAD data. It is recommended to organize the data under a unified root directory, for example:

${ONESCIENCE_DATASETS_DIR}/diffdock/
├── PDBBind_processed/
├── MOAD_processed/
└── splits/
    ├── timesplit_no_lig_overlap_train
    ├── timesplit_no_lig_overlap_val
    └── timesplit_test

Split files should be plain text files, with one complex name per line. Each name must correspond to a subdirectory name in the data directory.

Official OneScience Information

Citation and License

• —The original DiffDock code is licensed under the MIT License. This repository retains source attribution and is organized for OneScience Hugging Face automated runtime scenarios.
• —If you use DiffDock results in research, we recommend citing the original DiffDock and DiffDock-L papers, as well as relevant OneScience project information. Depending on the actual task, also add citations for datasets or tools such as PDBBind, MOAD, RDKit, OpenBabel, and GNINA.

@inproceedings{corso2023diffdock,
  title={DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking},
  author={Corso, Gabriele and St{\"a}rk, Hannes and Jing, Bowen and Barzilay, Regina and Jaakkola, Tommi},
  booktitle={International Conference on Learning Representations},
  year={2023}
}