Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,47 @@
 *.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
+*.pkl 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
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,234 @@
+---
+frameworks:
+- ""
+language:
+- en
+license: mit
+tags:
+- OneScience
+- protein-ligand molecular docking
+- generative
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">DiffDock</span>
+  </strong>
+</p>
+
+# Model Overview
+
+DiffDock is a diffusion model for protein-ligand molecular docking proposed by Corso et al. It supports single-complex molecular docking, batch molecular docking, evaluation, and related workflows.
+
+# Model Description
+
+DiffDock formulates docking as a generative modeling problem. Given a protein receptor and a small-molecule ligand, it jointly models the ligand's translational, rotational, and torsional degrees of freedom through diffusion processes on SO(3)/SE(3), and generates candidate binding conformations.
+
+Weights and datasets are not available at the moment. They will be uploaded to Hugging Face soon, and command-line downloads will be supported later.
+
+# Use Cases
+
+| Use case | Description |
+| :---: | :---: |
+| Score model training | Takes processed PDBBind or MOAD data as input and outputs a DiffDock score model checkpoint |
+| Single-complex molecular docking | Takes a protein PDB file and a ligand SMILES/SDF/MOL2 input, and outputs candidate ligand binding conformations in SDF format |
+| Batch molecular docking | Takes a CSV file containing proteins, ligands, and complex names, and outputs sampled conformations in batches |
+| Confidence rerank | Uses an additional confidence model to rank sampled conformations |
+| Dataset evaluation | Computes metrics such as RMSD for sampling results on a validation or test set, with optional GNINA integration |
+
+# 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**
+
+Common DiffDock dependencies include PyTorch, PyTorch Geometric, RDKit, OpenBabel, e3nn, torch-scatter, torch-cluster, NumPy, SciPy, tqdm, PyYAML, and others. GNINA energy minimization during dataset evaluation requires installing `gnina` separately.
+
+**Environment Checks**
+
+- NVIDIA GPU:
+
+```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 Package and Database
+
+```bash
+# If the model ID on the Hugging Face page uses different capitalization, use the actual published name.
+hf download --model OneScience-Sugon/diffdock --local-dir ./diffdock
+cd diffdock
+```
+
+### Training Weights and Datasets
+
+Weights and datasets are not available at the moment. They will be uploaded to Hugging Face soon, and command-line downloads will be supported later.
+
+### 3. Check Files in the Package
+
+The current package does not include pretrained weights. The command below should only show source code, configuration files, and example input files:
+
+```bash
+find . -maxdepth 3 -type f
+```
+
+If sampling or evaluation is required, first obtain a checkpoint through training, or place external DiffDock score/confidence weights in a local directory. Make sure the weight directory contains `model_parameters.yml`.
+
+### 4. Score Model Training
+
+Before training, update the data paths in `configs/training.yml` to local paths. Common fields include:
+
+- `data.pdbbind_dir`: processed PDBBind data directory.
+- `data.moad_dir`: processed MOAD data directory.
+- `data.split_train`: list of training-set complexes.
+- `data.split_val`: list of validation-set complexes.
+- `runtime.log_dir`: training output directory.
+
+Start training:
+
+```bash
+cd scripts
+bash train.sh
+```
+
+After training succeeds, the output directory usually contains:
+
+- `model_parameters.yml`: model architecture and training parameters.
+- `best_model.pt`: weights with the best validation loss.
+- `best_ema_model.pt`: EMA weights.
+- `best_inference_epoch_model.pt` or `best_ema_inference_epoch_model.pt`: weights saved when inference validation is enabled.
+- `last_model.pt`: training state from the final epoch.
+
+### 5. Molecular Docking Sampling
+
+Sampling requires a trained score model directory, for example:
+
+```text
+outputs/train/diffdock_cg_example/
+├── model_parameters.yml
+└── best_model.pt
+```
+
+Update the key fields in `configs/sampling.yml` to the actual paths:
+
+- `model.model_dir`: score model directory.
+- `model.ckpt`: score checkpoint file name.
+- `confidence.confidence_model_dir`: confidence model directory; set it to `null` when rerank is not used.
+- `input.protein_path`: protein PDB path. You can use `data/6o5u_protein_processed.pdb`.
+- `input.ligand_description`: SMILES string or ligand SDF/MOL2 path. You can use `data/6o5u_ligand.sdf`.
+- `runtime.out_dir`: sampling output directory.
+
+Start single-complex sampling:
+
+```bash
+cd scripts
+bash infer.sh
+```
+
+### 6. CSV Batch Sampling
+
+Set `input.protein_ligand_csv` in `configs/sampling.yml` to the CSV path, and set the single-complex fields to `null` as needed. The CSV is recommended to contain the following columns:
+
+```text
+complex_name,protein_path,ligand_description,protein_sequence
+```
+
+Here, `ligand_description` can be either a SMILES string or an SDF/MOL2 file path.
+
+### 7. Dataset Evaluation
+
+Before evaluation, update the dataset, model, and output directories in `configs/evaluate.yml` to actual paths:
+
+```bash
+python -m scripts.evaluate --config configs/evaluate.yml
+```
+
+If `gnina.gnina_minimize=true` is enabled, first make sure the following command can be run directly in the current environment:
+
+```bash
+gnina --help
+```
+
+### 8. Confidence Model Training
+
+The confidence training entry point is provided by the installed OneScience package. This Hugging Face package does not duplicate that module. After completing score model training, run the following in the OneScience environment:
+
+```bash
+python -m onescience.confidence.diffdock.confidence_train \
+    --original_model_dir outputs/train/diffdock_cg_example \
+    --data_dir /path/to/PDBBind_processed \
+    --split_train /path/to/splits/timesplit_no_lig_overlap_train \
+    --split_val /path/to/splits/timesplit_no_lig_overlap_val
+```
+
+# Data Preparation
+
+DiffDock training usually requires preprocessed PDBBind or MOAD data. It is recommended to organize the data under a unified root directory, for example:
+
+```text
+${ONESCIENCE_DATASETS_DIR}/diffdock/
+├── PDBBind_processed/
+├── MOAD_processed/
+└── splits/
+    ├── timesplit_no_lig_overlap_train
+    ├── timesplit_no_lig_overlap_val
+    └── timesplit_test
+```
+
+Split files should be plain text files, with one complex name per line. Each name must correspond to a subdirectory name in the data directory.
+
+# Official OneScience Information
+
+| 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 DiffDock code is licensed under the MIT License. This repository retains source attribution and is organized for OneScience Hugging Face automated runtime scenarios.
+- If you use DiffDock results in research, we recommend citing the original DiffDock and DiffDock-L papers, as well as relevant OneScience project information. Depending on the actual task, also add citations for datasets or tools such as PDBBind, MOAD, RDKit, OpenBabel, and GNINA.
+
+```bibtex
+@inproceedings{corso2023diffdock,
+  title={DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking},
+  author={Corso, Gabriele and St{\"a}rk, Hannes and Jing, Bowen and Barzilay, Regina and Jaakkola, Tommi},
+  booktitle={International Conference on Learning Representations},
+  year={2023}
+}
+```

__init__.py

@@ -0,0 +1 @@
+__all__ = []

configs/evaluate.yml

@@ -0,0 +1,101 @@
+# ============================================================================
+# DiffDock CGModel Dataset Evaluation Configuration
+#
+# IMPORTANT: Update all filesystem paths below to match your environment.
+#
+# Supported dataset values: pdbbind | moad | generalisation
+#
+# External dependencies (optional):
+#   gnina_minimize=true  → requires `gnina` executable on PATH
+#   RMSD computations     → requires `spyrmsd` (pip install spyrmsd)
+# ============================================================================
+
+runtime:
+  run_name: diffdock_eval_example
+  project: diffdock
+  device: auto
+  out_dir: examples/biosciences/diffdock/outputs/evaluate
+  wandb: false
+  num_cpu: null
+  restrict_cpu: false
+
+model:
+  model_dir: /public/home/liuyx19/modelscope/diffdock/outputs/train/diffdock_pdbbind_smoke100_val20_cpu
+  ckpt: best_model.pt
+  old_score_model: false  # must stay false; old_* score-model paths are intentionally not migrated
+  no_model: false
+  force_fixed_center_conv: false
+
+confidence:
+  confidence_model_dir: null
+  confidence_ckpt: best_model.pt
+  old_confidence_model: false  # must stay false; old confidence-model paths are intentionally not migrated
+
+data:
+  dataset: pdbbind  # supported: pdbbind | moad | generalisation; unsupported: pdbsidechain | distillation
+  cache_path: /public/home/liuyx19/modelscope/diffdock/cache
+  data_dir: /public/share/sugonhpcapp01/onestore/onedatasets/diffdock/datasets/PDBBind_processed
+  split_path: /public/share/sugonhpcapp01/onestore/onedatasets/diffdock/datasets/splits/timesplit_no_lig_overlap_val
+  split: val
+  limit_complexes: 0
+  num_workers: 1
+  chain_cutoff: null
+  protein_file: protein_processed
+  ligand_file: ligand
+  esm_embeddings_path: null
+  moad_esm_embeddings_sequences_path: null
+  not_knn_only_graph: false
+  include_miscellaneous_atoms: false
+  triple_training: false  # unsupported in the current migration; will fail fast if set true
+  unroll_clusters: false
+  remove_pdbbind: false
+  min_ligand_size: 0
+  max_receptor_size: null
+  remove_promiscuous_targets: null
+  matching_popsize: 40
+  matching_maxiter: 40
+
+sampling:
+  batch_size: 20
+  inference_steps: 20
+  actual_steps: null
+  samples_per_complex: 10
+  no_random: false
+  no_final_step_noise: true
+  ode: false
+  sigma_schedule: expbeta
+  inf_sched_alpha: 1.0
+  inf_sched_beta: 1.0
+  pocket_knowledge: false
+  no_random_pocket: false
+  pocket_tr_max: 3.0
+  pocket_cutoff: 5.0
+  different_schedules: false
+  resample_rdkit: false
+  skip_matching: false
+  initial_noise_std_proportion: -1.0
+  choose_residue: false
+  limit_failures: 5
+  tqdm: true
+  temp_sampling_tr: 1.0
+  temp_psi_tr: 0.0
+  temp_sigma_data_tr: 0.5
+  temp_sampling_rot: 1.0
+  temp_psi_rot: 0.0
+  temp_sigma_data_rot: 0.5
+  temp_sampling_tor: 1.0
+  temp_psi_tor: 0.0
+  temp_sigma_data_tor: 0.5
+
+outputs:
+  save_visualisation: false
+  save_complexes: false
+  complexes_save_path: null
+
+gnina:
+  gnina_minimize: false
+  gnina_path: gnina
+  gnina_full_dock: false
+  save_gnina_metrics: false
+  gnina_autobox_add: 4.0
+  gnina_poses_to_optimize: 1

configs/sampling.yml

@@ -0,0 +1,61 @@
+# ============================================================================
+# DiffDock CGModel Sampling Configuration
+#
+# IMPORTANT: Update all filesystem paths below to match your environment.
+#   - model_dir: must contain a CGModel checkpoint + model_parameters.yml
+#   - protein_path: path to receptor PDB
+#   - ligand_description: SMILES string or path to .sdf/.mol2
+#
+# CSV batch mode:
+#   Set protein_ligand_csv instead of the single-complex fields.
+#   CSV columns: complex_name, protein_path, ligand_description, protein_sequence
+#
+# Confidence rerank:
+#   Set confidence_model_dir to a CGModel-based confidence checkpoint dir.
+# ============================================================================
+
+runtime:
+  device: auto
+  loglevel: INFO
+  out_dir: examples/biosciences/diffdock/outputs/sample
+
+model:
+  model_dir: ${ONESCIENCE_DATASETS_DIR}/diffdock/score_model
+  ckpt: best_ema_inference_epoch_model.pt
+  old_score_model: false
+
+confidence:
+  confidence_model_dir: ${ONESCIENCE_DATASETS_DIR}/diffdock/confidence_model
+  confidence_ckpt: best_model_epoch75.pt
+  old_confidence_model: false
+
+input:
+  protein_ligand_csv: null
+  complex_name: 6o5u_test
+  protein_path: examples/biosciences/diffdock/data/6o5u_protein_processed.pdb
+  protein_sequence: null
+  ligand_description: examples/biosciences/diffdock/data/6o5u_ligand.sdf
+  lm_embeddings: null
+
+sampling:
+  samples_per_complex: 10
+  batch_size: 10
+  inference_steps: 20
+  actual_steps: null
+  sigma_schedule: expbeta
+  inf_sched_alpha: 1.0
+  inf_sched_beta: 1.0
+  no_random: false
+  no_final_step_noise: true
+  ode: false
+  choose_residue: false
+  initial_noise_std_proportion: 1.0
+  temp_sampling_tr: 1.0
+  temp_psi_tr: 0.0
+  temp_sigma_data_tr: 0.5
+  temp_sampling_rot: 1.0
+  temp_psi_rot: 0.0
+  temp_sigma_data_rot: 0.5
+  temp_sampling_tor: 1.0
+  temp_psi_tor: 0.0
+  temp_sigma_data_tor: 0.5

configs/training.yml

@@ -0,0 +1,150 @@
+# ============================================================================
+# DiffDock CGModel Training Configuration
+#
+# IMPORTANT: Update all filesystem paths below to match your environment.
+#   - cache_path, pdbbind_dir, moad_dir: dataset directories
+#   - split_train, split_val: text files listing complex names (one per line)
+#
+# Supported scope (current migration):
+#   dataset: pdbbind | moad | generalisation
+# Unsupported (will fail fast):
+#   dataset: pdbsidechain | distillation
+#   all_atoms: true  (AAModel not migrated)
+#   triple_training: true
+# ============================================================================
+
+runtime:
+  run_name: diffdock_cg_example
+  log_dir: examples/biosciences/diffdock/outputs/train
+  device: auto
+  seed: 0
+  cudnn_benchmark: false
+  wandb: false
+  project: diffdock
+
+data:
+  dataset: pdbbind  # supported: pdbbind | moad | generalisation
+  cache_path: ./cache
+  pdbbind_dir: ${ONESCIENCE_DATASETS_DIR}/diffdock/datasets/PDBBind_processed
+  moad_dir: ${ONESCIENCE_DATASETS_DIR}/diffdock/datasets/BindingMOAD_2020_processed
+  split_train: ${ONESCIENCE_DATASETS_DIR}/diffdock/datasets/splits/timesplit_no_lig_overlap_train
+  split_val: ${ONESCIENCE_DATASETS_DIR}/diffdock/datasets/splits/timesplit_no_lig_overlap_val
+  protein_file: protein_processed
+  limit_complexes: null
+  num_conformers: 1
+  num_workers: 1
+  num_dataloader_workers: 0
+  batch_size: 4
+  pin_memory: false
+  dataloader_drop_last: false
+  remove_hs: true
+  receptor_radius: 30
+  c_alpha_max_neighbors: 10
+  atom_radius: 5
+  atom_max_neighbors: 8
+  chain_cutoff: null
+  max_lig_size: null
+  matching_popsize: 20
+  matching_maxiter: 20
+  matching_tries: 1
+  not_knn_only_graph: false
+  include_miscellaneous_atoms: false
+  all_atoms: false  # must stay false in the current migration; only the CGModel main path is supported
+  triple_training: false  # unsupported in the current migration; will fail fast if set true
+  combined_training: false
+  double_val: false
+  train_multiplicity: 1
+  val_multiplicity: 1
+  max_receptor_size: null
+  remove_promiscuous_targets: null
+  min_ligand_size: 0
+  unroll_clusters: false
+  enforce_timesplit: false
+  crop_beyond: 20
+  moad_esm_embeddings_path: null
+  pdbbind_esm_embeddings_path: null
+  moad_esm_embeddings_sequences_path: null
+  esm_embeddings_model: null
+
+diffusion:
+  no_torsion: false
+  tr_sigma_min: 0.1
+  tr_sigma_max: 30.0
+  rot_sigma_min: 0.1
+  rot_sigma_max: 1.65
+  tor_sigma_min: 0.0314
+  tor_sigma_max: 3.14
+  sampling_alpha: 1.0
+  sampling_beta: 1.0
+  tr_weight: 0.33
+  rot_weight: 0.33
+  tor_weight: 0.33
+  backbone_loss_weight: 0.0
+  sidechain_loss_weight: 0.0
+
+model:
+  # old_score_model is intentionally unsupported in the current migration; training only covers the CGModel main path.
+  num_conv_layers: 2
+  max_radius: 5.0
+  scale_by_sigma: true
+  norm_by_sigma: false
+  ns: 16
+  nv: 4
+  distance_embed_dim: 32
+  cross_distance_embed_dim: 32
+  no_batch_norm: true
+  use_second_order_repr: false
+  cross_max_distance: 80
+  dynamic_max_cross: false
+  dropout: 0.0
+  smooth_edges: false
+  odd_parity: false
+  embedding_type: sinusoidal
+  sigma_embed_dim: 32
+  embedding_scale: 1000
+  no_aminoacid_identities: false
+  sh_lmax: 2
+  no_differentiate_convolutions: false
+  tp_weights_layers: 2
+  num_prot_emb_layers: 0
+  reduce_pseudoscalars: false
+  embed_also_ligand: true
+  depthwise_convolution: false
+  use_old_atom_encoder: false
+
+optimization:
+  n_epochs: 10
+  lr: 0.001
+  w_decay: 0.0
+  scheduler: null
+  scheduler_patience: 20
+  lr_start_factor: 0.001
+  warmup_dur: 4
+  use_ema: true
+  ema_rate: 0.999
+  restart_dir: null
+  restart_ckpt: last_model
+  restart_lr: null
+  pretrain_dir: null
+  pretrain_ckpt: null
+  save_model_freq: null
+
+validation:
+  test_sigma_intervals: false
+  inference_samples: 1
+  val_inference_freq: null
+  train_inference_freq: null
+  inference_steps: 20
+  num_inference_complexes: 20
+  inference_earlystop_metric: valinf_min_rmsds_lt2
+  inference_secondary_metric: null
+  inference_earlystop_goal: max
+
+# ----------------------------------------------------------------------------
+# NOTE: Confidence model training is a separate step after score-model training.
+# Run: python -m onescience.confidence.diffdock.confidence_train
+#            --original_model_dir <this_run_dir>
+#            --data_dir <pdbbind_dir> --split_train <...> --split_val <...>
+# There is no example-level confidence_train script; the module lives at
+# src/onescience/confidence/diffdock/confidence_train.py
+# ----------------------------------------------------------------------------

data/1a46_ligand.sdf

@@ -0,0 +1,179 @@
+1a46_ligand
+  -I-interpret- 
+
+ 85 88  0  0  0  0  0  0  0  0999 V2000
+   17.8330  -13.0420   21.6620 C   0  0  0  0  0
+   18.8870  -13.0710   20.5870 C   0  0  0  0  0
+   19.8510  -14.2200   21.1170 C   0  0  0  0  0
+   19.3270  -16.4440   22.1560 C   0  0  0  0  0
+   18.1340  -17.2300   22.7620 C   0  0  0  0  0
+   17.2230  -16.3290   23.5970 C   0  0  0  0  0
+   17.0320  -14.9230   23.0460 C   0  0  0  0  0
+   18.8520  -15.2420   21.4440 N   0  3  0  0  0
+   17.7750  -14.5090   22.0480 N   0  0  0  0  0
+   15.9850  -14.2900   23.3800 O   0  0  0  0  0
+   16.6380  -13.0610   20.7550 C   0  0  0  0  0
+   16.4620  -13.9620   19.8370 O   0  0  0  0  0
+   15.8090  -16.7300   23.6610 N   0  3  0  0  0
+   17.4150  -16.4170   25.1230 C   0  0  0  0  0
+   18.7640  -15.9840   25.5820 C   0  0  0  0  0
+   19.0510  -14.6340   25.7600 C   0  0  0  0  0
+   20.3910  -14.2520   26.0760 C   0  0  0  0  0
+   21.4290  -15.1780   26.2150 C   0  0  0  0  0
+   21.0990  -16.5480   26.0980 C   0  0  0  0  0
+   19.7890  -16.9510   25.7560 C   0  0  0  0  0
+   15.6470  -12.0890   20.7690 N   0  0  0  0  0
+   14.4940  -11.8920   19.9090 C   0  0  0  0  0
+   14.4960  -10.9450   18.7130 C   0  0  0  0  0
+   13.3800  -10.6840   18.0770 O   0  0  0  0  0
+   13.1950  -11.6150   20.6280 C   0  0  0  0  0
+   12.8670  -12.5040   21.7570 C   0  0  0  0  0
+   11.5610  -12.2200   22.4370 C   0  0  0  0  0
+   11.1700  -13.3510   23.3530 C   0  0  0  0  0
+   10.0380  -13.1110   24.2350 N   0  3  0  0  0
+   14.8040  -11.9210   16.4570 N   0  0  0  0  0
+   15.3450  -11.4350   17.5510 C   0  0  0  0  0
+   16.4740  -11.0890   17.7310 O   0  0  0  0  0
+   15.6510  -12.3330   15.3350 C   0  0  0  0  0
+   16.0390  -13.7960   15.2500 C   0  0  0  0  0
+   14.9560  -14.6030   14.5390 C   0  0  0  0  0
+   14.5990  -13.9990   13.1800 C   0  0  0  0  0
+   14.1680  -12.5610   13.3540 C   0  0  0  0  0
+   15.2770  -11.7400   13.9980 C   0  0  0  0  0
+   17.9332  -12.2994   22.4536 H   0  0  0  0  0
+   19.3882  -12.1140   20.4420 H   0  0  0  0  0
+   18.4882  -13.2617   19.5906 H   0  0  0  0  0
+   20.4926  -13.9283   21.9484 H   0  0  0  0  0
+   20.6127  -14.5392   20.4056 H   0  0  0  0  0
+   19.8508  -17.0880   21.4496 H   0  0  0  0  0
+   19.9921  -16.1358   22.9627 H   0  0  0  0  0
+   18.5327  -18.0092   23.4116 H   0  0  0  0  0
+   17.5467  -17.6450   21.9429 H   0  0  0  0  0
+   18.5389  -15.7277   20.6035 H   0  0  0  0  0
+   15.7428  -17.6818   24.0216 H   0  0  0  0  0
+   15.3044  -16.0949   24.2794 H   0  0  0  0  0
+   15.4029  -16.6903   22.7262 H   0  0  0  0  0
+   17.2937  -17.4623   25.4072 H   0  0  0  0  0
+   16.6848  -15.7509   25.5825 H   0  0  0  0  0
+   18.2682  -13.8821   25.6602 H   0  0  0  0  0
+   20.6133  -13.1939   26.2145 H   0  0  0  0  0
+   22.4528  -14.8565   26.4061 H   0  0  0  0  0
+   21.8654  -17.3029   26.2740 H   0  0  0  0  0
+   19.5640  -18.0094   25.6250 H   0  0  0  0  0
+   15.7457  -11.3948   21.5098 H   0  0  0  0  0
+   14.5905  -12.8910   19.4839 H   0  0  0  0  0
+   14.8425  -10.0689   19.2612 H   0  0  0  0  0
+   13.5584  -10.0751   17.3566 H   0  0  0  0  0
+   12.4050  -11.7585   19.8909 H   0  0  0  0  0
+   13.2901  -10.6141   21.0491 H   0  0  0  0  0
+   13.6465  -12.3595   22.5050 H   0  0  0  0  0
+   12.7942  -13.5124   21.3496 H   0  0  0  0  0
+   10.7892  -12.1043   21.6761 H   0  0  0  0  0
+   11.6663  -11.3113   23.0296 H   0  0  0  0  0
+   12.0278  -13.5229   24.0031 H   0  0  0  0  0
+   10.8774  -14.1769   22.7046 H   0  0  0  0  0
+    9.8690  -13.9413   24.8029 H   0  0  0  0  0
+   10.2441  -12.3181   24.8427 H   0  0  0  0  0
+    9.2101  -12.9059   23.6756 H   0  0  0  0  0
+   13.7904  -12.0118   16.3885 H   0  0  0  0  0
+   16.5871  -11.8550   15.6237 H   0  0  0  0  0
+   16.1623  -14.1864   16.2602 H   0  0  0  0  0
+   16.9681  -13.8812   14.6864 H   0  0  0  0  0
+   14.0613  -14.5994   15.1616 H   0  0  0  0  0
+   15.3317  -15.6133   14.3772 H   0  0  0  0  0
+   13.7819  -14.5683   12.7368 H   0  0  0  0  0
+   15.4725  -14.0364   12.5291 H   0  0  0  0  0
+   13.2893  -12.5323   13.9983 H   0  0  0  0  0
+   13.9420  -12.1402   12.3742 H   0  0  0  0  0
+   16.1510  -11.7459   13.3467 H   0  0  0  0  0
+   14.9268  -10.7183   14.1449 H   0  0  0  0  0
+  2  1  1  0  0  0
+  1  9  1  0  0  0
+  1 11  1  0  0  0
+  3  2  1  0  0  0
+  8  3  1  0  0  0
+  4  5  1  0  0  0
+  4  8  1  0  0  0
+  5  6  1  0  0  0
+  6  7  1  0  0  0
+  6 13  1  0  0  0
+  6 14  1  0  0  0
+  7  9  1  0  0  0
+  7 10  2  0  0  0
+  8  9  1  0  0  0
+ 11 12  2  0  0  0
+ 11 21  1  0  0  0
+ 14 15  1  0  0  0
+ 15 16  4  0  0  0
+ 15 20  4  0  0  0
+ 16 17  4  0  0  0
+ 17 18  4  0  0  0
+ 18 19  4  0  0  0
+ 19 20  4  0  0  0
+ 21 22  1  0  0  0
+ 22 23  1  0  0  0
+ 22 25  1  0  0  0
+ 23 24  1  0  0  0
+ 23 31  1  0  0  0
+ 25 26  1  0  0  0
+ 26 27  1  0  0  0
+ 27 28  1  0  0  0
+ 28 29  1  0  0  0
+ 31 30  1  0  0  0
+ 30 33  1  0  0  0
+ 31 32  2  0  0  0
+ 33 34  1  0  0  0
+ 33 38  1  0  0  0
+ 34 35  1  0  0  0
+ 35 36  1  0  0  0
+ 36 37  1  0  0  0
+ 37 38  1  0  0  0
+  1 39  1  0  0  0
+  2 40  1  0  0  0
+  2 41  1  0  0  0
+  3 42  1  0  0  0
+  3 43  1  0  0  0
+  4 44  1  0  0  0
+  4 45  1  0  0  0
+  5 46  1  0  0  0
+  5 47  1  0  0  0
+  8 48  1  0  0  0
+ 13 49  1  0  0  0
+ 13 50  1  0  0  0
+ 13 51  1  0  0  0
+ 14 52  1  0  0  0
+ 14 53  1  0  0  0
+ 16 54  1  0  0  0
+ 17 55  1  0  0  0
+ 18 56  1  0  0  0
+ 19 57  1  0  0  0
+ 20 58  1  0  0  0
+ 21 59  1  0  0  0
+ 22 60  1  0  0  0
+ 23 61  1  0  0  0
+ 24 62  1  0  0  0
+ 25 63  1  0  0  0
+ 25 64  1  0  0  0
+ 26 65  1  0  0  0
+ 26 66  1  0  0  0
+ 27 67  1  0  0  0
+ 27 68  1  0  0  0
+ 28 69  1  0  0  0
+ 28 70  1  0  0  0
+ 29 71  1  0  0  0
+ 29 72  1  0  0  0
+ 29 73  1  0  0  0
+ 30 74  1  0  0  0
+ 33 75  1  0  0  0
+ 34 76  1  0  0  0
+ 34 77  1  0  0  0
+ 35 78  1  0  0  0
+ 35 79  1  0  0  0
+ 36 80  1  0  0  0
+ 36 81  1  0  0  0
+ 37 82  1  0  0  0
+ 37 83  1  0  0  0
+ 38 84  1  0  0  0
+ 38 85  1  0  0  0
+M  END
+$$$$

data/6o5u_ligand.sdf

@@ -0,0 +1,154 @@
+6o5u_ligand
+  -I-interpret- 
+
+ 73 75  0  0  0  0  0  0  0  0999 V2000
+  108.6730  -18.4700   -7.8360 C   0  0  0  0  0
+  107.3110  -18.9740   -7.3680 C   0  0  0  0  0
+  107.4800  -20.3570   -6.7640 C   0  0  0  0  0
+  108.0820  -21.2810   -7.8150 C   0  0  0  0  0
+  109.3440  -20.6380   -8.3930 C   0  0  0  0  0
+  109.7960  -21.4700   -9.5890 C   0  0  0  0  0
+  112.5990  -15.8100   -5.7750 C   0  0  0  0  0
+  112.1040  -15.6660   -7.2080 C   0  0  0  0  0
+  110.7540  -16.3630   -7.2880 C   0  0  0  0  0
+  110.8240  -17.8460   -6.8990 C   0  0  0  0  0
+  111.5110  -18.0160   -5.5370 C   0  0  0  0  0
+  112.8430  -17.2770   -5.4680 C   0  0  0  0  0
+  112.8110  -13.7240   -8.4160 C   0  0  0  0  0
+  112.9750  -12.2640   -8.0080 C   0  0  0  0  0
+  111.6610  -11.5000   -8.1330 C   0  0  0  0  0
+  111.2120  -11.6530   -9.5690 C   0  0  0  0  0
+  111.1210  -13.1180   -9.9350 C   0  0  0  0  0
+  110.8030  -13.1430  -11.4220 C   0  0  0  0  0
+  110.9760  -20.8510  -10.1750 N   0  3  0  0  0
+  111.7470  -19.4450   -5.3220 N   0  3  0  0  0
+  113.8400  -15.0670   -5.6100 N   0  3  0  0  0
+  111.8410  -10.0810   -7.8770 N   0  3  0  0  0
+  109.1180  -19.3280   -8.8520 O   0  0  0  0  0
+  106.8170  -18.1010   -6.3680 O   0  0  0  0  0
+  106.1940  -20.8240   -6.3760 O   0  0  0  0  0
+  108.4250  -22.5370   -7.2480 O   0  0  0  0  0
+  109.5510  -18.4650   -6.7100 O   0  0  0  0  0
+  110.3630  -16.2700   -8.6430 O   0  0  0  0  0
+  111.9040  -14.2970   -7.4880 O   0  0  0  0  0
+  112.3420  -13.8090   -9.7580 O   0  0  0  0  0
+  113.3340  -12.2190   -6.6510 O   0  0  0  0  0
+  109.9130  -11.0980   -9.7100 O   0  0  0  0  0
+  110.7860  -14.4810  -11.8020 O   0  0  0  0  0
+  108.6303  -17.4570   -8.2362 H   0  0  0  0  0
+  106.6193  -19.0120   -8.2095 H   0  0  0  0  0
+  108.1401  -20.3313   -5.8970 H   0  0  0  0  0
+  107.3417  -21.4380   -8.5994 H   0  0  0  0  0
+  110.0905  -20.6009   -7.5997 H   0  0  0  0  0
+  108.9974  -21.5151  -10.3294 H   0  0  0  0  0
+  110.0340  -22.4829   -9.2644 H   0  0  0  0  0
+  111.8471  -15.4154   -5.0916 H   0  0  0  0  0
+  112.8209  -16.0908   -7.9107 H   0  0  0  0  0
+  110.0567  -15.8930   -6.5945 H   0  0  0  0  0
+  111.3562  -18.3049   -7.7322 H   0  0  0  0  0
+  110.8622  -17.5955   -4.7687 H   0  0  0  0  0
+  113.2743  -17.3813   -4.4724 H   0  0  0  0  0
+  113.5390  -17.6975   -6.1939 H   0  0  0  0  0
+  113.7563  -14.2662   -8.3935 H   0  0  0  0  0
+  113.7270  -11.8166   -8.6580 H   0  0  0  0  0
+  110.9429  -11.8905   -7.4119 H   0  0  0  0  0
+  111.9311  -11.1480  -10.2139 H   0  0  0  0  0
+  110.3778  -13.5999   -9.2998 H   0  0  0  0  0
+  111.5492  -12.5865  -11.9892 H   0  0  0  0  0
+  109.8423  -12.6689  -11.6228 H   0  0  0  0  0
+  111.2850  -21.4016  -10.9761 H   0  0  0  0  0
+  110.7465  -19.9049  -10.4794 H   0  0  0  0  0
+  111.7205  -20.8142   -9.4787 H   0  0  0  0  0
+  112.2040  -19.5810   -4.4203 H   0  0  0  0  0
+  110.8556  -19.9407   -5.3310 H   0  0  0  0  0
+  112.3452  -19.8052   -6.0655 H   0  0  0  0  0
+  114.1687  -15.1655   -4.6495 H   0  0  0  0  0
+  114.5431  -15.4352   -6.2507 H   0  0  0  0  0
+  113.6746  -14.0815   -5.8147 H   0  0  0  0  0
+  110.9443   -9.6037   -7.9692 H   0  0  0  0  0
+  112.2015   -9.9503   -6.9318 H   0  0  0  0  0
+  112.5031   -9.6965   -8.5509 H   0  0  0  0  0
+  106.7059  -17.2223   -6.7383 H   0  0  0  0  0
+  105.8266  -20.2335   -5.7143 H   0  0  0  0  0
+  107.6379  -22.9476   -6.8827 H   0  0  0  0  0
+  109.5104  -16.6957   -8.7588 H   0  0  0  0  0
+  113.4409  -11.3044   -6.3796 H   0  0  0  0  0
+  109.6229  -11.1925  -10.6202 H   0  0  0  0  0
+  110.5878  -14.5438  -12.7392 H   0  0  0  0  0
+  1  2  1  0  0  0
+  1 23  1  0  0  0
+ 27  1  1  0  0  0
+  2  3  1  0  0  0
+  2 24  1  0  0  0
+  3  4  1  0  0  0
+  3 25  1  0  0  0
+  5  4  1  0  0  0
+  4 26  1  0  0  0
+  5  6  1  0  0  0
+ 23  5  1  0  0  0
+  6 19  1  0  0  0
+  8  7  1  0  0  0
+  7 12  1  0  0  0
+  7 21  1  0  0  0
+  8  9  1  0  0  0
+ 29  8  1  0  0  0
+  9 10  1  0  0  0
+  9 28  1  0  0  0
+ 10 11  1  0  0  0
+ 10 27  1  0  0  0
+ 12 11  1  0  0  0
+ 11 20  1  0  0  0
+ 13 14  1  0  0  0
+ 13 29  1  0  0  0
+ 13 30  1  0  0  0
+ 14 15  1  0  0  0
+ 14 31  1  0  0  0
+ 15 16  1  0  0  0
+ 15 22  1  0  0  0
+ 16 17  1  0  0  0
+ 16 32  1  0  0  0
+ 17 18  1  0  0  0
+ 17 30  1  0  0  0
+ 18 33  1  0  0  0
+  1 34  1  0  0  0
+  2 35  1  0  0  0
+  3 36  1  0  0  0
+  4 37  1  0  0  0
+  5 38  1  0  0  0
+  6 39  1  0  0  0
+  6 40  1  0  0  0
+  7 41  1  0  0  0
+  8 42  1  0  0  0
+  9 43  1  0  0  0
+ 10 44  1  0  0  0
+ 11 45  1  0  0  0
+ 12 46  1  0  0  0
+ 12 47  1  0  0  0
+ 13 48  1  0  0  0
+ 14 49  1  0  0  0
+ 15 50  1  0  0  0
+ 16 51  1  0  0  0
+ 17 52  1  0  0  0
+ 18 53  1  0  0  0
+ 18 54  1  0  0  0
+ 19 55  1  0  0  0
+ 19 56  1  0  0  0
+ 19 57  1  0  0  0
+ 20 58  1  0  0  0
+ 20 59  1  0  0  0
+ 20 60  1  0  0  0
+ 21 61  1  0  0  0
+ 21 62  1  0  0  0
+ 21 63  1  0  0  0
+ 22 64  1  0  0  0
+ 22 65  1  0  0  0
+ 22 66  1  0  0  0
+ 24 67  1  0  0  0
+ 25 68  1  0  0  0
+ 26 69  1  0  0  0
+ 28 70  1  0  0  0
+ 31 71  1  0  0  0
+ 32 72  1  0  0  0
+ 33 73  1  0  0  0
+M  END
+$$$$

models/__init__.py

@@ -0,0 +1,34 @@
+__all__ = [
+    "CGModel",
+    "build_score_model",
+    "load_model_args",
+    "load_score_model",
+    "model_uses_lm_embeddings",
+]
+
+
+def __getattr__(name):
+    if name == "CGModel":
+        from .cg_model import CGModel
+
+        return CGModel
+    if name in {
+        "build_score_model",
+        "load_model_args",
+        "load_score_model",
+        "model_uses_lm_embeddings",
+    }:
+        from .score_wrapper import (
+            build_score_model,
+            load_model_args,
+            load_score_model,
+            model_uses_lm_embeddings,
+        )
+
+        return {
+            "build_score_model": build_score_model,
+            "load_model_args": load_model_args,
+            "load_score_model": load_score_model,
+            "model_uses_lm_embeddings": model_uses_lm_embeddings,
+        }[name]
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

models/aa_model.py

@@ -0,0 +1,721 @@
+from e3nn import o3
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch_cluster import radius, radius_graph
+from torch_geometric.utils import subgraph
+from torch_scatter import scatter_mean
+import numpy as np
+
+from onescience.datapipes.diffdock.process_mols import lig_feature_dims, rec_residue_feature_dims, rec_atom_feature_dims
+from onescience.utils.diffdock import so3, torus
+
+from .layers import GaussianSmearing, AtomEncoder
+from .tensor_layers import get_irrep_seq, TensorProductConvLayer
+
+try:
+    from esm.pretrained import load_model_and_alphabet
+except ImportError:
+    load_model_and_alphabet = None
+
+AGGREGATORS = {"mean": lambda x: torch.mean(x, dim=1),
+               "max": lambda x: torch.max(x, dim=1)[0],
+               "min": lambda x: torch.min(x, dim=1)[0],
+               "std": lambda x: torch.std(x, dim=1)}
+
+
+class AAModel(torch.nn.Module):
+    def __init__(self, t_to_sigma, device, timestep_emb_func, in_lig_edge_features=4, sigma_embed_dim=32, sh_lmax=2,
+                 ns=16, nv=4, num_conv_layers=2, lig_max_radius=5, rec_max_radius=30, cross_max_distance=250,
+                 center_max_distance=30, distance_embed_dim=32, cross_distance_embed_dim=32, no_torsion=False,
+                 scale_by_sigma=True, norm_by_sigma=True, use_second_order_repr=False, batch_norm=True,
+                 dynamic_max_cross=False, dropout=0.0, smooth_edges=False, odd_parity=False,
+                 separate_noise_schedule=False, lm_embedding_type=None, confidence_mode=False,
+                 confidence_dropout=0, confidence_no_batchnorm = False,
+                 asyncronous_noise_schedule=False, affinity_prediction=False, parallel=1,
+                 parallel_aggregators="mean max min std", num_confidence_outputs=1, atom_num_confidence_outputs=1, fixed_center_conv=False,
+                 no_aminoacid_identities=False, include_miscellaneous_atoms=False,
+                 differentiate_convolutions=True, tp_weights_layers=2, num_prot_emb_layers=0,
+                 reduce_pseudoscalars=False, embed_also_ligand=False, atom_confidence=False, sidechain_pred=False,
+                 depthwise_convolution=False, crop_beyond=None):
+        super(AAModel, self).__init__()
+        assert (not no_aminoacid_identities) or (lm_embedding_type is None), "no language model emb without identities"
+        assert not sidechain_pred, "sidechain prediction not implemented/makes sense for all atom model"
+        assert not depthwise_convolution, "depthwise convolution not implemented for all atom model"
+        if parallel > 1: assert affinity_prediction
+
+        self.t_to_sigma = t_to_sigma
+        self.in_lig_edge_features = in_lig_edge_features
+        sigma_embed_dim *= (3 if separate_noise_schedule else 1)
+        self.sigma_embed_dim = sigma_embed_dim
+        self.lig_max_radius = lig_max_radius
+        self.rec_max_radius = rec_max_radius
+        self.cross_max_distance = cross_max_distance
+        self.dynamic_max_cross = dynamic_max_cross
+        self.center_max_distance = center_max_distance
+        self.distance_embed_dim = distance_embed_dim
+        self.cross_distance_embed_dim = cross_distance_embed_dim
+        self.sh_irreps = o3.Irreps.spherical_harmonics(lmax=sh_lmax)
+        self.ns, self.nv = ns, nv
+        self.scale_by_sigma = scale_by_sigma
+        self.norm_by_sigma = norm_by_sigma
+        self.device = device
+        self.no_torsion = no_torsion
+        self.smooth_edges = smooth_edges
+        self.odd_parity = odd_parity
+        self.num_conv_layers = num_conv_layers
+        self.timestep_emb_func = timestep_emb_func
+        self.separate_noise_schedule = separate_noise_schedule
+        self.confidence_mode = confidence_mode
+        self.num_conv_layers = num_conv_layers
+        self.num_prot_emb_layers = num_prot_emb_layers
+        self.asyncronous_noise_schedule = asyncronous_noise_schedule
+        self.affinity_prediction = affinity_prediction
+        self.parallel, self.parallel_aggregators = parallel, parallel_aggregators.split(' ')
+        self.fixed_center_conv = fixed_center_conv
+        self.no_aminoacid_identities = no_aminoacid_identities
+        self.differentiate_convolutions = differentiate_convolutions
+        self.reduce_pseudoscalars = reduce_pseudoscalars
+        self.atom_confidence = atom_confidence
+        self.atom_num_confidence_outputs = atom_num_confidence_outputs
+        self.crop_beyond = crop_beyond
+
+        self.lm_embedding_type = lm_embedding_type
+        if lm_embedding_type is None:
+            lm_embedding_dim = 0
+        elif lm_embedding_type == "precomputed":
+            lm_embedding_dim=1280
+        else:
+            if load_model_and_alphabet is None:
+                raise ImportError(
+                    "esm is required when lm_embedding_type is not None and not 'precomputed'."
+                )
+            lm, alphabet = load_model_and_alphabet(lm_embedding_type)
+            self.batch_converter = alphabet.get_batch_converter()
+            lm.lm_head = torch.nn.Identity()
+            lm.contact_head = torch.nn.Identity()
+            lm_embedding_dim = lm.embed_dim
+            self.lm = lm
+
+        # embedding layers
+        atom_encoder_class = AtomEncoder
+        self.lig_node_embedding = atom_encoder_class(emb_dim=ns, feature_dims=lig_feature_dims, sigma_embed_dim=sigma_embed_dim)
+        self.lig_edge_embedding = nn.Sequential(nn.Linear(in_lig_edge_features + sigma_embed_dim + distance_embed_dim, ns),nn.ReLU(),nn.Dropout(dropout),nn.Linear(ns, ns))
+
+        self.rec_sigma_embedding = nn.Sequential(nn.Linear(sigma_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout), nn.Linear(ns, ns))
+        self.rec_node_embedding = atom_encoder_class(emb_dim=ns, feature_dims=rec_residue_feature_dims, sigma_embed_dim=0, lm_embedding_dim=lm_embedding_dim)
+        self.rec_edge_embedding = nn.Sequential(nn.Linear(distance_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout), nn.Linear(ns, ns))
+        self.atom_node_embedding = atom_encoder_class(emb_dim=ns, feature_dims=rec_atom_feature_dims, sigma_embed_dim=0)
+        self.atom_edge_embedding = nn.Sequential(nn.Linear(distance_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout), nn.Linear(ns, ns))
+
+        self.lr_edge_embedding = nn.Sequential(nn.Linear(sigma_embed_dim + cross_distance_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout),nn.Linear(ns, ns))
+        self.ar_edge_embedding = nn.Sequential(nn.Linear(distance_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout),nn.Linear(ns, ns))
+        self.la_edge_embedding = nn.Sequential(nn.Linear(sigma_embed_dim + cross_distance_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout),nn.Linear(ns, ns))
+
+        self.lig_distance_expansion = GaussianSmearing(0.0, lig_max_radius, distance_embed_dim)
+        self.rec_distance_expansion = GaussianSmearing(0.0, rec_max_radius, distance_embed_dim)
+        self.cross_distance_expansion = GaussianSmearing(0.0, cross_max_distance, cross_distance_embed_dim)
+
+        irrep_seq = get_irrep_seq(ns, nv, use_second_order_repr, reduce_pseudoscalars)
+        assert not include_miscellaneous_atoms, "currently not supported"
+
+        rec_emb_layers = []
+        for i in range(num_prot_emb_layers):
+            in_irreps = irrep_seq[min(i, len(irrep_seq) - 1)]
+            out_irreps = irrep_seq[min(i + 1, len(irrep_seq) - 1)]
+            layer = TensorProductConvLayer(
+                in_irreps=in_irreps,
+                sh_irreps=self.sh_irreps,
+                out_irreps=out_irreps,
+                n_edge_features=3 * ns,
+                hidden_features=3 * ns,
+                residual=True,
+                batch_norm=batch_norm,
+                dropout=dropout,
+                faster=sh_lmax == 1 and not use_second_order_repr,
+                tp_weights_layers=tp_weights_layers,
+                edge_groups=1 if not differentiate_convolutions else 4,
+            )
+            rec_emb_layers.append(layer)
+        self.rec_emb_layers = nn.ModuleList(rec_emb_layers)
+
+        self.embed_also_ligand = embed_also_ligand
+        if embed_also_ligand:
+            lig_emb_layers = []
+            for i in range(num_prot_emb_layers):
+                in_irreps = irrep_seq[min(i, len(irrep_seq) - 1)]
+                out_irreps = irrep_seq[min(i + 1, len(irrep_seq) - 1)]
+                layer = TensorProductConvLayer(
+                    in_irreps=in_irreps,
+                    sh_irreps=self.sh_irreps,
+                    out_irreps=out_irreps,
+                    n_edge_features=3 * ns,
+                    hidden_features=3 * ns,
+                    residual=True,
+                    batch_norm=batch_norm,
+                    dropout=dropout,
+                    faster=sh_lmax == 1 and not use_second_order_repr,
+                    tp_weights_layers=tp_weights_layers,
+                    edge_groups=1,
+                )
+                lig_emb_layers.append(layer)
+            self.lig_emb_layers = nn.ModuleList(lig_emb_layers)
+
+        # convolutional layers
+        conv_layers = []
+        for i in range(num_prot_emb_layers, num_prot_emb_layers + num_conv_layers):
+            in_irreps = irrep_seq[min(i, len(irrep_seq) - 1)]
+            out_irreps = irrep_seq[min(i + 1, len(irrep_seq) - 1)]
+            layer = TensorProductConvLayer(
+                in_irreps=in_irreps,
+                sh_irreps=self.sh_irreps,
+                out_irreps=out_irreps,
+                n_edge_features=3 * ns,
+                hidden_features=3 * ns,
+                residual=True,
+                batch_norm=batch_norm,
+                dropout=dropout,
+                faster=sh_lmax == 1 and not use_second_order_repr,
+                tp_weights_layers=tp_weights_layers,
+                edge_groups=1 if not differentiate_convolutions else (3 if i == num_prot_emb_layers + num_conv_layers - 1 else 9),
+            )
+            conv_layers.append(layer)
+        self.conv_layers = nn.ModuleList(conv_layers)
+
+        # confidence and affinity prediction layers
+        if self.confidence_mode:
+            if self.affinity_prediction:
+                if self.parallel > 1:
+                    output_confidence_dim = 1 + ns
+                else:
+                    output_confidence_dim = num_confidence_outputs + 1
+            else:
+                output_confidence_dim = num_confidence_outputs
+
+            input_size = ns + (nv if reduce_pseudoscalars else ns) if num_conv_layers + num_prot_emb_layers >= 3 else ns
+
+            if self.atom_confidence:
+                self.atom_confidence_predictor = nn.Sequential(
+                    nn.Linear(input_size, ns),
+                    nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(),
+                    nn.ReLU(),
+                    nn.Dropout(confidence_dropout),
+                    nn.Linear(ns, ns),
+                    nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(),
+                    nn.ReLU(),
+                    nn.Dropout(confidence_dropout),
+                    nn.Linear(ns, atom_num_confidence_outputs + ns)
+                )
+                input_size = ns
+
+            self.confidence_predictor = nn.Sequential(
+                nn.Linear(input_size, ns),
+                nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(),
+                nn.ReLU(),
+                nn.Dropout(confidence_dropout),
+                nn.Linear(ns, ns),
+                nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(),
+                nn.ReLU(),
+                nn.Dropout(confidence_dropout),
+                nn.Linear(ns, output_confidence_dim)
+            )
+
+            if self.parallel > 1:
+                self.affinity_predictor = nn.Sequential(
+                    nn.Linear(len(self.parallel_aggregators) * ns, ns),
+                    nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(),
+                    nn.ReLU(),
+                    nn.Dropout(confidence_dropout),
+                    nn.Linear(ns, ns),
+                    nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(),
+                    nn.ReLU(),
+                    nn.Dropout(confidence_dropout),
+                    nn.Linear(ns, 1)
+                )
+
+        else:
+            # convolution for translational and rotational scores
+            self.center_distance_expansion = GaussianSmearing(0.0, center_max_distance, distance_embed_dim)
+            self.center_edge_embedding = nn.Sequential(
+                nn.Linear(distance_embed_dim + sigma_embed_dim, ns),
+                nn.ReLU(),
+                nn.Dropout(dropout),
+                nn.Linear(ns, ns)
+            )
+
+            self.final_conv = TensorProductConvLayer(
+                in_irreps=self.conv_layers[-1].out_irreps,
+                sh_irreps=self.sh_irreps,
+                out_irreps=f'2x1o + 2x1e' if not self.odd_parity else '1x1o + 1x1e',
+                n_edge_features=2 * ns,
+                residual=False,
+                dropout=dropout,
+                batch_norm=batch_norm
+            )
+
+            self.tr_final_layer = nn.Sequential(nn.Linear(1 + sigma_embed_dim, ns),nn.Dropout(dropout), nn.ReLU(), nn.Linear(ns, 1))
+            self.rot_final_layer = nn.Sequential(nn.Linear(1 + sigma_embed_dim, ns),nn.Dropout(dropout), nn.ReLU(), nn.Linear(ns, 1))
+
+            if not no_torsion:
+                # convolution for torsional score
+                self.final_edge_embedding = nn.Sequential(
+                    nn.Linear(distance_embed_dim, ns),
+                    nn.ReLU(),
+                    nn.Dropout(dropout),
+                    nn.Linear(ns, ns)
+                )
+                self.final_tp_tor = o3.FullTensorProduct(self.sh_irreps, "2e")
+                self.tor_bond_conv = TensorProductConvLayer(
+                    in_irreps=self.conv_layers[-1].out_irreps,
+                    sh_irreps=self.final_tp_tor.irreps_out,
+                    out_irreps=f'{ns}x0o + {ns}x0e' if not self.odd_parity else f'{ns}x0o',
+                    n_edge_features=3 * ns,
+                    residual=False,
+                    dropout=dropout,
+                    batch_norm=batch_norm
+                )
+                self.tor_final_layer = nn.Sequential(
+                    nn.Linear(2 * ns if not self.odd_parity else ns, ns, bias=False),
+                    nn.Tanh(),
+                    nn.Dropout(dropout),
+                    nn.Linear(ns, 1, bias=False)
+                )
+
+    @staticmethod
+    def _resolve_edge_store(data, primary_key, fallback_key):
+        edge_types = getattr(data, "edge_types", ())
+        if primary_key in edge_types:
+            return data[primary_key]
+        if fallback_key in edge_types:
+            return data[fallback_key]
+        try:
+            return data[primary_key]
+        except Exception:
+            return data[fallback_key]
+
+    def _ligand_edge_store(self, data):
+        return self._resolve_edge_store(
+            data,
+            ("ligand", "ligand"),
+            ("ligand", "lig_bond", "ligand"),
+        )
+
+    def _receptor_edge_store(self, data):
+        return self._resolve_edge_store(
+            data,
+            ("receptor", "receptor"),
+            ("receptor", "rec_contact", "receptor"),
+        )
+
+    def _atom_edge_store(self, data):
+        return self._resolve_edge_store(
+            data,
+            ("atom", "atom"),
+            ("atom", "atom_contact", "atom"),
+        )
+
+    def _atom_receptor_edge_store(self, data):
+        return self._resolve_edge_store(
+            data,
+            ("atom", "receptor"),
+            ("atom", "atom_rec_contact", "receptor"),
+        )
+
+    def embedding(self, data):
+        receptor_edge_store = self._receptor_edge_store(data)
+        atom_edge_store = self._atom_edge_store(data)
+        atom_receptor_edge_store = self._atom_receptor_edge_store(data)
+        if not hasattr(data['receptor'], "rec_node_attr"):
+            if self.lm_embedding_type not in [None, 'precomputed']:
+                sequences = [s for l in data['receptor'].sequence for s in l]
+                if isinstance(sequences[0], list):
+                    sequences = [s for l in sequences for s in l]
+                sequences = [(i, s) for i, s in enumerate(sequences)]
+                batch_labels, batch_strs, batch_tokens = self.batch_converter(sequences)
+                out = self.lm(batch_tokens.to(data['receptor'].x.device), repr_layers=[self.lm.num_layers], return_contacts=False)
+                rec_lm_emb = torch.cat([t[:len(sequences[i][1])] for i, t in enumerate(out['representations'][self.lm.num_layers])], dim=0)
+                data['receptor'].x = torch.cat([data['receptor'].x, rec_lm_emb], dim=-1)
+
+            rec_node_attr, rec_edge_attr, rec_edge_sh, rec_edge_weight = self.build_rec_conv_graph(data)
+            rec_node_attr = self.rec_node_embedding(rec_node_attr)
+            rec_edge_attr = self.rec_edge_embedding(rec_edge_attr)
+
+            atom_node_attr, atom_edge_attr, atom_edge_sh, atom_edge_weight = self.build_atom_conv_graph(data)
+            atom_node_attr = self.atom_node_embedding(atom_node_attr)
+            atom_edge_attr = self.atom_edge_embedding(atom_edge_attr)
+
+            ar_edge_attr, ar_edge_sh, ar_edge_weight = self.build_cross_rec_conv_graph(data)
+            ar_edge_attr = self.ar_edge_embedding(ar_edge_attr)
+
+            rec_edge_index = receptor_edge_store.edge_index.clone()
+            atom_edge_index = atom_edge_store.edge_index.clone()
+            ar_edge_index = atom_receptor_edge_store.edge_index.clone()
+
+            node_attr = torch.cat([rec_node_attr, atom_node_attr], dim=0)
+            ar_edge_index[0] = ar_edge_index[0] + len(rec_node_attr)
+            edge_index = torch.cat([rec_edge_index, ar_edge_index, atom_edge_index + len(rec_node_attr), torch.flip(ar_edge_index, dims=[0])], dim=1)
+            edge_attr = torch.cat([rec_edge_attr, ar_edge_attr, atom_edge_attr, ar_edge_attr], dim=0)
+            edge_sh = torch.cat([rec_edge_sh, ar_edge_sh, atom_edge_sh, ar_edge_sh], dim=0)
+            edge_weight = torch.cat([rec_edge_weight, ar_edge_weight, atom_edge_weight, ar_edge_weight], dim=0) \
+                if torch.is_tensor(rec_edge_weight) else torch.ones((len(edge_index[0]), 1), device=edge_index.device)
+            s1, s2, s3 = len(rec_edge_index[0]), len(rec_edge_index[0]) + len(ar_edge_index[0]), len(rec_edge_index[0]) + len(ar_edge_index[0]) + len(atom_edge_index[0])
+
+            for l in range(len(self.rec_emb_layers)):
+                edge_attr_ = torch.cat(
+                    [edge_attr, node_attr[edge_index[0], :self.ns], node_attr[edge_index[1], :self.ns]], -1)
+                if self.differentiate_convolutions: edge_attr_ = [edge_attr_[:s1], edge_attr_[s1:s2], edge_attr_[s2:s3], edge_attr_[s3:]]
+                node_attr = self.rec_emb_layers[l](node_attr, edge_index, edge_attr_, edge_sh, edge_weight=edge_weight)
+
+
+            data['receptor'].rec_node_attr = node_attr[:len(rec_node_attr)]
+            receptor_edge_store.rec_edge_attr = rec_edge_attr
+            receptor_edge_store.edge_sh = rec_edge_sh
+            receptor_edge_store.edge_weight = rec_edge_weight
+
+            data['atom'].atom_node_attr = node_attr[len(rec_node_attr):]
+            atom_edge_store.atom_edge_attr = atom_edge_attr
+            atom_edge_store.edge_sh = atom_edge_sh
+            atom_edge_store.edge_weight = atom_edge_weight
+
+            atom_receptor_edge_store.edge_attr = ar_edge_attr
+            atom_receptor_edge_store.edge_sh = ar_edge_sh
+            atom_receptor_edge_store.edge_weight = ar_edge_weight
+
+        # receptor embedding
+        rec_sigma_emb = self.rec_sigma_embedding(self.timestep_emb_func(data.complex_t['tr']))
+        rec_node_attr = data['receptor'].rec_node_attr + 0
+        rec_node_attr[:, :self.ns] = rec_node_attr[:, :self.ns] + rec_sigma_emb[data['receptor'].batch]
+        rec_edge_attr = receptor_edge_store.rec_edge_attr + rec_sigma_emb[data['receptor'].batch[receptor_edge_store.edge_index[0]]]
+
+        # atom embedding
+        atom_node_attr = data['atom'].atom_node_attr + 0
+        atom_node_attr[:, :self.ns] = atom_node_attr[:, :self.ns] + rec_sigma_emb[data['atom'].batch]
+        atom_edge_attr = atom_edge_store.atom_edge_attr + rec_sigma_emb[data['atom'].batch[atom_edge_store.edge_index[0]]]
+
+        # atom-receptor embedding
+        ar_edge_attr = atom_receptor_edge_store.edge_attr + rec_sigma_emb[data['atom'].batch[atom_receptor_edge_store.edge_index[0]]]
+
+        # ligand embedding
+        lig_node_attr, lig_edge_index, lig_edge_attr, lig_edge_sh, lig_edge_weight = self.build_lig_conv_graph(data)
+        lig_node_attr = self.lig_node_embedding(lig_node_attr)
+        lig_edge_attr = self.lig_edge_embedding(lig_edge_attr)
+
+        if self.embed_also_ligand:
+            for l in range(len(self.lig_emb_layers)):
+                edge_attr_ = torch.cat([lig_edge_attr, lig_node_attr[lig_edge_index[0], :self.ns], lig_node_attr[lig_edge_index[1], :self.ns]], -1)
+                lig_node_attr = self.lig_emb_layers[l](lig_node_attr, lig_edge_index, edge_attr_, lig_edge_sh, edge_weight=lig_edge_weight)
+
+        else:
+            lig_node_attr = F.pad(lig_node_attr, (0, rec_node_attr.shape[-1] - lig_node_attr.shape[-1]))
+
+        return lig_node_attr, lig_edge_index, lig_edge_attr, lig_edge_sh, lig_edge_weight, \
+               rec_node_attr, receptor_edge_store.edge_index, rec_edge_attr, receptor_edge_store.edge_sh, receptor_edge_store.edge_weight, \
+               atom_node_attr, atom_edge_store.edge_index, atom_edge_attr, atom_edge_store.edge_sh, atom_edge_store.edge_weight, \
+               atom_receptor_edge_store.edge_index, ar_edge_attr, atom_receptor_edge_store.edge_sh, atom_receptor_edge_store.edge_weight
+
+    def forward(self, data):
+        if self.crop_beyond is not None:
+            # TODO missing filtering atoms
+            raise NotImplementedError
+            ligand_pos = data['ligand'].pos
+            receptor_pos = data['receptor'].pos
+            residues_to_keep = torch.any(torch.sum((ligand_pos.unsqueeze(0) - receptor_pos.unsqueeze(1)) ** 2, -1) < self.crop_beyond ** 2, dim=1)
+
+            data['receptor'].pos = data['receptor'].pos[residues_to_keep]
+            data['receptor'].x = data['receptor'].x[residues_to_keep]
+            data['receptor'].side_chain_vecs = data['receptor'].side_chain_vecs[residues_to_keep]
+            data['receptor', 'rec_contact', 'receptor'].edge_index = subgraph(residues_to_keep, data['receptor', 'rec_contact', 'receptor'].edge_index, relabel_nodes=True)[0]
+
+        if self.no_aminoacid_identities:
+            data['receptor'].x = data['receptor'].x * 0
+
+        if not self.confidence_mode:
+            tr_sigma, rot_sigma, tor_sigma = self.t_to_sigma(*[data.complex_t[noise_type] for noise_type in ['tr', 'rot', 'tor']])
+        else:
+            tr_sigma, rot_sigma, tor_sigma = [data.complex_t[noise_type] for noise_type in ['tr', 'rot', 'tor']]
+
+        lig_node_attr, lig_edge_index, lig_edge_attr, lig_edge_sh, lig_edge_weight, rec_node_attr, \
+            rec_edge_index, rec_edge_attr, rec_edge_sh, rec_edge_weight,\
+            atom_node_attr, atom_edge_index, atom_edge_attr, atom_edge_sh, atom_edge_weight, \
+            ar_edge_index, ar_edge_attr, ar_edge_sh, ar_edge_weight = self.embedding(data)
+
+        # build lig cross graph
+        cross_cutoff = (tr_sigma * 3 + 20).unsqueeze(1) if self.dynamic_max_cross else self.cross_max_distance
+        lr_edge_index, lr_edge_attr, lr_edge_sh, lr_edge_weight, la_edge_index, la_edge_attr, \
+            la_edge_sh, la_edge_weight = self.build_cross_lig_conv_graph(data, cross_cutoff)
+        lr_edge_attr= self.lr_edge_embedding(lr_edge_attr)
+        la_edge_attr = self.la_edge_embedding(la_edge_attr)
+
+        n_lig, n_rec = len(lig_node_attr), len(rec_node_attr)
+
+        node_attr = torch.cat([lig_node_attr, rec_node_attr, atom_node_attr], dim=0)
+        rec_edge_index, atom_edge_index, lr_edge_index, la_edge_index, ar_edge_index = rec_edge_index.clone(), atom_edge_index.clone(), lr_edge_index.clone(), la_edge_index.clone(), ar_edge_index.clone()
+        rec_edge_index[0], rec_edge_index[1] = rec_edge_index[0] + n_lig, rec_edge_index[1] + n_lig
+        atom_edge_index[0], atom_edge_index[1] = atom_edge_index[0] + n_lig + n_rec, atom_edge_index[1] + n_lig + n_rec
+        lr_edge_index[1] = lr_edge_index[1] + n_lig
+        la_edge_index[1] = la_edge_index[1] + n_lig + n_rec
+        ar_edge_index[0], ar_edge_index[1] = ar_edge_index[0] + n_lig + n_rec, ar_edge_index[1] + n_lig
+
+        edge_index = torch.cat([lig_edge_index, lr_edge_index, la_edge_index, rec_edge_index,
+                                torch.flip(lr_edge_index, dims=[0]), torch.flip(ar_edge_index, dims=[0]),
+                                atom_edge_index, torch.flip(la_edge_index, dims=[0]), ar_edge_index], dim=1)
+        edge_attr = torch.cat([lig_edge_attr, lr_edge_attr, la_edge_attr, rec_edge_attr, lr_edge_attr,
+                               ar_edge_attr, atom_edge_attr, la_edge_attr, ar_edge_attr], dim=0)
+        edge_sh = torch.cat([lig_edge_sh, lr_edge_sh, la_edge_sh, rec_edge_sh, lr_edge_sh, ar_edge_sh,
+                             atom_edge_sh, la_edge_sh, ar_edge_sh], dim=0)
+        edge_weight = torch.cat([lig_edge_weight, lr_edge_weight, la_edge_weight, rec_edge_weight, lr_edge_weight,
+                                 ar_edge_weight, atom_edge_weight, la_edge_weight, ar_edge_weight],
+                                dim=0) if torch.is_tensor(lig_edge_weight) else torch.ones((len(edge_index[0]), 1),
+                                                                                           device=edge_index.device)
+        s1, s2, s3, s4, s5, s6, s7, s8, _ = tuple(np.cumsum(list(map(len, [lig_edge_attr, lr_edge_attr, la_edge_attr,
+                                                                           rec_edge_attr, lr_edge_attr, ar_edge_attr, atom_edge_attr, la_edge_attr, ar_edge_attr]))).tolist())
+
+        for l in range(len(self.conv_layers)):
+            if l < len(self.conv_layers) - 1:
+                edge_attr_ = torch.cat([edge_attr, node_attr[edge_index[0], :self.ns], node_attr[edge_index[1], :self.ns]], -1)
+                if self.differentiate_convolutions: edge_attr_ = [edge_attr_[:s1], edge_attr_[s1:s2], edge_attr_[s2:s3], edge_attr_[s3:s4],
+                                                                  edge_attr_[s4:s5], edge_attr_[s5:s6], edge_attr_[s6:s7], edge_attr_[s7:s8], edge_attr_[s8:]]
+                node_attr = self.conv_layers[l](node_attr, edge_index, edge_attr_, edge_sh, edge_weight=edge_weight)
+            else:
+                edge_attr_ = torch.cat([edge_attr[:s3], node_attr[edge_index[0, :s3], :self.ns], node_attr[edge_index[1, :s3], :self.ns]], -1)
+                if self.differentiate_convolutions: edge_attr_ = [edge_attr_[:s1], edge_attr_[s1:s2], edge_attr_[s2:s3]]
+                node_attr = self.conv_layers[l](node_attr, edge_index[:, :s3], edge_attr_, edge_sh[:s3], edge_weight=edge_weight[:s3])
+
+        lig_node_attr = node_attr[:len(lig_node_attr)]
+
+        # confidence and affinity prediction
+        if self.confidence_mode:
+            scalar_lig_attr = torch.cat([lig_node_attr[:,:self.ns], lig_node_attr[:,-(self.nv if self.reduce_pseudoscalars else self.ns):] ], dim=1) \
+                if self.num_conv_layers + self.num_prot_emb_layers >= 3 else lig_node_attr[:,:self.ns]
+
+            if self.atom_confidence:
+                scalar_lig_attr = self.atom_confidence_predictor(scalar_lig_attr)
+                atom_confidence = scalar_lig_attr[:, :self.atom_num_confidence_outputs]
+                scalar_lig_attr = scalar_lig_attr[:, self.atom_num_confidence_outputs:]
+            else:
+                atom_confidence = torch.zeros((len(lig_node_attr),), device=lig_node_attr.device)
+
+            confidence = self.confidence_predictor(scatter_mean(scalar_lig_attr, data['ligand'].batch, dim=0)).squeeze(dim=-1)
+
+            if self.parallel > 1:
+                confidence, affinity = confidence[:, 0], confidence[:, 1:]
+                confidence = confidence.reshape(data.num_graphs, self.parallel)
+                affinity = affinity.reshape(data.num_graphs, self.parallel, -1)
+                affinity = torch.cat([AGGREGATORS[agg](affinity) for agg in self.parallel_aggregators], dim=-1)
+                affinity = self.affinity_predictor(affinity).squeeze(dim=-1)
+                confidence = confidence, affinity
+            return confidence, atom_confidence
+        assert self.parallel == 1
+
+        # compute translational and rotational score vectors
+        center_edge_index, center_edge_attr, center_edge_sh = self.build_center_conv_graph(data)
+        center_edge_attr = self.center_edge_embedding(center_edge_attr)
+        if self.fixed_center_conv:
+            center_edge_attr = torch.cat([center_edge_attr, lig_node_attr[center_edge_index[1], :self.ns]], -1)
+        else:
+            center_edge_attr = torch.cat([center_edge_attr, lig_node_attr[center_edge_index[0], :self.ns]], -1)
+        global_pred = self.final_conv(lig_node_attr, center_edge_index, center_edge_attr, center_edge_sh, out_nodes=data.num_graphs)
+
+        tr_pred = global_pred[:, :3] + (global_pred[:, 6:9] if not self.odd_parity else 0)
+        rot_pred = global_pred[:, 3:6] + (global_pred[:, 9:] if not self.odd_parity else 0)
+
+        if self.separate_noise_schedule:
+            data.graph_sigma_emb = torch.cat([self.timestep_emb_func(data.complex_t[noise_type]) for noise_type in ['tr', 'rot', 'tor']], dim=1)
+        elif self.asyncronous_noise_schedule:
+            data.graph_sigma_emb = self.timestep_emb_func(data.complex_t['t'])
+        else:  # tr rot and tor noise is all the same in this case
+            data.graph_sigma_emb = self.timestep_emb_func(data.complex_t['tr'])
+
+        # adjust the magniture of the score vectors
+        tr_norm = torch.linalg.vector_norm(tr_pred, dim=1).unsqueeze(1)
+        tr_pred = tr_pred / tr_norm * self.tr_final_layer(torch.cat([tr_norm, data.graph_sigma_emb], dim=1))
+
+        rot_norm = torch.linalg.vector_norm(rot_pred, dim=1).unsqueeze(1)
+        rot_pred = rot_pred / rot_norm * self.rot_final_layer(torch.cat([rot_norm, data.graph_sigma_emb], dim=1))
+
+        if self.scale_by_sigma:
+            tr_pred = tr_pred / tr_sigma.unsqueeze(1)
+            rot_pred = rot_pred * so3.score_norm(rot_sigma.cpu()).unsqueeze(1).to(data['ligand'].x.device)
+
+        if self.no_torsion or data['ligand'].edge_mask.sum() == 0: return tr_pred, rot_pred, torch.empty(0,device=self.device), None
+
+        # torsional components
+        tor_bonds, tor_edge_index, tor_edge_attr, tor_edge_sh, tor_edge_weight = self.build_bond_conv_graph(data)
+        tor_bond_vec = data['ligand'].pos[tor_bonds[1]] - data['ligand'].pos[tor_bonds[0]]
+        tor_bond_attr = lig_node_attr[tor_bonds[0]] + lig_node_attr[tor_bonds[1]]
+
+        tor_bonds_sh = o3.spherical_harmonics("2e", tor_bond_vec, normalize=True, normalization='component')
+        tor_edge_sh = self.final_tp_tor(tor_edge_sh, tor_bonds_sh[tor_edge_index[0]])
+
+        tor_edge_attr = torch.cat([tor_edge_attr, lig_node_attr[tor_edge_index[1], :self.ns],
+                                   tor_bond_attr[tor_edge_index[0], :self.ns]], -1)
+        tor_pred = self.tor_bond_conv(lig_node_attr, tor_edge_index, tor_edge_attr, tor_edge_sh,
+                                  out_nodes=data['ligand'].edge_mask.sum(), reduce='mean', edge_weight=tor_edge_weight)
+        tor_pred = self.tor_final_layer(tor_pred).squeeze(1)
+        ligand_edge_store = self._ligand_edge_store(data)
+        edge_sigma = tor_sigma[data['ligand'].batch][ligand_edge_store.edge_index[0]][data['ligand'].edge_mask]
+
+        if self.scale_by_sigma:
+            tor_pred = tor_pred * torch.sqrt(torch.tensor(torus.score_norm(edge_sigma.cpu().numpy())).float()
+                                             .to(data['ligand'].x.device))
+        return tr_pred, rot_pred, tor_pred, None
+
+    def get_edge_weight(self, edge_vec, max_norm):
+        if self.smooth_edges:
+            normalised_norm = torch.clip(edge_vec.norm(dim=-1) * np.pi / max_norm, max=np.pi)
+            return 0.5 * (torch.cos(normalised_norm) + 1.0).unsqueeze(-1)
+        return 1.0
+
+    def build_lig_conv_graph(self, data):
+        # build the graph between ligand atoms
+        if self.separate_noise_schedule:
+            data['ligand'].node_sigma_emb = torch.cat(
+                [self.timestep_emb_func(data['ligand'].node_t[noise_type]) for noise_type in ['tr', 'rot', 'tor']],
+                dim=1)
+        elif self.asyncronous_noise_schedule:
+            data['ligand'].node_sigma_emb = self.timestep_emb_func(data['ligand'].node_t['t'])
+        else:
+            data['ligand'].node_sigma_emb = self.timestep_emb_func(
+                data['ligand'].node_t['tr'])  # tr rot and tor noise is all the same
+
+        if self.parallel == 1:
+            radius_edges = radius_graph(data['ligand'].pos, self.lig_max_radius, data['ligand'].batch)
+        else:
+            batches = torch.zeros(data.num_graphs, device=data['ligand'].x.device).long()
+            batches = batches.index_add(0, data['ligand'].batch, torch.ones(len(data['ligand'].batch), device=data['ligand'].x.device).long())
+            outer_batches = data.num_graphs
+            b = [torch.ones(batches[i].item()//self.parallel, device=data['ligand'].x.device).long() * (self.parallel * i + j)
+                 for i in range(outer_batches) for j in range(self.parallel)]
+            data['ligand'].batch_parallel = torch.cat(b)
+            radius_edges = radius_graph(data['ligand'].pos, self.lig_max_radius, data['ligand'].batch_parallel)
+        ligand_edge_store = self._ligand_edge_store(data)
+        edge_index = torch.cat([ligand_edge_store.edge_index, radius_edges], 1).long()
+        edge_attr = torch.cat([
+            ligand_edge_store.edge_attr,
+            torch.zeros(radius_edges.shape[-1], self.in_lig_edge_features, device=data['ligand'].x.device)
+        ], 0)
+
+        edge_sigma_emb = data['ligand'].node_sigma_emb[edge_index[0].long()]
+        edge_attr = torch.cat([edge_attr, edge_sigma_emb], 1)
+        node_attr = torch.cat([data['ligand'].x, data['ligand'].node_sigma_emb], 1)
+
+        src, dst = edge_index
+        edge_vec = data['ligand'].pos[dst.long()] - data['ligand'].pos[src.long()]
+        edge_length_emb = self.lig_distance_expansion(edge_vec.norm(dim=-1))
+
+        edge_attr = torch.cat([edge_attr, edge_length_emb], 1)
+        edge_sh = o3.spherical_harmonics(self.sh_irreps, edge_vec, normalize=True, normalization='component')
+        edge_weight = self.get_edge_weight(edge_vec, self.lig_max_radius)
+
+        return node_attr, edge_index, edge_attr, edge_sh, edge_weight
+
+    def build_rec_conv_graph(self, data):
+        # build the graph between receptor residues
+        node_attr = data['receptor'].x
+
+        # this assumes the edges were already created in preprocessing since protein's structure is fixed
+        edge_index = self._receptor_edge_store(data).edge_index
+        src, dst = edge_index
+        edge_vec = data['receptor'].pos[dst.long()] - data['receptor'].pos[src.long()]
+
+        edge_attr = self.rec_distance_expansion(edge_vec.norm(dim=-1))
+        edge_sh = o3.spherical_harmonics(self.sh_irreps, edge_vec, normalize=True, normalization='component')
+        edge_weight = self.get_edge_weight(edge_vec, self.rec_max_radius)
+
+        return node_attr, edge_attr, edge_sh, edge_weight
+
+    def build_atom_conv_graph(self, data):
+        # build the graph between receptor atoms
+        node_attr = data['atom'].x
+
+        # this assumes the edges were already created in preprocessing since protein's structure is fixed
+        edge_index = self._atom_edge_store(data).edge_index
+        src, dst = edge_index
+        edge_vec = data['atom'].pos[dst.long()] - data['atom'].pos[src.long()]
+
+        edge_attr = self.lig_distance_expansion(edge_vec.norm(dim=-1))
+        edge_sh = o3.spherical_harmonics(self.sh_irreps, edge_vec, normalize=True, normalization='component')
+        edge_weight = self.get_edge_weight(edge_vec, self.lig_max_radius)
+
+        return node_attr, edge_attr, edge_sh, edge_weight
+
+    def build_cross_lig_conv_graph(self, data, lr_cross_distance_cutoff):
+        # build the cross edges between ligand atoms and receptor residues + atoms
+
+        # LIGAND to RECEPTOR
+        if torch.is_tensor(lr_cross_distance_cutoff):
+            # different cutoff for every graph
+            lr_edge_index = radius(data['receptor'].pos / lr_cross_distance_cutoff[data['receptor'].batch],
+                                data['ligand'].pos / lr_cross_distance_cutoff[data['ligand'].batch], 1,
+                                data['receptor'].batch, data['ligand'].batch, max_num_neighbors=10000)
+        else:
+            lr_edge_index = radius(data['receptor'].pos, data['ligand'].pos, lr_cross_distance_cutoff,
+                            data['receptor'].batch, data['ligand'].batch, max_num_neighbors=10000)
+
+        lr_edge_vec = data['receptor'].pos[lr_edge_index[1].long()] - data['ligand'].pos[lr_edge_index[0].long()]
+        lr_edge_length_emb = self.cross_distance_expansion(lr_edge_vec.norm(dim=-1))
+        lr_edge_sigma_emb = data['ligand'].node_sigma_emb[lr_edge_index[0].long()]
+        lr_edge_attr = torch.cat([lr_edge_sigma_emb, lr_edge_length_emb], 1)
+        lr_edge_sh = o3.spherical_harmonics(self.sh_irreps, lr_edge_vec, normalize=True, normalization='component')
+
+        cutoff_d = lr_cross_distance_cutoff[data['ligand'].batch[lr_edge_index[0]]].squeeze() \
+            if torch.is_tensor(lr_cross_distance_cutoff) else lr_cross_distance_cutoff
+        lr_edge_weight = self.get_edge_weight(lr_edge_vec, cutoff_d)
+
+        # LIGAND to ATOM
+        la_edge_index = radius(data['atom'].pos, data['ligand'].pos, self.lig_max_radius,
+                               data['atom'].batch, data['ligand'].batch, max_num_neighbors=10000)
+
+        la_edge_vec = data['atom'].pos[la_edge_index[1].long()] - data['ligand'].pos[la_edge_index[0].long()]
+        la_edge_length_emb = self.lig_distance_expansion(la_edge_vec.norm(dim=-1))
+        la_edge_sigma_emb = data['ligand'].node_sigma_emb[la_edge_index[0].long()]
+        la_edge_attr = torch.cat([la_edge_sigma_emb, la_edge_length_emb], 1)
+        la_edge_sh = o3.spherical_harmonics(self.sh_irreps, la_edge_vec, normalize=True, normalization='component')
+        la_edge_weight = self.get_edge_weight(la_edge_vec, self.lig_max_radius)
+
+        return lr_edge_index, lr_edge_attr, lr_edge_sh, lr_edge_weight, la_edge_index, la_edge_attr, \
+               la_edge_sh, la_edge_weight
+
+    def build_cross_rec_conv_graph(self, data):
+        # build the cross edges between ligan atoms, receptor residues and receptor atoms
+
+        # ATOM to RECEPTOR
+        ar_edge_index = self._atom_receptor_edge_store(data).edge_index
+        ar_edge_vec = data['receptor'].pos[ar_edge_index[1].long()] - data['atom'].pos[ar_edge_index[0].long()]
+        ar_edge_attr = self.rec_distance_expansion(ar_edge_vec.norm(dim=-1))
+        ar_edge_sh = o3.spherical_harmonics(self.sh_irreps, ar_edge_vec, normalize=True, normalization='component')
+        ar_edge_weight = 1
+
+        return ar_edge_attr, ar_edge_sh, ar_edge_weight
+
+    def build_center_conv_graph(self, data):
+        # build the filter for the convolution of the center with the ligand atoms
+        # for translational and rotational score
+        edge_index = torch.cat([data['ligand'].batch.unsqueeze(0), torch.arange(len(data['ligand'].batch)).to(data['ligand'].x.device).unsqueeze(0)], dim=0)
+
+        center_pos, count = torch.zeros((data.num_graphs, 3)).to(data['ligand'].x.device), torch.zeros((data.num_graphs, 3)).to(data['ligand'].x.device)
+        center_pos.index_add_(0, index=data['ligand'].batch, source=data['ligand'].pos)
+        center_pos = center_pos / torch.bincount(data['ligand'].batch).unsqueeze(1)
+
+        edge_vec = data['ligand'].pos[edge_index[1]] - center_pos[edge_index[0]]
+        edge_attr = self.center_distance_expansion(edge_vec.norm(dim=-1))
+        edge_sigma_emb = data['ligand'].node_sigma_emb[edge_index[1].long()]
+        edge_attr = torch.cat([edge_attr, edge_sigma_emb], 1)
+        edge_sh = o3.spherical_harmonics(self.sh_irreps, edge_vec, normalize=True, normalization='component')
+        return edge_index, edge_attr, edge_sh
+
+    def build_bond_conv_graph(self, data):
+        # build graph for the pseudotorque layer
+        bonds = self._ligand_edge_store(data).edge_index[:, data['ligand'].edge_mask].long()
+        bond_pos = (data['ligand'].pos[bonds[0]] + data['ligand'].pos[bonds[1]]) / 2
+        bond_batch = data['ligand'].batch[bonds[0]]
+        edge_index = radius(data['ligand'].pos, bond_pos, self.lig_max_radius, batch_x=data['ligand'].batch, batch_y=bond_batch)
+
+        edge_vec = data['ligand'].pos[edge_index[1]] - bond_pos[edge_index[0]]
+        edge_attr = self.lig_distance_expansion(edge_vec.norm(dim=-1))
+
+        edge_attr = self.final_edge_embedding(edge_attr)
+        edge_sh = o3.spherical_harmonics(self.sh_irreps, edge_vec, normalize=True, normalization='component')
+        edge_weight = self.get_edge_weight(edge_vec, self.lig_max_radius)
+
+        return bonds, edge_index, edge_attr, edge_sh, edge_weight

models/cg_model.py

@@ -0,0 +1,1070 @@
+import math
+import os
+
+import numpy as np
+import torch
+from e3nn import o3
+from e3nn.o3 import Linear
+from torch import nn
+from torch.nn import functional as F
+from torch_cluster import radius, radius_graph
+from torch_scatter import scatter, scatter_mean
+
+from onescience.datapipes.diffdock.process_mols import (
+    lig_feature_dims,
+    rec_atom_feature_dims,
+    rec_residue_feature_dims,
+)
+from onescience.utils.diffdock import so3, torus
+
+from .layers import AtomEncoder, GaussianSmearing
+from .tensor_layers import TensorProductConvLayer, get_irrep_seq
+
+
+def _env_flag(name):
+    return os.environ.get(name, "").lower() in {"1", "true", "yes", "on"}
+
+try:
+    from esm.pretrained import load_model_and_alphabet
+except ImportError:
+    load_model_and_alphabet = None
+
+
+class CGModel(torch.nn.Module):
+    def __init__(
+        self,
+        t_to_sigma,
+        device,
+        timestep_emb_func,
+        in_lig_edge_features=4,
+        sigma_embed_dim=32,
+        sh_lmax=2,
+        ns=16,
+        nv=4,
+        num_conv_layers=2,
+        lig_max_radius=5,
+        rec_max_radius=30,
+        cross_max_distance=250,
+        center_max_distance=30,
+        distance_embed_dim=32,
+        cross_distance_embed_dim=32,
+        no_torsion=False,
+        scale_by_sigma=True,
+        norm_by_sigma=True,
+        use_second_order_repr=False,
+        batch_norm=True,
+        dynamic_max_cross=False,
+        dropout=0.0,
+        smooth_edges=False,
+        odd_parity=False,
+        separate_noise_schedule=False,
+        lm_embedding_type=None,
+        confidence_mode=False,
+        confidence_dropout=0,
+        confidence_no_batchnorm=False,
+        asyncronous_noise_schedule=False,
+        affinity_prediction=False,
+        parallel=1,
+        parallel_aggregators="mean max min std",
+        num_confidence_outputs=1,
+        atom_num_confidence_outputs=1,
+        fixed_center_conv=False,
+        no_aminoacid_identities=False,
+        include_miscellaneous_atoms=False,
+        differentiate_convolutions=True,
+        tp_weights_layers=2,
+        num_prot_emb_layers=0,
+        reduce_pseudoscalars=False,
+        embed_also_ligand=False,
+        atom_confidence=False,
+        sidechain_pred=False,
+        depthwise_convolution=False,
+    ):
+        super(CGModel, self).__init__()
+        assert parallel == 1, "not implemented"
+        assert (not no_aminoacid_identities) or (
+            lm_embedding_type is None
+        ), "no language model emb without identities"
+        self.t_to_sigma = t_to_sigma
+        self.in_lig_edge_features = in_lig_edge_features
+        sigma_embed_dim *= 3 if separate_noise_schedule else 1
+        self.sigma_embed_dim = sigma_embed_dim
+        self.lig_max_radius = lig_max_radius
+        self.rec_max_radius = rec_max_radius
+        self.include_miscellaneous_atoms = include_miscellaneous_atoms
+        self.cross_max_distance = cross_max_distance
+        self.dynamic_max_cross = dynamic_max_cross
+        self.center_max_distance = center_max_distance
+        self.distance_embed_dim = distance_embed_dim
+        self.cross_distance_embed_dim = cross_distance_embed_dim
+        self.sh_irreps = o3.Irreps.spherical_harmonics(lmax=sh_lmax)
+        self.ns, self.nv = ns, nv
+        self.scale_by_sigma = scale_by_sigma
+        self.norm_by_sigma = norm_by_sigma
+        self.device = device
+        self.no_torsion = no_torsion
+        self.smooth_edges = smooth_edges
+        self.odd_parity = odd_parity
+        self.timestep_emb_func = timestep_emb_func
+        self.separate_noise_schedule = separate_noise_schedule
+        self.confidence_mode = confidence_mode
+        self.num_conv_layers = num_conv_layers
+        self.num_prot_emb_layers = num_prot_emb_layers
+        self.asyncronous_noise_schedule = asyncronous_noise_schedule
+        self.affinity_prediction = affinity_prediction
+        self.fixed_center_conv = fixed_center_conv
+        self.no_aminoacid_identities = no_aminoacid_identities
+        self.differentiate_convolutions = differentiate_convolutions
+        self.reduce_pseudoscalars = reduce_pseudoscalars
+        self.atom_confidence = atom_confidence
+        self.atom_num_confidence_outputs = atom_num_confidence_outputs
+        self.sidechain_pred = sidechain_pred
+
+        self.lm_embedding_type = lm_embedding_type
+        if lm_embedding_type is None:
+            lm_embedding_dim = 0
+        elif lm_embedding_type == "precomputed":
+            lm_embedding_dim = 1280
+        else:
+            if load_model_and_alphabet is None:
+                raise ImportError(
+                    "esm is required when lm_embedding_type is not None and not 'precomputed'."
+                )
+            lm, alphabet = load_model_and_alphabet(lm_embedding_type)
+            self.batch_converter = alphabet.get_batch_converter()
+            lm.lm_head = torch.nn.Identity()
+            lm.contact_head = torch.nn.Identity()
+            lm_embedding_dim = lm.embed_dim
+            self.lm = lm
+
+        atom_encoder_class = AtomEncoder
+
+        self.lig_node_embedding = atom_encoder_class(
+            emb_dim=ns,
+            feature_dims=lig_feature_dims,
+            sigma_embed_dim=sigma_embed_dim,
+        )
+        self.lig_edge_embedding = nn.Sequential(
+            nn.Linear(in_lig_edge_features + sigma_embed_dim + distance_embed_dim, ns),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(ns, ns),
+        )
+
+        self.rec_node_embedding = atom_encoder_class(
+            emb_dim=ns,
+            feature_dims=rec_residue_feature_dims,
+            sigma_embed_dim=0,
+            lm_embedding_dim=lm_embedding_dim,
+        )
+        self.rec_edge_embedding = nn.Sequential(
+            nn.Linear(distance_embed_dim, ns),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(ns, ns),
+        )
+        self.rec_sigma_embedding = nn.Sequential(
+            nn.Linear(sigma_embed_dim, ns),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(ns, ns),
+        )
+
+        if self.include_miscellaneous_atoms:
+            self.misc_atom_node_embedding = atom_encoder_class(
+                emb_dim=ns,
+                feature_dims=rec_atom_feature_dims,
+                sigma_embed_dim=sigma_embed_dim,
+            )
+            self.misc_atom_edge_embedding = nn.Sequential(
+                nn.Linear(sigma_embed_dim + distance_embed_dim, ns),
+                nn.ReLU(),
+                nn.Dropout(dropout),
+                nn.Linear(ns, ns),
+            )
+            self.ar_edge_embedding = nn.Sequential(
+                nn.Linear(sigma_embed_dim + distance_embed_dim, ns),
+                nn.ReLU(),
+                nn.Dropout(dropout),
+                nn.Linear(ns, ns),
+            )
+            self.la_edge_embedding = nn.Sequential(
+                nn.Linear(sigma_embed_dim + cross_distance_embed_dim, ns),
+                nn.ReLU(),
+                nn.Dropout(dropout),
+                nn.Linear(ns, ns),
+            )
+
+        self.cross_edge_embedding = nn.Sequential(
+            nn.Linear(sigma_embed_dim + cross_distance_embed_dim, ns),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(ns, ns),
+        )
+
+        self.lig_distance_expansion = GaussianSmearing(0.0, lig_max_radius, distance_embed_dim)
+        self.rec_distance_expansion = GaussianSmearing(0.0, rec_max_radius, distance_embed_dim)
+        self.cross_distance_expansion = GaussianSmearing(
+            0.0, cross_max_distance, cross_distance_embed_dim
+        )
+
+        irrep_seq = get_irrep_seq(ns, nv, use_second_order_repr, reduce_pseudoscalars)
+
+        assert not self.include_miscellaneous_atoms, "currently not supported"
+
+        rec_emb_layers = []
+        for i in range(num_prot_emb_layers):
+            in_irreps = irrep_seq[min(i, len(irrep_seq) - 1)]
+            out_irreps = irrep_seq[min(i + 1, len(irrep_seq) - 1)]
+            layer = TensorProductConvLayer(
+                in_irreps=in_irreps,
+                sh_irreps=self.sh_irreps,
+                out_irreps=out_irreps,
+                n_edge_features=3 * ns,
+                hidden_features=3 * ns,
+                residual=True,
+                batch_norm=batch_norm,
+                dropout=dropout,
+                faster=sh_lmax == 1 and not use_second_order_repr,
+                tp_weights_layers=tp_weights_layers,
+                edge_groups=1,
+                depthwise=depthwise_convolution,
+            )
+            rec_emb_layers.append(layer)
+        self.rec_emb_layers = nn.ModuleList(rec_emb_layers)
+
+        self.embed_also_ligand = embed_also_ligand
+        if embed_also_ligand:
+            lig_emb_layers = []
+            for i in range(num_prot_emb_layers):
+                in_irreps = irrep_seq[min(i, len(irrep_seq) - 1)]
+                out_irreps = irrep_seq[min(i + 1, len(irrep_seq) - 1)]
+                layer = TensorProductConvLayer(
+                    in_irreps=in_irreps,
+                    sh_irreps=self.sh_irreps,
+                    out_irreps=out_irreps,
+                    n_edge_features=3 * ns,
+                    hidden_features=3 * ns,
+                    residual=True,
+                    batch_norm=batch_norm,
+                    dropout=dropout,
+                    faster=sh_lmax == 1 and not use_second_order_repr,
+                    tp_weights_layers=tp_weights_layers,
+                    edge_groups=1,
+                    depthwise=depthwise_convolution,
+                )
+                lig_emb_layers.append(layer)
+            self.lig_emb_layers = nn.ModuleList(lig_emb_layers)
+
+        conv_layers = []
+        for i in range(num_prot_emb_layers, num_prot_emb_layers + num_conv_layers):
+            in_irreps = irrep_seq[min(i, len(irrep_seq) - 1)]
+            out_irreps = irrep_seq[min(i + 1, len(irrep_seq) - 1)]
+            layer = TensorProductConvLayer(
+                in_irreps=in_irreps,
+                sh_irreps=self.sh_irreps,
+                out_irreps=out_irreps,
+                n_edge_features=3 * ns,
+                hidden_features=3 * ns,
+                residual=True,
+                batch_norm=batch_norm,
+                dropout=dropout,
+                faster=sh_lmax == 1 and not use_second_order_repr,
+                tp_weights_layers=tp_weights_layers,
+                edge_groups=(
+                    1
+                    if not differentiate_convolutions
+                    else (2 if i == num_prot_emb_layers + num_conv_layers - 1 else 4)
+                ),
+                depthwise=depthwise_convolution,
+            )
+            conv_layers.append(layer)
+        self.conv_layers = nn.ModuleList(conv_layers)
+
+        if sidechain_pred:
+            self.sidechain_predictor = Linear(
+                irreps_in=irrep_seq[min(num_prot_emb_layers + num_conv_layers, len(irrep_seq) - 1)],
+                irreps_out="4x0e + 2x1e + 4x0o + 2x1o",
+                internal_weights=True,
+                shared_weights=True,
+            )
+
+        if self.confidence_mode:
+            input_size = (
+                ns + (nv if reduce_pseudoscalars else ns)
+                if num_conv_layers + num_prot_emb_layers >= 3
+                else ns
+            )
+
+            if self.atom_confidence:
+                self.atom_confidence_predictor = nn.Sequential(
+                    nn.Linear(input_size, ns),
+                    nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(),
+                    nn.ReLU(),
+                    nn.Dropout(confidence_dropout),
+                    nn.Linear(ns, ns),
+                    nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(),
+                    nn.ReLU(),
+                    nn.Dropout(confidence_dropout),
+                    nn.Linear(ns, atom_num_confidence_outputs + ns),
+                )
+                input_size = ns
+
+            self.confidence_predictor = nn.Sequential(
+                nn.Linear(input_size, ns),
+                nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(),
+                nn.ReLU(),
+                nn.Dropout(confidence_dropout),
+                nn.Linear(ns, ns),
+                nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(),
+                nn.ReLU(),
+                nn.Dropout(confidence_dropout),
+                nn.Linear(ns, num_confidence_outputs + (1 if self.affinity_prediction else 0)),
+            )
+
+        else:
+            self.center_distance_expansion = GaussianSmearing(
+                0.0, center_max_distance, distance_embed_dim
+            )
+            self.center_edge_embedding = nn.Sequential(
+                nn.Linear(distance_embed_dim + sigma_embed_dim, ns),
+                nn.ReLU(),
+                nn.Dropout(dropout),
+                nn.Linear(ns, ns),
+            )
+
+            self.final_conv = TensorProductConvLayer(
+                in_irreps=self.conv_layers[-1].out_irreps,
+                sh_irreps=self.sh_irreps,
+                out_irreps="2x1o + 2x1e" if not self.odd_parity else "1x1o + 1x1e",
+                n_edge_features=2 * ns,
+                residual=False,
+                dropout=dropout,
+                batch_norm=batch_norm,
+            )
+            self.tr_final_layer = nn.Sequential(
+                nn.Linear(1 + sigma_embed_dim, ns),
+                nn.Dropout(dropout),
+                nn.ReLU(),
+                nn.Linear(ns, 1),
+            )
+            self.rot_final_layer = nn.Sequential(
+                nn.Linear(1 + sigma_embed_dim, ns),
+                nn.Dropout(dropout),
+                nn.ReLU(),
+                nn.Linear(ns, 1),
+            )
+
+            if not no_torsion:
+                self.final_edge_embedding = nn.Sequential(
+                    nn.Linear(distance_embed_dim, ns),
+                    nn.ReLU(),
+                    nn.Dropout(dropout),
+                    nn.Linear(ns, ns),
+                )
+                self.final_tp_tor = o3.FullTensorProduct(self.sh_irreps, "2e")
+                self.tor_bond_conv = TensorProductConvLayer(
+                    in_irreps=self.conv_layers[-1].out_irreps,
+                    sh_irreps=self.final_tp_tor.irreps_out,
+                    out_irreps=f"{ns}x0o + {ns}x0e" if not self.odd_parity else f"{ns}x0o",
+                    n_edge_features=3 * ns,
+                    residual=False,
+                    dropout=dropout,
+                    batch_norm=batch_norm,
+                )
+                self.tor_final_layer = nn.Sequential(
+                    nn.Linear(2 * ns if not self.odd_parity else ns, ns, bias=False),
+                    nn.Tanh(),
+                    nn.Dropout(dropout),
+                    nn.Linear(ns, 1, bias=False),
+                )
+
+    @staticmethod
+    def _resolve_edge_store(data, primary_key, fallback_key):
+        edge_types = getattr(data, "edge_types", ())
+        if primary_key in edge_types:
+            return data[primary_key]
+        if fallback_key in edge_types:
+            return data[fallback_key]
+        try:
+            return data[primary_key]
+        except Exception:
+            return data[fallback_key]
+
+    def _ligand_edge_store(self, data):
+        return self._resolve_edge_store(
+            data,
+            ("ligand", "ligand"),
+            ("ligand", "lig_bond", "ligand"),
+        )
+
+    def _receptor_edge_store(self, data):
+        return self._resolve_edge_store(
+            data,
+            ("receptor", "receptor"),
+            ("receptor", "rec_contact", "receptor"),
+        )
+
+    def ligand_embedding(self, data):
+        lig_node_attr, lig_edge_index, lig_edge_attr, lig_edge_sh, lig_edge_weight = self.build_lig_conv_graph(
+            data
+        )
+        lig_node_attr = self.lig_node_embedding(lig_node_attr)
+        lig_edge_attr = self.lig_edge_embedding(lig_edge_attr)
+
+        assert self.embed_also_ligand, "otherwise reimplement padding"
+        for l in range(len(self.lig_emb_layers)):
+            edge_attr_ = torch.cat(
+                [
+                    lig_edge_attr,
+                    lig_node_attr[lig_edge_index[0], : self.ns],
+                    lig_node_attr[lig_edge_index[1], : self.ns],
+                ],
+                -1,
+            )
+            lig_node_attr = self.lig_emb_layers[l](
+                lig_node_attr,
+                lig_edge_index,
+                edge_attr_,
+                lig_edge_sh,
+                edge_weight=lig_edge_weight,
+            )
+
+        return lig_node_attr, lig_edge_index, lig_edge_attr, lig_edge_sh, lig_edge_weight
+
+    def embedding(self, data):
+        receptor_edge_store = self._receptor_edge_store(data)
+        if not hasattr(data["receptor"], "rec_node_attr"):
+            if self.lm_embedding_type not in [None, "precomputed"]:
+                sequences = [s for l in data["receptor"].sequence for s in l]
+                if isinstance(sequences[0], list):
+                    sequences = [s for l in sequences for s in l]
+                sequences = [(i, s) for i, s in enumerate(sequences)]
+                batch_labels, batch_strs, batch_tokens = self.batch_converter(sequences)
+                out = self.lm(
+                    batch_tokens.to(data["receptor"].x.device),
+                    repr_layers=[self.lm.num_layers],
+                    return_contacts=False,
+                )
+                rec_lm_emb = torch.cat(
+                    [
+                        t[: len(sequences[i][1])]
+                        for i, t in enumerate(out["representations"][self.lm.num_layers])
+                    ],
+                    dim=0,
+                )
+                data["receptor"].x = torch.cat([data["receptor"].x, rec_lm_emb], dim=-1)
+
+            rec_node_attr, rec_edge_attr, rec_edge_sh, rec_edge_weight = self.build_rec_conv_graph(data)
+            rec_node_attr = self.rec_node_embedding(rec_node_attr)
+            rec_edge_attr = self.rec_edge_embedding(rec_edge_attr)
+
+            for l in range(len(self.rec_emb_layers)):
+                edge_attr_ = torch.cat(
+                    [
+                        rec_edge_attr,
+                        rec_node_attr[receptor_edge_store.edge_index[0], : self.ns],
+                        rec_node_attr[receptor_edge_store.edge_index[1], : self.ns],
+                    ],
+                    -1,
+                )
+                rec_node_attr = self.rec_emb_layers[l](
+                    rec_node_attr,
+                    receptor_edge_store.edge_index,
+                    edge_attr_,
+                    rec_edge_sh,
+                    edge_weight=rec_edge_weight,
+                )
+
+            data["receptor"].rec_node_attr = rec_node_attr
+            receptor_edge_store.rec_edge_attr = rec_edge_attr
+            receptor_edge_store.edge_sh = rec_edge_sh
+            receptor_edge_store.edge_weight = rec_edge_weight
+
+        rec_sigma_emb = self.rec_sigma_embedding(self.timestep_emb_func(data.complex_t["tr"]))
+        rec_node_attr = data["receptor"].rec_node_attr + 0
+        rec_node_attr[:, : self.ns] = rec_node_attr[:, : self.ns] + rec_sigma_emb[data["receptor"].batch]
+        rec_edge_attr = receptor_edge_store.rec_edge_attr + rec_sigma_emb[
+            data["receptor"].batch[receptor_edge_store.edge_index[0]]
+        ]
+
+        lig_node_attr, lig_edge_index, lig_edge_attr, lig_edge_sh, lig_edge_weight = self.ligand_embedding(
+            data
+        )
+
+        return (
+            lig_node_attr,
+            lig_edge_index,
+            lig_edge_attr,
+            lig_edge_sh,
+            lig_edge_weight,
+            rec_node_attr,
+            receptor_edge_store.edge_index,
+            rec_edge_attr,
+            receptor_edge_store.edge_sh,
+            receptor_edge_store.edge_weight,
+        )
+
+    def forward(self, data):
+        if self.no_aminoacid_identities:
+            data["receptor"].x = data["receptor"].x * 0
+
+        if not self.confidence_mode:
+            tr_sigma, rot_sigma, tor_sigma = self.t_to_sigma(
+                *[data.complex_t[noise_type] for noise_type in ["tr", "rot", "tor"]]
+            )
+        else:
+            tr_sigma, rot_sigma, tor_sigma = [
+                data.complex_t[noise_type] for noise_type in ["tr", "rot", "tor"]
+            ]
+
+        (
+            lig_node_attr,
+            lig_edge_index,
+            lig_edge_attr,
+            lig_edge_sh,
+            lig_edge_weight,
+            rec_node_attr,
+            rec_edge_index,
+            rec_edge_attr,
+            rec_edge_sh,
+            rec_edge_weight,
+        ) = self.embedding(data)
+
+        if self.dynamic_max_cross:
+            cross_cutoff = (tr_sigma * 3 + 20).unsqueeze(1)
+        else:
+            cross_cutoff = self.cross_max_distance
+
+        lr_edge_index, lr_edge_attr, lr_edge_sh, rev_lr_edge_sh, lr_edge_weight = self.build_cross_conv_graph(
+            data, cross_cutoff
+        )
+        lr_edge_attr = self.cross_edge_embedding(lr_edge_attr)
+
+        node_attr = torch.cat([lig_node_attr, rec_node_attr], dim=0)
+        lr_edge_index[1] = lr_edge_index[1] + len(lig_node_attr)
+        edge_index = torch.cat(
+            [
+                lig_edge_index,
+                lr_edge_index,
+                rec_edge_index + len(lig_node_attr),
+                torch.flip(lr_edge_index, dims=[0]),
+            ],
+            dim=1,
+        )
+        edge_attr = torch.cat([lig_edge_attr, lr_edge_attr, rec_edge_attr, lr_edge_attr], dim=0)
+        edge_sh = torch.cat([lig_edge_sh, lr_edge_sh, rec_edge_sh, rev_lr_edge_sh], dim=0)
+        edge_weight = (
+            torch.cat([lig_edge_weight, lr_edge_weight, rec_edge_weight, lr_edge_weight], dim=0)
+            if torch.is_tensor(lig_edge_weight)
+            else torch.ones((len(edge_index[0]), 1), device=edge_index.device)
+        )
+        s1, s2, s3 = (
+            len(lig_edge_index[0]),
+            len(lig_edge_index[0]) + len(lr_edge_index[0]),
+            len(lig_edge_index[0]) + len(lr_edge_index[0]) + len(rec_edge_index[0]),
+        )
+
+        for l in range(len(self.conv_layers)):
+            if l < len(self.conv_layers) - 1:
+                edge_attr_ = torch.cat(
+                    [
+                        edge_attr,
+                        node_attr[edge_index[0], : self.ns],
+                        node_attr[edge_index[1], : self.ns],
+                    ],
+                    -1,
+                )
+                if self.differentiate_convolutions:
+                    edge_attr_ = [edge_attr_[:s1], edge_attr_[s1:s2], edge_attr_[s2:s3], edge_attr_[s3:]]
+                node_attr = self.conv_layers[l](
+                    node_attr,
+                    edge_index,
+                    edge_attr_,
+                    edge_sh,
+                    edge_weight=edge_weight,
+                )
+            else:
+                edge_attr_ = torch.cat(
+                    [
+                        edge_attr[:s2],
+                        node_attr[edge_index[0, :s2], : self.ns],
+                        node_attr[edge_index[1, :s2], : self.ns],
+                    ],
+                    -1,
+                )
+                if self.differentiate_convolutions:
+                    edge_attr_ = [edge_attr_[:s1], edge_attr_[s1:s2]]
+                node_attr = self.conv_layers[l](
+                    node_attr,
+                    edge_index[:, :s2],
+                    edge_attr_,
+                    edge_sh[:s2],
+                    edge_weight=edge_weight[:s2],
+                )
+
+        lig_node_attr = node_attr[: len(lig_node_attr)]
+
+        if self.confidence_mode:
+            scalar_lig_attr = (
+                torch.cat(
+                    [
+                        lig_node_attr[:, : self.ns],
+                        lig_node_attr[:, -(self.nv if self.reduce_pseudoscalars else self.ns) :],
+                    ],
+                    dim=1,
+                )
+                if self.num_conv_layers + self.num_prot_emb_layers >= 3
+                else lig_node_attr[:, : self.ns]
+            )
+
+            if self.atom_confidence:
+                scalar_lig_attr = self.atom_confidence_predictor(scalar_lig_attr)
+                atom_confidence = scalar_lig_attr[:, : self.atom_num_confidence_outputs]
+                scalar_lig_attr = scalar_lig_attr[:, self.atom_num_confidence_outputs :]
+            else:
+                atom_confidence = torch.zeros((len(lig_node_attr),), device=lig_node_attr.device)
+
+            confidence = self.confidence_predictor(
+                scatter_mean(scalar_lig_attr, data["ligand"].batch, dim=0)
+            ).squeeze(dim=-1)
+            return confidence, atom_confidence
+
+        center_edge_index, center_edge_attr, center_edge_sh = self.build_center_conv_graph(data)
+        center_edge_attr = self.center_edge_embedding(center_edge_attr)
+        if self.fixed_center_conv:
+            center_edge_attr = torch.cat(
+                [center_edge_attr, lig_node_attr[center_edge_index[1], : self.ns]], -1
+            )
+        else:
+            center_edge_attr = torch.cat(
+                [center_edge_attr, lig_node_attr[center_edge_index[0], : self.ns]], -1
+            )
+        global_pred = self.final_conv(
+            lig_node_attr,
+            center_edge_index,
+            center_edge_attr,
+            center_edge_sh,
+            out_nodes=data.num_graphs,
+        )
+
+        tr_pred = global_pred[:, :3] + (global_pred[:, 6:9] if not self.odd_parity else 0)
+        rot_pred = global_pred[:, 3:6] + (global_pred[:, 9:] if not self.odd_parity else 0)
+
+        if self.separate_noise_schedule:
+            data.graph_sigma_emb = torch.cat(
+                [self.timestep_emb_func(data.complex_t[noise_type]) for noise_type in ["tr", "rot", "tor"]],
+                dim=1,
+            )
+        elif self.asyncronous_noise_schedule:
+            data.graph_sigma_emb = self.timestep_emb_func(data.complex_t["t"])
+        else:
+            data.graph_sigma_emb = self.timestep_emb_func(data.complex_t["tr"])
+
+        tr_norm = torch.linalg.vector_norm(tr_pred, dim=1).unsqueeze(1)
+        tr_pred = tr_pred / tr_norm * self.tr_final_layer(
+            torch.cat([tr_norm, data.graph_sigma_emb], dim=1)
+        )
+        rot_norm = torch.linalg.vector_norm(rot_pred, dim=1).unsqueeze(1)
+        rot_pred = rot_pred / rot_norm * self.rot_final_layer(
+            torch.cat([rot_norm, data.graph_sigma_emb], dim=1)
+        )
+
+        if self.scale_by_sigma:
+            tr_pred = tr_pred / tr_sigma.unsqueeze(1)
+            rot_pred = rot_pred * so3.score_norm(rot_sigma.cpu()).unsqueeze(1).to(
+                data["ligand"].x.device
+            )
+
+        sidechain_pred = None
+        if self.sidechain_pred:
+            rec_node_attr = node_attr[len(lig_node_attr) :]
+            sidechain_pred = self.sidechain_predictor(rec_node_attr)
+            sidechain_pred = sidechain_pred[:, :10] + sidechain_pred[:, 10:]
+
+        if self.no_torsion or data["ligand"].edge_mask.sum() == 0:
+            return tr_pred, rot_pred, torch.empty(0, device=self.device), sidechain_pred
+
+        tor_bonds, tor_edge_index, tor_edge_attr, tor_edge_sh, tor_edge_weight = self.build_bond_conv_graph(
+            data
+        )
+        tor_bond_vec = data["ligand"].pos[tor_bonds[1]] - data["ligand"].pos[tor_bonds[0]]
+        tor_bond_attr = lig_node_attr[tor_bonds[0]] + lig_node_attr[tor_bonds[1]]
+
+        tor_bonds_sh = o3.spherical_harmonics(
+            "2e", tor_bond_vec, normalize=True, normalization="component"
+        )
+        tor_edge_sh = self.final_tp_tor(tor_edge_sh, tor_bonds_sh[tor_edge_index[0]])
+
+        tor_edge_attr = torch.cat(
+            [
+                tor_edge_attr,
+                lig_node_attr[tor_edge_index[1], : self.ns],
+                tor_bond_attr[tor_edge_index[0], : self.ns],
+            ],
+            -1,
+        )
+        tor_pred = self.tor_bond_conv(
+            lig_node_attr,
+            tor_edge_index,
+            tor_edge_attr,
+            tor_edge_sh,
+            out_nodes=data["ligand"].edge_mask.sum(),
+            reduce="mean",
+            edge_weight=tor_edge_weight,
+        )
+        tor_pred = self.tor_final_layer(tor_pred).squeeze(1)
+        ligand_edge_store = self._ligand_edge_store(data)
+        edge_sigma = tor_sigma[data["ligand"].batch][ligand_edge_store.edge_index[0]][
+            data["ligand"].edge_mask
+        ]
+
+        if self.scale_by_sigma:
+            tor_pred = tor_pred * torch.sqrt(
+                torch.tensor(torus.score_norm(edge_sigma.cpu().numpy())).float().to(
+                    data["ligand"].x.device
+                )
+            )
+        return tr_pred, rot_pred, tor_pred, sidechain_pred
+
+    def torsional_forward(self, data):
+        tor_sigma = self.t_to_sigma(data.complex_t["tor"])
+
+        lig_node_attr, lig_edge_index, lig_edge_attr, lig_edge_sh, lig_edge_weight = self.ligand_embedding(
+            data
+        )
+
+        if self.separate_noise_schedule:
+            data.graph_sigma_emb = torch.cat(
+                [self.timestep_emb_func(data.complex_t[noise_type]) for noise_type in ["tr", "rot", "tor"]],
+                dim=1,
+            )
+        elif self.asyncronous_noise_schedule:
+            data.graph_sigma_emb = self.timestep_emb_func(data.complex_t["t"])
+        else:
+            data.graph_sigma_emb = self.timestep_emb_func(data.complex_t["tr"])
+
+        tor_bonds, tor_edge_index, tor_edge_attr, tor_edge_sh, tor_edge_weight = self.build_bond_conv_graph(
+            data
+        )
+        tor_bond_vec = data["ligand"].pos[tor_bonds[1]] - data["ligand"].pos[tor_bonds[0]]
+        tor_bond_attr = lig_node_attr[tor_bonds[0]] + lig_node_attr[tor_bonds[1]]
+
+        tor_bonds_sh = o3.spherical_harmonics(
+            "2e", tor_bond_vec, normalize=True, normalization="component"
+        )
+        tor_edge_sh = self.final_tp_tor(tor_edge_sh, tor_bonds_sh[tor_edge_index[0]])
+
+        tor_edge_attr = torch.cat(
+            [
+                tor_edge_attr,
+                lig_node_attr[tor_edge_index[1], : self.ns],
+                tor_bond_attr[tor_edge_index[0], : self.ns],
+            ],
+            -1,
+        )
+        tor_pred = self.tor_bond_conv(
+            lig_node_attr,
+            tor_edge_index,
+            tor_edge_attr,
+            tor_edge_sh,
+            out_nodes=data["ligand"].edge_mask.sum(),
+            reduce="mean",
+            edge_weight=tor_edge_weight,
+        )
+        tor_pred = self.tor_final_layer(tor_pred).squeeze(1)
+        ligand_edge_store = self._ligand_edge_store(data)
+        edge_sigma = tor_sigma[data["ligand"].batch][ligand_edge_store.edge_index[0]][
+            data["ligand"].edge_mask
+        ]
+
+        if self.scale_by_sigma:
+            tor_pred = tor_pred * torch.sqrt(
+                torch.tensor(torus.score_norm(edge_sigma.cpu().numpy())).float().to(
+                    data["ligand"].x.device
+                )
+            )
+        return 0, 0, tor_pred, 0
+
+    def get_edge_weight(self, edge_vec, max_norm):
+        if self.smooth_edges:
+            normalised_norm = torch.clip(edge_vec.norm(dim=-1) * np.pi / max_norm, max=np.pi)
+            return 0.5 * (torch.cos(normalised_norm) + 1.0).unsqueeze(-1)
+        return 1.0
+
+    def build_lig_conv_graph(self, data):
+        if self.separate_noise_schedule:
+            data["ligand"].node_sigma_emb = torch.cat(
+                [self.timestep_emb_func(data["ligand"].node_t[noise_type]) for noise_type in ["tr", "rot", "tor"]],
+                dim=1,
+            )
+        elif self.asyncronous_noise_schedule:
+            data["ligand"].node_sigma_emb = self.timestep_emb_func(data["ligand"].node_t["t"])
+        else:
+            data["ligand"].node_sigma_emb = self.timestep_emb_func(data["ligand"].node_t["tr"])
+
+        ligand_edge_store = self._ligand_edge_store(data)
+        radius_edges = radius_graph(data["ligand"].pos, self.lig_max_radius, data["ligand"].batch)
+        edge_index = torch.cat([ligand_edge_store.edge_index, radius_edges], 1).long()
+        edge_attr = torch.cat(
+            [
+                ligand_edge_store.edge_attr,
+                torch.zeros(
+                    radius_edges.shape[-1],
+                    self.in_lig_edge_features,
+                    device=data["ligand"].x.device,
+                ),
+            ],
+            0,
+        )
+
+        edge_sigma_emb = data["ligand"].node_sigma_emb[edge_index[0].long()]
+        edge_attr = torch.cat([edge_attr, edge_sigma_emb], 1)
+        node_attr = torch.cat([data["ligand"].x, data["ligand"].node_sigma_emb], 1)
+
+        src, dst = edge_index
+        edge_vec = data["ligand"].pos[dst.long()] - data["ligand"].pos[src.long()]
+        edge_length_emb = self.lig_distance_expansion(edge_vec.norm(dim=-1))
+
+        edge_attr = torch.cat([edge_attr, edge_length_emb], 1)
+        edge_sh = o3.spherical_harmonics(
+            self.sh_irreps, edge_vec, normalize=True, normalization="component"
+        )
+        edge_weight = self.get_edge_weight(edge_vec, self.lig_max_radius)
+
+        return node_attr, edge_index, edge_attr, edge_sh, edge_weight
+
+    def build_rec_conv_graph(self, data):
+        assert not self.separate_noise_schedule or self.asyncronous_noise_schedule, "removed support in this function"
+        node_attr = data["receptor"].x
+
+        receptor_edge_store = self._receptor_edge_store(data)
+        edge_index = receptor_edge_store.edge_index
+        src, dst = edge_index
+        edge_vec = data["receptor"].pos[dst.long()] - data["receptor"].pos[src.long()]
+
+        edge_length_emb = self.rec_distance_expansion(edge_vec.norm(dim=-1))
+        edge_attr = edge_length_emb
+        edge_sh = o3.spherical_harmonics(
+            self.sh_irreps, edge_vec, normalize=True, normalization="component"
+        )
+        edge_weight = self.get_edge_weight(edge_vec, self.rec_max_radius)
+
+        return node_attr, edge_attr, edge_sh, edge_weight
+
+    def build_misc_atom_conv_graph(self, data):
+        if self.separate_noise_schedule:
+            data["misc_atom"].node_sigma_emb = torch.cat(
+                [self.timestep_emb_func(data["misc_atom"].node_t[noise_type]) for noise_type in ["tr", "rot", "tor"]],
+                dim=1,
+            )
+        elif self.asyncronous_noise_schedule:
+            data["misc_atom"].node_sigma_emb = self.timestep_emb_func(data["misc_atom"].node_t["t"])
+        else:
+            data["misc_atom"].node_sigma_emb = self.timestep_emb_func(data["misc_atom"].node_t["tr"])
+        node_attr = torch.cat([data["misc_atom"].x, data["misc_atom"].node_sigma_emb], 1)
+
+        edge_index = data["misc_atom", "misc_atom"].edge_index
+        src, dst = edge_index
+        edge_vec = data["misc_atom"].pos[dst.long()] - data["misc_atom"].pos[src.long()]
+
+        edge_length_emb = self.lig_distance_expansion(edge_vec.norm(dim=-1))
+        edge_sigma_emb = data["misc_atom"].node_sigma_emb[edge_index[0].long()]
+        edge_attr = torch.cat([edge_sigma_emb, edge_length_emb], 1)
+        edge_sh = o3.spherical_harmonics(
+            self.sh_irreps, edge_vec, normalize=True, normalization="component"
+        )
+        edge_weight = self.get_edge_weight(edge_vec, self.lig_max_radius)
+
+        return node_attr, edge_index, edge_attr, edge_sh, edge_weight
+
+    def build_cross_conv_graph(self, data, cross_distance_cutoff):
+        radius_on_cpu = _env_flag("DIFFDOCK_RADIUS_ON_CPU")
+        if torch.is_tensor(cross_distance_cutoff):
+            receptor_pos = data["receptor"].pos
+            ligand_pos = data["ligand"].pos
+            receptor_batch = data["receptor"].batch
+            ligand_batch = data["ligand"].batch
+            cutoff = cross_distance_cutoff
+            if radius_on_cpu:
+                receptor_pos = receptor_pos.detach().cpu()
+                ligand_pos = ligand_pos.detach().cpu()
+                receptor_batch = receptor_batch.detach().cpu()
+                ligand_batch = ligand_batch.detach().cpu()
+                cutoff = cutoff.detach().cpu()
+            edge_index = radius(
+                receptor_pos / cutoff[receptor_batch],
+                ligand_pos / cutoff[ligand_batch],
+                1,
+                receptor_batch,
+                ligand_batch,
+                max_num_neighbors=10000,
+            )
+        else:
+            receptor_pos = data["receptor"].pos
+            ligand_pos = data["ligand"].pos
+            receptor_batch = data["receptor"].batch
+            ligand_batch = data["ligand"].batch
+            if radius_on_cpu:
+                receptor_pos = receptor_pos.detach().cpu()
+                ligand_pos = ligand_pos.detach().cpu()
+                receptor_batch = receptor_batch.detach().cpu()
+                ligand_batch = ligand_batch.detach().cpu()
+            edge_index = radius(
+                receptor_pos,
+                ligand_pos,
+                cross_distance_cutoff,
+                receptor_batch,
+                ligand_batch,
+                max_num_neighbors=10000,
+            )
+        if radius_on_cpu:
+            edge_index = edge_index.to(data["ligand"].pos.device)
+
+        src, dst = edge_index
+        edge_vec = data["receptor"].pos[dst.long()] - data["ligand"].pos[src.long()]
+
+        edge_length_emb = self.cross_distance_expansion(edge_vec.norm(dim=-1))
+        edge_sigma_emb = data["ligand"].node_sigma_emb[src.long()]
+        edge_attr = torch.cat([edge_sigma_emb, edge_length_emb], 1)
+        edge_sh = o3.spherical_harmonics(
+            self.sh_irreps, edge_vec, normalize=True, normalization="component"
+        )
+        rev_edge_sh = o3.spherical_harmonics(
+            self.sh_irreps, -edge_vec, normalize=True, normalization="component"
+        )
+
+        cutoff_d = (
+            cross_distance_cutoff[data["ligand"].batch[src]].squeeze()
+            if torch.is_tensor(cross_distance_cutoff)
+            else cross_distance_cutoff
+        )
+        edge_weight = self.get_edge_weight(edge_vec, cutoff_d)
+
+        return edge_index, edge_attr, edge_sh, rev_edge_sh, edge_weight
+
+    def build_misc_cross_conv_graph(self, data, lr_cross_distance_cutoff):
+        if torch.is_tensor(lr_cross_distance_cutoff):
+            lr_edge_index = radius(
+                data["receptor"].pos / lr_cross_distance_cutoff[data["receptor"].batch],
+                data["ligand"].pos / lr_cross_distance_cutoff[data["ligand"].batch],
+                1,
+                data["receptor"].batch,
+                data["ligand"].batch,
+                max_num_neighbors=10000,
+            )
+        else:
+            lr_edge_index = radius(
+                data["receptor"].pos,
+                data["ligand"].pos,
+                lr_cross_distance_cutoff,
+                data["receptor"].batch,
+                data["ligand"].batch,
+                max_num_neighbors=10000,
+            )
+
+        lr_edge_vec = data["receptor"].pos[lr_edge_index[1].long()] - data["ligand"].pos[lr_edge_index[0].long()]
+        lr_edge_length_emb = self.cross_distance_expansion(lr_edge_vec.norm(dim=-1))
+        lr_edge_sigma_emb = data["ligand"].node_sigma_emb[lr_edge_index[0].long()]
+        lr_edge_attr = torch.cat([lr_edge_sigma_emb, lr_edge_length_emb], 1)
+        lr_edge_sh = o3.spherical_harmonics(
+            self.sh_irreps, lr_edge_vec, normalize=True, normalization="component"
+        )
+
+        cutoff_d = (
+            lr_cross_distance_cutoff[data["ligand"].batch[lr_edge_index[0]]].squeeze()
+            if torch.is_tensor(lr_cross_distance_cutoff)
+            else lr_cross_distance_cutoff
+        )
+        lr_edge_weight = self.get_edge_weight(lr_edge_vec, cutoff_d)
+
+        la_edge_index = radius(
+            data["misc_atom"].pos,
+            data["ligand"].pos,
+            self.lig_max_radius,
+            data["misc_atom"].batch,
+            data["ligand"].batch,
+            max_num_neighbors=10000,
+        )
+
+        la_edge_vec = data["misc_atom"].pos[la_edge_index[1].long()] - data["ligand"].pos[la_edge_index[0].long()]
+        la_edge_length_emb = self.cross_distance_expansion(la_edge_vec.norm(dim=-1))
+        la_edge_sigma_emb = data["ligand"].node_sigma_emb[la_edge_index[0].long()]
+        la_edge_attr = torch.cat([la_edge_sigma_emb, la_edge_length_emb], 1)
+        la_edge_sh = o3.spherical_harmonics(
+            self.sh_irreps, la_edge_vec, normalize=True, normalization="component"
+        )
+        la_edge_weight = self.get_edge_weight(la_edge_vec, self.lig_max_radius)
+
+        ar_edge_index = data["misc_atom", "receptor"].edge_index
+        ar_edge_vec = data["receptor"].pos[ar_edge_index[1].long()] - data["misc_atom"].pos[ar_edge_index[0].long()]
+        ar_edge_length_emb = self.rec_distance_expansion(ar_edge_vec.norm(dim=-1))
+        ar_edge_sigma_emb = data["misc_atom"].node_sigma_emb[ar_edge_index[0].long()]
+        ar_edge_attr = torch.cat([ar_edge_sigma_emb, ar_edge_length_emb], 1)
+        ar_edge_sh = o3.spherical_harmonics(
+            self.sh_irreps, ar_edge_vec, normalize=True, normalization="component"
+        )
+        ar_edge_weight = 1
+
+        return (
+            lr_edge_index,
+            lr_edge_attr,
+            lr_edge_sh,
+            lr_edge_weight,
+            la_edge_index,
+            la_edge_attr,
+            la_edge_sh,
+            la_edge_weight,
+            ar_edge_index,
+            ar_edge_attr,
+            ar_edge_sh,
+            ar_edge_weight,
+        )
+
+    def build_center_conv_graph(self, data):
+        edge_index = torch.cat(
+            [
+                data["ligand"].batch.unsqueeze(0),
+                torch.arange(len(data["ligand"].batch)).to(data["ligand"].x.device).unsqueeze(0),
+            ],
+            dim=0,
+        )
+
+        center_pos, count = torch.zeros((data.num_graphs, 3)).to(data["ligand"].x.device), torch.zeros(
+            (data.num_graphs, 3)
+        ).to(data["ligand"].x.device)
+        center_pos.index_add_(0, index=data["ligand"].batch, source=data["ligand"].pos)
+        center_pos = center_pos / torch.bincount(data["ligand"].batch).unsqueeze(1)
+
+        edge_vec = data["ligand"].pos[edge_index[1]] - center_pos[edge_index[0]]
+        edge_attr = self.center_distance_expansion(edge_vec.norm(dim=-1))
+        edge_sigma_emb = data["ligand"].node_sigma_emb[edge_index[1].long()]
+        edge_attr = torch.cat([edge_attr, edge_sigma_emb], 1)
+        edge_sh = o3.spherical_harmonics(
+            self.sh_irreps, edge_vec, normalize=True, normalization="component"
+        )
+        return edge_index, edge_attr, edge_sh
+
+    def build_bond_conv_graph(self, data):
+        ligand_edge_store = self._ligand_edge_store(data)
+        bonds = ligand_edge_store.edge_index[:, data["ligand"].edge_mask].long()
+        bond_pos = (data["ligand"].pos[bonds[0]] + data["ligand"].pos[bonds[1]]) / 2
+        bond_batch = data["ligand"].batch[bonds[0]]
+        edge_index = radius(
+            data["ligand"].pos,
+            bond_pos,
+            self.lig_max_radius,
+            batch_x=data["ligand"].batch,
+            batch_y=bond_batch,
+        )
+
+        edge_vec = data["ligand"].pos[edge_index[1]] - bond_pos[edge_index[0]]
+        edge_attr = self.lig_distance_expansion(edge_vec.norm(dim=-1))
+
+        edge_attr = self.final_edge_embedding(edge_attr)
+        edge_sh = o3.spherical_harmonics(
+            self.sh_irreps, edge_vec, normalize=True, normalization="component"
+        )
+        edge_weight = self.get_edge_weight(edge_vec, self.lig_max_radius)
+
+        return bonds, edge_index, edge_attr, edge_sh, edge_weight

models/layers.py

@@ -0,0 +1,128 @@
+import torch
+from torch import nn
+
+ACTIVATIONS = {
+    "relu": nn.ReLU,
+    "silu": nn.SiLU,
+}
+
+
+def FCBlock(in_dim, hidden_dim, out_dim, layers, dropout, activation="relu"):
+    activation = ACTIVATIONS[activation]
+    assert layers >= 2
+    sequential = [nn.Linear(in_dim, hidden_dim), activation(), nn.Dropout(dropout)]
+    for i in range(layers - 2):
+        sequential += [nn.Linear(hidden_dim, hidden_dim), activation(), nn.Dropout(dropout)]
+    sequential += [nn.Linear(hidden_dim, out_dim)]
+    return nn.Sequential(*sequential)
+
+
+class GaussianSmearing(torch.nn.Module):
+    # used to embed the edge distances
+    def __init__(self, start=0.0, stop=5.0, num_gaussians=50):
+        super().__init__()
+        offset = torch.linspace(start, stop, num_gaussians)
+        self.coeff = -0.5 / (offset[1] - offset[0]).item() ** 2
+        self.register_buffer("offset", offset)
+
+    def forward(self, dist):
+        dist = dist.view(-1, 1) - self.offset.view(1, -1)
+        return torch.exp(self.coeff * torch.pow(dist, 2))
+
+
+class AtomEncoder(torch.nn.Module):
+    def __init__(self, emb_dim, feature_dims, sigma_embed_dim, lm_embedding_dim=0):
+        """
+
+        Parameters
+        ----------
+        emb_dim
+        feature_dims
+            first element of feature_dims tuple is a list with the length of each categorical feature,
+            and the second is the number of scalar features
+        sigma_embed_dim
+        lm_embedding_dim
+        """
+        super(AtomEncoder, self).__init__()
+        self.atom_embedding_list = torch.nn.ModuleList()
+        self.num_categorical_features = len(feature_dims[0])
+        self.additional_features_dim = feature_dims[1] + sigma_embed_dim + lm_embedding_dim
+        for i, dim in enumerate(feature_dims[0]):
+            emb = torch.nn.Embedding(dim, emb_dim)
+            torch.nn.init.xavier_uniform_(emb.weight.data)
+            self.atom_embedding_list.append(emb)
+
+        if self.additional_features_dim > 0:
+            self.additional_features_embedder = torch.nn.Linear(
+                self.additional_features_dim + emb_dim, emb_dim
+            )
+
+    def forward(self, x):
+        x_embedding = 0
+        assert x.shape[1] == self.num_categorical_features + self.additional_features_dim
+        for i in range(self.num_categorical_features):
+            x_embedding += self.atom_embedding_list[i](x[:, i].long())
+
+        if self.additional_features_dim > 0:
+            x_embedding = self.additional_features_embedder(
+                torch.cat([x_embedding, x[:, self.num_categorical_features :]], axis=1)
+            )
+        return x_embedding
+
+
+class OldAtomEncoder(torch.nn.Module):
+    def __init__(self, emb_dim, feature_dims, sigma_embed_dim, lm_embedding_type=None):
+        """
+
+        Parameters
+        ----------
+        emb_dim
+        feature_dims
+            first element of feature_dims tuple is a list with the length of each categorical feature,
+            and the second is the number of scalar features
+        sigma_embed_dim
+        lm_embedding_type
+        """
+        super(OldAtomEncoder, self).__init__()
+        self.atom_embedding_list = torch.nn.ModuleList()
+        self.num_categorical_features = len(feature_dims[0])
+        self.num_scalar_features = feature_dims[1] + sigma_embed_dim
+        self.lm_embedding_type = lm_embedding_type
+        for i, dim in enumerate(feature_dims[0]):
+            emb = torch.nn.Embedding(dim, emb_dim)
+            torch.nn.init.xavier_uniform_(emb.weight.data)
+            self.atom_embedding_list.append(emb)
+
+        if self.num_scalar_features > 0:
+            self.linear = torch.nn.Linear(self.num_scalar_features, emb_dim)
+        if self.lm_embedding_type is not None:
+            if self.lm_embedding_type == "esm":
+                self.lm_embedding_dim = 1280
+            else:
+                raise ValueError(
+                    "LM Embedding type was not correctly determined. LM embedding type: ",
+                    self.lm_embedding_type,
+                )
+            self.lm_embedding_layer = torch.nn.Linear(self.lm_embedding_dim + emb_dim, emb_dim)
+
+    def forward(self, x):
+        x_embedding = 0
+        if self.lm_embedding_type is not None:
+            assert (
+                x.shape[1]
+                == self.num_categorical_features + self.num_scalar_features + self.lm_embedding_dim
+            )
+        else:
+            assert x.shape[1] == self.num_categorical_features + self.num_scalar_features
+        for i in range(self.num_categorical_features):
+            x_embedding += self.atom_embedding_list[i](x[:, i].long())
+
+        if self.num_scalar_features > 0:
+            x_embedding += self.linear(
+                x[:, self.num_categorical_features : self.num_categorical_features + self.num_scalar_features]
+            )
+        if self.lm_embedding_type is not None:
+            x_embedding = self.lm_embedding_layer(
+                torch.cat([x_embedding, x[:, -self.lm_embedding_dim :]], axis=1)
+            )
+        return x_embedding

models/old_aa_model.py

@@ -0,0 +1,582 @@
+from e3nn import o3
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch_cluster import radius, radius_graph
+from torch_scatter import scatter, scatter_mean
+import numpy as np
+from e3nn.nn import BatchNorm
+
+from onescience.datapipes.diffdock.process_mols import lig_feature_dims, rec_residue_feature_dims, rec_atom_feature_dims
+from onescience.utils.diffdock import so3, torus
+
+from .layers import GaussianSmearing, OldAtomEncoder, AtomEncoder
+from .tensor_layers import OldTensorProductConvLayer
+
+AGGREGATORS = {"mean": lambda x: torch.mean(x, dim=1),
+               "max": lambda x: torch.max(x, dim=1)[0],
+               "min": lambda x: torch.min(x, dim=1)[0],
+               "std": lambda x: torch.std(x, dim=1)}
+
+
+class AAOldModel(torch.nn.Module):
+    def __init__(self, t_to_sigma, device, timestep_emb_func, in_lig_edge_features=4, sigma_embed_dim=32, sh_lmax=2,
+                 ns=16, nv=4, num_conv_layers=2, lig_max_radius=5, rec_max_radius=30, cross_max_distance=250,
+                 center_max_distance=30, distance_embed_dim=32, cross_distance_embed_dim=32, no_torsion=False,
+                 scale_by_sigma=True, norm_by_sigma=True, use_second_order_repr=False, batch_norm=True,
+                 dynamic_max_cross=False, dropout=0.0, smooth_edges=False, odd_parity=False,
+                 separate_noise_schedule=False, lm_embedding_type=None, confidence_mode=False,
+                 confidence_dropout=0, confidence_no_batchnorm = False,
+                 asyncronous_noise_schedule=False, affinity_prediction=False, parallel=1,
+                 parallel_aggregators="mean max min std", num_confidence_outputs=1, fixed_center_conv=False,
+                 no_aminoacid_identities=False, include_miscellaneous_atoms=False, use_old_atom_encoder=False):
+        super(AAOldModel, self).__init__()
+        assert (not no_aminoacid_identities) or (lm_embedding_type is None), "no language model emb without identities"
+        if parallel > 1: assert affinity_prediction
+        self.t_to_sigma = t_to_sigma
+        self.in_lig_edge_features = in_lig_edge_features
+        sigma_embed_dim *= (3 if separate_noise_schedule else 1)
+        self.sigma_embed_dim = sigma_embed_dim
+        self.lig_max_radius = lig_max_radius
+        self.rec_max_radius = rec_max_radius
+        self.cross_max_distance = cross_max_distance
+        self.dynamic_max_cross = dynamic_max_cross
+        self.center_max_distance = center_max_distance
+        self.distance_embed_dim = distance_embed_dim
+        self.cross_distance_embed_dim = cross_distance_embed_dim
+        self.sh_irreps = o3.Irreps.spherical_harmonics(lmax=sh_lmax)
+        self.ns, self.nv = ns, nv
+        self.scale_by_sigma = scale_by_sigma
+        self.norm_by_sigma = norm_by_sigma
+        self.device = device
+        self.no_torsion = no_torsion
+        self.smooth_edges = smooth_edges
+        self.odd_parity = odd_parity
+        self.num_conv_layers = num_conv_layers
+        self.timestep_emb_func = timestep_emb_func
+        self.separate_noise_schedule = separate_noise_schedule
+        self.confidence_mode = confidence_mode
+        self.num_conv_layers = num_conv_layers
+        self.asyncronous_noise_schedule = asyncronous_noise_schedule
+        self.affinity_prediction = affinity_prediction
+        self.parallel, self.parallel_aggregators = parallel, parallel_aggregators.split(' ')
+        self.fixed_center_conv = fixed_center_conv
+        self.no_aminoacid_identities = no_aminoacid_identities
+
+        # embedding layers
+        atom_encoder_class = OldAtomEncoder if use_old_atom_encoder else AtomEncoder
+        lm_embedding_dim = 0 if lm_embedding_type is None else 1280
+        self.lig_node_embedding = atom_encoder_class(emb_dim=ns, feature_dims=lig_feature_dims, sigma_embed_dim=sigma_embed_dim)
+        self.lig_edge_embedding = nn.Sequential(nn.Linear(in_lig_edge_features + sigma_embed_dim + distance_embed_dim, ns),nn.ReLU(),nn.Dropout(dropout),nn.Linear(ns, ns))
+
+        self.rec_node_embedding = (
+            atom_encoder_class(
+                emb_dim=ns,
+                feature_dims=rec_residue_feature_dims,
+                sigma_embed_dim=sigma_embed_dim,
+                lm_embedding_type=lm_embedding_type,
+            )
+            if use_old_atom_encoder
+            else atom_encoder_class(
+                emb_dim=ns,
+                feature_dims=rec_residue_feature_dims,
+                sigma_embed_dim=sigma_embed_dim,
+                lm_embedding_dim=lm_embedding_dim,
+            )
+        )
+        self.rec_edge_embedding = nn.Sequential(nn.Linear(sigma_embed_dim + distance_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout),nn.Linear(ns, ns))
+
+        self.atom_node_embedding = atom_encoder_class(emb_dim=ns, feature_dims=rec_atom_feature_dims, sigma_embed_dim=sigma_embed_dim)
+        self.atom_edge_embedding = nn.Sequential(nn.Linear(sigma_embed_dim + distance_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout),nn.Linear(ns, ns))
+
+        self.lr_edge_embedding = nn.Sequential(nn.Linear(sigma_embed_dim + cross_distance_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout),nn.Linear(ns, ns))
+        self.ar_edge_embedding = nn.Sequential(nn.Linear(sigma_embed_dim + distance_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout),nn.Linear(ns, ns))
+        self.la_edge_embedding = nn.Sequential(nn.Linear(sigma_embed_dim + cross_distance_embed_dim, ns), nn.ReLU(), nn.Dropout(dropout),nn.Linear(ns, ns))
+
+        self.lig_distance_expansion = GaussianSmearing(0.0, lig_max_radius, distance_embed_dim)
+        self.rec_distance_expansion = GaussianSmearing(0.0, rec_max_radius, distance_embed_dim)
+        self.cross_distance_expansion = GaussianSmearing(0.0, cross_max_distance, cross_distance_embed_dim)
+
+        if use_second_order_repr:
+            irrep_seq = [
+                f'{ns}x0e',
+                f'{ns}x0e + {nv}x1o + {nv}x2e',
+                f'{ns}x0e + {nv}x1o + {nv}x2e + {nv}x1e + {nv}x2o',
+                f'{ns}x0e + {nv}x1o + {nv}x2e + {nv}x1e + {nv}x2o + {ns}x0o'
+            ]
+        else:
+            irrep_seq = [
+                f'{ns}x0e',
+                f'{ns}x0e + {nv}x1o',
+                f'{ns}x0e + {nv}x1o + {nv}x1e',
+                f'{ns}x0e + {nv}x1o + {nv}x1e + {ns}x0o'
+            ]
+
+        # convolutional layers
+        conv_layers = []
+        for i in range(num_conv_layers):
+            in_irreps = irrep_seq[min(i, len(irrep_seq) - 1)]
+            out_irreps = irrep_seq[min(i + 1, len(irrep_seq) - 1)]
+            parameters = {
+                'in_irreps': in_irreps,
+                'sh_irreps': self.sh_irreps,
+                'out_irreps': out_irreps,
+                'n_edge_features': 3 * ns,
+                'residual': False,
+                'batch_norm': batch_norm,
+                'dropout': dropout
+            }
+
+            for _ in range(9): # 3 intra & 6 inter per each layer
+                conv_layers.append(OldTensorProductConvLayer(**parameters))
+
+        self.conv_layers = nn.ModuleList(conv_layers)
+
+        # confidence and affinity prediction layers
+        if self.confidence_mode:
+            if self.affinity_prediction:
+                if self.parallel > 1:
+                    output_confidence_dim = 1 + ns
+                else:
+                    output_confidence_dim = num_confidence_outputs +1
+            else:
+                output_confidence_dim = num_confidence_outputs
+
+            self.confidence_predictor = nn.Sequential(
+                nn.Linear(2 * self.ns if num_conv_layers >= 3 else self.ns, ns),
+                nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(),
+                nn.ReLU(),
+                nn.Dropout(confidence_dropout),
+                nn.Linear(ns, ns),
+                nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(),
+                nn.ReLU(),
+                nn.Dropout(confidence_dropout),
+                nn.Linear(ns, output_confidence_dim)
+            )
+
+            if self.parallel > 1:
+                self.affinity_predictor = nn.Sequential(
+                    nn.Linear(len(self.parallel_aggregators) * ns, ns),
+                    nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(),
+                    nn.ReLU(),
+                    nn.Dropout(confidence_dropout),
+                    nn.Linear(ns, ns),
+                    nn.BatchNorm1d(ns) if not confidence_no_batchnorm else nn.Identity(),
+                    nn.ReLU(),
+                    nn.Dropout(confidence_dropout),
+                    nn.Linear(ns, 1)
+                )
+
+        else:
+            # convolution for translational and rotational scores
+            self.center_distance_expansion = GaussianSmearing(0.0, center_max_distance, distance_embed_dim)
+            self.center_edge_embedding = nn.Sequential(
+                nn.Linear(distance_embed_dim + sigma_embed_dim, ns),
+                nn.ReLU(),
+                nn.Dropout(dropout),
+                nn.Linear(ns, ns)
+            )
+
+            self.final_conv = OldTensorProductConvLayer(
+                in_irreps=self.conv_layers[-1].out_irreps,
+                sh_irreps=self.sh_irreps,
+                out_irreps=f'2x1o + 2x1e' if not self.odd_parity else '1x1o + 1x1e',
+                n_edge_features=2 * ns,
+                residual=False,
+                dropout=dropout,
+                batch_norm=batch_norm
+            )
+
+            self.tr_final_layer = nn.Sequential(nn.Linear(1 + sigma_embed_dim, ns),nn.Dropout(dropout), nn.ReLU(), nn.Linear(ns, 1))
+            self.rot_final_layer = nn.Sequential(nn.Linear(1 + sigma_embed_dim, ns),nn.Dropout(dropout), nn.ReLU(), nn.Linear(ns, 1))
+
+            if not no_torsion:
+                # convolution for torsional score
+                self.final_edge_embedding = nn.Sequential(
+                    nn.Linear(distance_embed_dim, ns),
+                    nn.ReLU(),
+                    nn.Dropout(dropout),
+                    nn.Linear(ns, ns)
+                )
+                self.final_tp_tor = o3.FullTensorProduct(self.sh_irreps, "2e")
+                self.tor_bond_conv = OldTensorProductConvLayer(
+                    in_irreps=self.conv_layers[-1].out_irreps,
+                    sh_irreps=self.final_tp_tor.irreps_out,
+                    out_irreps=f'{ns}x0o + {ns}x0e' if not self.odd_parity else f'{ns}x0o',
+                    n_edge_features=3 * ns,
+                    residual=False,
+                    dropout=dropout,
+                    batch_norm=batch_norm
+                )
+                self.tor_final_layer = nn.Sequential(
+                    nn.Linear(2 * ns if not self.odd_parity else ns, ns, bias=False),
+                    nn.Tanh(),
+                    nn.Dropout(dropout),
+                    nn.Linear(ns, 1, bias=False)
+                )
+
+    @staticmethod
+    def _resolve_edge_store(data, primary_key, fallback_key):
+        edge_types = getattr(data, "edge_types", ())
+        if primary_key in edge_types:
+            return data[primary_key]
+        if fallback_key in edge_types:
+            return data[fallback_key]
+        try:
+            return data[primary_key]
+        except Exception:
+            return data[fallback_key]
+
+    def _ligand_edge_store(self, data):
+        return self._resolve_edge_store(
+            data,
+            ("ligand", "ligand"),
+            ("ligand", "lig_bond", "ligand"),
+        )
+
+    def _receptor_edge_store(self, data):
+        return self._resolve_edge_store(
+            data,
+            ("receptor", "receptor"),
+            ("receptor", "rec_contact", "receptor"),
+        )
+
+    def _atom_edge_store(self, data):
+        return self._resolve_edge_store(
+            data,
+            ("atom", "atom"),
+            ("atom", "atom_contact", "atom"),
+        )
+
+    def _atom_receptor_edge_store(self, data):
+        return self._resolve_edge_store(
+            data,
+            ("atom", "receptor"),
+            ("atom", "atom_rec_contact", "receptor"),
+        )
+
+    def forward(self, data):
+        if self.no_aminoacid_identities:
+            data['receptor'].x = data['receptor'].x * 0
+
+        if not self.confidence_mode:
+            tr_sigma, rot_sigma, tor_sigma = self.t_to_sigma(*[data.complex_t[noise_type] for noise_type in ['tr', 'rot', 'tor']])
+        else:
+            tr_sigma, rot_sigma, tor_sigma = [data.complex_t[noise_type] for noise_type in ['tr', 'rot', 'tor']]
+
+        # build ligand graph
+        lig_node_attr, lig_edge_index, lig_edge_attr, lig_edge_sh, lig_edge_weight = self.build_lig_conv_graph(data)
+        lig_node_attr = self.lig_node_embedding(lig_node_attr)
+        lig_edge_attr = self.lig_edge_embedding(lig_edge_attr)
+
+        # build receptor graph
+        rec_node_attr, rec_edge_index, rec_edge_attr, rec_edge_sh, rec_edge_weight = self.build_rec_conv_graph(data)
+        rec_node_attr = self.rec_node_embedding(rec_node_attr)
+        rec_edge_attr = self.rec_edge_embedding(rec_edge_attr)
+
+        # build atom graph
+        atom_node_attr, atom_edge_index, atom_edge_attr, atom_edge_sh, atom_edge_weight = self.build_atom_conv_graph(data)
+        atom_node_attr = self.atom_node_embedding(atom_node_attr)
+        atom_edge_attr = self.atom_edge_embedding(atom_edge_attr)
+
+        # build cross graph
+        cross_cutoff = (tr_sigma * 3 + 20).unsqueeze(1) if self.dynamic_max_cross else self.cross_max_distance
+        lr_edge_index, lr_edge_attr, lr_edge_sh, lr_edge_weight, la_edge_index, la_edge_attr, \
+            la_edge_sh, la_edge_weight, ar_edge_index, ar_edge_attr, ar_edge_sh, ar_edge_weight = \
+            self.build_cross_conv_graph(data, cross_cutoff)
+        lr_edge_attr= self.lr_edge_embedding(lr_edge_attr)
+        la_edge_attr = self.la_edge_embedding(la_edge_attr)
+        ar_edge_attr = self.ar_edge_embedding(ar_edge_attr)
+
+        for l in range(self.num_conv_layers):
+            # LIGAND updates
+            lig_edge_attr_ = torch.cat([lig_edge_attr, lig_node_attr[lig_edge_index[0], :self.ns], lig_node_attr[lig_edge_index[1], :self.ns]], -1)
+            lig_update = self.conv_layers[9*l](lig_node_attr, lig_edge_index, lig_edge_attr_, lig_edge_sh, edge_weight=lig_edge_weight)
+
+            lr_edge_attr_ = torch.cat([lr_edge_attr, lig_node_attr[lr_edge_index[0], :self.ns], rec_node_attr[lr_edge_index[1], :self.ns]], -1)
+            lr_update = self.conv_layers[9*l+1](rec_node_attr, lr_edge_index, lr_edge_attr_, lr_edge_sh,
+                                                out_nodes=lig_node_attr.shape[0], edge_weight=lr_edge_weight)
+
+            la_edge_attr_ = torch.cat([la_edge_attr, lig_node_attr[la_edge_index[0], :self.ns], atom_node_attr[la_edge_index[1], :self.ns]], -1)
+            la_update = self.conv_layers[9*l+2](atom_node_attr, la_edge_index, la_edge_attr_, la_edge_sh,
+                                                out_nodes=lig_node_attr.shape[0], edge_weight=la_edge_weight)
+
+            if l != self.num_conv_layers-1:  # last layer optimisation
+
+                # ATOM UPDATES
+                atom_edge_attr_ = torch.cat([atom_edge_attr, atom_node_attr[atom_edge_index[0], :self.ns], atom_node_attr[atom_edge_index[1], :self.ns]], -1)
+                atom_update = self.conv_layers[9*l+3](atom_node_attr, atom_edge_index, atom_edge_attr_, atom_edge_sh, edge_weight=atom_edge_weight)
+
+                al_edge_attr_ = torch.cat([la_edge_attr, atom_node_attr[la_edge_index[1], :self.ns], lig_node_attr[la_edge_index[0], :self.ns]], -1)
+                al_update = self.conv_layers[9*l+4](lig_node_attr, torch.flip(la_edge_index, dims=[0]), al_edge_attr_,
+                                                    la_edge_sh, out_nodes=atom_node_attr.shape[0], edge_weight=la_edge_weight)
+
+                ar_edge_attr_ = torch.cat([ar_edge_attr, atom_node_attr[ar_edge_index[0], :self.ns], rec_node_attr[ar_edge_index[1], :self.ns]],-1)
+                ar_update = self.conv_layers[9*l+5](rec_node_attr, ar_edge_index, ar_edge_attr_, ar_edge_sh, out_nodes=atom_node_attr.shape[0], edge_weight=ar_edge_weight)
+
+                # RECEPTOR updates
+                rec_edge_attr_ = torch.cat([rec_edge_attr, rec_node_attr[rec_edge_index[0], :self.ns], rec_node_attr[rec_edge_index[1], :self.ns]], -1)
+                rec_update = self.conv_layers[9*l+6](rec_node_attr, rec_edge_index, rec_edge_attr_, rec_edge_sh, edge_weight=rec_edge_weight)
+
+                rl_edge_attr_ = torch.cat([lr_edge_attr, rec_node_attr[lr_edge_index[1], :self.ns], lig_node_attr[lr_edge_index[0], :self.ns]], -1)
+                rl_update = self.conv_layers[9*l+7](lig_node_attr, torch.flip(lr_edge_index, dims=[0]), rl_edge_attr_,
+                                                    lr_edge_sh, out_nodes=rec_node_attr.shape[0], edge_weight=lr_edge_weight)
+
+                ra_edge_attr_ = torch.cat([ar_edge_attr, rec_node_attr[ar_edge_index[1], :self.ns], atom_node_attr[ar_edge_index[0], :self.ns]], -1)
+                ra_update = self.conv_layers[9*l+8](atom_node_attr, torch.flip(ar_edge_index, dims=[0]), ra_edge_attr_,
+                                                    ar_edge_sh, out_nodes=rec_node_attr.shape[0], edge_weight=ar_edge_weight)
+
+            # padding original features and update features with residual updates
+            lig_node_attr = F.pad(lig_node_attr, (0, lig_update.shape[-1] - lig_node_attr.shape[-1]))
+            lig_node_attr = lig_node_attr + lig_update + la_update + lr_update
+
+            if l != self.num_conv_layers - 1:  # last layer optimisation
+                atom_node_attr = F.pad(atom_node_attr, (0, atom_update.shape[-1] - atom_node_attr.shape[-1]))
+                atom_node_attr = atom_node_attr + atom_update + al_update + ar_update
+                rec_node_attr = F.pad(rec_node_attr, (0, rec_update.shape[-1] - rec_node_attr.shape[-1]))
+                rec_node_attr = rec_node_attr + rec_update + ra_update + rl_update
+
+        # confidence and affinity prediction
+        if self.confidence_mode:
+            scalar_lig_attr = torch.cat([lig_node_attr[:,:self.ns],lig_node_attr[:,-self.ns:]], dim=1) if self.num_conv_layers >= 3 else lig_node_attr[:,:self.ns]
+            confidence = self.confidence_predictor(scatter_mean(scalar_lig_attr, data['ligand'].batch if self.parallel == 1 else data['ligand'].batch_parallel, dim=0)).squeeze(dim=-1)
+
+            if self.parallel > 1:
+                confidence, affinity = confidence[:, 0], confidence[:, 1:]
+                confidence = confidence.reshape(data.num_graphs, self.parallel)
+                affinity = affinity.reshape(data.num_graphs, self.parallel, -1)
+                affinity = torch.cat([AGGREGATORS[agg](affinity) for agg in self.parallel_aggregators], dim=-1)
+                affinity = self.affinity_predictor(affinity).squeeze(dim=-1)
+                confidence = confidence, affinity
+            return confidence
+        assert self.parallel == 1
+
+        # compute translational and rotational score vectors
+        center_edge_index, center_edge_attr, center_edge_sh = self.build_center_conv_graph(data)
+        center_edge_attr = self.center_edge_embedding(center_edge_attr)
+        if self.fixed_center_conv:
+            center_edge_attr = torch.cat([center_edge_attr, lig_node_attr[center_edge_index[1], :self.ns]], -1)
+        else:
+            center_edge_attr = torch.cat([center_edge_attr, lig_node_attr[center_edge_index[0], :self.ns]], -1)
+        global_pred = self.final_conv(lig_node_attr, center_edge_index, center_edge_attr, center_edge_sh, out_nodes=data.num_graphs)
+
+        tr_pred = global_pred[:, :3] + (global_pred[:, 6:9] if not self.odd_parity else 0)
+        rot_pred = global_pred[:, 3:6] + (global_pred[:, 9:] if not self.odd_parity else 0)
+
+        if self.separate_noise_schedule:
+            data.graph_sigma_emb = torch.cat([self.timestep_emb_func(data.complex_t[noise_type]) for noise_type in ['tr', 'rot', 'tor']], dim=1)
+        elif self.asyncronous_noise_schedule:
+            data.graph_sigma_emb = self.timestep_emb_func(data.complex_t['t'])
+        else:  # tr rot and tor noise is all the same in this case
+            data.graph_sigma_emb = self.timestep_emb_func(data.complex_t['tr'])
+
+        # adjust the magniture of the score vectors
+        tr_norm = torch.linalg.vector_norm(tr_pred, dim=1).unsqueeze(1)
+        tr_pred = tr_pred / tr_norm * self.tr_final_layer(torch.cat([tr_norm, data.graph_sigma_emb], dim=1))
+
+        rot_norm = torch.linalg.vector_norm(rot_pred, dim=1).unsqueeze(1)
+        rot_pred = rot_pred / rot_norm * self.rot_final_layer(torch.cat([rot_norm, data.graph_sigma_emb], dim=1))
+
+        if self.scale_by_sigma:
+            tr_pred = tr_pred / tr_sigma.unsqueeze(1)
+            rot_pred = rot_pred * so3.score_norm(rot_sigma.cpu()).unsqueeze(1).to(data['ligand'].x.device)
+
+        if self.no_torsion or data['ligand'].edge_mask.sum() == 0: return tr_pred, rot_pred, torch.empty(0,device=self.device)
+
+        # torsional components
+        tor_bonds, tor_edge_index, tor_edge_attr, tor_edge_sh, tor_edge_weight = self.build_bond_conv_graph(data)
+        tor_bond_vec = data['ligand'].pos[tor_bonds[1]] - data['ligand'].pos[tor_bonds[0]]
+        tor_bond_attr = lig_node_attr[tor_bonds[0]] + lig_node_attr[tor_bonds[1]]
+
+        tor_bonds_sh = o3.spherical_harmonics("2e", tor_bond_vec, normalize=True, normalization='component')
+        tor_edge_sh = self.final_tp_tor(tor_edge_sh, tor_bonds_sh[tor_edge_index[0]])
+
+        tor_edge_attr = torch.cat([tor_edge_attr, lig_node_attr[tor_edge_index[1], :self.ns],
+                                   tor_bond_attr[tor_edge_index[0], :self.ns]], -1)
+        tor_pred = self.tor_bond_conv(lig_node_attr, tor_edge_index, tor_edge_attr, tor_edge_sh,
+                                  out_nodes=data['ligand'].edge_mask.sum(), reduce='mean', edge_weight=tor_edge_weight)
+        tor_pred = self.tor_final_layer(tor_pred).squeeze(1)
+        ligand_edge_store = self._ligand_edge_store(data)
+        edge_sigma = tor_sigma[data['ligand'].batch][ligand_edge_store.edge_index[0]][data['ligand'].edge_mask]
+
+        if self.scale_by_sigma:
+            tor_pred = tor_pred * torch.sqrt(torch.tensor(torus.score_norm(edge_sigma.cpu().numpy())).float()
+                                             .to(data['ligand'].x.device))
+        return tr_pred, rot_pred, tor_pred
+
+    def get_edge_weight(self, edge_vec, max_norm):
+        if self.smooth_edges:
+            normalised_norm = torch.clip(edge_vec.norm(dim=-1) * np.pi / max_norm, max=np.pi)
+            return 0.5 * (torch.cos(normalised_norm) + 1.0).unsqueeze(-1)
+        return 1.0
+
+    def build_lig_conv_graph(self, data):
+        # build the graph between ligand atoms
+        if self.separate_noise_schedule:
+            data['ligand'].node_sigma_emb = torch.cat(
+                [self.timestep_emb_func(data['ligand'].node_t[noise_type]) for noise_type in ['tr', 'rot', 'tor']],
+                dim=1)
+        elif self.asyncronous_noise_schedule:
+            data['ligand'].node_sigma_emb = self.timestep_emb_func(data['ligand'].node_t['t'])
+        else:
+            data['ligand'].node_sigma_emb = self.timestep_emb_func(
+                data['ligand'].node_t['tr'])  # tr rot and tor noise is all the same
+
+        if self.parallel == 1:
+            radius_edges = radius_graph(data['ligand'].pos, self.lig_max_radius, data['ligand'].batch)
+        else:
+            batches = torch.zeros(data.num_graphs, device=data['ligand'].x.device).long()
+            batches = batches.index_add(0, data['ligand'].batch, torch.ones(len(data['ligand'].batch), device=data['ligand'].x.device).long())
+            outer_batches = data.num_graphs
+            b = [torch.ones(batches[i].item()//self.parallel, device=data['ligand'].x.device).long() * (self.parallel * i + j)
+                 for i in range(outer_batches) for j in range(self.parallel)]
+            data['ligand'].batch_parallel = torch.cat(b)
+            radius_edges = radius_graph(data['ligand'].pos, self.lig_max_radius, data['ligand'].batch_parallel)
+        ligand_edge_store = self._ligand_edge_store(data)
+        edge_index = torch.cat([ligand_edge_store.edge_index, radius_edges], 1).long()
+        edge_attr = torch.cat([
+            ligand_edge_store.edge_attr,
+            torch.zeros(radius_edges.shape[-1], self.in_lig_edge_features, device=data['ligand'].x.device)
+        ], 0)
+
+        edge_sigma_emb = data['ligand'].node_sigma_emb[edge_index[0].long()]
+        edge_attr = torch.cat([edge_attr, edge_sigma_emb], 1)
+        node_attr = torch.cat([data['ligand'].x, data['ligand'].node_sigma_emb], 1)
+
+        src, dst = edge_index
+        edge_vec = data['ligand'].pos[dst.long()] - data['ligand'].pos[src.long()]
+        edge_length_emb = self.lig_distance_expansion(edge_vec.norm(dim=-1))
+
+        edge_attr = torch.cat([edge_attr, edge_length_emb], 1)
+        edge_sh = o3.spherical_harmonics(self.sh_irreps, edge_vec, normalize=True, normalization='component')
+        edge_weight = self.get_edge_weight(edge_vec, self.lig_max_radius)
+
+        return node_attr, edge_index, edge_attr, edge_sh, edge_weight
+
+    def build_rec_conv_graph(self, data):
+        # build the graph between receptor residues
+        if self.separate_noise_schedule:
+            data['receptor'].node_sigma_emb = torch.cat(
+                [self.timestep_emb_func(data['receptor'].node_t[noise_type]) for noise_type in ['tr', 'rot', 'tor']],
+                dim=1)
+        elif self.asyncronous_noise_schedule:
+            data['receptor'].node_sigma_emb = self.timestep_emb_func(data['receptor'].node_t['t'])
+        else:
+            data['receptor'].node_sigma_emb = self.timestep_emb_func(
+                data['receptor'].node_t['tr'])  # tr rot and tor noise is all the same
+        node_attr = torch.cat([data['receptor'].x, data['receptor'].node_sigma_emb], 1)
+
+        # this assumes the edges were already created in preprocessing since protein's structure is fixed
+        edge_index = self._receptor_edge_store(data).edge_index
+        src, dst = edge_index
+        edge_vec = data['receptor'].pos[dst.long()] - data['receptor'].pos[src.long()]
+        #assert torch.all(edge_vec.norm(dim=-1) < self.rec_max_radius)
+
+        edge_length_emb = self.rec_distance_expansion(edge_vec.norm(dim=-1))
+        edge_sigma_emb = data['receptor'].node_sigma_emb[edge_index[0].long()]
+        edge_attr = torch.cat([edge_sigma_emb, edge_length_emb], 1)
+        edge_sh = o3.spherical_harmonics(self.sh_irreps, edge_vec, normalize=True, normalization='component')
+        edge_weight = self.get_edge_weight(edge_vec, self.rec_max_radius)
+
+        return node_attr, edge_index, edge_attr, edge_sh, edge_weight
+
+    def build_atom_conv_graph(self, data):
+        # build the graph between receptor atoms
+        if self.separate_noise_schedule:
+            data['atom'].node_sigma_emb = torch.cat([self.timestep_emb_func(data['atom'].node_t[noise_type]) for noise_type in ['tr', 'rot', 'tor']],dim=1)
+        elif self.asyncronous_noise_schedule:
+            data['atom'].node_sigma_emb = self.timestep_emb_func(data['atom'].node_t['t'])
+        else:
+            data['atom'].node_sigma_emb = self.timestep_emb_func(data['atom'].node_t['tr'])  # tr rot and tor noise is all the same
+        node_attr = torch.cat([data['atom'].x, data['atom'].node_sigma_emb], 1)
+
+        # this assumes the edges were already created in preprocessing since protein's structure is fixed
+        edge_index = self._atom_edge_store(data).edge_index
+        src, dst = edge_index
+        edge_vec = data['atom'].pos[dst.long()] - data['atom'].pos[src.long()]
+
+        edge_length_emb = self.lig_distance_expansion(edge_vec.norm(dim=-1))
+        edge_sigma_emb = data['atom'].node_sigma_emb[edge_index[0].long()]
+        edge_attr = torch.cat([edge_sigma_emb, edge_length_emb], 1)
+        edge_sh = o3.spherical_harmonics(self.sh_irreps, edge_vec, normalize=True, normalization='component')
+        edge_weight = self.get_edge_weight(edge_vec, self.lig_max_radius)
+
+        return node_attr, edge_index, edge_attr, edge_sh, edge_weight
+
+    def build_cross_conv_graph(self, data, lr_cross_distance_cutoff):
+        # build the cross edges between ligan atoms, receptor residues and receptor atoms
+
+        # LIGAND to RECEPTOR
+        if torch.is_tensor(lr_cross_distance_cutoff):
+            # different cutoff for every graph
+            lr_edge_index = radius(data['receptor'].pos / lr_cross_distance_cutoff[data['receptor'].batch],
+                                data['ligand'].pos / lr_cross_distance_cutoff[data['ligand'].batch], 1,
+                                data['receptor'].batch, data['ligand'].batch, max_num_neighbors=10000)
+        else:
+            lr_edge_index = radius(data['receptor'].pos, data['ligand'].pos, lr_cross_distance_cutoff,
+                            data['receptor'].batch, data['ligand'].batch, max_num_neighbors=10000)
+
+        lr_edge_vec = data['receptor'].pos[lr_edge_index[1].long()] - data['ligand'].pos[lr_edge_index[0].long()]
+        lr_edge_length_emb = self.cross_distance_expansion(lr_edge_vec.norm(dim=-1))
+        lr_edge_sigma_emb = data['ligand'].node_sigma_emb[lr_edge_index[0].long()]
+        lr_edge_attr = torch.cat([lr_edge_sigma_emb, lr_edge_length_emb], 1)
+        lr_edge_sh = o3.spherical_harmonics(self.sh_irreps, lr_edge_vec, normalize=True, normalization='component')
+
+        cutoff_d = lr_cross_distance_cutoff[data['ligand'].batch[lr_edge_index[0]]].squeeze() \
+            if torch.is_tensor(lr_cross_distance_cutoff) else lr_cross_distance_cutoff
+        lr_edge_weight = self.get_edge_weight(lr_edge_vec, cutoff_d)
+
+        # LIGAND to ATOM
+        la_edge_index = radius(data['atom'].pos, data['ligand'].pos, self.lig_max_radius,
+                               data['atom'].batch, data['ligand'].batch, max_num_neighbors=10000)
+
+        la_edge_vec = data['atom'].pos[la_edge_index[1].long()] - data['ligand'].pos[la_edge_index[0].long()]
+        la_edge_length_emb = self.cross_distance_expansion(la_edge_vec.norm(dim=-1))
+        la_edge_sigma_emb = data['ligand'].node_sigma_emb[la_edge_index[0].long()]
+        la_edge_attr = torch.cat([la_edge_sigma_emb, la_edge_length_emb], 1)
+        la_edge_sh = o3.spherical_harmonics(self.sh_irreps, la_edge_vec, normalize=True, normalization='component')
+        la_edge_weight = self.get_edge_weight(la_edge_vec, self.lig_max_radius)
+
+        # ATOM to RECEPTOR
+        ar_edge_index = self._atom_receptor_edge_store(data).edge_index
+        ar_edge_vec = data['receptor'].pos[ar_edge_index[1].long()] - data['atom'].pos[ar_edge_index[0].long()]
+        ar_edge_length_emb = self.rec_distance_expansion(ar_edge_vec.norm(dim=-1))
+        ar_edge_sigma_emb = data['atom'].node_sigma_emb[ar_edge_index[0].long()]
+        ar_edge_attr = torch.cat([ar_edge_sigma_emb, ar_edge_length_emb], 1)
+        ar_edge_sh = o3.spherical_harmonics(self.sh_irreps, ar_edge_vec, normalize=True, normalization='component')
+        ar_edge_weight = 1
+
+        return lr_edge_index, lr_edge_attr, lr_edge_sh, lr_edge_weight, la_edge_index, la_edge_attr, \
+               la_edge_sh, la_edge_weight, ar_edge_index, ar_edge_attr, ar_edge_sh, ar_edge_weight
+
+    def build_center_conv_graph(self, data):
+        # build the filter for the convolution of the center with the ligand atoms
+        # for translational and rotational score
+        edge_index = torch.cat([data['ligand'].batch.unsqueeze(0), torch.arange(len(data['ligand'].batch)).to(data['ligand'].x.device).unsqueeze(0)], dim=0)
+
+        center_pos, count = torch.zeros((data.num_graphs, 3)).to(data['ligand'].x.device), torch.zeros((data.num_graphs, 3)).to(data['ligand'].x.device)
+        center_pos.index_add_(0, index=data['ligand'].batch, source=data['ligand'].pos)
+        center_pos = center_pos / torch.bincount(data['ligand'].batch).unsqueeze(1)
+
+        edge_vec = data['ligand'].pos[edge_index[1]] - center_pos[edge_index[0]]
+        edge_attr = self.center_distance_expansion(edge_vec.norm(dim=-1))
+        edge_sigma_emb = data['ligand'].node_sigma_emb[edge_index[1].long()]
+        edge_attr = torch.cat([edge_attr, edge_sigma_emb], 1)
+        edge_sh = o3.spherical_harmonics(self.sh_irreps, edge_vec, normalize=True, normalization='component')
+        return edge_index, edge_attr, edge_sh
+
+    def build_bond_conv_graph(self, data):
+        # build graph for the pseudotorque layer
+        bonds = self._ligand_edge_store(data).edge_index[:, data['ligand'].edge_mask].long()
+        bond_pos = (data['ligand'].pos[bonds[0]] + data['ligand'].pos[bonds[1]]) / 2
+        bond_batch = data['ligand'].batch[bonds[0]]
+        edge_index = radius(data['ligand'].pos, bond_pos, self.lig_max_radius, batch_x=data['ligand'].batch, batch_y=bond_batch)
+
+        edge_vec = data['ligand'].pos[edge_index[1]] - bond_pos[edge_index[0]]
+        edge_attr = self.lig_distance_expansion(edge_vec.norm(dim=-1))
+
+        edge_attr = self.final_edge_embedding(edge_attr)
+        edge_sh = o3.spherical_harmonics(self.sh_irreps, edge_vec, normalize=True, normalization='component')
+        edge_weight = self.get_edge_weight(edge_vec, self.lig_max_radius)
+
+        return bonds, edge_index, edge_attr, edge_sh, edge_weight

models/score_wrapper.py

@@ -0,0 +1,219 @@
+from argparse import Namespace
+from functools import partial
+from pathlib import Path
+
+import torch
+import yaml
+from torch_geometric.nn.data_parallel import DataParallel
+
+from onescience.utils.diffdock.diffusion_utils import get_timestep_embedding, t_to_sigma as t_to_sigma_compl
+from onescience.utils.diffdock.utils import ExponentialMovingAverage
+
+from .aa_model import AAModel
+from .cg_model import CGModel
+from .old_aa_model import AAOldModel
+
+
+_LM_EMBEDDING_KEYS = (
+    "moad_esm_embeddings_path",
+    "pdbbind_esm_embeddings_path",
+    "pdbsidechain_esm_embeddings_path",
+    "esm_embeddings_path",
+    "esm_embeddings_model",
+)
+
+
+def load_model_args(model_dir):
+    model_dir = Path(model_dir)
+    config_path = model_dir / "model_parameters.yml"
+    with config_path.open("r", encoding="utf-8") as handle:
+        config = yaml.full_load(handle) or {}
+    return Namespace(**config)
+
+
+def model_uses_lm_embeddings(model_args):
+    return any(getattr(model_args, key, None) is not None for key in _LM_EMBEDDING_KEYS)
+
+
+def _has_arg(args, name):
+    try:
+        return name in args
+    except TypeError:
+        return hasattr(args, name)
+
+
+def _get_arg(args, name, default=None):
+    if isinstance(args, dict):
+        return args.get(name, default)
+    if _has_arg(args, name):
+        return getattr(args, name)
+    return default
+
+
+def get_model(args, device, t_to_sigma, no_parallel=False, confidence_mode=False, old=False):
+    timestep_emb_func = get_timestep_embedding(
+        embedding_type=_get_arg(args, "embedding_type", "sinusoidal"),
+        embedding_dim=args.sigma_embed_dim,
+        embedding_scale=_get_arg(args, "embedding_scale", 10000),
+    )
+
+    all_atoms = _get_arg(args, "all_atoms", False)
+    if old and not all_atoms:
+        raise NotImplementedError(
+            "The old coarse-grained DiffDock model path is not migrated yet. "
+            "Use old=True only with all_atoms=True, or migrate old_cg_model.py first."
+        )
+
+    lm_embedding_type = None
+    if (
+        _get_arg(args, "moad_esm_embeddings_path") is not None
+        or _get_arg(args, "pdbbind_esm_embeddings_path") is not None
+        or _get_arg(args, "pdbsidechain_esm_embeddings_path") is not None
+        or _get_arg(args, "esm_embeddings_path") is not None
+    ):
+        lm_embedding_type = "precomputed"
+    if _get_arg(args, "esm_embeddings_model") is not None:
+        lm_embedding_type = args.esm_embeddings_model
+
+    if old:
+        model_class = AAOldModel
+    elif all_atoms:
+        model_class = AAModel
+    else:
+        model_class = CGModel
+
+    model_kwargs = dict(
+        t_to_sigma=t_to_sigma,
+        device=device,
+        no_torsion=args.no_torsion,
+        timestep_emb_func=timestep_emb_func,
+        num_conv_layers=args.num_conv_layers,
+        lig_max_radius=args.max_radius,
+        scale_by_sigma=args.scale_by_sigma,
+        sigma_embed_dim=args.sigma_embed_dim,
+        norm_by_sigma=_get_arg(args, "norm_by_sigma", False),
+        ns=args.ns,
+        nv=args.nv,
+        distance_embed_dim=args.distance_embed_dim,
+        cross_distance_embed_dim=args.cross_distance_embed_dim,
+        batch_norm=not args.no_batch_norm,
+        dropout=args.dropout,
+        use_second_order_repr=args.use_second_order_repr,
+        cross_max_distance=args.cross_max_distance,
+        dynamic_max_cross=args.dynamic_max_cross,
+        smooth_edges=_get_arg(args, "smooth_edges", False),
+        odd_parity=_get_arg(args, "odd_parity", False),
+        lm_embedding_type=lm_embedding_type,
+        confidence_mode=confidence_mode,
+        confidence_dropout=_get_arg(args, "confidence_dropout", 0.0),
+        confidence_no_batchnorm=_get_arg(args, "confidence_no_batchnorm", False),
+        affinity_prediction=_get_arg(args, "affinity_prediction", False),
+        parallel=_get_arg(args, "parallel", 1),
+        num_confidence_outputs=(
+            len(args.rmsd_classification_cutoff) + 1
+            if isinstance(_get_arg(args, "rmsd_classification_cutoff"), list)
+            else 1
+        ),
+        atom_num_confidence_outputs=(
+            len(args.atom_rmsd_classification_cutoff) + 1
+            if isinstance(_get_arg(args, "atom_rmsd_classification_cutoff"), list)
+            else 1
+        ),
+        parallel_aggregators=_get_arg(args, "parallel_aggregators", ""),
+        fixed_center_conv=not _get_arg(args, "not_fixed_center_conv", False),
+        no_aminoacid_identities=_get_arg(args, "no_aminoacid_identities", False),
+        include_miscellaneous_atoms=_get_arg(args, "include_miscellaneous_atoms", False),
+        sh_lmax=_get_arg(args, "sh_lmax", 2),
+        differentiate_convolutions=not _get_arg(args, "no_differentiate_convolutions", False),
+        tp_weights_layers=_get_arg(args, "tp_weights_layers", 2),
+        num_prot_emb_layers=_get_arg(args, "num_prot_emb_layers", 0),
+        reduce_pseudoscalars=_get_arg(args, "reduce_pseudoscalars", False),
+        embed_also_ligand=_get_arg(args, "embed_also_ligand", False),
+        atom_confidence=_get_arg(args, "atom_confidence_loss_weight", 0.0) > 0.0,
+        sidechain_pred=(
+            (_has_arg(args, "sidechain_loss_weight") and args.sidechain_loss_weight > 0)
+            or (_has_arg(args, "backbone_loss_weight") and args.backbone_loss_weight > 0)
+        ),
+        depthwise_convolution=_get_arg(args, "depthwise_convolution", False),
+    )
+    if model_class is AAModel:
+        model_kwargs["crop_beyond"] = _get_arg(args, "crop_beyond", None)
+    elif model_class is AAOldModel:
+        for key in (
+            "atom_num_confidence_outputs",
+            "differentiate_convolutions",
+            "tp_weights_layers",
+            "num_prot_emb_layers",
+            "reduce_pseudoscalars",
+            "embed_also_ligand",
+            "atom_confidence",
+            "sidechain_pred",
+            "depthwise_convolution",
+        ):
+            model_kwargs.pop(key, None)
+        model_kwargs["lm_embedding_type"] = (
+            "esm" if _get_arg(args, "esm_embeddings_path") is not None else None
+        )
+        model_kwargs["use_old_atom_encoder"] = _get_arg(args, "use_old_atom_encoder", True)
+
+    model = model_class(**model_kwargs)
+
+    if device.type == "cuda" and not no_parallel and _get_arg(args, "dataset") != "torsional":
+        model = DataParallel(model)
+    model.to(device)
+    return model
+
+
+def build_score_model(
+    model_args,
+    device,
+    no_parallel=False,
+    confidence_mode=False,
+    old=False,
+):
+    t_to_sigma = partial(t_to_sigma_compl, args=model_args)
+    model = get_model(
+        model_args,
+        device,
+        t_to_sigma=t_to_sigma,
+        no_parallel=no_parallel,
+        confidence_mode=confidence_mode,
+        old=old,
+    )
+    return model, t_to_sigma
+
+
+def load_score_model(
+    model_dir,
+    ckpt,
+    device=None,
+    no_parallel=True,
+    confidence_mode=False,
+    old=False,
+    strict=True,
+):
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model_args = load_model_args(model_dir)
+    model, t_to_sigma = build_score_model(
+        model_args=model_args,
+        device=device,
+        no_parallel=no_parallel,
+        confidence_mode=confidence_mode,
+        old=old,
+    )
+
+    checkpoint_path = Path(model_dir) / ckpt
+    state_dict = torch.load(checkpoint_path, map_location=torch.device("cpu"))
+    if isinstance(state_dict, dict) and "model" in state_dict and "optimizer" in state_dict:
+        model.load_state_dict(state_dict["model"], strict=strict)
+        if "ema_weights" in state_dict and getattr(model_args, "ema_rate", None) is not None:
+            ema_weights = ExponentialMovingAverage(model.parameters(), decay=model_args.ema_rate)
+            ema_weights.load_state_dict(state_dict["ema_weights"], device=device)
+            ema_weights.copy_to(model.parameters())
+    else:
+        model.load_state_dict(state_dict, strict=strict)
+    model = model.to(device)
+    model.eval()
+    return model, model_args, t_to_sigma

models/tensor_layers.py

@@ -0,0 +1,459 @@
+import logging
+import os
+from typing import List, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from e3nn import o3
+from e3nn.nn import BatchNorm
+from e3nn.o3 import Linear, TensorProduct
+from torch_scatter import scatter, scatter_mean
+
+from .layers import FCBlock
+
+
+def get_irrep_seq(ns, nv, use_second_order_repr, reduce_pseudoscalars):
+    if use_second_order_repr:
+        irrep_seq = [
+            f"{ns}x0e",
+            f"{ns}x0e + {nv}x1o + {nv}x2e",
+            f"{ns}x0e + {nv}x1o + {nv}x2e + {nv}x1e + {nv}x2o",
+            f"{ns}x0e + {nv}x1o + {nv}x2e + {nv}x1e + {nv}x2o + {nv if reduce_pseudoscalars else ns}x0o",
+        ]
+    else:
+        irrep_seq = [
+            f"{ns}x0e",
+            f"{ns}x0e + {nv}x1o",
+            f"{ns}x0e + {nv}x1o + {nv}x1e",
+            f"{ns}x0e + {nv}x1o + {nv}x1e + {nv if reduce_pseudoscalars else ns}x0o",
+        ]
+    return irrep_seq
+
+
+def irrep_to_size(irrep):
+    irreps = irrep.split(" + ")
+    size = 0
+    for ir in irreps:
+        m, (l, p) = ir.split("x")
+        size += int(m) * (2 * int(l) + 1)
+    return size
+
+
+class FasterTensorProduct(torch.nn.Module):
+    # Implemented by Bowen Jing
+    def __init__(self, in_irreps, sh_irreps, out_irreps, **kwargs):
+        super().__init__()
+        assert (
+            o3.Irreps(sh_irreps) == o3.Irreps("1x0e+1x1o")
+        ), "sh_irreps don't look like 1st order spherical harmonics"
+        self.in_irreps = o3.Irreps(in_irreps)
+        self.out_irreps = o3.Irreps(out_irreps)
+
+        in_muls = {"0e": 0, "1o": 0, "1e": 0, "0o": 0}
+        out_muls = {"0e": 0, "1o": 0, "1e": 0, "0o": 0}
+        for (m, ir) in self.in_irreps:
+            in_muls[str(ir)] = m
+        for (m, ir) in self.out_irreps:
+            out_muls[str(ir)] = m
+
+        self.weight_shapes = {
+            "0e": (in_muls["0e"] + in_muls["1o"], out_muls["0e"]),
+            "1o": (in_muls["0e"] + in_muls["1o"] + in_muls["1e"], out_muls["1o"]),
+            "1e": (in_muls["1o"] + in_muls["1e"] + in_muls["0o"], out_muls["1e"]),
+            "0o": (in_muls["1e"] + in_muls["0o"], out_muls["0o"]),
+        }
+        self.weight_numel = sum(a * b for (a, b) in self.weight_shapes.values())
+
+    def forward(self, in_, sh, weight):
+        in_dict, out_dict = {}, {"0e": [], "1o": [], "1e": [], "0o": []}
+        for (m, ir), sl in zip(self.in_irreps, self.in_irreps.slices()):
+            in_dict[str(ir)] = in_[..., sl]
+            if ir[0] == 1:
+                in_dict[str(ir)] = in_dict[str(ir)].reshape(
+                    list(in_dict[str(ir)].shape)[:-1] + [-1, 3]
+                )
+        sh_0e, sh_1o = sh[..., 0], sh[..., 1:]
+        if "0e" in in_dict:
+            out_dict["0e"].append(in_dict["0e"] * sh_0e.unsqueeze(-1))
+            out_dict["1o"].append(in_dict["0e"].unsqueeze(-1) * sh_1o.unsqueeze(-2))
+        if "1o" in in_dict:
+            out_dict["0e"].append((in_dict["1o"] * sh_1o.unsqueeze(-2)).sum(-1) / np.sqrt(3))
+            out_dict["1o"].append(in_dict["1o"] * sh_0e.unsqueeze(-1).unsqueeze(-1))
+            out_dict["1e"].append(
+                torch.linalg.cross(in_dict["1o"], sh_1o.unsqueeze(-2), dim=-1) / np.sqrt(2)
+            )
+        if "1e" in in_dict:
+            out_dict["1o"].append(
+                torch.linalg.cross(in_dict["1e"], sh_1o.unsqueeze(-2), dim=-1) / np.sqrt(2)
+            )
+            out_dict["1e"].append(in_dict["1e"] * sh_0e.unsqueeze(-1).unsqueeze(-1))
+            out_dict["0o"].append((in_dict["1e"] * sh_1o.unsqueeze(-2)).sum(-1) / np.sqrt(3))
+        if "0o" in in_dict:
+            out_dict["1e"].append(in_dict["0o"].unsqueeze(-1) * sh_1o.unsqueeze(-2))
+            out_dict["0o"].append(in_dict["0o"] * sh_0e.unsqueeze(-1))
+
+        weight_dict = {}
+        start = 0
+        for key in self.weight_shapes:
+            in_, out = self.weight_shapes[key]
+            weight_dict[key] = weight[..., start : start + in_ * out].reshape(
+                list(weight.shape)[:-1] + [in_, out]
+            ) / np.sqrt(in_)
+            start += in_ * out
+
+        if out_dict["0e"]:
+            out_dict["0e"] = torch.cat(out_dict["0e"], dim=-1)
+            out_dict["0e"] = torch.matmul(
+                out_dict["0e"].unsqueeze(-2), weight_dict["0e"]
+            ).squeeze(-2)
+
+        if out_dict["1o"]:
+            out_dict["1o"] = torch.cat(out_dict["1o"], dim=-2)
+            out_dict["1o"] = (
+                out_dict["1o"].unsqueeze(-2) * weight_dict["1o"].unsqueeze(-1)
+            ).sum(-3)
+            out_dict["1o"] = out_dict["1o"].reshape(list(out_dict["1o"].shape)[:-2] + [-1])
+
+        if out_dict["1e"]:
+            out_dict["1e"] = torch.cat(out_dict["1e"], dim=-2)
+            out_dict["1e"] = (
+                out_dict["1e"].unsqueeze(-2) * weight_dict["1e"].unsqueeze(-1)
+            ).sum(-3)
+            out_dict["1e"] = out_dict["1e"].reshape(list(out_dict["1e"].shape)[:-2] + [-1])
+
+        if out_dict["0o"]:
+            out_dict["0o"] = torch.cat(out_dict["0o"], dim=-1)
+            out_dict["0o"] = torch.matmul(
+                out_dict["0o"].unsqueeze(-2), weight_dict["0o"]
+            ).squeeze(-2)
+
+        out = []
+        for _, ir in self.out_irreps:
+            out.append(out_dict[str(ir)])
+        return torch.cat(out, dim=-1)
+
+
+def tp_scatter_simple(
+    tp,
+    fc_layer,
+    node_attr,
+    edge_index,
+    edge_attr,
+    edge_sh,
+    out_nodes=None,
+    reduce="mean",
+    edge_weight=1.0,
+):
+    """
+    Perform TensorProduct + scatter operation, aka graph convolution.
+
+    This function is only for edge_groups == 1. For multiple edge groups, and for larger graphs,
+    use tp_scatter_multigroup instead.
+    """
+
+    assert isinstance(
+        edge_attr, torch.Tensor
+    ), "This function is only for a single edge group, so edge_attr must be a tensor and not a list."
+
+    _device = node_attr.device
+    _dtype = node_attr.dtype
+    edge_src, edge_dst = edge_index
+    out_irreps = fc_layer(edge_attr).to(_device).to(_dtype)
+    out_irreps.mul_(edge_weight)
+    tp = tp(node_attr[edge_dst], edge_sh, out_irreps)
+    out_nodes = out_nodes or node_attr.shape[0]
+    out = scatter(tp, edge_src, dim=0, dim_size=out_nodes, reduce=reduce)
+    return out
+
+
+def tp_scatter_multigroup(
+    tp: o3.TensorProduct,
+    fc_layer: Union[nn.Module, nn.ModuleList],
+    node_attr: torch.Tensor,
+    edge_index: torch.Tensor,
+    edge_attr_groups: List[torch.Tensor],
+    edge_sh: torch.Tensor,
+    out_nodes=None,
+    reduce="mean",
+    edge_weight=1.0,
+):
+    """
+    Perform TensorProduct + scatter operation, aka graph convolution.
+
+    To keep the peak memory usage reasonably low, this function does not concatenate the edge_attr_groups.
+    Rather, we sum the output of the tensor product for each edge group, and then divide by the number of edges
+    """
+
+    assert isinstance(edge_attr_groups, list), "This function is only for a list of edge groups"
+    assert reduce in {"mean", "sum"}, "Only 'mean' and 'sum' are supported for reduce"
+
+    _device = node_attr.device
+    _dtype = node_attr.dtype
+    edge_src, edge_dst = edge_index
+    edge_attr_lengths = [_edge_attr.shape[0] for _edge_attr in edge_attr_groups]
+    total_rows = sum(edge_attr_lengths)
+    assert (
+        total_rows == edge_index.shape[1]
+    ), "Sum of edge_attr_groups must be equal to edge_index.shape[1]"
+    num_edge_groups = len(edge_attr_groups)
+    edge_weight_is_indexable = hasattr(edge_weight, "__getitem__")
+
+    out_nodes = out_nodes or node_attr.shape[0]
+    total_output_dim = sum([x.dim for x in tp.irreps_out])
+    final_out = torch.zeros((out_nodes, total_output_dim), device=_device, dtype=_dtype)
+    div_factors = torch.zeros(out_nodes, device=_device, dtype=_dtype)
+
+    cur_start = 0
+    for ii in range(num_edge_groups):
+        cur_length = edge_attr_lengths[ii]
+        cur_end = cur_start + cur_length
+        cur_edge_range = slice(cur_start, cur_end)
+        cur_edge_src, cur_edge_dst = edge_src[cur_edge_range], edge_dst[cur_edge_range]
+
+        cur_fc = fc_layer[ii] if isinstance(fc_layer, nn.ModuleList) else fc_layer
+        cur_out_irreps = cur_fc(edge_attr_groups[ii])
+        if edge_weight_is_indexable:
+            cur_out_irreps.mul_(edge_weight[cur_edge_range])
+        else:
+            cur_out_irreps.mul_(edge_weight)
+
+        summand = tp(node_attr[cur_edge_dst, :], edge_sh[cur_edge_range, :], cur_out_irreps)
+        final_out += scatter(summand, cur_edge_src, dim=0, dim_size=out_nodes, reduce="sum")
+        div_factors += torch.bincount(cur_edge_src, minlength=out_nodes)
+
+        cur_start = cur_end
+
+        del cur_out_irreps, summand
+
+    if reduce == "mean":
+        div_factors = torch.clamp(div_factors, torch.finfo(_dtype).eps)
+        final_out = final_out / div_factors[:, None]
+
+    return final_out
+
+
+class TensorProductConvLayer(torch.nn.Module):
+    def __init__(
+        self,
+        in_irreps,
+        sh_irreps,
+        out_irreps,
+        n_edge_features,
+        residual=True,
+        batch_norm=True,
+        dropout=0.0,
+        hidden_features=None,
+        faster=False,
+        edge_groups=1,
+        tp_weights_layers=2,
+        activation="relu",
+        depthwise=False,
+    ):
+        super(TensorProductConvLayer, self).__init__()
+        self.in_irreps = in_irreps
+        self.out_irreps = out_irreps
+        self.sh_irreps = sh_irreps
+        self.residual = residual
+        self.edge_groups = edge_groups
+        self.out_size = irrep_to_size(out_irreps)
+        self.depthwise = depthwise
+        if hidden_features is None:
+            hidden_features = n_edge_features
+
+        if depthwise:
+            in_irreps = o3.Irreps(in_irreps)
+            sh_irreps = o3.Irreps(sh_irreps)
+            out_irreps = o3.Irreps(out_irreps)
+
+            irreps_mid = []
+            instructions = []
+            for i, (mul, ir_in) in enumerate(in_irreps):
+                for j, (_, ir_edge) in enumerate(sh_irreps):
+                    for ir_out in ir_in * ir_edge:
+                        if ir_out in out_irreps:
+                            k = len(irreps_mid)
+                            irreps_mid.append((mul, ir_out))
+                            instructions.append((i, j, k, "uvu", True))
+
+            irreps_mid = o3.Irreps(irreps_mid)
+            irreps_mid, p, _ = irreps_mid.sort()
+
+            instructions = [
+                (i_in1, i_in2, p[i_out], mode, train)
+                for i_in1, i_in2, i_out, mode, train in instructions
+            ]
+
+            self.tp = TensorProduct(
+                in_irreps,
+                sh_irreps,
+                irreps_mid,
+                instructions,
+                shared_weights=False,
+                internal_weights=False,
+            )
+
+            self.linear_2 = Linear(
+                irreps_in=irreps_mid.simplify(),
+                irreps_out=out_irreps,
+                internal_weights=True,
+                shared_weights=True,
+            )
+
+        else:
+            if faster:
+                print("Faster Tensor Product")
+                self.tp = FasterTensorProduct(in_irreps, sh_irreps, out_irreps)
+            else:
+                self.tp = o3.FullyConnectedTensorProduct(
+                    in_irreps, sh_irreps, out_irreps, shared_weights=False
+                )
+
+        if edge_groups == 1:
+            self.fc = FCBlock(
+                n_edge_features,
+                hidden_features,
+                self.tp.weight_numel,
+                tp_weights_layers,
+                dropout,
+                activation,
+            )
+        else:
+            self.fc = [
+                FCBlock(
+                    n_edge_features,
+                    hidden_features,
+                    self.tp.weight_numel,
+                    tp_weights_layers,
+                    dropout,
+                    activation,
+                )
+                for _ in range(edge_groups)
+            ]
+            self.fc = nn.ModuleList(self.fc)
+
+        self.batch_norm = BatchNorm(out_irreps) if batch_norm else None
+
+    def forward(
+        self,
+        node_attr,
+        edge_index,
+        edge_attr,
+        edge_sh,
+        out_nodes=None,
+        reduce="mean",
+        edge_weight=1.0,
+    ):
+        if edge_index.shape[1] == 0 and node_attr.shape[0] == 0:
+            raise ValueError("No edges and no nodes")
+
+        _dtype = node_attr.dtype
+        if edge_index.shape[1] == 0:
+            out = torch.zeros((node_attr.shape[0], self.out_size), dtype=_dtype, device=node_attr.device)
+        else:
+            if self.edge_groups == 1:
+                out = tp_scatter_simple(
+                    self.tp,
+                    self.fc,
+                    node_attr,
+                    edge_index,
+                    edge_attr,
+                    edge_sh,
+                    out_nodes,
+                    reduce,
+                    edge_weight,
+                )
+            else:
+                out = tp_scatter_multigroup(
+                    self.tp,
+                    self.fc,
+                    node_attr,
+                    edge_index,
+                    edge_attr,
+                    edge_sh,
+                    out_nodes,
+                    reduce,
+                    edge_weight,
+                )
+
+            if self.depthwise:
+                out = self.linear_2(out)
+
+            if self.batch_norm:
+                out = self.batch_norm(out)
+
+        if self.residual:
+            padded = F.pad(node_attr, (0, out.shape[-1] - node_attr.shape[-1]))
+            out = out + padded
+
+        out = out.to(_dtype)
+        return out
+
+
+class OldTensorProductConvLayer(torch.nn.Module):
+    def __init__(
+        self,
+        in_irreps,
+        sh_irreps,
+        out_irreps,
+        n_edge_features,
+        residual=True,
+        batch_norm=True,
+        dropout=0.0,
+        hidden_features=None,
+    ):
+        super(OldTensorProductConvLayer, self).__init__()
+        self.in_irreps = in_irreps
+        self.out_irreps = out_irreps
+        self.sh_irreps = sh_irreps
+        self.residual = residual
+        if hidden_features is None:
+            hidden_features = n_edge_features
+
+        self.tp = tp = o3.FullyConnectedTensorProduct(
+            in_irreps, sh_irreps, out_irreps, shared_weights=False
+        )
+
+        self.fc = nn.Sequential(
+            nn.Linear(n_edge_features, hidden_features),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_features, tp.weight_numel),
+        )
+        self.batch_norm = BatchNorm(out_irreps) if batch_norm else None
+
+    def forward(
+        self, node_attr, edge_index, edge_attr, edge_sh, out_nodes=None, reduce="mean", edge_weight=1.0
+    ):
+        edge_chunk_size = 100_000
+        num_edges = edge_attr.shape[0]
+        num_chunks = (
+            (num_edges // edge_chunk_size)
+            if num_edges % edge_chunk_size == 0
+            else (num_edges // edge_chunk_size) + 1
+        )
+        edge_ranges = np.array_split(np.arange(num_edges), num_chunks)
+        edge_attr_groups = [edge_attr[cur_range] for cur_range in edge_ranges]
+
+        out = tp_scatter_multigroup(
+            self.tp,
+            self.fc,
+            node_attr,
+            edge_index,
+            edge_attr_groups,
+            edge_sh,
+            out_nodes,
+            reduce,
+            edge_weight,
+        )
+
+        if self.residual:
+            padded = F.pad(node_attr, (0, out.shape[-1] - node_attr.shape[-1]))
+            out = out + padded
+
+        if self.batch_norm:
+            out = self.batch_norm(out)
+
+        out = out.to(node_attr.dtype)
+        return out

scripts/__init__.py

@@ -0,0 +1 @@
+__all__ = []

scripts/evaluate.py

@@ -0,0 +1,5 @@
+from onescience.utils.diffdock.evaluate import main
+
+
+if __name__ == "__main__":
+    main()

scripts/infer.sh

@@ -0,0 +1,147 @@
+#!/bin/bash
+set -euo pipefail
+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"
+export LD_LIBRARY_PATH=${ROCM_PATH}/opencl/lib:$LD_LIBRARY_PATH
+
+SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+EXAMPLE_DIR=$(cd "${SCRIPT_DIR}/.." && pwd)
+REPO_ROOT=$(cd "${SCRIPT_DIR}/../../../.." && pwd)
+
+# source "${REPO_ROOT}/env.sh"
+
+if [[ -n "${ROCM_PATH:-}" && -f "${ROCM_PATH}/cuda/env.sh" ]]; then
+  source "${ROCM_PATH}/cuda/env.sh"
+fi
+
+export PYTHONPATH="${REPO_ROOT}/src:${REPO_ROOT}:${PYTHONPATH:-}"
+export HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-0}"
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-${HIP_VISIBLE_DEVICES}}"
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}"
+export DIFFDOCK_RADIUS_ON_CPU="${RADIUS_ON_CPU:-false}"
+
+DIFFDOCK_DATA_ROOT="${DIFFDOCK_DATA_ROOT:-${ONESCIENCE_DATASETS_DIR}/diffdock}"
+export TORCH_HOME="${TORCH_HOME:-${DIFFDOCK_DATA_ROOT}/torch_home}"
+
+SCORE_MODEL_DIR="${SCORE_MODEL_DIR:-${DIFFDOCK_DATA_ROOT}/score_model}"
+SCORE_CKPT="${SCORE_CKPT:-best_ema_inference_epoch_model.pt}"
+CONFIDENCE_MODEL_DIR="${CONFIDENCE_MODEL_DIR:-${DIFFDOCK_DATA_ROOT}/confidence_model}"
+CONFIDENCE_CKPT="${CONFIDENCE_CKPT:-best_model_epoch75.pt}"
+ENABLE_CONFIDENCE="${ENABLE_CONFIDENCE:-true}"
+OLD_CONFIDENCE_MODEL="${OLD_CONFIDENCE_MODEL:-true}"
+
+OUT_DIR="${OUT_DIR:-${EXAMPLE_DIR}/outputs/scnet_inference}"
+CONFIG_PATH="${CONFIG_PATH:-${OUT_DIR}/inference_config.yml}"
+mkdir -p "${OUT_DIR}"
+OUT_DIR=$(cd "${OUT_DIR}" && pwd)
+CONFIG_DIR=$(dirname "${CONFIG_PATH}")
+CONFIG_NAME=$(basename "${CONFIG_PATH}")
+mkdir -p "${CONFIG_DIR}"
+CONFIG_PATH=$(cd "${CONFIG_DIR}" && pwd)/"${CONFIG_NAME}"
+
+DEFAULT_SHARED_CSV="${DIFFDOCK_DATA_ROOT}/datasets/inferdata/protein_ligand_example.csv"
+PROTEIN_LIGAND_CSV="${PROTEIN_LIGAND_CSV:-}"
+if [[ -z "${PROTEIN_LIGAND_CSV}" && -f "${DEFAULT_SHARED_CSV}" ]]; then
+  PROTEIN_LIGAND_CSV="${DEFAULT_SHARED_CSV}"
+fi
+COMPLEX_NAME="${COMPLEX_NAME:-6o5u_test}"
+PROTEIN_PATH="${PROTEIN_PATH:-${EXAMPLE_DIR}/data/6o5u_protein_processed.pdb}"
+PROTEIN_SEQUENCE="${PROTEIN_SEQUENCE:-}"
+LIGAND_DESCRIPTION="${LIGAND_DESCRIPTION:-${EXAMPLE_DIR}/data/6o5u_ligand.sdf}"
+
+DEVICE="${DEVICE:-auto}"
+SAMPLES_PER_COMPLEX="${SAMPLES_PER_COMPLEX:-10}"
+BATCH_SIZE="${BATCH_SIZE:-10}"
+INFERENCE_STEPS="${INFERENCE_STEPS:-20}"
+ACTUAL_STEPS="${ACTUAL_STEPS:-}"
+NO_RANDOM="${NO_RANDOM:-false}"
+NO_FINAL_STEP_NOISE="${NO_FINAL_STEP_NOISE:-true}"
+CROP_BEYOND="${CROP_BEYOND:-}"
+
+yaml_value() {
+  if [[ -z "${1:-}" || "${1}" == "null" ]]; then
+    printf "null"
+  else
+    local value
+    value=$(printf "%s" "$1" | sed "s/'/''/g")
+    printf "'%s'" "$value"
+  fi
+}
+
+yaml_bool() {
+  if [[ "${1,,}" == "true" ]]; then
+    printf "true"
+  else
+    printf "false"
+  fi
+}
+
+if [[ "${ENABLE_CONFIDENCE,,}" == "true" ]]; then
+  CONFIDENCE_MODEL_VALUE=$(yaml_value "${CONFIDENCE_MODEL_DIR}")
+else
+  CONFIDENCE_MODEL_VALUE="null"
+fi
+
+cat > "${CONFIG_PATH}" <<EOF
+runtime:
+  device: $(yaml_value "${DEVICE}")
+  loglevel: INFO
+  out_dir: $(yaml_value "${OUT_DIR}")
+
+model:
+  model_dir: $(yaml_value "${SCORE_MODEL_DIR}")
+  ckpt: $(yaml_value "${SCORE_CKPT}")
+  old_score_model: false
+
+confidence:
+  confidence_model_dir: ${CONFIDENCE_MODEL_VALUE}
+  confidence_ckpt: $(yaml_value "${CONFIDENCE_CKPT}")
+  old_confidence_model: $(yaml_bool "${OLD_CONFIDENCE_MODEL}")
+
+input:
+  protein_ligand_csv: $(yaml_value "${PROTEIN_LIGAND_CSV}")
+  complex_name: $(yaml_value "${COMPLEX_NAME}")
+  protein_path: $(yaml_value "${PROTEIN_PATH}")
+  protein_sequence: $(yaml_value "${PROTEIN_SEQUENCE}")
+  ligand_description: $(yaml_value "${LIGAND_DESCRIPTION}")
+  lm_embeddings: null
+  crop_beyond: $(yaml_value "${CROP_BEYOND}")
+
+sampling:
+  samples_per_complex: ${SAMPLES_PER_COMPLEX}
+  batch_size: ${BATCH_SIZE}
+  inference_steps: ${INFERENCE_STEPS}
+  actual_steps: $(yaml_value "${ACTUAL_STEPS}")
+  sigma_schedule: expbeta
+  inf_sched_alpha: 1.0
+  inf_sched_beta: 1.0
+  no_random: $(yaml_bool "${NO_RANDOM}")
+  no_final_step_noise: $(yaml_bool "${NO_FINAL_STEP_NOISE}")
+  ode: false
+  choose_residue: false
+  initial_noise_std_proportion: 1.0
+  temp_sampling_tr: 1.0
+  temp_psi_tr: 0.0
+  temp_sigma_data_tr: 0.5
+  temp_sampling_rot: 1.0
+  temp_psi_rot: 0.0
+  temp_sigma_data_rot: 0.5
+  temp_sampling_tor: 1.0
+  temp_psi_tor: 0.0
+  temp_sigma_data_tor: 0.5
+EOF
+
+echo "DiffDock inference config: ${CONFIG_PATH}"
+echo "Output directory: ${OUT_DIR}"
+echo "Score model: ${SCORE_MODEL_DIR}/${SCORE_CKPT}"
+echo "Confidence rerank: ${ENABLE_CONFIDENCE}"
+echo "TORCH_HOME: ${TORCH_HOME}"
+echo "Radius on CPU: ${DIFFDOCK_RADIUS_ON_CPU}"
+if [[ -n "${PROTEIN_LIGAND_CSV}" ]]; then
+  echo "Input CSV: ${PROTEIN_LIGAND_CSV}"
+else
+  echo "Single input: ${PROTEIN_PATH} + ${LIGAND_DESCRIPTION}"
+fi
+
+cd "${REPO_ROOT}"
+python "${SCRIPT_DIR}/sample_diffdock.py" --config "${CONFIG_PATH}"

scripts/infer_slurm.sbatch

@@ -0,0 +1,26 @@
+#!/bin/bash
+#SBATCH -J diffdock_infer
+#SBATCH -p largedev
+#SBATCH -N 1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=dcu:1
+#SBATCH --cpus-per-task=8
+#SBATCH --time=02:00:00
+#SBATCH -o slurm-%j.out
+#SBATCH -e slurm-%j.err
+
+set -euo pipefail
+
+echo "START TIME: $(date)"
+echo "SLURM_JOB_ID: ${SLURM_JOB_ID:-none}"
+echo "HOSTNAME: $(hostname)"
+
+module purge || true
+module load sghpc-mpi-gcc/25.8 || true
+
+SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+mkdir -p "${SCRIPT_DIR}/../outputs"
+
+srun bash "${SCRIPT_DIR}/infer.sh"
+
+echo "END TIME: $(date)"

scripts/sample_diffdock.py

@@ -0,0 +1,395 @@
+import sys
+from pathlib import Path
+
+DIR = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(DIR))
+import argparse
+import copy
+import csv
+import os
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import torch
+import yaml
+from rdkit.Chem import RemoveAllHs
+from torch_geometric.loader import DataLoader
+
+from onescience.datapipes.diffdock.process_mols import write_mol_with_coords
+from onescience.utils.diffdock.diffusion_utils import get_t_schedule
+from onescience.utils.diffdock.inference_utils import InferenceDataset, set_nones
+from onescience.utils.diffdock.logging_utils import configure_logger, get_logger
+from onescience.utils.diffdock.sampling import randomize_position, sampling
+from onescience.utils.diffdock.validation import validate_sampling_entrypoint
+
+try:
+    from models.score_wrapper import load_model_args, load_score_model, model_uses_lm_embeddings
+except ImportError:
+    from models.score_wrapper import load_model_args, load_score_model, model_uses_lm_embeddings
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", required=True, help="Path to the sampling YAML config.")
+    return parser.parse_args()
+
+def _resolve_env_vars(obj):
+    if isinstance(obj, str):
+        return os.path.expandvars(obj)
+    if isinstance(obj, dict):
+        return {k: _resolve_env_vars(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [_resolve_env_vars(v) for v in obj]
+    return obj
+
+# def load_config(config_path):
+#     with open(config_path, "r", encoding="utf-8") as handle:
+#         return yaml.safe_load(handle) or {}
+def load_config(config_path):
+    with open(config_path, "r", encoding="utf-8") as handle:
+        return _resolve_env_vars(yaml.safe_load(handle) or {})
+
+
+def flatten_config(config):
+    flat = {}
+    for key, value in config.items():
+        if isinstance(value, dict):
+            flat.update(value)
+        else:
+            flat[key] = value
+    return flat
+
+
+def to_namespace(config):
+    return SimpleNamespace(**config)
+
+
+def resolve_device(device_name):
+    if device_name in {None, "auto"}:
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return torch.device(device_name)
+
+
+def load_inputs(args):
+    if args.protein_ligand_csv is not None:
+        with open(args.protein_ligand_csv, "r", encoding="utf-8", newline="") as handle:
+            rows = list(csv.DictReader(handle))
+        complex_names = set_nones([row.get("complex_name") for row in rows])
+        protein_paths = set_nones([row.get("protein_path") for row in rows])
+        protein_sequences = set_nones([row.get("protein_sequence") for row in rows])
+        ligand_descriptions = set_nones([row.get("ligand_description") for row in rows])
+    else:
+        complex_names = [args.complex_name or "complex_0"]
+        protein_paths = [args.protein_path]
+        protein_sequences = [args.protein_sequence]
+        ligand_descriptions = [args.ligand_description]
+
+    complex_names = [name if name is not None else f"complex_{idx}" for idx, name in enumerate(complex_names)]
+    return complex_names, protein_paths, protein_sequences, ligand_descriptions
+
+
+def ensure_output_dirs(out_dir, complex_names):
+    os.makedirs(out_dir, exist_ok=True)
+    for name in complex_names:
+        os.makedirs(os.path.join(out_dir, name), exist_ok=True)
+
+
+def get_ligand_mol(complex_graph):
+    mol = complex_graph.mol
+    return mol[0] if isinstance(mol, (list, tuple)) else mol
+
+
+def resolve_lm_embeddings_flag(args, model_args):
+    lm_embeddings = getattr(args, "lm_embeddings", None)
+    if lm_embeddings is None:
+        return model_uses_lm_embeddings(model_args)
+    return lm_embeddings
+
+
+def build_inference_dataset(
+    args,
+    model_args,
+    *,
+    complex_names,
+    protein_paths,
+    protein_sequences,
+    ligand_descriptions,
+    lm_embeddings,
+):
+    return InferenceDataset(
+        out_dir=args.out_dir,
+        complex_names=complex_names,
+        protein_files=protein_paths,
+        ligand_descriptions=ligand_descriptions,
+        protein_sequences=protein_sequences,
+        lm_embeddings=lm_embeddings,
+        receptor_radius=model_args.receptor_radius,
+        remove_hs=model_args.remove_hs,
+        c_alpha_max_neighbors=model_args.c_alpha_max_neighbors,
+        all_atoms=model_args.all_atoms,
+        atom_radius=model_args.atom_radius,
+        atom_max_neighbors=model_args.atom_max_neighbors,
+        knn_only_graph=not getattr(model_args, "not_knn_only_graph", False),
+    )
+
+
+def inference_graph_signature(model_args, lm_embeddings):
+    return (
+        getattr(model_args, "receptor_radius", None),
+        getattr(model_args, "remove_hs", None),
+        getattr(model_args, "c_alpha_max_neighbors", None),
+        getattr(model_args, "all_atoms", None),
+        getattr(model_args, "atom_radius", None),
+        getattr(model_args, "atom_max_neighbors", None),
+        not getattr(model_args, "not_knn_only_graph", False),
+        bool(lm_embeddings),
+    )
+
+
+def extract_confidence_scores(confidence, confidence_model_args):
+    confidence_scores = confidence
+    if isinstance(getattr(confidence_model_args, "rmsd_classification_cutoff", None), list):
+        confidence_scores = confidence_scores[:, 0]
+    return np.asarray(confidence_scores.detach().cpu().numpy()).reshape(-1)
+
+
+def load_optional_confidence_model(args, device, logger):
+    confidence_model_dir = getattr(args, "confidence_model_dir", None)
+    if confidence_model_dir is None:
+        return None, None
+
+    confidence_model_dir = Path(confidence_model_dir)
+    confidence_checkpoint_path = confidence_model_dir / args.confidence_ckpt
+    if not confidence_checkpoint_path.exists():
+        raise FileNotFoundError(
+            f"Confidence checkpoint not found: {confidence_checkpoint_path}. "
+            "This example does not auto-download models."
+        )
+
+    confidence_model_args_preview = load_model_args(confidence_model_dir)
+    validate_sampling_entrypoint(
+        confidence_model_args_preview,
+        context="DiffDock sampling confidence-model checkpoint",
+        include_confidence=True,
+        confidence_mode=True,
+    )
+
+    confidence_model, confidence_model_args, _ = load_score_model(
+        model_dir=confidence_model_dir,
+        ckpt=args.confidence_ckpt,
+        device=device,
+        no_parallel=True,
+        confidence_mode=True,
+        old=getattr(args, "old_confidence_model", False),
+    )
+    if hasattr(args, "crop_beyond"):
+        confidence_model_args.crop_beyond = args.crop_beyond
+    logger.info("Loaded optional confidence model from %s", confidence_model_dir)
+    return confidence_model, confidence_model_args
+
+
+def main():
+    parsed = parse_args()
+    raw_config = load_config(parsed.config)
+    args = to_namespace(flatten_config(raw_config))
+    device = resolve_device(getattr(args, "device", "auto"))
+    validate_sampling_entrypoint(
+        args,
+        include_confidence=(
+            getattr(args, "confidence_model_dir", None) is not None
+            or getattr(args, "old_confidence_model", False)
+        ),
+    )
+
+    configure_logger(getattr(args, "loglevel", "INFO"))
+    logger = get_logger()
+
+    model_dir = Path(args.model_dir)
+    checkpoint_path = model_dir / args.ckpt
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(
+            f"Checkpoint not found: {checkpoint_path}. This example does not auto-download models."
+        )
+    score_model_args_preview = load_model_args(model_dir)
+    validate_sampling_entrypoint(
+        score_model_args_preview,
+        context="DiffDock sampling score-model checkpoint",
+    )
+
+    model, score_model_args, t_to_sigma = load_score_model(
+        model_dir=model_dir,
+        ckpt=args.ckpt,
+        device=device,
+        no_parallel=True,
+        old=getattr(args, "old_score_model", False),
+    )
+    if hasattr(args, "crop_beyond"):
+        score_model_args.crop_beyond = args.crop_beyond
+    logger.info("DiffDock sampling will run on %s", device)
+
+    confidence_model, confidence_model_args = load_optional_confidence_model(args, device, logger)
+
+    complex_names, protein_paths, protein_sequences, ligand_descriptions = load_inputs(args)
+    ensure_output_dirs(args.out_dir, complex_names)
+
+    score_lm_embeddings = resolve_lm_embeddings_flag(args, score_model_args)
+    test_dataset = build_inference_dataset(
+        args,
+        score_model_args,
+        complex_names=complex_names,
+        protein_paths=protein_paths,
+        protein_sequences=protein_sequences,
+        ligand_descriptions=ligand_descriptions,
+        lm_embeddings=score_lm_embeddings,
+    )
+    test_loader = DataLoader(dataset=test_dataset, batch_size=1, shuffle=False)
+
+    confidence_loader = None
+    if confidence_model is not None:
+        confidence_lm_embeddings = resolve_lm_embeddings_flag(args, confidence_model_args)
+        confidence_needs_independent_graph = (
+            inference_graph_signature(score_model_args, score_lm_embeddings)
+            != inference_graph_signature(confidence_model_args, confidence_lm_embeddings)
+        )
+        if confidence_needs_independent_graph:
+            logger.info(
+                "Confidence rerank requires independent inference graphs; building a separate confidence dataset."
+            )
+            confidence_dataset = build_inference_dataset(
+                args,
+                confidence_model_args,
+                complex_names=complex_names,
+                protein_paths=protein_paths,
+                protein_sequences=protein_sequences,
+                ligand_descriptions=ligand_descriptions,
+                lm_embeddings=confidence_lm_embeddings,
+            )
+            confidence_loader = DataLoader(dataset=confidence_dataset, batch_size=1, shuffle=False)
+        else:
+            logger.info("Confidence rerank will reuse the score-model inference graphs.")
+
+    tr_schedule = get_t_schedule(
+        sigma_schedule=args.sigma_schedule,
+        inference_steps=args.inference_steps,
+        inf_sched_alpha=args.inf_sched_alpha,
+        inf_sched_beta=args.inf_sched_beta,
+    )
+
+    failures = 0
+    skipped = 0
+    num_samples = args.samples_per_complex
+    test_ds_size = len(test_dataset)
+    logger.info("Size of test dataset: %s", test_ds_size)
+
+    if confidence_loader is None:
+        loader_iter = ((orig_complex_graph, None) for orig_complex_graph in test_loader)
+    else:
+        loader_iter = zip(test_loader, confidence_loader)
+
+    for idx, (orig_complex_graph, confidence_orig_complex_graph) in enumerate(loader_iter):
+        if not orig_complex_graph.success[0]:
+            skipped += 1
+            logger.warning(
+                "Skipping %s because preprocessing failed.",
+                test_dataset.complex_names[idx],
+            )
+            continue
+        if confidence_orig_complex_graph is not None and not confidence_orig_complex_graph.success[0]:
+            skipped += 1
+            logger.warning(
+                "Skipping %s because confidence preprocessing failed.",
+                test_dataset.complex_names[idx],
+            )
+            continue
+
+        try:
+            data_list = [copy.deepcopy(orig_complex_graph) for _ in range(num_samples)]
+            confidence_data_list = None
+            if confidence_orig_complex_graph is not None:
+                confidence_data_list = [copy.deepcopy(confidence_orig_complex_graph) for _ in range(num_samples)]
+            randomize_position(
+                data_list,
+                score_model_args.no_torsion,
+                args.no_random,
+                score_model_args.tr_sigma_max,
+                initial_noise_std_proportion=args.initial_noise_std_proportion,
+                choose_residue=args.choose_residue,
+            )
+
+            ligand = get_ligand_mol(orig_complex_graph)
+            data_list, confidence = sampling(
+                data_list=data_list,
+                model=model,
+                inference_steps=args.actual_steps if args.actual_steps is not None else args.inference_steps,
+                tr_schedule=tr_schedule,
+                rot_schedule=tr_schedule,
+                tor_schedule=tr_schedule,
+                device=device,
+                t_to_sigma=t_to_sigma,
+                model_args=score_model_args,
+                no_random=args.no_random,
+                ode=args.ode,
+                confidence_model=confidence_model,
+                confidence_data_list=confidence_data_list,
+                confidence_model_args=confidence_model_args,
+                batch_size=args.batch_size,
+                no_final_step_noise=args.no_final_step_noise,
+                temp_sampling=[
+                    args.temp_sampling_tr,
+                    args.temp_sampling_rot,
+                    args.temp_sampling_tor,
+                ],
+                temp_psi=[
+                    args.temp_psi_tr,
+                    args.temp_psi_rot,
+                    args.temp_psi_tor,
+                ],
+                temp_sigma_data=[
+                    args.temp_sigma_data_tr,
+                    args.temp_sigma_data_rot,
+                    args.temp_sigma_data_tor,
+                ],
+            )
+
+            ligand_positions = np.asarray(
+                [
+                    complex_graph["ligand"].pos.cpu().numpy() + orig_complex_graph.original_center.cpu().numpy()
+                    for complex_graph in data_list
+                ]
+            )
+
+            rerank_order = np.arange(len(ligand_positions))
+            confidence_scores = None
+            if confidence is not None:
+                confidence_scores = extract_confidence_scores(confidence, confidence_model_args)
+                confidence_scores = np.nan_to_num(confidence_scores, nan=-1e-6)
+                rerank_order = np.argsort(confidence_scores)[::-1]
+                logger.info(
+                    "Applied confidence rerank for %s. Ranked confidences: %s",
+                    complex_names[idx],
+                    np.array2string(confidence_scores[rerank_order], precision=4),
+                )
+
+            write_dir = os.path.join(args.out_dir, complex_names[idx])
+            for rank, sample_idx in enumerate(rerank_order, start=1):
+                pos = ligand_positions[sample_idx]
+                mol_pred = copy.deepcopy(ligand)
+                if score_model_args.remove_hs:
+                    mol_pred = RemoveAllHs(mol_pred)
+                filename = f"rank{rank}.sdf"
+                if confidence_scores is not None:
+                    filename = f"rank{rank}_conf{confidence_scores[sample_idx]:.4f}.sdf"
+                write_mol_with_coords(mol_pred, pos, os.path.join(write_dir, filename))
+
+        except Exception as exc:
+            logger.warning("Failed on %s with error: %s", complex_names[idx], exc)
+            failures += 1
+
+    logger.info("Failed for %s / %s complexes.", failures, test_ds_size)
+    logger.info("Skipped %s / %s complexes.", skipped, test_ds_size)
+    logger.info("Results saved in %s", args.out_dir)
+
+
+if __name__ == "__main__":
+    main()

scripts/train.sh

@@ -0,0 +1,231 @@
+#!/bin/bash
+set -euo pipefail
+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"
+export LD_LIBRARY_PATH=${ROCM_PATH}/opencl/lib:$LD_LIBRARY_PATH
+
+SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+EXAMPLE_DIR=$(cd "${SCRIPT_DIR}/.." && pwd)
+REPO_ROOT=$(cd "${SCRIPT_DIR}/../../../.." && pwd)
+
+#source "${REPO_ROOT}/env.sh"
+
+if [[ -n "${ROCM_PATH:-}" && -f "${ROCM_PATH}/cuda/env.sh" ]]; then
+  source "${ROCM_PATH}/cuda/env.sh"
+fi
+
+export PYTHONPATH="${REPO_ROOT}/src:${REPO_ROOT}:${PYTHONPATH:-}"
+export HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-0}"
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-${HIP_VISIBLE_DEVICES}}"
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}"
+
+DIFFDOCK_DATA_ROOT="${DIFFDOCK_DATA_ROOT:-${ONESCIENCE_DATASETS_DIR}/diffdock}"
+DIFFDOCK_DATASETS_DIR="${DIFFDOCK_DATASETS_DIR:-${DIFFDOCK_DATA_ROOT}/datasets}"
+export TORCH_HOME="${TORCH_HOME:-${DIFFDOCK_DATA_ROOT}/torch_home}"
+
+DATASET="${DATASET:-pdbbind}"
+RUN_NAME="${RUN_NAME:-diffdock_${DATASET}_scnet}"
+LOG_DIR="${LOG_DIR:-${EXAMPLE_DIR}/outputs/train}"
+CACHE_PATH="${CACHE_PATH:-${EXAMPLE_DIR}/outputs/cache}"
+CONFIG_PATH="${CONFIG_PATH:-${LOG_DIR}/${RUN_NAME}_train_config.yml}"
+mkdir -p "${LOG_DIR}" "${CACHE_PATH}"
+LOG_DIR=$(cd "${LOG_DIR}" && pwd)
+CACHE_PATH=$(cd "${CACHE_PATH}" && pwd)
+CONFIG_DIR=$(dirname "${CONFIG_PATH}")
+CONFIG_NAME=$(basename "${CONFIG_PATH}")
+mkdir -p "${CONFIG_DIR}"
+CONFIG_PATH=$(cd "${CONFIG_DIR}" && pwd)/"${CONFIG_NAME}"
+
+PDBBIND_DIR="${PDBBIND_DIR:-${DIFFDOCK_DATASETS_DIR}/PDBBind_processed}"
+MOAD_DIR="${MOAD_DIR:-${DIFFDOCK_DATASETS_DIR}/BindingMOAD_2020_processed}"
+SPLIT_TRAIN="${SPLIT_TRAIN:-${DIFFDOCK_DATASETS_DIR}/splits/timesplit_no_lig_overlap_train}"
+SPLIT_VAL="${SPLIT_VAL:-${DIFFDOCK_DATASETS_DIR}/splits/timesplit_no_lig_overlap_val}"
+
+DEVICE="${DEVICE:-auto}"
+if [[ "$DEVICE" == "auto" ]]; then
+    if python -c "import torch; print(torch.cuda.is_available())" | grep -q True; then
+        DEVICE="cuda"
+    else
+        DEVICE="cpu"
+    fi
+fi
+echo "****** DEVICE ******: $DEVICE"
+
+SEED="${SEED:-0}"
+BATCH_SIZE="${BATCH_SIZE:-4}"
+N_EPOCHS="${N_EPOCHS:-10}"
+LR="${LR:-0.001}"
+NUM_WORKERS="${NUM_WORKERS:-1}"
+NUM_DATALOADER_WORKERS="${NUM_DATALOADER_WORKERS:-0}"
+LIMIT_COMPLEXES="${LIMIT_COMPLEXES:-null}"
+SAVE_MODEL_FREQ="${SAVE_MODEL_FREQ:-null}"
+VAL_INFERENCE_FREQ="${VAL_INFERENCE_FREQ:-null}"
+TRAIN_INFERENCE_FREQ="${TRAIN_INFERENCE_FREQ:-null}"
+WANDB="${WANDB:-false}"
+NO_TORSION="${NO_TORSION:-false}"
+NO_BATCH_NORM="${NO_BATCH_NORM:-true}"
+
+PDBBIND_ESM_EMBEDDINGS_PATH="${PDBBIND_ESM_EMBEDDINGS_PATH:-}"
+MOAD_ESM_EMBEDDINGS_PATH="${MOAD_ESM_EMBEDDINGS_PATH:-}"
+MOAD_ESM_EMBEDDINGS_SEQUENCES_PATH="${MOAD_ESM_EMBEDDINGS_SEQUENCES_PATH:-}"
+ESM_EMBEDDINGS_MODEL="${ESM_EMBEDDINGS_MODEL:-}"
+
+yaml_value() {
+  if [[ -z "${1:-}" || "${1}" == "null" ]]; then
+    printf "null"
+  else
+    local value
+    value=$(printf "%s" "$1" | sed "s/'/''/g")
+    printf "'%s'" "$value"
+  fi
+}
+
+yaml_bool() {
+  if [[ "${1,,}" == "true" ]]; then
+    printf "true"
+  else
+    printf "false"
+  fi
+}
+
+cat > "${CONFIG_PATH}" <<EOF
+runtime:
+  run_name: $(yaml_value "${RUN_NAME}")
+  log_dir: $(yaml_value "${LOG_DIR}")
+  device: $(yaml_value "${DEVICE}")
+  seed: ${SEED}
+  cudnn_benchmark: false
+  wandb: $(yaml_bool "${WANDB}")
+  project: diffdock
+
+data:
+  dataset: $(yaml_value "${DATASET}")
+  cache_path: $(yaml_value "${CACHE_PATH}")
+  pdbbind_dir: $(yaml_value "${PDBBIND_DIR}")
+  moad_dir: $(yaml_value "${MOAD_DIR}")
+  split_train: $(yaml_value "${SPLIT_TRAIN}")
+  split_val: $(yaml_value "${SPLIT_VAL}")
+  protein_file: protein_processed
+  limit_complexes: ${LIMIT_COMPLEXES}
+  num_conformers: 1
+  num_workers: ${NUM_WORKERS}
+  num_dataloader_workers: ${NUM_DATALOADER_WORKERS}
+  batch_size: ${BATCH_SIZE}
+  pin_memory: false
+  dataloader_drop_last: false
+  remove_hs: true
+  receptor_radius: 30
+  c_alpha_max_neighbors: 10
+  atom_radius: 5
+  atom_max_neighbors: 8
+  chain_cutoff: null
+  max_lig_size: null
+  matching_popsize: 20
+  matching_maxiter: 20
+  matching_tries: 1
+  not_knn_only_graph: false
+  include_miscellaneous_atoms: false
+  all_atoms: false
+  triple_training: false
+  combined_training: false
+  double_val: false
+  train_multiplicity: 1
+  val_multiplicity: 1
+  max_receptor_size: null
+  remove_promiscuous_targets: null
+  min_ligand_size: 0
+  unroll_clusters: false
+  enforce_timesplit: false
+  crop_beyond: 20
+  moad_esm_embeddings_path: $(yaml_value "${MOAD_ESM_EMBEDDINGS_PATH}")
+  pdbbind_esm_embeddings_path: $(yaml_value "${PDBBIND_ESM_EMBEDDINGS_PATH}")
+  moad_esm_embeddings_sequences_path: $(yaml_value "${MOAD_ESM_EMBEDDINGS_SEQUENCES_PATH}")
+  esm_embeddings_model: $(yaml_value "${ESM_EMBEDDINGS_MODEL}")
+
+diffusion:
+  no_torsion: $(yaml_bool "${NO_TORSION}")
+  tr_sigma_min: 0.1
+  tr_sigma_max: 30.0
+  rot_sigma_min: 0.1
+  rot_sigma_max: 1.65
+  tor_sigma_min: 0.0314
+  tor_sigma_max: 3.14
+  sampling_alpha: 1.0
+  sampling_beta: 1.0
+  tr_weight: 0.33
+  rot_weight: 0.33
+  tor_weight: 0.33
+  backbone_loss_weight: 0.0
+  sidechain_loss_weight: 0.0
+
+model:
+  num_conv_layers: 2
+  max_radius: 5.0
+  scale_by_sigma: true
+  norm_by_sigma: false
+  ns: 16
+  nv: 4
+  distance_embed_dim: 32
+  cross_distance_embed_dim: 32
+  no_batch_norm: $(yaml_bool "${NO_BATCH_NORM}")
+  use_second_order_repr: false
+  cross_max_distance: 80
+  dynamic_max_cross: false
+  dropout: 0.0
+  smooth_edges: false
+  odd_parity: false
+  embedding_type: sinusoidal
+  sigma_embed_dim: 32
+  embedding_scale: 1000
+  no_aminoacid_identities: false
+  sh_lmax: 2
+  no_differentiate_convolutions: false
+  tp_weights_layers: 2
+  num_prot_emb_layers: 0
+  reduce_pseudoscalars: false
+  embed_also_ligand: true
+  depthwise_convolution: false
+  use_old_atom_encoder: false
+
+optimization:
+  n_epochs: ${N_EPOCHS}
+  lr: ${LR}
+  w_decay: 0.0
+  scheduler: null
+  scheduler_patience: 20
+  lr_start_factor: 0.001
+  warmup_dur: 4
+  use_ema: true
+  ema_rate: 0.999
+  restart_dir: null
+  restart_ckpt: last_model
+  restart_lr: null
+  pretrain_dir: null
+  pretrain_ckpt: null
+  save_model_freq: ${SAVE_MODEL_FREQ}
+
+validation:
+  test_sigma_intervals: false
+  inference_samples: 1
+  val_inference_freq: ${VAL_INFERENCE_FREQ}
+  train_inference_freq: ${TRAIN_INFERENCE_FREQ}
+  inference_steps: 20
+  num_inference_complexes: 20
+  inference_earlystop_metric: valinf_min_rmsds_lt2
+  inference_secondary_metric: null
+  inference_earlystop_goal: max
+EOF
+
+echo "DiffDock training config: ${CONFIG_PATH}"
+echo "Dataset mode: ${DATASET}"
+echo "PDBBind dir: ${PDBBIND_DIR}"
+echo "MOAD dir: ${MOAD_DIR}"
+echo "Train split: ${SPLIT_TRAIN}"
+echo "Val split: ${SPLIT_VAL}"
+echo "Cache path: ${CACHE_PATH}"
+echo "Run output: ${LOG_DIR}/${RUN_NAME}"
+echo "TORCH_HOME: ${TORCH_HOME}"
+echo "Batch normalization disabled: ${NO_BATCH_NORM}"
+
+cd "${REPO_ROOT}"
+python "${SCRIPT_DIR}/train_diffdock.py" --config "${CONFIG_PATH}"

scripts/train_demo.sh

@@ -0,0 +1,199 @@
+#!/bin/bash
+set -euo pipefail
+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"
+export LD_LIBRARY_PATH=${ROCM_PATH}/opencl/lib:$LD_LIBRARY_PATH
+
+SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+EXAMPLE_DIR=$(cd "${SCRIPT_DIR}/.." && pwd)
+REPO_ROOT=$(cd "${SCRIPT_DIR}/../../../.." && pwd)
+# echo "*** REPO_ROOT : ${REPO_ROOT}"
+# source "${REPO_ROOT}/env.sh"
+
+if [[ -n "${ROCM_PATH:-}" && -f "${ROCM_PATH}/cuda/env.sh" ]]; then
+  source "${ROCM_PATH}/cuda/env.sh"
+fi
+
+export PYTHONPATH="${REPO_ROOT}/src:${REPO_ROOT}:${PYTHONPATH:-}"
+export HIP_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES:-0}"
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-${HIP_VISIBLE_DEVICES}}"
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}"
+
+DIFFDOCK_DATA_ROOT="${DIFFDOCK_DATA_ROOT:-${ONESCIENCE_DATASETS_DIR}/diffdock}"
+DIFFDOCK_DATASETS_DIR="${DIFFDOCK_DATASETS_DIR:-${DIFFDOCK_DATA_ROOT}/datasets}"
+export TORCH_HOME="${TORCH_HOME:-${DIFFDOCK_DATA_ROOT}/torch_home}"
+
+RUN_NAME="${RUN_NAME:-diffdock_pdbbind_smoke100_val20_cpu}"
+LOG_DIR="${LOG_DIR:-${EXAMPLE_DIR}/outputs/train}"
+CACHE_PATH="${CACHE_PATH:-${EXAMPLE_DIR}/outputs/cache_pdbbind_smoke100_val20}"
+CONFIG_PATH="${CONFIG_PATH:-${LOG_DIR}/${RUN_NAME}_train_config.yml}"
+
+DEVICE="${DEVICE:-cuda:0}"
+BATCH_SIZE="${BATCH_SIZE:-2}"
+N_EPOCHS="${N_EPOCHS:-1}"
+NUM_WORKERS="${NUM_WORKERS:-1}"
+NUM_DATALOADER_WORKERS="${NUM_DATALOADER_WORKERS:-0}"
+LR="${LR:-0.001}"
+SEED="${SEED:-0}"
+
+PDBBIND_DIR="${PDBBIND_DIR:-${DIFFDOCK_DATASETS_DIR}/PDBBind_processed}"
+SPLIT_TRAIN="${SPLIT_TRAIN:-${DIFFDOCK_DATASETS_DIR}/splits/timesplit_no_lig_overlap_train_smoke100}"
+SPLIT_VAL="${SPLIT_VAL:-${DIFFDOCK_DATASETS_DIR}/splits/timesplit_no_lig_overlap_val_smoke20}"
+
+mkdir -p "${LOG_DIR}" "${CACHE_PATH}"
+LOG_DIR=$(cd "${LOG_DIR}" && pwd)
+CACHE_PATH=$(cd "${CACHE_PATH}" && pwd)
+CONFIG_DIR=$(dirname "${CONFIG_PATH}")
+CONFIG_NAME=$(basename "${CONFIG_PATH}")
+mkdir -p "${CONFIG_DIR}"
+CONFIG_PATH=$(cd "${CONFIG_DIR}" && pwd)/"${CONFIG_NAME}"
+
+yaml_value() {
+  if [[ -z "${1:-}" || "${1}" == "null" ]]; then
+    printf "null"
+  else
+    local value
+    value=$(printf "%s" "$1" | sed "s/'/''/g")
+    printf "'%s'" "$value"
+  fi
+}
+
+cat > "${CONFIG_PATH}" <<EOF
+runtime:
+  run_name: $(yaml_value "${RUN_NAME}")
+  log_dir: $(yaml_value "${LOG_DIR}")
+  device: $(yaml_value "${DEVICE}")
+  seed: ${SEED}
+  cudnn_benchmark: false
+  wandb: false
+  project: diffdock
+
+data:
+  dataset: 'pdbbind'
+  cache_path: $(yaml_value "${CACHE_PATH}")
+  pdbbind_dir: $(yaml_value "${PDBBIND_DIR}")
+  moad_dir: $(yaml_value "${PDBBIND_DIR}")
+  split_train: $(yaml_value "${SPLIT_TRAIN}")
+  split_val: $(yaml_value "${SPLIT_VAL}")
+  protein_file: protein_processed
+  limit_complexes: null
+  num_conformers: 1
+  num_workers: ${NUM_WORKERS}
+  num_dataloader_workers: ${NUM_DATALOADER_WORKERS}
+  batch_size: ${BATCH_SIZE}
+  pin_memory: false
+  dataloader_drop_last: false
+  remove_hs: true
+  receptor_radius: 30
+  c_alpha_max_neighbors: 10
+  atom_radius: 5
+  atom_max_neighbors: 8
+  chain_cutoff: null
+  max_lig_size: null
+  matching_popsize: 20
+  matching_maxiter: 20
+  matching_tries: 1
+  not_knn_only_graph: false
+  include_miscellaneous_atoms: false
+  all_atoms: false
+  triple_training: false
+  combined_training: false
+  double_val: false
+  train_multiplicity: 1
+  val_multiplicity: 1
+  max_receptor_size: null
+  remove_promiscuous_targets: null
+  min_ligand_size: 0
+  unroll_clusters: false
+  enforce_timesplit: false
+  crop_beyond: 20
+  moad_esm_embeddings_path: null
+  pdbbind_esm_embeddings_path: null
+  moad_esm_embeddings_sequences_path: null
+  esm_embeddings_model: null
+
+diffusion:
+  no_torsion: false
+  tr_sigma_min: 0.1
+  tr_sigma_max: 30.0
+  rot_sigma_min: 0.1
+  rot_sigma_max: 1.65
+  tor_sigma_min: 0.0314
+  tor_sigma_max: 3.14
+  sampling_alpha: 1.0
+  sampling_beta: 1.0
+  tr_weight: 0.33
+  rot_weight: 0.33
+  tor_weight: 0.33
+  backbone_loss_weight: 0.0
+  sidechain_loss_weight: 0.0
+
+model:
+  num_conv_layers: 2
+  max_radius: 5.0
+  scale_by_sigma: true
+  norm_by_sigma: false
+  ns: 16
+  nv: 4
+  distance_embed_dim: 32
+  cross_distance_embed_dim: 32
+  no_batch_norm: true
+  use_second_order_repr: false
+  cross_max_distance: 80
+  dynamic_max_cross: false
+  dropout: 0.0
+  smooth_edges: false
+  odd_parity: false
+  embedding_type: sinusoidal
+  sigma_embed_dim: 32
+  embedding_scale: 1000
+  no_aminoacid_identities: false
+  sh_lmax: 2
+  no_differentiate_convolutions: false
+  tp_weights_layers: 2
+  num_prot_emb_layers: 0
+  reduce_pseudoscalars: false
+  embed_also_ligand: true
+  depthwise_convolution: false
+  use_old_atom_encoder: false
+
+optimization:
+  n_epochs: ${N_EPOCHS}
+  lr: ${LR}
+  w_decay: 0.0
+  scheduler: null
+  scheduler_patience: 20
+  lr_start_factor: 0.001
+  warmup_dur: 4
+  use_ema: true
+  ema_rate: 0.999
+  restart_dir: null
+  restart_ckpt: last_model
+  restart_lr: null
+  pretrain_dir: null
+  pretrain_ckpt: null
+  save_model_freq: null
+
+validation:
+  test_sigma_intervals: false
+  inference_samples: 1
+  val_inference_freq: null
+  train_inference_freq: null
+  inference_steps: 20
+  num_inference_complexes: 20
+  inference_earlystop_metric: valinf_min_rmsds_lt2
+  inference_secondary_metric: null
+  inference_earlystop_goal: max
+EOF
+
+echo "DiffDock smoke training config: ${CONFIG_PATH}"
+echo "PDBBind dir: ${PDBBIND_DIR}"
+echo "Train split: ${SPLIT_TRAIN}"
+echo "Val split: ${SPLIT_VAL}"
+echo "Cache path: ${CACHE_PATH}"
+echo "Run output: ${LOG_DIR}/${RUN_NAME}"
+echo "TORCH_HOME: ${TORCH_HOME}"
+
+cd "${REPO_ROOT}"
+python "${SCRIPT_DIR}/train_diffdock.py" --config "${CONFIG_PATH}"
+

scripts/train_diffdock.py

@@ -0,0 +1,404 @@
+import sys
+from pathlib import Path
+
+DIR = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(DIR))
+
+import argparse
+import copy
+import math
+import os
+import random
+from functools import partial
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import torch
+import yaml
+
+from onescience.datapipes.diffdock.loader import construct_loader
+from onescience.utils.diffdock.diffusion_utils import t_to_sigma as t_to_sigma_compl
+from onescience.utils.diffdock.training import (
+    inference_epoch_fix,
+    loss_function,
+    test_epoch,
+    train_epoch,
+)
+from onescience.utils.diffdock.utils import (
+    ExponentialMovingAverage,
+    get_optimizer_and_scheduler,
+    save_yaml_file,
+)
+from onescience.utils.diffdock.validation import validate_training_entrypoint
+
+try:
+    from models.score_wrapper import build_score_model
+except ImportError:
+    from models.score_wrapper import build_score_model
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", required=True, help="Path to the training YAML config.")
+    return parser.parse_args()
+
+def _resolve_env_vars(obj):
+    if isinstance(obj, str):
+        return os.path.expandvars(obj)
+    if isinstance(obj, dict):
+        return {k: _resolve_env_vars(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [_resolve_env_vars(v) for v in obj]
+    return obj
+
+# def load_config(config_path):
+#     with open(config_path, "r", encoding="utf-8") as handle:
+#         return yaml.safe_load(handle) or {}
+def load_config(config_path):
+    with open(config_path, "r", encoding="utf-8") as handle:
+        return _resolve_env_vars(yaml.safe_load(handle) or {})
+
+
+def flatten_config(config):
+    flat = {}
+    for key, value in config.items():
+        if isinstance(value, dict):
+            flat.update(value)
+        else:
+            flat[key] = value
+    return flat
+
+
+def to_namespace(config):
+    return SimpleNamespace(**config)
+
+
+def resolve_device(device_name):
+    if device_name in {None, "auto"}:
+        return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    return torch.device(device_name)
+
+
+def state_dict_for_save(model):
+    return model.module.state_dict() if hasattr(model, "module") else model.state_dict()
+
+
+def model_target(model):
+    return model.module if hasattr(model, "module") else model
+
+
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def maybe_init_wandb(args):
+    if not args.wandb:
+        return None
+    try:
+        import wandb
+    except ImportError as exc:
+        raise ImportError("wandb is enabled in the config but is not installed.") from exc
+    wandb.init(project=args.project, name=args.run_name, config=vars(args))
+    return wandb
+
+
+def maybe_load_restart(args, model, optimizer, ema_weights):
+    if args.restart_dir is None:
+        return
+
+    checkpoint_path = Path(args.restart_dir) / f"{args.restart_ckpt}.pt"
+    try:
+        checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
+        if args.restart_lr is not None:
+            checkpoint["optimizer"]["param_groups"][0]["lr"] = args.restart_lr
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        model_target(model).load_state_dict(checkpoint["model"], strict=True)
+        ema_weights.load_state_dict(checkpoint["ema_weights"], device=args.device)
+        print("Restarting from epoch", checkpoint["epoch"])
+    except Exception as exc:
+        print("Exception", exc)
+        checkpoint = torch.load(Path(args.restart_dir) / "best_model.pt", map_location=torch.device("cpu"))
+        model_target(model).load_state_dict(checkpoint, strict=True)
+        print("Due to exception had to take the best epoch and no optimiser")
+
+
+def maybe_load_pretrain(args, model):
+    if args.pretrain_dir is None:
+        return
+    checkpoint = torch.load(
+        Path(args.pretrain_dir) / f"{args.pretrain_ckpt}.pt",
+        map_location=torch.device("cpu"),
+    )
+    if isinstance(checkpoint, dict) and "model" in checkpoint and "optimizer" in checkpoint:
+        checkpoint = checkpoint["model"]
+    model_target(model).load_state_dict(checkpoint, strict=True)
+    print("Using pretrained model", str(Path(args.pretrain_dir) / f"{args.pretrain_ckpt}.pt"))
+
+
+def train(args, model, optimizer, scheduler, ema_weights, train_loader, val_loader, t_to_sigma, run_dir, val_dataset2):
+    loss_fn = partial(
+        loss_function,
+        tr_weight=args.tr_weight,
+        rot_weight=args.rot_weight,
+        tor_weight=args.tor_weight,
+        no_torsion=args.no_torsion,
+        backbone_weight=args.backbone_loss_weight,
+        sidechain_weight=args.sidechain_loss_weight,
+    )
+
+    best_val_loss = math.inf
+    best_val_inference_value = math.inf if args.inference_earlystop_goal == "min" else 0
+    best_val_secondary_value = math.inf if args.inference_earlystop_goal == "min" else 0
+    best_epoch = 0
+    best_val_inference_epoch = 0
+
+    freeze_params = 0
+    scheduler_mode = args.inference_earlystop_goal if args.val_inference_freq is not None else "min"
+    if args.scheduler == "layer_linear_warmup":
+        freeze_params = args.warmup_dur * (args.num_conv_layers + 2) - 1
+        print("Freezing some parameters until epoch {}".format(freeze_params))
+
+    wandb_run = maybe_init_wandb(args)
+
+    print("Starting training...")
+    for epoch in range(args.n_epochs):
+        if epoch % 5 == 0:
+            print("Run name:", args.run_name)
+
+        if args.scheduler == "layer_linear_warmup" and (epoch + 1) % args.warmup_dur == 0:
+            step = (epoch + 1) // args.warmup_dur
+            if step < args.num_conv_layers + 2:
+                print("New unfreezing step")
+                optimizer, scheduler = get_optimizer_and_scheduler(
+                    args,
+                    model,
+                    step=step,
+                    scheduler_mode=scheduler_mode,
+                )
+            elif step == args.num_conv_layers + 2:
+                print("Unfreezing all parameters")
+                optimizer, scheduler = get_optimizer_and_scheduler(
+                    args,
+                    model,
+                    step=step,
+                    scheduler_mode=scheduler_mode,
+                )
+                ema_weights = ExponentialMovingAverage(model.parameters(), decay=args.ema_rate)
+        elif args.scheduler == "linear_warmup" and epoch == args.warmup_dur:
+            print("Moving to plateu scheduler")
+            optimizer, scheduler = get_optimizer_and_scheduler(
+                args,
+                model,
+                step=1,
+                scheduler_mode=scheduler_mode,
+                optimizer=optimizer,
+            )
+
+        logs = {}
+        train_losses = train_epoch(
+            model,
+            train_loader,
+            optimizer,
+            args.device,
+            t_to_sigma,
+            loss_fn,
+            ema_weights if epoch > freeze_params else None,
+        )
+        print(
+            "Epoch {}: Training loss {:.4f}  tr {:.4f}   rot {:.4f}   tor {:.4f}   sc {:.4f}  lr {:.4f}".format(
+                epoch,
+                train_losses["loss"],
+                train_losses["tr_loss"],
+                train_losses["rot_loss"],
+                train_losses["tor_loss"],
+                train_losses["sidechain_loss"],
+                optimizer.param_groups[0]["lr"],
+            )
+        )
+
+        if epoch > freeze_params:
+            ema_weights.store(model.parameters())
+            if args.use_ema:
+                ema_weights.copy_to(model.parameters())
+
+        val_losses = test_epoch(model, val_loader, args.device, t_to_sigma, loss_fn, args.test_sigma_intervals)
+        print(
+            "Epoch {}: Validation loss {:.4f}  tr {:.4f}   rot {:.4f}   tor {:.4f}   sc {:.4f}".format(
+                epoch,
+                val_losses["loss"],
+                val_losses["tr_loss"],
+                val_losses["rot_loss"],
+                val_losses["tor_loss"],
+                val_losses["sidechain_loss"],
+            )
+        )
+
+        if args.val_inference_freq is not None and (epoch + 1) % args.val_inference_freq == 0:
+            inf_dataset = [
+                val_loader.dataset.get(i)
+                for i in range(min(args.num_inference_complexes, val_loader.dataset.__len__()))
+            ]
+            inf_metrics = inference_epoch_fix(model, inf_dataset, args.device, t_to_sigma, args)
+            print(
+                "Epoch {}: Val inference rmsds_lt2 {:.3f} rmsds_lt5 {:.3f} min_rmsds_lt2 {:.3f} min_rmsds_lt5 {:.3f}".format(
+                    epoch,
+                    inf_metrics["rmsds_lt2"],
+                    inf_metrics["rmsds_lt5"],
+                    inf_metrics["min_rmsds_lt2"],
+                    inf_metrics["min_rmsds_lt5"],
+                )
+            )
+            logs.update({"valinf_" + k: v for k, v in inf_metrics.items()})
+
+        if args.double_val and args.val_inference_freq is not None and (epoch + 1) % args.val_inference_freq == 0:
+            inf_dataset = [
+                val_dataset2.get(i)
+                for i in range(min(args.num_inference_complexes, val_dataset2.__len__()))
+            ]
+            inf_metrics2 = inference_epoch_fix(model, inf_dataset, args.device, t_to_sigma, args)
+            print(
+                "Epoch {}: Val inference on second validation rmsds_lt2 {:.3f} rmsds_lt5 {:.3f} min_rmsds_lt2 {:.3f} min_rmsds_lt5 {:.3f}".format(
+                    epoch,
+                    inf_metrics2["rmsds_lt2"],
+                    inf_metrics2["rmsds_lt5"],
+                    inf_metrics2["min_rmsds_lt2"],
+                    inf_metrics2["min_rmsds_lt5"],
+                )
+            )
+            logs.update({"valinf2_" + k: v for k, v in inf_metrics2.items()})
+            logs.update({"valinfcomb_" + k: (v + inf_metrics[k]) / 2 for k, v in inf_metrics2.items()})
+
+        if args.train_inference_freq is not None and (epoch + 1) % args.train_inference_freq == 0:
+            inf_dataset = [
+                train_loader.dataset.get(i)
+                for i in range(min(min(args.num_inference_complexes, 300), train_loader.dataset.__len__()))
+            ]
+            inf_metrics = inference_epoch_fix(model, inf_dataset, args.device, t_to_sigma, args)
+            print(
+                "Epoch {}: Train inference rmsds_lt2 {:.3f} rmsds_lt5 {:.3f} min_rmsds_lt2 {:.3f} min_rmsds_lt5 {:.3f}".format(
+                    epoch,
+                    inf_metrics["rmsds_lt2"],
+                    inf_metrics["rmsds_lt5"],
+                    inf_metrics["min_rmsds_lt2"],
+                    inf_metrics["min_rmsds_lt5"],
+                )
+            )
+            logs.update({"traininf_" + k: v for k, v in inf_metrics.items()})
+
+        if epoch > freeze_params:
+            if not args.use_ema:
+                ema_weights.copy_to(model.parameters())
+            ema_state_dict = copy.deepcopy(state_dict_for_save(model))
+            ema_weights.restore(model.parameters())
+        else:
+            ema_state_dict = copy.deepcopy(state_dict_for_save(model))
+
+        if wandb_run is not None:
+            logs.update({"train_" + k: v for k, v in train_losses.items()})
+            logs.update({"val_" + k: v for k, v in val_losses.items()})
+            logs["current_lr"] = optimizer.param_groups[0]["lr"]
+            wandb_run.log(logs, step=epoch + 1)
+
+        model_state_dict = state_dict_for_save(model)
+        if args.inference_earlystop_metric in logs and (
+            (args.inference_earlystop_goal == "min" and logs[args.inference_earlystop_metric] <= best_val_inference_value)
+            or (args.inference_earlystop_goal == "max" and logs[args.inference_earlystop_metric] >= best_val_inference_value)
+        ):
+            best_val_inference_value = logs[args.inference_earlystop_metric]
+            best_val_inference_epoch = epoch
+            torch.save(model_state_dict, os.path.join(run_dir, "best_inference_epoch_model.pt"))
+            if epoch > freeze_params:
+                torch.save(ema_state_dict, os.path.join(run_dir, "best_ema_inference_epoch_model.pt"))
+
+        if args.inference_secondary_metric is not None and args.inference_secondary_metric in logs and (
+            (args.inference_earlystop_goal == "min" and logs[args.inference_secondary_metric] <= best_val_secondary_value)
+            or (args.inference_earlystop_goal == "max" and logs[args.inference_secondary_metric] >= best_val_secondary_value)
+        ):
+            best_val_secondary_value = logs[args.inference_secondary_metric]
+            if epoch > freeze_params:
+                torch.save(ema_state_dict, os.path.join(run_dir, "best_ema_secondary_epoch_model.pt"))
+
+        if val_losses["loss"] <= best_val_loss:
+            best_val_loss = val_losses["loss"]
+            best_epoch = epoch
+            torch.save(model_state_dict, os.path.join(run_dir, "best_model.pt"))
+            if epoch > freeze_params:
+                torch.save(ema_state_dict, os.path.join(run_dir, "best_ema_model.pt"))
+
+        if args.save_model_freq is not None and (epoch + 1) % args.save_model_freq == 0:
+            best_model_path = os.path.join(run_dir, "best_model.pt")
+            if os.path.exists(best_model_path):
+                torch.save(torch.load(best_model_path, map_location=torch.device("cpu")), os.path.join(run_dir, f"epoch{epoch + 1}_best_model.pt"))
+
+        if scheduler:
+            if epoch < freeze_params or (args.scheduler == "linear_warmup" and epoch < args.warmup_dur):
+                scheduler.step()
+            elif args.val_inference_freq is not None:
+                scheduler.step(best_val_inference_value)
+            else:
+                scheduler.step(val_losses["loss"])
+
+        torch.save(
+            {
+                "epoch": epoch,
+                "model": model_state_dict,
+                "optimizer": optimizer.state_dict(),
+                "ema_weights": ema_weights.state_dict(),
+            },
+            os.path.join(run_dir, "last_model.pt"),
+        )
+
+    print("Best Validation Loss {} on Epoch {}".format(best_val_loss, best_epoch))
+    print("Best inference metric {} on Epoch {}".format(best_val_inference_value, best_val_inference_epoch))
+
+
+def main():
+    parsed = parse_args()
+    raw_config = load_config(parsed.config)
+    flat_config = flatten_config(raw_config)
+    args = to_namespace(flat_config)
+
+    if getattr(args, "run_name", None) in {None, ""}:
+        args.run_name = Path(parsed.config).stem
+
+    args.device = resolve_device(getattr(args, "device", "auto"))
+    if getattr(args, "cudnn_benchmark", False) and args.device.type == "cuda":
+        torch.backends.cudnn.benchmark = True
+    set_seed(getattr(args, "seed", 0))
+    validate_training_entrypoint(args)
+
+    assert args.inference_earlystop_goal in {"max", "min"}
+    if args.val_inference_freq is not None and args.scheduler is not None:
+        assert args.scheduler_patience > args.val_inference_freq
+
+    run_dir = os.path.join(args.log_dir, args.run_name)
+    saved_args = vars(args).copy()
+    saved_args["device"] = str(args.device)
+    save_yaml_file(os.path.join(run_dir, "model_parameters.yml"), saved_args)
+
+    t_to_sigma = partial(t_to_sigma_compl, args=args)
+    train_loader, val_loader, val_dataset2 = construct_loader(args, t_to_sigma, args.device)
+    model, _ = build_score_model(args, args.device, no_parallel=False)
+    optimizer, scheduler = get_optimizer_and_scheduler(
+        args,
+        model,
+        scheduler_mode=args.inference_earlystop_goal if args.val_inference_freq is not None else "min",
+    )
+    ema_weights = ExponentialMovingAverage(model.parameters(), decay=args.ema_rate)
+
+    maybe_load_restart(args, model, optimizer, ema_weights)
+    maybe_load_pretrain(args, model)
+
+    numel = sum(p.numel() for p in model.parameters())
+    print("Model with", numel, "parameters")
+
+    train(args, model, optimizer, scheduler, ema_weights, train_loader, val_loader, t_to_sigma, run_dir, val_dataset2)
+
+
+if __name__ == "__main__":
+    main()

scripts/train_slurm.sbatch

@@ -0,0 +1,26 @@
+#!/bin/bash
+#SBATCH -J diffdock_train
+#SBATCH -p largedev
+#SBATCH -N 1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=dcu:1
+#SBATCH --cpus-per-task=8
+#SBATCH --time=24:00:00
+#SBATCH -o slurm-train-%j.out
+#SBATCH -e slurm-train-%j.err
+
+set -euo pipefail
+
+echo "START TIME: $(date)"
+echo "SLURM_JOB_ID: ${SLURM_JOB_ID:-none}"
+echo "HOSTNAME: $(hostname)"
+
+module purge || true
+module load sghpc-mpi-gcc/25.8 || true
+
+SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+mkdir -p "${SCRIPT_DIR}/../outputs"
+
+srun bash "${SCRIPT_DIR}/train.sh"
+
+echo "END TIME: $(date)"