README.md

---
frameworks:
- ""
language:
- en
license: mit
tags:
- OneScience
- affinity prediction
- molecule generation
---
<p align="center">
  <strong>
    <span style="font-size: 30px;">TargetDiff</span>
  </strong>
</p>

# 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](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)

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

```bash
nvidia-smi
```

- Hygon DCU:

```bash
hy-smi
```

## Quick Start

### 1. Install the Runtime Environment

```bash
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.

```bash
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

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

```bash
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
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

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

```bash
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

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

```bash
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
```

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

```bash
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
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:**

```bash
# 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 Å:

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

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

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

```bash
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.

```bibtex
@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}
}
```