Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,50 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
+ 
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+ 
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,387 @@
+---
+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 ModelScope 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 ModelScope 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}
+}
+```

configs/prop/pdbbind_general_egnn.yml

@@ -0,0 +1,39 @@
+model:
+  hidden_channels: 256
+  encoder:
+    name: egnn
+    num_layers: 6
+    hidden_dim: 256
+    edge_dim: 0
+    num_r_gaussian: 64
+    act_fn: relu
+    norm: False
+    cutoff: 10.0
+    knn: 48 # kNN Graph
+
+train:
+  seed: 2021
+  batch_size: 4
+  num_workers: 4
+  max_epochs: 200
+  report_iter: 200
+  val_freq: 1
+  pos_noise_std: 0.1
+  max_grad_norm: 10.
+  optimizer:
+    type: adam
+    lr: 1.e-4
+    weight_decay: 0
+    beta1: 0.99
+    beta2: 0.999
+  scheduler:
+    type: plateau
+    factor: 0.6
+    patience: 10
+    min_lr: 1.e-5
+
+dataset:
+  name: pdbbind
+  heavy_only: True
+  path: ${ONESCIENCE_DATASETS_DIR}/targetdiff/data/pdbbind_v2020/pocket_10_general
+  split: ${ONESCIENCE_DATASETS_DIR}/targetdiff/data/pdbbind_v2020/pocket_10_general/split.pt

configs/prop/pdbbind_general_egnn_enc_final_h.yml

@@ -0,0 +1,44 @@
+model:
+  hidden_channels: 256
+  enc_ligand_dim: 0
+  enc_node_dim: 128
+  enc_graph_dim: 0
+  enc_feature_type: final_h
+  encoder:
+    name: egnn_enc
+    num_layers: 6
+    hidden_dim: 256
+    edge_dim: 0
+    num_r_gaussian: 64
+    act_fn: relu
+    norm: False
+    cutoff: 10.0
+    knn: 48 # kNN Graph
+
+train:
+  seed: 2021
+  batch_size: 4
+  num_workers: 4
+  max_epochs: 200
+  report_iter: 200
+  val_freq: 1
+  pos_noise_std: 0.1
+  max_grad_norm: 10.
+  optimizer:
+    type: adam
+    lr: 1.e-4
+    weight_decay: 0
+    beta1: 0.99
+    beta2: 0.999
+  scheduler:
+    type: plateau
+    factor: 0.6
+    patience: 10
+    min_lr: 1.e-5
+
+dataset:
+  name: pdbbind
+  heavy_only: True
+  path: ${ONESCIENCE_DATASETS_DIR}/targetdiff/data/pdbbind_v2020/pocket_10_general
+  split: ${ONESCIENCE_DATASETS_DIR}/targetdiff/data/pdbbind_v2020/pocket_10_general/split.pt
+  emb_path: ${ONESCIENCE_DATASETS_DIR}/targetdiff/data/pdbbind_v2020/pdbbind_v2020_all_match_emb.pt

configs/sampling.yml

@@ -0,0 +1,10 @@
+model:
+   checkpoint: /public/home/liuyx19/onescience-examples/models/biosciences/TargetDiff/weight/pretrained_diffusion.pt
+
+sample:
+  seed: 2021
+  num_samples: 100
+  num_steps: 1000
+  pos_only: False
+  center_pos_mode: protein
+  sample_num_atoms: prior

configs/training.yml

@@ -0,0 +1,64 @@
+data:
+  name: pl
+  path: ${ONESCIENCE_DATASETS_DIR}/targetdiff/data/crossdocked_v1.1_rmsd1.0_pocket10
+  split: ${ONESCIENCE_DATASETS_DIR}/targetdiff/data/crossdocked_pocket10_pose_split.pt
+  transform:
+    ligand_atom_mode: add_aromatic
+    random_rot: False
+
+model:
+  model_mean_type: C0  # ['noise', 'C0']
+  beta_schedule: sigmoid
+  beta_start: 1.e-7
+  beta_end: 2.e-3
+  v_beta_schedule: cosine
+  v_beta_s: 0.01
+  num_diffusion_timesteps: 1000
+  loss_v_weight: 100.
+  sample_time_method: symmetric  # ['importance', 'symmetric']
+
+  time_emb_dim: 0
+  time_emb_mode: simple
+  center_pos_mode: protein
+
+  node_indicator: True
+  model_type: uni_o2
+  num_blocks: 1
+  num_layers: 9
+  hidden_dim: 128
+  n_heads: 16
+  edge_feat_dim: 4  # edge type feat
+  num_r_gaussian: 20
+  knn: 32 # !
+  num_node_types: 8
+  act_fn: relu
+  norm: True
+  cutoff_mode: knn  # [radius, none]
+  ew_net_type: global  # [r, m, none]
+  num_x2h: 1
+  num_h2x: 1
+  r_max: 10.
+  x2h_out_fc: False
+  sync_twoup: False
+
+train:
+  seed: 2021
+  batch_size: 4
+  num_workers: 4
+  n_acc_batch: 1
+  max_iters: 10000000
+  val_freq: 100
+  pos_noise_std: 0.1
+  max_grad_norm: 8.0
+  bond_loss_weight: 1.0
+  optimizer:
+    type: adam
+    lr: 5.e-4
+    weight_decay: 0
+    beta1: 0.95
+    beta2: 0.999
+  scheduler:
+    type: plateau
+    factor: 0.6
+    patience: 10
+    min_lr: 1.e-6

examples/3ug2_ligand.sdf

@@ -0,0 +1,136 @@
+3ug2_ligand
+
+Created by X-TOOL on Fri Sep 26 17:34:55 2014
+ 56 59  0  0  0  0  0  0  0  0999 V2000
+    3.4040   46.1170   14.5070  C 0  0  0  3  0  4
+    4.3860   44.9380   14.3750  C 0  0  0  3  0  4
+    3.6900   43.6880   14.5540  O 0  0  0  1  0  2
+    3.0800   43.4440   15.8490  C 0  0  0  3  0  4
+    2.7940   44.7300   16.6440  C 0  0  0  3  0  4
+    2.5950   45.8870   15.7230  N 0  3  0  2  0  4
+    1.5560   46.8550   16.0410  C 0  0  0  3  0  4
+    1.9860   47.4510   17.3890  C 0  0  0  3  0  4
+    0.7930   47.7880   18.2910  C 0  0  0  3  0  4
+   -0.0010   48.8270   17.6860  O 0  0  0  1  0  2
+   -0.5460   49.8840   18.3700  C 0  0  0  1  0  3
+   -1.1750   50.8290   17.5810  C 0  0  0  1  0  3
+   -1.7490   51.9430   18.1780  C 0  0  0  2  0  3
+   -1.2170   50.6790   16.2200  O 0  0  0  1  0  2
+   -2.3380   51.2600   15.5080  C 0  0  0  4  0  4
+   -0.5160   50.0310   19.7780  C 0  0  0  2  0  3
+   -1.1300   51.1410   20.3560  C 0  0  0  1  0  3
+   -1.7470   52.0880   19.5480  C 0  0  0  1  0  3
+   -2.3440   53.1660   20.0860  N 0  0  0  1  0  2
+   -2.3130   53.3880   21.3930  C 0  0  0  2  0  3
+   -1.7230   52.5050   22.1960  N 0  0  0  1  0  2
+   -1.1370   51.3830   21.7250  C 0  0  0  1  0  3
+   -0.5500   50.5160   22.5680  N 0  0  0  2  0  3
+   -0.7890   50.3540   23.9050  C 0  0  0  1  0  3
+   -2.0680   50.5030   24.4450  C 0  0  0  2  0  3
+   -2.2860   50.3360   25.8070  C 0  0  0  1  0  3
+   -3.8770   50.5280   26.6050 Cl 0  0  0  1  0  1
+   -1.2480   49.9960   26.6110  C 0  0  0  1  0  3
+   -1.5490   49.8410   27.9240  F 0  0  0  1  0  1
+    0.0430   49.8020   26.0870  C 0  0  0  2  0  3
+    0.2670   49.9660   24.7340  C 0  0  0  2  0  3
+    2.7500   46.1619   13.6237  H 0  0  0  1  0  1
+    3.9608   47.0611   14.5993  H 0  0  0  1  0  1
+    5.1698   45.0304   15.1412  H 0  0  0  1  0  1
+    4.8457   44.9581   13.3759  H 0  0  0  1  0  1
+    3.7592   42.8126   16.4406  H 0  0  0  1  0  1
+    2.1295   42.9131   15.6919  H 0  0  0  1  0  1
+    3.6448   44.9401   17.3088  H 0  0  0  1  0  1
+    1.8845   44.5874   17.2461  H 0  0  0  1  0  1
+    1.9433   45.3395   15.1831  H 0  0  0  1  0  1
+    0.5774   46.3604   16.1290  H 0  0  0  1  0  1
+    1.5036   47.6369   15.2691  H 0  0  0  1  0  1
+    2.5570   48.3722   17.2010  H 0  0  0  1  0  1
+    2.6259   46.7225   17.9085  H 0  0  0  1  0  1
+    1.1608   48.1326   19.2687  H 0  0  0  1  0  1
+    0.1740   46.8891   18.4282  H 0  0  0  1  0  1
+   -2.2041   52.7077   17.5590  H 0  0  0  1  0  1
+   -2.2330   51.0584   14.4317  H 0  0  0  1  0  1
+   -2.3565   52.3467   15.6774  H 0  0  0  1  0  1
+   -3.2752   50.8157   15.8745  H 0  0  0  1  0  1
+   -0.0226   49.2909   20.3976  H 0  0  0  1  0  1
+   -2.7652   54.2829   21.8051  H 0  0  0  1  0  1
+    0.1659   49.8995   22.1578  H 0  0  0  1  0  1
+   -2.8992   50.7515   23.7951  H 0  0  0  1  0  1
+    0.8599   49.5245   26.7434  H 0  0  0  1  0  1
+    1.2536   49.7952   24.3187  H 0  0  0  1  0  1
+  1  2  1  0  0  1
+  1  6  1  0  0  1
+  2  3  1  0  0  1
+  3  4  1  0  0  1
+  4  5  1  0  0  1
+  5  6  1  0  0  1
+  6  7  1  0  0  2
+  7  8  1  0  0  2
+  8  9  1  0  0  2
+  9 10  1  0  0  2
+ 10 11  1  0  0  2
+ 11 12  4  0  0  1
+ 11 16  4  0  0  1
+ 12 13  4  0  0  1
+ 12 14  1  0  0  2
+ 13 18  4  0  0  1
+ 14 15  1  0  0  2
+ 16 17  4  0  0  1
+ 17 18  4  0  0  1
+ 17 22  4  0  0  1
+ 18 19  4  0  0  1
+ 19 20  4  0  0  1
+ 20 21  4  0  0  1
+ 21 22  4  0  0  1
+ 22 23  1  0  0  2
+ 23 24  1  0  0  2
+ 24 25  4  0  0  1
+ 24 31  4  0  0  1
+ 25 26  4  0  0  1
+ 26 27  1  0  0  2
+ 26 28  4  0  0  1
+ 28 29  1  0  0  2
+ 28 30  4  0  0  1
+ 30 31  4  0  0  1
+  1 32  1  0  0  2
+  1 33  1  0  0  2
+  2 34  1  0  0  2
+  2 35  1  0  0  2
+  4 36  1  0  0  2
+  4 37  1  0  0  2
+  5 38  1  0  0  2
+  5 39  1  0  0  2
+  6 40  1  0  0  2
+  7 41  1  0  0  2
+  7 42  1  0  0  2
+  8 43  1  0  0  2
+  8 44  1  0  0  2
+  9 45  1  0  0  2
+  9 46  1  0  0  2
+ 13 47  1  0  0  2
+ 15 48  1  0  0  2
+ 15 49  1  0  0  2
+ 15 50  1  0  0  2
+ 16 51  1  0  0  2
+ 20 52  1  0  0  2
+ 23 53  1  0  0  2
+ 25 54  1  0  0  2
+ 30 55  1  0  0  2
+ 31 56  1  0  0  2
+M  END
+> <MOLECULAR_FORMULA>
+C22H25N4O3FCl
+
+> <MOLECULAR_WEIGHT>
+447.7
+
+> <NUM_HB_ATOMS>
+7  
+
+> <NUM_ROTOR>
+6  
+
+> <XLOGP2>
+3.77 
+
+$$$$

models/__init__.py

@@ -0,0 +1,2 @@
+"""TargetDiff model components."""
+

models/common.py

@@ -0,0 +1,212 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.nn import knn_graph
+
+
+class GaussianSmearing(nn.Module):
+    def __init__(self, start=0.0, stop=5.0, num_gaussians=50, fixed_offset=True):
+        super(GaussianSmearing, self).__init__()
+        self.start = start
+        self.stop = stop
+        self.num_gaussians = num_gaussians
+        if fixed_offset:
+            # customized offset
+            offset = torch.tensor([0, 1, 1.25, 1.5, 1.75, 2, 2.25, 2.5, 2.75, 3, 3.5, 4, 4.5, 5, 5.5, 6, 7, 8, 9, 10])
+        else:
+            offset = torch.linspace(start, stop, num_gaussians)
+        self.coeff = -0.5 / (offset[1] - offset[0]).item() ** 2
+        self.register_buffer('offset', offset)
+
+    def __repr__(self):
+        return f'GaussianSmearing(start={self.start}, stop={self.stop}, num_gaussians={self.num_gaussians})'
+
+    def forward(self, dist):
+        dist = dist.view(-1, 1) - self.offset.view(1, -1)
+        return torch.exp(self.coeff * torch.pow(dist, 2))
+
+
+class AngleExpansion(nn.Module):
+    def __init__(self, start=1.0, stop=5.0, half_expansion=10):
+        super(AngleExpansion, self).__init__()
+        l_mul = 1. / torch.linspace(stop, start, half_expansion)
+        r_mul = torch.linspace(start, stop, half_expansion)
+        coeff = torch.cat([l_mul, r_mul], dim=-1)
+        self.register_buffer('coeff', coeff)
+
+    def forward(self, angle):
+        return torch.cos(angle.view(-1, 1) * self.coeff.view(1, -1))
+
+
+class Swish(nn.Module):
+    def __init__(self):
+        super(Swish, self).__init__()
+        self.beta = nn.Parameter(torch.tensor(1.0))
+
+    def forward(self, x):
+        return x * torch.sigmoid(self.beta * x)
+
+
+NONLINEARITIES = {
+    "tanh": nn.Tanh(),
+    "relu": nn.ReLU(),
+    "softplus": nn.Softplus(),
+    "elu": nn.ELU(),
+    "swish": Swish(),
+    'silu': nn.SiLU()
+}
+
+
+class MLP(nn.Module):
+    """MLP with the same hidden dim across all layers."""
+
+    def __init__(self, in_dim, out_dim, hidden_dim, num_layer=2, norm=True, act_fn='relu', act_last=False):
+        super().__init__()
+        layers = []
+        for layer_idx in range(num_layer):
+            if layer_idx == 0:
+                layers.append(nn.Linear(in_dim, hidden_dim))
+            elif layer_idx == num_layer - 1:
+                layers.append(nn.Linear(hidden_dim, out_dim))
+            else:
+                layers.append(nn.Linear(hidden_dim, hidden_dim))
+            if layer_idx < num_layer - 1 or act_last:
+                if norm:
+                    layers.append(nn.LayerNorm(hidden_dim))
+                layers.append(NONLINEARITIES[act_fn])
+        self.net = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.net(x)
+
+
+def outer_product(*vectors):
+    for index, vector in enumerate(vectors):
+        if index == 0:
+            out = vector.unsqueeze(-1)
+        else:
+            out = out * vector.unsqueeze(1)
+            out = out.view(out.shape[0], -1).unsqueeze(-1)
+    return out.squeeze()
+
+
+def get_h_dist(dist_metric, hi, hj):
+    if dist_metric == 'euclidean':
+        h_dist = torch.sum((hi - hj) ** 2, -1, keepdim=True)
+        return h_dist
+    elif dist_metric == 'cos_sim':
+        hi_norm = torch.norm(hi, p=2, dim=-1, keepdim=True)
+        hj_norm = torch.norm(hj, p=2, dim=-1, keepdim=True)
+        h_dist = torch.sum(hi * hj, -1, keepdim=True) / (hi_norm * hj_norm)
+        return h_dist, hj_norm
+
+
+def get_r_feat(r, r_exp_func, node_type=None, edge_index=None, mode='basic'):
+    if mode == 'origin':
+        r_feat = r
+    elif mode == 'basic':
+        r_feat = r_exp_func(r)
+    elif mode == 'sparse':
+        src, dst = edge_index
+        nt_src = node_type[src]  # [n_edges, 8]
+        nt_dst = node_type[dst]
+        r_exp = r_exp_func(r)
+        r_feat = outer_product(nt_src, nt_dst, r_exp)
+    else:
+        raise ValueError(mode)
+    return r_feat
+
+
+def compose_context(h_protein, h_ligand, pos_protein, pos_ligand, batch_protein, batch_ligand):
+    # previous version has problems when ligand atom types are fixed
+    # (due to sorting randomly in case of same element)
+
+    batch_ctx = torch.cat([batch_protein, batch_ligand], dim=0)
+    # sort_idx = batch_ctx.argsort()
+    sort_idx = torch.sort(batch_ctx, stable=True).indices
+
+    mask_ligand = torch.cat([
+        torch.zeros([batch_protein.size(0)], device=batch_protein.device).bool(),
+        torch.ones([batch_ligand.size(0)], device=batch_ligand.device).bool(),
+    ], dim=0)[sort_idx]
+
+    batch_ctx = batch_ctx[sort_idx]
+    h_ctx = torch.cat([h_protein, h_ligand], dim=0)[sort_idx]  # (N_protein+N_ligand, H)
+    pos_ctx = torch.cat([pos_protein, pos_ligand], dim=0)[sort_idx]  # (N_protein+N_ligand, 3)
+
+    return h_ctx, pos_ctx, batch_ctx, mask_ligand
+
+
+def compose_context_prop(h_protein, h_ligand, pos_protein, pos_ligand, batch_protein, batch_ligand):
+    batch_ctx = torch.cat([batch_protein, batch_ligand], dim=0)
+    sort_idx = batch_ctx.argsort()
+
+    mask_protein = torch.cat([
+        torch.ones([batch_protein.size(0)], device=batch_protein.device).bool(),
+        torch.zeros([batch_ligand.size(0)], device=batch_ligand.device).bool(),
+    ], dim=0)[sort_idx]
+
+    batch_ctx = batch_ctx[sort_idx]
+    h_ctx = torch.cat([h_protein, h_ligand], dim=0)[sort_idx]       # (N_protein+N_ligand, H)
+    pos_ctx = torch.cat([pos_protein, pos_ligand], dim=0)[sort_idx]  # (N_protein+N_ligand, 3)
+
+    return h_ctx, pos_ctx, batch_ctx
+
+
+class ShiftedSoftplus(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.shift = torch.log(torch.tensor(2.0)).item()
+
+    def forward(self, x):
+        return F.softplus(x) - self.shift
+
+
+def hybrid_edge_connection(ligand_pos, protein_pos, k, ligand_index, protein_index):
+    # fully-connected for ligand atoms
+    dst = torch.repeat_interleave(ligand_index, len(ligand_index))
+    src = ligand_index.repeat(len(ligand_index))
+    mask = dst != src
+    dst, src = dst[mask], src[mask]
+    ll_edge_index = torch.stack([src, dst])
+
+    # knn for ligand-protein edges
+    ligand_protein_pos_dist = torch.unsqueeze(ligand_pos, 1) - torch.unsqueeze(protein_pos, 0)
+    ligand_protein_pos_dist = torch.norm(ligand_protein_pos_dist, p=2, dim=-1)
+    knn_p_idx = torch.topk(ligand_protein_pos_dist, k=k, largest=False, dim=1).indices
+    knn_p_idx = protein_index[knn_p_idx]
+    knn_l_idx = torch.unsqueeze(ligand_index, 1)
+    knn_l_idx = knn_l_idx.repeat(1, k)
+    pl_edge_index = torch.stack([knn_p_idx, knn_l_idx], dim=0)
+    pl_edge_index = pl_edge_index.view(2, -1)
+    return ll_edge_index, pl_edge_index
+
+
+def batch_hybrid_edge_connection(x, k, mask_ligand, batch, add_p_index=False):
+    batch_size = batch.max().item() + 1
+    batch_ll_edge_index, batch_pl_edge_index, batch_p_edge_index = [], [], []
+    with torch.no_grad():
+        for i in range(batch_size):
+            ligand_index = ((batch == i) & (mask_ligand == 1)).nonzero()[:, 0]
+            protein_index = ((batch == i) & (mask_ligand == 0)).nonzero()[:, 0]
+            ligand_pos, protein_pos = x[ligand_index], x[protein_index]
+            ll_edge_index, pl_edge_index = hybrid_edge_connection(
+                ligand_pos, protein_pos, k, ligand_index, protein_index)
+            batch_ll_edge_index.append(ll_edge_index)
+            batch_pl_edge_index.append(pl_edge_index)
+            if add_p_index:
+                all_pos = torch.cat([protein_pos, ligand_pos], 0)
+                p_edge_index = knn_graph(all_pos, k=k, flow='source_to_target')
+                p_edge_index = p_edge_index[:, p_edge_index[1] < len(protein_pos)]
+                p_src, p_dst = p_edge_index
+                all_index = torch.cat([protein_index, ligand_index], 0)
+                p_edge_index = torch.stack([all_index[p_src], all_index[p_dst]], 0)
+                batch_p_edge_index.append(p_edge_index)
+
+    if add_p_index:
+        edge_index = [torch.cat([ll, pl, p], -1) for ll, pl, p in zip(
+            batch_ll_edge_index, batch_pl_edge_index, batch_p_edge_index)]
+    else:
+        edge_index = [torch.cat([ll, pl], -1) for ll, pl in zip(batch_ll_edge_index, batch_pl_edge_index)]
+    edge_index = torch.cat(edge_index, -1)
+    return edge_index

models/egnn.py

@@ -0,0 +1,133 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_scatter import scatter_sum
+from torch_geometric.nn import radius_graph, knn_graph
+from .common import GaussianSmearing, MLP, batch_hybrid_edge_connection, NONLINEARITIES
+
+
+class EnBaseLayer(nn.Module):
+    def __init__(self, hidden_dim, edge_feat_dim, num_r_gaussian, update_x=True, act_fn='silu', norm=False):
+        super().__init__()
+        self.r_min = 0.
+        self.r_max = 10.
+        self.hidden_dim = hidden_dim
+        self.num_r_gaussian = num_r_gaussian
+        self.edge_feat_dim = edge_feat_dim
+        self.update_x = update_x
+        self.act_fn = act_fn
+        self.norm = norm
+        if num_r_gaussian > 1:
+            self.distance_expansion = GaussianSmearing(self.r_min, self.r_max, num_gaussians=num_r_gaussian)
+        self.edge_mlp = MLP(2 * hidden_dim + edge_feat_dim + num_r_gaussian, hidden_dim, hidden_dim,
+                            num_layer=2, norm=norm, act_fn=act_fn, act_last=True)
+        self.edge_inf = nn.Sequential(nn.Linear(hidden_dim, 1), nn.Sigmoid())
+        if self.update_x:
+            # self.x_mlp = MLP(hidden_dim, 1, hidden_dim, num_layer=2, norm=norm, act_fn=act_fn)
+            x_mlp = [nn.Linear(hidden_dim, hidden_dim), NONLINEARITIES[act_fn]]
+            layer = nn.Linear(hidden_dim, 1, bias=False)
+            torch.nn.init.xavier_uniform_(layer.weight, gain=0.001)
+            x_mlp.append(layer)
+            x_mlp.append(nn.Tanh())
+            self.x_mlp = nn.Sequential(*x_mlp)
+
+        self.node_mlp = MLP(2 * hidden_dim, hidden_dim, hidden_dim, num_layer=2, norm=norm, act_fn=act_fn)
+
+    def forward(self, h, x, edge_index, mask_ligand, edge_attr=None):
+        src, dst = edge_index
+        hi, hj = h[dst], h[src]
+        # \phi_e in Eq(3)
+        rel_x = x[dst] - x[src]
+        d_sq = torch.sum(rel_x ** 2, -1, keepdim=True)
+        if self.num_r_gaussian > 1:
+            d_feat = self.distance_expansion(torch.sqrt(d_sq + 1e-8))
+        else:
+            d_feat = d_sq
+        if edge_attr is not None:
+            edge_feat = torch.cat([d_feat, edge_attr], -1)
+        else:
+            edge_feat = d_sq
+
+        mij = self.edge_mlp(torch.cat([hi, hj, edge_feat], -1))
+        eij = self.edge_inf(mij)
+        mi = scatter_sum(mij * eij, dst, dim=0, dim_size=h.shape[0])
+
+        # h update in Eq(6)
+        h = h + self.node_mlp(torch.cat([mi, h], -1))
+        if self.update_x:
+            # x update in Eq(4)
+            xi, xj = x[dst], x[src]
+            # (xi - xj) / (\|xi - xj\| + C) to make it more stable
+            delta_x = scatter_sum((xi - xj) / (torch.sqrt(d_sq + 1e-8) + 1) * self.x_mlp(mij), dst, dim=0)
+            x = x + delta_x * mask_ligand[:, None]  # only ligand positions will be updated
+
+        return h, x
+
+
+class EGNN(nn.Module):
+    def __init__(self, num_layers, hidden_dim, edge_feat_dim, num_r_gaussian, k=32, cutoff=10.0, cutoff_mode='knn',
+                 update_x=True, act_fn='silu', norm=False):
+        super().__init__()
+        # Build the network
+        self.num_layers = num_layers
+        self.hidden_dim = hidden_dim
+        self.edge_feat_dim = edge_feat_dim
+        self.num_r_gaussian = num_r_gaussian
+        self.update_x = update_x
+        self.act_fn = act_fn
+        self.norm = norm
+        self.k = k
+        self.cutoff = cutoff
+        self.cutoff_mode = cutoff_mode
+        self.distance_expansion = GaussianSmearing(stop=cutoff, num_gaussians=num_r_gaussian)
+        self.net = self._build_network()
+
+    def _build_network(self):
+        # Equivariant layers
+        layers = []
+        for l_idx in range(self.num_layers):
+            layer = EnBaseLayer(self.hidden_dim, self.edge_feat_dim, self.num_r_gaussian,
+                                update_x=self.update_x, act_fn=self.act_fn, norm=self.norm)
+            layers.append(layer)
+        return nn.ModuleList(layers)
+
+    # todo: refactor
+    def _connect_edge(self, x, mask_ligand, batch):
+        # if self.cutoff_mode == 'radius':
+        #     edge_index = radius_graph(x, r=self.r, batch=batch, flow='source_to_target')
+        if self.cutoff_mode == 'knn':
+            edge_index = knn_graph(x, k=self.k, batch=batch, flow='source_to_target')
+        elif self.cutoff_mode == 'hybrid':
+            edge_index = batch_hybrid_edge_connection(
+                x, k=self.k, mask_ligand=mask_ligand, batch=batch, add_p_index=True)
+        else:
+            raise ValueError(f'Not supported cutoff mode: {self.cutoff_mode}')
+        return edge_index
+
+    # todo: refactor
+    @staticmethod
+    def _build_edge_type(edge_index, mask_ligand):
+        src, dst = edge_index
+        edge_type = torch.zeros(len(src)).to(edge_index)
+        n_src = mask_ligand[src] == 1
+        n_dst = mask_ligand[dst] == 1
+        edge_type[n_src & n_dst] = 0
+        edge_type[n_src & ~n_dst] = 1
+        edge_type[~n_src & n_dst] = 2
+        edge_type[~n_src & ~n_dst] = 3
+        edge_type = F.one_hot(edge_type, num_classes=4)
+        return edge_type
+
+    def forward(self, h, x, mask_ligand, batch, return_all=False):
+        all_x = [x]
+        all_h = [h]
+        for l_idx, layer in enumerate(self.net):
+            edge_index = self._connect_edge(x, mask_ligand, batch)
+            edge_type = self._build_edge_type(edge_index, mask_ligand)
+            h, x = layer(h, x, edge_index, mask_ligand, edge_attr=edge_type)
+            all_x.append(x)
+            all_h.append(h)
+        outputs = {'x': x, 'h': h}
+        if return_all:
+            outputs.update({'all_x': all_x, 'all_h': all_h})
+        return outputs

models/molopt_score_model.py

@@ -0,0 +1,708 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_scatter import scatter_sum, scatter_mean
+from tqdm.auto import tqdm
+
+from .common import compose_context, ShiftedSoftplus
+from .egnn import EGNN
+from .uni_transformer import UniTransformerO2TwoUpdateGeneral
+
+
+def get_refine_net(refine_net_type, config):
+    if refine_net_type == 'uni_o2':
+        refine_net = UniTransformerO2TwoUpdateGeneral(
+            num_blocks=config.num_blocks,
+            num_layers=config.num_layers,
+            hidden_dim=config.hidden_dim,
+            n_heads=config.n_heads,
+            k=config.knn,
+            edge_feat_dim=config.edge_feat_dim,
+            num_r_gaussian=config.num_r_gaussian,
+            num_node_types=config.num_node_types,
+            act_fn=config.act_fn,
+            norm=config.norm,
+            cutoff_mode=config.cutoff_mode,
+            ew_net_type=config.ew_net_type,
+            num_x2h=config.num_x2h,
+            num_h2x=config.num_h2x,
+            r_max=config.r_max,
+            x2h_out_fc=config.x2h_out_fc,
+            sync_twoup=config.sync_twoup
+        )
+    elif refine_net_type == 'egnn':
+        refine_net = EGNN(
+            num_layers=config.num_layers,
+            hidden_dim=config.hidden_dim,
+            edge_feat_dim=config.edge_feat_dim,
+            num_r_gaussian=1,
+            k=config.knn,
+            cutoff_mode=config.cutoff_mode
+        )
+    else:
+        raise ValueError(refine_net_type)
+    return refine_net
+
+
+def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps):
+    def sigmoid(x):
+        return 1 / (np.exp(-x) + 1)
+
+    if beta_schedule == "quad":
+        betas = (
+                np.linspace(
+                    beta_start ** 0.5,
+                    beta_end ** 0.5,
+                    num_diffusion_timesteps,
+                    dtype=np.float64,
+                )
+                ** 2
+        )
+    elif beta_schedule == "linear":
+        betas = np.linspace(
+            beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64
+        )
+    elif beta_schedule == "const":
+        betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64)
+    elif beta_schedule == "jsd":  # 1/T, 1/(T-1), 1/(T-2), ..., 1
+        betas = 1.0 / np.linspace(
+            num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64
+        )
+    elif beta_schedule == "sigmoid":
+        betas = np.linspace(-6, 6, num_diffusion_timesteps)
+        betas = sigmoid(betas) * (beta_end - beta_start) + beta_start
+    else:
+        raise NotImplementedError(beta_schedule)
+    assert betas.shape == (num_diffusion_timesteps,)
+    return betas
+
+
+def cosine_beta_schedule(timesteps, s=0.008):
+    """
+    cosine schedule
+    as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
+    """
+    steps = timesteps + 1
+    x = np.linspace(0, steps, steps)
+    alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
+    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+    alphas = (alphas_cumprod[1:] / alphas_cumprod[:-1])
+
+    alphas = np.clip(alphas, a_min=0.001, a_max=1.)
+
+    # Use sqrt of this, so the alpha in our paper is the alpha_sqrt from the
+    # Gaussian diffusion in Ho et al.
+    alphas = np.sqrt(alphas)
+    return alphas
+
+
+def get_distance(pos, edge_index):
+    return (pos[edge_index[0]] - pos[edge_index[1]]).norm(dim=-1)
+
+
+def to_torch_const(x):
+    x = torch.from_numpy(x).float()
+    x = nn.Parameter(x, requires_grad=False)
+    return x
+
+
+def center_pos(protein_pos, ligand_pos, batch_protein, batch_ligand, mode='protein'):
+    if mode == 'none':
+        offset = 0.
+        pass
+    elif mode == 'protein':
+        offset = scatter_mean(protein_pos, batch_protein, dim=0)
+        protein_pos = protein_pos - offset[batch_protein]
+        ligand_pos = ligand_pos - offset[batch_ligand]
+    else:
+        raise NotImplementedError
+    return protein_pos, ligand_pos, offset
+
+
+# %% categorical diffusion related
+def index_to_log_onehot(x, num_classes):
+    assert x.max().item() < num_classes, f'Error: {x.max().item()} >= {num_classes}'
+    x_onehot = F.one_hot(x, num_classes)
+    # permute_order = (0, -1) + tuple(range(1, len(x.size())))
+    # x_onehot = x_onehot.permute(permute_order)
+    log_x = torch.log(x_onehot.float().clamp(min=1e-30))
+    return log_x
+
+
+def log_onehot_to_index(log_x):
+    return log_x.argmax(1)
+
+
+def categorical_kl(log_prob1, log_prob2):
+    kl = (log_prob1.exp() * (log_prob1 - log_prob2)).sum(dim=1)
+    return kl
+
+
+def log_categorical(log_x_start, log_prob):
+    return (log_x_start.exp() * log_prob).sum(dim=1)
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+    """
+    KL divergence between normal distributions parameterized by mean and log-variance.
+    """
+    kl = 0.5 * (-1.0 + logvar2 - logvar1 + torch.exp(logvar1 - logvar2) + (mean1 - mean2) ** 2 * torch.exp(-logvar2))
+    return kl.sum(-1)
+
+
+def log_normal(values, means, log_scales):
+    var = torch.exp(log_scales * 2)
+    log_prob = -((values - means) ** 2) / (2 * var) - log_scales - np.log(np.sqrt(2 * np.pi))
+    return log_prob.sum(-1)
+
+
+def log_sample_categorical(logits):
+    uniform = torch.rand_like(logits)
+    gumbel_noise = -torch.log(-torch.log(uniform + 1e-30) + 1e-30)
+    sample_index = (gumbel_noise + logits).argmax(dim=-1)
+    # sample_onehot = F.one_hot(sample, self.num_classes)
+    # log_sample = index_to_log_onehot(sample, self.num_classes)
+    return sample_index
+
+
+def log_1_min_a(a):
+    return np.log(1 - np.exp(a) + 1e-40)
+
+
+def log_add_exp(a, b):
+    maximum = torch.max(a, b)
+    return maximum + torch.log(torch.exp(a - maximum) + torch.exp(b - maximum))
+
+
+# %%
+
+
+# Time embedding
+class SinusoidalPosEmb(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        device = x.device
+        half_dim = self.dim // 2
+        emb = np.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
+        emb = x[:, None] * emb[None, :]
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+# Model
+class ScorePosNet3D(nn.Module):
+
+    def __init__(self, config, protein_atom_feature_dim, ligand_atom_feature_dim):
+        super().__init__()
+        self.config = config
+
+        # variance schedule
+        self.model_mean_type = config.model_mean_type  # ['noise', 'C0']
+        self.loss_v_weight = config.loss_v_weight
+        # self.v_mode = config.v_mode
+        # assert self.v_mode == 'categorical'
+        # self.v_net_type = getattr(config, 'v_net_type', 'mlp')
+        # self.bond_loss = getattr(config, 'bond_loss', False)
+        # self.bond_net_type = getattr(config, 'bond_net_type', 'pre_att')
+        # self.loss_bond_weight = getattr(config, 'loss_bond_weight', 0.)
+        # self.loss_non_bond_weight = getattr(config, 'loss_non_bond_weight', 0.)
+
+        self.sample_time_method = config.sample_time_method  # ['importance', 'symmetric']
+        # self.loss_pos_type = config.loss_pos_type  # ['mse', 'kl']
+        # print(f'Loss pos mode {self.loss_pos_type} applied!')
+        # print(f'Loss bond net type: {self.bond_net_type} '
+        #       f'bond weight: {self.loss_bond_weight} non bond weight: {self.loss_non_bond_weight}')
+
+        if config.beta_schedule == 'cosine':
+            alphas = cosine_beta_schedule(config.num_diffusion_timesteps, config.pos_beta_s) ** 2
+            # print('cosine pos alpha schedule applied!')
+            betas = 1. - alphas
+        else:
+            betas = get_beta_schedule(
+                beta_schedule=config.beta_schedule,
+                beta_start=config.beta_start,
+                beta_end=config.beta_end,
+                num_diffusion_timesteps=config.num_diffusion_timesteps,
+            )
+            alphas = 1. - betas
+        alphas_cumprod = np.cumprod(alphas, axis=0)
+        alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
+
+        self.betas = to_torch_const(betas)
+        self.num_timesteps = self.betas.size(0)
+        self.alphas_cumprod = to_torch_const(alphas_cumprod)
+        self.alphas_cumprod_prev = to_torch_const(alphas_cumprod_prev)
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.sqrt_alphas_cumprod = to_torch_const(np.sqrt(alphas_cumprod))
+        self.sqrt_one_minus_alphas_cumprod = to_torch_const(np.sqrt(1. - alphas_cumprod))
+        self.sqrt_recip_alphas_cumprod = to_torch_const(np.sqrt(1. / alphas_cumprod))
+        self.sqrt_recipm1_alphas_cumprod = to_torch_const(np.sqrt(1. / alphas_cumprod - 1))
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
+        self.posterior_mean_c0_coef = to_torch_const(betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))
+        self.posterior_mean_ct_coef = to_torch_const(
+            (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod))
+        # log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+        self.posterior_var = to_torch_const(posterior_variance)
+        self.posterior_logvar = to_torch_const(np.log(np.append(self.posterior_var[1], self.posterior_var[1:])))
+
+        # atom type diffusion schedule in log space
+        if config.v_beta_schedule == 'cosine':
+            alphas_v = cosine_beta_schedule(self.num_timesteps, config.v_beta_s)
+            # print('cosine v alpha schedule applied!')
+        else:
+            raise NotImplementedError
+        log_alphas_v = np.log(alphas_v)
+        log_alphas_cumprod_v = np.cumsum(log_alphas_v)
+        self.log_alphas_v = to_torch_const(log_alphas_v)
+        self.log_one_minus_alphas_v = to_torch_const(log_1_min_a(log_alphas_v))
+        self.log_alphas_cumprod_v = to_torch_const(log_alphas_cumprod_v)
+        self.log_one_minus_alphas_cumprod_v = to_torch_const(log_1_min_a(log_alphas_cumprod_v))
+
+        self.register_buffer('Lt_history', torch.zeros(self.num_timesteps))
+        self.register_buffer('Lt_count', torch.zeros(self.num_timesteps))
+
+        # model definition
+        self.hidden_dim = config.hidden_dim
+        self.num_classes = ligand_atom_feature_dim
+        if self.config.node_indicator:
+            emb_dim = self.hidden_dim - 1
+        else:
+            emb_dim = self.hidden_dim
+
+        # atom embedding
+        self.protein_atom_emb = nn.Linear(protein_atom_feature_dim, emb_dim)
+
+        # center pos
+        self.center_pos_mode = config.center_pos_mode  # ['none', 'protein']
+
+        # time embedding
+        self.time_emb_dim = config.time_emb_dim
+        self.time_emb_mode = config.time_emb_mode  # ['simple', 'sin']
+        if self.time_emb_dim > 0:
+            if self.time_emb_mode == 'simple':
+                self.ligand_atom_emb = nn.Linear(ligand_atom_feature_dim + 1, emb_dim)
+            elif self.time_emb_mode == 'sin':
+                self.time_emb = nn.Sequential(
+                    SinusoidalPosEmb(self.time_emb_dim),
+                    nn.Linear(self.time_emb_dim, self.time_emb_dim * 4),
+                    nn.GELU(),
+                    nn.Linear(self.time_emb_dim * 4, self.time_emb_dim)
+                )
+                self.ligand_atom_emb = nn.Linear(ligand_atom_feature_dim + self.time_emb_dim, emb_dim)
+            else:
+                raise NotImplementedError
+        else:
+            self.ligand_atom_emb = nn.Linear(ligand_atom_feature_dim, emb_dim)
+
+        self.refine_net_type = config.model_type
+        self.refine_net = get_refine_net(self.refine_net_type, config)
+        self.v_inference = nn.Sequential(
+            nn.Linear(self.hidden_dim, self.hidden_dim),
+            ShiftedSoftplus(),
+            nn.Linear(self.hidden_dim, ligand_atom_feature_dim),
+        )
+
+    def forward(self, protein_pos, protein_v, batch_protein, init_ligand_pos, init_ligand_v, batch_ligand,
+                time_step=None, return_all=False, fix_x=False):
+
+        batch_size = batch_protein.max().item() + 1
+        init_ligand_v = F.one_hot(init_ligand_v, self.num_classes).float()
+        # time embedding
+        if self.time_emb_dim > 0:
+            if self.time_emb_mode == 'simple':
+                input_ligand_feat = torch.cat([
+                    init_ligand_v,
+                    (time_step / self.num_timesteps)[batch_ligand].unsqueeze(-1)
+                ], -1)
+            elif self.time_emb_mode == 'sin':
+                time_feat = self.time_emb(time_step)
+                input_ligand_feat = torch.cat([init_ligand_v, time_feat], -1)
+            else:
+                raise NotImplementedError
+        else:
+            input_ligand_feat = init_ligand_v
+
+        h_protein = self.protein_atom_emb(protein_v)
+        init_ligand_h = self.ligand_atom_emb(input_ligand_feat)
+
+        if self.config.node_indicator:
+            h_protein = torch.cat([h_protein, torch.zeros(len(h_protein), 1).to(h_protein)], -1)
+            init_ligand_h = torch.cat([init_ligand_h, torch.ones(len(init_ligand_h), 1).to(h_protein)], -1)
+
+        h_all, pos_all, batch_all, mask_ligand = compose_context(
+            h_protein=h_protein,
+            h_ligand=init_ligand_h,
+            pos_protein=protein_pos,
+            pos_ligand=init_ligand_pos,
+            batch_protein=batch_protein,
+            batch_ligand=batch_ligand,
+        )
+
+        outputs = self.refine_net(h_all, pos_all, mask_ligand, batch_all, return_all=return_all, fix_x=fix_x)
+        final_pos, final_h = outputs['x'], outputs['h']
+        final_ligand_pos, final_ligand_h = final_pos[mask_ligand], final_h[mask_ligand]
+        final_ligand_v = self.v_inference(final_ligand_h)
+
+        preds = {
+            'pred_ligand_pos': final_ligand_pos,
+            'pred_ligand_v': final_ligand_v,
+            'final_h': final_h,
+            'final_ligand_h': final_ligand_h
+        }
+        if return_all:
+            final_all_pos, final_all_h = outputs['all_x'], outputs['all_h']
+            final_all_ligand_pos = [pos[mask_ligand] for pos in final_all_pos]
+            final_all_ligand_v = [self.v_inference(h[mask_ligand]) for h in final_all_h]
+            preds.update({
+                'layer_pred_ligand_pos': final_all_ligand_pos,
+                'layer_pred_ligand_v': final_all_ligand_v
+            })
+        return preds
+
+    # atom type diffusion process
+    def q_v_pred_one_timestep(self, log_vt_1, t, batch):
+        # q(vt | vt-1)
+        log_alpha_t = extract(self.log_alphas_v, t, batch)
+        log_1_min_alpha_t = extract(self.log_one_minus_alphas_v, t, batch)
+
+        # alpha_t * vt + (1 - alpha_t) 1 / K
+        log_probs = log_add_exp(
+            log_vt_1 + log_alpha_t,
+            log_1_min_alpha_t - np.log(self.num_classes)
+        )
+        return log_probs
+
+    def q_v_pred(self, log_v0, t, batch):
+        # compute q(vt | v0)
+        log_cumprod_alpha_t = extract(self.log_alphas_cumprod_v, t, batch)
+        log_1_min_cumprod_alpha = extract(self.log_one_minus_alphas_cumprod_v, t, batch)
+
+        log_probs = log_add_exp(
+            log_v0 + log_cumprod_alpha_t,
+            log_1_min_cumprod_alpha - np.log(self.num_classes)
+        )
+        return log_probs
+
+    def q_v_sample(self, log_v0, t, batch):
+        log_qvt_v0 = self.q_v_pred(log_v0, t, batch)
+        sample_index = log_sample_categorical(log_qvt_v0)
+        log_sample = index_to_log_onehot(sample_index, self.num_classes)
+        return sample_index, log_sample
+
+    # atom type generative process
+    def q_v_posterior(self, log_v0, log_vt, t, batch):
+        # q(vt-1 | vt, v0) = q(vt | vt-1, x0) * q(vt-1 | x0) / q(vt | x0)
+        t_minus_1 = t - 1
+        # Remove negative values, will not be used anyway for final decoder
+        t_minus_1 = torch.where(t_minus_1 < 0, torch.zeros_like(t_minus_1), t_minus_1)
+        log_qvt1_v0 = self.q_v_pred(log_v0, t_minus_1, batch)
+        unnormed_logprobs = log_qvt1_v0 + self.q_v_pred_one_timestep(log_vt, t, batch)
+        log_vt1_given_vt_v0 = unnormed_logprobs - torch.logsumexp(unnormed_logprobs, dim=-1, keepdim=True)
+        return log_vt1_given_vt_v0
+
+    def kl_v_prior(self, log_x_start, batch):
+        num_graphs = batch.max().item() + 1
+        log_qxT_prob = self.q_v_pred(log_x_start, t=[self.num_timesteps - 1] * num_graphs, batch=batch)
+        log_half_prob = -torch.log(self.num_classes * torch.ones_like(log_qxT_prob))
+        kl_prior = categorical_kl(log_qxT_prob, log_half_prob)
+        kl_prior = scatter_mean(kl_prior, batch, dim=0)
+        return kl_prior
+
+    def _predict_x0_from_eps(self, xt, eps, t, batch):
+        pos0_from_e = extract(self.sqrt_recip_alphas_cumprod, t, batch) * xt - \
+                      extract(self.sqrt_recipm1_alphas_cumprod, t, batch) * eps
+        return pos0_from_e
+
+    def q_pos_posterior(self, x0, xt, t, batch):
+        # Compute the mean and variance of the diffusion posterior q(x_{t-1} | x_t, x_0)
+        pos_model_mean = extract(self.posterior_mean_c0_coef, t, batch) * x0 + \
+                         extract(self.posterior_mean_ct_coef, t, batch) * xt
+        return pos_model_mean
+
+    def kl_pos_prior(self, pos0, batch):
+        num_graphs = batch.max().item() + 1
+        a_pos = extract(self.alphas_cumprod, [self.num_timesteps - 1] * num_graphs, batch)  # (num_ligand_atoms, 1)
+        pos_model_mean = a_pos.sqrt() * pos0
+        pos_log_variance = torch.log((1.0 - a_pos).sqrt())
+        kl_prior = normal_kl(torch.zeros_like(pos_model_mean), torch.zeros_like(pos_log_variance),
+                             pos_model_mean, pos_log_variance)
+        kl_prior = scatter_mean(kl_prior, batch, dim=0)
+        return kl_prior
+
+    def sample_time(self, num_graphs, device, method):
+        if method == 'importance':
+            if not (self.Lt_count > 10).all():
+                return self.sample_time(num_graphs, device, method='symmetric')
+
+            Lt_sqrt = torch.sqrt(self.Lt_history + 1e-10) + 0.0001
+            Lt_sqrt[0] = Lt_sqrt[1]  # Overwrite decoder term with L1.
+            pt_all = Lt_sqrt / Lt_sqrt.sum()
+
+            time_step = torch.multinomial(pt_all, num_samples=num_graphs, replacement=True)
+            pt = pt_all.gather(dim=0, index=time_step)
+            return time_step, pt
+
+        elif method == 'symmetric':
+            time_step = torch.randint(
+                0, self.num_timesteps, size=(num_graphs // 2 + 1,), device=device)
+            time_step = torch.cat(
+                [time_step, self.num_timesteps - time_step - 1], dim=0)[:num_graphs]
+            pt = torch.ones_like(time_step).float() / self.num_timesteps
+            return time_step, pt
+
+        else:
+            raise ValueError
+
+    def compute_pos_Lt(self, pos_model_mean, x0, xt, t, batch):
+        # fixed pos variance
+        pos_log_variance = extract(self.posterior_logvar, t, batch)
+        pos_true_mean = self.q_pos_posterior(x0=x0, xt=xt, t=t, batch=batch)
+        kl_pos = normal_kl(pos_true_mean, pos_log_variance, pos_model_mean, pos_log_variance)
+        kl_pos = kl_pos / np.log(2.)
+
+        decoder_nll_pos = -log_normal(x0, means=pos_model_mean, log_scales=0.5 * pos_log_variance)
+        assert kl_pos.shape == decoder_nll_pos.shape
+        mask = (t == 0).float()[batch]
+        loss_pos = scatter_mean(mask * decoder_nll_pos + (1. - mask) * kl_pos, batch, dim=0)
+        return loss_pos
+
+    def compute_v_Lt(self, log_v_model_prob, log_v0, log_v_true_prob, t, batch):
+        kl_v = categorical_kl(log_v_true_prob, log_v_model_prob)  # [num_atoms, ]
+        decoder_nll_v = -log_categorical(log_v0, log_v_model_prob)  # L0
+        assert kl_v.shape == decoder_nll_v.shape
+        mask = (t == 0).float()[batch]
+        loss_v = scatter_mean(mask * decoder_nll_v + (1. - mask) * kl_v, batch, dim=0)
+        return loss_v
+
+    def get_diffusion_loss(
+            self, protein_pos, protein_v, batch_protein, ligand_pos, ligand_v, batch_ligand, time_step=None
+    ):
+        num_graphs = batch_protein.max().item() + 1
+        protein_pos, ligand_pos, _ = center_pos(
+            protein_pos, ligand_pos, batch_protein, batch_ligand, mode=self.center_pos_mode)
+
+        # 1. sample noise levels
+        if time_step is None:
+            time_step, pt = self.sample_time(num_graphs, protein_pos.device, self.sample_time_method)
+        else:
+            pt = torch.ones_like(time_step).float() / self.num_timesteps
+        a = self.alphas_cumprod.index_select(0, time_step)  # (num_graphs, )
+
+        # 2. perturb pos and v
+        a_pos = a[batch_ligand].unsqueeze(-1)  # (num_ligand_atoms, 1)
+        pos_noise = torch.zeros_like(ligand_pos)
+        pos_noise.normal_()
+        # Xt = a.sqrt() * X0 + (1-a).sqrt() * eps
+        ligand_pos_perturbed = a_pos.sqrt() * ligand_pos + (1.0 - a_pos).sqrt() * pos_noise  # pos_noise * std
+        # Vt = a * V0 + (1-a) / K
+        log_ligand_v0 = index_to_log_onehot(ligand_v, self.num_classes)
+        ligand_v_perturbed, log_ligand_vt = self.q_v_sample(log_ligand_v0, time_step, batch_ligand)
+
+        # 3. forward-pass NN, feed perturbed pos and v, output noise
+        preds = self(
+            protein_pos=protein_pos,
+            protein_v=protein_v,
+            batch_protein=batch_protein,
+
+            init_ligand_pos=ligand_pos_perturbed,
+            init_ligand_v=ligand_v_perturbed,
+            batch_ligand=batch_ligand,
+            time_step=time_step
+        )
+
+        pred_ligand_pos, pred_ligand_v = preds['pred_ligand_pos'], preds['pred_ligand_v']
+        pred_pos_noise = pred_ligand_pos - ligand_pos_perturbed
+        # atom position
+        if self.model_mean_type == 'noise':
+            pos0_from_e = self._predict_x0_from_eps(
+                xt=ligand_pos_perturbed, eps=pred_pos_noise, t=time_step, batch=batch_ligand)
+            pos_model_mean = self.q_pos_posterior(
+                x0=pos0_from_e, xt=ligand_pos_perturbed, t=time_step, batch=batch_ligand)
+        elif self.model_mean_type == 'C0':
+            pos_model_mean = self.q_pos_posterior(
+                x0=pred_ligand_pos, xt=ligand_pos_perturbed, t=time_step, batch=batch_ligand)
+        else:
+            raise ValueError
+
+        # atom pos loss
+        if self.model_mean_type == 'C0':
+            target, pred = ligand_pos, pred_ligand_pos
+        elif self.model_mean_type == 'noise':
+            target, pred = pos_noise, pred_pos_noise
+        else:
+            raise ValueError
+        loss_pos = scatter_mean(((pred - target) ** 2).sum(-1), batch_ligand, dim=0)
+        loss_pos = torch.mean(loss_pos)
+
+        # atom type loss
+        log_ligand_v_recon = F.log_softmax(pred_ligand_v, dim=-1)
+        log_v_model_prob = self.q_v_posterior(log_ligand_v_recon, log_ligand_vt, time_step, batch_ligand)
+        log_v_true_prob = self.q_v_posterior(log_ligand_v0, log_ligand_vt, time_step, batch_ligand)
+        kl_v = self.compute_v_Lt(log_v_model_prob=log_v_model_prob, log_v0=log_ligand_v0,
+                                 log_v_true_prob=log_v_true_prob, t=time_step, batch=batch_ligand)
+        loss_v = torch.mean(kl_v)
+        loss = loss_pos + loss_v * self.loss_v_weight
+
+        return {
+            'loss_pos': loss_pos,
+            'loss_v': loss_v,
+            'loss': loss,
+            'x0': ligand_pos,
+            'pred_ligand_pos': pred_ligand_pos,
+            'pred_ligand_v': pred_ligand_v,
+            'pred_pos_noise': pred_pos_noise,
+            'ligand_v_recon': F.softmax(pred_ligand_v, dim=-1)
+        }
+
+    @torch.no_grad()
+    def likelihood_estimation(
+            self, protein_pos, protein_v, batch_protein, ligand_pos, ligand_v, batch_ligand, time_step
+    ):
+        protein_pos, ligand_pos, _ = center_pos(
+            protein_pos, ligand_pos, batch_protein, batch_ligand, mode='protein')
+        assert (time_step == self.num_timesteps).all() or (time_step < self.num_timesteps).all()
+        if (time_step == self.num_timesteps).all():
+            kl_pos_prior = self.kl_pos_prior(ligand_pos, batch_ligand)
+            log_ligand_v0 = index_to_log_onehot(batch_ligand, self.num_classes)
+            kl_v_prior = self.kl_v_prior(log_ligand_v0, batch_ligand)
+            return kl_pos_prior, kl_v_prior
+
+        # perturb pos and v
+        a = self.alphas_cumprod.index_select(0, time_step)  # (num_graphs, )
+        a_pos = a[batch_ligand].unsqueeze(-1)  # (num_ligand_atoms, 1)
+        pos_noise = torch.zeros_like(ligand_pos)
+        pos_noise.normal_()
+        # Xt = a.sqrt() * X0 + (1-a).sqrt() * eps
+        ligand_pos_perturbed = a_pos.sqrt() * ligand_pos + (1.0 - a_pos).sqrt() * pos_noise  # pos_noise * std
+        # Vt = a * V0 + (1-a) / K
+        log_ligand_v0 = index_to_log_onehot(ligand_v, self.num_classes)
+        ligand_v_perturbed, log_ligand_vt = self.q_v_sample(log_ligand_v0, time_step, batch_ligand)
+
+        preds = self(
+            protein_pos=protein_pos,
+            protein_v=protein_v,
+            batch_protein=batch_protein,
+
+            init_ligand_pos=ligand_pos_perturbed,
+            init_ligand_v=ligand_v_perturbed,
+            batch_ligand=batch_ligand,
+            time_step=time_step
+        )
+
+        pred_ligand_pos, pred_ligand_v = preds['pred_ligand_pos'], preds['pred_ligand_v']
+        if self.model_mean_type == 'C0':
+            pos_model_mean = self.q_pos_posterior(
+                x0=pred_ligand_pos, xt=ligand_pos_perturbed, t=time_step, batch=batch_ligand)
+        else:
+            raise ValueError
+
+        # atom type
+        log_ligand_v_recon = F.log_softmax(pred_ligand_v, dim=-1)
+        log_v_model_prob = self.q_v_posterior(log_ligand_v_recon, log_ligand_vt, time_step, batch_ligand)
+        log_v_true_prob = self.q_v_posterior(log_ligand_v0, log_ligand_vt, time_step, batch_ligand)
+
+        # t = [T-1, ... , 0]
+        kl_pos = self.compute_pos_Lt(pos_model_mean=pos_model_mean, x0=ligand_pos,
+                                     xt=ligand_pos_perturbed, t=time_step, batch=batch_ligand)
+        kl_v = self.compute_v_Lt(log_v_model_prob=log_v_model_prob, log_v0=log_ligand_v0,
+                                 log_v_true_prob=log_v_true_prob, t=time_step, batch=batch_ligand)
+        return kl_pos, kl_v
+
+    @torch.no_grad()
+    def fetch_embedding(self, protein_pos, protein_v, batch_protein, ligand_pos, ligand_v, batch_ligand):
+        preds = self(
+            protein_pos=protein_pos,
+            protein_v=protein_v,
+            batch_protein=batch_protein,
+
+            init_ligand_pos=ligand_pos,
+            init_ligand_v=ligand_v,
+            batch_ligand=batch_ligand,
+            fix_x=True
+        )
+        return preds
+
+    @torch.no_grad()
+    def sample_diffusion(self, protein_pos, protein_v, batch_protein,
+                         init_ligand_pos, init_ligand_v, batch_ligand,
+                         num_steps=None, center_pos_mode=None, pos_only=False):
+
+        if num_steps is None:
+            num_steps = self.num_timesteps
+        num_graphs = batch_protein.max().item() + 1
+
+        protein_pos, init_ligand_pos, offset = center_pos(
+            protein_pos, init_ligand_pos, batch_protein, batch_ligand, mode=center_pos_mode)
+
+        pos_traj, v_traj = [], []
+        v0_pred_traj, vt_pred_traj = [], []
+        ligand_pos, ligand_v = init_ligand_pos, init_ligand_v
+        # time sequence
+        time_seq = list(reversed(range(self.num_timesteps - num_steps, self.num_timesteps)))
+        for i in tqdm(time_seq, desc='sampling', total=len(time_seq)):
+            t = torch.full(size=(num_graphs,), fill_value=i, dtype=torch.long, device=protein_pos.device)
+            preds = self(
+                protein_pos=protein_pos,
+                protein_v=protein_v,
+                batch_protein=batch_protein,
+
+                init_ligand_pos=ligand_pos,
+                init_ligand_v=ligand_v,
+                batch_ligand=batch_ligand,
+                time_step=t
+            )
+            # Compute posterior mean and variance
+            if self.model_mean_type == 'noise':
+                pred_pos_noise = preds['pred_ligand_pos'] - ligand_pos
+                pos0_from_e = self._predict_x0_from_eps(xt=ligand_pos, eps=pred_pos_noise, t=t, batch=batch_ligand)
+                v0_from_e = preds['pred_ligand_v']
+            elif self.model_mean_type == 'C0':
+                pos0_from_e = preds['pred_ligand_pos']
+                v0_from_e = preds['pred_ligand_v']
+            else:
+                raise ValueError
+
+            pos_model_mean = self.q_pos_posterior(x0=pos0_from_e, xt=ligand_pos, t=t, batch=batch_ligand)
+            pos_log_variance = extract(self.posterior_logvar, t, batch_ligand)
+            # no noise when t == 0
+            nonzero_mask = (1 - (t == 0).float())[batch_ligand].unsqueeze(-1)
+            ligand_pos_next = pos_model_mean + nonzero_mask * (0.5 * pos_log_variance).exp() * torch.randn_like(
+                ligand_pos)
+            ligand_pos = ligand_pos_next
+
+            if not pos_only:
+                log_ligand_v_recon = F.log_softmax(v0_from_e, dim=-1)
+                log_ligand_v = index_to_log_onehot(ligand_v, self.num_classes)
+                log_model_prob = self.q_v_posterior(log_ligand_v_recon, log_ligand_v, t, batch_ligand)
+                ligand_v_next = log_sample_categorical(log_model_prob)
+
+                v0_pred_traj.append(log_ligand_v_recon.clone().cpu())
+                vt_pred_traj.append(log_model_prob.clone().cpu())
+                ligand_v = ligand_v_next
+
+            ori_ligand_pos = ligand_pos + offset[batch_ligand]
+            pos_traj.append(ori_ligand_pos.clone().cpu())
+            v_traj.append(ligand_v.clone().cpu())
+
+        ligand_pos = ligand_pos + offset[batch_ligand]
+        return {
+            'pos': ligand_pos,
+            'v': ligand_v,
+            'pos_traj': pos_traj,
+            'v_traj': v_traj,
+            'v0_traj': v0_pred_traj,
+            'vt_traj': vt_pred_traj
+        }
+
+
+def extract(coef, t, batch):
+    out = coef[t][batch]
+    return out.unsqueeze(-1)

models/property_pred/__init__.py

@@ -0,0 +1,2 @@
+"""TargetDiff property prediction models."""
+

models/property_pred/prop_egnn.py

@@ -0,0 +1,83 @@
+import torch
+import torch.nn as nn
+from torch_scatter import scatter_sum
+from torch_geometric.nn import radius_graph, knn_graph
+from ..common import GaussianSmearing, MLP
+
+
+class EnBaseLayer(nn.Module):
+    def __init__(self, hidden_dim, edge_feat_dim, num_r_gaussian, update_x=True, act_fn='relu', norm=False):
+        super().__init__()
+        self.r_min = 0.
+        self.r_max = 10. ** 2
+        self.hidden_dim = hidden_dim
+        self.num_r_gaussian = num_r_gaussian
+        self.edge_feat_dim = edge_feat_dim
+        self.update_x = update_x
+        self.act_fn = act_fn
+        self.norm = norm
+        if num_r_gaussian > 1:
+            self.r_expansion = GaussianSmearing(self.r_min, self.r_max, num_gaussians=num_r_gaussian, fixed_offset=False)
+        self.edge_mlp = MLP(2 * hidden_dim + edge_feat_dim + num_r_gaussian, hidden_dim, hidden_dim,
+                            num_layer=2, norm=norm, act_fn=act_fn, act_last=True)
+        self.edge_inf = nn.Sequential(nn.Linear(hidden_dim, 1), nn.Sigmoid())
+        if self.update_x:
+            self.x_mlp = MLP(hidden_dim, 1, hidden_dim, num_layer=2, norm=norm, act_fn=act_fn)
+        self.node_mlp = MLP(2 * hidden_dim, hidden_dim, hidden_dim, num_layer=2, norm=norm, act_fn=act_fn)
+
+    def forward(self, h, edge_index, edge_attr):
+        dst, src = edge_index
+        hi, hj = h[dst], h[src]
+        # \phi_e in Eq(3)
+        mij = self.edge_mlp(torch.cat([edge_attr, hi, hj], -1))
+        eij = self.edge_inf(mij)
+        mi = scatter_sum(mij * eij, dst, dim=0, dim_size=h.shape[0])
+
+        # h update in Eq(6)
+        # h = h + self.node_mlp(torch.cat([mi, h], -1))
+        output = self.node_mlp(torch.cat([mi, h], -1))
+        # if self.update_x:
+        #     # x update in Eq(4)
+        #     xi, xj = x[dst], x[src]
+        #     delta_x = scatter_sum((xi - xj) * self.x_mlp(mij), dst, dim=0)
+        #     x = x + delta_x
+
+        return output
+
+
+class EnEquiEncoder(nn.Module):
+    def __init__(self, num_layers, hidden_dim, edge_feat_dim, num_r_gaussian, k=32, cutoff=10.0,
+                 update_x=True, act_fn='relu', norm=False):
+        super().__init__()
+        # Build the network
+        self.num_layers = num_layers
+        self.hidden_dim = hidden_dim
+        self.edge_feat_dim = edge_feat_dim
+        self.num_r_gaussian = num_r_gaussian
+        self.update_x = update_x
+        self.act_fn = act_fn
+        self.norm = norm
+        self.k = k
+        self.cutoff = cutoff
+        self.distance_expansion = GaussianSmearing(stop=cutoff, num_gaussians=num_r_gaussian, fixed_offset=False)
+        self.net = self._build_network()
+
+    def _build_network(self):
+        # Equivariant layers
+        layers = []
+        for l_idx in range(self.num_layers):
+            layer = EnBaseLayer(self.hidden_dim, self.edge_feat_dim, self.num_r_gaussian,
+                                update_x=self.update_x, act_fn=self.act_fn, norm=self.norm)
+            layers.append(layer)
+        return nn.ModuleList(layers)
+
+    def forward(self, node_attr, pos, batch):
+        # edge_index = radius_graph(pos, self.cutoff, batch=batch, loop=False)
+        edge_index = knn_graph(pos, k=self.k, batch=batch, flow='target_to_source')
+        edge_length = torch.norm(pos[edge_index[0]] - pos[edge_index[1]], dim=1)
+        edge_attr = self.distance_expansion(edge_length)
+
+        h = node_attr
+        for interaction in self.net:
+            h = h + interaction(h, edge_index, edge_attr)
+        return h

models/property_pred/prop_model.py

@@ -0,0 +1,212 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_scatter import scatter
+
+from .prop_egnn import EnEquiEncoder
+from ..common import compose_context_prop, ShiftedSoftplus
+
+
+def get_encoder(config):
+    if config.name == 'egnn' or config.name == 'egnn_enc':
+        net = EnEquiEncoder(
+            num_layers=config.num_layers,
+            edge_feat_dim=config.edge_dim,
+            hidden_dim=config.hidden_dim,
+            num_r_gaussian=config.num_r_gaussian,
+            act_fn=config.act_fn,
+            norm=config.norm,
+            update_x=False,
+            k=config.knn,
+            cutoff=config.cutoff,
+        )
+    else:
+        raise ValueError(config.name)
+    return net
+
+
+class PropPredNet(nn.Module):
+    def __init__(self, config, protein_atom_feature_dim, ligand_atom_feature_dim, output_dim=3):
+        super(PropPredNet, self).__init__()
+        self.config = config
+        self.hidden_dim = config.hidden_channels
+        self.output_dim = output_dim
+        self.protein_atom_emb = nn.Linear(protein_atom_feature_dim, self.hidden_dim)
+        self.ligand_atom_emb = nn.Linear(ligand_atom_feature_dim, self.hidden_dim)
+
+        # self.mean = target_mean
+        # self.std = target_std
+        # self.register_buffer('target_mean', target_mean)
+        # self.register_buffer('target_std', target_std)
+        self.encoder = get_encoder(config.encoder)
+        self.out_block = nn.Sequential(
+            nn.Linear(self.hidden_dim, self.hidden_dim),
+            ShiftedSoftplus(),
+            nn.Linear(self.hidden_dim, output_dim),
+        )
+
+    def forward(self, protein_pos, protein_atom_feature, ligand_pos, ligand_atom_feature, batch_protein, batch_ligand,
+                output_kind):
+        h_protein = self.protein_atom_emb(protein_atom_feature)
+        h_ligand = self.ligand_atom_emb(ligand_atom_feature)
+
+        h_ctx, pos_ctx, batch_ctx = compose_context_prop(
+            h_protein=h_protein,
+            h_ligand=h_ligand,
+            pos_protein=protein_pos,
+            pos_ligand=ligand_pos,
+            batch_protein=batch_protein,
+            batch_ligand=batch_ligand,
+        )
+
+        h_ctx = self.encoder(
+            node_attr=h_ctx,
+            pos=pos_ctx,
+            batch=batch_ctx,
+        )  # (N_p+N_l, H)
+
+        # Aggregate messages
+        pre_out = scatter(h_ctx, index=batch_ctx, dim=0, reduce='sum')  # (N, F)
+        output = self.out_block(pre_out)  # (N, C)
+        if output_kind is not None:
+            output_mask = F.one_hot(output_kind - 1, self.output_dim)
+            output = torch.sum(output * output_mask, dim=-1, keepdim=True)
+        return output
+
+    def get_loss(self, batch, pos_noise_std, return_pred=False):
+        protein_noise = torch.randn_like(batch.protein_pos) * pos_noise_std
+        ligand_noise = torch.randn_like(batch.ligand_pos) * pos_noise_std
+        pred = self(
+            protein_pos=batch.protein_pos + protein_noise,
+            protein_atom_feature=batch.protein_atom_feature.float(),
+            ligand_pos=batch.ligand_pos + ligand_noise,
+            ligand_atom_feature=batch.ligand_atom_feature_full.float(),
+            batch_protein=batch.protein_element_batch,
+            batch_ligand=batch.ligand_element_batch,
+            output_kind=batch.kind,
+            # output_kind=None
+        )
+        # pred = pred * y_std + y_mean
+        loss_func = nn.MSELoss()
+        loss = loss_func(pred.view(-1), batch.y)
+        if return_pred:
+            return loss, pred
+        else:
+            return loss
+
+
+class PropPredNetEnc(nn.Module):
+    def __init__(self, config, protein_atom_feature_dim, ligand_atom_feature_dim,
+                 enc_ligand_dim, enc_node_dim, enc_graph_dim, enc_feature_type=None, output_dim=1):
+        super(PropPredNetEnc, self).__init__()
+        self.config = config
+        self.hidden_dim = config.hidden_channels
+        self.output_dim = output_dim
+        self.enc_ligand_dim = enc_ligand_dim
+        self.enc_node_dim = enc_node_dim
+        self.enc_graph_dim = enc_graph_dim
+        self.enc_feature_type = enc_feature_type
+
+        self.protein_atom_emb = nn.Linear(protein_atom_feature_dim, self.hidden_dim)
+        self.ligand_atom_emb = nn.Linear(ligand_atom_feature_dim + enc_ligand_dim, self.hidden_dim)
+        # self.mean = target_mean
+        # self.std = target_std
+        # self.register_buffer('target_mean', target_mean)
+        # self.register_buffer('target_std', target_std)
+        self.encoder = get_encoder(config.encoder)
+        if self.enc_node_dim > 0:
+            self.enc_node_layer = nn.Sequential(
+                nn.Linear(self.hidden_dim + self.enc_node_dim, self.hidden_dim),
+                nn.ReLU(),
+                nn.Linear(self.hidden_dim, self.hidden_dim),
+            )
+
+        self.out_block = nn.Sequential(
+            nn.Linear(self.hidden_dim + self.enc_graph_dim, self.hidden_dim),
+            ShiftedSoftplus(),
+            nn.Linear(self.hidden_dim, output_dim),
+        )
+
+    def forward(self, protein_pos, protein_atom_feature, ligand_pos, ligand_atom_feature, batch_protein, batch_ligand,
+                output_kind, enc_ligand_feature, enc_node_feature, enc_graph_feature):
+        h_protein = self.protein_atom_emb(protein_atom_feature)
+        if enc_ligand_feature is not None:
+            ligand_atom_feature = torch.cat([ligand_atom_feature, enc_ligand_feature], dim=-1)
+        h_ligand = self.ligand_atom_emb(ligand_atom_feature)
+
+        h_ctx, pos_ctx, batch_ctx = compose_context_prop(
+            h_protein=h_protein,
+            h_ligand=h_ligand,
+            pos_protein=protein_pos,
+            pos_ligand=ligand_pos,
+            batch_protein=batch_protein,
+            batch_ligand=batch_ligand,
+        )
+
+        h_ctx = self.encoder(
+            node_attr=h_ctx,
+            pos=pos_ctx,
+            batch=batch_ctx,
+        )  # (N_p+N_l, H)
+
+        if enc_node_feature is not None:
+            h_ctx = torch.cat([h_ctx, enc_node_feature], dim=-1)
+            h_ctx = self.enc_node_layer(h_ctx)
+
+        # Aggregate messages
+        pre_out = scatter(h_ctx, index=batch_ctx, dim=0, reduce='sum')  # (N, F)
+        if enc_graph_feature is not None:
+            pre_out = torch.cat([pre_out, enc_graph_feature], dim=-1)
+
+        output = self.out_block(pre_out)  # (N, C)
+        if output_kind is not None:
+            output_mask = F.one_hot(output_kind - 1, self.output_dim)
+            output = torch.sum(output * output_mask, dim=-1, keepdim=True)
+        return output
+
+    def get_loss(self, batch, pos_noise_std, return_pred=False):
+        protein_noise = torch.randn_like(batch.protein_pos) * pos_noise_std
+        ligand_noise = torch.randn_like(batch.ligand_pos) * pos_noise_std
+
+        # add features
+        enc_ligand_feature, enc_node_feature, enc_graph_feature = None, None, None
+        if self.enc_feature_type == 'nll_all':
+            enc_graph_feature = batch.nll_all  # [num_graphs, 22]
+        elif self.enc_feature_type == 'nll':
+            enc_graph_feature = batch.nll  # [num_graphs, 20]
+        elif self.enc_feature_type == 'final_h':
+            enc_node_feature = batch.final_h  # [num_pl_atoms, 128]
+        elif self.enc_feature_type == 'pred_ligand_v':
+            enc_ligand_feature = batch.pred_ligand_v  # [num_l_atoms, 13]
+        elif self.enc_feature_type == 'pred_v_entropy_pre':
+            enc_ligand_feature = batch.pred_v_entropy   # [num_l_atoms, 1]
+        elif self.enc_feature_type == 'pred_v_entropy_post':
+            enc_graph_feature = scatter(batch.pred_v_entropy, index=batch.ligand_element_batch, dim=0, reduce='sum')   # [num_graphs, 1]
+        elif self.enc_feature_type == 'full':
+            enc_graph_feature = torch.cat(
+                [batch.nll_all, scatter(batch.pred_v_entropy, index=batch.ligand_element_batch, dim=0, reduce='sum')], dim=-1)
+            enc_node_feature = batch.final_h
+            enc_ligand_feature = torch.cat([batch.pred_ligand_v, batch.pred_v_entropy], -1)
+        else:
+            raise NotImplementedError
+
+        pred = self(
+            protein_pos=batch.protein_pos + protein_noise,
+            protein_atom_feature=batch.protein_atom_feature.float(),
+            ligand_pos=batch.ligand_pos + ligand_noise,
+            ligand_atom_feature=batch.ligand_atom_feature_full.float(),
+            batch_protein=batch.protein_element_batch,
+            batch_ligand=batch.ligand_element_batch,
+            output_kind=batch.kind,
+            # output_kind=None,
+            enc_ligand_feature=enc_ligand_feature,
+            enc_node_feature=enc_node_feature,
+            enc_graph_feature=enc_graph_feature
+        )
+        # pred = pred * y_std + y_mean
+        loss_func = nn.MSELoss()
+        loss = loss_func(pred.view(-1), batch.y)
+        if return_pred:
+            return loss, pred
+        else:
+            return loss

models/uni_transformer.py

@@ -0,0 +1,328 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.nn import radius_graph, knn_graph
+from torch_scatter import scatter_softmax, scatter_sum
+
+from .common import GaussianSmearing, MLP, batch_hybrid_edge_connection, outer_product
+
+
+class BaseX2HAttLayer(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim, n_heads, edge_feat_dim, r_feat_dim,
+                 act_fn='relu', norm=True, ew_net_type='r', out_fc=True):
+        super().__init__()
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.output_dim = output_dim
+        self.n_heads = n_heads
+        self.act_fn = act_fn
+        self.edge_feat_dim = edge_feat_dim
+        self.r_feat_dim = r_feat_dim
+        self.ew_net_type = ew_net_type
+        self.out_fc = out_fc
+
+        # attention key func
+        kv_input_dim = input_dim * 2 + edge_feat_dim + r_feat_dim
+        self.hk_func = MLP(kv_input_dim, output_dim, hidden_dim, norm=norm, act_fn=act_fn)
+
+        # attention value func
+        self.hv_func = MLP(kv_input_dim, output_dim, hidden_dim, norm=norm, act_fn=act_fn)
+
+        # attention query func
+        self.hq_func = MLP(input_dim, output_dim, hidden_dim, norm=norm, act_fn=act_fn)
+        if ew_net_type == 'r':
+            self.ew_net = nn.Sequential(nn.Linear(r_feat_dim, 1), nn.Sigmoid())
+        elif ew_net_type == 'm':
+            self.ew_net = nn.Sequential(nn.Linear(output_dim, 1), nn.Sigmoid())
+
+        if self.out_fc:
+            self.node_output = MLP(2 * hidden_dim, hidden_dim, hidden_dim, norm=norm, act_fn=act_fn)
+
+    def forward(self, h, r_feat, edge_feat, edge_index, e_w=None):
+        N = h.size(0)
+        src, dst = edge_index
+        hi, hj = h[dst], h[src]
+
+        # multi-head attention
+        # decide inputs of k_func and v_func
+        kv_input = torch.cat([r_feat, hi, hj], -1)
+        if edge_feat is not None:
+            kv_input = torch.cat([edge_feat, kv_input], -1)
+
+        # compute k
+        k = self.hk_func(kv_input).view(-1, self.n_heads, self.output_dim // self.n_heads)
+        # compute v
+        v = self.hv_func(kv_input)
+
+        if self.ew_net_type == 'r':
+            e_w = self.ew_net(r_feat)
+        elif self.ew_net_type == 'm':
+            e_w = self.ew_net(v[..., :self.hidden_dim])
+        elif e_w is not None:
+            e_w = e_w.view(-1, 1)
+        else:
+            e_w = 1.
+        v = v * e_w
+        v = v.view(-1, self.n_heads, self.output_dim // self.n_heads)
+
+        # compute q
+        q = self.hq_func(h).view(-1, self.n_heads, self.output_dim // self.n_heads)
+
+        # compute attention weights
+        alpha = scatter_softmax((q[dst] * k / np.sqrt(k.shape[-1])).sum(-1), dst, dim=0,
+                                dim_size=N)  # [num_edges, n_heads]
+
+        # perform attention-weighted message-passing
+        m = alpha.unsqueeze(-1) * v  # (E, heads, H_per_head)
+        output = scatter_sum(m, dst, dim=0, dim_size=N)  # (N, heads, H_per_head)
+        output = output.view(-1, self.output_dim)
+        if self.out_fc:
+            output = self.node_output(torch.cat([output, h], -1))
+
+        output = output + h
+        return output
+
+
+class BaseH2XAttLayer(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim, n_heads, edge_feat_dim, r_feat_dim,
+                 act_fn='relu', norm=True, ew_net_type='r'):
+        super().__init__()
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.output_dim = output_dim
+        self.n_heads = n_heads
+        self.edge_feat_dim = edge_feat_dim
+        self.r_feat_dim = r_feat_dim
+        self.act_fn = act_fn
+        self.ew_net_type = ew_net_type
+
+        kv_input_dim = input_dim * 2 + edge_feat_dim + r_feat_dim
+
+        self.xk_func = MLP(kv_input_dim, output_dim, hidden_dim, norm=norm, act_fn=act_fn)
+        self.xv_func = MLP(kv_input_dim, self.n_heads, hidden_dim, norm=norm, act_fn=act_fn)
+        self.xq_func = MLP(input_dim, output_dim, hidden_dim, norm=norm, act_fn=act_fn)
+        if ew_net_type == 'r':
+            self.ew_net = nn.Sequential(nn.Linear(r_feat_dim, 1), nn.Sigmoid())
+
+    def forward(self, h, rel_x, r_feat, edge_feat, edge_index, e_w=None):
+        N = h.size(0)
+        src, dst = edge_index
+        hi, hj = h[dst], h[src]
+
+        # multi-head attention
+        # decide inputs of k_func and v_func
+        kv_input = torch.cat([r_feat, hi, hj], -1)
+        if edge_feat is not None:
+            kv_input = torch.cat([edge_feat, kv_input], -1)
+
+        k = self.xk_func(kv_input).view(-1, self.n_heads, self.output_dim // self.n_heads)
+        v = self.xv_func(kv_input)
+        if self.ew_net_type == 'r':
+            e_w = self.ew_net(r_feat)
+        elif self.ew_net_type == 'm':
+            e_w = 1.
+        elif e_w is not None:
+            e_w = e_w.view(-1, 1)
+        else:
+            e_w = 1.
+        v = v * e_w
+
+        v = v.unsqueeze(-1) * rel_x.unsqueeze(1)  # (xi - xj) [n_edges, n_heads, 3]
+        q = self.xq_func(h).view(-1, self.n_heads, self.output_dim // self.n_heads)
+
+        # Compute attention weights
+        alpha = scatter_softmax((q[dst] * k / np.sqrt(k.shape[-1])).sum(-1), dst, dim=0, dim_size=N)  # (E, heads)
+
+        # Perform attention-weighted message-passing
+        m = alpha.unsqueeze(-1) * v  # (E, heads, 3)
+        output = scatter_sum(m, dst, dim=0, dim_size=N)  # (N, heads, 3)
+        return output.mean(1)  # [num_nodes, 3]
+
+
+class AttentionLayerO2TwoUpdateNodeGeneral(nn.Module):
+    def __init__(self, hidden_dim, n_heads, num_r_gaussian, edge_feat_dim, act_fn='relu', norm=True,
+                 num_x2h=1, num_h2x=1, r_min=0., r_max=10., num_node_types=8,
+                 ew_net_type='r', x2h_out_fc=True, sync_twoup=False):
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.n_heads = n_heads
+        self.edge_feat_dim = edge_feat_dim
+        self.num_r_gaussian = num_r_gaussian
+        self.norm = norm
+        self.act_fn = act_fn
+        self.num_x2h = num_x2h
+        self.num_h2x = num_h2x
+        self.r_min, self.r_max = r_min, r_max
+        self.num_node_types = num_node_types
+        self.ew_net_type = ew_net_type
+        self.x2h_out_fc = x2h_out_fc
+        self.sync_twoup = sync_twoup
+
+        self.distance_expansion = GaussianSmearing(self.r_min, self.r_max, num_gaussians=num_r_gaussian)
+
+        self.x2h_layers = nn.ModuleList()
+        for i in range(self.num_x2h):
+            self.x2h_layers.append(
+                BaseX2HAttLayer(hidden_dim, hidden_dim, hidden_dim, n_heads, edge_feat_dim,
+                                r_feat_dim=num_r_gaussian * 4,
+                                act_fn=act_fn, norm=norm,
+                                ew_net_type=self.ew_net_type, out_fc=self.x2h_out_fc)
+            )
+        self.h2x_layers = nn.ModuleList()
+        for i in range(self.num_h2x):
+            self.h2x_layers.append(
+                BaseH2XAttLayer(hidden_dim, hidden_dim, hidden_dim, n_heads, edge_feat_dim,
+                                r_feat_dim=num_r_gaussian * 4,
+                                act_fn=act_fn, norm=norm,
+                                ew_net_type=self.ew_net_type)
+            )
+
+    def forward(self, h, x, edge_attr, edge_index, mask_ligand, e_w=None, fix_x=False):
+        src, dst = edge_index
+        if self.edge_feat_dim > 0:
+            edge_feat = edge_attr  # shape: [#edges_in_batch, #bond_types]
+        else:
+            edge_feat = None
+
+        rel_x = x[dst] - x[src]
+        dist = torch.norm(rel_x, p=2, dim=-1, keepdim=True)
+
+        h_in = h
+        # 4 separate distance embedding for p-p, p-l, l-p, l-l
+        for i in range(self.num_x2h):
+            dist_feat = self.distance_expansion(dist)
+            dist_feat = outer_product(edge_attr, dist_feat)
+            h_out = self.x2h_layers[i](h_in, dist_feat, edge_feat, edge_index, e_w=e_w)
+            h_in = h_out
+        x2h_out = h_in
+
+        new_h = h if self.sync_twoup else x2h_out
+        for i in range(self.num_h2x):
+            dist_feat = self.distance_expansion(dist)
+            dist_feat = outer_product(edge_attr, dist_feat)
+            delta_x = self.h2x_layers[i](new_h, rel_x, dist_feat, edge_feat, edge_index, e_w=e_w)
+            if not fix_x:
+                x = x + delta_x * mask_ligand[:, None]  # only ligand positions will be updated
+            rel_x = x[dst] - x[src]
+            dist = torch.norm(rel_x, p=2, dim=-1, keepdim=True)
+
+        return x2h_out, x
+
+
+class UniTransformerO2TwoUpdateGeneral(nn.Module):
+    def __init__(self, num_blocks, num_layers, hidden_dim, n_heads=1, k=32,
+                 num_r_gaussian=50, edge_feat_dim=0, num_node_types=8, act_fn='relu', norm=True,
+                 cutoff_mode='radius', ew_net_type='r',
+                 num_init_x2h=1, num_init_h2x=0, num_x2h=1, num_h2x=1, r_max=10., x2h_out_fc=True, sync_twoup=False):
+        super().__init__()
+        # Build the network
+        self.num_blocks = num_blocks
+        self.num_layers = num_layers
+        self.hidden_dim = hidden_dim
+        self.n_heads = n_heads
+        self.num_r_gaussian = num_r_gaussian
+        self.edge_feat_dim = edge_feat_dim
+        self.act_fn = act_fn
+        self.norm = norm
+        self.num_node_types = num_node_types
+        # radius graph / knn graph
+        self.cutoff_mode = cutoff_mode  # [radius, none]
+        self.k = k
+        self.ew_net_type = ew_net_type  # [r, m, none]
+
+        self.num_x2h = num_x2h
+        self.num_h2x = num_h2x
+        self.num_init_x2h = num_init_x2h
+        self.num_init_h2x = num_init_h2x
+        self.r_max = r_max
+        self.x2h_out_fc = x2h_out_fc
+        self.sync_twoup = sync_twoup
+        self.distance_expansion = GaussianSmearing(0., r_max, num_gaussians=num_r_gaussian)
+        if self.ew_net_type == 'global':
+            self.edge_pred_layer = MLP(num_r_gaussian, 1, hidden_dim)
+
+        self.init_h_emb_layer = self._build_init_h_layer()
+        self.base_block = self._build_share_blocks()
+
+    def __repr__(self):
+        return f'UniTransformerO2(num_blocks={self.num_blocks}, num_layers={self.num_layers}, n_heads={self.n_heads}, ' \
+               f'act_fn={self.act_fn}, norm={self.norm}, cutoff_mode={self.cutoff_mode}, ew_net_type={self.ew_net_type}, ' \
+               f'init h emb: {self.init_h_emb_layer.__repr__()} \n' \
+               f'base block: {self.base_block.__repr__()} \n' \
+               f'edge pred layer: {self.edge_pred_layer.__repr__() if hasattr(self, "edge_pred_layer") else "None"}) '
+
+    def _build_init_h_layer(self):
+        layer = AttentionLayerO2TwoUpdateNodeGeneral(
+            self.hidden_dim, self.n_heads, self.num_r_gaussian, self.edge_feat_dim, act_fn=self.act_fn, norm=self.norm,
+            num_x2h=self.num_init_x2h, num_h2x=self.num_init_h2x, r_max=self.r_max, num_node_types=self.num_node_types,
+            ew_net_type=self.ew_net_type, x2h_out_fc=self.x2h_out_fc, sync_twoup=self.sync_twoup,
+        )
+        return layer
+
+    def _build_share_blocks(self):
+        # Equivariant layers
+        base_block = []
+        for l_idx in range(self.num_layers):
+            layer = AttentionLayerO2TwoUpdateNodeGeneral(
+                self.hidden_dim, self.n_heads, self.num_r_gaussian, self.edge_feat_dim, act_fn=self.act_fn,
+                norm=self.norm,
+                num_x2h=self.num_x2h, num_h2x=self.num_h2x, r_max=self.r_max, num_node_types=self.num_node_types,
+                ew_net_type=self.ew_net_type, x2h_out_fc=self.x2h_out_fc, sync_twoup=self.sync_twoup,
+            )
+            base_block.append(layer)
+        return nn.ModuleList(base_block)
+
+    def _connect_edge(self, x, mask_ligand, batch):
+        if self.cutoff_mode == 'radius':
+            edge_index = radius_graph(x, r=self.r, batch=batch, flow='source_to_target')
+        elif self.cutoff_mode == 'knn':
+            edge_index = knn_graph(x, k=self.k, batch=batch, flow='source_to_target')
+        elif self.cutoff_mode == 'hybrid':
+            edge_index = batch_hybrid_edge_connection(
+                x, k=self.k, mask_ligand=mask_ligand, batch=batch, add_p_index=True)
+        else:
+            raise ValueError(f'Not supported cutoff mode: {self.cutoff_mode}')
+        return edge_index
+
+    @staticmethod
+    def _build_edge_type(edge_index, mask_ligand):
+        src, dst = edge_index
+        edge_type = torch.zeros(len(src)).to(edge_index)
+        n_src = mask_ligand[src] == 1
+        n_dst = mask_ligand[dst] == 1
+        edge_type[n_src & n_dst] = 0
+        edge_type[n_src & ~n_dst] = 1
+        edge_type[~n_src & n_dst] = 2
+        edge_type[~n_src & ~n_dst] = 3
+        edge_type = F.one_hot(edge_type, num_classes=4)
+        return edge_type
+
+    def forward(self, h, x, mask_ligand, batch, return_all=False, fix_x=False):
+
+        all_x = [x]
+        all_h = [h]
+
+        for b_idx in range(self.num_blocks):
+            edge_index = self._connect_edge(x, mask_ligand, batch)
+            src, dst = edge_index
+
+            # edge type (dim: 4)
+            edge_type = self._build_edge_type(edge_index, mask_ligand)
+            if self.ew_net_type == 'global':
+                dist = torch.norm(x[dst] - x[src], p=2, dim=-1, keepdim=True)
+                dist_feat = self.distance_expansion(dist)
+                logits = self.edge_pred_layer(dist_feat)
+                e_w = torch.sigmoid(logits)
+            else:
+                e_w = None
+
+            for l_idx, layer in enumerate(self.base_block):
+                h, x = layer(h, x, edge_type, edge_index, mask_ligand, e_w=e_w, fix_x=fix_x)
+            all_x.append(x)
+            all_h.append(h)
+
+        outputs = {'x': x, 'h': h}
+        if return_all:
+            outputs.update({'all_x': all_x, 'all_h': all_h})
+        return outputs

scripts/__init__.py

@@ -0,0 +1,2 @@
+"""TargetDiff example scripts."""
+

scripts/batch_sample_diffusion.sh

@@ -0,0 +1,21 @@
+TOTAL_TASKS=100
+BATCH_SIZE=50
+
+if [ $# != 5 ]; then
+    echo "Error: 5 arguments required."
+    exit 1
+fi
+
+CONFIG_FILE=$1
+RESULT_PATH=$2
+NODE_ALL=$3
+NODE_THIS=$4
+START_IDX=$5
+
+for ((i=$START_IDX;i<$TOTAL_TASKS;i++)); do
+    NODE_TARGET=$(($i % $NODE_ALL))
+    if [ $NODE_TARGET == $NODE_THIS ]; then
+        echo "Task ${i} assigned to this worker (${NODE_THIS})"
+        python -m scripts.sample_diffusion ${CONFIG_FILE} -i ${i} --batch_size ${BATCH_SIZE} --result_path ${RESULT_PATH}
+    fi
+done

scripts/data_preparation/__init__.py

@@ -0,0 +1,2 @@
+"""TargetDiff data preparation scripts."""
+

scripts/data_preparation/clean_crossdocked.py

@@ -0,0 +1,59 @@
+import os
+import shutil
+import gzip
+import pickle
+import argparse
+from tqdm.auto import tqdm
+
+TYPES_FILENAME = 'types/it2_tt_v1.1_completeset_train0.types'  # 'types/it2_tt_completeset_train0.types'
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--source', type=str, default='./data/CrossDocked2020')
+    parser.add_argument('--dest', type=str, required=True)
+    parser.add_argument('--rmsd_thr', type=float, default=1.0)
+    args = parser.parse_args()
+
+    os.makedirs(args.dest, exist_ok=False)
+    types_path = os.path.join(args.source, TYPES_FILENAME)
+
+    index = []
+    with open(types_path, 'r') as f:
+        for ln in tqdm(f.readlines()):
+            _, _, rmsd, protein_fn, ligand_fn, _ = ln.split()
+            rmsd = float(rmsd)
+            if rmsd > args.rmsd_thr: 
+                continue
+
+            ligand_id = int(ligand_fn[ligand_fn.rfind('_')+1:ligand_fn.rfind('.')])
+
+            protein_fn = protein_fn[:protein_fn.rfind('_')] + '.pdb'
+            # For CrossDocked v1.0
+            # ligand_raw_fn = ligand_fn[:ligand_fn.rfind('_')] + '.sdf'
+            ligand_raw_fn = ligand_fn[:ligand_fn.rfind('_')] + '.sdf.gz'
+            protein_path = os.path.join(args.source, protein_fn)
+            ligand_raw_path = os.path.join(args.source, ligand_raw_fn)
+            if not (os.path.exists(protein_path) and os.path.exists(ligand_raw_path)):
+                continue
+
+            # For CrossDocked v1.0
+            # with open(ligand_raw_path, 'r') as f:
+            with gzip.open(ligand_raw_path, 'rt') as f:
+                ligand_sdf = f.read().split('$$$$\n')[ligand_id]
+            ligand_save_fn = ligand_fn[:ligand_fn.rfind('.')] + '.sdf'  # include ligand id
+
+            protein_dest = os.path.join(args.dest, protein_fn)
+            ligand_dest = os.path.join(args.dest, ligand_save_fn)
+            os.makedirs(os.path.dirname(protein_dest), exist_ok=True)
+            os.makedirs(os.path.dirname(ligand_dest), exist_ok=True)
+            shutil.copyfile(protein_path, protein_dest)
+            with open(ligand_dest, 'w') as f:
+                f.write(ligand_sdf)
+
+            index.append((protein_fn, ligand_save_fn, rmsd))
+
+    index_path = os.path.join(args.dest, 'index.pkl')
+    with open(index_path, 'wb') as f:
+        pickle.dump(index, f)
+
+    print(f'Done processing {len(index)} protein-ligand pairs in total.\n Processed files in {args.dest}.')

scripts/data_preparation/extract_pockets.py

@@ -0,0 +1,76 @@
+import os
+import argparse
+import multiprocessing as mp
+import pickle
+import shutil
+from functools import partial
+
+from tqdm.auto import tqdm
+
+from onescience.utils.targetdiff.data import PDBProtein, parse_sdf_file
+
+
+def load_item(item, path):
+    pdb_path = os.path.join(path, item[0])
+    sdf_path = os.path.join(path, item[1])
+    with open(pdb_path, 'r') as f:
+        pdb_block = f.read()
+    with open(sdf_path, 'r') as f:
+        sdf_block = f.read()
+    return pdb_block, sdf_block
+
+
+def process_item(item, args):
+    try:
+        pdb_block, sdf_block = load_item(item, args.source)
+        protein = PDBProtein(pdb_block)
+        # ligand = parse_sdf_block(sdf_block)
+        ligand = parse_sdf_file(os.path.join(args.source, item[1]))
+
+        pdb_block_pocket = protein.residues_to_pdb_block(
+            protein.query_residues_ligand(ligand, args.radius)
+        )
+        
+        ligand_fn = item[1]
+        pocket_fn = ligand_fn[:-4] + '_pocket%d.pdb' % args.radius
+        ligand_dest = os.path.join(args.dest, ligand_fn)
+        pocket_dest = os.path.join(args.dest, pocket_fn)
+        os.makedirs(os.path.dirname(ligand_dest), exist_ok=True)
+
+        shutil.copyfile(
+            src=os.path.join(args.source, ligand_fn),
+            dst=os.path.join(args.dest, ligand_fn)
+        )
+        with open(pocket_dest, 'w') as f:
+            f.write(pdb_block_pocket)
+        return pocket_fn, ligand_fn, item[0], item[2]  # item[0]: original protein filename; item[2]: rmsd.
+    except Exception:
+        print('Exception occurred.', item)
+        return None, item[1], item[0], item[2]
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--source', type=str, default='./data/crossdocked_subset')
+    parser.add_argument('--dest', type=str, required=True)
+    parser.add_argument('--radius', type=int, default=10)
+    parser.add_argument('--num_workers', type=int, default=16)
+    args = parser.parse_args()
+
+    os.makedirs(args.dest, exist_ok=False)
+    with open(os.path.join(args.source, 'index.pkl'), 'rb') as f:
+        index = pickle.load(f)
+
+    pool = mp.Pool(args.num_workers)
+    index_pocket = []
+    for item_pocket in tqdm(pool.imap_unordered(partial(process_item, args=args), index), total=len(index)):
+        index_pocket.append(item_pocket)
+    # index_pocket = pool.map(partial(process_item, args=args), index)
+    pool.close()
+
+    index_path = os.path.join(args.dest, 'index.pkl')
+    with open(index_path, 'wb') as f:
+        pickle.dump(index_pocket, f)
+
+    print('Done. %d protein-ligand pairs in total.' % len(index))
+    

scripts/data_preparation/split_pl_dataset.py

@@ -0,0 +1,112 @@
+import os
+import argparse
+import random
+import torch
+from tqdm.auto import tqdm
+
+from torch.utils.data import Subset
+from onescience.datapipes.targetdiff.pl_pair_dataset import PocketLigandPairDataset
+
+
+def get_chain_name(fn):
+    return os.path.basename(fn)[:6]
+
+
+def get_pdb_name(fn):
+    return os.path.basename(fn)[:4]
+
+
+def get_unique_pockets(dataset, raw_id, used_pdb, num_pockets):
+    # only save first encountered id for unseen pdbs
+    unique_id = []
+    pdb_visited = set()
+    for idx in tqdm(raw_id, 'Filter'):
+        pdb_name = get_pdb_name(dataset[idx].ligand_filename)
+        if pdb_name not in used_pdb and pdb_name not in pdb_visited:
+            unique_id.append(idx)
+            pdb_visited.add(pdb_name)
+
+    print('Number of Pairs: %d' % len(unique_id))
+    print('Number of PDBs:  %d' % len(pdb_visited))
+
+    random.Random(args.seed).shuffle(unique_id)
+    unique_id = unique_id[:num_pockets]
+    print('Number of selected: %d' % len(unique_id))
+    return unique_id, pdb_visited.union(used_pdb)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--path', type=str, default='./data/crossdocked_v1.1_rmsd1.0_pocket10')
+    parser.add_argument('--dest', type=str, default='./data/crossdocked_pocket10_pose_split.pt')
+    parser.add_argument('--fixed_split', type=str, default='./data/split_by_name.pt')
+    parser.add_argument('--train', type=int, default=100000)
+    parser.add_argument('--val', type=int, default=1000)
+    parser.add_argument('--test', type=int, default=20000)
+    parser.add_argument('--val_num_pockets', type=int, default=-1)
+    parser.add_argument('--test_num_pockets', type=int, default=100)
+    parser.add_argument('--seed', type=int, default=2021)
+    args = parser.parse_args()
+
+    dataset = PocketLigandPairDataset(args.path)
+    print('Load dataset successfully!')
+    if args.fixed_split:
+        fixed_split = torch.load(args.fixed_split)
+        print('Load fixed split successfully!')
+        name_id_dict = {}
+        for idx, data in enumerate(tqdm(dataset, desc='Indexing')):
+            name_id_dict[data.protein_filename + data.ligand_filename] = idx
+
+        selected_ids = {'train': [], 'test': []}
+        for split in ['train', 'test']:
+            print(f'Selecting {split} split...')
+            for fn in fixed_split[split]:
+                if (fn[0] + fn[1]) in name_id_dict:
+                    selected_ids[split].append(name_id_dict[fn[0] + fn[1]])
+                else:
+                    print(f'Warning: data with PDB fn {fn[0]} and ligand fn {fn[1]} not found!')
+        train_id, val_id, test_id = selected_ids['train'], [], selected_ids['test']
+
+    else:
+        allowed_elements = {1, 6, 7, 8, 9, 15, 16, 17}
+        elements = {i: set() for i in range(90)}
+        for i, data in enumerate(tqdm(dataset, desc='Filter')):
+            for e in data.ligand_element:
+                elements[e.item()].add(i)
+
+        all_id = set(range(len(dataset)))
+        blocked_id = set().union(*[
+            elements[i] for i in elements.keys() if i not in allowed_elements
+        ])
+
+        allowed_id = list(all_id - blocked_id)
+        random.Random(args.seed).shuffle(allowed_id)
+        print('Allowed: %d' % len(allowed_id))
+
+        train_id = allowed_id[:args.train]
+        train_set = Subset(dataset, indices=train_id)
+        train_pdb = {get_pdb_name(d.ligand_filename) for d in tqdm(train_set)}
+        print('train pdb: ', train_pdb)
+
+        if args.val_num_pockets == -1:
+            # not group by pocket
+            val_id = allowed_id[args.train: args.train + args.val]
+            used_pdb = train_pdb
+        else:
+            raw_val_id = allowed_id[args.train: args.train + args.val]
+            val_id, used_pdb = get_unique_pockets(dataset, raw_val_id, train_pdb, args.val_num_pockets)
+
+        if args.test_num_pockets == -1:
+            test_id = allowed_id[args.train + args.val: args.train + args.val + args.test]
+        else:
+            raw_test_id = allowed_id[args.train + args.val: args.train + args.val + args.test]
+            test_id, used_pdb = get_unique_pockets(dataset, raw_test_id, used_pdb, args.test_num_pockets)
+
+    torch.save({
+        'train': train_id,
+        'val': val_id,
+        'test': test_id,
+    }, args.dest)
+
+    print('Train %d, Validation %d, Test %d.' % (len(train_id), len(val_id), len(test_id)))
+    print('Done.')

scripts/dock_baseline.py

@@ -0,0 +1,59 @@
+import argparse
+import os
+import torch
+from tqdm.auto import tqdm
+from onescience.utils.targetdiff.evaluation.docking_qvina import QVinaDockingTask
+from onescience.utils.targetdiff.evaluation.docking_vina import VinaDockingTask
+import multiprocessing as mp
+
+
+def dock_pocket_samples(pocket_samples):
+    ligand_fn = pocket_samples[0]['ligand_filename']
+    print('Start docking pocket: %s' % ligand_fn)
+    pocket_results = []
+    for idx, s in enumerate(tqdm(pocket_samples, desc='docking %d' % os.getpid())):
+        try:
+            if args.docking_mode == 'qvina':
+                vina_task = QVinaDockingTask.from_generated_mol(
+                    s['mol'], s['ligand_filename'], protein_root=args.protein_root, size_factor=args.dock_size_factor)
+                vina_results = vina_task.run_sync()
+            elif args.docking_mode == 'vina_score':
+                vina_task = VinaDockingTask.from_generated_mol(
+                    s['mol'], s['ligand_filename'], protein_root=args.protein_root)
+                score_only_results = vina_task.run(mode='score_only', exhaustiveness=args.exhaustiveness)
+                minimize_results = vina_task.run(mode='minimize', exhaustiveness=args.exhaustiveness)
+                vina_results = {
+                    'score_only': score_only_results,
+                    'minimize': minimize_results
+                }
+            else:
+                raise ValueError
+        except:
+            print('Error at %d of %s' % (idx, ligand_fn))
+            vina_results = None
+        pocket_results.append({**s, 'vina': vina_results})
+    return pocket_results
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('sample_path', type=str)
+    parser.add_argument('-o', '--out', type=str, default=None)
+    parser.add_argument('-n', '--num_processes', type=int, default=10)
+    parser.add_argument('--protein_root', type=str, default='./data/crossdocked_v1.1_rmsd1.0')
+    parser.add_argument('--dock_size_factor', type=float, default=None)
+    parser.add_argument('--exhaustiveness', type=int, default=16)
+    parser.add_argument('--docking_mode', type=str, default='vina_score',
+                        choices=['none', 'qvina', 'vina_score'])
+    args = parser.parse_args()
+
+    samples = torch.load(args.sample_path)
+    with mp.Pool(args.num_processes) as p:
+        docked_samples = p.map(dock_pocket_samples, samples)
+    if args.out is None:
+        dir_name = os.path.dirname(args.sample_path)
+        baseline_name = os.path.basename(args.sample_path).split('_')[0]
+        out_path = os.path.join(dir_name, baseline_name + '_test_docked.pt')
+    else:
+        out_path = args.out
+    torch.save(docked_samples, out_path)

scripts/dock_testset.py

@@ -0,0 +1,76 @@
+import argparse
+import os
+
+from rdkit import Chem
+import torch
+from tqdm.auto import tqdm
+from onescience.utils.targetdiff import misc
+from onescience.utils.targetdiff.evaluation import scoring_func
+from onescience.utils.targetdiff.evaluation.docking_qvina import QVinaDockingTask
+from onescience.datapipes.targetdiff import get_dataset
+from easydict import EasyDict
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-d', '--dataset', type=str, default='./data/crossdocked_v1.1_rmsd1.0_pocket10')
+    parser.add_argument('-s', '--split', type=str, default='./data/crossdocked_pocket10_pose_split.pt')
+    parser.add_argument('-o', '--out', type=str, default=None)
+    parser.add_argument('--protein_root', type=str, default='./data/crossdocked_v1.1_rmsd1.0')
+    parser.add_argument('--ligand_root', type=str, default='./data/crossdocked_v1.1_rmsd1.0_pocket10')
+    parser.add_argument('--use_uff', type=eval, default=True)
+    parser.add_argument('--size_factor', type=float, default=1.2)
+    args = parser.parse_args()
+
+    logger = misc.get_logger('docking')
+    logger.info(args)
+
+    # Load dataset
+    dataset, subsets = get_dataset(
+        config=EasyDict({
+            'name': 'pl',
+            'path': args.dataset,
+            'split': args.split
+        })
+    )
+    train_set, test_set = subsets['train'], subsets['test']
+    logger.info(f'Successfully load the dataset (size: {len(test_set)})!')
+
+    # Dock
+    logger.info('Start docking...')
+    results = []
+    for i, data in enumerate(tqdm(test_set)):
+        mol = next(iter(Chem.SDMolSupplier(os.path.join(args.ligand_root, data.ligand_filename))))
+        # try:
+        chem_results = scoring_func.get_chem(mol)
+        vina_task = QVinaDockingTask.from_original_data(
+            data,
+            ligand_root=args.ligand_root,
+            protein_root=args.protein_root,
+            use_uff=args.use_uff,
+            size_factor=args.size_factor
+        )
+        vina_results = vina_task.run_sync()
+        # except:
+        #     logger.warning('Error #%d' % i)
+        #     continue
+
+        results.append({
+            'mol': mol,
+            'smiles': data.ligand_smiles,
+            'ligand_filename': data.ligand_filename,
+            'chem_results': chem_results,
+            'vina': vina_results
+        })
+
+    # Save
+    if args.out is None:
+        split_name = os.path.basename(args.split)
+        split_name = split_name[:split_name.rfind('.')]
+        docked_name = f'{split_name}_test_docked_uff_{args.use_uff}_size_{args.size_factor}.pt'
+        out_path = os.path.join(os.path.dirname(args.dataset), docked_name)
+    else:
+        out_path = args.out
+    logger.info('Num docked: %d' % len(results))
+    logger.info('Saving results to %s' % out_path)
+    torch.save(results, out_path)

scripts/evaluate_diffusion.py

@@ -0,0 +1,208 @@
+import argparse
+import os
+
+import numpy as np
+from rdkit import Chem
+from rdkit import RDLogger
+import torch
+from tqdm.auto import tqdm
+from glob import glob
+from collections import Counter
+
+from onescience.utils.targetdiff.evaluation import eval_atom_type, scoring_func, analyze, eval_bond_length
+from onescience.utils.targetdiff import misc, reconstruct, transforms
+from onescience.utils.targetdiff.evaluation.docking_qvina import QVinaDockingTask
+from onescience.utils.targetdiff.evaluation.docking_vina import VinaDockingTask
+
+
+def print_dict(d, logger):
+    for k, v in d.items():
+        if v is not None:
+            logger.info(f'{k}:\t{v:.4f}')
+        else:
+            logger.info(f'{k}:\tNone')
+
+
+def print_ring_ratio(all_ring_sizes, logger):
+    for ring_size in range(3, 10):
+        n_mol = 0
+        for counter in all_ring_sizes:
+            if ring_size in counter:
+                n_mol += 1
+        logger.info(f'ring size: {ring_size} ratio: {n_mol / len(all_ring_sizes):.3f}')
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('sample_path', type=str)
+    parser.add_argument('--verbose', type=eval, default=False)
+    parser.add_argument('--eval_step', type=int, default=-1)
+    parser.add_argument('--eval_num_examples', type=int, default=None)
+    parser.add_argument('--save', type=eval, default=True)
+    parser.add_argument('--protein_root', type=str, default='./data/crossdocked_v1.1_rmsd1.0')
+    parser.add_argument('--atom_enc_mode', type=str, default='add_aromatic')
+    parser.add_argument('--docking_mode', type=str, choices=['qvina', 'vina_score', 'vina_dock', 'none'])
+    parser.add_argument('--exhaustiveness', type=int, default=16)
+    args = parser.parse_args()
+
+    result_path = os.path.join(args.sample_path, 'eval_results')
+    os.makedirs(result_path, exist_ok=True)
+    logger = misc.get_logger('evaluate', log_dir=result_path)
+    if not args.verbose:
+        RDLogger.DisableLog('rdApp.*')
+
+    # Load generated data
+    results_fn_list = glob(os.path.join(args.sample_path, '*result_*.pt'))
+    results_fn_list = sorted(results_fn_list, key=lambda x: int(os.path.basename(x)[:-3].split('_')[-1]))
+    if args.eval_num_examples is not None:
+        results_fn_list = results_fn_list[:args.eval_num_examples]
+    num_examples = len(results_fn_list)
+    logger.info(f'Load generated data done! {num_examples} examples in total.')
+
+    num_samples = 0
+    all_mol_stable, all_atom_stable, all_n_atom = 0, 0, 0
+    n_recon_success, n_eval_success, n_complete = 0, 0, 0
+    results = []
+    all_pair_dist, all_bond_dist = [], []
+    all_atom_types = Counter()
+    success_pair_dist, success_atom_types = [], Counter()
+    for example_idx, r_name in enumerate(tqdm(results_fn_list, desc='Eval')):
+        r = torch.load(r_name)  # ['data', 'pred_ligand_pos', 'pred_ligand_v', 'pred_ligand_pos_traj', 'pred_ligand_v_traj']
+        all_pred_ligand_pos = r['pred_ligand_pos_traj']  # [num_samples, num_steps, num_atoms, 3]
+        all_pred_ligand_v = r['pred_ligand_v_traj']
+        num_samples += len(all_pred_ligand_pos)
+
+        for sample_idx, (pred_pos, pred_v) in enumerate(zip(all_pred_ligand_pos, all_pred_ligand_v)):
+            pred_pos, pred_v = pred_pos[args.eval_step], pred_v[args.eval_step]
+
+            # stability check
+            pred_atom_type = transforms.get_atomic_number_from_index(pred_v, mode=args.atom_enc_mode)
+            all_atom_types += Counter(pred_atom_type)
+            r_stable = analyze.check_stability(pred_pos, pred_atom_type)
+            all_mol_stable += r_stable[0]
+            all_atom_stable += r_stable[1]
+            all_n_atom += r_stable[2]
+
+            pair_dist = eval_bond_length.pair_distance_from_pos_v(pred_pos, pred_atom_type)
+            all_pair_dist += pair_dist
+
+            # reconstruction
+            try:
+                pred_aromatic = transforms.is_aromatic_from_index(pred_v, mode=args.atom_enc_mode)
+                mol = reconstruct.reconstruct_from_generated(pred_pos, pred_atom_type, pred_aromatic)
+                smiles = Chem.MolToSmiles(mol)
+            except reconstruct.MolReconsError:
+                if args.verbose:
+                    logger.warning('Reconstruct failed %s' % f'{example_idx}_{sample_idx}')
+                continue
+            n_recon_success += 1
+
+            if '.' in smiles:
+                continue
+            n_complete += 1
+
+            # chemical and docking check
+            try:
+                chem_results = scoring_func.get_chem(mol)
+                if args.docking_mode == 'qvina':
+                    vina_task = QVinaDockingTask.from_generated_mol(
+                        mol, r['data'].ligand_filename, protein_root=args.protein_root)
+                    vina_results = vina_task.run_sync()
+                elif args.docking_mode in ['vina_score', 'vina_dock']:
+                    vina_task = VinaDockingTask.from_generated_mol(
+                        mol, r['data'].ligand_filename, protein_root=args.protein_root)
+                    score_only_results = vina_task.run(mode='score_only', exhaustiveness=args.exhaustiveness)
+                    minimize_results = vina_task.run(mode='minimize', exhaustiveness=args.exhaustiveness)
+                    vina_results = {
+                        'score_only': score_only_results,
+                        'minimize': minimize_results
+                    }
+                    if args.docking_mode == 'vina_dock':
+                        docking_results = vina_task.run(mode='dock', exhaustiveness=args.exhaustiveness)
+                        vina_results['dock'] = docking_results
+                else:
+                    vina_results = None
+
+                n_eval_success += 1
+            except:
+                if args.verbose:
+                    logger.warning('Evaluation failed for %s' % f'{example_idx}_{sample_idx}')
+                continue
+
+            # now we only consider complete molecules as success
+            bond_dist = eval_bond_length.bond_distance_from_mol(mol)
+            all_bond_dist += bond_dist
+
+            success_pair_dist += pair_dist
+            success_atom_types += Counter(pred_atom_type)
+
+            results.append({
+                'mol': mol,
+                'smiles': smiles,
+                'ligand_filename': r['data'].ligand_filename,
+                'pred_pos': pred_pos,
+                'pred_v': pred_v,
+                'chem_results': chem_results,
+                'vina': vina_results
+            })
+    logger.info(f'Evaluate done! {num_samples} samples in total.')
+
+    fraction_mol_stable = all_mol_stable / num_samples
+    fraction_atm_stable = all_atom_stable / all_n_atom
+    fraction_recon = n_recon_success / num_samples
+    fraction_eval = n_eval_success / num_samples
+    fraction_complete = n_complete / num_samples
+    validity_dict = {
+        'mol_stable': fraction_mol_stable,
+        'atm_stable': fraction_atm_stable,
+        'recon_success': fraction_recon,
+        'eval_success': fraction_eval,
+        'complete': fraction_complete
+    }
+    print_dict(validity_dict, logger)
+
+    c_bond_length_profile = eval_bond_length.get_bond_length_profile(all_bond_dist)
+    c_bond_length_dict = eval_bond_length.eval_bond_length_profile(c_bond_length_profile)
+    logger.info('JS bond distances of complete mols: ')
+    print_dict(c_bond_length_dict, logger)
+
+    success_pair_length_profile = eval_bond_length.get_pair_length_profile(success_pair_dist)
+    success_js_metrics = eval_bond_length.eval_pair_length_profile(success_pair_length_profile)
+    print_dict(success_js_metrics, logger)
+
+    atom_type_js = eval_atom_type.eval_atom_type_distribution(success_atom_types)
+    logger.info('Atom type JS: %.4f' % atom_type_js)
+
+    if args.save:
+        eval_bond_length.plot_distance_hist(success_pair_length_profile,
+                                            metrics=success_js_metrics,
+                                            save_path=os.path.join(result_path, f'pair_dist_hist_{args.eval_step}.png'))
+
+    logger.info('Number of reconstructed mols: %d, complete mols: %d, evaluated mols: %d' % (
+        n_recon_success, n_complete, len(results)))
+
+    qed = [r['chem_results']['qed'] for r in results]
+    sa = [r['chem_results']['sa'] for r in results]
+    logger.info('QED:   Mean: %.3f Median: %.3f' % (np.mean(qed), np.median(qed)))
+    logger.info('SA:    Mean: %.3f Median: %.3f' % (np.mean(sa), np.median(sa)))
+    if args.docking_mode == 'qvina':
+        vina = [r['vina'][0]['affinity'] for r in results]
+        logger.info('Vina:  Mean: %.3f Median: %.3f' % (np.mean(vina), np.median(vina)))
+    elif args.docking_mode in ['vina_dock', 'vina_score']:
+        vina_score_only = [r['vina']['score_only'][0]['affinity'] for r in results]
+        vina_min = [r['vina']['minimize'][0]['affinity'] for r in results]
+        logger.info('Vina Score:  Mean: %.3f Median: %.3f' % (np.mean(vina_score_only), np.median(vina_score_only)))
+        logger.info('Vina Min  :  Mean: %.3f Median: %.3f' % (np.mean(vina_min), np.median(vina_min)))
+        if args.docking_mode == 'vina_dock':
+            vina_dock = [r['vina']['dock'][0]['affinity'] for r in results]
+            logger.info('Vina Dock :  Mean: %.3f Median: %.3f' % (np.mean(vina_dock), np.median(vina_dock)))
+
+    # check ring distribution
+    print_ring_ratio([r['chem_results']['ring_size'] for r in results], logger)
+
+    if args.save:
+        torch.save({
+            'stability': validity_dict,
+            'bond_length': all_bond_dist,
+            'all_results': results
+        }, os.path.join(result_path, f'metrics_{args.eval_step}.pt'))

scripts/evaluate_from_meta.py

@@ -0,0 +1,138 @@
+import argparse
+import os
+
+import numpy as np
+from rdkit import RDLogger
+import torch
+from tqdm.auto import tqdm
+from copy import deepcopy
+
+from onescience.utils.targetdiff import misc
+from onescience.utils.targetdiff.evaluation import scoring_func
+from onescience.utils.targetdiff.evaluation.docking_qvina import QVinaDockingTask
+from onescience.utils.targetdiff.evaluation.docking_vina import VinaDockingTask
+from multiprocessing import Pool
+from functools import partial
+from glob import glob
+
+
+def eval_single_datapoint(index, id, args):
+    if isinstance(index, dict):
+        # reference set
+        index = [index]
+
+    ligand_filename = index[0]['ligand_filename']
+    num_samples = len(index[:100])
+    results = []
+    n_eval_success = 0
+    for sample_idx, sample_dict in enumerate(tqdm(index[:num_samples], desc='Eval', total=num_samples)):
+        mol = sample_dict['mol']
+        smiles = sample_dict['smiles']
+        if '.' in smiles:
+            continue
+        # chemical and docking check
+        try:
+            chem_results = scoring_func.get_chem(mol)
+            if 'vina' in sample_dict:
+                vina_results = sample_dict['vina']
+            else:
+                if args.docking_mode == 'qvina':
+                    vina_task = QVinaDockingTask.from_generated_mol(mol, ligand_filename, protein_root=args.protein_root)
+                    vina_results = vina_task.run_sync()
+                elif args.docking_mode == 'vina':
+                    vina_task = VinaDockingTask.from_generated_mol(mol, ligand_filename, protein_root=args.protein_root)
+                    vina_results = vina_task.run(mode='dock')
+                elif args.docking_mode in ['vina_full', 'vina_score']:
+                    vina_task = VinaDockingTask.from_generated_mol(deepcopy(mol),
+                                                                   ligand_filename, protein_root=args.protein_root)
+                    score_only_results = vina_task.run(mode='score_only', exhaustiveness=args.exhaustiveness)
+                    minimize_results = vina_task.run(mode='minimize', exhaustiveness=args.exhaustiveness)
+                    vina_results = {
+                        'score_only': score_only_results,
+                        'minimize': minimize_results
+                    }
+                    if args.docking_mode == 'vina_full':
+                        dock_results = vina_task.run(mode='dock', exhaustiveness=args.exhaustiveness)
+                        vina_results.update({
+                            'dock': dock_results,
+                        })
+
+                elif args.docking_mode == 'none':
+                    vina_results = None
+                else:
+                    raise NotImplementedError
+            n_eval_success += 1
+        except Exception as e:
+            logger.warning('Evaluation failed for %s' % f'{sample_idx}')
+            print(str(e))
+            continue
+
+        results.append({
+            **sample_dict,
+            'chem_results': chem_results,
+            'vina': vina_results
+        })
+    logger.info(f'Evaluate No {id} done! {num_samples} samples in total. {n_eval_success} eval success!')
+    torch.save(results, os.path.join(args.result_path, f'eval_{id:03d}_{os.path.basename(ligand_filename[:-4])}.pt'))
+    return results
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('meta_file', type=str)  # 'sampling_results/targetdiff_vina_docked.pt'
+    parser.add_argument('-n', '--eval_num_examples', type=int, default=100)
+    parser.add_argument('--verbose', type=eval, default=False)
+    parser.add_argument('--protein_root', type=str, default='./data/crossdocked_v1.1_rmsd1.0')
+    parser.add_argument('--docking_mode', type=str, default='vina_full',
+                        choices=['none', 'qvina', 'vina', 'vina_full', 'vina_score'])
+    parser.add_argument('--exhaustiveness', type=int, default=32)
+    parser.add_argument('--num_workers', type=int, default=4)
+    parser.add_argument('--result_path', type=str, required=True)
+    parser.add_argument('--aggregate_meta', type=eval, default=False)
+    args = parser.parse_args()
+
+    if args.result_path:
+        os.makedirs(args.result_path, exist_ok=True)
+    logger = misc.get_logger('evaluate', args.result_path)
+    logger.info(args)
+    if not args.verbose:
+        RDLogger.DisableLog('rdApp.*')
+
+    if args.aggregate_meta:
+        meta_file_list = sorted(glob(os.path.join(args.meta_file, '*/result.pt')))
+        print(f'There are {len(meta_file_list)} files to aggregate')
+        test_index = []
+        for f in tqdm(meta_file_list, desc='Load meta files'):
+            test_index.append(torch.load(f))
+    else:
+        test_index = torch.load(args.meta_file)
+        if isinstance(test_index[0], dict):  # single datapoint sampling result
+            test_index = [test_index]
+
+    testset_results = []
+    with Pool(args.num_workers) as p:
+        for r in tqdm(p.starmap(partial(eval_single_datapoint, args=args),
+                                zip(test_index[:args.eval_num_examples], list(range(args.eval_num_examples)))),
+                      total=args.eval_num_examples, desc='Overall Eval'):
+            testset_results.append(r)
+
+    if args.result_path:
+        torch.save(testset_results, os.path.join(args.result_path, f'eval_all.pt'))
+
+    qed = [x['chem_results']['qed'] for r in testset_results for x in r]
+    sa = [x['chem_results']['sa'] for r in testset_results for x in r]
+    num_atoms = [len(x['pred_pos']) for r in testset_results for x in r]
+    logger.info('QED:   Mean: %.3f Median: %.3f' % (np.mean(qed), np.median(qed)))
+    logger.info('SA:    Mean: %.3f Median: %.3f' % (np.mean(sa), np.median(sa)))
+    logger.info('Num atoms:   Mean: %.3f Median: %.3f' % (np.mean(num_atoms), np.median(num_atoms)))
+    if args.docking_mode in ['vina', 'qvina']:
+        vina = [x['vina'][0]['affinity'] for r in testset_results for x in r]
+        logger.info('Vina:  Mean: %.3f Median: %.3f' % (np.mean(vina), np.median(vina)))
+    elif args.docking_mode in ['vina_full', 'vina_score']:
+        vina_score_only = [x['vina']['score_only'][0]['affinity'] for r in testset_results for x in r]
+        vina_min = [x['vina']['minimize'][0]['affinity'] for r in testset_results for x in r]
+        logger.info('Vina Score:  Mean: %.3f Median: %.3f' % (np.mean(vina_score_only), np.median(vina_score_only)))
+        logger.info('Vina Min  :  Mean: %.3f Median: %.3f' % (np.mean(vina_min), np.median(vina_min)))
+        if args.docking_mode == 'vina_full':
+            vina_dock = [x['vina']['dock'][0]['affinity'] for r in testset_results for x in r]
+            logger.info('Vina Dock :  Mean: %.3f Median: %.3f' % (np.mean(vina_dock), np.median(vina_dock)))

scripts/inference.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+source ${ROCM_PATH}/cuda/env.sh
+export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/lib:$LD_LIBRARY_PATH"
+export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/python3.11/site-packages/fastpt/torch/lib:$LD_LIBRARY_PATH"
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+#source ${PROJECT_ROOT}/env.sh
+
+echo "ONESCIENCE_PATH:" $PROJECT_ROOT
+echo ${ONESCIENCE_DATASETS_DIR}
+echo ${ONESCIENCE_MODELS_DIR}
+
+#cd ${PROJECT_ROOT}/examples/biosciences/targetdiff
+
+CKPT_PATH=${ONESCIENCE_MODELS_DIR}/targetdiff/pretrained_models/egnn_pdbbind_v2016.pt
+PROTEIN_PATH=${ONESCIENCE_DATASETS_DIR}/targetdiff/examples/3ug2_protein.pdb
+LIGAND_PATH=${ONESCIENCE_DATASETS_DIR}/targetdiff/examples/3ug2_ligand.sdf
+KIND=${4:-Kd}
+DEVICE=${5:-cuda}
+
+python scripts/property_prediction/fixed_inference.py \
+    --ckpt_path ${CKPT_PATH} \
+    --protein_path ${PROTEIN_PATH} \
+    --ligand_path ${LIGAND_PATH} \
+    --kind ${KIND} \
+    --device ${DEVICE}

scripts/likelihood_est_diffusion.py

@@ -0,0 +1,198 @@
+import argparse
+import os
+import pickle
+
+import numpy as np
+import torch
+from torch_geometric.data import Batch
+from torch_geometric.transforms import Compose
+from tqdm.auto import tqdm
+
+import onescience.utils.targetdiff.misc as misc
+import onescience.utils.targetdiff.transforms as trans
+from onescience.datapipes.targetdiff import get_dataset
+from onescience.datapipes.targetdiff.pl_data import FOLLOW_BATCH
+from models.molopt_score_model import ScorePosNet3D
+
+
+def data_likelihood_estimation(model, data, time_steps, batch_size=1, device='cuda:0'):
+    num_timesteps = len(time_steps)
+    num_batch = int(np.ceil(num_timesteps / batch_size))
+    all_kl_pos, all_kl_v = [], []
+
+    cur_i = 0
+    # t in [T-1, ..., 0]
+    for i in range(num_batch):
+        n_data = batch_size if i < num_batch - 1 else num_timesteps - batch_size * (num_batch - 1)
+        batch = Batch.from_data_list([data.clone() for _ in range(n_data)], follow_batch=FOLLOW_BATCH).to(device)
+        time_step = time_steps[cur_i:cur_i + n_data]
+
+        kl_pos, kl_v = model.likelihood_estimation(
+            protein_pos=batch.protein_pos,
+            protein_v=batch.protein_atom_feature.float(),
+            batch_protein=batch.protein_element_batch,
+
+            ligand_pos=batch.ligand_pos,
+            ligand_v=batch.ligand_atom_feature_full,
+            batch_ligand=batch.ligand_element_batch,
+
+            time_step=time_step
+        )
+        all_kl_pos.append(kl_pos)
+        all_kl_v.append(kl_v)
+        cur_i += n_data
+
+    # prior
+    batch = Batch.from_data_list([data.clone() for _ in range(1)], follow_batch=FOLLOW_BATCH).to(device)
+    time_step = torch.tensor([model.num_timesteps], device=device)
+    kl_pos_prior, kl_v_prior = model.likelihood_estimation(
+        protein_pos=batch.protein_pos,
+        protein_v=batch.protein_atom_feature.float(),
+        batch_protein=batch.protein_element_batch,
+
+        ligand_pos=batch.ligand_pos,
+        ligand_v=batch.ligand_atom_feature_full,
+        batch_ligand=batch.ligand_element_batch,
+
+        time_step=time_step
+    )
+    all_kl_pos, all_kl_v = torch.cat(all_kl_pos), torch.cat(all_kl_v)
+    sum_kl_pos, sum_kl_v = model.num_timesteps * torch.mean(all_kl_pos), model.num_timesteps * torch.mean(all_kl_v)
+    all_kl_pos, all_kl_v = torch.cat([all_kl_pos, kl_pos_prior]), torch.cat([all_kl_v, kl_v_prior])
+    sum_kl_pos += kl_pos_prior[0]
+    sum_kl_v += kl_v_prior[0]
+    return all_kl_pos.cpu(), all_kl_v.cpu(), sum_kl_pos.item(), sum_kl_v.item()
+
+
+def get_dataset_result(dset, affinity_info):
+    valid_id = []
+    for data_id in tqdm(range(len(dset)), desc='Filtering data'):
+        data = dset[data_id]
+        ligand_fn_key = data.ligand_filename[:-4]
+        pk = affinity_info[ligand_fn_key]['pk']
+        if pk > 0:
+            valid_id.append(data_id)
+    print(f'There are {len(valid_id)} examples with valid pK in total.')
+
+    all_results = []
+    for data_id in tqdm(valid_id, desc='Evaluating'):
+        data = dset[data_id]
+        # likelihoods
+        time_steps = torch.tensor(list(range(0, 1000, 100)), device=args.device)
+        all_kl_pos, all_kl_v, sum_kl_pos, sum_kl_v = data_likelihood_estimation(
+            model, data, time_steps, batch_size=args.batch_size, device=args.device)
+        kl = sum_kl_pos + sum_kl_v
+
+        # embedding
+        batch = Batch.from_data_list([data.clone() for _ in range(1)], follow_batch=FOLLOW_BATCH).to(args.device)
+        preds = model.fetch_embedding(
+            protein_pos=batch.protein_pos,
+            protein_v=batch.protein_atom_feature.float(),
+            batch_protein=batch.protein_element_batch,
+
+            ligand_pos=batch.ligand_pos,
+            ligand_v=batch.ligand_atom_feature_full,
+            batch_ligand=batch.ligand_element_batch,
+        )
+
+        # gather results
+        ligand_fn_key = data.ligand_filename[:-4]
+        result = {
+            'idx': data_id,
+            **affinity_info[ligand_fn_key],
+            'kl_pos': all_kl_pos,
+            'kl_v': all_kl_v,
+            'nll': kl,
+            'pred_ligand_v': preds['pred_ligand_v'].cpu(),
+            'final_h': preds['final_h'].cpu(),
+            'final_ligand_h': preds['final_ligand_h'].cpu()
+        }
+        all_results.append(result)
+    return all_results
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--device', type=str, default='cuda:0')
+    parser.add_argument('--config', type=str, default='configs/sampling/final_diffusion/aromatic_176k.yml')
+    parser.add_argument('--affinity_path', type=str, default='data/affinity_info.pkl')
+    parser.add_argument('--index_path', type=str, default='data/crossdocked_v1.1_rmsd1.0_pocket10/index.pkl')
+    parser.add_argument('--batch_size', type=int, default=4)
+    parser.add_argument('--result_path', type=str, default='./outputs_embedding')
+
+    args = parser.parse_args()
+
+    logger = misc.get_logger('evaluate')
+
+    if os.path.exists(args.affinity_path):
+        with open(args.affinity_path, 'rb') as f:
+            affinity_info = pickle.load(f)
+    else:
+        # collect index
+        with open(args.index_path, 'rb') as f:
+            index = pickle.load(f)
+        affinity_info = {}
+        for pdb_file, sdf_file, rmsd in index:
+            affinity_info[sdf_file[:-4]] = {'rmsd': rmsd}
+
+        # fetch reference vina score / binding affinity
+        types_path = 'data/CrossDocked2020/types/it2_tt_v1.1_completeset_train0.types'
+        with open(types_path, 'r') as f:
+            for ln in tqdm(f.readlines()):
+                # <label> <pK> <RMSD to crystal> <Receptor> <Ligand> # <Autodock Vina score>
+                _, pk, rmsd, protein_fn, ligand_fn, vina = ln.split()
+                ligand_raw_fn = ligand_fn[:ligand_fn.rfind('.')]
+                if ligand_raw_fn in affinity_info:
+                    affinity_info[ligand_raw_fn].update({
+                        'pk': float(pk),
+                        'vina': float(vina[1:])
+                    })
+
+        # save affinity info
+        with open(args.affinity_path, 'wb') as f:
+            pickle.dump(affinity_info, f)
+
+    # Load config
+    config = misc.load_config(args.config)
+    logger.info(config)
+    misc.seed_all(config.sample.seed)
+
+    # Load checkpoint
+    ckpt = torch.load(config.model.checkpoint, map_location=args.device)
+
+    # Transforms
+    protein_featurizer = trans.FeaturizeProteinAtom()
+    ligand_featurizer = trans.FeaturizeLigandAtom(ckpt['config'].data.transform.ligand_atom_mode)
+    transform_list = [
+        protein_featurizer,
+        ligand_featurizer,
+        trans.FeaturizeLigandBond(),
+    ]
+    if ckpt['config'].data.transform.random_rot:
+        transform_list.append(trans.RandomRotation())
+    transform = Compose(transform_list)
+
+    # Load dataset
+    dataset, subsets = get_dataset(
+        config=ckpt['config'].data,
+        transform=transform
+    )
+    train_set, test_set = subsets['train'], subsets['test']
+    logger.info(f'Successfully load the dataset (size: {len(test_set)})!')
+
+    # Load model
+    model = ScorePosNet3D(
+        ckpt['config'].model,
+        protein_atom_feature_dim=protein_featurizer.feature_dim,
+        ligand_atom_feature_dim=ligand_featurizer.feature_dim
+    ).to(args.device)
+    model.load_state_dict(ckpt['model'], strict=False if 'train_config' in config.model else True)
+    logger.info(f'Successfully load the model! {config.model.checkpoint}')
+
+    # filter data with valid pK, compute likelihood and embedding
+    valid_train_results = get_dataset_result(train_set, affinity_info)
+    valid_test_results = get_dataset_result(test_set, affinity_info)
+
+    os.makedirs(args.result_path, exist_ok=True)
+    torch.save(valid_train_results, os.path.join(args.result_path, 'crossdocked_train.pt'))
+    torch.save(valid_test_results, os.path.join(args.result_path, 'crossdocked_test.pt'))

scripts/property_prediction/__init__.py

@@ -0,0 +1,2 @@
+"""TargetDiff property prediction scripts."""
+

scripts/property_prediction/eval_prop.py

@@ -0,0 +1,93 @@
+import argparse
+from tqdm.auto import tqdm
+import torch
+import torch.utils.tensorboard
+from torch_geometric.transforms import Compose
+
+from onescience.datapipes.targetdiff import get_dataset
+import onescience.utils.targetdiff.transforms_prop as utils_trans
+import onescience.utils.targetdiff.misc as utils_misc
+import numpy as np
+from onescience.datapipes.targetdiff.protein_ligand import KMAP
+from scripts.property_prediction.local_misc_prop import get_model, get_dataloader, get_eval_scores
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--ckpt_path', type=str)
+    parser.add_argument('--device', type=str, default='cuda')
+    parser.add_argument('--seed', type=int, default=2021)
+    args = parser.parse_args()
+    utils_misc.seed_all(args.seed)
+
+    # Logging
+    logger = utils_misc.get_logger('eval')
+    logger.info(args)
+
+    # Load config
+    logger.info(f'Loading model from {args.ckpt_path}')
+    ckpt_restore = torch.load(args.ckpt_path, map_location=torch.device('cpu'))
+    config = ckpt_restore['config']
+    logger.info(f'ckpt_config: {config}')
+
+    # Transforms
+    protein_featurizer = utils_trans.FeaturizeProteinAtom()
+    ligand_featurizer = utils_trans.FeaturizeLigandAtom()
+    transform = Compose([
+        protein_featurizer,
+        ligand_featurizer,
+    ])
+
+    # Load model
+    model = get_model(config, protein_featurizer.feature_dim, ligand_featurizer.feature_dim)
+    model.load_state_dict(ckpt_restore['model'])
+    model = model.to(args.device)
+    # print(model)
+    # logger.info(f'# trainable parameters: {utils_misc.count_parameters(model) / 1e6:.4f} M')
+    model.eval()
+
+    # Datasets and loaders
+    # config.dataset.path = './data/pdbbind_v2020/pocket_10_refined'
+    # config.dataset.split = './data/pdbbind_v2020/pocket_10_refined/split.pt'
+    logger.info('Loading dataset...')
+    dataset, subsets = get_dataset(
+        config=config.dataset,
+        transform=transform,
+        heavy_only=config.dataset.get('heavy_only', False)
+    )
+    train_set, val_set, test_set = subsets['train'], subsets['val'], subsets['test']
+    logger.info(f'Train set: {len(train_set)} Val set: {len(val_set)} Test set: {len(test_set)}')
+    train_loader, val_loader, test_loader = get_dataloader(train_set, val_set, test_set, config)
+
+    def validate(epoch, data_loader, prefix='Test'):
+        sum_loss, sum_n = 0, 0
+        ytrue_arr, ypred_arr = [], []
+        y_kind = []
+        with torch.no_grad():
+            model.eval()
+            for batch in tqdm(data_loader, desc=prefix):
+                batch = batch.to(args.device)
+                loss, pred = model.get_loss(batch, pos_noise_std=0., return_pred=True)
+                sum_loss += loss.item() * len(batch.y)
+                sum_n += len(batch.y)
+                ypred_arr.append(pred.view(-1))
+                ytrue_arr.append(batch.y)
+                y_kind.append(batch.kind)
+        avg_loss = sum_loss / sum_n
+        logger.info('[%s] Epoch %03d | Loss %.6f' % (
+            prefix, epoch, avg_loss,
+        ))
+        ypred_arr = torch.cat(ypred_arr).cpu().numpy().astype(np.float64)
+        ytrue_arr = torch.cat(ytrue_arr).cpu().numpy().astype(np.float64)
+        y_kind = torch.cat(y_kind).cpu().numpy()
+        rmse = get_eval_scores(ypred_arr, ytrue_arr, logger)
+        for k, v in KMAP.items():
+            get_eval_scores(ypred_arr[y_kind == v], ytrue_arr[y_kind == v], logger, prefix=k)
+        return avg_loss
+
+    test_loss = validate(ckpt_restore['epoch'], test_loader)
+    print('Test loss: ', test_loss)
+
+
+if __name__ == '__main__':
+    main()

scripts/property_prediction/extract_pockets.py

@@ -0,0 +1,121 @@
+import argparse
+import multiprocessing as mp
+import os
+import pickle
+from functools import partial
+
+from rdkit import RDLogger
+from tqdm.auto import tqdm
+
+from onescience.datapipes.targetdiff.protein_ligand import parse_sdf_file_mol, read_mol, KMAP
+from onescience.utils.targetdiff.data import PDBProtein
+
+RDLogger.DisableLog('rdApp.*')
+
+
+def parse_pdbbind_index_file(raw_path, subset='refined'):
+    all_index = []
+    version = int(raw_path.rstrip('/')[-4:])
+    assert version >= 2016
+    if subset == 'refined':
+        data_path = os.path.join(raw_path, f'refined-set')
+        index_path = os.path.join(data_path, 'index', f'INDEX_refined_data.{version}')
+    elif subset == 'general':
+        data_path = os.path.join(raw_path, f'general-set-except-refined')
+        index_path = os.path.join(data_path, 'index', f'INDEX_general_PL_data.{version}')
+    else:
+        raise ValueError(subset)
+
+    all_files = os.listdir(data_path)
+    with open(index_path, 'r') as f:
+        lines = f.readlines()
+    for line in lines:
+        if line.startswith('#'): continue
+        index, res, year, pka, kv = line.split('//')[0].strip().split()
+        kind = [v for k, v in KMAP.items() if k in kv]
+        assert len(kind) == 1
+        if index in all_files:
+            all_index.append([index, res, year, pka, kind[0]])
+    return all_index
+
+
+def process_item(item, args):
+    pdb_idx, res, year, pka, kind = item
+    ligand_path = None
+    try:
+        if args.subset == 'refined':
+            pdb_path = os.path.join(args.source, 'refined-set', pdb_idx)
+        elif args.subset == 'general':
+            pdb_path = os.path.join(args.source, 'general-set-except-refined', pdb_idx)
+        else:
+            raise ValueError(args.subset)
+
+        protein_path = os.path.join(pdb_path, f'{pdb_idx}_protein.pdb')
+        ligand_sdf_path = os.path.join(pdb_path, f'{pdb_idx}_ligand.sdf')
+        ligand_mol2_path = os.path.join(pdb_path, f'{pdb_idx}_ligand.mol2')
+        mol, problem, ligand_path = read_mol(ligand_sdf_path, ligand_mol2_path)
+        if problem:
+            print('Read mol error.', item)
+            return None, ligand_path, res, pka, kind
+
+        protein = PDBProtein(protein_path)
+        # ligand = parse_sdf_file_mol(ligand_path, heavy_only=True)
+        ligand = parse_sdf_file_mol(ligand_path, heavy_only=False)
+        if args.dest is None:
+            pocket_path = os.path.join(pdb_path, f'{pdb_idx}_pocket{args.radius}.pdb')
+        else:
+            pocket_dir = os.path.join(args.dest, f'pocket_{args.radius}_{args.subset}', 'pockets')
+            os.makedirs(pocket_dir, exist_ok=True)
+            pocket_path = os.path.join(pocket_dir, f'{pdb_idx}_pocket{args.radius}.pdb')
+        if not os.path.exists(pocket_path):
+            pdb_block_pocket = protein.residues_to_pdb_block(
+                protein.query_residues_ligand(ligand, args.radius)
+            )
+            with open(pocket_path, 'w') as f:
+                f.write(pdb_block_pocket)
+        return pocket_path, ligand_path, res, pka, kind
+
+    except Exception as e:
+        print('Exception occured.', item, repr(e))
+        return None, ligand_path, res, pka, kind
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--source', type=str, default='./data/pdbbind_v2016')
+    parser.add_argument('--dest', type=str, default=None)
+    parser.add_argument('--fixed_sdf_dir', type=str, default='./data/pdbbind_v2016/fixed_sdf_files')
+    parser.add_argument('--subset', type=str, default='refined')
+    parser.add_argument('--refined_index_pkl', type=str, default=None)
+    parser.add_argument('--radius', type=int, default=10)
+    parser.add_argument('--num_workers', type=int, default=16)
+    args = parser.parse_args()
+
+    index = parse_pdbbind_index_file(args.source, args.subset)
+    # if not os.path.exists(args.fixed_sdf_dir):
+    #     os.makedirs(args.fixed_sdf_dir)
+
+    pool = mp.Pool(args.num_workers)
+    index_pocket = []
+    for item_pocket in tqdm(pool.imap_unordered(partial(process_item, args=args), index), total=len(index)):
+        index_pocket.append(item_pocket)
+    pool.close()
+
+    valid_index_pocket = []
+    for index in index_pocket:
+        if index[0] is not None:
+            valid_index_pocket.append(index)
+
+    output_root = args.source if args.dest is None else args.dest
+    save_path = os.path.join(output_root, f'pocket_{args.radius}_{args.subset}')
+    if not os.path.exists(save_path):
+        os.makedirs(save_path)
+    index_path = os.path.join(save_path, 'index.pkl')
+
+    if args.subset == 'general' and args.refined_index_pkl is not None:
+        with open(args.refined_index_pkl, 'rb') as f:
+            refined_index = pickle.load(f)
+        valid_index_pocket += refined_index
+    with open(index_path, 'wb') as f:
+        pickle.dump(valid_index_pocket, f)
+    print('Done. %d protein-ligand pairs in total.' % len(valid_index_pocket))

scripts/property_prediction/fixed_inference.py

@@ -0,0 +1,225 @@
+import argparse
+
+import torch
+import torch.utils.tensorboard
+from torch.utils.data import Dataset
+from torch_geometric.loader import DataLoader
+from tqdm.auto import tqdm
+from torch_geometric.transforms import Compose
+
+import onescience.utils.targetdiff.misc as utils_misc
+from onescience.datapipes.targetdiff.protein_ligand import KMAP, parse_sdf_file_mol
+from onescience.datapipes.targetdiff.pl_data import ProteinLigandData, torchify_dict
+from onescience.utils.targetdiff.data import PDBProtein
+import onescience.utils.targetdiff.transforms_prop as utils_trans
+from scripts.property_prediction.local_misc_prop import get_model
+
+
+class InferenceDataset(Dataset):
+    def __init__(self, data_list):
+        super().__init__()
+        self.data_list = data_list
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, idx):
+        data = self.data_list[idx]
+        return data
+
+
+def convert_data(pdb_path, ligand_path, transform, radius=10, pocket=False, heavy_only=False):
+    ligand_dict = parse_sdf_file_mol(ligand_path, heavy_only=heavy_only)
+    if not pocket:
+        protein = PDBProtein(pdb_path)
+        pdb_block_pocket = protein.residues_to_pdb_block(
+            protein.query_residues_ligand(ligand_dict, radius)
+        )
+        pocket_dict = PDBProtein(pdb_block_pocket).to_dict_atom()
+    else:
+        pocket_dict = PDBProtein(pdb_path).to_dict_atom()
+
+    data = ProteinLigandData.from_protein_ligand_dicts(
+        protein_dict=torchify_dict(pocket_dict),
+        ligand_dict=torchify_dict(ligand_dict),
+    )
+    data.protein_filename = pdb_path
+    data.ligand_filename = ligand_path
+    assert data.protein_pos.size(0) > 0
+    if transform is not None:
+        data = transform(data)
+    return data
+
+
+def adjust_ligand_features_for_compatibility(batch, expected_ligand_dim=30):
+    """
+    调整ligand特征以匹配模型期望的维度
+    对于RDKit版本差异导致的特征维度不匹配问题
+    """
+    if batch.ligand_atom_feature_full.size(-1) == expected_ligand_dim:
+        # 特征维度已经匹配，不需要调整
+        return batch
+    
+    current_dim = batch.ligand_atom_feature_full.size(-1)
+    print(f"Ligand features dimension mismatch: got {current_dim}, expected {expected_ligand_dim}")
+    
+    if current_dim > expected_ligand_dim:
+        # 当前特征维度更大，截取前面的部分
+        print(f"Truncating ligand features from {current_dim} to {expected_ligand_dim}")
+        batch.ligand_atom_feature_full = batch.ligand_atom_feature_full[:, :expected_ligand_dim]
+    elif current_dim < expected_ligand_dim:
+        # 当前特征维度较小，用零填充
+        print(f"Padding ligand features from {current_dim} to {expected_ligand_dim}")
+        padding = torch.zeros(batch.ligand_atom_feature_full.size(0), 
+                              expected_ligand_dim - current_dim,
+                              dtype=batch.ligand_atom_feature_full.dtype,
+                              device=batch.ligand_atom_feature_full.device)
+        batch.ligand_atom_feature_full = torch.cat([batch.ligand_atom_feature_full, padding], dim=-1)
+    
+    return batch
+
+
+def load_model_compatible(ckpt_restore, protein_featurizer, ligand_featurizer):
+    """
+    智能加载模型，兼容不同版本的RDKit导致的参数形状差异
+    """
+    config = ckpt_restore['config']
+    
+    # 首先尝试按原始方式加载模型
+    try:
+        model = get_model(config, protein_featurizer.feature_dim, ligand_featurizer.feature_dim)
+        model.load_state_dict(ckpt_restore['model'])
+        return model
+    except RuntimeError as e:
+        if "size mismatch" in str(e) and "ligand_atom_emb.weight" in str(e):
+            print(f"Found size mismatch error: {e}")
+            
+            # 从checkpoint中推断原始的ligand特征维度
+            original_ligand_feature_dim = None
+            for key, tensor in ckpt_restore['model'].items():
+                if 'ligand_atom_emb.weight' in key:
+                    original_ligand_feature_dim = tensor.shape[1]
+                    break
+            
+            if original_ligand_feature_dim is not None:
+                print(f"Detected original ligand feature dimension: {original_ligand_feature_dim}")
+                
+                # 创建具有正确维度的模型
+                # 注意：这里我们传入原始的ligand featurizer维度，但稍后会调整输入数据
+                model = get_model(config, protein_featurizer.feature_dim, original_ligand_feature_dim)
+                
+                # 智能加载状态字典，只加载形状匹配的参数
+                model_dict = model.state_dict()
+                pretrained_dict = {}
+                skipped_params = []
+                
+                for k, v in ckpt_restore['model'].items():
+                    if k in model_dict:
+                        if model_dict[k].shape == v.shape:
+                            pretrained_dict[k] = v
+                        else:
+                            print(f"Shape mismatch for {k}: checkpoint {v.shape} vs model {model_dict[k].shape}")
+                            skipped_params.append(k)
+                    else:
+                        print(f"Parameter {k} not found in current model")
+                        skipped_params.append(k)
+                
+                print(f"Loaded {len(pretrained_dict)} parameters, skipped {len(skipped_params)} parameters")
+                
+                # 更新模型状态字典
+                model_dict.update(pretrained_dict)
+                model.load_state_dict(model_dict)
+                
+                print("Model loaded successfully with compatibility adjustments")
+                return model
+            else:
+                raise e
+        else:
+            raise e
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--ckpt_path', type=str)
+    parser.add_argument('--protein_path', type=str)
+    parser.add_argument('--ligand_path', type=str)
+    parser.add_argument('--kind', type=str, default='Ki', choices=['Ki', 'Kd', 'IC50'])
+    parser.add_argument('--device', type=str, default='cuda')
+    parser.add_argument('--seed', type=int, default=2021)
+    args = parser.parse_args()
+    utils_misc.seed_all(args.seed)
+
+    # Logging
+    logger = utils_misc.get_logger('eval')
+    logger.info(args)
+
+    # Load config
+    logger.info(f'Loading model from {args.ckpt_path}')
+    ckpt_restore = torch.load(args.ckpt_path, map_location=torch.device('cpu'), weights_only=False)
+    config = ckpt_restore['config']
+    logger.info(f'ckpt_config: {config}')
+
+    # Transforms
+    protein_featurizer = utils_trans.FeaturizeProteinAtom()
+    ligand_featurizer = utils_trans.FeaturizeLigandAtom()
+    transform = Compose([
+        protein_featurizer,
+        ligand_featurizer,
+    ])
+
+    # Load model with compatibility fix
+    model = load_model_compatible(ckpt_restore, protein_featurizer, ligand_featurizer)
+    
+    # Get the expected ligand feature dimension from the loaded model
+    expected_ligand_dim = None
+    for name, param in model.named_parameters():
+        if 'ligand_atom_emb.weight' in name:
+            expected_ligand_dim = param.shape[1]  # input dimension
+            break
+    
+    if expected_ligand_dim is None:
+        # If we can't find the dimension from parameters, it means we couldn't load that layer
+        # Use the value detected during model loading
+        expected_ligand_dim = 30  # Default based on error message
+    
+    model = model.to(args.device)
+    logger.info(f'# trainable parameters: {utils_misc.count_parameters(model) / 1e6:.4f} M')
+    model.eval()
+
+    test_data = convert_data(args.protein_path, args.ligand_path, transform,
+                             heavy_only=config.dataset.get('heavy_only', False))
+    test_data.kind = KMAP[args.kind]
+    test_set = InferenceDataset([test_data])
+    test_loader = DataLoader(test_set, batch_size=1, shuffle=False,
+                             follow_batch=['protein_element', 'ligand_element'])
+
+    with torch.no_grad():
+        model.eval()
+        for batch in tqdm(test_loader, desc='Inference'):
+            batch = batch.to(args.device)
+            
+            # Adjust ligand features to match the expected dimension
+            batch = adjust_ligand_features_for_compatibility(batch, expected_ligand_dim)
+            
+            pred = model(
+                protein_pos=batch.protein_pos,
+                protein_atom_feature=batch.protein_atom_feature.float(),
+                ligand_pos=batch.ligand_pos,
+                ligand_atom_feature=batch.ligand_atom_feature_full.float(),
+                batch_protein=batch.protein_element_batch,
+                batch_ligand=batch.ligand_element_batch,
+                output_kind=batch.kind
+            )
+
+            print(f'PDB ID: {batch.protein_filename[0]} '
+                  f'Prediction: {args.kind}={unit_transform(pred.cpu().squeeze()):.2e} m')
+
+
+def unit_transform(pka):
+    # pka = -log10 Kd / Ki
+    affinity = torch.pow(10, -pka.cpu().squeeze())
+    return affinity
+
+
+if __name__ == '__main__':
+    main()

scripts/property_prediction/inference.py

@@ -0,0 +1,124 @@
+import argparse
+
+import torch
+import torch.utils.tensorboard
+from torch.utils.data import Dataset
+from torch_geometric.loader import DataLoader
+from tqdm.auto import tqdm
+from torch_geometric.transforms import Compose
+
+import onescience.utils.targetdiff.misc as utils_misc
+from onescience.datapipes.targetdiff.protein_ligand import KMAP, parse_sdf_file_mol
+from onescience.datapipes.targetdiff.pl_data import ProteinLigandData, torchify_dict
+from onescience.utils.targetdiff.data import PDBProtein
+import onescience.utils.targetdiff.transforms_prop as utils_trans
+from scripts.property_prediction.local_misc_prop import get_model
+
+
+class InferenceDataset(Dataset):
+    def __init__(self, data_list):
+        super().__init__()
+        self.data_list = data_list
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, idx):
+        data = self.data_list[idx]
+        return data
+
+
+def convert_data(pdb_path, ligand_path, transform, radius=10, pocket=False, heavy_only=False):
+    ligand_dict = parse_sdf_file_mol(ligand_path, heavy_only=heavy_only)
+    if not pocket:
+        protein = PDBProtein(pdb_path)
+        pdb_block_pocket = protein.residues_to_pdb_block(
+            protein.query_residues_ligand(ligand_dict, radius)
+        )
+        pocket_dict = PDBProtein(pdb_block_pocket).to_dict_atom()
+    else:
+        pocket_dict = PDBProtein(pdb_path).to_dict_atom()
+
+    data = ProteinLigandData.from_protein_ligand_dicts(
+        protein_dict=torchify_dict(pocket_dict),
+        ligand_dict=torchify_dict(ligand_dict),
+    )
+    data.protein_filename = pdb_path
+    data.ligand_filename = ligand_path
+    assert data.protein_pos.size(0) > 0
+    if transform is not None:
+        data = transform(data)
+    return data
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--ckpt_path', type=str)
+    parser.add_argument('--protein_path', type=str)
+    parser.add_argument('--ligand_path', type=str)
+    parser.add_argument('--kind', type=str, default='Ki', choices=['Ki', 'Kd', 'IC50'])
+    parser.add_argument('--device', type=str, default='cuda')
+    parser.add_argument('--seed', type=int, default=2021)
+    args = parser.parse_args()
+    utils_misc.seed_all(args.seed)
+
+    # Logging
+    logger = utils_misc.get_logger('eval')
+    logger.info(args)
+
+    # Load config
+    logger.info(f'Loading model from {args.ckpt_path}')
+    ckpt_restore = torch.load(args.ckpt_path, map_location=torch.device('cpu'),weights_only=
+False)
+    config = ckpt_restore['config']
+    logger.info(f'ckpt_config: {config}')
+
+    # Transforms
+    protein_featurizer = utils_trans.FeaturizeProteinAtom()
+    ligand_featurizer = utils_trans.FeaturizeLigandAtom()
+    transform = Compose([
+        protein_featurizer,
+        ligand_featurizer,
+    ])
+
+    # Load model
+    model = get_model(config, protein_featurizer.feature_dim, ligand_featurizer.feature_dim)
+    model.load_state_dict(ckpt_restore['model'])
+    model = model.to(args.device)
+    # print(model)
+    logger.info(f'# trainable parameters: {utils_misc.count_parameters(model) / 1e6:.4f} M')
+    model.eval()
+
+    test_data = convert_data(args.protein_path, args.ligand_path, transform,
+                             heavy_only=config.dataset.get('heavy_only', False))
+    test_data.kind = KMAP[args.kind]
+    test_set = InferenceDataset([test_data])
+    test_loader = DataLoader(test_set, batch_size=1, shuffle=False,
+                             follow_batch=['protein_element', 'ligand_element'])
+
+    with torch.no_grad():
+        model.eval()
+        for batch in tqdm(test_loader, desc='Inference'):
+            batch = batch.to(args.device)
+            pred = model(
+                protein_pos=batch.protein_pos,
+                protein_atom_feature=batch.protein_atom_feature.float(),
+                ligand_pos=batch.ligand_pos,
+                ligand_atom_feature=batch.ligand_atom_feature_full.float(),
+                batch_protein=batch.protein_element_batch,
+                batch_ligand=batch.ligand_element_batch,
+                output_kind=batch.kind
+            )
+
+            print(f'PDB ID: {batch.protein_filename[0]} '
+                  f'Prediction: {args.kind}={unit_transform(pred.cpu().squeeze()):.2e} m')
+
+
+def unit_transform(pka):
+    # pka = -log10 Kd / Ki
+    affinity = torch.pow(10, -pka.cpu().squeeze())
+    return affinity
+
+
+if __name__ == '__main__':
+    main()

scripts/property_prediction/local_misc_prop.py

@@ -0,0 +1,64 @@
+import numpy as np
+from scipy.stats import pearsonr, spearmanr
+from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
+from torch_geometric.loader import DataLoader
+
+from models.property_pred.prop_model import PropPredNet, PropPredNetEnc
+
+
+def get_eval_scores(ypred_arr, ytrue_arr, logger, prefix='All'):
+    if len(ypred_arr) == 0:
+        return None
+    rmse = np.sqrt(mean_squared_error(ytrue_arr, ypred_arr))
+    mae = mean_absolute_error(ytrue_arr, ypred_arr)
+    r2 = r2_score(ytrue_arr, ypred_arr)
+    pearson, ppval = pearsonr(ytrue_arr, ypred_arr)
+    spearman, spval = spearmanr(ytrue_arr, ypred_arr)
+    mean = np.mean(ypred_arr)
+    std = np.std(ypred_arr)
+    logger.info("Evaluation Summary:")
+    logger.info(
+        "[%4s] num: %3d, RMSE: %.3f, MAE: %.3f, "
+        "R^2 score: %.3f, Pearson: %.3f, Spearman: %.3f, mean/std: %.3f/%.3f" % (
+            prefix, len(ypred_arr), rmse, mae, r2, pearson, spearman, mean, std))
+    return rmse
+
+
+def get_dataloader(train_set, val_set, test_set, config):
+    follow_batch = ['protein_element', 'ligand_element']
+    collate_exclude_keys = ['ligand_nbh_list']
+    train_loader = DataLoader(
+        train_set,
+        batch_size=config.train.batch_size,
+        shuffle=True,
+        num_workers=config.train.num_workers,
+        follow_batch=follow_batch,
+        exclude_keys=collate_exclude_keys
+    )
+    val_loader = DataLoader(val_set, config.train.batch_size, shuffle=False,
+                            follow_batch=follow_batch, exclude_keys=collate_exclude_keys)
+    test_loader = DataLoader(test_set, config.train.batch_size, shuffle=False,
+                             follow_batch=follow_batch, exclude_keys=collate_exclude_keys)
+    return train_loader, val_loader, test_loader
+
+
+def get_model(config, protein_atom_feat_dim, ligand_atom_feat_dim):
+    if config.model.encoder.name == 'egnn_enc':
+        model = PropPredNetEnc(
+            config.model,
+            protein_atom_feature_dim=protein_atom_feat_dim,
+            ligand_atom_feature_dim=ligand_atom_feat_dim,
+            enc_ligand_dim=config.model.enc_ligand_dim,
+            enc_node_dim=config.model.enc_node_dim,
+            enc_graph_dim=config.model.enc_graph_dim,
+            enc_feature_type=config.model.enc_feature_type,
+            output_dim=1
+        )
+    else:
+        model = PropPredNet(
+            config.model,
+            protein_atom_feature_dim=protein_atom_feat_dim,
+            ligand_atom_feature_dim=ligand_atom_feat_dim,
+            output_dim=3
+        )
+    return model

scripts/property_prediction/pdbbind_split.py

@@ -0,0 +1,66 @@
+import os
+import pickle
+import random
+import argparse
+import torch
+import numpy as np
+
+
+def coretest_split(index_path, test_path, val_ratio=0.1, val_num=None):
+    with open(index_path, 'rb') as f:
+        index = pickle.load(f)
+    test_ids = [f for f in os.listdir(test_path) if len(f) == 4]
+    all_ids = [os.path.basename(i[0])[:4] for i in index]
+    print('Original test set size: ', len(test_ids))
+    test_index = [all_ids.index(test_id) for test_id in test_ids if test_id in all_ids]
+    train_val_index = list(set(range(len(all_ids))) - set(test_index))
+    assert len(train_val_index) == len(all_ids) - len(test_index)
+    random.shuffle(train_val_index)
+    if val_num is not None:
+        n_val = val_num
+    else:
+        n_val = int(len(train_val_index) * val_ratio)
+    val_index = train_val_index[:n_val]
+    train_index = train_val_index[n_val:]
+    return train_index, val_index, test_index
+
+
+def time_split(index_path):
+    valid_ids = np.loadtxt("./data/pdbbind_v2020/timesplit_no_lig_overlap_val", dtype=str)
+    test_ids = np.loadtxt("./data/pdbbind_v2020/timesplit_test", dtype=str)
+
+    with open(index_path, 'rb') as f:
+        index = pickle.load(f)
+    all_ids = [os.path.basename(i[0])[:4] for i in index]
+    val_index = [all_ids.index(val_id) for val_id in valid_ids if val_id in all_ids]
+    test_index = [all_ids.index(test_id) for test_id in test_ids if test_id in all_ids]
+    train_index = list(set(range(len(all_ids))) - set(test_index) - set(val_index))
+    return train_index, val_index, test_index
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--index_path', type=str, default='./data/pdbbind_v2016/pocket_10/index.pkl')
+    parser.add_argument('--split_mode', type=str, choices=['coreset', 'time'], default='coreset')
+    parser.add_argument('--test_path', type=str, default='./data/pdbbind_v2016/coreset')
+    parser.add_argument('--val_ratio', type=float, default=0.1)
+    parser.add_argument('--val_num', type=int, default=None)
+    parser.add_argument('--save_path', type=str, default='./data/pdbbind_v2016/pocket_10/split.pt')
+    parser.add_argument('--seed', type=int, default=2021)
+    args = parser.parse_args()
+    random.seed(args.seed)
+
+    if args.split_mode == 'coreset':
+        train_index, val_index, test_index = coretest_split(args.index_path, args.test_path,
+                                                            args.val_ratio, args.val_num)
+    elif args.split_mode == 'time':
+        train_index, val_index, test_index = time_split(args.index_path)
+    else:
+        raise ValueError(args.split_mode)
+    torch.save({
+        'train': train_index,
+        'val': val_index,
+        'test': test_index
+    }, args.save_path)
+    print('Train %d, Validation %d, Test %d.' % (len(train_index), len(val_index), len(test_index)))
+    print('Done.')

scripts/property_prediction/train_prop.py

@@ -0,0 +1,229 @@
+import os
+import shutil
+import argparse
+from tqdm.auto import tqdm
+import torch
+from torch.nn.utils import clip_grad_norm_
+import torch.utils.tensorboard
+import yaml
+from torch_geometric.transforms import Compose
+
+from onescience.datapipes.targetdiff import get_dataset
+import onescience.utils.targetdiff.transforms_prop as utils_trans
+import onescience.utils.targetdiff.misc as utils_misc
+from onescience.utils.targetdiff.train import get_scheduler, get_optimizer
+import numpy as np
+from onescience.datapipes.targetdiff.protein_ligand import KMAP
+from scripts.property_prediction.local_misc_prop import get_model, get_dataloader, get_eval_scores
+
+
+REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..', '..', '..'))
+MODELS_SNAPSHOT_SRC = os.path.join(REPO_ROOT, 'src', 'onescience', 'models', 'targetdiff')
+
+
+def parse_override_value(raw_value, old_value):
+    parsed_value = yaml.safe_load(raw_value)
+    if old_value is None:
+        return parsed_value
+    if isinstance(old_value, bool):
+        if isinstance(parsed_value, bool):
+            return parsed_value
+        return str(parsed_value).lower() in ('1', 'true', 'yes', 'y')
+    if isinstance(old_value, tuple):
+        if isinstance(parsed_value, str):
+            return tuple(item.strip() for item in parsed_value.split(','))
+        return tuple(parsed_value)
+    if isinstance(old_value, list):
+        if isinstance(parsed_value, str):
+            return [item.strip() for item in parsed_value.split(',')]
+        return list(parsed_value)
+    return type(old_value)(parsed_value)
+
+
+def apply_config_overrides(config, overrides):
+    if len(overrides) % 2 != 0:
+        raise ValueError('Config overrides must use "--key value" pairs.')
+
+    for key_arg, raw_value in zip(overrides[::2], overrides[1::2]):
+        if not key_arg.startswith('--'):
+            raise ValueError(f'Config override key must start with "--": {key_arg}')
+
+        key_path = key_arg[2:]
+        parts = key_path.split('.')
+        node = config
+        for part in parts[:-1]:
+            if part not in node:
+                raise KeyError(f'Unknown config override: {key_path}')
+            node = node[part]
+
+        leaf = parts[-1]
+        if leaf not in node:
+            raise KeyError(f'Unknown config override: {key_path}')
+        node[leaf] = parse_override_value(raw_value, node[leaf])
+
+    return config
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('config', type=str)
+    parser.add_argument('--device', type=str, default='cuda')
+    parser.add_argument('--logdir', type=str, default='./logs')
+    parser.add_argument('--tag', type=str, default='')
+    args, config_overrides = parser.parse_known_args()
+
+    # Load configs
+    config = utils_misc.load_config(args.config)
+    config = apply_config_overrides(config, config_overrides)
+    config_name = os.path.basename(args.config)[:os.path.basename(args.config).rfind('.')]
+    utils_misc.seed_all(config.train.seed)
+
+    # Logging
+    log_dir = utils_misc.get_new_log_dir(args.logdir, prefix=config_name, tag=args.tag)
+    ckpt_dir = os.path.join(log_dir, 'checkpoints')
+    os.makedirs(ckpt_dir, exist_ok=True)
+    logger = utils_misc.get_logger('train', log_dir)
+    writer = torch.utils.tensorboard.SummaryWriter(log_dir)
+    logger.info(args)
+    logger.info(config)
+    shutil.copyfile(args.config, os.path.join(log_dir, os.path.basename(args.config)))
+    local_models_dir = './models'
+    if os.path.isdir(local_models_dir):
+        shutil.copytree(local_models_dir, os.path.join(log_dir, 'models'))
+    elif os.path.isdir(MODELS_SNAPSHOT_SRC):
+        shutil.copytree(MODELS_SNAPSHOT_SRC, os.path.join(log_dir, 'models'))
+    else:
+        logger.warning('Skip model source snapshot: no ./models or migrated TargetDiff models directory found.')
+
+    # Transforms
+    protein_featurizer = utils_trans.FeaturizeProteinAtom()
+    ligand_featurizer = utils_trans.FeaturizeLigandAtom()
+    transform = Compose([
+        protein_featurizer,
+        ligand_featurizer,
+    ])
+
+    # Datasets and loaders
+    logger.info('Loading dataset...')
+    dataset, subsets = get_dataset(
+        config=config.dataset,
+        transform=transform,
+        emb_path=config.dataset.emb_path if 'emb_path' in config.dataset else None,
+        heavy_only=config.dataset.heavy_only
+    )
+    train_set, val_set, test_set = subsets['train'], subsets['val'], subsets['test']
+    logger.info(f'Train set: {len(train_set)} Val set: {len(val_set)} Test set: {len(test_set)}')
+    train_loader, val_loader, test_loader = get_dataloader(train_set, val_set, test_set, config)
+    # Model
+    logger.info('Building model...')
+    model = get_model(config, protein_featurizer.feature_dim, ligand_featurizer.feature_dim)
+    model = model.to(args.device)
+    logger.info(f'# trainable parameters: {utils_misc.count_parameters(model) / 1e6:.4f} M')
+
+    # Optimizer and scheduler
+    optimizer = get_optimizer(config.train.optimizer, model)
+    scheduler = get_scheduler(config.train.scheduler, optimizer)
+
+    def train(epoch):
+        model.train()
+        optimizer.zero_grad()
+        it = 0
+        num_it = len(train_loader)
+        for batch in tqdm(train_loader, dynamic_ncols=True, desc=f'Epoch {epoch}', position=1):
+            it += 1
+            batch = batch.to(args.device)
+            # compute loss
+            loss = model.get_loss(batch, pos_noise_std=config.train.pos_noise_std)
+            loss.backward()
+            orig_grad_norm = clip_grad_norm_(model.parameters(), config.train.max_grad_norm)
+            optimizer.step()
+            optimizer.zero_grad()
+
+            if it % config.train.report_iter == 0:
+                logger.info('[Train] Epoch %03d Iter %04d | Loss %.6f | Lr %.4f * 1e-3' % (
+                    epoch, it, loss.item(), optimizer.param_groups[0]['lr'] * 1000
+                ))
+
+            writer.add_scalar('train/loss', loss, it + epoch * num_it)
+            writer.add_scalar('train/lr', optimizer.param_groups[0]['lr'], it + epoch * num_it)
+            writer.add_scalar('train/grad', orig_grad_norm, it + epoch * num_it)
+            writer.flush()
+
+    def validate(epoch, data_loader, scheduler, writer, prefix='Validate'):
+        sum_loss, sum_n = 0, 0
+        ytrue_arr, ypred_arr = [], []
+        y_kind = []
+        with torch.no_grad():
+            model.eval()
+            for batch in tqdm(data_loader, desc=prefix):
+                batch = batch.to(args.device)
+                loss, pred = model.get_loss(batch, pos_noise_std=0., return_pred=True)
+                sum_loss += loss.item() * len(batch.y)
+                sum_n += len(batch.y)
+                ypred_arr.append(pred.view(-1))
+                ytrue_arr.append(batch.y)
+                y_kind.append(batch.kind)
+        avg_loss = sum_loss / sum_n
+        logger.info('[%s] Epoch %03d | Loss %.6f' % (
+            prefix, epoch, avg_loss,
+        ))
+        ypred_arr = torch.cat(ypred_arr).cpu().numpy().astype(np.float64)
+        ytrue_arr = torch.cat(ytrue_arr).cpu().numpy().astype(np.float64)
+        y_kind = torch.cat(y_kind).cpu().numpy()
+        rmse = get_eval_scores(ypred_arr, ytrue_arr, logger)
+        for k, v in KMAP.items():
+            get_eval_scores(ypred_arr[y_kind == v], ytrue_arr[y_kind == v], logger, prefix=k)
+
+        if scheduler:
+            if config.train.scheduler.type == 'plateau':
+                scheduler.step(avg_loss)
+            elif config.train.scheduler.type == 'warmup_plateau':
+                scheduler.step_ReduceLROnPlateau(avg_loss)
+            else:
+                scheduler.step()
+
+        if writer:
+            writer.add_scalar('val/loss', avg_loss, epoch)
+            writer.add_scalar('val/rmse', rmse, epoch)
+            writer.flush()
+
+        return avg_loss
+
+    try:
+        best_val_loss = float('inf')
+        best_val_epoch = 0
+        patience = 0
+        for epoch in range(1, config.train.max_epochs + 1):
+            # with torch.autograd.detect_anomaly():
+            train(epoch)
+            if epoch % config.train.val_freq == 0 or epoch == config.train.max_epochs:
+                val_loss = validate(epoch, val_loader, scheduler, writer)
+                validate(epoch, test_loader, scheduler=None, writer=None, prefix='Test')
+
+                if val_loss < best_val_loss:
+                    patience = 0
+                    best_val_loss = val_loss
+                    best_val_epoch = epoch
+                    logger.info(f'Best val achieved at epoch {epoch}, val loss: {best_val_loss:.3f}')
+                    logger.info(f'Eval on Test set:')
+                    validate(epoch, test_loader, scheduler=None, writer=None, prefix='Test')
+                    ckpt_path = os.path.join(ckpt_dir, '%d.pt' % epoch)
+                    torch.save({
+                        'config': config,
+                        'model': model.state_dict(),
+                        'optimizer': optimizer.state_dict(),
+                        'scheduler': scheduler.state_dict(),
+                        'epoch': epoch,
+                    }, ckpt_path)
+                    logger.info(f'Model {log_dir}/{epoch}.pt saved!')
+                else:
+                    patience += 1
+                    logger.info(f'Val loss does not improve, patience: {patience} '
+                                f'(Best val loss: {best_val_loss:.3f} at epoch {best_val_epoch})')
+
+    except KeyboardInterrupt:
+        logger.info('Terminating...')
+
+
+if __name__ == '__main__':
+    main()

scripts/sample_diffusion.py

@@ -0,0 +1,188 @@
+import argparse
+import os
+import shutil
+import time
+
+import numpy as np
+import torch
+from torch_geometric.data import Batch
+from torch_geometric.transforms import Compose
+from torch_scatter import scatter_sum, scatter_mean
+from tqdm.auto import tqdm
+
+import onescience.utils.targetdiff.misc as misc
+import onescience.utils.targetdiff.transforms as trans
+from onescience.datapipes.targetdiff import get_dataset
+from onescience.datapipes.targetdiff.pl_data import FOLLOW_BATCH
+from models.molopt_score_model import ScorePosNet3D, log_sample_categorical
+from onescience.utils.targetdiff.evaluation import atom_num
+
+
+def unbatch_v_traj(ligand_v_traj, n_data, ligand_cum_atoms):
+    all_step_v = [[] for _ in range(n_data)]
+    for v in ligand_v_traj:  # step_i
+        v_array = v.cpu().numpy()
+        for k in range(n_data):
+            all_step_v[k].append(v_array[ligand_cum_atoms[k]:ligand_cum_atoms[k + 1]])
+    all_step_v = [np.stack(step_v) for step_v in all_step_v]  # num_samples * [num_steps, num_atoms_i]
+    return all_step_v
+
+
+def sample_diffusion_ligand(model, data, num_samples, batch_size=16, device='cuda:0',
+                            num_steps=None, pos_only=False, center_pos_mode='protein',
+                            sample_num_atoms='prior'):
+    all_pred_pos, all_pred_v = [], []
+    all_pred_pos_traj, all_pred_v_traj = [], []
+    all_pred_v0_traj, all_pred_vt_traj = [], []
+    time_list = []
+    num_batch = int(np.ceil(num_samples / batch_size))
+    current_i = 0
+    for i in tqdm(range(num_batch)):
+        n_data = batch_size if i < num_batch - 1 else num_samples - batch_size * (num_batch - 1)
+        batch = Batch.from_data_list([data.clone() for _ in range(n_data)], follow_batch=FOLLOW_BATCH).to(device)
+
+        t1 = time.time()
+        with torch.no_grad():
+            batch_protein = batch.protein_element_batch
+            if sample_num_atoms == 'prior':
+                pocket_size = atom_num.get_space_size(data.protein_pos.detach().cpu().numpy())
+                ligand_num_atoms = [atom_num.sample_atom_num(pocket_size).astype(int) for _ in range(n_data)]
+                batch_ligand = torch.repeat_interleave(torch.arange(n_data), torch.tensor(ligand_num_atoms)).to(device)
+            elif sample_num_atoms == 'range':
+                ligand_num_atoms = list(range(current_i + 1, current_i + n_data + 1))
+                batch_ligand = torch.repeat_interleave(torch.arange(n_data), torch.tensor(ligand_num_atoms)).to(device)
+            elif sample_num_atoms == 'ref':
+                batch_ligand = batch.ligand_element_batch
+                ligand_num_atoms = scatter_sum(torch.ones_like(batch_ligand), batch_ligand, dim=0).tolist()
+            else:
+                raise ValueError
+
+            # init ligand pos
+            center_pos = scatter_mean(batch.protein_pos, batch_protein, dim=0)
+            batch_center_pos = center_pos[batch_ligand]
+            init_ligand_pos = batch_center_pos + torch.randn_like(batch_center_pos)
+
+            # init ligand v
+            if pos_only:
+                init_ligand_v = batch.ligand_atom_feature_full
+            else:
+                uniform_logits = torch.zeros(len(batch_ligand), model.num_classes).to(device)
+                init_ligand_v = log_sample_categorical(uniform_logits)
+
+            r = model.sample_diffusion(
+                protein_pos=batch.protein_pos,
+                protein_v=batch.protein_atom_feature.float(),
+                batch_protein=batch_protein,
+
+                init_ligand_pos=init_ligand_pos,
+                init_ligand_v=init_ligand_v,
+                batch_ligand=batch_ligand,
+                num_steps=num_steps,
+                pos_only=pos_only,
+                center_pos_mode=center_pos_mode
+            )
+            ligand_pos, ligand_v, ligand_pos_traj, ligand_v_traj = r['pos'], r['v'], r['pos_traj'], r['v_traj']
+            ligand_v0_traj, ligand_vt_traj = r['v0_traj'], r['vt_traj']
+            # unbatch pos
+            ligand_cum_atoms = np.cumsum([0] + ligand_num_atoms)
+            ligand_pos_array = ligand_pos.cpu().numpy().astype(np.float64)
+            all_pred_pos += [ligand_pos_array[ligand_cum_atoms[k]:ligand_cum_atoms[k + 1]] for k in
+                             range(n_data)]  # num_samples * [num_atoms_i, 3]
+
+            all_step_pos = [[] for _ in range(n_data)]
+            for p in ligand_pos_traj:  # step_i
+                p_array = p.cpu().numpy().astype(np.float64)
+                for k in range(n_data):
+                    all_step_pos[k].append(p_array[ligand_cum_atoms[k]:ligand_cum_atoms[k + 1]])
+            all_step_pos = [np.stack(step_pos) for step_pos in
+                            all_step_pos]  # num_samples * [num_steps, num_atoms_i, 3]
+            all_pred_pos_traj += [p for p in all_step_pos]
+
+            # unbatch v
+            ligand_v_array = ligand_v.cpu().numpy()
+            all_pred_v += [ligand_v_array[ligand_cum_atoms[k]:ligand_cum_atoms[k + 1]] for k in range(n_data)]
+
+            all_step_v = unbatch_v_traj(ligand_v_traj, n_data, ligand_cum_atoms)
+            all_pred_v_traj += [v for v in all_step_v]
+
+            if not pos_only:
+                all_step_v0 = unbatch_v_traj(ligand_v0_traj, n_data, ligand_cum_atoms)
+                all_pred_v0_traj += [v for v in all_step_v0]
+                all_step_vt = unbatch_v_traj(ligand_vt_traj, n_data, ligand_cum_atoms)
+                all_pred_vt_traj += [v for v in all_step_vt]
+        t2 = time.time()
+        time_list.append(t2 - t1)
+        current_i += n_data
+    return all_pred_pos, all_pred_v, all_pred_pos_traj, all_pred_v_traj, all_pred_v0_traj, all_pred_vt_traj, time_list
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('config', type=str)
+    parser.add_argument('-i', '--data_id', type=int)
+    parser.add_argument('--device', type=str, default='cuda:0')
+    parser.add_argument('--batch_size', type=int, default=100)
+    parser.add_argument('--result_path', type=str, default='./outputs')
+    args = parser.parse_args()
+
+    logger = misc.get_logger('sampling')
+
+    # Load config
+    config = misc.load_config(args.config)
+    logger.info(config)
+    misc.seed_all(config.sample.seed)
+
+    # Load checkpoint
+    ckpt = torch.load(config.model.checkpoint, map_location=args.device)
+    logger.info(f"Training Config: {ckpt['config']}")
+
+    # Transforms
+    protein_featurizer = trans.FeaturizeProteinAtom()
+    ligand_atom_mode = ckpt['config'].data.transform.ligand_atom_mode
+    ligand_featurizer = trans.FeaturizeLigandAtom(ligand_atom_mode)
+    transform = Compose([
+        protein_featurizer,
+        ligand_featurizer,
+        trans.FeaturizeLigandBond(),
+    ])
+
+    # Load dataset
+    dataset, subsets = get_dataset(
+        config=ckpt['config'].data,
+        transform=transform
+    )
+    train_set, test_set = subsets['train'], subsets['test']
+    logger.info(f'Successfully load the dataset (size: {len(test_set)})!')
+
+    # Load model
+    model = ScorePosNet3D(
+        ckpt['config'].model,
+        protein_atom_feature_dim=protein_featurizer.feature_dim,
+        ligand_atom_feature_dim=ligand_featurizer.feature_dim
+    ).to(args.device)
+    model.load_state_dict(ckpt['model'])
+    logger.info(f'Successfully load the model! {config.model.checkpoint}')
+
+    data = test_set[args.data_id]
+    pred_pos, pred_v, pred_pos_traj, pred_v_traj, pred_v0_traj, pred_vt_traj, time_list = sample_diffusion_ligand(
+        model, data, config.sample.num_samples,
+        batch_size=args.batch_size, device=args.device,
+        num_steps=config.sample.num_steps,
+        pos_only=config.sample.pos_only,
+        center_pos_mode=config.sample.center_pos_mode,
+        sample_num_atoms=config.sample.sample_num_atoms
+    )
+    result = {
+        'data': data,
+        'pred_ligand_pos': pred_pos,
+        'pred_ligand_v': pred_v,
+        'pred_ligand_pos_traj': pred_pos_traj,
+        'pred_ligand_v_traj': pred_v_traj,
+        'time': time_list
+    }
+    logger.info('Sample done!')
+
+    result_path = args.result_path
+    os.makedirs(result_path, exist_ok=True)
+    shutil.copyfile(args.config, os.path.join(result_path, 'sample.yml'))
+    torch.save(result, os.path.join(result_path, f'result_{args.data_id}.pt'))