README.md

metadata

frameworks:
  - ''
language:
  - en
license: mit
tags:
  - OneScience
  - affinity prediction
  - molecule generation

TargetDiff

Model Overview

TargetDiff is a bioinformatics model for target-aware molecule generation and protein-ligand affinity prediction.

The original paper is 3D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Prediction (ICLR 2023).

Model Description

TargetDiff is based on a 3D equivariant diffusion network. Given a protein binding pocket, it generates candidate small molecules and can use an EGNN property prediction branch to predict protein-ligand complex affinity.

The current package organizes TargetDiff example scripts, a snapshot of the model implementation, sample inputs, and pretrained weights in the same directory. Datasets will be uploaded later.

Use Cases

Use case	Description
Protein-ligand affinity prediction	Takes a protein PDB file and ligand SDF file as input, and outputs molar concentration predictions for Ki, Kd, or IC50
Target-aware small molecule generation	Takes a protein binding pocket PDB file as input, and outputs generated molecule sample.pt files and SDF files for molecules that can be reconstructed
Diffusion model training	Trains the TargetDiff molecule generation model using CrossDocked2020 pocket data
Property prediction training	Trains an EGNN affinity prediction model using PDBbind data
Generated result evaluation	Computes metrics such as stability, reconstruction success rate, QED, SA, and optional Vina docking metrics for sampling results

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

First enter an environment where the OneScience bioinformatics dependencies have been installed.

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

hf download --model OneScience-Sugon/targetdiff --local-dir ./targetdiff
cd targetdiff

Training Weights and Datasets

Training weights are already included in the weights folder and can be used directly after downloading the model package. Datasets will be uploaded to Hugging Face soon, and command-line downloads will be supported later.

3. Affinity Prediction

3.1 Affinity Prediction Training

bash scripts/train_prop.sh

This script automatically performs the following steps:

1. Extract binding pockets from the PDBbind refined set.
2. Split the training, validation, and test sets according to the coreset.
3. Train an EGNN-based protein-ligand binding affinity prediction model.

3.2 Affinity Prediction Evaluation

Use the trained affinity prediction model to evaluate on the test set. The expected official metrics on PDBBind v2016 are:

RMSE	MAE	R²	Pearson	Spearman
1.316	1.031	0.633	0.797	0.782

Run:

export PYTHONPATH=../../../src:$PYTHONPATH
python scripts/property_prediction/eval_prop.py \
    --ckpt_path ${ONESCIENCE_MODELS_DIR}/targetdiff/pretrained_models/egnn_pdbbind_v2016.pt \
    --device cuda

3.3 Affinity Prediction Inference

bash scripts/inference.sh

By default, this script performs affinity prediction on the example protein-ligand pair 3ug2, using the default weights and example data:

• Model weights: ${ONESCIENCE_MODELS_DIR}/targetdiff/pretrained_models/egnn_pdbbind_v2016.pt
• Protein: ${ONESCIENCE_DATASETS_DIR}/targetdiff/examples/3ug2_protein.pdb
• Ligand: ${ONESCIENCE_DATASETS_DIR}/targetdiff/examples/3ug2_ligand.sdf
• Affinity type: default Kd
• Compute device: default cuda

4. Molecular Sampling

4.1 Single Test-Set Sample

export PYTHONPATH=../../../src:$PYTHONPATH
python -m scripts.sample_diffusion configs/sampling.yml -i 0 --batch_size 50 --result_path ./outputs

This script reads the diffusion model checkpoint from configs/sampling.yml, generates candidate ligand molecules for the i-th pocket in the test set, and saves the result as result_i.pt.

4.2 Multi-GPU Batch Sampling

bash scripts/batch_sample_diffusion.sh configs/sampling.yml outputs 4 0 0

Parameter descriptions:

Parameter	Description
$1	Sampling configuration file path
$2	Result output directory
$3	Total number of worker nodes
$4	Current node index, starting from 0
$5	Starting data index

Example of multi-GPU parallel sampling:

CUDA_VISIBLE_DEVICES=0 bash scripts/batch_sample_diffusion.sh configs/sampling.yml outputs 4 0 0 &
CUDA_VISIBLE_DEVICES=1 bash scripts/batch_sample_diffusion.sh configs/sampling.yml outputs 4 1 0 &
CUDA_VISIBLE_DEVICES=2 bash scripts/batch_sample_diffusion.sh configs/sampling.yml outputs 4 2 0 &
CUDA_VISIBLE_DEVICES=3 bash scripts/batch_sample_diffusion.sh configs/sampling.yml outputs 4 3 0 &
wait

The script internally fixes TOTAL_TASKS=100 and BATCH_SIZE=50, and assigns samples to nodes by index modulo.

4.3 Sampling for a Custom PDB Pocket

export PYTHONPATH=../../../src:$PYTHONPATH
python -m scripts.sample_for_pocket configs/sampling.yml \
    --pdb_path /path/to/pocket.pdb \
    --result_path ./outputs_pdb \
    --num_samples 5 \
    --batch_size 1

Parameter descriptions:

Parameter	Required	Description
config	Yes	Sampling configuration file path
--pdb_path	Yes	Protein pocket PDB file path. A 10 Å pocket is recommended.
--result_path	No	Result output directory, default ./outputs_pdb
--num_samples	No	Number of molecules to sample, read from the configuration file by default
--batch_size	No	Batch size, default 100
--device	No	Compute device, default cuda:0

Sampling results are saved as outputs_pdb/sample.pt, and successfully reconstructed molecules are additionally written to outputs_pdb/sdf/.

5. Generated Molecule Evaluation

5.1 Evaluation from Sampling Results

If docking evaluation is required (vina_score / vina_dock / qvina), also install:

pip install meeko==0.1.dev3 scipy pdb2pqr vina==1.2.2
python -m pip install git+https://github.com/Valdes-Tresanco-MS/AutoDockTools_py3

export PYTHONPATH=../../../src:$PYTHONPATH
python scripts/evaluate_diffusion.py ./outputs --docking_mode vina_score --protein_root /path/to/protein_root

Parameter descriptions:

Parameter	Required	Description
sample_path	Yes	Sampling result directory containing result_*.pt files
--docking_mode	Yes	Docking mode. Options: none, vina_score, vina_dock, qvina
--protein_root	No	Root directory of the original protein files, used for docking evaluation
--eval_step	No	Which sampling step to evaluate, default -1 (the final step)
--eval_num_examples	No	Number of samples to evaluate, default all
--exhaustiveness	No	Docking search intensity, default 16
--save	No	Whether to save evaluation results, default True

Supported docking modes:

Mode	Description
none	Computes only metrics such as validity, uniqueness, and novelty, without docking
vina_score	Uses AutoDock Vina to score generated molecules
vina_dock	Uses AutoDock Vina to redock generated molecules
qvina	Uses QuickVina for docking

The first run in vina_score or vina_dock mode may take some time to prepare pdbqt and pqr files.

5.2 Evaluation from Meta Files

The official project provides sampled and docked meta files, including TargetDiff and baselines such as CVAE, AR, and Pocket2Mol. These files can be downloaded and evaluated directly:

Meta file	Corresponding paper
crossdocked_test_vina_docked.pt	Original test-set docking results
cvae_vina_docked.pt	liGAN
ar_vina_docked.pt	AR
pocket2mol_vina_docked.pt	Pocket2Mol
targetdiff_vina_docked.pt	TargetDiff

Official meta file download URL: https://drive.google.com/drive/folders/19imu-mlwrjnQhgbXpwsLgA17s1Rv70YS?usp=share_link

Evaluation command:

export PYTHONPATH=../../../src:$PYTHONPATH
python scripts/evaluate_from_meta.py sampling_results/targetdiff_vina_docked.pt --result_path eval_targetdiff

Parameter descriptions:

Parameter	Required	Description
meta_file	Yes	.pt file containing sampling and docking results
--result_path	No	Evaluation result output directory, default eval_results

---

6. Diffusion Model Training

bash train_diffusion.sh

By default, this reads configs/training.yml, trains on the CrossDocked2020 pocket dataset, and saves logs and checkpoints to ./logs_diffusion/.

Custom configuration or parameter overrides:

# Specify a custom configuration file
bash train_diffusion.sh configs/custom_training.yml

# Override training parameters from the command line
bash train_diffusion.sh --train.batch_size 8 --train.max_iters 500000

Main training parameters:

Parameter	Default value	Description
data.path	${ONESCIENCE_DATASETS_DIR}/targetdiff/data/crossdocked_v1.1_rmsd1.0_pocket10	Training data directory
data.split	${ONESCIENCE_DATASETS_DIR}/targetdiff/data/crossdocked_pocket10_pose_split.pt	Training/validation/test split file
train.batch_size	4	Batch size
train.max_iters	10000000	Maximum number of iterations
train.lr	5.e-4	Learning rate
logdir	./logs_diffusion	Log output directory

---

7. Data Preprocessing

7.1 CrossDocked2020 Data Preprocessing

To process CrossDocked2020 data from scratch, follow these steps:

1. Download CrossDocked2020 v1.1 and save it to data/CrossDocked2020.

2. Filter samples with RMSD < 1 Å:

   export PYTHONPATH=../../../src:$PYTHONPATH
   python scripts/data_preparation/clean_crossdocked.py \
       --source data/CrossDocked2020 \
       --dest data/crossdocked_v1.1_rmsd1.0 \
       --rmsd_thr 1.0
   

3. Extract 10 Å binding pockets from proteins:

   python scripts/data_preparation/extract_pockets.py \
       --source data/crossdocked_v1.1_rmsd1.0 \
       --dest data/crossdocked_v1.1_rmsd1.0_pocket10
   

4. Split the training and test sets:

   python scripts/data_preparation/split_pl_dataset.py \
       --path data/crossdocked_v1.1_rmsd1.0_pocket10 \
       --dest data/crossdocked_pocket10_pose_split.pt \
       --fixed_split data/split_by_name.pt
   

7.2 PDBbind Data Preprocessing

The affinity prediction training script train_prop.sh automatically performs pocket extraction and dataset splitting. To run these steps separately, use the following commands:

export PYTHONPATH=../../../src:$PYTHONPATH

python scripts/property_prediction/extract_pockets.py \
    --source data/pdbbind_v2020 \
    --dest data/pdbbind_v2020_processed \
    --subset refined \
    --num_workers 16

python scripts/property_prediction/pdbbind_split.py \
    --split_mode coreset \
    --index_path data/pdbbind_v2020_processed/pocket_10_refined/index.pkl \
    --test_path data/pdbbind_v2016/coreset \
    --save_path data/pdbbind_v2020_processed/pocket_10_refined/split.pt

---

Official OneScience 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 original TargetDiff 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 TargetDiff results in research, we recommend citing the original TargetDiff paper and relevant OneScience project information. Depending on the actual task, also add citations for datasets or tools such as CrossDocked2020, PDBbind, RDKit, OpenBabel, and Vina/QVina.

@inproceedings{guan3d,
  title={3D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Prediction},
  author={Guan, Jiaqi and Qian, Wesley Wei and Peng, Xingang and Su, Yufeng and Peng, Jian and Ma, Jianzhu},
  booktitle={International Conference on Learning Representations},
  year={2023}
}