+
+# FlowDock
+
+
+
+
+
+ + + + + +
+
+
+
+## Description
+
+Official implementation of "FlowDock: Geometric Flow Matching for Generative Protein-Ligand Docking and Affinity Prediction".
+
+## Contents
+
+- [Installation](#installation)
+- [How to prepare data for FlowDock](#how-to-prepare-data-for-flowdock)
+- [How to train FlowDock](#how-to-train-flowdock)
+- [How to evaluate FlowDock](#how-to-evaluate-flowdock)
+- [How to create comparative plots of evaluation results](#how-to-create-comparative-plots-of-evaluation-results)
+- [How to predict new protein-ligand complex structures and their affinities using FlowDock](#how-to-predict-new-protein-ligand-complex-structures-using-flowdock)
+- [For developers](#for-developers)
+- [Acknowledgements](#acknowledgements)
+- [Citing this work](#citing-this-work)
+
+## Installation
+
+
+
+
++ + + + + +
+
+
+
+Install Mamba
+
+```bash
+wget "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh"
+bash Miniforge3-$(uname)-$(uname -m).sh # accept all terms and install to the default location
+rm Miniforge3-$(uname)-$(uname -m).sh # (optionally) remove installer after using it
+source ~/.bashrc # alternatively, one can restart their shell session to achieve the same result
+```
+
+Install dependencies
+
+```bash
+# clone project
+git clone https://github.com/BioinfoMachineLearning/FlowDock
+cd FlowDock
+
+# create conda environment
+mamba env create -f environments/flowdock_environment.yaml
+conda activate FlowDock # NOTE: one still needs to use `conda` to (de)activate environments
+pip3 install -e . # install local project as package
+```
+
+Download checkpoints
+
+```bash
+# pretrained NeuralPLexer weights
+cd checkpoints/
+wget https://zenodo.org/records/10373581/files/neuralplexermodels_downstream_datasets_predictions.zip
+unzip neuralplexermodels_downstream_datasets_predictions.zip
+rm neuralplexermodels_downstream_datasets_predictions.zip
+cd ../
+```
+
+```bash
+# pretrained FlowDock weights
+cd checkpoints/
+wget https://zenodo.org/records/10373582/files/flowdock_weights.zip
+unzip flowdock_weights.zip
+rm flowdock_weights.zip
+cd ../
+```
+
+
+
+## How to prepare data for `FlowDock`
+
+
+
+**NOTE:** The following steps are only necessary if one wants to manually process each of the following datasets.
+Otherwise, preprocessed versions of each dataset can be found on [Zenodo](https://zenodo.org/records/10373582).
+
+Download data
+
+```bash
+# fetch preprocessed PDBBind, Binding MOAD & DockGen, as well as van der Mers (vdM) datasets
+cd data/
+
+wget https://zenodo.org/record/6408497/files/PDBBind.zip
+wget https://zenodo.org/records/10656052/files/BindingMOAD_2020_processed.tar
+wget https://zenodo.org/records/10656052/files/DockGen.tar
+wget https://files.ipd.uw.edu/pub/training_sets/pdb_2021aug02.tar.gz
+
+unzip PDBBind.zip
+tar -xf BindingMOAD_2020_processed.tar
+tar -xf DockGen.tar
+tar -xzf pdb_2021aug02.tar.gz
+
+rm PDBBind.zip BindingMOAD_2020_processed.tar DockGen.tar pdb_2021aug02.tar.gz
+
+mkdir pdbbind/ moad/ pdbsidechain/
+mv PDBBind_processed/ pdbbind/
+mv BindingMOAD_2020_processed/ moad/
+mv pdb_2021aug02/ pdbsidechain/
+
+cd ../
+```
+
+### Generating ESM2 embeddings for each protein
+
+To generate the ESM2 embeddings for the protein inputs,
+first create all the corresponding FASTA files for each protein sequence
+
+```bash
+python flowdock/data/components/esm_embedding_preparation.py --dataset pdbbind --data_dir data/pdbbind/PDBBind_processed/ --out_file data/pdbbind/pdbbind_sequences.fasta
+python flowdock/data/components/esm_embedding_preparation.py --dataset moad --data_dir data/moad/BindingMOAD_2020_processed/pdb_protein/ --out_file data/moad/moad_sequences.fasta
+python flowdock/data/components/esm_embedding_preparation.py --dataset dockgen --data_dir data/DockGen/processed_files/ --out_file data/DockGen/dockgen_sequences.fasta
+python flowdock/data/components/esm_embedding_preparation.py --dataset pdbsidechain --data_dir data/pdbsidechain/pdb_2021aug02/pdb/ --out_file data/pdbsidechain/pdbsidechain_sequences.fasta
+```
+
+Then, generate all ESM2 embeddings in batch using the ESM repository's helper script
+
+```bash
+python flowdock/data/components/esm_embedding_extraction.py esm2_t33_650M_UR50D data/pdbbind/pdbbind_sequences.fasta data/pdbbind/embeddings_output --repr_layers 33 --include per_tok --truncation_seq_length 4096 --cuda_device_index 0
+python flowdock/data/components/esm_embedding_extraction.py esm2_t33_650M_UR50D data/moad/moad_sequences.fasta data/moad/embeddings_output --repr_layers 33 --include per_tok --truncation_seq_length 4096 --cuda_device_index 0
+python flowdock/data/components/esm_embedding_extraction.py esm2_t33_650M_UR50D data/DockGen/dockgen_sequences.fasta data/DockGen/embeddings_output --repr_layers 33 --include per_tok --truncation_seq_length 4096 --cuda_device_index 0
+python flowdock/data/components/esm_embedding_extraction.py esm2_t33_650M_UR50D data/pdbsidechain/pdbsidechain_sequences.fasta data/pdbsidechain/embeddings_output --repr_layers 33 --include per_tok --truncation_seq_length 4096 --cuda_device_index 0
+```
+
+### Predicting apo protein structures using ESMFold
+
+To generate the apo version of each protein structure,
+first create ESMFold-ready versions of the combined FASTA files
+prepared above by the script `esm_embedding_preparation.py`
+for the PDBBind, Binding MOAD, DockGen, and PDBSidechain datasets, respectively
+
+```bash
+python flowdock/data/components/esmfold_sequence_preparation.py dataset=pdbbind
+python flowdock/data/components/esmfold_sequence_preparation.py dataset=moad
+python flowdock/data/components/esmfold_sequence_preparation.py dataset=dockgen
+python flowdock/data/components/esmfold_sequence_preparation.py dataset=pdbsidechain
+```
+
+Then, predict each apo protein structure using ESMFold's batch
+inference script
+
+```bash
+# Note: Having a CUDA-enabled device available when running this script is highly recommended
+python flowdock/data/components/esmfold_batch_structure_prediction.py -i data/pdbbind/pdbbind_esmfold_sequences.fasta -o data/pdbbind/pdbbind_esmfold_structures --cuda-device-index 0 --skip-existing
+python flowdock/data/components/esmfold_batch_structure_prediction.py -i data/moad/moad_esmfold_sequences.fasta -o data/moad/moad_esmfold_structures --cuda-device-index 0 --skip-existing
+python flowdock/data/components/esmfold_batch_structure_prediction.py -i data/DockGen/dockgen_esmfold_sequences.fasta -o data/DockGen/dockgen_esmfold_structures --cuda-device-index 0 --skip-existing
+python flowdock/data/components/esmfold_batch_structure_prediction.py -i data/pdbsidechain/pdbsidechain_esmfold_sequences.fasta -o data/pdbsidechain/pdbsidechain_esmfold_structures --cuda-device-index 0 --skip-existing
+```
+
+Align each apo protein structure to its corresponding
+holo protein structure counterpart in PDBBind, Binding MOAD, and PDBSidechain,
+taking ligand conformations into account during each alignment
+
+```bash
+python flowdock/data/components/esmfold_apo_to_holo_alignment.py dataset=pdbbind num_workers=1
+python flowdock/data/components/esmfold_apo_to_holo_alignment.py dataset=moad num_workers=1
+python flowdock/data/components/esmfold_apo_to_holo_alignment.py dataset=dockgen num_workers=1
+python flowdock/data/components/esmfold_apo_to_holo_alignment.py dataset=pdbsidechain num_workers=1
+```
+
+Lastly, assess the apo-to-holo alignments in terms of statistics and structural metrics
+to enable runtime-dynamic dataset filtering using such information
+
+```bash
+python flowdock/data/components/esmfold_apo_to_holo_assessment.py dataset=pdbbind usalign_exec_path=$MY_USALIGN_EXEC_PATH
+python flowdock/data/components/esmfold_apo_to_holo_assessment.py dataset=moad usalign_exec_path=$MY_USALIGN_EXEC_PATH
+python flowdock/data/components/esmfold_apo_to_holo_assessment.py dataset=dockgen usalign_exec_path=$MY_USALIGN_EXEC_PATH
+python flowdock/data/components/esmfold_apo_to_holo_assessment.py dataset=pdbsidechain usalign_exec_path=$MY_USALIGN_EXEC_PATH
+```
+
+
+
+## How to train `FlowDock`
+
+
+
+Train model with default configuration
+
+```bash
+# train on CPU
+python flowdock/train.py trainer=cpu
+
+# train on GPU
+python flowdock/train.py trainer=gpu
+```
+
+Train model with chosen experiment configuration from [configs/experiment/](configs/experiment/)
+
+```bash
+python flowdock/train.py experiment=experiment_name.yaml
+```
+
+For example, reproduce `FlowDock`'s default model training run
+
+```bash
+python flowdock/train.py experiment=flowdock_fm
+```
+
+**Note:** You can override any parameter from command line like this
+
+```bash
+python flowdock/train.py experiment=flowdock_fm trainer.max_epochs=20 data.batch_size=8
+```
+
+For example, override parameters to finetune `FlowDock`'s pretrained weights using a new dataset
+
+```bash
+python flowdock/train.py experiment=flowdock_fm data=my_new_datamodule ckpt_path=checkpoints/esmfold_prior_paper_weights.ckpt
+```
+
+
+
+## How to evaluate `FlowDock`
+
+
+
+To reproduce `FlowDock`'s evaluation results for structure prediction, please refer to its documentation in the [PoseBench](https://github.com/BioinfoMachineLearning/PoseBench?tab=readme-ov-file#how-to-run-inference-with-flowdock) GitHub repository.
+
+To reproduce `FlowDock`'s evaluation results for binding affinity prediction using the PDBBind dataset
+
+```bash
+python flowdock/eval.py data.test_datasets=[pdbbind] ckpt_path=checkpoints/esmfold_prior_paper_weights_EMA.ckpt trainer=gpu
+... # re-run two more times to gather triplicate results
+```
+
+
+
+## How to create comparative plots of structure prediction results
+
+
+
+```bash
+jupyter notebook notebooks/structure_prediction_results_plotting.ipynb
+```
+
+
+
+## How to predict new protein-ligand complex structures and their affinities using `FlowDock`
+
+
+
+For example, generate new protein-ligand complexes for a pair of protein sequence and ligand SMILES strings such as those of the PDBBind 2020 test target `6i67`
+
+```bash
+python flowdock/sample.py ckpt_path=checkpoints/esmfold_prior_paper_weights_EMA.ckpt model.cfg.prior_type=esmfold sampling_task=batched_structure_sampling input_receptor='YNKIVHLLVAEPEKIYAMPDPTVPDSDIKALTTLCDLADRELVVIIGWAKHIPGFSTLSLADQMSLLQSAWMEILILGVVYRSLFEDELVYADDYIMDEDQSKLAGLLDLNNAILQLVKKYKSMKLEKEEFVTLKAIALANSDSMHIEDVEAVQKLQDVLHEALQDYEAGQHMEDPRRAGKMLMTLPLLRQTSTKAVQHFYNKLEGKVPMHKLFLEMLEAKV' input_ligand='"c1cc2c(cc1O)CCCC2"' input_template=data/pdbbind/pdbbind_holo_aligned_esmfold_structures/6i67_holo_aligned_esmfold_protein.pdb sample_id='6i67' out_path='./6i67_sampled_structures/' n_samples=5 chunk_size=5 num_steps=40 sampler=VDODE sampler_eta=1.0 start_time='1.0' use_template=true separate_pdb=true visualize_sample_trajectories=true auxiliary_estimation_only=false esmfold_chunk_size=null trainer=gpu
+```
+
+Or, for example, generate new protein-ligand complexes for pairs of protein sequences and (multi-)ligand SMILES strings (delimited via `|`) such as those of the CASP15 target `T1152`
+
+```bash
+python flowdock/sample.py ckpt_path=checkpoints/esmfold_prior_paper_weights_EMA.ckpt model.cfg.prior_type=esmfold sampling_task=batched_structure_sampling input_receptor='MYTVKPGDTMWKIAVKYQIGISEIIAANPQIKNPNLIYPGQKINIP|MYTVKPGDTMWKIAVKYQIGISEIIAANPQIKNPNLIYPGQKINIP|MYTVKPGDTMWKIAVKYQIGISEIIAANPQIKNPNLIYPGQKINIPN' input_ligand='"CC(=O)NC1C(O)OC(CO)C(OC2OC(CO)C(OC3OC(CO)C(O)C(O)C3NC(C)=O)C(O)C2NC(C)=O)C1O"' input_template=data/test_cases/predicted_structures/T1152.pdb sample_id='T1152' out_path='./T1152_sampled_structures/' n_samples=5 chunk_size=5 num_steps=40 sampler=VDODE sampler_eta=1.0 start_time='1.0' use_template=true separate_pdb=true visualize_sample_trajectories=true auxiliary_estimation_only=false esmfold_chunk_size=null trainer=gpu
+```
+
+If you do not already have a template protein structure available for your target of interest, set `input_template=null` to instead have the sampling script predict the ESMFold structure of your provided `input_protein` sequence before running the sampling pipeline. For more information regarding the input arguments available for sampling, please refer to the config at `configs/sample.yaml`.
+
+
+
+## For developers
+
+
+
+Set up `pre-commit` (one time only) for automatic code linting and formatting upon each `git commit`
+
+```bash
+pre-commit install
+```
+
+Manually reformat all files in the project, as desired
+
+```bash
+pre-commit run -a
+```
+
+Update dependencies in a `*_environment.yml` file
+
+```bash
+mamba env export > env.yaml # e.g., run this after installing new dependencies locally
+diff environments/flowdock_environment.yaml env.yaml # note the differences and copy accepted changes back into e.g., `environments/flowdock_environment.yaml`
+rm env.yaml # clean up temporary environment file
+```
+
+
+
+## Acknowledgements
+
+`FlowDock` builds upon the source code and data from the following projects:
+
+- [DiffDock](https://github.com/gcorso/DiffDock)
+- [lightning-hydra-template](https://github.com/ashleve/lightning-hydra-template)
+- [NeuralPLexer](https://github.com/zrqiao/NeuralPLexer)
+
+We thank all their contributors and maintainers!
+
+## Citing this work
+
+If you use the code associated with this repository or otherwise find this work useful, please cite:
+
+```bibtex
+@article{morehead2024flowdock,
+ title={FlowDock: Geometric Flow Matching for Generative Protein-Ligand Docking and Affinity Prediction},
+ author={Morehead, Alex and Cheng, Jianlin},
+ year={2024}
+}
+```
diff --git a/forks/FlowDock/configs/__init__.py b/forks/FlowDock/configs/__init__.py
new file mode 100644
index 00000000..56bf7f4a
--- /dev/null
+++ b/forks/FlowDock/configs/__init__.py
@@ -0,0 +1 @@
+# this file is needed here to include configs when building project as a package
diff --git a/forks/FlowDock/configs/callbacks/default.yaml b/forks/FlowDock/configs/callbacks/default.yaml
new file mode 100644
index 00000000..47225285
--- /dev/null
+++ b/forks/FlowDock/configs/callbacks/default.yaml
@@ -0,0 +1,21 @@
+defaults:
+ - ema
+ - last_model_checkpoint
+ - learning_rate_monitor
+ - model_checkpoint
+ - model_summary
+ - rich_progress_bar
+ - _self_
+
+last_model_checkpoint:
+ dirpath: ${paths.output_dir}/checkpoints
+ filename: "last"
+ monitor: null
+ verbose: True
+ auto_insert_metric_name: False
+ every_n_epochs: 1
+ save_on_train_epoch_end: True
+ enable_version_counter: False
+
+model_summary:
+ max_depth: -1
diff --git a/forks/FlowDock/configs/callbacks/early_stopping.yaml b/forks/FlowDock/configs/callbacks/early_stopping.yaml
new file mode 100644
index 00000000..c826c8d5
--- /dev/null
+++ b/forks/FlowDock/configs/callbacks/early_stopping.yaml
@@ -0,0 +1,15 @@
+# https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.EarlyStopping.html
+
+early_stopping:
+ _target_: lightning.pytorch.callbacks.EarlyStopping
+ monitor: ??? # quantity to be monitored, must be specified !!!
+ min_delta: 0. # minimum change in the monitored quantity to qualify as an improvement
+ patience: 3 # number of checks with no improvement after which training will be stopped
+ verbose: False # verbosity mode
+ mode: "min" # "max" means higher metric value is better, can be also "min"
+ strict: True # whether to crash the training if monitor is not found in the validation metrics
+ check_finite: True # when set True, stops training when the monitor becomes NaN or infinite
+ stopping_threshold: null # stop training immediately once the monitored quantity reaches this threshold
+ divergence_threshold: null # stop training as soon as the monitored quantity becomes worse than this threshold
+ check_on_train_epoch_end: null # whether to run early stopping at the end of the training epoch
+ # log_rank_zero_only: False # this keyword argument isn't available in stable version
diff --git a/forks/FlowDock/configs/callbacks/ema.yaml b/forks/FlowDock/configs/callbacks/ema.yaml
new file mode 100644
index 00000000..5bff7bba
--- /dev/null
+++ b/forks/FlowDock/configs/callbacks/ema.yaml
@@ -0,0 +1,10 @@
+# https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/callbacks/ema.py
+
+# Maintains an exponential moving average (EMA) of model weights.
+# Look at the above link for more detailed information regarding the original implementation.
+ema:
+ _target_: flowdock.models.components.callbacks.ema.EMA
+ decay: 0.999
+ validate_original_weights: false
+ every_n_steps: 4
+ cpu_offload: false
diff --git a/forks/FlowDock/configs/callbacks/last_model_checkpoint.yaml b/forks/FlowDock/configs/callbacks/last_model_checkpoint.yaml
new file mode 100644
index 00000000..0f780d02
--- /dev/null
+++ b/forks/FlowDock/configs/callbacks/last_model_checkpoint.yaml
@@ -0,0 +1,21 @@
+# https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html
+
+last_model_checkpoint:
+ # NOTE: this is a direct copy of `model_checkpoint`,
+ # which is necessary to make to work around the
+ # key-duplication limitations of YAML config files
+ _target_: flowdock.models.components.callbacks.ema.EMAModelCheckpoint
+ dirpath: null # directory to save the model file
+ filename: null # checkpoint filename
+ monitor: null # name of the logged metric which determines when model is improving
+ verbose: False # verbosity mode
+ save_last: null # additionally always save an exact copy of the last checkpoint to a file last.ckpt
+ save_top_k: 1 # save k best models (determined by above metric)
+ mode: "min" # "max" means higher metric value is better, can be also "min"
+ auto_insert_metric_name: True # when True, the checkpoints filenames will contain the metric name
+ save_weights_only: False # if True, then only the model’s weights will be saved
+ every_n_train_steps: null # number of training steps between checkpoints
+ train_time_interval: null # checkpoints are monitored at the specified time interval
+ every_n_epochs: null # number of epochs between checkpoints
+ save_on_train_epoch_end: null # whether to run checkpointing at the end of the training epoch or the end of validation
+ enable_version_counter: True # enables versioning for checkpoint names
diff --git a/forks/FlowDock/configs/callbacks/learning_rate_monitor.yaml b/forks/FlowDock/configs/callbacks/learning_rate_monitor.yaml
new file mode 100644
index 00000000..fa8ecd6b
--- /dev/null
+++ b/forks/FlowDock/configs/callbacks/learning_rate_monitor.yaml
@@ -0,0 +1,7 @@
+# https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.LearningRateMonitor.html
+
+learning_rate_monitor:
+ _target_: lightning.pytorch.callbacks.LearningRateMonitor
+ logging_interval: null # set to `epoch` or `step` to log learning rate of all optimizers at the same interval, or set to `null` to log at individual interval according to the interval key of each scheduler
+ log_momentum: false # whether to also log the momentum values of the optimizer, if the optimizer has the `momentum` or `betas` attribute
+ log_weight_decay: false # whether to also log the weight decay values of the optimizer, if the optimizer has the `weight_decay` attribute
diff --git a/forks/FlowDock/configs/callbacks/model_checkpoint.yaml b/forks/FlowDock/configs/callbacks/model_checkpoint.yaml
new file mode 100644
index 00000000..8443e2d4
--- /dev/null
+++ b/forks/FlowDock/configs/callbacks/model_checkpoint.yaml
@@ -0,0 +1,18 @@
+# https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html
+
+model_checkpoint:
+ _target_: flowdock.models.components.callbacks.ema.EMAModelCheckpoint
+ dirpath: null # directory to save the model file
+ filename: "best" # checkpoint filename
+ monitor: val_sampling/ligand_hit_score_2A_epoch # name of the logged metric which determines when model is improving
+ verbose: True # verbosity mode
+ save_last: False # additionally always save an exact copy of the last checkpoint to a file last.ckpt
+ save_top_k: 1 # save k best models (determined by above metric)
+ mode: "max" # "max" means higher metric value is better, can be also "min"
+ auto_insert_metric_name: True # when True, the checkpoints filenames will contain the metric name
+ save_weights_only: False # if True, then only the model’s weights will be saved
+ every_n_train_steps: null # number of training steps between checkpoints
+ train_time_interval: null # checkpoints are monitored at the specified time interval
+ every_n_epochs: null # number of epochs between checkpoints
+ save_on_train_epoch_end: null # whether to run checkpointing at the end of the training epoch or the end of validation
+ enable_version_counter: False # enables versioning for checkpoint names
diff --git a/forks/FlowDock/configs/callbacks/model_summary.yaml b/forks/FlowDock/configs/callbacks/model_summary.yaml
new file mode 100644
index 00000000..b75981d8
--- /dev/null
+++ b/forks/FlowDock/configs/callbacks/model_summary.yaml
@@ -0,0 +1,5 @@
+# https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.RichModelSummary.html
+
+model_summary:
+ _target_: lightning.pytorch.callbacks.RichModelSummary
+ max_depth: 1 # the maximum depth of layer nesting that the summary will include
diff --git a/forks/FlowDock/configs/callbacks/none.yaml b/forks/FlowDock/configs/callbacks/none.yaml
new file mode 100644
index 00000000..e69de29b
diff --git a/forks/FlowDock/configs/callbacks/rich_progress_bar.yaml b/forks/FlowDock/configs/callbacks/rich_progress_bar.yaml
new file mode 100644
index 00000000..de6f1ccb
--- /dev/null
+++ b/forks/FlowDock/configs/callbacks/rich_progress_bar.yaml
@@ -0,0 +1,4 @@
+# https://lightning.ai/docs/pytorch/latest/api/lightning.pytorch.callbacks.RichProgressBar.html
+
+rich_progress_bar:
+ _target_: lightning.pytorch.callbacks.RichProgressBar
diff --git a/forks/FlowDock/configs/data/combined.yaml b/forks/FlowDock/configs/data/combined.yaml
new file mode 100644
index 00000000..9218348e
--- /dev/null
+++ b/forks/FlowDock/configs/data/combined.yaml
@@ -0,0 +1,100 @@
+_target_: flowdock.data.combined_datamodule.CombinedDataModule
+data_dir: ${paths.data_dir}
+train_datasets: ["pdbbind", "moad"] # NOTE: must be any non-empty combination of `("pdbbind", "moad", "pdbsidechain")`
+test_datasets: ["pdbbind", "moad"] # NOTE: must be any non-empty combination of `("pdbbind", "moad")`
+batch_size: 16 # Needs to be divisible by the number of devices (e.g., if in a distributed setup)
+num_workers: 4
+pin_memory: True
+# overfitting arguments
+overfitting_example_name: null # NOTE: if specified, only this example will be used during model overfitting
+# model arguments
+n_protein_patches: 96
+n_lig_patches: 32
+epoch_frac: 1.0
+edge_crop_size: 400000
+# split arguments
+split_train: ${paths.data_dir}/splits/timesplit_no_lig_overlap_train
+split_val: ${paths.data_dir}/splits/timesplit_no_lig_overlap_val
+split_test: ${paths.data_dir}/splits/timesplit_test
+split_time: ${paths.data_dir}/splits/pdbids_2019
+double_val: false
+# general dataset arguments
+cache_path: ${paths.data_dir}/cache
+limit_complexes: 0
+all_atoms: false
+max_lig_size: null
+min_multi_lig_distance: 0.5 # NOTE: measured in Angstrom
+remove_hs: true
+esm_embeddings_model: null
+include_miscellaneous_atoms: false
+train_multiplicity: 1
+val_multiplicity: 1
+max_receptor_size: null
+remove_promiscuous_targets: null
+min_ligand_size: 0
+unroll_clusters: false # NOTE: this value affects MOAD training and validation splits, but the MOAD test split is always unrolled to support DockGen
+enforce_timesplit: true
+merge_clusters: 1
+crop_beyond: 20
+protein_file: protein_processed
+ligand_file: ligand
+min_protein_length: 50
+max_protein_length: 750
+# PDBBind dataset arguments
+pdbbind_dir: ${paths.data_dir}/pdbbind/PDBBind_processed
+pdbbind_esm_embeddings_path: ${paths.data_dir}/pdbbind/embeddings_output
+pdbbind_apo_protein_structure_dir: ${paths.data_dir}/pdbbind/pdbbind_holo_aligned_esmfold_structures
+pdbbind_split_train: ${paths.data_dir}/splits/timesplit_no_lig_overlap_train
+pdbbind_split_val: ${paths.data_dir}/splits/timesplit_no_lig_overlap_val
+pdbbind_a2h_assessment_csv_filepath: ${paths.data_dir}/pdbbind/pdbbind_apo_to_holo_assessment.csv
+pdbbind_filter_using_a2h_assessment: true # if `pdbbind_a2h_assessment_csv_filepath` is provided, filter the dataset using the assessment
+pdbbind_a2h_min_tmscore: 0.7
+pdbbind_a2h_max_rmsd: 5.0
+pdbbind_a2h_min_protein_length: 50
+pdbbind_a2h_max_protein_length: 750
+pdbbind_a2h_min_ligand_length: 1
+pdbbind_a2h_max_ligand_length: 300
+pdbbind_binding_affinity_values_path: ${paths.data_dir}/pdbbind/pdbbind_binding_affinity_data/binding_affinity_values.txt # path to the text file containing PDBBind's binding affinity values
+# Binding MOAD dataset arguments
+moad_dir: ${paths.data_dir}/moad/BindingMOAD_2020_processed
+moad_dockgen_dir: ${paths.data_dir}/DockGen/processed_files
+moad_esm_embeddings_path: ${paths.data_dir}/moad/embeddings_output
+moad_dockgen_esm_embeddings_path: ${paths.data_dir}/DockGen/embeddings_output
+moad_esm_embeddings_sequences_path: ${paths.data_dir}/moad/moad_sequences.fasta
+moad_dockgen_esm_embeddings_sequences_path: ${paths.data_dir}/DockGen/dockgen_sequences.fasta
+moad_apo_protein_structure_dir: ${paths.data_dir}/moad/moad_holo_aligned_esmfold_structures
+moad_dockgen_apo_protein_structure_dir: ${paths.data_dir}/DockGen/dockgen_holo_aligned_esmfold_structures
+moad_clusters_filepath: ${paths.data_dir}/splits/MOAD_generalisation_splits.pkl
+moad_a2h_assessment_csv_filepath: ${paths.data_dir}/moad/moad_apo_to_holo_assessment.csv
+moad_dockgen_a2h_assessment_csv_filepath: ${paths.data_dir}/DockGen/dockgen_apo_to_holo_assessment.csv
+moad_filter_using_a2h_assessment: true # if `moad_a2h_assessment_csv_filepath` is provided, filter the dataset using the assessment
+moad_a2h_min_tmscore: 0.7
+moad_a2h_max_rmsd: 5.0
+moad_a2h_min_protein_length: 50
+moad_a2h_max_protein_length: 750
+moad_a2h_min_ligand_length: 1
+moad_a2h_max_ligand_length: 300
+remove_pdbbind: false
+moad_binding_affinity_values_path: ${paths.data_dir}/moad/moad_binding_affinity_data/binding_affinity_values.txt # path to the text file containing Binding MOAD's binding affinity values
+moad_map_binding_affinities_to_superligands: false # if `true`, map experimental binding affinities to the preprocessed Binding MOAD superligands in a potentially-noisy manner (i.e., this mapping may not always be correct)
+# van der Mers dataset arguments
+pdbsidechain_dir: ${paths.data_dir}/pdbsidechain/pdb_2021aug02
+pdbsidechain_esm_embeddings_path: ${paths.data_dir}/pdbsidechain/embeddings_output
+pdbsidechain_esm_embeddings_sequences_path: ${paths.data_dir}/pdbsidechain/pdbsidechain_sequences.fasta
+pdbsidechain_apo_protein_structure_dir: ${paths.data_dir}/pdbsidechain/pdbsidechain_holo_aligned_esmfold_structures
+pdbsidechain_metadata_dir: ${paths.data_dir}/pdbsidechain
+pdbsidechain_a2h_assessment_csv_filepath: ${paths.data_dir}/pdbsidechain/pdbsidechain_apo_to_holo_assessment.csv
+pdbsidechain_filter_using_a2h_assessment: true # if `pdbsidechain_a2h_assessment_csv_filepath` is provided, filter the dataset using the assessment
+pdbsidechain_a2h_min_tmscore: 0.8
+pdbsidechain_a2h_max_rmsd: 4.0
+pdbsidechain_a2h_min_protein_length: 50
+pdbsidechain_a2h_max_protein_length: 750
+pdbsidechain_postprocess_min_protein_length: 100
+pdbsidechain_postprocess_max_protein_length: 400
+vandermers_max_dist: 5.0
+vandermers_buffer_residue_num: 7
+vandermers_min_contacts: 5
+vandermers_max_surrogate_binding_affinity: 15.0
+vandermers_second_ligand_max_closeness: 10.0
+vandermers_extract_second_ligand: false
+vandermers_use_prob_as_surrogate_binding_affinity: false
diff --git a/forks/FlowDock/configs/data/esmfold_apo_to_holo_alignment.yaml b/forks/FlowDock/configs/data/esmfold_apo_to_holo_alignment.yaml
new file mode 100644
index 00000000..fee2771e
--- /dev/null
+++ b/forks/FlowDock/configs/data/esmfold_apo_to_holo_alignment.yaml
@@ -0,0 +1,5 @@
+dataset: pdbbind # the dataset to use - NOTE: must be one of (`pdbbind`, `moad`, `dockgen`, `pdbsidechain`)
+data_dir: ${oc.env:PROJECT_ROOT}/data # where the processed PDBBind, Binding MOAD, and van der Mers datasets are placed
+esmfold_structures_dir: ${oc.env:PROJECT_ROOT}/data/${resolve_dataset_path_dirname:${dataset}}/${dataset}_esmfold_structures # where the ESMFold-predicted protein structures are placed
+output_dir: ${oc.env:PROJECT_ROOT}/data/${resolve_dataset_path_dirname:${dataset}}/${dataset}_holo_aligned_esmfold_structures # where the holo-aligned ESMFold apo structures should be stored
+num_workers: 1 # number of CPU workers for parallel processing
diff --git a/forks/FlowDock/configs/data/esmfold_apo_to_holo_assessment.yaml b/forks/FlowDock/configs/data/esmfold_apo_to_holo_assessment.yaml
new file mode 100644
index 00000000..79f80406
--- /dev/null
+++ b/forks/FlowDock/configs/data/esmfold_apo_to_holo_assessment.yaml
@@ -0,0 +1,8 @@
+dataset: pdbbind # the dataset to use - NOTE: must be one of (`pdbbind`, `moad`, `dockgen`, `pdbsidechain`)
+data_dir: ${oc.env:PROJECT_ROOT}/data # where the processed datasets are placed
+aligned_dir: ${oc.env:PROJECT_ROOT}/data/${resolve_dataset_path_dirname:${dataset}}/${dataset}_holo_aligned_esmfold_structures # where the holo-aligned ESMFold apo structures are stored
+usalign_exec_path: ??? # path to a local copy of a precompiled USalign executable and its source files
+preprocessed_graph_dirs: null # if provided, all (dataset split) directories from which to load processed `HeteroData` graphs for apo-to-holo alignment assessment
+# preprocessed_graph_dirs: [data/cache/PDBBind_limit4_INDEXtimesplit_no_lig_overlap_train_maxLigSizeNone_H0_chainCutoffNone_esmEmbeddings_full_chainOrd_minProteinLength50_maxProteinLength750,] # note: an example of where preprocessed graphs may reside locally
+output_filepath: ${oc.env:PROJECT_ROOT}/data/${resolve_dataset_path_dirname:${dataset}}/${dataset}_apo_to_holo_assessment.csv # where to save the assessment results
+erroneous_output_filepath: ${oc.env:PROJECT_ROOT}/data/${resolve_dataset_path_dirname:${dataset}}/${dataset}_apo_to_holo_assessment_erroneous_file_inputs.csv # where to save the assessment errors
diff --git a/forks/FlowDock/configs/data/esmfold_sequence_preparation.yaml b/forks/FlowDock/configs/data/esmfold_sequence_preparation.yaml
new file mode 100644
index 00000000..0d6a5a18
--- /dev/null
+++ b/forks/FlowDock/configs/data/esmfold_sequence_preparation.yaml
@@ -0,0 +1,3 @@
+dataset: pdbbind # the dataset to use - NOTE: must be one of (`pdbbind`, `moad`, `dockgen`, `pdbsidechain`)
+input_fasta_file: ${oc.env:PROJECT_ROOT}/data/${resolve_dataset_path_dirname:${dataset}}/${dataset}_sequences.fasta # the input FASTA file to modify
+output_fasta_file: ${oc.env:PROJECT_ROOT}/data/${resolve_dataset_path_dirname:${dataset}}/${dataset}_esmfold_sequences.fasta # the input FASTA file to modify
diff --git a/forks/FlowDock/configs/debug/default.yaml b/forks/FlowDock/configs/debug/default.yaml
new file mode 100644
index 00000000..1886902b
--- /dev/null
+++ b/forks/FlowDock/configs/debug/default.yaml
@@ -0,0 +1,35 @@
+# @package _global_
+
+# default debugging setup, runs 1 full epoch
+# other debugging configs can inherit from this one
+
+# overwrite task name so debugging logs are stored in separate folder
+task_name: "debug"
+
+# disable callbacks and loggers during debugging
+callbacks: null
+logger: null
+
+extras:
+ ignore_warnings: False
+ enforce_tags: False
+
+# sets level of all command line loggers to 'DEBUG'
+# https://hydra.cc/docs/tutorials/basic/running_your_app/logging/
+hydra:
+ job_logging:
+ root:
+ level: DEBUG
+
+ # use this to also set hydra loggers to 'DEBUG'
+ # verbose: True
+
+trainer:
+ max_epochs: 1
+ accelerator: cpu # debuggers don't like gpus
+ devices: 1 # debuggers don't like multiprocessing
+ detect_anomaly: true # raise exception if NaN or +/-inf is detected in any tensor
+
+data:
+ num_workers: 0 # debuggers don't like multiprocessing
+ pin_memory: False # disable gpu memory pin
diff --git a/forks/FlowDock/configs/debug/fdr.yaml b/forks/FlowDock/configs/debug/fdr.yaml
new file mode 100644
index 00000000..7f2d34fa
--- /dev/null
+++ b/forks/FlowDock/configs/debug/fdr.yaml
@@ -0,0 +1,9 @@
+# @package _global_
+
+# runs 1 train, 1 validation and 1 test step
+
+defaults:
+ - default
+
+trainer:
+ fast_dev_run: true
diff --git a/forks/FlowDock/configs/debug/limit.yaml b/forks/FlowDock/configs/debug/limit.yaml
new file mode 100644
index 00000000..514d77fb
--- /dev/null
+++ b/forks/FlowDock/configs/debug/limit.yaml
@@ -0,0 +1,12 @@
+# @package _global_
+
+# uses only 1% of the training data and 5% of validation/test data
+
+defaults:
+ - default
+
+trainer:
+ max_epochs: 3
+ limit_train_batches: 0.01
+ limit_val_batches: 0.05
+ limit_test_batches: 0.05
diff --git a/forks/FlowDock/configs/debug/overfit.yaml b/forks/FlowDock/configs/debug/overfit.yaml
new file mode 100644
index 00000000..9906586a
--- /dev/null
+++ b/forks/FlowDock/configs/debug/overfit.yaml
@@ -0,0 +1,13 @@
+# @package _global_
+
+# overfits to 3 batches
+
+defaults:
+ - default
+
+trainer:
+ max_epochs: 20
+ overfit_batches: 3
+
+# model ckpt and early stopping need to be disabled during overfitting
+callbacks: null
diff --git a/forks/FlowDock/configs/debug/profiler.yaml b/forks/FlowDock/configs/debug/profiler.yaml
new file mode 100644
index 00000000..2bd7da87
--- /dev/null
+++ b/forks/FlowDock/configs/debug/profiler.yaml
@@ -0,0 +1,12 @@
+# @package _global_
+
+# runs with execution time profiling
+
+defaults:
+ - default
+
+trainer:
+ max_epochs: 1
+ profiler: "simple"
+ # profiler: "advanced"
+ # profiler: "pytorch"
diff --git a/forks/FlowDock/configs/environment/default.yaml b/forks/FlowDock/configs/environment/default.yaml
new file mode 100644
index 00000000..758b22bc
--- /dev/null
+++ b/forks/FlowDock/configs/environment/default.yaml
@@ -0,0 +1,2 @@
+defaults:
+ - _self_
diff --git a/forks/FlowDock/configs/environment/lightning.yaml b/forks/FlowDock/configs/environment/lightning.yaml
new file mode 100644
index 00000000..a5e70f45
--- /dev/null
+++ b/forks/FlowDock/configs/environment/lightning.yaml
@@ -0,0 +1 @@
+_target_: lightning.fabric.plugins.environments.LightningEnvironment
diff --git a/forks/FlowDock/configs/environment/slurm.yaml b/forks/FlowDock/configs/environment/slurm.yaml
new file mode 100644
index 00000000..55a5522a
--- /dev/null
+++ b/forks/FlowDock/configs/environment/slurm.yaml
@@ -0,0 +1,3 @@
+_target_: lightning.fabric.plugins.environments.SLURMEnvironment
+auto_requeue: true
+requeue_signal: null
diff --git a/forks/FlowDock/configs/eval.yaml b/forks/FlowDock/configs/eval.yaml
new file mode 100644
index 00000000..b27fbc37
--- /dev/null
+++ b/forks/FlowDock/configs/eval.yaml
@@ -0,0 +1,46 @@
+# @package _global_
+
+defaults:
+ - data: combined # choose datamodule with `test_dataloader()` for evaluation
+ - model: flowdock_fm
+ - logger: null
+ - strategy: default
+ - trainer: default
+ - paths: default
+ - extras: default
+ - hydra: default
+ - environment: default
+ - _self_
+
+task_name: "eval"
+
+tags: ["eval", "combined", "flowdock_fm"]
+
+# passing checkpoint path is necessary for evaluation
+ckpt_path: ???
+
+# model arguments
+model:
+ cfg:
+ mol_encoder:
+ from_pretrained: false
+ protein_encoder:
+ from_pretrained: false
+ relational_reasoning:
+ from_pretrained: false
+ contact_predictor:
+ from_pretrained: false
+ score_head:
+ from_pretrained: false
+ confidence:
+ from_pretrained: false
+ affinity:
+ from_pretrained: false
+ task:
+ freeze_mol_encoder: true
+ freeze_protein_encoder: false
+ freeze_relational_reasoning: false
+ freeze_contact_predictor: false
+ freeze_score_head: false
+ freeze_confidence: true
+ freeze_affinity: false
diff --git a/forks/FlowDock/configs/experiment/flowdock_fm.yaml b/forks/FlowDock/configs/experiment/flowdock_fm.yaml
new file mode 100644
index 00000000..b14b3e27
--- /dev/null
+++ b/forks/FlowDock/configs/experiment/flowdock_fm.yaml
@@ -0,0 +1,35 @@
+# @package _global_
+
+# to execute this experiment run:
+# python train.py experiment=flowdock_fm
+
+defaults:
+ - override /data: combined
+ - override /model: flowdock_fm
+ - override /callbacks: default
+ - override /trainer: default
+
+# all parameters below will be merged with parameters from default configurations set above
+# this allows you to overwrite only specified parameters
+
+tags: ["flowdock_fm", "combined_dataset"]
+
+seed: 496
+
+trainer:
+ max_epochs: 300
+ check_val_every_n_epoch: 5 # NOTE: we increase this since validation steps involve full model sampling and evaluation
+ reload_dataloaders_every_n_epochs: 1
+
+model:
+ optimizer:
+ lr: 2e-4
+ compile: false
+
+data:
+ batch_size: 16
+
+logger:
+ wandb:
+ tags: ${tags}
+ group: "FlowDock-FM"
diff --git a/forks/FlowDock/configs/extras/default.yaml b/forks/FlowDock/configs/extras/default.yaml
new file mode 100644
index 00000000..b9c6b622
--- /dev/null
+++ b/forks/FlowDock/configs/extras/default.yaml
@@ -0,0 +1,8 @@
+# disable python warnings if they annoy you
+ignore_warnings: False
+
+# ask user for tags if none are provided in the config
+enforce_tags: True
+
+# pretty print config tree at the start of the run using Rich library
+print_config: True
diff --git a/forks/FlowDock/configs/hparams_search/combined_optuna.yaml b/forks/FlowDock/configs/hparams_search/combined_optuna.yaml
new file mode 100644
index 00000000..43960de1
--- /dev/null
+++ b/forks/FlowDock/configs/hparams_search/combined_optuna.yaml
@@ -0,0 +1,50 @@
+# @package _global_
+
+# example hyperparameter optimization of some experiment with Optuna:
+# python train.py -m hparams_search=mnist_optuna experiment=example
+
+defaults:
+ - override /hydra/sweeper: optuna
+
+# choose metric which will be optimized by Optuna
+# make sure this is the correct name of some metric logged in lightning module!
+optimized_metric: "val/loss"
+
+# here we define Optuna hyperparameter search
+# it optimizes for value returned from function with @hydra.main decorator
+# docs: https://hydra.cc/docs/next/plugins/optuna_sweeper
+hydra:
+ mode: "MULTIRUN" # set hydra to multirun by default if this config is attached
+
+ sweeper:
+ _target_: hydra_plugins.hydra_optuna_sweeper.optuna_sweeper.OptunaSweeper
+
+ # storage URL to persist optimization results
+ # for example, you can use SQLite if you set 'sqlite:///example.db'
+ storage: null
+
+ # name of the study to persist optimization results
+ study_name: null
+
+ # number of parallel workers
+ n_jobs: 1
+
+ # 'minimize' or 'maximize' the objective
+ direction: minimize
+
+ # total number of runs that will be executed
+ n_trials: 20
+
+ # choose Optuna hyperparameter sampler
+ # you can choose bayesian sampler (tpe), random search (without optimization), grid sampler, and others
+ # docs: https://optuna.readthedocs.io/en/stable/reference/samplers.html
+ sampler:
+ _target_: optuna.samplers.TPESampler
+ seed: 1234
+ n_startup_trials: 10 # number of random sampling runs before optimization starts
+
+ # define hyperparameter search space
+ params:
+ model.optimizer.lr: interval(0.0001, 0.1)
+ data.batch_size: choice(2, 4, 8, 16)
+ model.net.hidden_dim: choice(64, 128, 256)
diff --git a/forks/FlowDock/configs/hydra/default.yaml b/forks/FlowDock/configs/hydra/default.yaml
new file mode 100644
index 00000000..a61e9b3a
--- /dev/null
+++ b/forks/FlowDock/configs/hydra/default.yaml
@@ -0,0 +1,19 @@
+# https://hydra.cc/docs/configure_hydra/intro/
+
+# enable color logging
+defaults:
+ - override hydra_logging: colorlog
+ - override job_logging: colorlog
+
+# output directory, generated dynamically on each run
+run:
+ dir: ${paths.log_dir}/${task_name}/runs/${now:%Y-%m-%d}_${now:%H-%M-%S}
+sweep:
+ dir: ${paths.log_dir}/${task_name}/multiruns/${now:%Y-%m-%d}_${now:%H-%M-%S}
+ subdir: ${hydra.job.num}
+
+job_logging:
+ handlers:
+ file:
+ # Incorporates fix from https://github.com/facebookresearch/hydra/pull/2242
+ filename: ${hydra.runtime.output_dir}/${task_name}.log
diff --git a/forks/FlowDock/configs/local/.gitkeep b/forks/FlowDock/configs/local/.gitkeep
new file mode 100644
index 00000000..e69de29b
diff --git a/forks/FlowDock/configs/logger/aim.yaml b/forks/FlowDock/configs/logger/aim.yaml
new file mode 100644
index 00000000..8f9f6ada
--- /dev/null
+++ b/forks/FlowDock/configs/logger/aim.yaml
@@ -0,0 +1,28 @@
+# https://aimstack.io/
+
+# example usage in lightning module:
+# https://github.com/aimhubio/aim/blob/main/examples/pytorch_lightning_track.py
+
+# open the Aim UI with the following command (run in the folder containing the `.aim` folder):
+# `aim up`
+
+aim:
+ _target_: aim.pytorch_lightning.AimLogger
+ repo: ${paths.root_dir} # .aim folder will be created here
+ # repo: "aim://ip_address:port" # can instead provide IP address pointing to Aim remote tracking server which manages the repo, see https://aimstack.readthedocs.io/en/latest/using/remote_tracking.html#
+
+ # aim allows to group runs under experiment name
+ experiment: null # any string, set to "default" if not specified
+
+ train_metric_prefix: "train/"
+ val_metric_prefix: "val/"
+ test_metric_prefix: "test/"
+
+ # sets the tracking interval in seconds for system usage metrics (CPU, GPU, memory, etc.)
+ system_tracking_interval: 10 # set to null to disable system metrics tracking
+
+ # enable/disable logging of system params such as installed packages, git info, env vars, etc.
+ log_system_params: true
+
+ # enable/disable tracking console logs (default value is true)
+ capture_terminal_logs: false # set to false to avoid infinite console log loop issue https://github.com/aimhubio/aim/issues/2550
diff --git a/forks/FlowDock/configs/logger/comet.yaml b/forks/FlowDock/configs/logger/comet.yaml
new file mode 100644
index 00000000..9de43235
--- /dev/null
+++ b/forks/FlowDock/configs/logger/comet.yaml
@@ -0,0 +1,12 @@
+# https://www.comet.ml
+
+comet:
+ _target_: lightning.pytorch.loggers.comet.CometLogger
+ api_key: ${oc.env:COMET_API_TOKEN} # api key is loaded from environment variable
+ save_dir: "${paths.output_dir}"
+ project_name: "FlowDock_FM"
+ rest_api_key: null
+ # experiment_name: ""
+ experiment_key: null # set to resume experiment
+ offline: False
+ prefix: ""
diff --git a/forks/FlowDock/configs/logger/csv.yaml b/forks/FlowDock/configs/logger/csv.yaml
new file mode 100644
index 00000000..fa028e9c
--- /dev/null
+++ b/forks/FlowDock/configs/logger/csv.yaml
@@ -0,0 +1,7 @@
+# csv logger built in lightning
+
+csv:
+ _target_: lightning.pytorch.loggers.csv_logs.CSVLogger
+ save_dir: "${paths.output_dir}"
+ name: "csv/"
+ prefix: ""
diff --git a/forks/FlowDock/configs/logger/many_loggers.yaml b/forks/FlowDock/configs/logger/many_loggers.yaml
new file mode 100644
index 00000000..dd586800
--- /dev/null
+++ b/forks/FlowDock/configs/logger/many_loggers.yaml
@@ -0,0 +1,9 @@
+# train with many loggers at once
+
+defaults:
+ # - comet
+ - csv
+ # - mlflow
+ # - neptune
+ - tensorboard
+ - wandb
diff --git a/forks/FlowDock/configs/logger/mlflow.yaml b/forks/FlowDock/configs/logger/mlflow.yaml
new file mode 100644
index 00000000..f8fb7e68
--- /dev/null
+++ b/forks/FlowDock/configs/logger/mlflow.yaml
@@ -0,0 +1,12 @@
+# https://mlflow.org
+
+mlflow:
+ _target_: lightning.pytorch.loggers.mlflow.MLFlowLogger
+ # experiment_name: ""
+ # run_name: ""
+ tracking_uri: ${paths.log_dir}/mlflow/mlruns # run `mlflow ui` command inside the `logs/mlflow/` dir to open the UI
+ tags: null
+ # save_dir: "./mlruns"
+ prefix: ""
+ artifact_location: null
+ # run_id: ""
diff --git a/forks/FlowDock/configs/logger/neptune.yaml b/forks/FlowDock/configs/logger/neptune.yaml
new file mode 100644
index 00000000..874a3fc3
--- /dev/null
+++ b/forks/FlowDock/configs/logger/neptune.yaml
@@ -0,0 +1,9 @@
+# https://neptune.ai
+
+neptune:
+ _target_: lightning.pytorch.loggers.neptune.NeptuneLogger
+ api_key: ${oc.env:NEPTUNE_API_TOKEN} # api key is loaded from environment variable
+ project: username/FlowDock_FM
+ # name: ""
+ log_model_checkpoints: True
+ prefix: ""
diff --git a/forks/FlowDock/configs/logger/tensorboard.yaml b/forks/FlowDock/configs/logger/tensorboard.yaml
new file mode 100644
index 00000000..2bd31f6d
--- /dev/null
+++ b/forks/FlowDock/configs/logger/tensorboard.yaml
@@ -0,0 +1,10 @@
+# https://www.tensorflow.org/tensorboard/
+
+tensorboard:
+ _target_: lightning.pytorch.loggers.tensorboard.TensorBoardLogger
+ save_dir: "${paths.output_dir}/tensorboard/"
+ name: null
+ log_graph: False
+ default_hp_metric: True
+ prefix: ""
+ # version: ""
diff --git a/forks/FlowDock/configs/logger/wandb.yaml b/forks/FlowDock/configs/logger/wandb.yaml
new file mode 100644
index 00000000..c66c04ef
--- /dev/null
+++ b/forks/FlowDock/configs/logger/wandb.yaml
@@ -0,0 +1,16 @@
+# https://wandb.ai
+
+wandb:
+ _target_: lightning.pytorch.loggers.wandb.WandbLogger
+ # name: "" # name of the run (normally generated by wandb)
+ save_dir: "${paths.output_dir}"
+ offline: False
+ id: null # pass correct id to resume experiment!
+ anonymous: null # enable anonymous logging
+ project: "FlowDock_FM"
+ log_model: False # upload lightning ckpts
+ prefix: "" # a string to put at the beginning of metric keys
+ entity: "bml-lab" # set to name of your wandb team
+ group: ""
+ tags: []
+ job_type: ""
diff --git a/forks/FlowDock/configs/model/flowdock_fm.yaml b/forks/FlowDock/configs/model/flowdock_fm.yaml
new file mode 100644
index 00000000..2812281f
--- /dev/null
+++ b/forks/FlowDock/configs/model/flowdock_fm.yaml
@@ -0,0 +1,148 @@
+_target_: flowdock.models.flowdock_fm_module.FlowDockFMLitModule
+
+net:
+ _target_: flowdock.models.components.flowdock.FlowDock
+ _partial_: true
+
+optimizer:
+ _target_: torch.optim.Adam
+ _partial_: true
+ lr: 2e-4
+ weight_decay: 0.0
+
+scheduler:
+ _target_: torch.optim.lr_scheduler.CosineAnnealingWarmRestarts
+ _partial_: true
+ T_0: ${int_divide:${trainer.max_epochs},15}
+ T_mult: 2
+ eta_min: 1e-8
+ verbose: true
+
+# compile model for faster training with pytorch 2.0
+compile: false
+
+# model arguments
+cfg:
+ mol_encoder:
+ node_channels: 512
+ pair_channels: 64
+ n_atom_encodings: 23
+ n_bond_encodings: 4
+ n_atom_pos_encodings: 6
+ n_stereo_encodings: 14
+ n_attention_heads: 8
+ attention_head_dim: 8
+ hidden_dim: 2048
+ max_path_integral_length: 6
+ n_transformer_stacks: 8
+ n_heads: 8
+ n_patches: ${data.n_lig_patches}
+ checkpoint_file: ${oc.env:PROJECT_ROOT}/checkpoints/neuralplexermodels_downstream_datasets_predictions/models/complex_structure_prediction.ckpt
+ megamolbart: null
+ from_pretrained: true
+
+ protein_encoder:
+ use_esm_embedding: true
+ esm_version: esm2_t33_650M_UR50D
+ esm_repr_layer: 33
+ residue_dim: 512
+ plm_embed_dim: 1280
+ n_aa_types: 21
+ atom_padding_dim: 37
+ n_atom_types: 4 # [C, N, O, S]
+ n_patches: ${data.n_protein_patches}
+ n_attention_heads: 8
+ scalar_dim: 16
+ point_dim: 4
+ pair_dim: 64
+ n_heads: 8
+ head_dim: 8
+ max_residue_degree: 32
+ n_encoder_stacks: 2
+ from_pretrained: true
+
+ relational_reasoning:
+ from_pretrained: true
+
+ contact_predictor:
+ n_stacks: 4
+ dropout: 0.01
+ from_pretrained: true
+
+ score_head:
+ fiber_dim: 64
+ hidden_dim: 512
+ n_stacks: 4
+ max_atom_degree: 8
+ from_pretrained: true
+
+ confidence:
+ enabled: true # whether the confidence prediction head is to be used e.g., during inference
+ fiber_dim: ${..score_head.fiber_dim}
+ hidden_dim: ${..score_head.hidden_dim}
+ n_stacks: ${..score_head.n_stacks}
+ from_pretrained: true
+
+ affinity:
+ enabled: true # whether the affinity prediction head is to be used e.g., during inference
+ fiber_dim: ${..score_head.fiber_dim}
+ hidden_dim: ${..score_head.hidden_dim}
+ n_stacks: ${..score_head.n_stacks}
+ ligand_pooling: sum # NOTE: must be a value in (`sum`, `mean`)
+ dropout: 0.01
+ from_pretrained: false
+
+ latent_model: default
+ prior_type: esmfold # NOTE: must be a value in (`gaussian`, `harmonic`, `esmfold`)
+
+ task:
+ pretrained: null
+ ligands: true
+ epoch_frac: ${data.epoch_frac}
+ label_name: null
+ sequence_crop_size: 1600
+ edge_crop_size: ${data.edge_crop_size} # NOTE: for dynamic batching via `max_n_edges`
+ max_masking_rate: 0.0
+ n_modes: 8
+ dropout: 0.01
+ # pretraining: true
+ freeze_mol_encoder: true
+ freeze_protein_encoder: false
+ freeze_relational_reasoning: false
+ freeze_contact_predictor: true
+ freeze_score_head: false
+ freeze_confidence: true
+ freeze_affinity: false
+ use_template: true
+ use_plddt: false
+ block_contact_decoding_scheme: "beam"
+ frozen_ligand_backbone: false
+ frozen_protein_backbone: false
+ single_protein_batch: true
+ contact_loss_weight: 0.2
+ global_score_loss_weight: 0.2
+ ligand_score_loss_weight: 0.1
+ clash_loss_weight: 10.0
+ local_distgeom_loss_weight: 10.0
+ drmsd_loss_weight: 2.0
+ distogram_loss_weight: 0.05
+ plddt_loss_weight: 1.0
+ affinity_loss_weight: 0.1
+ aux_batch_freq: 10 # NOTE: e.g., `10` means that auxiliary estimation losses will be calculated every 10th batch
+ global_max_sigma: 5.0
+ internal_max_sigma: 2.0
+ detect_covalent: true
+ # runtime configs
+ float32_matmul_precision: highest
+ # sampling configs
+ constrained_inpainting: false
+ visualize_generated_samples: true
+ # testing configs
+ loss_mode: auxiliary_estimation # NOTE: must be one of (`structure_prediction`, `auxiliary_estimation`, `auxiliary_estimation_without_structure_prediction`)
+ num_steps: 20
+ sampler: VDODE # NOTE: must be one of (`ODE`, `VDODE`)
+ sampler_eta: 1.0 # NOTE: this corresponds to the variance diminishing factor for the `VDODE` sampler, which offers a trade-off between exploration (1.0) and exploitation (> 1.0)
+ start_time: 1.0
+ eval_structure_prediction: false # whether to evaluate structure prediction performance (`true`) or instead only binding affinity performance (`false`) when running a test epoch
+ # overfitting configs
+ overfitting_example_name: ${data.overfitting_example_name}
diff --git a/forks/FlowDock/configs/paths/default.yaml b/forks/FlowDock/configs/paths/default.yaml
new file mode 100644
index 00000000..3898b289
--- /dev/null
+++ b/forks/FlowDock/configs/paths/default.yaml
@@ -0,0 +1,21 @@
+# path to root directory
+# this requires PROJECT_ROOT environment variable to exist
+# you can replace it with "." if you want the root to be the current working directory
+root_dir: ${oc.env:PROJECT_ROOT}
+
+# path to data directory
+data_dir: ${paths.root_dir}/data/
+
+# path to logging directory
+log_dir: ${paths.root_dir}/logs/
+
+# path to output directory, created dynamically by hydra
+# path generation pattern is specified in `configs/hydra/default.yaml`
+# use it to store all files generated during the run, like ckpts and metrics
+output_dir: ${hydra:runtime.output_dir}
+
+# path to working directory
+work_dir: ${hydra:runtime.cwd}
+
+# path to the directory containing the model checkpoints
+ckpt_dir: ${paths.root_dir}/checkpoints/
diff --git a/forks/FlowDock/configs/sample.yaml b/forks/FlowDock/configs/sample.yaml
new file mode 100644
index 00000000..6dda0aea
--- /dev/null
+++ b/forks/FlowDock/configs/sample.yaml
@@ -0,0 +1,75 @@
+# @package _global_
+
+defaults:
+ - data: combined # NOTE: this will not be referenced during sampling
+ - model: flowdock_fm
+ - logger: null
+ - strategy: default
+ - trainer: default
+ - paths: default
+ - extras: default
+ - hydra: default
+ - environment: default
+ - _self_
+
+task_name: "sample"
+
+tags: ["sample", "combined", "flowdock_fm"]
+
+# passing checkpoint path is necessary for sampling
+ckpt_path: ???
+
+# sampling arguments
+sampling_task: batched_structure_sampling # NOTE: must be one of (`batched_structure_sampling`)
+sample_id: null # optional identifier for the sampling run
+input_receptor: null # NOTE: must be either a protein sequence string (with chains separated by `|`) or a path to a PDB file (from which protein chain sequences will be parsed)
+input_ligand: null # NOTE: must be either a ligand SMILES string (with chains/fragments separated by `|`) or a path to a ligand SDF file (from which ligand SMILES will be parsed)
+input_template: null # path to a protein PDB file to use as a starting protein template for sampling (with an ESMFold prior model)
+out_path: ??? # path to which to save the output PDB and SDF files
+n_samples: 5 # number of structures to sample
+chunk_size: 5 # number of structures to concurrently sample within each batch segment - NOTE: `n_samples` should be evenly divisible by `chunk_size` to produce the expected number of outputs
+num_steps: 40 # number of sampling steps to perform
+latent_model: null # if provided, the type of latent model to use
+sampler: VDODE # sampling algorithm to use - NOTE: must be one of (`ODE`, `VDODE`)
+sampler_eta: 1.0 # the variance diminishing factor for the `VDODE` sampler - NOTE: offers a trade-off between exploration (1.0) and exploitation (> 1.0)
+start_time: "1.0" # time at which to start sampling
+max_chain_encoding_k: -1 # maximum number of chains to encode in the chain encoding
+exact_prior: false # whether to use the "ground-truth" binding site for sampling, if available
+prior_type: esmfold # the type of prior to use for sampling - NOTE: must be one of (`gaussian`, `harmonic`, `esmfold`)
+discard_ligand: false # whether to discard a given input ligand during sampling
+discard_sdf_coords: true # whether to discard the input ligand's 3D structure during sampling, if available
+detect_covalent: false # whether to detect covalent bonds between the input receptor and ligand
+use_template: true # whether to use the input protein template for sampling if one is provided
+separate_pdb: true # whether to save separate PDB files for each sampled structure instead of simply a single PDB file
+rank_outputs_by_confidence: true # whether to rank the sampled structures by estimated confidence
+plddt_ranking_type: ligand # the type of plDDT ranking to apply to generated samples - NOTE: must be one of (`protein`, `ligand`, `protein_ligand`)
+visualize_sample_trajectories: false # whether to visualize the generated samples' trajectories
+auxiliary_estimation_only: false # whether to only estimate auxiliary outputs (e.g., confidence, affinity) for the input (generated) samples (potentially derived from external sources)
+csv_path: null # if provided, the CSV file (with columns `id`, `input_receptor`, `input_ligand`, and `input_template`) from which to parse input receptors and ligands for sampling, overriding the `input_receptor` and `input_ligand` arguments in the process and ignoring the `input_template` for now
+esmfold_chunk_size: null # chunks axial attention computation to reduce memory usage from O(L^2) to O(L); equivalent to running a for loop over chunks of of each dimension; lower values will result in lower memory usage at the cost of speed; recommended values: 128, 64, 32
+
+# model arguments
+model:
+ cfg:
+ mol_encoder:
+ from_pretrained: false
+ protein_encoder:
+ from_pretrained: false
+ relational_reasoning:
+ from_pretrained: false
+ contact_predictor:
+ from_pretrained: false
+ score_head:
+ from_pretrained: false
+ confidence:
+ from_pretrained: false
+ affinity:
+ from_pretrained: false
+ task:
+ freeze_mol_encoder: true
+ freeze_protein_encoder: false
+ freeze_relational_reasoning: false
+ freeze_contact_predictor: false
+ freeze_score_head: false
+ freeze_confidence: true
+ freeze_affinity: false
diff --git a/forks/FlowDock/configs/strategy/ddp.yaml b/forks/FlowDock/configs/strategy/ddp.yaml
new file mode 100644
index 00000000..14933f3a
--- /dev/null
+++ b/forks/FlowDock/configs/strategy/ddp.yaml
@@ -0,0 +1,4 @@
+_target_: lightning.pytorch.strategies.DDPStrategy
+static_graph: false
+gradient_as_bucket_view: false
+find_unused_parameters: true
diff --git a/forks/FlowDock/configs/strategy/ddp_spawn.yaml b/forks/FlowDock/configs/strategy/ddp_spawn.yaml
new file mode 100644
index 00000000..e51e4f23
--- /dev/null
+++ b/forks/FlowDock/configs/strategy/ddp_spawn.yaml
@@ -0,0 +1,5 @@
+_target_: lightning.pytorch.strategies.DDPStrategy
+static_graph: false
+gradient_as_bucket_view: false
+find_unused_parameters: true
+start_method: spawn
diff --git a/forks/FlowDock/configs/strategy/deepspeed.yaml b/forks/FlowDock/configs/strategy/deepspeed.yaml
new file mode 100644
index 00000000..3c05b4c2
--- /dev/null
+++ b/forks/FlowDock/configs/strategy/deepspeed.yaml
@@ -0,0 +1,5 @@
+_target_: lightning.pytorch.strategies.DeepSpeedStrategy
+stage: 2
+offload_optimizer: false
+allgather_bucket_size: 200_000_000
+reduce_bucket_size: 200_000_000
diff --git a/forks/FlowDock/configs/strategy/default.yaml b/forks/FlowDock/configs/strategy/default.yaml
new file mode 100644
index 00000000..758b22bc
--- /dev/null
+++ b/forks/FlowDock/configs/strategy/default.yaml
@@ -0,0 +1,2 @@
+defaults:
+ - _self_
diff --git a/forks/FlowDock/configs/strategy/fsdp.yaml b/forks/FlowDock/configs/strategy/fsdp.yaml
new file mode 100644
index 00000000..12e29e86
--- /dev/null
+++ b/forks/FlowDock/configs/strategy/fsdp.yaml
@@ -0,0 +1,12 @@
+_target_: lightning.pytorch.strategies.FSDPStrategy
+sharding_strategy: ${resolve_variable:torch.distributed.fsdp.ShardingStrategy.FULL_SHARD}
+cpu_offload: null
+activation_checkpointing: null
+mixed_precision:
+ _target_: torch.distributed.fsdp.MixedPrecision
+ param_dtype: null
+ reduce_dtype: null
+ buffer_dtype: null
+ keep_low_precision_grads: false
+ cast_forward_inputs: false
+ cast_root_forward_inputs: true
diff --git a/forks/FlowDock/configs/strategy/optimized_ddp.yaml b/forks/FlowDock/configs/strategy/optimized_ddp.yaml
new file mode 100644
index 00000000..b8b4b312
--- /dev/null
+++ b/forks/FlowDock/configs/strategy/optimized_ddp.yaml
@@ -0,0 +1,4 @@
+_target_: lightning.pytorch.strategies.DDPStrategy
+static_graph: true
+gradient_as_bucket_view: true
+find_unused_parameters: false
diff --git a/forks/FlowDock/configs/train.yaml b/forks/FlowDock/configs/train.yaml
new file mode 100644
index 00000000..d2205a40
--- /dev/null
+++ b/forks/FlowDock/configs/train.yaml
@@ -0,0 +1,51 @@
+# @package _global_
+
+# specify here default configuration
+# order of defaults determines the order in which configs override each other
+defaults:
+ - _self_
+ - data: combined
+ - model: flowdock_fm
+ - callbacks: default
+ - logger: null # set logger here or use command line (e.g. `python train.py logger=tensorboard`)
+ - strategy: default
+ - trainer: default
+ - paths: default
+ - extras: default
+ - hydra: default
+ - environment: default
+
+ # experiment configs allow for version control of specific hyperparameters
+ # e.g. best hyperparameters for given model and datamodule
+ - experiment: null
+
+ # config for hyperparameter optimization
+ - hparams_search: null
+
+ # optional local config for machine/user specific settings
+ # it's optional since it doesn't need to exist and is excluded from version control
+ - optional local: default
+
+ # debugging config (enable through command line, e.g. `python train.py debug=default)
+ - debug: null
+
+# task name, determines output directory path
+task_name: "train"
+
+# tags to help you identify your experiments
+# you can overwrite this in experiment configs
+# overwrite from command line with `python train.py tags="[first_tag, second_tag]"`
+tags: ["train", "combined", "flowdock_fm"]
+
+# set False to skip model training
+train: True
+
+# evaluate on test set, using best model weights achieved during training
+# lightning chooses best weights based on the metric specified in checkpoint callback
+test: False
+
+# simply provide checkpoint path to resume training
+ckpt_path: null
+
+# seed for random number generators in pytorch, numpy and python.random
+seed: null
diff --git a/forks/FlowDock/configs/trainer/cpu.yaml b/forks/FlowDock/configs/trainer/cpu.yaml
new file mode 100644
index 00000000..b7d6767e
--- /dev/null
+++ b/forks/FlowDock/configs/trainer/cpu.yaml
@@ -0,0 +1,5 @@
+defaults:
+ - default
+
+accelerator: cpu
+devices: 1
diff --git a/forks/FlowDock/configs/trainer/ddp.yaml b/forks/FlowDock/configs/trainer/ddp.yaml
new file mode 100644
index 00000000..ab8f8900
--- /dev/null
+++ b/forks/FlowDock/configs/trainer/ddp.yaml
@@ -0,0 +1,9 @@
+defaults:
+ - default
+
+strategy: ddp
+
+accelerator: gpu
+devices: 4
+num_nodes: 1
+sync_batchnorm: True
diff --git a/forks/FlowDock/configs/trainer/ddp_sim.yaml b/forks/FlowDock/configs/trainer/ddp_sim.yaml
new file mode 100644
index 00000000..8404419e
--- /dev/null
+++ b/forks/FlowDock/configs/trainer/ddp_sim.yaml
@@ -0,0 +1,7 @@
+defaults:
+ - default
+
+# simulate DDP on CPU, useful for debugging
+accelerator: cpu
+devices: 2
+strategy: ddp_spawn
diff --git a/forks/FlowDock/configs/trainer/ddp_spawn.yaml b/forks/FlowDock/configs/trainer/ddp_spawn.yaml
new file mode 100644
index 00000000..021150aa
--- /dev/null
+++ b/forks/FlowDock/configs/trainer/ddp_spawn.yaml
@@ -0,0 +1,9 @@
+defaults:
+ - default
+
+strategy: ddp_spawn
+
+accelerator: gpu
+devices: 4
+num_nodes: 1
+sync_batchnorm: True
diff --git a/forks/FlowDock/configs/trainer/default.yaml b/forks/FlowDock/configs/trainer/default.yaml
new file mode 100644
index 00000000..ca4a10bd
--- /dev/null
+++ b/forks/FlowDock/configs/trainer/default.yaml
@@ -0,0 +1,29 @@
+_target_: lightning.pytorch.trainer.Trainer
+
+default_root_dir: ${paths.output_dir}
+
+min_epochs: 1 # prevents early stopping
+max_epochs: 10
+
+accelerator: cpu
+devices: 1
+
+# mixed precision for extra speed-up
+# precision: 16
+
+# perform a validation loop every N training epochs
+check_val_every_n_epoch: 1
+
+# set True to to ensure deterministic results
+# makes training slower but gives more reproducibility than just setting seeds
+deterministic: False
+
+# determine the frequency of how often to reload the dataloaders
+reload_dataloaders_every_n_epochs: 1
+
+# if `gradient_clip_val` is not `null`, gradients will be norm-clipped during training
+gradient_clip_algorithm: norm
+gradient_clip_val: 1.0
+
+# if `num_sanity_val_steps` is > 0, then specifically that many validation steps will be run during the first call to `trainer.fit`
+num_sanity_val_steps: 0
diff --git a/forks/FlowDock/configs/trainer/gpu.yaml b/forks/FlowDock/configs/trainer/gpu.yaml
new file mode 100644
index 00000000..b2389510
--- /dev/null
+++ b/forks/FlowDock/configs/trainer/gpu.yaml
@@ -0,0 +1,5 @@
+defaults:
+ - default
+
+accelerator: gpu
+devices: 1
diff --git a/forks/FlowDock/configs/trainer/mps.yaml b/forks/FlowDock/configs/trainer/mps.yaml
new file mode 100644
index 00000000..1ecf6d5c
--- /dev/null
+++ b/forks/FlowDock/configs/trainer/mps.yaml
@@ -0,0 +1,5 @@
+defaults:
+ - default
+
+accelerator: mps
+devices: 1
diff --git a/forks/FlowDock/environments/flowdock_environment.yaml b/forks/FlowDock/environments/flowdock_environment.yaml
new file mode 100644
index 00000000..a74a7f8c
--- /dev/null
+++ b/forks/FlowDock/environments/flowdock_environment.yaml
@@ -0,0 +1,448 @@
+name: FlowDock
+channels:
+ - pyg
+ - iopath
+ - pytorch
+ - senyan.dev
+ - nvidia
+ - defaults
+ - conda-forge
+dependencies:
+ - _libgcc_mutex=0.1=conda_forge
+ - _openmp_mutex=4.5=2_kmp_llvm
+ - aiohttp=3.9.3=py39hd1e30aa_1
+ - aiosignal=1.3.1=pyhd8ed1ab_0
+ - antlr-python-runtime=4.9.3=pyhd8ed1ab_1
+ - anyio=4.2.0=py39h06a4308_0
+ - argon2-cffi=23.1.0=pyhd8ed1ab_0
+ - argon2-cffi-bindings=21.2.0=py39hd1e30aa_4
+ - arpack=3.7.0=hdefa2d7_2
+ - arrow=1.2.3=py39h06a4308_1
+ - astunparse=1.6.3=pyhd8ed1ab_0
+ - async-lru=2.0.4=pyhd8ed1ab_0
+ - async-timeout=4.0.3=pyhd8ed1ab_0
+ - attrs=22.1.0=py39h06a4308_0
+ - babel=2.14.0=pyhd8ed1ab_0
+ - backcall=0.2.0=pyhd3eb1b0_0
+ - backoff=2.2.1=pyhd8ed1ab_0
+ - beautifulsoup4=4.12.2=py39h06a4308_0
+ - blas=2.116=mkl
+ - blas-devel=3.9.0=16_linux64_mkl
+ - bleach=6.1.0=pyhd8ed1ab_0
+ - blessed=1.19.1=pyhe4f9e05_2
+ - blosc=1.21.5=h0f2a231_0
+ - boost=1.84.0=h8003fee_2
+ - boost-cpp=1.84.0=h44aadfe_2
+ - bottleneck=1.3.7=py39ha9d4c09_0
+ - brotli=1.0.9=h166bdaf_9
+ - brotli-bin=1.0.9=h166bdaf_9
+ - brotli-python=1.0.9=py39h5a03fae_9
+ - bson=0.5.9=py_0
+ - bzip2=1.0.8=h7f98852_4
+ - c-ares=1.25.0=hd590300_0
+ - c-blosc2=2.12.0=hb4ffafa_0
+ - ca-certificates=2024.2.2=hbcca054_0
+ - cachecontrol=0.13.1=pyhd8ed1ab_0
+ - cachecontrol-with-filecache=0.13.1=pyhd8ed1ab_0
+ - cached-property=1.5.2=hd8ed1ab_1
+ - cached_property=1.5.2=pyha770c72_1
+ - cachetools=5.3.2=pyhd8ed1ab_0
+ - cairo=1.18.0=h3faef2a_0
+ - certifi=2024.2.2=pyhd8ed1ab_0
+ - cffi=1.15.1=py39h5eee18b_3
+ - cfgv=3.3.1=py39h06a4308_0
+ - charset-normalizer=3.2.0=pyhd8ed1ab_0
+ - cleo=2.0.1=py39h06a4308_0
+ - click=8.0.4=py39h06a4308_0
+ - colorama=0.4.6=pyhd8ed1ab_0
+ - comm=0.2.1=pyhd8ed1ab_0
+ - contourpy=1.2.0=py39hdb19cb5_0
+ - crashtest=0.4.1=py39h06a4308_0
+ - croniter=1.3.7=py39h06a4308_0
+ - cryptography=41.0.2=py39h22a60cf_0
+ - cuda-cudart=11.8.89=0
+ - cuda-cupti=11.8.87=0
+ - cuda-libraries=11.8.0=0
+ - cuda-nvrtc=11.8.89=0
+ - cuda-nvtx=11.8.86=0
+ - cuda-runtime=11.8.0=0
+ - cudatoolkit=11.8.0=h4ba93d1_12
+ - cudatoolkit-dev=11.8.0=h1fa729e_6
+ - curl=8.1.2=h409715c_0
+ - cycler=0.11.0=pyhd3eb1b0_0
+ - cython=3.0.7=py39h3d6467e_0
+ - dateutils=0.6.12=py_0
+ - dbus=1.13.18=hb2f20db_0
+ - debugpy=1.8.0=py39h3d6467e_1
+ - decorator=5.1.1=pyhd8ed1ab_0
+ - deepdiff=6.3.1=pyhd8ed1ab_0
+ - defusedxml=0.7.1=pyhd8ed1ab_0
+ - distlib=0.3.7=pyhd8ed1ab_0
+ - dulwich=0.21.3=py39h5eee18b_0
+ - entrypoints=0.4=pyhd8ed1ab_0
+ - exceptiongroup=1.0.4=py39h06a4308_0
+ - expat=2.6.2=h59595ed_0
+ - ffmpeg=4.3=hf484d3e_0
+ - fftw=3.3.10=nompi_hc118613_108
+ - filelock=3.12.2=pyhd8ed1ab_0
+ - font-ttf-dejavu-sans-mono=2.37=hab24e00_0
+ - font-ttf-inconsolata=3.000=h77eed37_0
+ - font-ttf-source-code-pro=2.038=h77eed37_0
+ - font-ttf-ubuntu=0.83=h77eed37_1
+ - fontconfig=2.14.2=h14ed4e7_0
+ - fonts-conda-ecosystem=1=0
+ - fonts-conda-forge=1=0
+ - fonttools=4.25.0=pyhd3eb1b0_0
+ - fqdn=1.5.1=pyhd8ed1ab_0
+ - freetype=2.12.1=hca18f0e_1
+ - frozenlist=1.4.1=py39hd1e30aa_0
+ - fvcore=0.1.5.post20221221=pyhd8ed1ab_0
+ - glib=2.80.0=hf2295e7_1
+ - glib-tools=2.80.0=hde27a5a_1
+ - gmp=6.2.1=h58526e2_0
+ - gmpy2=2.1.2=py39h376b7d2_1
+ - gnutls=3.6.13=h85f3911_1
+ - greenlet=3.0.1=py39h6a678d5_0
+ - h11=0.12.0=pyhd3eb1b0_0
+ - hdf4=4.2.15=h9772cbc_5
+ - hdf5=1.14.1=nompi_h4f84152_100
+ - hydra-core=1.3.2=pyhd8ed1ab_0
+ - icu=73.2=h59595ed_0
+ - identify=2.5.5=py39h06a4308_0
+ - idna=3.4=pyhd8ed1ab_0
+ - importlib-metadata=6.0.0=py39h06a4308_0
+ - importlib-resources=3.3.1=pyhd8ed1ab_0
+ - importlib_metadata=6.0.0=hd3eb1b0_0
+ - importlib_resources=5.2.0=pyhd3eb1b0_1
+ - iniconfig=1.1.1=pyhd3eb1b0_0
+ - inquirer=3.1.3=pyhd8ed1ab_0
+ - iopath=0.1.9=py39
+ - ipykernel=6.28.0=py39h06a4308_0
+ - ipython=8.15.0=py39h06a4308_0
+ - ipywidgets=8.1.1=pyhd8ed1ab_0
+ - isoduration=20.11.0=pyhd8ed1ab_0
+ - itsdangerous=2.0.1=pyhd3eb1b0_0
+ - jaraco.classes=3.2.1=pyhd3eb1b0_0
+ - jedi=0.18.1=py39h06a4308_1
+ - jeepney=0.7.1=pyhd3eb1b0_0
+ - jinja2=3.1.2=pyhd8ed1ab_1
+ - joblib=1.3.2=pyhd8ed1ab_0
+ - jpeg=9e=h166bdaf_2
+ - json5=0.9.14=pyhd8ed1ab_0
+ - jsonpointer=2.4=py39hf3d152e_3
+ - jsonschema=4.17.3=py39h06a4308_0
+ - jsonschema-with-format-nongpl=4.17.3=pyhd8ed1ab_0
+ - jupyter-lsp=2.2.1=pyhd8ed1ab_0
+ - jupyter_client=8.6.0=pyhd8ed1ab_0
+ - jupyter_core=5.7.1=py39hf3d152e_0
+ - jupyter_events=0.6.3=pyhd8ed1ab_1
+ - jupyter_server=2.10.0=pyhd8ed1ab_0
+ - jupyter_server_terminals=0.5.1=pyhd8ed1ab_0
+ - jupyterlab=4.0.10=pyhd8ed1ab_0
+ - jupyterlab_pygments=0.3.0=pyhd8ed1ab_0
+ - jupyterlab_server=2.24.0=pyhd8ed1ab_0
+ - jupyterlab_widgets=3.0.9=pyhd8ed1ab_0
+ - keyring=24.2.0=py39hf3d152e_0
+ - keyutils=1.6.1=h166bdaf_0
+ - kiwisolver=1.4.4=py39h6a678d5_0
+ - krb5=1.20.1=h81ceb04_0
+ - lame=3.100=h166bdaf_1003
+ - lcms2=2.12=hddcbb42_0
+ - ld_impl_linux-64=2.40=h41732ed_0
+ - lerc=4.0.0=h27087fc_0
+ - libaec=1.1.2=h59595ed_1
+ - libblas=3.9.0=16_linux64_mkl
+ - libboost=1.84.0=h8013b2b_2
+ - libboost-devel=1.84.0=h00ab1b0_2
+ - libboost-headers=1.84.0=ha770c72_2
+ - libboost-python=1.84.0=py39hda80f44_2
+ - libboost-python-devel=1.84.0=py39h8003fee_2
+ - libbrotlicommon=1.0.9=h166bdaf_9
+ - libbrotlidec=1.0.9=h166bdaf_9
+ - libbrotlienc=1.0.9=h166bdaf_9
+ - libcblas=3.9.0=16_linux64_mkl
+ - libcublas=11.11.3.6=0
+ - libcufft=10.9.0.58=0
+ - libcufile=1.7.1.12=0
+ - libcurand=10.3.3.129=0
+ - libcurl=8.1.2=h409715c_0
+ - libcusolver=11.4.1.48=0
+ - libcusparse=11.7.5.86=0
+ - libdeflate=1.17=h0b41bf4_0
+ - libedit=3.1.20191231=he28a2e2_2
+ - libev=4.33=hd590300_2
+ - libexpat=2.6.2=h59595ed_0
+ - libffi=3.4.2=h7f98852_5
+ - libgcc-ng=13.1.0=he5830b7_0
+ - libgfortran-ng=13.1.0=h69a702a_0
+ - libgfortran5=13.1.0=h15d22d2_0
+ - libglib=2.80.0=hf2295e7_1
+ - libgomp=13.1.0=he5830b7_0
+ - libhwloc=2.4.1=h6746aa3_0
+ - libiconv=1.17=h166bdaf_0
+ - libjpeg-turbo=2.0.0=h9bf148f_0
+ - liblapack=3.9.0=16_linux64_mkl
+ - liblapacke=3.9.0=16_linux64_mkl
+ - libnghttp2=1.58.0=h47da74e_1
+ - libnpp=11.8.0.86=0
+ - libnsl=2.0.0=h7f98852_0
+ - libnvjpeg=11.9.0.86=0
+ - libpng=1.6.39=h753d276_0
+ - libsodium=1.0.18=h36c2ea0_1
+ - libsqlite=3.42.0=h2797004_0
+ - libssh2=1.11.0=h0841786_0
+ - libstdcxx-ng=13.1.0=hfd8a6a1_0
+ - libtiff=4.2.0=hf544144_3
+ - libuuid=2.38.1=h0b41bf4_0
+ - libwebp-base=1.3.1=hd590300_0
+ - libxcb=1.15=h0b41bf4_0
+ - libxcrypt=4.4.36=hd590300_1
+ - libxml2=2.12.6=h232c23b_1
+ - libxslt=1.1.39=h76b75d6_0
+ - libzip=1.10.1=h2629f0a_3
+ - libzlib=1.2.13=hd590300_5
+ - llvm-openmp=15.0.7=h0cdce71_0
+ - lxml=4.8.0=py39hb9d737c_3
+ - lz4-c=1.9.3=h9c3ff4c_1
+ - lzo=2.10=h516909a_1000
+ - markdown-it-py=2.2.0=py39h06a4308_1
+ - markupsafe=2.1.1=py39h7f8727e_0
+ - matplotlib-base=3.8.0=py39he9076e7_1
+ - matplotlib-inline=0.1.6=pyhd8ed1ab_0
+ - mda-xdrlib=0.2.0=pyhd8ed1ab_0
+ - mdtraj=1.9.9=py39h67e0f1a_1
+ - mdurl=0.1.0=py39h06a4308_0
+ - mendeleev=0.15.0=pyhc1e730c_0
+ - mistune=3.0.2=pyhd8ed1ab_0
+ - mkl=2022.1.0=h84fe81f_915
+ - mkl-devel=2022.1.0=ha770c72_916
+ - mkl-include=2022.1.0=h84fe81f_915
+ - more-itertools=8.12.0=pyhd3eb1b0_0
+ - mpc=1.3.1=hfe3b2da_0
+ - mpfr=4.2.0=hb012696_0
+ - mpmath=1.3.0=py39h06a4308_0
+ - msgpack-python=1.0.3=py39hd09550d_0
+ - multidict=6.0.5=py39hd1e30aa_0
+ - munkres=1.1.4=pyh9f0ad1d_0
+ - nbclient=0.8.0=pyhd8ed1ab_0
+ - nbconvert-core=7.14.1=pyhd8ed1ab_0
+ - nbformat=5.9.2=pyhd8ed1ab_0
+ - ncurses=6.4=hcb278e6_0
+ - nest-asyncio=1.5.8=pyhd8ed1ab_0
+ - nettle=3.6=he412f7d_0
+ - networkx=3.1=py39h06a4308_0
+ - nodeenv=1.7.0=py39h06a4308_0
+ - notebook=7.0.6=pyhd8ed1ab_0
+ - notebook-shim=0.2.3=pyhd8ed1ab_0
+ - numexpr=2.8.0=mkl_py39h25e7801_2
+ - ocl-icd=2.3.1=h7f98852_0
+ - ocl-icd-system=1.0.0=1
+ - olefile=0.47=pyhd8ed1ab_0
+ - omegaconf=2.3.0=pyhd8ed1ab_0
+ - openff-amber-ff-ports=0.0.4=pyhca7485f_0
+ - openff-forcefields=2023.11.0=pyhca7485f_0
+ - openff-interchange-base=0.3.18=pyhd8ed1ab_0
+ - openff-models=0.1.1=pyhca7485f_0
+ - openff-toolkit-base=0.14.5=pyhd8ed1ab_1
+ - openff-units=0.2.1=pyh1a96a4e_0
+ - openff-utilities=0.1.12=pyhd8ed1ab_0
+ - openh264=2.1.1=h780b84a_0
+ - openjpeg=2.4.0=hb52868f_1
+ - openmm=8.1.0=py39he81762f_1
+ - openssl=3.2.1=hd590300_1
+ - ordered-set=4.1.0=py39h06a4308_0
+ - orjson=3.8.8=py39h52d8a92_0
+ - overrides=7.4.0=pyhd8ed1ab_0
+ - packaging=23.1=py39h06a4308_0
+ - packmol=20.010=h86c2bf4_0
+ - pandas=2.1.4=py39h1128e8f_0
+ - pandocfilters=1.5.0=pyhd8ed1ab_0
+ - panedr=0.8.0=pyhd8ed1ab_0
+ - parmed=3.4.4=py39h227be39_0
+ - parso=0.8.3=pyhd8ed1ab_0
+ - pcre=8.45=h295c915_0
+ - pcre2=10.43=hcad00b1_0
+ - pdbfixer=1.9=pyh1a96a4e_0
+ - perl=5.32.1=7_hd590300_perl5
+ - pexpect=4.8.0=pyhd3eb1b0_3
+ - pickleshare=0.7.5=py_1003
+ - pillow=8.2.0=py39hf95b381_1
+ - pint=0.21=pyhd8ed1ab_0
+ - pip=23.2.1=pyhd8ed1ab_0
+ - pixman=0.43.2=h59595ed_0
+ - pkginfo=1.9.6=py39h06a4308_0
+ - platformdirs=3.10.0=py39h06a4308_0
+ - pluggy=1.0.0=py39h06a4308_1
+ - poetry=1.6.1=linux_pyha804496_0
+ - poetry-core=1.7.0=pyhd8ed1ab_0
+ - poetry-plugin-export=1.5.0=pyhd8ed1ab_0
+ - portalocker=2.8.2=py39hf3d152e_1
+ - pre-commit=3.2.1=py39h06a4308_0
+ - prometheus_client=0.19.0=pyhd8ed1ab_0
+ - prompt-toolkit=3.0.43=py39h06a4308_0
+ - psutil=5.9.0=py39h5eee18b_0
+ - pthread-stubs=0.4=h36c2ea0_1001
+ - ptyprocess=0.7.0=pyhd3eb1b0_2
+ - pure_eval=0.2.2=pyhd8ed1ab_0
+ - py-cpuinfo=9.0.0=pyhd8ed1ab_0
+ - pycairo=1.25.0=py39hc92de75_1
+ - pycparser=2.21=pyhd3eb1b0_0
+ - pydantic=1.10.8=py39h5eee18b_0
+ - pyedr=0.8.0=pyhd8ed1ab_0
+ - pyfiglet=0.8.post1=py_0
+ - pyg=2.5.2=py39_torch_2.2.0_cu118
+ - pygments=2.15.1=py39h06a4308_1
+ - pyjwt=2.4.0=py39h06a4308_0
+ - pynvml=11.5.0=pyhd8ed1ab_0
+ - pyparsing=3.0.9=py39h06a4308_0
+ - pyproject_hooks=1.0.0=py39h06a4308_0
+ - pyrsistent=0.18.0=py39heee7806_0
+ - pysocks=1.7.1=pyha2e5f31_6
+ - pytables=3.8.0=py39hb8e3aad_2
+ - pytest=7.4.0=py39h06a4308_0
+ - python=3.9.17=h0755675_0_cpython
+ - python-build=0.10.0=py39h06a4308_0
+ - python-constraint=1.4.0=py_0
+ - python-dateutil=2.8.2=pyhd3eb1b0_0
+ - python-editor=1.0.4=pyhd3eb1b0_0
+ - python-fastjsonschema=2.19.1=pyhd8ed1ab_0
+ - python-installer=0.7.0=pyhd8ed1ab_0
+ - python-json-logger=2.0.7=pyhd8ed1ab_0
+ - python-multipart=0.0.6=pyhd8ed1ab_0
+ - python-tzdata=2023.4=pyhd8ed1ab_0
+ - python_abi=3.9=3_cp39
+ - pytorch=2.2.1=py3.9_cuda11.8_cudnn8.7.0_0
+ - pytorch-cuda=11.8=h7e8668a_5
+ - pytorch-mutex=1.0=cuda
+ - pytorch-scatter=2.1.2=py39_torch_2.2.0_cu118
+ - pytz=2023.3.post1=py39h06a4308_0
+ - pyyaml=6.0.1=py39hd1e30aa_0
+ - pyzmq=25.1.2=py39h6a678d5_0
+ - rapidfuzz=2.13.7=py39h417a72b_0
+ - rdkit=2024.03.1=py39h6cc1c65_0
+ - readchar=4.0.5=pyhd8ed1ab_0
+ - readline=8.2=h8228510_1
+ - reportlab=3.5.68=py39he59360d_1
+ - requests=2.31.0=pyhd8ed1ab_0
+ - requests-toolbelt=1.0.0=py39h06a4308_0
+ - rfc3339-validator=0.1.4=pyhd8ed1ab_0
+ - rfc3986-validator=0.1.1=pyh9f0ad1d_0
+ - rich=13.3.5=py39h06a4308_0
+ - rocm-smi=5.6.0=h59595ed_1
+ - scikit-learn=1.4.1.post1=py39ha22ef79_0
+ - scipy=1.12.0=py39h474f0d3_2
+ - secretstorage=3.3.1=py39h06a4308_1
+ - send2trash=1.8.2=pyh41d4057_0
+ - setuptools=63.4.1=py39h06a4308_0
+ - shellingham=1.5.0=py39h06a4308_0
+ - six=1.16.0=pyhd3eb1b0_1
+ - smirnoff99frosst=1.1.0=pyh44b312d_0
+ - snappy=1.1.10=h9fff704_0
+ - sniffio=1.2.0=py39h06a4308_1
+ - soupsieve=2.4=py39h06a4308_0
+ - sqlalchemy=2.0.25=py39h5eee18b_0
+ - stack_data=0.6.2=pyhd8ed1ab_0
+ - starlette=0.27.0=pyhd8ed1ab_0
+ - starsessions=1.3.0=pyhd8ed1ab_0
+ - sympy=1.12=pypyh9d50eac_103
+ - tabulate=0.9.0=pyhd8ed1ab_1
+ - tbb=2021.7.0=h924138e_0
+ - termcolor=2.4.0=pyhd8ed1ab_0
+ - terminado=0.18.0=pyh0d859eb_0
+ - threadpoolctl=3.4.0=pyhc1e730c_0
+ - tinycss2=1.2.1=pyhd8ed1ab_0
+ - tinydb=4.8.0=pyhd8ed1ab_0
+ - tk=8.6.12=h27826a3_0
+ - tomli=2.0.1=py39h06a4308_0
+ - tomlkit=0.12.1=pyha770c72_0
+ - torchaudio=2.2.1=py39_cu118
+ - torchtriton=2.2.0=py39
+ - torchvision=0.17.1=py39_cu118
+ - tornado=6.3.3=py39h5eee18b_0
+ - tqdm=4.66.1=pyhd8ed1ab_0
+ - traitlets=5.7.1=py39h06a4308_0
+ - trove-classifiers=2023.3.9=py39h06a4308_0
+ - typing-extensions=4.8.0=hd8ed1ab_0
+ - typing_extensions=4.8.0=pyha770c72_0
+ - typing_utils=0.1.0=pyhd8ed1ab_0
+ - tzdata=2023c=h71feb2d_0
+ - ukkonen=1.0.1=py39hdb19cb5_0
+ - unicodedata2=15.1.0=py39hd1e30aa_0
+ - uri-template=1.3.0=pyhd8ed1ab_0
+ - urllib3=2.1.0=py39h06a4308_1
+ - uvicorn=0.20.0=py39h06a4308_0
+ - validators=0.22.0=pyhd8ed1ab_0
+ - virtualenv=20.24.3=pyhd8ed1ab_0
+ - wcwidth=0.2.5=pyhd3eb1b0_0
+ - webcolors=1.13=pyhd8ed1ab_0
+ - webencodings=0.5.1=pyhd8ed1ab_2
+ - websocket-client=0.58.0=py39h06a4308_4
+ - websockets=10.4=py39h5eee18b_1
+ - wheel=0.41.1=pyhd8ed1ab_0
+ - widgetsnbextension=4.0.9=pyhd8ed1ab_0
+ - xmltodict=0.13.0=pyhd8ed1ab_0
+ - xorg-kbproto=1.0.7=h7f98852_1002
+ - xorg-libice=1.1.1=hd590300_0
+ - xorg-libsm=1.2.4=h7391055_0
+ - xorg-libx11=1.8.7=h8ee46fc_0
+ - xorg-libxau=1.0.11=hd590300_0
+ - xorg-libxdmcp=1.1.3=h7f98852_0
+ - xorg-libxext=1.3.4=h0b41bf4_2
+ - xorg-libxrender=0.9.11=hd590300_0
+ - xorg-renderproto=0.11.1=h7f98852_1002
+ - xorg-xextproto=7.3.0=h0b41bf4_1003
+ - xorg-xproto=7.0.31=h7f98852_1007
+ - xz=5.2.6=h166bdaf_0
+ - yacs=0.1.8=pyhd8ed1ab_0
+ - yaml=0.2.5=h7b6447c_0
+ - yarl=1.9.4=py39hd1e30aa_0
+ - zeromq=4.3.5=h59595ed_0
+ - zipp=3.11.0=py39h06a4308_0
+ - zlib=1.2.13=hd590300_5
+ - zlib-ng=2.0.7=h0b41bf4_0
+ - zstd=1.5.5=hfc55251_0
+ - pip:
+ - absl-py==2.1.0
+ - appdirs==1.4.4
+ - asttokens==2.4.1
+ - beartype==0.17.2
+ - biopandas==0.4.1
+ - biopython>=1.79
+ - colorlog==6.8.2
+ - contextlib2==21.6.0
+ - deprecated==1.2.14
+ - git+https://github.com/NVIDIA/dllogger.git@0540a43971f4a8a16693a9de9de73c1072020769
+ - dm-tree==0.1.8
+ - docker-pycreds==0.4.0
+ - einops==0.7.0
+ - executing==2.0.1
+ - fair-esm==2.0.0
+ - fairscale==0.4.13
+ - fastcore==1.5.29
+ - fsspec==2024.2.0
+ - gitdb==4.0.11
+ - gitpython==3.1.42
+ - hydra-colorlog==1.2.0
+ - lightning==2.2.2
+ - lightning-utilities==0.11.1
+ - lovely-numpy==0.2.11
+ - lovely-tensors==0.1.15
+ - ml-collections==0.1.1
+ - msgpack-numpy==0.4.8
+ - numpy==1.23.5
+ - git+https://github.com/amorehead/openfold.git@fe1275099639bf7e617e09ef24d6af778647dd64
+ - posebusters==0.2.13
+ - git+https://git@github.com/zrqiao/power_spherical.git@290b1630c5f84e3bb0d61711046edcf6e47200d4
+ - prody==2.4.1
+ - protobuf==4.25.3
+ - python-dotenv==1.0.1
+ - pytorch-lightning==2.2.2
+ - git+https://github.com/facebookresearch/pytorch3d.git@3da7703c5ac10039645966deddffe8db52eab8c5
+ - rootutils==1.0.7
+ - seaborn==0.13.2
+ - sentry-sdk==1.40.5
+ - setproctitle==1.3.3
+ - smmap==5.0.1
+ - torchmetrics==1.3.1
+ - wandb==0.17.0
+ - wrapt==1.16.0
diff --git a/forks/FlowDock/flowdock/__init__.py b/forks/FlowDock/flowdock/__init__.py
new file mode 100644
index 00000000..26bca07a
--- /dev/null
+++ b/forks/FlowDock/flowdock/__init__.py
@@ -0,0 +1,120 @@
+import importlib
+import os
+
+from beartype.typing import Any
+from omegaconf import OmegaConf
+
+METHOD_TITLE_MAPPING = {
+ "diffdock": "DiffDock",
+ "flowdock": "FlowDock",
+ "neuralplexer": "NeuralPLexer",
+}
+
+STANDARDIZED_DIR_METHODS = ["diffdock"]
+
+
+def resolve_omegaconf_variable(variable_path: str) -> Any:
+ """Resolve an OmegaConf variable path to its value."""
+ # split the string into parts using the dot separator
+ parts = variable_path.rsplit(".", 1)
+
+ # get the module name from the first part of the path
+ module_name = parts[0]
+
+ # dynamically import the module using the module name
+ try:
+ module = importlib.import_module(module_name)
+ # use the imported module to get the requested attribute value
+ attribute = getattr(module, parts[1])
+ except Exception:
+ module = importlib.import_module(".".join(module_name.split(".")[:-1]))
+ inner_module = ".".join(module_name.split(".")[-1:])
+ # use the imported module to get the requested attribute value
+ attribute = getattr(getattr(module, inner_module), parts[1])
+
+ return attribute
+
+
+def resolve_dataset_path_dirname(dataset: str) -> str:
+ """Resolve the dataset path directory name based on the dataset's name.
+
+ :param dataset: Name of the dataset.
+ :return: Directory name for the dataset path.
+ """
+ return "DockGen" if dataset == "dockgen" else dataset
+
+
+def resolve_method_input_csv_path(method: str, dataset: str) -> str:
+ """Resolve the input CSV path for a given method.
+
+ :param method: The method name.
+ :param dataset: The dataset name.
+ :return: The input CSV path for the given method.
+ """
+ if method in STANDARDIZED_DIR_METHODS or method in ["flowdock", "neuralplexer"]:
+ return os.path.join(
+ "forks",
+ METHOD_TITLE_MAPPING.get(method, method),
+ "inference",
+ f"{method}_{dataset}_inputs.csv",
+ )
+ else:
+ raise ValueError(f"Invalid method: {method}")
+
+
+def resolve_method_title(method: str) -> str:
+ """Resolve the method title for a given method.
+
+ :param method: The method name.
+ :return: The method title for the given method.
+ """
+ return METHOD_TITLE_MAPPING.get(method, method)
+
+
+def resolve_method_output_dir(
+ method: str,
+ dataset: str,
+ repeat_index: int,
+) -> str:
+ """Resolve the output directory for a given method.
+
+ :param method: The method name.
+ :param dataset: The dataset name.
+ :param repeat_index: The repeat index for the method.
+ :return: The output directory for the given method.
+ """
+ if method in STANDARDIZED_DIR_METHODS or method in ["flowdock", "neuralplexer"]:
+ return os.path.join(
+ "forks",
+ METHOD_TITLE_MAPPING.get(method, method),
+ "inference",
+ f"{method}_{dataset}_output{'s' if method in ['flowdock', 'neuralplexer'] else ''}_{repeat_index}",
+ )
+ else:
+ raise ValueError(f"Invalid method: {method}")
+
+
+def register_custom_omegaconf_resolvers():
+ """Register custom OmegaConf resolvers."""
+ OmegaConf.register_new_resolver(
+ "resolve_variable", lambda variable_path: resolve_omegaconf_variable(variable_path)
+ )
+ OmegaConf.register_new_resolver(
+ "resolve_dataset_path_dirname", lambda dataset: resolve_dataset_path_dirname(dataset)
+ )
+ OmegaConf.register_new_resolver(
+ "resolve_method_input_csv_path",
+ lambda method, dataset: resolve_method_input_csv_path(method, dataset),
+ )
+ OmegaConf.register_new_resolver(
+ "resolve_method_title", lambda method: resolve_method_title(method)
+ )
+ OmegaConf.register_new_resolver(
+ "resolve_method_output_dir",
+ lambda method, dataset, repeat_index: resolve_method_output_dir(
+ method, dataset, repeat_index
+ ),
+ )
+ OmegaConf.register_new_resolver(
+ "int_divide", lambda dividend, divisor: int(dividend) // int(divisor)
+ )
diff --git a/forks/FlowDock/flowdock/data/__init__.py b/forks/FlowDock/flowdock/data/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/forks/FlowDock/flowdock/data/combined_datamodule.py b/forks/FlowDock/flowdock/data/combined_datamodule.py
new file mode 100644
index 00000000..a143aade
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/combined_datamodule.py
@@ -0,0 +1,629 @@
+import numpy as np
+import rootutils
+import torch
+from beartype.typing import Any, Dict, List, Literal, Optional
+from lightning import LightningDataModule
+from torch.utils.data import DataLoader, Dataset
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock.data.components.combined_dataset import CombinedDataset
+from flowdock.data.components.moad_dataset import BindingMOADDataset
+from flowdock.data.components.mol_features import collate_samples
+from flowdock.data.components.pdbbind_dataset import PDBBindDataset
+from flowdock.data.components.pdbsidechain_dataset import PDBSidechainDataset
+from flowdock.utils import RankedLogger
+from flowdock.utils.data_utils import (
+ parse_moad_binding_affinity_data_file,
+ parse_pdbbind_binding_affinity_data_file,
+)
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+DATA_PHASE = Literal["train", "val", "test"]
+VALID_TRAIN_DATASETS = ["pdbbind", "moad", "pdbsidechain"]
+VALID_TEST_DATASETS = ["pdbbind", "moad"]
+
+
+class CombinedDataModule(LightningDataModule):
+ """`LightningDataModule` for combined training and test datasets.
+
+ A `LightningDataModule` implements 7 key methods:
+
+ ```python
+ def prepare_data(self):
+ # Things to do on 1 GPU/TPU (not on every GPU/TPU in DDP).
+ # Download data, pre-process, split, save to disk, etc...
+
+ def setup(self, stage):
+ # Things to do on every process in DDP.
+ # Load data, set variables, etc...
+
+ def train_dataloader(self):
+ # return train dataloader
+
+ def val_dataloader(self):
+ # return validation dataloader
+
+ def test_dataloader(self):
+ # return test dataloader
+
+ def predict_dataloader(self):
+ # return predict dataloader
+
+ def teardown(self, stage):
+ # Called on every process in DDP.
+ # Clean up after fit or test.
+ ```
+
+ This allows you to share a full dataset without explaining how to download,
+ split, transform and process the data.
+
+ Read the docs:
+ https://lightning.ai/docs/pytorch/latest/data/datamodule.html
+ """
+
+ def __init__(
+ self,
+ data_dir: str = "data/",
+ train_datasets: List[str] = ["pdbbind", "moad", "pdbsidechain"],
+ test_datasets: List[str] = ["pdbbind", "moad"],
+ batch_size: int = 16,
+ num_workers: int = 0,
+ pin_memory: bool = False,
+ stage: Optional[str] = None,
+ **kwargs: Any,
+ ) -> None:
+ """Initialize a `CombinedDataModule`.
+
+ :param data_dir: The data directory. Defaults to `"data/"`.
+ :param train_datasets: The datasets to use for training. Defaults to `["pdbbind", "moad", "pdbsidechain"]`.
+ :param test_datasets: The datasets to use for testing. Defaults to `["pdbbind", "moad"]`.
+ :param batch_size: The batch size. Defaults to `16`.
+ :param num_workers: The number of workers. Defaults to `0`.
+ :param pin_memory: Whether to pin memory. Defaults to `False`.
+ :param stage: The stage to setup. Either `"fit"`, `"validate"`, `"test"`, or `"predict"`. Defaults to ``None``.
+ """
+ super().__init__()
+ load_train_val_datasets = stage is None or (
+ stage is not None and stage in ["fit", "validate"]
+ )
+
+ # this line allows to access init params with 'self.hparams' attribute
+ # also ensures init params will be stored in ckpt
+ self.save_hyperparameters(logger=False)
+
+ # parse binding affinity values
+ self.pdbbind_binding_affinity_values_dict = (
+ parse_pdbbind_binding_affinity_data_file(
+ self.hparams.pdbbind_binding_affinity_values_path
+ )
+ if "pdbbind" in train_datasets
+ else None
+ )
+ self.moad_binding_affinity_values_dict = (
+ parse_moad_binding_affinity_data_file(self.hparams.moad_binding_affinity_values_path)
+ if "moad" in train_datasets
+ and self.hparams.moad_map_binding_affinities_to_superligands
+ else None
+ )
+
+ # prepare for dataset(s) to be loaded into rank-shared memory (e.g., when setting `trainer.strategy = "ddp_spawn"`)
+ invalid_trainsets = [
+ dataset
+ for dataset in self.hparams.train_datasets
+ if dataset not in VALID_TRAIN_DATASETS
+ ]
+ if invalid_trainsets:
+ raise ValueError(f"Invalid dataset(s) specified for training: {invalid_trainsets}")
+ if any(dataset not in VALID_TEST_DATASETS for dataset in self.hparams.test_datasets):
+ raise ValueError(
+ f"Invalid dataset(s) specified for testing: {self.hparams.test_datasets}"
+ )
+
+ if load_train_val_datasets and "pdbbind" in self.hparams.train_datasets:
+ pdbbind_train = PDBBindDataset(
+ root=self.hparams.pdbbind_dir,
+ cache_path=self.hparams.cache_path,
+ split_path=self.hparams.split_train,
+ esm_embeddings_path=self.hparams.pdbbind_esm_embeddings_path,
+ apo_protein_structure_dir=self.hparams.pdbbind_apo_protein_structure_dir,
+ protein_file=self.hparams.protein_file,
+ require_ligand=self.hparams.overfitting_example_name is not None,
+ limit_complexes=self.hparams.limit_complexes,
+ remove_hs=self.hparams.remove_hs,
+ max_lig_size=self.hparams.max_lig_size,
+ num_workers=self.hparams.num_workers,
+ include_miscellaneous_atoms=self.hparams.include_miscellaneous_atoms,
+ min_protein_length=self.hparams.min_protein_length,
+ max_protein_length=self.hparams.max_protein_length,
+ a2h_assessment_csv_filepath=self.hparams.pdbbind_a2h_assessment_csv_filepath,
+ filter_using_a2h_assessment=self.hparams.pdbbind_filter_using_a2h_assessment,
+ a2h_min_tmscore=self.hparams.pdbbind_a2h_min_tmscore,
+ a2h_max_rmsd=self.hparams.pdbbind_a2h_max_rmsd,
+ a2h_min_protein_length=self.hparams.pdbbind_a2h_min_protein_length,
+ a2h_max_protein_length=self.hparams.pdbbind_a2h_max_protein_length,
+ a2h_min_ligand_length=self.hparams.pdbbind_a2h_min_ligand_length,
+ a2h_max_ligand_length=self.hparams.pdbbind_a2h_max_ligand_length,
+ binding_affinity_values_dict=self.pdbbind_binding_affinity_values_dict,
+ n_lig_patches=self.hparams.n_lig_patches,
+ )
+ pdbbind_val = PDBBindDataset(
+ root=self.hparams.pdbbind_dir,
+ cache_path=self.hparams.cache_path,
+ split_path=self.hparams.split_val,
+ esm_embeddings_path=self.hparams.pdbbind_esm_embeddings_path,
+ apo_protein_structure_dir=self.hparams.pdbbind_apo_protein_structure_dir,
+ protein_file=self.hparams.protein_file,
+ require_ligand=True,
+ limit_complexes=self.hparams.limit_complexes,
+ remove_hs=self.hparams.remove_hs,
+ max_lig_size=self.hparams.max_lig_size,
+ num_workers=self.hparams.num_workers,
+ include_miscellaneous_atoms=self.hparams.include_miscellaneous_atoms,
+ min_protein_length=self.hparams.min_protein_length,
+ max_protein_length=self.hparams.max_protein_length,
+ a2h_assessment_csv_filepath=self.hparams.pdbbind_a2h_assessment_csv_filepath,
+ filter_using_a2h_assessment=self.hparams.pdbbind_filter_using_a2h_assessment,
+ a2h_min_tmscore=self.hparams.pdbbind_a2h_min_tmscore,
+ a2h_max_rmsd=self.hparams.pdbbind_a2h_max_rmsd,
+ a2h_min_protein_length=self.hparams.pdbbind_a2h_min_protein_length,
+ a2h_max_protein_length=self.hparams.pdbbind_a2h_max_protein_length,
+ a2h_min_ligand_length=self.hparams.pdbbind_a2h_min_ligand_length,
+ a2h_max_ligand_length=self.hparams.pdbbind_a2h_max_ligand_length,
+ binding_affinity_values_dict=self.pdbbind_binding_affinity_values_dict,
+ n_lig_patches=self.hparams.n_lig_patches,
+ )
+ if load_train_val_datasets and "moad" in self.hparams.train_datasets:
+ moad_train = BindingMOADDataset(
+ root=self.hparams.moad_dir,
+ dockgen_dir=self.hparams.moad_dockgen_dir,
+ clusters_filepath=self.hparams.moad_clusters_filepath,
+ cache_path=self.hparams.cache_path,
+ split="train",
+ max_receptor_size=self.hparams.max_receptor_size,
+ remove_promiscuous_targets=self.hparams.remove_promiscuous_targets,
+ min_ligand_size=self.hparams.min_ligand_size,
+ multiplicity=self.hparams.train_multiplicity,
+ unroll_clusters=self.hparams.unroll_clusters,
+ esm_embeddings_sequences_path=self.hparams.moad_esm_embeddings_sequences_path,
+ esm_embeddings_path=self.hparams.moad_esm_embeddings_path,
+ apo_protein_structure_dir=self.hparams.moad_apo_protein_structure_dir,
+ require_ligand=self.hparams.overfitting_example_name is not None,
+ enforce_timesplit=self.hparams.enforce_timesplit,
+ limit_complexes=self.hparams.limit_complexes,
+ remove_hs=self.hparams.remove_hs,
+ max_lig_size=self.hparams.max_lig_size,
+ min_multi_lig_distance=self.hparams.min_multi_lig_distance,
+ include_miscellaneous_atoms=self.hparams.include_miscellaneous_atoms,
+ pdbbind_split_train=self.hparams.pdbbind_split_train,
+ pdbbind_split_val=self.hparams.pdbbind_split_val,
+ remove_pdbbind=self.hparams.remove_pdbbind,
+ split_time=self.hparams.split_time,
+ min_protein_length=self.hparams.min_protein_length,
+ max_protein_length=self.hparams.max_protein_length,
+ a2h_assessment_csv_filepath=self.hparams.moad_a2h_assessment_csv_filepath,
+ filter_using_a2h_assessment=self.hparams.moad_filter_using_a2h_assessment,
+ a2h_min_tmscore=self.hparams.moad_a2h_min_tmscore,
+ a2h_max_rmsd=self.hparams.moad_a2h_max_rmsd,
+ a2h_min_protein_length=self.hparams.moad_a2h_min_protein_length,
+ a2h_max_protein_length=self.hparams.moad_a2h_max_protein_length,
+ a2h_min_ligand_length=self.hparams.moad_a2h_min_ligand_length,
+ a2h_max_ligand_length=self.hparams.moad_a2h_max_ligand_length,
+ binding_affinity_values_dict=self.moad_binding_affinity_values_dict,
+ n_lig_patches=self.hparams.n_lig_patches,
+ )
+ moad_val = BindingMOADDataset(
+ root=self.hparams.moad_dir,
+ dockgen_dir=self.hparams.moad_dockgen_dir,
+ clusters_filepath=self.hparams.moad_clusters_filepath,
+ cache_path=self.hparams.cache_path,
+ split="val",
+ multiplicity=self.hparams.val_multiplicity,
+ remove_promiscuous_targets=self.hparams.remove_promiscuous_targets,
+ dockgen_esm_embeddings_path=self.hparams.moad_dockgen_esm_embeddings_path,
+ dockgen_esm_embeddings_sequences_path=self.hparams.moad_dockgen_esm_embeddings_sequences_path,
+ dockgen_apo_protein_structure_dir=self.hparams.moad_dockgen_apo_protein_structure_dir,
+ unroll_clusters=self.hparams.unroll_clusters,
+ require_ligand=True,
+ limit_complexes=self.hparams.limit_complexes,
+ remove_hs=self.hparams.remove_hs,
+ num_workers=self.hparams.num_workers,
+ include_miscellaneous_atoms=self.hparams.include_miscellaneous_atoms,
+ pdbbind_split_train=self.hparams.pdbbind_split_train,
+ pdbbind_split_val=self.hparams.pdbbind_split_val,
+ remove_pdbbind=self.hparams.remove_pdbbind,
+ split_time=self.hparams.split_time,
+ dockgen_a2h_assessment_csv_filepath=self.hparams.moad_dockgen_a2h_assessment_csv_filepath,
+ filter_using_a2h_assessment=self.hparams.moad_filter_using_a2h_assessment,
+ # NOTE: the following `a2h` parameters were deduced directly from the DockGen dataset's `a2h` assessment plots
+ a2h_min_tmscore=0.6,
+ a2h_max_rmsd=6.0,
+ a2h_min_protein_length=50,
+ a2h_max_protein_length=2000,
+ a2h_min_ligand_length=1,
+ a2h_max_ligand_length=300,
+ binding_affinity_values_dict=self.moad_binding_affinity_values_dict,
+ n_lig_patches=self.hparams.n_lig_patches,
+ )
+ if load_train_val_datasets and "pdbsidechain" in self.hparams.train_datasets:
+ pdbsidechain_train = PDBSidechainDataset(
+ root=self.hparams.pdbsidechain_dir,
+ cache_path=self.hparams.cache_path,
+ split="train",
+ esm_embeddings_path=self.hparams.pdbsidechain_esm_embeddings_path,
+ esm_embeddings_sequences_path=self.hparams.pdbsidechain_esm_embeddings_sequences_path,
+ apo_protein_structure_dir=self.hparams.pdbsidechain_apo_protein_structure_dir,
+ pdbsidechain_metadata_dir=self.hparams.pdbsidechain_metadata_dir,
+ multiplicity=self.hparams.train_multiplicity,
+ limit_complexes=self.hparams.limit_complexes,
+ remove_hs=self.hparams.remove_hs,
+ num_workers=self.hparams.num_workers,
+ vandermers_max_dist=self.hparams.vandermers_max_dist,
+ vandermers_buffer_residue_num=self.hparams.vandermers_buffer_residue_num,
+ vandermers_min_contacts=self.hparams.vandermers_min_contacts,
+ vandermers_max_surrogate_binding_affinity=self.hparams.vandermers_max_surrogate_binding_affinity,
+ vandermers_second_ligand_max_closeness=self.hparams.vandermers_second_ligand_max_closeness,
+ vandermers_extract_second_ligand=self.hparams.vandermers_extract_second_ligand,
+ merge_clusters=self.hparams.merge_clusters,
+ min_protein_length=self.hparams.min_protein_length,
+ max_protein_length=self.hparams.max_protein_length,
+ vandermers_use_prob_as_surrogate_binding_affinity=self.hparams.vandermers_use_prob_as_surrogate_binding_affinity,
+ a2h_assessment_csv_filepath=self.hparams.pdbsidechain_a2h_assessment_csv_filepath,
+ filter_using_a2h_assessment=self.hparams.pdbsidechain_filter_using_a2h_assessment,
+ a2h_min_tmscore=self.hparams.pdbsidechain_a2h_min_tmscore,
+ a2h_max_rmsd=self.hparams.pdbsidechain_a2h_max_rmsd,
+ a2h_min_protein_length=self.hparams.pdbsidechain_a2h_min_protein_length,
+ a2h_max_protein_length=self.hparams.pdbsidechain_a2h_max_protein_length,
+ postprocess_min_protein_length=self.hparams.pdbsidechain_postprocess_min_protein_length,
+ postprocess_max_protein_length=self.hparams.pdbsidechain_postprocess_max_protein_length,
+ n_lig_patches=self.hparams.n_lig_patches,
+ )
+ pdbsidechain_val = PDBSidechainDataset(
+ root=self.hparams.pdbsidechain_dir,
+ cache_path=self.hparams.cache_path,
+ split="val",
+ esm_embeddings_path=self.hparams.pdbsidechain_esm_embeddings_path,
+ esm_embeddings_sequences_path=self.hparams.pdbsidechain_esm_embeddings_sequences_path,
+ apo_protein_structure_dir=self.hparams.pdbsidechain_apo_protein_structure_dir,
+ pdbsidechain_metadata_dir=self.hparams.pdbsidechain_metadata_dir,
+ multiplicity=self.hparams.val_multiplicity,
+ limit_complexes=self.hparams.limit_complexes,
+ remove_hs=self.hparams.remove_hs,
+ num_workers=self.hparams.num_workers,
+ vandermers_max_dist=self.hparams.vandermers_max_dist,
+ vandermers_buffer_residue_num=self.hparams.vandermers_buffer_residue_num,
+ vandermers_min_contacts=self.hparams.vandermers_min_contacts,
+ vandermers_max_surrogate_binding_affinity=self.hparams.vandermers_max_surrogate_binding_affinity,
+ vandermers_second_ligand_max_closeness=self.hparams.vandermers_second_ligand_max_closeness,
+ vandermers_extract_second_ligand=self.hparams.vandermers_extract_second_ligand,
+ merge_clusters=self.hparams.merge_clusters,
+ min_protein_length=self.hparams.min_protein_length,
+ max_protein_length=self.hparams.max_protein_length,
+ vandermers_use_prob_as_surrogate_binding_affinity=self.hparams.vandermers_use_prob_as_surrogate_binding_affinity,
+ a2h_assessment_csv_filepath=self.hparams.pdbsidechain_a2h_assessment_csv_filepath,
+ filter_using_a2h_assessment=self.hparams.pdbsidechain_filter_using_a2h_assessment,
+ a2h_min_tmscore=self.hparams.pdbsidechain_a2h_min_tmscore,
+ a2h_max_rmsd=self.hparams.pdbsidechain_a2h_max_rmsd,
+ a2h_min_protein_length=self.hparams.pdbsidechain_a2h_min_protein_length,
+ a2h_max_protein_length=self.hparams.pdbsidechain_a2h_max_protein_length,
+ postprocess_min_protein_length=self.hparams.pdbsidechain_postprocess_min_protein_length,
+ postprocess_max_protein_length=self.hparams.pdbsidechain_postprocess_max_protein_length,
+ n_lig_patches=self.hparams.n_lig_patches,
+ )
+ if "pdbbind" in self.hparams.test_datasets:
+ pdbbind_test = PDBBindDataset(
+ root=self.hparams.pdbbind_dir,
+ limit_complexes=self.hparams.limit_complexes,
+ cache_path=self.hparams.cache_path,
+ split_path=self.hparams.split_test,
+ remove_hs=self.hparams.remove_hs,
+ max_lig_size=None,
+ esm_embeddings_path=self.hparams.pdbbind_esm_embeddings_path,
+ apo_protein_structure_dir=self.hparams.pdbbind_apo_protein_structure_dir,
+ require_ligand=True,
+ num_workers=self.hparams.num_workers,
+ protein_file=self.hparams.protein_file,
+ ligand_file=self.hparams.ligand_file,
+ include_miscellaneous_atoms=self.hparams.include_miscellaneous_atoms,
+ min_protein_length=self.hparams.min_protein_length,
+ max_protein_length=self.hparams.max_protein_length,
+ is_test_dataset=True,
+ binding_affinity_values_dict=self.pdbbind_binding_affinity_values_dict,
+ n_lig_patches=self.hparams.n_lig_patches,
+ )
+ if "moad" in self.hparams.test_datasets:
+ moad_test = BindingMOADDataset(
+ root=self.hparams.moad_dir,
+ dockgen_dir=self.hparams.moad_dockgen_dir,
+ clusters_filepath=self.hparams.moad_clusters_filepath,
+ limit_complexes=self.hparams.limit_complexes,
+ cache_path=self.hparams.cache_path,
+ split="test",
+ remove_hs=self.hparams.remove_hs,
+ dockgen_esm_embeddings_path=self.hparams.moad_dockgen_esm_embeddings_path,
+ dockgen_esm_embeddings_sequences_path=self.hparams.moad_dockgen_esm_embeddings_sequences_path,
+ dockgen_apo_protein_structure_dir=self.hparams.moad_dockgen_apo_protein_structure_dir,
+ require_ligand=True,
+ num_workers=self.hparams.num_workers,
+ include_miscellaneous_atoms=self.hparams.include_miscellaneous_atoms,
+ pdbbind_split_train=self.hparams.pdbbind_split_train,
+ pdbbind_split_val=self.hparams.pdbbind_split_val,
+ split_time=self.hparams.split_time,
+ unroll_clusters=True,
+ remove_pdbbind=self.hparams.remove_pdbbind,
+ remove_promiscuous_targets=self.hparams.remove_promiscuous_targets,
+ no_randomness=True,
+ dockgen_a2h_assessment_csv_filepath=self.hparams.moad_dockgen_a2h_assessment_csv_filepath,
+ is_test_dataset=True,
+ binding_affinity_values_dict=self.moad_binding_affinity_values_dict,
+ n_lig_patches=self.hparams.n_lig_patches,
+ )
+ if not ("pdbbind" in self.hparams.test_datasets or "moad" in self.hparams.test_datasets):
+ raise ValueError(f"Invalid test dataset(s) specified: {self.hparams.test_datasets}")
+
+ trainset, valset, valset_secondary = None, None, None
+ if load_train_val_datasets and self.hparams.overfitting_example_name is not None:
+ log.info(
+ "Overfitting on a single batch. Establishing a fixed (uncombined) training set."
+ )
+ if "moad" in self.hparams.train_datasets:
+ trainset = moad_train
+ valset = moad_train
+ log.info("Overfitting using Binding MOAD training set.")
+ elif "pdbbind" in self.hparams.train_datasets:
+ trainset = pdbbind_train
+ valset = pdbbind_train
+ log.info("Overfitting using PDBBind training set.")
+ elif "pdbsidechain" in self.hparams.train_datasets:
+ trainset = pdbsidechain_train
+ valset = pdbsidechain_train
+ log.info("Overfitting using van der Mers training set.")
+
+ elif load_train_val_datasets and len(self.hparams.train_datasets) == 3:
+ trainset1 = pdbbind_train
+ trainset2 = moad_train
+ trainset3 = pdbsidechain_train
+ trainset = CombinedDataset(CombinedDataset(trainset2, trainset1), trainset3)
+ valset = moad_val
+ if self.hparams.double_val:
+ valset_secondary = pdbbind_val
+
+ elif load_train_val_datasets and len(self.hparams.train_datasets) == 2:
+ if "pdbbind" in self.hparams.train_datasets and "moad" in self.hparams.train_datasets:
+ trainset = CombinedDataset(
+ moad_train,
+ pdbbind_train,
+ )
+ valset = moad_val
+ if self.hparams.double_val:
+ valset_secondary = pdbbind_val
+
+ elif (
+ load_train_val_datasets
+ and "pdbbind" in self.hparams.train_datasets
+ and "pdbsidechain" in self.hparams.train_datasets
+ ):
+ trainset = CombinedDataset(
+ pdbbind_train,
+ pdbsidechain_train,
+ )
+ valset = pdbbind_val
+ if self.hparams.double_val:
+ valset_secondary = pdbsidechain_val
+
+ elif (
+ load_train_val_datasets
+ and "moad" in self.hparams.train_datasets
+ and "pdbsidechain" in self.hparams.train_datasets
+ ):
+ trainset = CombinedDataset(
+ moad_train,
+ pdbsidechain_train,
+ )
+ valset = moad_val
+ if self.hparams.double_val:
+ valset_secondary = pdbsidechain_val
+
+ elif load_train_val_datasets and len(self.hparams.train_datasets) == 1:
+ if "pdbbind" in self.hparams.train_datasets:
+ trainset = pdbbind_train
+ valset = pdbbind_val
+ elif "moad" in self.hparams.train_datasets:
+ trainset = moad_train
+ valset = moad_val
+ elif "pdbsidechain" in self.hparams.train_datasets:
+ trainset = pdbsidechain_train
+ valset = pdbsidechain_val
+
+ elif load_train_val_datasets:
+ raise ValueError("No datasets specified for training.")
+
+ testset, testset_secondary = None, None
+ if "pdbbind" in self.hparams.test_datasets:
+ testset = pdbbind_test
+ if "moad" in self.hparams.test_datasets:
+ testset_secondary = moad_test
+
+ (
+ self.data_train,
+ self.data_val,
+ self.data_val_secondary,
+ self.data_test,
+ self.data_test_secondary,
+ ) = (
+ trainset,
+ valset,
+ valset_secondary,
+ testset,
+ testset_secondary,
+ )
+
+ def prepare_data(self) -> None:
+ """Download data if needed. Lightning ensures that `self.prepare_data()` is called only
+ within a single process on CPU, so you can safely add your downloading logic within. In
+ case of multi-node training, the execution of this hook depends upon
+ `self.prepare_data_per_node()`.
+
+ Do not use it to assign state (self.x = y).
+ """
+ pass
+
+ def setup(self, stage: Optional[str] = None) -> None:
+ """Load data. Set variables: `self.data_train`, `self.data_val`, `self.data_test`.
+
+ This method is called by Lightning before `trainer.fit()`, `trainer.validate()`, `trainer.test()`, and
+ `trainer.predict()`, so be careful not to execute things like random split twice! Also, it is called after
+ `self.prepare_data()` and there is a barrier in between which ensures that all the processes proceed to
+ `self.setup()` once the data is prepared and available for use.
+
+ :param stage: The stage to setup. Either `"fit"`, `"validate"`, `"test"`, or `"predict"`. Defaults to ``None``.
+ """
+ pass
+
+ @staticmethod
+ def dynamic_batching_by_max_edge_count(
+ x: Dict[str, Any], max_n_edges: int, max_batch_size: int
+ ) -> Any:
+ """Dynamically batch by maximum edge count.
+
+ :param x: The input graph data. If `None`, skip the current batch.
+ :param max_n_edges: The maximum number of edges.
+ :param max_batch_size: The maximum batch size.
+ :return: The batched data.
+ """
+ if "num_u" in x["metadata"].keys():
+ num_edges_upperbound = (
+ x["metadata"]["num_a"] * 128 + x["metadata"]["num_i"] * 8 + 160**2
+ )
+ else:
+ num_edges_upperbound = x["metadata"]["num_a"] * 128 + 160**2
+ batch_size = max(1, min(max_n_edges // num_edges_upperbound, max_batch_size))
+ return collate_samples([x] * batch_size)
+
+ def get_dataloader(
+ self,
+ phase: DATA_PHASE,
+ dataset: Dataset,
+ **kwargs: Dict[str, Any],
+ ) -> DataLoader[Any]:
+ """Create a dataloader from a dataset.
+
+ :param phase: The phase of the dataset. Either `"train"`, `"val"`, or `"test"`.
+ :param dataset: The dataset.
+ :param kwargs: Additional keyword arguments to pass to the dataloader.
+ :return: The dataloader.
+ """
+ if phase == "train":
+ batch_size = self.hparams.batch_size
+ epoch_frac = self.hparams.epoch_frac
+ else:
+ batch_size = 1
+ epoch_frac = 1
+ sampled_indices = np.random.choice(
+ len(dataset),
+ int(len(dataset) * epoch_frac),
+ replace=False,
+ )
+ if phase == "val":
+ sampled_indices = np.repeat(sampled_indices, self.trainer.world_size)
+ subdataset = torch.utils.data.Subset(dataset, sampled_indices)
+ return DataLoader(
+ subdataset,
+ batch_size=None,
+ collate_fn=lambda x: self.dynamic_batching_by_max_edge_count(
+ x, self.hparams.edge_crop_size, batch_size
+ ),
+ num_workers=self.hparams.num_workers,
+ pin_memory=self.hparams.pin_memory,
+ **kwargs,
+ )
+
+ def train_dataloader(self) -> DataLoader[Any]:
+ """Create and return the train dataloader.
+
+ :return: The train dataloader.
+ """
+ return self.get_dataloader(
+ "train",
+ dataset=self.data_train,
+ shuffle=True,
+ )
+
+ def val_dataloader(self) -> List[DataLoader[Any]]:
+ """Create and return the validation dataloaders.
+
+ :return: The validation dataloaders.
+ """
+ val_dataloaders = []
+ if self.data_val is not None:
+ val_dataloaders.append(
+ self.get_dataloader(
+ "val",
+ dataset=self.data_val,
+ shuffle=False,
+ )
+ )
+ if self.data_val_secondary is not None:
+ val_dataloaders.append(
+ self.get_dataloader(
+ "val",
+ dataset=self.data_val_secondary,
+ shuffle=False,
+ )
+ )
+ assert len(val_dataloaders) > 0, "No validation datasets to be loaded."
+ return val_dataloaders
+
+ def test_dataloader(self) -> List[DataLoader[Any]]:
+ """Create and return the test dataloaders.
+
+ :return: The test dataloaders.
+ """
+ test_dataloaders = []
+ if self.data_test is not None:
+ test_dataloaders.append(
+ self.get_dataloader(
+ "test",
+ dataset=self.data_test,
+ shuffle=False,
+ )
+ )
+ if self.data_test_secondary is not None:
+ test_dataloaders.append(
+ self.get_dataloader(
+ "test",
+ dataset=self.data_test_secondary,
+ shuffle=False,
+ )
+ )
+ assert len(test_dataloaders) > 0, "No test datasets to be loaded."
+ return test_dataloaders
+
+ def teardown(self, stage: Optional[str] = None) -> None:
+ """Lightning hook for cleaning up after `trainer.fit()`, `trainer.validate()`,
+ `trainer.test()`, and `trainer.predict()`.
+
+ :param stage: The stage being torn down. Either `"fit"`, `"validate"`, `"test"`, or `"predict"`.
+ Defaults to ``None``.
+ """
+ pass
+
+ def state_dict(self) -> Dict[Any, Any]:
+ """Called when saving a checkpoint. Implement to generate and save the datamodule state.
+
+ :return: A dictionary containing the datamodule state that you want to save.
+ """
+ return {}
+
+ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+ """Called when loading a checkpoint. Implement to reload datamodule state given datamodule
+ `state_dict()`.
+
+ :param state_dict: The datamodule state returned by `self.state_dict()`.
+ """
+ pass
+
+
+if __name__ == "__main__":
+ _ = CombinedDataModule()
diff --git a/forks/FlowDock/flowdock/data/components/__init__.py b/forks/FlowDock/flowdock/data/components/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/forks/FlowDock/flowdock/data/components/chemical/20AA_template_peptide.pdb b/forks/FlowDock/flowdock/data/components/chemical/20AA_template_peptide.pdb
new file mode 100644
index 00000000..f624f1de
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/20AA_template_peptide.pdb
@@ -0,0 +1,189 @@
+HEADER 18-OCT-22
+TITLE ESMFOLD V1 PREDICTION FOR INPUT
+REMARK 1
+REMARK 1 REFERENCE 1
+REMARK 1 AUTH ZEMING LIN, HALIL AKIN, ROSHAN RAO, BRIAN HIE, ZHONGKAI ZHU,
+REMARK 1 AUTH 2 WENTING LU, NIKITA SMETANIN, ROBERT VERKUIL, ORI KABELI,
+REMARK 1 AUTH 3 YANIV SHMUELI, ALLAN DOS SANTOS COSTA,
+REMARK 1 AUTH 4 MARYAM FAZEL-ZARANDI, TOM SERCU, SALVATORE CANDIDO,
+REMARK 1 AUTH 5 ALEXANDER RIVES
+REMARK 1 TITL EVOLUTIONARY-SCALE PREDICTION OF ATOMIC LEVEL PROTEIN
+REMARK 1 TITL 2 STRUCTURE WITH A LANGUAGE MODEL
+REMARK 1 REF
+REMARK 1 REFN
+REMARK 1 PMID
+REMARK 1 DOI 10.1101/2022.07.20.500902
+REMARK 1
+REMARK 1 LICENSE AND DISCLAIMERS
+REMARK 1 ESM METAGENOMIC ATLAS DATA IS AVAILABLE UNDER
+REMARK 1 A CC-BY-4.0 LICENSE FOR ACADEMIC AND COMMERCIAL USE.
+REMARK 1 COPYRIGHT (C) META PLATFORMS, INC. ALL RIGHTS RESERVED.
+REMARK 1 USE OF THE ESM METAGENOMIC ATLAS DATA IS SUBJECT
+REMARK 1 TO THE META OPEN SOURCE TERMS OF USE AND PRIVACY POLICY.
+ATOM 1 N ALA A 1 -1.670 -7.378 10.112 1.00 0.56 N
+ATOM 2 CA ALA A 1 -0.600 -6.435 9.799 1.00 0.55 C
+ATOM 3 C ALA A 1 -0.467 -5.374 10.888 1.00 0.57 C
+ATOM 4 CB ALA A 1 -0.853 -5.774 8.446 1.00 0.52 C
+ATOM 5 O ALA A 1 -1.445 -4.709 11.239 1.00 0.55 O
+ATOM 6 N ARG A 2 0.292 -5.500 12.146 1.00 0.62 N
+ATOM 7 CA ARG A 2 0.421 -4.657 13.330 1.00 0.61 C
+ATOM 8 C ARG A 2 0.561 -3.188 12.944 1.00 0.62 C
+ATOM 9 CB ARG A 2 1.621 -5.094 14.173 1.00 0.58 C
+ATOM 10 O ARG A 2 1.519 -2.807 12.270 1.00 0.60 O
+ATOM 11 CG ARG A 2 1.836 -4.253 15.421 1.00 0.56 C
+ATOM 12 CD ARG A 2 3.029 -4.740 16.231 1.00 0.56 C
+ATOM 13 NE ARG A 2 3.017 -4.199 17.587 1.00 0.49 N
+ATOM 14 NH1 ARG A 2 5.124 -4.971 18.136 1.00 0.38 N
+ATOM 15 NH2 ARG A 2 3.896 -3.790 19.669 1.00 0.33 N
+ATOM 16 CZ ARG A 2 4.012 -4.322 18.461 1.00 0.55 C
+ATOM 17 N ASN A 3 -0.437 -2.382 12.444 1.00 0.60 N
+ATOM 18 CA ASN A 3 -0.632 -0.970 12.134 1.00 0.60 C
+ATOM 19 C ASN A 3 -0.016 -0.071 13.202 1.00 0.60 C
+ATOM 20 CB ASN A 3 -2.121 -0.657 11.971 1.00 0.56 C
+ATOM 21 O ASN A 3 -0.481 -0.047 14.343 1.00 0.58 O
+ATOM 22 CG ASN A 3 -2.695 -1.202 10.678 1.00 0.55 C
+ATOM 23 ND2 ASN A 3 -4.016 -1.320 10.618 1.00 0.53 N
+ATOM 24 OD1 ASN A 3 -1.956 -1.516 9.741 1.00 0.55 O
+ATOM 25 N ASP A 4 1.274 0.021 13.423 1.00 0.57 N
+ATOM 26 CA ASP A 4 2.172 1.003 14.023 1.00 0.57 C
+ATOM 27 C ASP A 4 1.607 2.415 13.894 1.00 0.58 C
+ATOM 28 CB ASP A 4 3.557 0.929 13.375 1.00 0.55 C
+ATOM 29 O ASP A 4 1.552 2.970 12.795 1.00 0.57 O
+ATOM 30 CG ASP A 4 3.557 0.167 12.062 1.00 0.55 C
+ATOM 31 OD1 ASP A 4 2.586 -0.569 11.783 1.00 0.55 O
+ATOM 32 OD2 ASP A 4 4.539 0.301 11.300 1.00 0.56 O
+ATOM 33 N CYS A 5 0.256 2.708 14.247 1.00 0.59 N
+ATOM 34 CA CYS A 5 -0.533 3.860 14.670 1.00 0.59 C
+ATOM 35 C CYS A 5 0.336 5.107 14.776 1.00 0.59 C
+ATOM 36 CB CYS A 5 -1.209 3.582 16.012 1.00 0.56 C
+ATOM 37 O CYS A 5 -0.124 6.152 15.241 1.00 0.58 O
+ATOM 38 SG CYS A 5 -1.228 1.835 16.470 1.00 0.56 S
+ATOM 39 N GLN A 6 1.254 5.296 13.796 1.00 0.59 N
+ATOM 40 CA GLN A 6 1.747 6.669 13.812 1.00 0.58 C
+ATOM 41 C GLN A 6 1.029 7.524 12.772 1.00 0.59 C
+ATOM 42 CB GLN A 6 3.256 6.701 13.566 1.00 0.56 C
+ATOM 43 O GLN A 6 0.710 7.046 11.681 1.00 0.57 O
+ATOM 44 CG GLN A 6 4.086 6.308 14.781 1.00 0.55 C
+ATOM 45 CD GLN A 6 5.575 6.284 14.491 1.00 0.56 C
+ATOM 46 NE2 GLN A 6 6.373 5.991 15.513 1.00 0.47 N
+ATOM 47 OE1 GLN A 6 6.005 6.525 13.359 1.00 0.57 O
+ATOM 48 N GLU A 7 0.211 8.541 13.186 1.00 0.61 N
+ATOM 49 CA GLU A 7 -0.316 9.705 12.480 1.00 0.61 C
+ATOM 50 C GLU A 7 0.397 9.912 11.147 1.00 0.61 C
+ATOM 51 CB GLU A 7 -0.187 10.962 13.345 1.00 0.56 C
+ATOM 52 O GLU A 7 1.626 10.003 11.101 1.00 0.58 O
+ATOM 53 CG GLU A 7 -0.985 12.149 12.827 1.00 0.54 C
+ATOM 54 CD GLU A 7 -0.902 13.369 13.731 1.00 0.56 C
+ATOM 55 OE1 GLU A 7 -0.258 13.288 14.802 1.00 0.54 O
+ATOM 56 OE2 GLU A 7 -1.486 14.414 13.367 1.00 0.48 O
+ATOM 57 N GLY A 8 -0.011 9.320 10.025 1.00 0.62 N
+ATOM 58 CA GLY A 8 0.256 9.579 8.619 1.00 0.61 C
+ATOM 59 C GLY A 8 0.720 8.349 7.863 1.00 0.62 C
+ATOM 60 O GLY A 8 1.049 8.429 6.678 1.00 0.60 O
+ATOM 61 N HIS A 9 0.589 7.093 8.240 1.00 0.60 N
+ATOM 62 CA HIS A 9 1.089 5.919 7.533 1.00 0.60 C
+ATOM 63 C HIS A 9 -0.056 5.013 7.091 1.00 0.60 C
+ATOM 64 CB HIS A 9 2.066 5.138 8.414 1.00 0.55 C
+ATOM 65 O HIS A 9 -0.684 4.350 7.919 1.00 0.57 O
+ATOM 66 CG HIS A 9 3.343 5.868 8.686 1.00 0.54 C
+ATOM 67 CD2 HIS A 9 3.795 7.065 8.242 1.00 0.53 C
+ATOM 68 ND1 HIS A 9 4.328 5.363 9.507 1.00 0.51 N
+ATOM 69 CE1 HIS A 9 5.333 6.221 9.557 1.00 0.48 C
+ATOM 70 NE2 HIS A 9 5.035 7.262 8.798 1.00 0.43 N
+ATOM 71 N ILE A 10 -0.948 5.551 6.156 1.00 0.65 N
+ATOM 72 CA ILE A 10 -1.786 4.506 5.579 1.00 0.64 C
+ATOM 73 C ILE A 10 -1.238 4.102 4.212 1.00 0.64 C
+ATOM 74 CB ILE A 10 -3.256 4.965 5.452 1.00 0.61 C
+ATOM 75 O ILE A 10 -1.049 4.950 3.336 1.00 0.61 O
+ATOM 76 CG1 ILE A 10 -3.784 5.444 6.809 1.00 0.55 C
+ATOM 77 CG2 ILE A 10 -4.127 3.838 4.892 1.00 0.56 C
+ATOM 78 CD1 ILE A 10 -5.171 4.920 7.154 1.00 0.55 C
+ATOM 79 N LEU A 11 -0.433 3.092 4.004 1.00 0.62 N
+ATOM 80 CA LEU A 11 0.013 2.300 2.863 1.00 0.61 C
+ATOM 81 C LEU A 11 -1.020 2.331 1.742 1.00 0.62 C
+ATOM 82 CB LEU A 11 0.281 0.854 3.287 1.00 0.59 C
+ATOM 83 O LEU A 11 -2.169 1.931 1.940 1.00 0.60 O
+ATOM 84 CG LEU A 11 1.484 0.627 4.203 1.00 0.57 C
+ATOM 85 CD1 LEU A 11 1.444 -0.779 4.793 1.00 0.55 C
+ATOM 86 CD2 LEU A 11 2.786 0.856 3.443 1.00 0.55 C
+ATOM 87 N LYS A 12 -0.950 3.372 0.786 1.00 0.66 N
+ATOM 88 CA LYS A 12 -1.589 3.550 -0.514 1.00 0.65 C
+ATOM 89 C LYS A 12 -1.379 2.328 -1.403 1.00 0.66 C
+ATOM 90 CB LYS A 12 -1.053 4.803 -1.210 1.00 0.62 C
+ATOM 91 O LYS A 12 -0.240 1.947 -1.684 1.00 0.63 O
+ATOM 92 CG LYS A 12 -1.358 6.098 -0.473 1.00 0.61 C
+ATOM 93 CD LYS A 12 -1.020 7.317 -1.321 1.00 0.59 C
+ATOM 94 CE LYS A 12 -1.187 8.611 -0.536 1.00 0.54 C
+ATOM 95 NZ LYS A 12 -0.898 9.811 -1.376 1.00 0.46 N
+ATOM 96 N MET A 13 -2.296 1.272 -1.465 1.00 0.65 N
+ATOM 97 CA MET A 13 -2.344 0.175 -2.427 1.00 0.64 C
+ATOM 98 C MET A 13 -2.899 0.650 -3.766 1.00 0.65 C
+ATOM 99 CB MET A 13 -3.193 -0.977 -1.888 1.00 0.61 C
+ATOM 100 O MET A 13 -3.887 1.384 -3.808 1.00 0.62 O
+ATOM 101 CG MET A 13 -2.648 -1.593 -0.610 1.00 0.59 C
+ATOM 102 SD MET A 13 -3.209 -3.324 -0.369 1.00 0.56 S
+ATOM 103 CE MET A 13 -4.949 -3.040 0.060 1.00 0.53 C
+ATOM 104 N PHE A 14 -2.006 0.639 -4.962 1.00 0.68 N
+ATOM 105 CA PHE A 14 -2.298 0.765 -6.385 1.00 0.67 C
+ATOM 106 C PHE A 14 -3.288 -0.304 -6.831 1.00 0.68 C
+ATOM 107 CB PHE A 14 -1.011 0.667 -7.210 1.00 0.65 C
+ATOM 108 O PHE A 14 -3.279 -1.422 -6.310 1.00 0.66 O
+ATOM 109 CG PHE A 14 -0.135 1.887 -7.116 1.00 0.64 C
+ATOM 110 CD1 PHE A 14 -0.449 3.041 -7.823 1.00 0.63 C
+ATOM 111 CD2 PHE A 14 1.003 1.879 -6.320 1.00 0.63 C
+ATOM 112 CE1 PHE A 14 0.360 4.172 -7.738 1.00 0.62 C
+ATOM 113 CE2 PHE A 14 1.816 3.006 -6.230 1.00 0.60 C
+ATOM 114 CZ PHE A 14 1.493 4.150 -6.941 1.00 0.63 C
+ATOM 115 N PRO A 15 -4.409 0.029 -7.509 1.00 0.69 N
+ATOM 116 CA PRO A 15 -5.414 -0.890 -8.049 1.00 0.67 C
+ATOM 117 C PRO A 15 -4.804 -2.179 -8.593 1.00 0.69 C
+ATOM 118 CB PRO A 15 -6.067 -0.079 -9.171 1.00 0.66 C
+ATOM 119 O PRO A 15 -3.776 -2.141 -9.274 1.00 0.66 O
+ATOM 120 CG PRO A 15 -5.389 1.253 -9.124 1.00 0.64 C
+ATOM 121 CD PRO A 15 -4.211 1.156 -8.197 1.00 0.63 C
+ATOM 122 N SER A 16 -4.944 -3.465 -8.076 1.00 0.68 N
+ATOM 123 CA SER A 16 -4.431 -4.782 -8.440 1.00 0.68 C
+ATOM 124 C SER A 16 -4.671 -5.080 -9.916 1.00 0.68 C
+ATOM 125 CB SER A 16 -5.080 -5.867 -7.580 1.00 0.65 C
+ATOM 126 O SER A 16 -4.031 -5.964 -10.488 1.00 0.69 O
+ATOM 127 OG SER A 16 -6.471 -5.951 -7.841 1.00 0.62 O
+ATOM 128 N THR A 17 -5.603 -4.552 -10.484 1.00 0.70 N
+ATOM 129 CA THR A 17 -5.932 -4.880 -11.866 1.00 0.68 C
+ATOM 130 C THR A 17 -4.834 -4.403 -12.812 1.00 0.69 C
+ATOM 131 CB THR A 17 -7.278 -4.258 -12.283 1.00 0.66 C
+ATOM 132 O THR A 17 -4.814 -4.776 -13.987 1.00 0.67 O
+ATOM 133 CG2 THR A 17 -8.430 -4.852 -11.480 1.00 0.62 C
+ATOM 134 OG1 THR A 17 -7.233 -2.844 -12.059 1.00 0.63 O
+ATOM 135 N TRP A 18 -3.961 -3.539 -12.286 1.00 0.67 N
+ATOM 136 CA TRP A 18 -2.831 -3.137 -13.116 1.00 0.69 C
+ATOM 137 C TRP A 18 -1.774 -4.235 -13.167 1.00 0.66 C
+ATOM 138 CB TRP A 18 -2.211 -1.839 -12.590 1.00 0.63 C
+ATOM 139 O TRP A 18 -0.883 -4.209 -14.019 1.00 0.64 O
+ATOM 140 CG TRP A 18 -3.043 -0.619 -12.847 1.00 0.59 C
+ATOM 141 CD1 TRP A 18 -4.352 -0.582 -13.241 1.00 0.53 C
+ATOM 142 CD2 TRP A 18 -2.620 0.743 -12.733 1.00 0.54 C
+ATOM 143 CE2 TRP A 18 -3.725 1.554 -13.072 1.00 0.53 C
+ATOM 144 CE3 TRP A 18 -1.412 1.357 -12.375 1.00 0.63 C
+ATOM 145 NE1 TRP A 18 -4.768 0.722 -13.378 1.00 0.62 N
+ATOM 146 CH2 TRP A 18 -2.465 3.524 -12.710 1.00 0.58 C
+ATOM 147 CZ2 TRP A 18 -3.658 2.950 -13.064 1.00 0.62 C
+ATOM 148 CZ3 TRP A 18 -1.347 2.746 -12.368 1.00 0.53 C
+ATOM 149 N TYR A 19 -1.742 -5.095 -12.150 1.00 0.64 N
+ATOM 150 CA TYR A 19 -0.725 -6.135 -12.259 1.00 0.64 C
+ATOM 151 C TYR A 19 -1.296 -7.390 -12.910 1.00 0.64 C
+ATOM 152 CB TYR A 19 -0.154 -6.476 -10.880 1.00 0.62 C
+ATOM 153 O TYR A 19 -2.167 -8.051 -12.340 1.00 0.65 O
+ATOM 154 CG TYR A 19 0.527 -5.313 -10.200 1.00 0.60 C
+ATOM 155 CD1 TYR A 19 1.860 -5.009 -10.467 1.00 0.58 C
+ATOM 156 CD2 TYR A 19 -0.160 -4.516 -9.291 1.00 0.58 C
+ATOM 157 CE1 TYR A 19 2.493 -3.939 -9.842 1.00 0.57 C
+ATOM 158 CE2 TYR A 19 0.462 -3.444 -8.660 1.00 0.55 C
+ATOM 159 OH TYR A 19 2.407 -2.103 -8.320 1.00 0.60 O
+ATOM 160 CZ TYR A 19 1.786 -3.164 -8.942 1.00 0.56 C
+ATOM 161 N VAL A 20 -2.060 -7.165 -13.824 1.00 0.65 N
+ATOM 162 CA VAL A 20 -2.336 -8.324 -14.666 1.00 0.66 C
+ATOM 163 C VAL A 20 -1.173 -8.552 -15.629 1.00 0.64 C
+ATOM 164 CB VAL A 20 -3.655 -8.154 -15.453 1.00 0.61 C
+ATOM 165 O VAL A 20 -0.527 -7.598 -16.069 1.00 0.63 O
+ATOM 166 CG1 VAL A 20 -4.137 -9.498 -15.997 1.00 0.53 C
+ATOM 167 CG2 VAL A 20 -4.725 -7.517 -14.569 1.00 0.54 C
diff --git a/forks/FlowDock/flowdock/data/components/chemical/ALA.pdb b/forks/FlowDock/flowdock/data/components/chemical/ALA.pdb
new file mode 100644
index 00000000..8de631eb
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/ALA.pdb
@@ -0,0 +1,5 @@
+ATOM 1 N ALA A 1 -1.670 -7.378 10.112 1.00 0.56 N
+ATOM 2 CA ALA A 1 -0.600 -6.435 9.799 1.00 0.55 C
+ATOM 3 C ALA A 1 -0.467 -5.374 10.888 1.00 0.57 C
+ATOM 4 CB ALA A 1 -0.853 -5.774 8.446 1.00 0.52 C
+ATOM 5 O ALA A 1 -1.445 -4.709 11.239 1.00 0.55 O
diff --git a/forks/FlowDock/flowdock/data/components/chemical/ARG.pdb b/forks/FlowDock/flowdock/data/components/chemical/ARG.pdb
new file mode 100644
index 00000000..86d39b44
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/ARG.pdb
@@ -0,0 +1,11 @@
+ATOM 6 N ARG A 2 0.292 -5.500 12.146 1.00 0.62 N
+ATOM 7 CA ARG A 2 0.421 -4.657 13.330 1.00 0.61 C
+ATOM 8 C ARG A 2 0.561 -3.188 12.944 1.00 0.62 C
+ATOM 9 CB ARG A 2 1.621 -5.094 14.173 1.00 0.58 C
+ATOM 10 O ARG A 2 1.519 -2.807 12.270 1.00 0.60 O
+ATOM 11 CG ARG A 2 1.836 -4.253 15.421 1.00 0.56 C
+ATOM 12 CD ARG A 2 3.029 -4.740 16.231 1.00 0.56 C
+ATOM 13 NE ARG A 2 3.017 -4.199 17.587 1.00 0.49 N
+ATOM 14 NH1 ARG A 2 5.124 -4.971 18.136 1.00 0.38 N
+ATOM 15 NH2 ARG A 2 3.896 -3.790 19.669 1.00 0.33 N
+ATOM 16 CZ ARG A 2 4.012 -4.322 18.461 1.00 0.55 C
diff --git a/forks/FlowDock/flowdock/data/components/chemical/ASN.pdb b/forks/FlowDock/flowdock/data/components/chemical/ASN.pdb
new file mode 100644
index 00000000..496dc76e
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/ASN.pdb
@@ -0,0 +1,8 @@
+ATOM 17 N ASN A 3 -0.437 -2.382 12.444 1.00 0.60 N
+ATOM 18 CA ASN A 3 -0.632 -0.970 12.134 1.00 0.60 C
+ATOM 19 C ASN A 3 -0.016 -0.071 13.202 1.00 0.60 C
+ATOM 20 CB ASN A 3 -2.121 -0.657 11.971 1.00 0.56 C
+ATOM 21 O ASN A 3 -0.481 -0.047 14.343 1.00 0.58 O
+ATOM 22 CG ASN A 3 -2.695 -1.202 10.678 1.00 0.55 C
+ATOM 23 ND2 ASN A 3 -4.016 -1.320 10.618 1.00 0.53 N
+ATOM 24 OD1 ASN A 3 -1.956 -1.516 9.741 1.00 0.55 O
diff --git a/forks/FlowDock/flowdock/data/components/chemical/ASP.pdb b/forks/FlowDock/flowdock/data/components/chemical/ASP.pdb
new file mode 100644
index 00000000..7d4cf534
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/ASP.pdb
@@ -0,0 +1,8 @@
+ATOM 25 N ASP A 4 1.274 0.021 13.423 1.00 0.57 N
+ATOM 26 CA ASP A 4 2.172 1.003 14.023 1.00 0.57 C
+ATOM 27 C ASP A 4 1.607 2.415 13.894 1.00 0.58 C
+ATOM 28 CB ASP A 4 3.557 0.929 13.375 1.00 0.55 C
+ATOM 29 O ASP A 4 1.552 2.970 12.795 1.00 0.57 O
+ATOM 30 CG ASP A 4 3.557 0.167 12.062 1.00 0.55 C
+ATOM 31 OD1 ASP A 4 2.586 -0.569 11.783 1.00 0.55 O
+ATOM 32 OD2 ASP A 4 4.539 0.301 11.300 1.00 0.56 O
diff --git a/forks/FlowDock/flowdock/data/components/chemical/CYS.pdb b/forks/FlowDock/flowdock/data/components/chemical/CYS.pdb
new file mode 100644
index 00000000..32ae97df
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/CYS.pdb
@@ -0,0 +1,6 @@
+ATOM 33 N CYS A 5 0.256 2.708 14.247 1.00 0.59 N
+ATOM 34 CA CYS A 5 -0.533 3.860 14.670 1.00 0.59 C
+ATOM 35 C CYS A 5 0.336 5.107 14.776 1.00 0.59 C
+ATOM 36 CB CYS A 5 -1.209 3.582 16.012 1.00 0.56 C
+ATOM 37 O CYS A 5 -0.124 6.152 15.241 1.00 0.58 O
+ATOM 38 SG CYS A 5 -1.228 1.835 16.470 1.00 0.56 S
diff --git a/forks/FlowDock/flowdock/data/components/chemical/GLN.pdb b/forks/FlowDock/flowdock/data/components/chemical/GLN.pdb
new file mode 100644
index 00000000..e815cc66
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/GLN.pdb
@@ -0,0 +1,9 @@
+ATOM 39 N GLN A 6 1.254 5.296 13.796 1.00 0.59 N
+ATOM 40 CA GLN A 6 1.747 6.669 13.812 1.00 0.58 C
+ATOM 41 C GLN A 6 1.029 7.524 12.772 1.00 0.59 C
+ATOM 42 CB GLN A 6 3.256 6.701 13.566 1.00 0.56 C
+ATOM 43 O GLN A 6 0.710 7.046 11.681 1.00 0.57 O
+ATOM 44 CG GLN A 6 4.086 6.308 14.781 1.00 0.55 C
+ATOM 45 CD GLN A 6 5.575 6.284 14.491 1.00 0.56 C
+ATOM 46 NE2 GLN A 6 6.373 5.991 15.513 1.00 0.47 N
+ATOM 47 OE1 GLN A 6 6.005 6.525 13.359 1.00 0.57 O
diff --git a/forks/FlowDock/flowdock/data/components/chemical/GLU.pdb b/forks/FlowDock/flowdock/data/components/chemical/GLU.pdb
new file mode 100644
index 00000000..94a89ab4
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/GLU.pdb
@@ -0,0 +1,9 @@
+ATOM 48 N GLU A 7 0.211 8.541 13.186 1.00 0.61 N
+ATOM 49 CA GLU A 7 -0.316 9.705 12.480 1.00 0.61 C
+ATOM 50 C GLU A 7 0.397 9.912 11.147 1.00 0.61 C
+ATOM 51 CB GLU A 7 -0.187 10.962 13.345 1.00 0.56 C
+ATOM 52 O GLU A 7 1.626 10.003 11.101 1.00 0.58 O
+ATOM 53 CG GLU A 7 -0.985 12.149 12.827 1.00 0.54 C
+ATOM 54 CD GLU A 7 -0.902 13.369 13.731 1.00 0.56 C
+ATOM 55 OE1 GLU A 7 -0.258 13.288 14.802 1.00 0.54 O
+ATOM 56 OE2 GLU A 7 -1.486 14.414 13.367 1.00 0.48 O
diff --git a/forks/FlowDock/flowdock/data/components/chemical/GLY.pdb b/forks/FlowDock/flowdock/data/components/chemical/GLY.pdb
new file mode 100644
index 00000000..f366c8bf
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/GLY.pdb
@@ -0,0 +1,4 @@
+ATOM 57 N GLY A 8 -0.011 9.320 10.025 1.00 0.62 N
+ATOM 58 CA GLY A 8 0.256 9.579 8.619 1.00 0.61 C
+ATOM 59 C GLY A 8 0.720 8.349 7.863 1.00 0.62 C
+ATOM 60 O GLY A 8 1.049 8.429 6.678 1.00 0.60 O
diff --git a/forks/FlowDock/flowdock/data/components/chemical/HIS.pdb b/forks/FlowDock/flowdock/data/components/chemical/HIS.pdb
new file mode 100644
index 00000000..7815c433
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/HIS.pdb
@@ -0,0 +1,10 @@
+ATOM 61 N HIS A 9 0.589 7.093 8.240 1.00 0.60 N
+ATOM 62 CA HIS A 9 1.089 5.919 7.533 1.00 0.60 C
+ATOM 63 C HIS A 9 -0.056 5.013 7.091 1.00 0.60 C
+ATOM 64 CB HIS A 9 2.066 5.138 8.414 1.00 0.55 C
+ATOM 65 O HIS A 9 -0.684 4.350 7.919 1.00 0.57 O
+ATOM 66 CG HIS A 9 3.343 5.868 8.686 1.00 0.54 C
+ATOM 67 CD2 HIS A 9 3.795 7.065 8.242 1.00 0.53 C
+ATOM 68 ND1 HIS A 9 4.328 5.363 9.507 1.00 0.51 N
+ATOM 69 CE1 HIS A 9 5.333 6.221 9.557 1.00 0.48 C
+ATOM 70 NE2 HIS A 9 5.035 7.262 8.798 1.00 0.43 N
diff --git a/forks/FlowDock/flowdock/data/components/chemical/ILE.pdb b/forks/FlowDock/flowdock/data/components/chemical/ILE.pdb
new file mode 100644
index 00000000..4d8d6249
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/ILE.pdb
@@ -0,0 +1,8 @@
+ATOM 71 N ILE A 10 -0.948 5.551 6.156 1.00 0.65 N
+ATOM 72 CA ILE A 10 -1.786 4.506 5.579 1.00 0.64 C
+ATOM 73 C ILE A 10 -1.238 4.102 4.212 1.00 0.64 C
+ATOM 74 CB ILE A 10 -3.256 4.965 5.452 1.00 0.61 C
+ATOM 75 O ILE A 10 -1.049 4.950 3.336 1.00 0.61 O
+ATOM 76 CG1 ILE A 10 -3.784 5.444 6.809 1.00 0.55 C
+ATOM 77 CG2 ILE A 10 -4.127 3.838 4.892 1.00 0.56 C
+ATOM 78 CD1 ILE A 10 -5.171 4.920 7.154 1.00 0.55 C
diff --git a/forks/FlowDock/flowdock/data/components/chemical/LEU.pdb b/forks/FlowDock/flowdock/data/components/chemical/LEU.pdb
new file mode 100644
index 00000000..5ec058f1
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/LEU.pdb
@@ -0,0 +1,8 @@
+ATOM 79 N LEU A 11 -0.433 3.092 4.004 1.00 0.62 N
+ATOM 80 CA LEU A 11 0.013 2.300 2.863 1.00 0.61 C
+ATOM 81 C LEU A 11 -1.020 2.331 1.742 1.00 0.62 C
+ATOM 82 CB LEU A 11 0.281 0.854 3.287 1.00 0.59 C
+ATOM 83 O LEU A 11 -2.169 1.931 1.940 1.00 0.60 O
+ATOM 84 CG LEU A 11 1.484 0.627 4.203 1.00 0.57 C
+ATOM 85 CD1 LEU A 11 1.444 -0.779 4.793 1.00 0.55 C
+ATOM 86 CD2 LEU A 11 2.786 0.856 3.443 1.00 0.55 C
diff --git a/forks/FlowDock/flowdock/data/components/chemical/LYS.pdb b/forks/FlowDock/flowdock/data/components/chemical/LYS.pdb
new file mode 100644
index 00000000..54c97246
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/LYS.pdb
@@ -0,0 +1,9 @@
+ATOM 87 N LYS A 12 -0.950 3.372 0.786 1.00 0.66 N
+ATOM 88 CA LYS A 12 -1.589 3.550 -0.514 1.00 0.65 C
+ATOM 89 C LYS A 12 -1.379 2.328 -1.403 1.00 0.66 C
+ATOM 90 CB LYS A 12 -1.053 4.803 -1.210 1.00 0.62 C
+ATOM 91 O LYS A 12 -0.240 1.947 -1.684 1.00 0.63 O
+ATOM 92 CG LYS A 12 -1.358 6.098 -0.473 1.00 0.61 C
+ATOM 93 CD LYS A 12 -1.020 7.317 -1.321 1.00 0.59 C
+ATOM 94 CE LYS A 12 -1.187 8.611 -0.536 1.00 0.54 C
+ATOM 95 NZ LYS A 12 -0.898 9.811 -1.376 1.00 0.46 N
diff --git a/forks/FlowDock/flowdock/data/components/chemical/MET.pdb b/forks/FlowDock/flowdock/data/components/chemical/MET.pdb
new file mode 100644
index 00000000..ad1dcdd1
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/MET.pdb
@@ -0,0 +1,13 @@
+REMARK 1 AUTH 2 WENTING LU, NIKITA SMETANIN, ROBERT VERKUIL, ORI KABELI,
+REMARK 1 ESM METAGENOMIC ATLAS DATA IS AVAILABLE UNDER
+REMARK 1 COPYRIGHT (C) META PLATFORMS, INC. ALL RIGHTS RESERVED.
+REMARK 1 USE OF THE ESM METAGENOMIC ATLAS DATA IS SUBJECT
+REMARK 1 TO THE META OPEN SOURCE TERMS OF USE AND PRIVACY POLICY.
+ATOM 96 N MET A 13 -2.296 1.272 -1.465 1.00 0.65 N
+ATOM 97 CA MET A 13 -2.344 0.175 -2.427 1.00 0.64 C
+ATOM 98 C MET A 13 -2.899 0.650 -3.766 1.00 0.65 C
+ATOM 99 CB MET A 13 -3.193 -0.977 -1.888 1.00 0.61 C
+ATOM 100 O MET A 13 -3.887 1.384 -3.808 1.00 0.62 O
+ATOM 101 CG MET A 13 -2.648 -1.593 -0.610 1.00 0.59 C
+ATOM 102 SD MET A 13 -3.209 -3.324 -0.369 1.00 0.56 S
+ATOM 103 CE MET A 13 -4.949 -3.040 0.060 1.00 0.53 C
diff --git a/forks/FlowDock/flowdock/data/components/chemical/PHE.pdb b/forks/FlowDock/flowdock/data/components/chemical/PHE.pdb
new file mode 100644
index 00000000..2fc74289
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/PHE.pdb
@@ -0,0 +1,11 @@
+ATOM 104 N PHE A 14 -2.006 0.639 -4.962 1.00 0.68 N
+ATOM 105 CA PHE A 14 -2.298 0.765 -6.385 1.00 0.67 C
+ATOM 106 C PHE A 14 -3.288 -0.304 -6.831 1.00 0.68 C
+ATOM 107 CB PHE A 14 -1.011 0.667 -7.210 1.00 0.65 C
+ATOM 108 O PHE A 14 -3.279 -1.422 -6.310 1.00 0.66 O
+ATOM 109 CG PHE A 14 -0.135 1.887 -7.116 1.00 0.64 C
+ATOM 110 CD1 PHE A 14 -0.449 3.041 -7.823 1.00 0.63 C
+ATOM 111 CD2 PHE A 14 1.003 1.879 -6.320 1.00 0.63 C
+ATOM 112 CE1 PHE A 14 0.360 4.172 -7.738 1.00 0.62 C
+ATOM 113 CE2 PHE A 14 1.816 3.006 -6.230 1.00 0.60 C
+ATOM 114 CZ PHE A 14 1.493 4.150 -6.941 1.00 0.63 C
diff --git a/forks/FlowDock/flowdock/data/components/chemical/PRO.pdb b/forks/FlowDock/flowdock/data/components/chemical/PRO.pdb
new file mode 100644
index 00000000..b741e659
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/PRO.pdb
@@ -0,0 +1,8 @@
+REMARK 1 TITL EVOLUTIONARY-SCALE PREDICTION OF ATOMIC LEVEL PROTEIN
+ATOM 115 N PRO A 15 -4.409 0.029 -7.509 1.00 0.69 N
+ATOM 116 CA PRO A 15 -5.414 -0.890 -8.049 1.00 0.67 C
+ATOM 117 C PRO A 15 -4.804 -2.179 -8.593 1.00 0.69 C
+ATOM 118 CB PRO A 15 -6.067 -0.079 -9.171 1.00 0.66 C
+ATOM 119 O PRO A 15 -3.776 -2.141 -9.274 1.00 0.66 O
+ATOM 120 CG PRO A 15 -5.389 1.253 -9.124 1.00 0.64 C
+ATOM 121 CD PRO A 15 -4.211 1.156 -8.197 1.00 0.63 C
diff --git a/forks/FlowDock/flowdock/data/components/chemical/SER.pdb b/forks/FlowDock/flowdock/data/components/chemical/SER.pdb
new file mode 100644
index 00000000..24e59dc3
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/SER.pdb
@@ -0,0 +1,8 @@
+REMARK 1 AUTH 4 MARYAM FAZEL-ZARANDI, TOM SERCU, SALVATORE CANDIDO,
+REMARK 1 COPYRIGHT (C) META PLATFORMS, INC. ALL RIGHTS RESERVED.
+ATOM 122 N SER A 16 -4.944 -3.465 -8.076 1.00 0.68 N
+ATOM 123 CA SER A 16 -4.431 -4.782 -8.440 1.00 0.68 C
+ATOM 124 C SER A 16 -4.671 -5.080 -9.916 1.00 0.68 C
+ATOM 125 CB SER A 16 -5.080 -5.867 -7.580 1.00 0.65 C
+ATOM 126 O SER A 16 -4.031 -5.964 -10.488 1.00 0.69 O
+ATOM 127 OG SER A 16 -6.471 -5.951 -7.841 1.00 0.62 O
diff --git a/forks/FlowDock/flowdock/data/components/chemical/THR.pdb b/forks/FlowDock/flowdock/data/components/chemical/THR.pdb
new file mode 100644
index 00000000..de104108
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/THR.pdb
@@ -0,0 +1,7 @@
+ATOM 128 N THR A 17 -5.603 -4.552 -10.484 1.00 0.70 N
+ATOM 129 CA THR A 17 -5.932 -4.880 -11.866 1.00 0.68 C
+ATOM 130 C THR A 17 -4.834 -4.403 -12.812 1.00 0.69 C
+ATOM 131 CB THR A 17 -7.278 -4.258 -12.283 1.00 0.66 C
+ATOM 132 O THR A 17 -4.814 -4.776 -13.987 1.00 0.67 O
+ATOM 133 CG2 THR A 17 -8.430 -4.852 -11.480 1.00 0.62 C
+ATOM 134 OG1 THR A 17 -7.233 -2.844 -12.059 1.00 0.63 O
diff --git a/forks/FlowDock/flowdock/data/components/chemical/TRP.pdb b/forks/FlowDock/flowdock/data/components/chemical/TRP.pdb
new file mode 100644
index 00000000..8fa157e7
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/TRP.pdb
@@ -0,0 +1,14 @@
+ATOM 135 N TRP A 18 -3.961 -3.539 -12.286 1.00 0.67 N
+ATOM 136 CA TRP A 18 -2.831 -3.137 -13.116 1.00 0.69 C
+ATOM 137 C TRP A 18 -1.774 -4.235 -13.167 1.00 0.66 C
+ATOM 138 CB TRP A 18 -2.211 -1.839 -12.590 1.00 0.63 C
+ATOM 139 O TRP A 18 -0.883 -4.209 -14.019 1.00 0.64 O
+ATOM 140 CG TRP A 18 -3.043 -0.619 -12.847 1.00 0.59 C
+ATOM 141 CD1 TRP A 18 -4.352 -0.582 -13.241 1.00 0.53 C
+ATOM 142 CD2 TRP A 18 -2.620 0.743 -12.733 1.00 0.54 C
+ATOM 143 CE2 TRP A 18 -3.725 1.554 -13.072 1.00 0.53 C
+ATOM 144 CE3 TRP A 18 -1.412 1.357 -12.375 1.00 0.63 C
+ATOM 145 NE1 TRP A 18 -4.768 0.722 -13.378 1.00 0.62 N
+ATOM 146 CH2 TRP A 18 -2.465 3.524 -12.710 1.00 0.58 C
+ATOM 147 CZ2 TRP A 18 -3.658 2.950 -13.064 1.00 0.62 C
+ATOM 148 CZ3 TRP A 18 -1.347 2.746 -12.368 1.00 0.53 C
diff --git a/forks/FlowDock/flowdock/data/components/chemical/TYR.pdb b/forks/FlowDock/flowdock/data/components/chemical/TYR.pdb
new file mode 100644
index 00000000..27e7b785
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/TYR.pdb
@@ -0,0 +1,12 @@
+ATOM 149 N TYR A 19 -1.742 -5.095 -12.150 1.00 0.64 N
+ATOM 150 CA TYR A 19 -0.725 -6.135 -12.259 1.00 0.64 C
+ATOM 151 C TYR A 19 -1.296 -7.390 -12.910 1.00 0.64 C
+ATOM 152 CB TYR A 19 -0.154 -6.476 -10.880 1.00 0.62 C
+ATOM 153 O TYR A 19 -2.167 -8.051 -12.340 1.00 0.65 O
+ATOM 154 CG TYR A 19 0.527 -5.313 -10.200 1.00 0.60 C
+ATOM 155 CD1 TYR A 19 1.860 -5.009 -10.467 1.00 0.58 C
+ATOM 156 CD2 TYR A 19 -0.160 -4.516 -9.291 1.00 0.58 C
+ATOM 157 CE1 TYR A 19 2.493 -3.939 -9.842 1.00 0.57 C
+ATOM 158 CE2 TYR A 19 0.462 -3.444 -8.660 1.00 0.55 C
+ATOM 159 OH TYR A 19 2.407 -2.103 -8.320 1.00 0.60 O
+ATOM 160 CZ TYR A 19 1.786 -3.164 -8.942 1.00 0.56 C
diff --git a/forks/FlowDock/flowdock/data/components/chemical/VAL.pdb b/forks/FlowDock/flowdock/data/components/chemical/VAL.pdb
new file mode 100644
index 00000000..e0d10be1
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/VAL.pdb
@@ -0,0 +1,7 @@
+ATOM 161 N VAL A 20 -2.060 -7.165 -13.824 1.00 0.65 N
+ATOM 162 CA VAL A 20 -2.336 -8.324 -14.666 1.00 0.66 C
+ATOM 163 C VAL A 20 -1.173 -8.552 -15.629 1.00 0.64 C
+ATOM 164 CB VAL A 20 -3.655 -8.154 -15.453 1.00 0.61 C
+ATOM 165 O VAL A 20 -0.527 -7.598 -16.069 1.00 0.63 O
+ATOM 166 CG1 VAL A 20 -4.137 -9.498 -15.997 1.00 0.53 C
+ATOM 167 CG2 VAL A 20 -4.725 -7.517 -14.569 1.00 0.54 C
diff --git a/forks/FlowDock/flowdock/data/components/chemical/standard_AA_featurization_test.ipynb b/forks/FlowDock/flowdock/data/components/chemical/standard_AA_featurization_test.ipynb
new file mode 100644
index 00000000..fea90295
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/chemical/standard_AA_featurization_test.ipynb
@@ -0,0 +1,135 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "a790013e",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from rdkit.Chem import AllChem"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "5addae52",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "m = AllChem.MolFromMolFile(\"20AA_peptide.mol\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "421115e7",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "AllChem.Compute2DCoords(m)\n",
+ "AllChem.Draw.MolToImage(m)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "a164b9b1",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "print(AllChem.MolToMolBlock(m))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "d6d21d12",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "m = AllChem.MolFromPDBFile(\"ARG.pdb\")\n",
+ "AllChem.Compute2DCoords(m)\n",
+ "AllChem.Draw.MolToImage(m)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "9b604ff2",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "m = AllChem.MolFromPDBFile(\"ALA.pdb\")\n",
+ "AllChem.Compute2DCoords(m)\n",
+ "AllChem.Draw.MolToImage(m)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "83b2fc70",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "ms = []\n",
+ "for code in [\n",
+ " \"ALA\",\n",
+ " \"ARG\",\n",
+ " \"ASN\",\n",
+ " \"ASP\",\n",
+ " \"CYS\",\n",
+ " \"GLN\",\n",
+ " \"GLU\",\n",
+ " \"GLY\",\n",
+ " \"HIS\",\n",
+ " \"ILE\",\n",
+ " \"LEU\",\n",
+ " \"LYS\",\n",
+ " \"MET\",\n",
+ " \"PHE\",\n",
+ " \"PRO\",\n",
+ " \"SER\",\n",
+ " \"THR\",\n",
+ " \"TRP\",\n",
+ " \"TYR\",\n",
+ " \"VAL\",\n",
+ "]:\n",
+ " m = AllChem.MolFromPDBFile(f\"{code}.pdb\")\n",
+ " m = AllChem.RemoveHs(m, updateExplicitCount=True)\n",
+ " AllChem.Compute2DCoords(m)\n",
+ " ms.append(m)\n",
+ "AllChem.Draw.MolsToGridImage(ms)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "aef1f6da",
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3 (ipykernel)",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.8.5"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/forks/FlowDock/flowdock/data/components/combined_dataset.py b/forks/FlowDock/flowdock/data/components/combined_dataset.py
new file mode 100644
index 00000000..57b43a94
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/combined_dataset.py
@@ -0,0 +1,27 @@
+from beartype.typing import Any, List
+from torch.utils.data import Dataset
+
+
+class CombinedDataset(Dataset):
+ """A `Dataset` that combines two datasets."""
+
+ def __init__(self, dataset1: Dataset, dataset2: Dataset):
+ """Initialize a `CombinedDataset`."""
+ super().__init__()
+ self.dataset1 = dataset1
+ self.dataset2 = dataset2
+
+ def __len__(self):
+ """Return the length of the combined dataset."""
+ return len(self.dataset1) + len(self.dataset2)
+
+ def __getitem__(self, idx: int):
+ """Return the example at the given index."""
+ if idx < len(self.dataset1):
+ return self.dataset1[idx]
+ else:
+ return self.dataset2[idx - len(self.dataset1)]
+
+ def add_examples(self, new_example_list: List[Any]):
+ """Add new examples to the first dataset."""
+ self.dataset1.add_examples(new_example_list)
diff --git a/forks/FlowDock/flowdock/data/components/constants.py b/forks/FlowDock/flowdock/data/components/constants.py
new file mode 100644
index 00000000..8c6a6fac
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/constants.py
@@ -0,0 +1,439 @@
+# From: https://github.com/gcorso/DiffDock
+# Significant contribution from Ben Fry and Nick Polizzi
+
+three_to_one = {
+ "ALA": "A",
+ "ARG": "R",
+ "ASN": "N",
+ "ASP": "D",
+ "CYS": "C",
+ "GLN": "Q",
+ "GLU": "E",
+ "GLY": "G",
+ "HIS": "H",
+ "ILE": "I",
+ "LEU": "L",
+ "LYS": "K",
+ "MET": "M",
+ "MSE": "M", # MSE this is almost the same AA as MET. The sulfur is just replaced by Selen
+ "PHE": "F",
+ "PRO": "P",
+ "PYL": "O",
+ "SER": "S",
+ "SEC": "U",
+ "THR": "T",
+ "TRP": "W",
+ "TYR": "Y",
+ "VAL": "V",
+ "ASX": "B",
+ "GLX": "Z",
+ "XAA": "X",
+ "XLE": "J",
+}
+
+
+aa_name2aa_idx = {
+ "ALA": 0,
+ "ARG": 1,
+ "ASN": 2,
+ "ASP": 3,
+ "CYS": 4,
+ "GLU": 5,
+ "GLN": 6,
+ "GLY": 7,
+ "HIS": 8,
+ "ILE": 9,
+ "LEU": 10,
+ "LYS": 11,
+ "MET": 12,
+ "PHE": 13,
+ "PRO": 14,
+ "SER": 15,
+ "THR": 16,
+ "TRP": 17,
+ "TYR": 18,
+ "VAL": 19,
+ "MSE": 12,
+}
+
+
+aa_short2long = {
+ "C": "CYS",
+ "D": "ASP",
+ "S": "SER",
+ "Q": "GLN",
+ "K": "LYS",
+ "I": "ILE",
+ "P": "PRO",
+ "T": "THR",
+ "F": "PHE",
+ "N": "ASN",
+ "G": "GLY",
+ "H": "HIS",
+ "L": "LEU",
+ "R": "ARG",
+ "W": "TRP",
+ "A": "ALA",
+ "V": "VAL",
+ "E": "GLU",
+ "Y": "TYR",
+ "M": "MET",
+}
+
+
+aa_short2aa_idx = {
+ aa_short: aa_name2aa_idx[aa_long] for aa_short, aa_long in aa_short2long.items()
+}
+aa_idx2aa_short = {aa_idx: aa_short for aa_short, aa_idx in aa_short2aa_idx.items()}
+aa_long2short = {aa_long: aa_short for aa_short, aa_long in aa_short2long.items()}
+aa_long2short["MSE"] = "M"
+
+chi = {
+ "C": {1: ("N", "CA", "CB", "SG")},
+ "D": {
+ 1: ("N", "CA", "CB", "CG"),
+ 2: ("CA", "CB", "CG", "OD1"),
+ },
+ "E": {
+ 1: ("N", "CA", "CB", "CG"),
+ 2: ("CA", "CB", "CG", "CD"),
+ 3: ("CB", "CG", "CD", "OE1"),
+ },
+ "F": {
+ 1: ("N", "CA", "CB", "CG"),
+ 2: ("CA", "CB", "CG", "CD1"),
+ },
+ "H": {
+ 1: ("N", "CA", "CB", "CG"),
+ 2: ("CA", "CB", "CG", "ND1"),
+ },
+ "I": {
+ 1: ("N", "CA", "CB", "CG1"),
+ 2: ("CA", "CB", "CG1", "CD1"),
+ },
+ "K": {
+ 1: ("N", "CA", "CB", "CG"),
+ 2: ("CA", "CB", "CG", "CD"),
+ 3: ("CB", "CG", "CD", "CE"),
+ 4: ("CG", "CD", "CE", "NZ"),
+ },
+ "L": {
+ 1: ("N", "CA", "CB", "CG"),
+ 2: ("CA", "CB", "CG", "CD1"),
+ },
+ "M": {
+ 1: ("N", "CA", "CB", "CG"),
+ 2: ("CA", "CB", "CG", "SD"),
+ 3: ("CB", "CG", "SD", "CE"),
+ },
+ "N": {
+ 1: ("N", "CA", "CB", "CG"),
+ 2: ("CA", "CB", "CG", "OD1"),
+ },
+ "P": {
+ 1: ("N", "CA", "CB", "CG"),
+ 2: ("CA", "CB", "CG", "CD"),
+ },
+ "Q": {
+ 1: ("N", "CA", "CB", "CG"),
+ 2: ("CA", "CB", "CG", "CD"),
+ 3: ("CB", "CG", "CD", "OE1"),
+ },
+ "R": {
+ 1: ("N", "CA", "CB", "CG"),
+ 2: ("CA", "CB", "CG", "CD"),
+ 3: ("CB", "CG", "CD", "NE"),
+ 4: ("CG", "CD", "NE", "CZ"),
+ },
+ "S": {
+ 1: ("N", "CA", "CB", "OG"),
+ },
+ "T": {
+ 1: ("N", "CA", "CB", "OG1"),
+ },
+ "V": {
+ 1: ("N", "CA", "CB", "CG1"),
+ },
+ "W": {
+ 1: ("N", "CA", "CB", "CG"),
+ 2: ("CA", "CB", "CG", "CD1"),
+ },
+ "Y": {
+ 1: ("N", "CA", "CB", "CG"),
+ 2: ("CA", "CB", "CG", "CD1"),
+ },
+}
+
+
+atom_order = {
+ "G": ["N", "CA", "C", "O"],
+ "A": ["N", "CA", "C", "O", "CB"],
+ "S": ["N", "CA", "C", "O", "CB", "OG"],
+ "C": ["N", "CA", "C", "O", "CB", "SG"],
+ "T": ["N", "CA", "C", "O", "CB", "OG1", "CG2"],
+ "P": ["N", "CA", "C", "O", "CB", "CG", "CD"],
+ "V": ["N", "CA", "C", "O", "CB", "CG1", "CG2"],
+ "M": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE"],
+ "N": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2"],
+ "I": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1"],
+ "L": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2"],
+ "D": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2"],
+ "E": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2"],
+ "K": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ"],
+ "Q": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2"],
+ "H": ["N", "CA", "C", "O", "CB", "CG", "ND1", "CD2", "CE1", "NE2"],
+ "F": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ"],
+ "R": ["N", "CA", "C", "O", "CB", "CG", "CD", "NE", "CZ", "NH1", "NH2"],
+ "Y": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "OH"],
+ "W": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE2", "CE3", "NE1", "CZ2", "CZ3", "CH2"],
+ "X": ["N", "CA", "C", "O"],
+} # unknown amino acid
+
+
+amino_acid_smiles = {
+ "PHE": "[NH3+]CC(=O)N[C@@H](Cc1ccccc1)C(=O)NCC(=O)O",
+ "MET": "CSCC[C@H](NC(=O)C[NH3+])C(=O)NCC(=O)O",
+ "TYR": "[NH3+]CC(=O)N[C@@H](Cc1ccc(O)cc1)C(=O)NCC(=O)O",
+ "ILE": "CC[C@H](C)[C@H](NC(=O)C[NH3+])C(=O)NCC(=O)O",
+ "TRP": "[NH3+]CC(=O)N[C@@H](Cc1c[nH]c2ccccc12)C(=O)NCC(=O)O",
+ "THR": "C[C@@H](O)[C@H](NC(=O)C[NH3+])C(=O)NCC(=O)O",
+ "CYS": "[NH3+]CC(=O)N[C@@H](CS)C(=O)NCC(=O)O",
+ "ALA": "C[C@H](NC(=O)C[NH3+])C(=O)NCC(=O)O",
+ "LYS": "[NH3+]CCCC[C@H](NC(=O)C[NH3+])C(=O)NCC(=O)O",
+ "PRO": "[NH3+]CC(=O)N1CCC[C@H]1C(=O)NCC(=O)O",
+ "LEU": "CC(C)C[C@H](NC(=O)C[NH3+])C(=O)NCC(=O)O",
+ "GLY": "[NH3+]CC(=O)NCC(=O)NCC(=O)O",
+ "ASP": "[NH3+]CC(=O)N[C@@H](CC(=O)O)C(=O)NCC(=O)O",
+ "HIS": "[NH3+]CC(=O)N[C@@H](Cc1c[nH]c[nH+]1)C(=O)NCC(=O)O",
+ "VAL": "CC(C)[C@H](NC(=O)C[NH3+])C(=O)NCC(=O)O",
+ "SER": "[NH3+]CC(=O)N[C@@H](CO)C(=O)NCC(=O)O",
+ "ARG": "NC(=[NH2+])NCCC[C@H](NC(=O)C[NH3+])C(=O)NCC(=O)O",
+ "GLU": "[NH3+]CC(=O)N[C@@H](CCC(=O)O)C(=O)NCC(=O)O",
+ "GLN": "NC(=O)CC[C@H](NC(=O)C[NH3+])C(=O)NCC(=O)O",
+ "ASN": "NC(=O)C[C@H](NC(=O)C[NH3+])C(=O)NCC(=O)O",
+}
+
+cg_rdkit_indices = {
+ "PHE": {
+ 4: "N",
+ 5: "CA",
+ 13: "C",
+ 14: "O",
+ 6: "CB",
+ 7: "CG",
+ 8: "CD1",
+ 12: "CD2",
+ 9: "CE1",
+ 11: "CE2",
+ 10: "CZ",
+ },
+ "MET": {5: "N", 4: "CA", 10: "C", 11: "O", 3: "CB", 2: "CG", 1: "SD", 0: "CE"},
+ "TYR": {
+ 4: "N",
+ 5: "CA",
+ 14: "C",
+ 15: "O",
+ 6: "CB",
+ 7: "CG",
+ 8: "CD1",
+ 13: "CD2",
+ 9: "CE1",
+ 12: "CE2",
+ 10: "CZ",
+ 11: "OH",
+ },
+ "ILE": {5: "N", 4: "CA", 10: "C", 11: "O", 2: "CB", 1: "CG1", 3: "CG2", 0: "CD1"},
+ "TRP": {
+ 4: "N",
+ 5: "CA",
+ 16: "C",
+ 17: "O",
+ 6: "CB",
+ 7: "CG",
+ 8: "CD1",
+ 15: "CD2",
+ 9: "NE1",
+ 10: "CE2",
+ 14: "CE3",
+ 11: "CZ2",
+ 13: "CZ3",
+ 12: "CH2",
+ },
+ "THR": {4: "N", 3: "CA", 9: "C", 10: "O", 1: "CB", 2: "OG1", 0: "CG2"},
+ "CYS": {4: "N", 5: "CA", 8: "C", 9: "O", 6: "CB", 7: "SG"},
+ "ALA": {2: "N", 1: "CA", 7: "C", 8: "O", 0: "CB"},
+ "LYS": {6: "N", 5: "CA", 11: "C", 12: "O", 4: "CB", 3: "CG", 2: "CD", 1: "CE", 0: "NZ"},
+ "PRO": {4: "N", 8: "CA", 9: "C", 10: "O", 7: "CB", 6: "CG", 5: "CD"},
+ "LEU": {5: "N", 4: "CA", 10: "C", 11: "O", 3: "CB", 1: "CG", 0: "CD1", 2: "CD2"},
+ "GLY": {4: "N", 5: "CA", 6: "C", 7: "O"},
+ "ASP": {4: "N", 5: "CA", 10: "C", 11: "O", 6: "CB", 7: "CG", 8: "OD1", 9: "OD2"},
+ "HIS": {
+ 4: "N",
+ 5: "CA",
+ 12: "C",
+ 13: "O",
+ 6: "CB",
+ 7: "CG",
+ 11: "ND1",
+ 8: "CD2",
+ 10: "CE1",
+ 9: "NE2",
+ },
+ "VAL": {4: "N", 3: "CA", 9: "C", 10: "O", 1: "CB", 0: "CG1", 2: "CG2"},
+ "SER": {4: "N", 5: "CA", 8: "C", 9: "O", 6: "CB", 7: "OG"},
+ "ARG": {
+ 8: "N",
+ 7: "CA",
+ 13: "C",
+ 14: "O",
+ 6: "CB",
+ 5: "CG",
+ 4: "CD",
+ 3: "NE",
+ 1: "CZ",
+ 0: "NH1",
+ 2: "NH2",
+ },
+ "GLU": {4: "N", 5: "CA", 11: "C", 12: "O", 6: "CB", 7: "CG", 8: "CD", 9: "OE1", 10: "OE2"},
+ "GLN": {6: "N", 5: "CA", 11: "C", 12: "O", 4: "CB", 3: "CG", 1: "CD", 2: "OE1", 0: "NE2"},
+ "ASN": {5: "N", 4: "CA", 10: "C", 11: "O", 3: "CB", 1: "CG", 2: "OD1", 0: "ND2"},
+}
+
+aa_to_cg_indices = {
+ aa_long2short[x]: [atom_order[aa_long2short[x]].index(aname) for aname in index_dict.values()]
+ for x, index_dict in cg_rdkit_indices.items()
+}
+
+af2_latest_excluded_ligs = [
+ "144",
+ "15P",
+ "1PE",
+ "2F2",
+ "2JC",
+ "3HR",
+ "3SY",
+ "7N5",
+ "7PE",
+ "9JE",
+ "AAE",
+ "ABA",
+ "ACE",
+ "ACN",
+ "ACT",
+ "ACY",
+ "AZI",
+ "BAM",
+ "BCN",
+ "BCT",
+ "BDN",
+ "BEN",
+ "BME",
+ "BO3",
+ "BTB",
+ "BTC",
+ "BU1",
+ "C8E",
+ "CAD",
+ "CAQ",
+ "CBM",
+ "CCN",
+ "CIT",
+ "CL",
+ "CLR",
+ "CM",
+ "CMO",
+ "CO3",
+ "CPT",
+ "CXS",
+ "D10",
+ "DEP",
+ "DIO",
+ "DMS",
+ "DN",
+ "DOD",
+ "DOX",
+ "EDO",
+ "EEE",
+ "EGL",
+ "EOH",
+ "EOX",
+ "EPE",
+ "ETF",
+ "FCY",
+ "FJO",
+ "FLC",
+ "FMT",
+ "FW5",
+ "GOL",
+ "GSH",
+ "GTT",
+ "GYF",
+ "HED",
+ "IHP",
+ "IHS",
+ "IMD",
+ "IOD",
+ "IPA",
+ "IPH",
+ "LDA",
+ "MB3",
+ "MEG",
+ "MES",
+ "MLA",
+ "MLI",
+ "MOH",
+ "MPD",
+ "MRD",
+ "MSE",
+ "MYR",
+ "N",
+ "NA",
+ "NH2",
+ "NH4",
+ "NHE",
+ "NO3",
+ "O4B",
+ "OHE",
+ "OLA",
+ "OLC",
+ "OMB",
+ "OME",
+ "OXA",
+ "P6G",
+ "PE3",
+ "PE4",
+ "PEG",
+ "PEO",
+ "PEP",
+ "PG0",
+ "PG4",
+ "PGE",
+ "PGR",
+ "PLM",
+ "PO4",
+ "POL",
+ "POP",
+ "PVO",
+ "SAR",
+ "SCN",
+ "SEO",
+ "SEP",
+ "SIN",
+ "SO4",
+ "SPD",
+ "SPM",
+ "SR",
+ "STE",
+ "STO",
+ "STU",
+ "TAR",
+ "TBU",
+ "TME",
+ "TPO",
+ "TRS",
+ "UNK",
+ "UNL",
+ "UNX",
+ "UPL",
+ "URE",
+]
diff --git a/forks/FlowDock/flowdock/data/components/esm_embedding_extraction.py b/forks/FlowDock/flowdock/data/components/esm_embedding_extraction.py
new file mode 100644
index 00000000..d1cf3d5c
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/esm_embedding_extraction.py
@@ -0,0 +1,188 @@
+#!/usr/bin/env python3 -u
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import os
+import pathlib
+
+import rootutils
+import torch
+from esm import FastaBatchedDataset, MSATransformer, pretrained
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+# ------------------------------------------------------------------------------------ #
+# the setup_root above is equivalent to:
+# - adding project root dir to PYTHONPATH
+# (so you don't need to force user to install project as a package)
+# (necessary before importing any local modules e.g. `from flowdock import utils`)
+# - setting up PROJECT_ROOT environment variable
+# (which is used as a base for paths in "configs/paths/default.yaml")
+# (this way all filepaths are the same no matter where you run the code)
+# - loading environment variables from ".env" in root dir
+#
+# you can remove it if you:
+# 1. either install project as a package or move entry files to project root dir
+# 2. set `root_dir` to "." in "configs/paths/default.yaml"
+#
+# more info: https://github.com/ashleve/rootutils
+# ------------------------------------------------------------------------------------ #
+
+from flowdock.utils import RankedLogger
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+def create_parser():
+ """Create the `ArgumentParser` responsible for orchestrating ESM2 sequence embedding
+ generation."""
+ parser = argparse.ArgumentParser(
+ description="Extract per-token representations and model outputs for sequences in a FASTA file" # noqa
+ )
+
+ parser.add_argument(
+ "model_location",
+ type=str,
+ help="PyTorch model file OR name of pretrained model to download (see README for models)",
+ )
+ parser.add_argument(
+ "fasta_file",
+ type=pathlib.Path,
+ help="FASTA file on which to extract representations",
+ )
+ parser.add_argument(
+ "output_dir",
+ type=pathlib.Path,
+ help="output directory for extracted representations",
+ )
+
+ parser.add_argument("--toks_per_batch", type=int, default=4096, help="maximum batch size")
+ parser.add_argument(
+ "--repr_layers",
+ type=int,
+ default=[-1],
+ nargs="+",
+ help="layers indices from which to extract representations (0 to num_layers, inclusive)",
+ )
+ parser.add_argument(
+ "--include",
+ type=str,
+ nargs="+",
+ choices=["mean", "per_tok", "bos", "contacts"],
+ help="specify which representations to return",
+ required=True,
+ )
+ parser.add_argument(
+ "--truncation_seq_length",
+ type=int,
+ default=1022,
+ help="truncate sequences longer than the given value",
+ )
+ parser.add_argument(
+ "--skip-existing", help="Skip inference for existing embeddings", action="store_true"
+ )
+ parser.add_argument(
+ "--cuda_device_index",
+ type=int,
+ default=0,
+ help="if CUDA is available, onto which CUDA device to load all Tensors",
+ )
+
+ parser.add_argument("--nogpu", action="store_true", help="Do not use GPU even if available")
+ return parser
+
+
+def run(args: argparse.Namespace):
+ """Generate all ESM2 sequence embeddings for a given FASTA file input.
+
+ Args: An ArgParse `Namespace` object.
+ """
+ model, alphabet = pretrained.load_model_and_alphabet(args.model_location)
+ model.eval()
+ if isinstance(model, MSATransformer):
+ raise ValueError(
+ "This script currently does not handle models with MSA input (MSA Transformer)."
+ )
+ if torch.cuda.is_available() and not args.nogpu:
+ model = model.to(f"cuda:{args.cuda_device_index}")
+ log.info("Transferred model to GPU")
+
+ dataset = FastaBatchedDataset.from_file(args.fasta_file)
+ batches = dataset.get_batch_indices(args.toks_per_batch, extra_toks_per_seq=1)
+ data_loader = torch.utils.data.DataLoader(
+ dataset,
+ collate_fn=alphabet.get_batch_converter(args.truncation_seq_length),
+ batch_sampler=batches,
+ )
+ log.info(f"Read {args.fasta_file} with {len(dataset)} sequences")
+
+ args.output_dir.mkdir(parents=True, exist_ok=True)
+ return_contacts = "contacts" in args.include
+
+ assert all(
+ -(model.num_layers + 1) <= i <= model.num_layers for i in args.repr_layers
+ ), "Invalid layer index."
+ repr_layers = [(i + model.num_layers + 1) % (model.num_layers + 1) for i in args.repr_layers]
+
+ with torch.no_grad():
+ for batch_idx, (labels, strs, toks) in enumerate(data_loader):
+ if args.skip_existing and all(
+ os.path.exists(args.output_dir / f"{label}.pt") for label in labels
+ ):
+ continue
+ log.info(
+ f"Processing {batch_idx + 1} of {len(batches)} batches ({toks.size(0)} sequences)"
+ )
+ if torch.cuda.is_available() and not args.nogpu:
+ toks = toks.to(device=f"cuda:{args.cuda_device_index}", non_blocking=True)
+
+ out = model(toks, repr_layers=repr_layers, return_contacts=return_contacts)
+
+ representations = {
+ layer: t.to(device="cpu") for layer, t in out["representations"].items()
+ }
+ if return_contacts:
+ contacts = out["contacts"].to(device="cpu")
+
+ for i, label in enumerate(labels):
+ args.output_file = args.output_dir / f"{label}.pt"
+ args.output_file.parent.mkdir(parents=True, exist_ok=True)
+ result = {"label": label}
+ truncate_len = min(args.truncation_seq_length, len(strs[i]))
+ # Call clone on tensors to ensure tensors are not views into a larger representation
+ # See https://github.com/pytorch/pytorch/issues/1995
+ if "per_tok" in args.include:
+ result["representations"] = {
+ layer: t[i, 1 : truncate_len + 1].clone()
+ for layer, t in representations.items()
+ }
+ if "mean" in args.include:
+ result["mean_representations"] = {
+ layer: t[i, 1 : truncate_len + 1].mean(0).clone()
+ for layer, t in representations.items()
+ }
+ if "bos" in args.include:
+ result["bos_representations"] = {
+ layer: t[i, 0].clone() for layer, t in representations.items()
+ }
+ if return_contacts:
+ result["contacts"] = contacts[i, :truncate_len, :truncate_len].clone()
+
+ torch.save(
+ result,
+ args.output_file,
+ )
+
+
+def main():
+ """Generate all ESM2 sequence embeddings for a given FASTA file input."""
+ parser = create_parser()
+ args = parser.parse_args()
+ run(args)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/forks/FlowDock/flowdock/data/components/esm_embedding_preparation.py b/forks/FlowDock/flowdock/data/components/esm_embedding_preparation.py
new file mode 100644
index 00000000..cb269bbe
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/esm_embedding_preparation.py
@@ -0,0 +1,194 @@
+# Adapted from: https://github.com/gcorso/DiffDock
+
+import os
+from argparse import ArgumentParser
+
+import pandas as pd
+import rootutils
+from Bio import SeqIO
+from Bio.PDB import PDBParser
+from Bio.Seq import Seq
+from Bio.SeqRecord import SeqRecord
+from tqdm import tqdm
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock.utils import RankedLogger
+from flowdock.utils.data_utils import (
+ convert_protein_pts_to_pdb,
+ pdb_filepath_to_protein,
+)
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+parser = ArgumentParser()
+parser.add_argument("--out_file", type=str, default="data/pdbbind/pdbbind_sequences.fasta")
+parser.add_argument(
+ "--dataset",
+ type=str,
+ choices=["pdbbind", "moad", "dockgen", "pdbsidechain"],
+ default="pdbbind",
+)
+parser.add_argument(
+ "--data_dir",
+ type=str,
+ default="data/pdbbind/PDBBind_processed/",
+)
+parser.add_argument(
+ "--pdbbind_split_filepath",
+ type=str,
+ default=None,
+ help="Path to a CSV file containing the PDB codes to include in the PDBBind dataset.",
+)
+parser.add_argument(
+ "--pdbsidechain_metadata_csv_path",
+ type=str,
+ default="data/pdbsidechain/pdb_2021aug02/list.csv",
+)
+parser.add_argument(
+ "--skip_extracting_existing_pdbsidechain_pdb_files",
+ action="store_true",
+ help="If True, will avoid overwriting the PDB files already generated from the corresponding vdM PyTorch files.",
+)
+args = parser.parse_args()
+
+if not os.path.exists(args.data_dir):
+ raise ValueError(f"Provided data directory {args.data_dir} does not exist.")
+
+biopython_parser = PDBParser()
+
+data_dir = args.data_dir
+
+if args.dataset == "pdbbind":
+ sequences = []
+ ids = []
+
+ if args.pdbbind_split_filepath is not None and os.path.exists(args.pdbbind_split_filepath):
+ split_pdb_codes = set(pd.read_csv(args.pdbbind_split_filepath, header=None).iloc[:, 0])
+ names = [item for item in os.listdir(data_dir) if item in split_pdb_codes]
+ else:
+ names = os.listdir(data_dir)
+ for name in tqdm(names):
+ if name == ".DS_Store":
+ continue
+ if os.path.exists(os.path.join(data_dir, name, f"{name}_protein_processed.pdb")):
+ rec_pdb_path = os.path.join(data_dir, name, f"{name}_protein_processed.pdb")
+ else:
+ rec_pdb_path = os.path.join(data_dir, name, f"{name}_protein.pdb")
+ try:
+ protein = pdb_filepath_to_protein(rec_pdb_path)
+ except Exception as e:
+ log.warning(f"Error in parsing {rec_pdb_path} due to: {e}")
+ continue
+ for i, seq in enumerate(protein.letter_sequences):
+ sequences.append(seq[1])
+ ids.append(f"{name}_chain_{i}")
+ records = []
+ for index, seq in zip(ids, sequences):
+ record = SeqRecord(Seq(seq), str(index))
+ record.description = ""
+ records.append(record)
+ SeqIO.write(records, args.out_file, "fasta")
+
+elif args.dataset == "moad":
+ names = os.listdir(data_dir)
+ names = [n[:6] for n in names]
+ sequences = []
+ ids = []
+
+ for name in tqdm(names):
+ if name == ".DS_Store":
+ continue
+ if not os.path.exists(os.path.join(data_dir, f"{name}_protein.pdb")):
+ log.warning(f"We are skipping {name} because there was no {name}_protein.pdb")
+ continue
+ rec_pdb_path = os.path.join(data_dir, f"{name}_protein.pdb")
+ try:
+ protein = pdb_filepath_to_protein(rec_pdb_path)
+ except Exception as e:
+ log.warning(f"Error in parsing {rec_pdb_path} due to: {e}")
+ continue
+ for i, seq in enumerate(protein.letter_sequences):
+ sequences.append(seq[1])
+ ids.append(f"{name}_chain_{i}")
+ records = []
+ for index, seq in zip(ids, sequences):
+ record = SeqRecord(Seq(seq), str(index))
+ record.description = ""
+ records.append(record)
+ SeqIO.write(records, args.out_file, "fasta")
+
+elif args.dataset == "dockgen":
+ names = os.listdir(data_dir)
+ sequences = []
+ ids = []
+
+ for name in tqdm(names):
+ if name == ".DS_Store":
+ continue
+ if not os.path.exists(os.path.join(data_dir, name, f"{name}_protein_processed.pdb")):
+ log.warning(
+ f"We are skipping {name} because there was no {name}_protein_processed.pdb"
+ )
+ continue
+ rec_pdb_path = os.path.join(data_dir, name, f"{name}_protein_processed.pdb")
+ try:
+ protein = pdb_filepath_to_protein(rec_pdb_path)
+ except Exception as e:
+ log.warning(f"Error in parsing {rec_pdb_path} due to: {e}")
+ continue
+ for i, seq in enumerate(protein.letter_sequences):
+ sequences.append(seq[1])
+ ids.append(f"{name}_chain_{i}")
+ records = []
+ for index, seq in zip(ids, sequences):
+ record = SeqRecord(Seq(seq), str(index))
+ record.description = ""
+ records.append(record)
+ SeqIO.write(records, args.out_file, "fasta")
+
+elif args.dataset == "pdbsidechain":
+ metadata = pd.read_csv(args.pdbsidechain_metadata_csv_path)
+ sequences = []
+ ids = []
+ for row in tqdm(metadata.itertuples()):
+ name = row.CHAINID
+ chain_id = name.split("_")[1]
+ rec_pt_path = os.path.join(data_dir, name[1:3], f"{name}.pt")
+ rec_pdb_path = os.path.join(data_dir, name[1:3], f"{name}.pdb")
+ if not os.path.exists(rec_pt_path):
+ log.warning(f"We are skipping {name} because there were no files matching `{name}.pt`")
+ continue
+ if (
+ not os.path.exists(rec_pdb_path)
+ or not args.skip_extracting_existing_pdbsidechain_pdb_files
+ ):
+ try:
+ convert_protein_pts_to_pdb([rec_pt_path], rec_pdb_path)
+ if not os.path.exists(rec_pdb_path):
+ log.warning(f"Failed to convert PyTorch {rec_pt_path} to PDB {rec_pdb_path}")
+ continue
+
+ except Exception as e:
+ log.warning(f"Error in converting {rec_pt_path} to {rec_pdb_path} due to: {e}")
+ continue
+ try:
+ protein = pdb_filepath_to_protein(rec_pdb_path)
+ except Exception as e:
+ log.warning(f"Error in parsing {rec_pdb_path} due to: {e}")
+ continue
+ if not (
+ len(protein.letter_sequences) > 0
+ and len(protein.letter_sequences[0]) == 3
+ and protein.letter_sequences[0][1] is not None
+ ):
+ log.warning(f"We are skipping {name} because the extracted sequence was `None`")
+ continue
+ sequences.append(protein.letter_sequences[0][1])
+ ids.append(f"{name.split('_')[0]}_chain_{chain_id}")
+ records = []
+ for index, seq in zip(ids, sequences):
+ record = SeqRecord(Seq(seq), str(index))
+ record.description = ""
+ records.append(record)
+ SeqIO.write(records, args.out_file, "fasta")
diff --git a/forks/FlowDock/flowdock/data/components/esmfold_apo_to_holo_alignment.py b/forks/FlowDock/flowdock/data/components/esmfold_apo_to_holo_alignment.py
new file mode 100644
index 00000000..59ee32b9
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/esmfold_apo_to_holo_alignment.py
@@ -0,0 +1,503 @@
+import glob
+import os
+import warnings
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import hydra
+import numpy as np
+import rootutils
+from beartype import beartype
+from beartype.typing import List, Optional, Tuple, Union
+from Bio.PDB import PDBParser
+from Bio.PDB.Model import Model
+from Bio.PDB.PDBExceptions import PDBConstructionWarning
+from biopandas.pdb import PandasPdb
+from omegaconf import DictConfig, open_dict
+from rdkit import Chem
+from rdkit.Chem import AllChem, RemoveHs
+from scipy import spatial as spa
+from scipy.optimize import Bounds, minimize
+from scipy.spatial.transform import Rotation
+from tqdm import tqdm
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+# ------------------------------------------------------------------------------------ #
+# the setup_root above is equivalent to:
+# - adding project root dir to PYTHONPATH
+# (so you don't need to force user to install project as a package)
+# (necessary before importing any local modules e.g. `from flowdock import utils`)
+# - setting up PROJECT_ROOT environment variable
+# (which is used as a base for paths in "configs/paths/default.yaml")
+# (this way all filepaths are the same no matter where you run the code)
+# - loading environment variables from ".env" in root dir
+#
+# you can remove it if you:
+# 1. either install project as a package or move entry files to project root dir
+# 2. set `root_dir` to "." in "configs/paths/default.yaml"
+#
+# more info: https://github.com/ashleve/rootutils
+# ------------------------------------------------------------------------------------ #
+
+from flowdock import register_custom_omegaconf_resolvers
+from flowdock.utils import RankedLogger
+from flowdock.utils.data_utils import (
+ combine_molecules,
+ convert_protein_pts_to_pdb,
+ pdb_filepath_to_protein,
+)
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+biopython_parser = PDBParser()
+
+
+@beartype
+def read_molecule(
+ molecule_file: str, sanitize: bool = False, calc_charges: bool = False, remove_hs: bool = False
+) -> Optional[Chem.Mol]:
+ """Load an RDKit molecule from a given filepath."""
+ if molecule_file.endswith(".mol2"):
+ mol = Chem.MolFromMol2File(molecule_file, sanitize=False, removeHs=False)
+ elif molecule_file.endswith(".mol"):
+ mol = Chem.MolFromMolFile(molecule_file, sanitize=False, removeHs=False)
+ elif molecule_file.endswith(".sdf"):
+ supplier = Chem.SDMolSupplier(molecule_file, sanitize=False, removeHs=False)
+ mol = supplier[0]
+ elif molecule_file.endswith(".pdbqt"):
+ with open(molecule_file) as file:
+ pdbqt_data = file.readlines()
+ pdb_block = ""
+ for line in pdbqt_data:
+ pdb_block += f"{line[:66]}\n"
+ mol = Chem.MolFromPDBBlock(pdb_block, sanitize=False, removeHs=False)
+ elif molecule_file.endswith(".pdb"):
+ mol = Chem.MolFromPDBFile(molecule_file, sanitize=False, removeHs=False)
+ else:
+ raise ValueError(
+ "Expected the format of `molecule_file` to be "
+ f"one of `.mol`, `.mol2`, `.sdf`, `.pdbqt`, or `.pdb`, yet got {molecule_file}"
+ )
+
+ try:
+ if sanitize or calc_charges:
+ Chem.SanitizeMol(mol)
+
+ if calc_charges:
+ # compute Gasteiger charges on the molecule
+ try:
+ AllChem.ComputeGasteigerCharges(mol)
+ except Exception:
+ warnings.warn("Unable to compute charges for the input molecule.")
+
+ if remove_hs:
+ mol = Chem.RemoveHs(mol, sanitize=sanitize)
+ except Exception as e:
+ log.warning(f"RDKit was unable to read the molecule due to the error: {e}")
+ return None
+
+ return mol
+
+
+@beartype
+def read_mols(
+ dataset_dir: str,
+ name: str,
+ remove_hs: bool = False,
+ sanitize: bool = True,
+) -> List[Chem.Mol]:
+ """Load RDKit `Mol` objects corresponding to a dataset's ligand name."""
+ ligs = []
+ for file in os.listdir(os.path.join(dataset_dir, name)):
+ if file.endswith(".mol2") and "rdkit" not in file:
+ lig = read_molecule(
+ os.path.join(dataset_dir, name, file), remove_hs=remove_hs, sanitize=sanitize
+ )
+ if lig is None and os.path.exists(
+ os.path.join(dataset_dir, name, file[:-4] + ".sdf")
+ ): # read sdf file if mol2 file cannot be sanitized
+ log.info(
+ "Using the .mol2 file failed. We found a .sdf file instead and are trying to use that. Be aware that the .sdf files from PDBBind 2020 are missing chirality tags, although we will do our best to impute such information automatically using RDKit. Reference: https://www.blopig.com/blog/2021/09/watch-out-when-using-pdbbind."
+ )
+ lig = read_molecule(
+ os.path.join(dataset_dir, name, file[:-4] + ".sdf"),
+ remove_hs=remove_hs,
+ sanitize=sanitize,
+ )
+ Chem.rdmolops.AssignAtomChiralTagsFromStructure(lig)
+ if lig is not None:
+ ligs.append(lig)
+ return ligs
+
+
+@beartype
+def parse_pdb_from_path(path: str) -> Model:
+ """Load a BioPython structure from a given PDB filepath."""
+ with warnings.catch_warnings():
+ warnings.filterwarnings("ignore", category=PDBConstructionWarning)
+ structure = biopython_parser.get_structure("random_id", path)
+ rec = structure[0]
+ return rec
+
+
+@beartype
+def parse_and_read_moad_mols(
+ dataset_path: str,
+ pdb_id: str,
+) -> Optional[List[Optional[Chem.Mol]]]:
+ """Parse (to separate PDB files) and read ligands from a given PDB filepath, for the Binding
+ MOAD dataset.
+
+ :param dataset_path: Path to the directory containing protein-ligand complex PDB files.
+ :param pdb_id: PDB ID of the protein-ligand complex.
+ :return: Optional list of optional RDKit molecules.
+ """
+ lig_mols = []
+ lig_pdb_files = glob.glob(
+ os.path.join(dataset_path, "pdb_superligand", f"{pdb_id}_superlig_*.pdb")
+ )
+ if len(lig_pdb_files) > 0:
+ for lig_pdb_file in lig_pdb_files:
+ lig_mol = read_molecule(lig_pdb_file)
+ lig_mols.append(RemoveHs(lig_mol) if lig_mol is not None else None)
+ return lig_mols if len(lig_mols) else None
+
+
+@beartype
+def parse_and_read_dockgen_mols(
+ dataset_path: str,
+ pdb_id: str,
+) -> Optional[List[Optional[Chem.Mol]]]:
+ """Parse (to separate PDB files) and read ligands from a given PDB filepath, for the DockGen
+ dataset.
+
+ :param dataset_path: Path to the directory containing protein-ligand complex PDB files.
+ :param pdb_id: PDB ID of the protein-ligand complex.
+ :return: Optional list of optional RDKit molecules.
+ """
+ lig_mols = []
+ lig_pdb_files = glob.glob(os.path.join(dataset_path, pdb_id, f"{pdb_id}_ligand*.pdb"))
+ if len(lig_pdb_files) > 0:
+ for lig_pdb_file in lig_pdb_files:
+ lig_mol = read_molecule(lig_pdb_file)
+ lig_mols.append(RemoveHs(lig_mol, sanitize=False) if lig_mol is not None else None)
+ return lig_mols if len(lig_mols) else None
+
+
+@beartype
+def align_prediction(
+ smoothing_factor: Optional[Union[float, np.ndarray]],
+ dataset_calpha_coords: np.ndarray,
+ esmfold_calpha_coords: np.ndarray,
+ dataset_ligand_coords: Optional[np.ndarray],
+ return_rotation: bool = False,
+) -> Union[Tuple[Rotation, np.ndarray, np.ndarray], Union[float, np.ndarray]]:
+ """Perform an alignment of apo and holo protein structures and (optionally) ligand coordinates
+ using an (optionally) optimized smoothing factor.
+
+ :param smoothing_factor: Smoothing factor controlling the alignment.
+ :param dataset_calpha_coords: Array of Ca atom coordinates for a dataset's protein structure.
+ :param esmfold_calpha_coords: Array of Ca atom coordinates for a dataset's protein structure.
+ :param dataset_ligand_coords: Optional array of ligand coordinates from a dataset.
+ :param return_rotation: Whether to return the rotation matrix and centroids (default: `False`).
+ :return: If return_rotation is `True`, returns a tuple containing rotation matrix (`Rotation`), centroid of CA atoms for a dataset protein (`np.ndarray`),
+ and centroid of CA atoms for ESMFold (`np.ndarray`). If return_rotation is `False`, returns the inverse root mean square error of reciprocal distances (`float`).
+ """
+ if dataset_ligand_coords is not None:
+ dataset_dists = spa.distance.cdist(dataset_calpha_coords, dataset_ligand_coords)
+ weights = np.exp(-1 * smoothing_factor * np.amin(dataset_dists, axis=1))
+ dataset_calpha_centroid = np.sum(
+ np.expand_dims(weights, axis=1) * dataset_calpha_coords, axis=0
+ ) / np.sum(weights)
+ esmfold_calpha_centroid = np.sum(
+ np.expand_dims(weights, axis=1) * esmfold_calpha_coords, axis=0
+ ) / np.sum(weights)
+ else:
+ weights = None
+ dataset_calpha_centroid = np.mean(dataset_calpha_coords, axis=0)
+ esmfold_calpha_centroid = np.mean(esmfold_calpha_coords, axis=0)
+ centered_dataset_calpha_coords = dataset_calpha_coords - dataset_calpha_centroid
+ centered_esmfold_calpha_coords = esmfold_calpha_coords - esmfold_calpha_centroid
+
+ rotation, _ = spa.transform.Rotation.align_vectors(
+ centered_dataset_calpha_coords, centered_esmfold_calpha_coords, weights
+ )
+ if return_rotation:
+ return rotation, dataset_calpha_centroid, esmfold_calpha_centroid
+
+ if dataset_ligand_coords is not None:
+ centered_dataset_ligand_coords = dataset_ligand_coords - dataset_calpha_centroid
+ aligned_esmfold_calpha_coords = rotation.apply(centered_esmfold_calpha_coords)
+ aligned_esmfold_dataset_dists = spa.distance.cdist(
+ aligned_esmfold_calpha_coords, centered_dataset_ligand_coords
+ )
+ inv_r_rmse = np.sqrt(
+ np.mean(((1 / dataset_dists) - (1 / aligned_esmfold_dataset_dists)) ** 2)
+ )
+ else:
+ inv_r_rmse = np.nan
+ return inv_r_rmse
+
+
+@beartype
+def get_alignment_rotation(
+ pdb_id: str,
+ dataset_protein_path: str,
+ esmfold_protein_path: str,
+ dataset: str,
+ dataset_path: str,
+) -> Tuple[Optional[Rotation], Optional[np.ndarray], Optional[np.ndarray]]:
+ """Calculate the alignment rotation between apo and holo protein structures and their ligand
+ coordinates.
+
+ :param pdb_id: PDB ID of the protein-ligand complex.
+ :param dataset_protein_path: Filepath to the PDB file of the protein structure from a dataset.
+ :param esmfold_protein_path: Filepath to the PDB file of the protein structure from ESMFold.
+ :param dataset: Name of the dataset.
+ :param dataset_path: Filepath to the PDB file containing ligand coordinates.
+ :return: A tuple containing rotation matrix (Optional[Rotation]), centroid of Ca atoms for a
+ dataset protein (Optional[np.ndarray]), and centroid of Ca atoms for ESMFold
+ (Optional[np.ndarray]).
+ """
+ try:
+ dataset_protein = pdb_filepath_to_protein(dataset_protein_path)
+ dataset_calpha_coords = dataset_protein.atom_positions[:, 1, :]
+ except Exception as e:
+ log.warning(
+ f"Unable to parse dataset Ca protein structure for PDB ID {pdb_id} due to the error: {e}. Skipping..."
+ )
+ return None, None, None
+ try:
+ esmfold_protein = pdb_filepath_to_protein(esmfold_protein_path)
+ esmfold_calpha_coords = esmfold_protein.atom_positions[:, 1, :]
+ except Exception as e:
+ log.warning(
+ f"Unable to parse ESMFold Ca protein structure for PDB ID {pdb_id} due to the error: {e}. Skipping..."
+ )
+ return None, None, None
+
+ if dataset == "pdbbind":
+ try:
+ dataset_ligand = read_mols(dataset_path, pdb_id, remove_hs=True, sanitize=True)[0]
+ except Exception as e:
+ try:
+ log.warning(
+ f"Unable to parse PDBBind ligand structure for PDB ID {pdb_id} with sanitization. Trying to load without sanitization..."
+ )
+ dataset_ligand = read_mols(dataset_path, pdb_id, remove_hs=True, sanitize=False)[0]
+ except Exception as e:
+ log.warning(
+ f"Unable to parse PDBBind ligand structure for PDB ID {pdb_id} due to the error: {e}. Skipping..."
+ )
+ return None, None, None
+ elif dataset in ["moad", "dockgen"]:
+ try:
+ dataset_ligands = (
+ parse_and_read_dockgen_mols(
+ dataset_path=dataset_path,
+ pdb_id=pdb_id,
+ )
+ if dataset == "dockgen"
+ else parse_and_read_moad_mols(
+ dataset_path=dataset_path,
+ pdb_id=pdb_id,
+ )
+ )
+ except Exception as e:
+ log.warning(
+ f"Unable to parse {dataset} ligand structures for PDB ID {pdb_id} due to the error: {e}. Skipping..."
+ )
+ return None, None, None
+ if len(dataset_ligands) > 0:
+ try:
+ dataset_ligand = combine_molecules(dataset_ligands)
+ except Exception as e:
+ log.warning(
+ f"Unable to combine {dataset} ligand structures for PDB ID {pdb_id} due to the error: {e}. Skipping..."
+ )
+ return None, None, None
+ else:
+ log.warning(
+ f"Unable to parse any {dataset} ligand structures for PDB ID {pdb_id}. Skipping..."
+ )
+ return None, None, None
+ elif dataset == "pdbsidechain":
+ # NOTE: the van der Mers dataset does not contain artificial ligands initially
+ dataset_ligand = None
+ else:
+ raise ValueError(f"Dataset {dataset} is not supported.")
+
+ try:
+ if dataset_ligand is not None:
+ dataset_ligand_coords = dataset_ligand.GetConformer().GetPositions()
+ else:
+ dataset_ligand_coords = None
+ except Exception as e:
+ log.warning(
+ f"Unable to extract dataset ligand structure for PDB ID {pdb_id} due to the error: {e}. Skipping..."
+ )
+ return None, None, None
+
+ if dataset_calpha_coords.shape != esmfold_calpha_coords.shape:
+ log.warning(
+ f"Dataset and ESMFold protein structures differ for PDB ID {pdb_id}. Skipping due to shape mismatch: dataset {dataset_calpha_coords.shape} vs ESMFold {esmfold_calpha_coords.shape}",
+ )
+ return None, None, None
+
+ try:
+ if dataset_ligand_coords is not None:
+ res = minimize(
+ align_prediction,
+ [0.1],
+ bounds=Bounds([0.0], [1.0]),
+ args=(dataset_calpha_coords, esmfold_calpha_coords, dataset_ligand_coords),
+ tol=1e-8,
+ )
+ smoothing_factor = res.x
+ else:
+ smoothing_factor = 0.1
+
+ rotation, dataset_calpha_centroid, esmfold_calpha_centroid = align_prediction(
+ smoothing_factor,
+ dataset_calpha_coords,
+ esmfold_calpha_coords,
+ dataset_ligand_coords,
+ return_rotation=True,
+ )
+ except Exception as e:
+ log.warning(
+ f"Unable to align protein structures for PDB ID {pdb_id} due to the error: {e}. Skipping..."
+ )
+ return None, None, None
+
+ return rotation, dataset_calpha_centroid, esmfold_calpha_centroid
+
+
+@beartype
+def align_apo_structure_to_holo_structure(
+ cfg: DictConfig, filename: str, atom_df_name: str = "ATOM"
+):
+ """Align a given ESMFold apo structure to its corresponding holo structure.
+
+ :param cfg: Hydra config for the alignment.
+ :param filename: Filename of the ESMFold apo structure.
+ :param atom_df_name: Name of the atom DataFrame derived from the corresponding PDB file input.
+ """
+ pdb_id = Path(filename.split("_")[0]).stem if cfg.dataset == "pdbbind" else Path(filename).stem
+ esm_protein_filename = os.path.join(cfg.esmfold_structures_dir, f"{pdb_id}.pdb")
+ if cfg.dataset == "pdbsidechain":
+ processed_pt_filename = os.path.join(cfg.data_dir, pdb_id[1:3], f"{pdb_id}.pt")
+ if not os.path.exists(processed_pt_filename):
+ log.info(
+ f"Unable to find processed protein chain structures for PDB ID {pdb_id}. Skipping..."
+ )
+ return
+ processed_protein_filename = os.path.join(cfg.data_dir, pdb_id[1:3], f"{pdb_id}.pdb")
+ if not os.path.exists(processed_protein_filename):
+ try:
+ convert_protein_pts_to_pdb([processed_pt_filename], processed_protein_filename)
+ except Exception as e:
+ log.warning(
+ f"Unable to convert protein chain structures to PDB format for PDB ID {pdb_id} due to the error: {e}. Skipping..."
+ )
+ return
+ assert os.path.exists(
+ processed_protein_filename
+ ), f"Processed `pdbsidechain` file {processed_protein_filename} does not exist."
+ else:
+ processed_protein_name = (
+ f"{pdb_id}_protein_processed.pdb"
+ if cfg.dataset in ["pdbbind", "dockgen"]
+ else f"{pdb_id}_protein.pdb"
+ )
+ processed_protein_filename = os.path.join(
+ (os.path.join(cfg.data_dir, "pdb_protein") if cfg.dataset == "moad" else cfg.data_dir),
+ (pdb_id if cfg.dataset in ["pdbbind", "dockgen"] else ""),
+ processed_protein_name,
+ )
+ esm_protein_output_filename = os.path.join(
+ cfg.output_dir, f"{pdb_id}_holo_aligned_esmfold_protein.pdb"
+ )
+
+ rotation, dataset_calpha_centroid, esmfold_calpha_centroid = get_alignment_rotation(
+ pdb_id=pdb_id,
+ dataset_protein_path=processed_protein_filename,
+ esmfold_protein_path=esm_protein_filename,
+ dataset=cfg.dataset,
+ dataset_path=cfg.data_dir,
+ )
+
+ if any(
+ [item is None for item in [rotation, dataset_calpha_centroid, esmfold_calpha_centroid]]
+ ):
+ return
+
+ ppdb_esmfold = PandasPdb().read_pdb(esm_protein_filename)
+ ppdb_esmfold_pre_rot = (
+ ppdb_esmfold.df[atom_df_name][["x_coord", "y_coord", "z_coord"]]
+ .to_numpy()
+ .squeeze()
+ .astype(np.float32)
+ )
+ ppdb_esmfold_aligned = (
+ rotation.apply(ppdb_esmfold_pre_rot - esmfold_calpha_centroid) + dataset_calpha_centroid
+ )
+
+ ppdb_esmfold.df[atom_df_name][["x_coord", "y_coord", "z_coord"]] = ppdb_esmfold_aligned
+ ppdb_esmfold.to_pdb(path=esm_protein_output_filename, records=[atom_df_name], gz=False)
+
+
+@hydra.main(
+ version_base="1.3",
+ config_path="../../../configs/data",
+ config_name="esmfold_apo_to_holo_alignment.yaml",
+)
+def main(cfg: DictConfig):
+ """Align all ESMFold apo structures to their corresponding holo structures.
+
+ :param cfg: Hydra config for the alignments.
+ """
+ with open_dict(cfg):
+ if cfg.dataset == "pdbbind":
+ cfg.data_dir = os.path.join(cfg.data_dir, cfg.dataset, "PDBBind_processed")
+ elif cfg.dataset == "moad":
+ cfg.data_dir = os.path.join(cfg.data_dir, cfg.dataset, "BindingMOAD_2020_processed")
+ elif cfg.dataset == "dockgen":
+ cfg.data_dir = os.path.join(cfg.data_dir, "DockGen", "processed_files")
+ elif cfg.dataset == "pdbsidechain":
+ cfg.data_dir = os.path.join(cfg.data_dir, cfg.dataset, "pdb_2021aug02", "pdb")
+ else:
+ raise ValueError(f"Dataset {cfg.dataset} is not supported.")
+ output_dir = cfg.output_dir
+ os.makedirs(output_dir, exist_ok=True)
+ structure_file_inputs = [
+ file
+ for file in os.listdir(cfg.esmfold_structures_dir)
+ if not os.path.exists(
+ os.path.join(cfg.output_dir, f"{Path(file).stem}_holo_aligned_esmfold_protein.pdb")
+ )
+ ]
+ pbar = tqdm(
+ structure_file_inputs,
+ desc="Submitting apo-to-holo alignment tasks for parallel processing",
+ total=len(structure_file_inputs),
+ )
+
+ # process files in parallel
+ with ThreadPoolExecutor(max_workers=cfg.num_workers) as executor:
+ futures = []
+ for filename in pbar:
+ futures.append(executor.submit(align_apo_structure_to_holo_structure, cfg, filename))
+
+ # wait for all tasks to complete
+ for future in tqdm(
+ futures,
+ desc="Aligning each ESMFold apo structure to its corresponding holo structure",
+ total=len(futures),
+ ):
+ future.result()
+
+
+if __name__ == "__main__":
+ register_custom_omegaconf_resolvers()
+ main()
diff --git a/forks/FlowDock/flowdock/data/components/esmfold_apo_to_holo_assessment.py b/forks/FlowDock/flowdock/data/components/esmfold_apo_to_holo_assessment.py
new file mode 100644
index 00000000..6cfcd713
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/esmfold_apo_to_holo_assessment.py
@@ -0,0 +1,361 @@
+import glob
+import os
+import tempfile
+from pathlib import Path
+
+import hydra
+import numpy as np
+import pandas as pd
+import rootutils
+from beartype import beartype
+from beartype.typing import List
+from Bio.PDB import PDBParser
+from omegaconf import DictConfig, open_dict
+from rdkit import Chem
+from tqdm import tqdm
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock import register_custom_omegaconf_resolvers
+from flowdock.utils import RankedLogger
+from flowdock.utils.data_utils import load_hetero_data_graph_as_pdb_file_pair
+from flowdock.utils.metric_utils import calculate_usalign_metrics
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+@beartype
+def remove_duplicate_filenames(file_paths: List[str]) -> List[str]:
+ """Remove duplicate filenames (not filepaths) from a list of file paths.
+
+ :param file_paths: List of file paths.
+ :return: List of file paths with duplicate filenames removed.
+ """
+ filename_dict = {}
+ for path in file_paths:
+ filename = os.path.split(path)[1]
+ if filename not in filename_dict:
+ filename_dict[filename] = path
+ return list(filename_dict.values())
+
+
+@beartype
+def count_num_residues_in_pdb_file(pdb_filepath: str) -> int:
+ """Count the number of Ca atoms (i.e., residues) in a PDB file.
+
+ :param pdb_filepath: Path to PDB file.
+ :return: Number of Ca atoms (i.e., residues) in the PDB file.
+ """
+ parser = PDBParser()
+ structure = parser.get_structure("protein", pdb_filepath)
+ count = 0
+ for model in structure:
+ for chain in model:
+ for residue in chain:
+ if "CA" in residue:
+ count += 1
+ return count
+
+
+@hydra.main(
+ version_base="1.3",
+ config_path="../../../configs/data",
+ config_name="esmfold_apo_to_holo_assessment.yaml",
+)
+@beartype
+def main(cfg: DictConfig):
+ """Assessment the alignments between all ESMFold apo structures and their corresponding holo
+ structures while also collecting ligand statistics.
+
+ :param cfg: Hydra config for the assessment.
+ """
+ with open_dict(cfg):
+ if cfg.dataset == "pdbbind":
+ cfg.data_dir = os.path.join(cfg.data_dir, "pdbbind", "PDBBind_processed")
+ elif cfg.dataset == "moad":
+ cfg.data_dir = os.path.join(cfg.data_dir, "moad", "BindingMOAD_2020_processed")
+ elif cfg.dataset == "dockgen":
+ cfg.data_dir = os.path.join(cfg.data_dir, "DockGen", "processed_files")
+ elif cfg.dataset == "pdbsidechain":
+ cfg.data_dir = os.path.join(cfg.data_dir, "pdbsidechain", "pdb_2021aug02", "pdb")
+ else:
+ raise ValueError(f"Dataset {cfg.dataset} is not supported.")
+ assert os.path.exists(cfg.data_dir), f"Data directory {cfg.data_dir} does not exist."
+ use_preprocessed_graphs = cfg.preprocessed_graph_dirs is not None and all(
+ os.path.exists(dir) for dir in cfg.preprocessed_graph_dirs
+ )
+ if use_preprocessed_graphs:
+ structure_file_inputs = []
+ for dir in cfg.preprocessed_graph_dirs:
+ heterograph_filepaths = sorted(glob.glob(os.path.join(dir, "heterographs*.pkl")))
+ rdkit_ligand_filepaths = sorted(glob.glob(os.path.join(dir, "rdkit_ligands*.pkl")))
+ for heterograph_filepath, rdkit_ligand_filepath in zip(
+ heterograph_filepaths, rdkit_ligand_filepaths
+ ):
+ assert (
+ Path(heterograph_filepath).stem.split("heterographs")[1]
+ == Path(rdkit_ligand_filepath).stem.split("rdkit_ligands")[1]
+ ), f"File indices do not match for protein-ligand file pair: {heterograph_filepath} and {rdkit_ligand_filepath}"
+ structure_file_inputs.append(((heterograph_filepath, ""), [rdkit_ligand_filepath]))
+ temp_dir_path = Path(tempfile.mkdtemp(prefix="apo_holo_alignment_assessment_pdbs_"))
+ else:
+ if cfg.dataset == "pdbbind":
+ structure_file_inputs = [
+ (
+ (
+ os.path.join(cfg.aligned_dir, f"{id}_holo_aligned_esmfold_protein.pdb"),
+ os.path.join(cfg.data_dir, id, f"{id}_protein_processed.pdb"),
+ ),
+ [os.path.join(cfg.data_dir, id, f"{id}_ligand.sdf")],
+ )
+ for id in os.listdir(cfg.data_dir)
+ if os.path.exists(
+ os.path.join(cfg.aligned_dir, f"{id}_holo_aligned_esmfold_protein.pdb")
+ )
+ and (
+ os.path.exists(os.path.join(cfg.data_dir, id, f"{id}_ligand.sdf"))
+ or os.path.exists(os.path.join(cfg.data_dir, id, f"{id}_ligand.mol2"))
+ )
+ ]
+ elif cfg.dataset == "moad":
+ structure_file_inputs = []
+ for filename in os.listdir(Path(cfg.data_dir) / "pdb_protein"):
+ ligand_file_inputs = [
+ os.path.join(
+ cfg.data_dir,
+ "pdb_superligand",
+ f"{Path(filename).stem[:6]}_superlig_{i}.pdb",
+ )
+ for i in range(10)
+ if os.path.exists(
+ os.path.join(
+ cfg.data_dir,
+ "pdb_superligand",
+ f"{Path(filename).stem[:6]}_superlig_{i}.pdb",
+ )
+ )
+ ]
+ if os.path.exists(
+ os.path.join(
+ cfg.aligned_dir,
+ f"{Path(filename).stem[:6]}_holo_aligned_esmfold_protein.pdb",
+ )
+ ) and len(ligand_file_inputs):
+ structure_file_inputs.append(
+ (
+ (
+ os.path.join(
+ cfg.aligned_dir,
+ f"{Path(filename).stem[:6]}_holo_aligned_esmfold_protein.pdb",
+ ),
+ os.path.join(
+ cfg.data_dir,
+ "pdb_protein",
+ f"{Path(filename).stem[:6]}_protein.pdb",
+ ),
+ ),
+ ligand_file_inputs,
+ )
+ )
+ elif cfg.dataset == "dockgen":
+ structure_file_inputs = []
+ for item in os.listdir(cfg.data_dir):
+ ligand_file_inputs = list(
+ glob.glob(
+ os.path.join(
+ cfg.data_dir,
+ item,
+ f"{item}_ligand*.pdb",
+ )
+ )
+ )
+ if os.path.exists(
+ os.path.join(
+ cfg.aligned_dir,
+ f"{item}_holo_aligned_esmfold_protein.pdb",
+ )
+ ) and len(ligand_file_inputs):
+ structure_file_inputs.append(
+ (
+ (
+ os.path.join(
+ cfg.aligned_dir,
+ f"{item}_holo_aligned_esmfold_protein.pdb",
+ ),
+ os.path.join(
+ cfg.data_dir,
+ item,
+ f"{item}_protein_processed.pdb",
+ ),
+ ),
+ ligand_file_inputs,
+ )
+ )
+ elif cfg.dataset == "pdbsidechain":
+ structure_file_inputs = []
+ for id in os.listdir(cfg.data_dir):
+ for filename in os.listdir(os.path.join(cfg.data_dir, id)):
+ full_id = "_".join(Path(filename).stem.split("_")[:2])
+ if filename.endswith(".pdb") and os.path.exists(
+ os.path.join(
+ cfg.aligned_dir, f"{full_id}_holo_aligned_esmfold_protein.pdb"
+ )
+ ):
+ structure_file_inputs.append(
+ (
+ (
+ os.path.join(
+ cfg.aligned_dir,
+ f"{full_id}_holo_aligned_esmfold_protein.pdb",
+ ),
+ os.path.join(cfg.data_dir, id, filename),
+ ),
+ [],
+ )
+ )
+ flags = [
+ "-mol",
+ "prot", # note: align only protein chains
+ "-ter",
+ "0", # note: biological unit alignment
+ "-split",
+ "0", # note: treat a whole file as a single chain
+ "-se", # note: calculate TMscore without alignment
+ ]
+ (
+ valid_file_inputs,
+ erroneous_file_inputs,
+ tm_scores,
+ rmsds,
+ apo_lengths,
+ holo_lengths,
+ ligands_num_atoms,
+ ) = (
+ [],
+ [],
+ [],
+ [],
+ [],
+ [],
+ [],
+ )
+ for file_input in tqdm(
+ structure_file_inputs,
+ desc="Calculating structural metrics (as well as ligand statistics) for holo-aligned apo and holo structure pairs",
+ ):
+ try:
+ flat_file_input = (*file_input[0], file_input[1])
+ if use_preprocessed_graphs:
+ pair_dict = load_hetero_data_graph_as_pdb_file_pair(
+ (flat_file_input[0], flat_file_input[2]), temp_dir_path=temp_dir_path
+ )
+ pair = (
+ pair_dict["apo_protein"]["filepath"],
+ pair_dict["holo_protein"]["filepath"],
+ )
+ apo_length, holo_length = (
+ pair_dict["apo_protein"]["length"],
+ pair_dict["holo_protein"]["length"],
+ )
+ ligand_num_atoms = pair_dict["ligand"]["num_atoms_per_mol_frag"]
+ else:
+ pair, rdkit_ligand_filepaths = file_input
+ apo_length, holo_length = (
+ count_num_residues_in_pdb_file(pair[0]),
+ count_num_residues_in_pdb_file(pair[1]),
+ )
+ lig_mol_frags = []
+ for rdkit_ligand_filepath in rdkit_ligand_filepaths:
+ ligand = None
+ try:
+ ligand = (
+ Chem.MolFromPDBFile(rdkit_ligand_filepath, sanitize=False)
+ if rdkit_ligand_filepath.endswith(".pdb")
+ else Chem.MolFromMolFile(rdkit_ligand_filepath, sanitize=False)
+ )
+ except Exception as e:
+ ligand = None
+ if ligand is None:
+ try:
+ ligand = Chem.MolFromMol2File(
+ rdkit_ligand_filepath.replace(".sdf", ".mol2"), sanitize=False
+ )
+ except Exception as e:
+ ligand = None
+ if ligand is None:
+ raise ValueError(f"Could not load ligand from {rdkit_ligand_filepath}")
+ lig_mol_frags.extend(
+ Chem.GetMolFrags(ligand, asMols=True, sanitizeFrags=False)
+ )
+ ligand_num_atoms = [lig.GetNumAtoms() for lig in lig_mol_frags]
+ tm_score_metrics = calculate_usalign_metrics(*pair, cfg.usalign_exec_path, flags=flags)
+ except Exception as e:
+ log.error(f"Error calculating structural metrics for {file_input} due to: {e}")
+ erroneous_file_inputs.append(flat_file_input)
+ continue
+ if "TM-score_2" not in tm_score_metrics or "RMSD" not in tm_score_metrics:
+ log.error(f"Error calculating structural metrics for {file_input}. Skipping...")
+ erroneous_file_inputs.append(flat_file_input)
+ continue
+ valid_file_inputs.append(flat_file_input)
+ tm_scores.append(tm_score_metrics["TM-score_2"])
+ rmsds.append(tm_score_metrics["RMSD"])
+ apo_lengths.append(apo_length)
+ holo_lengths.append(holo_length)
+ ligands_num_atoms.append(ligand_num_atoms)
+ tm_scores = np.array(tm_scores)
+ rmsds = np.array(rmsds)
+ apo_lengths = np.array(apo_lengths)
+ holo_lengths = np.array(holo_lengths)
+ ligands_num_atoms_array = np.array(
+ [np.array(num_atoms).mean() for num_atoms in ligands_num_atoms]
+ )
+ log.info(f"TM-score mean: {tm_scores.mean():.3f} +/- {tm_scores.std():.3f}")
+ # ≈ 0.862 +/- 0.200 protein TMscore for PDBBind when assessing with raw input files
+ # ≈ 0.803 +/- 0.257 protein TMscore for Binding MOAD when assessing with raw input files
+ # ≈ 0.736 +/- 0.294 protein TMscore for DockGen when assessing with raw input files
+ # ≈ 0.824 +/- 0.239 protein TMscore for PDBSidechain when assessing with raw input files
+ log.info(f"RMSD mean: {rmsds.mean():.3f} +/- {rmsds.std():.3f}")
+ # ≈ 2.165 +/- 1.830 protein RMSD for PDBBind when assessing with raw input files
+ # ≈ 2.203 +/- 1.959 protein RMSD for Binding MOAD when assessing with raw input files
+ # ≈ 3.202 +/- 2.683 protein RMSD for DockGen when assessing with raw input files
+ # ≈ 2.269 +/- 1.906 protein RMSD for PDBSidechain when assessing with raw input files
+ log.info(f"Apo length mean: {apo_lengths.mean():.3f} +/- {apo_lengths.std():.3f}")
+ # ≈ 316.076 +/- 149.825 apo residues for PDBBind when assessing with raw input files
+ # ≈ 339.254 +/- 153.889 apo residues for Binding MOAD when assessing with raw input files
+ # ≈ 462.712 +/- 300.881 apo residues for DockGen when assessing with raw input files
+ # ≈ 249.091 +/- 145.546 apo residues for PDBSidechain when assessing with raw input files
+ log.info(f"Holo length mean: {holo_lengths.mean():.3f} +/- {holo_lengths.std():.3f}")
+ # ≈ 316.699 +/- 149.910 holo residues for PDBBind when assessing with raw input files
+ # ≈ 340.776 +/- 202.198 holo residues for Binding MOAD when assessing with raw input files
+ # ≈ 463.791 +/- 301.404 holo residues for DockGen when assessing with raw input files
+ # ≈ 249.235 +/- 145.551 holo residues for PDBSidechain when assessing with raw input files
+ log.info(
+ f"Ligand num atoms: {ligands_num_atoms_array.mean():.3f} +/- {ligands_num_atoms_array.std():.3f}"
+ )
+ # ≈ 62.062 +/- 45.614 ligand atoms for PDBBind when assessing with raw input files
+ # ≈ 23.939 +/- 15.876 ligand atoms for Binding MOAD when assessing with raw input files
+ # ≈ 26.097 +/- 15.048 ligand atoms for DockGen when assessing with raw input files
+ # ≈ nan +/- nan ligand atoms for PDBSidechain when assessing with raw input files
+ pd.DataFrame(
+ {
+ "Filepath": valid_file_inputs,
+ "TM-score": tm_scores,
+ "RMSD": rmsds,
+ "Apo_Length": apo_lengths,
+ "Holo_Length": holo_lengths,
+ "Ligand_Num_Atoms": [
+ ",".join([str(num) for num in num_atoms]) for num_atoms in ligands_num_atoms
+ ],
+ }
+ ).to_csv(cfg.output_filepath, index=False)
+ log.info(f"Saved TM-score and RMSD results to {cfg.output_filepath}")
+ pd.DataFrame({"Filepath": erroneous_file_inputs}).to_csv(
+ cfg.erroneous_output_filepath, index=False
+ )
+ log.info(f"Saved erroneous file inputs to {cfg.erroneous_output_filepath}")
+
+
+if __name__ == "__main__":
+ register_custom_omegaconf_resolvers()
+ main()
diff --git a/forks/FlowDock/flowdock/data/components/esmfold_batch_structure_prediction.py b/forks/FlowDock/flowdock/data/components/esmfold_batch_structure_prediction.py
new file mode 100644
index 00000000..164368fa
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/esmfold_batch_structure_prediction.py
@@ -0,0 +1,264 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import argparse
+import logging
+import sys
+from pathlib import Path
+from timeit import default_timer as timer
+
+import beartype.typing as T
+import esm
+import torch
+from beartype import beartype
+from esm.data import read_fasta
+
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+formatter = logging.Formatter(
+ "%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+ datefmt="%y/%m/%d %H:%M:%S",
+)
+
+console_handler = logging.StreamHandler(sys.stdout)
+console_handler.setLevel(logging.INFO)
+console_handler.setFormatter(formatter)
+logger.addHandler(console_handler)
+
+
+PathLike = T.Union[str, Path]
+
+
+@beartype
+def enable_cpu_offloading(model: torch.nn.Module) -> torch.nn.Module:
+ """Enable CPU offloading for a PyTorch `Module`.
+
+ :param model: The `Module` for which to enable CPU offloading.
+ :return: A `Module` with CPU offloading enabled.
+ """
+ from torch.distributed.fsdp import CPUOffload, FullyShardedDataParallel
+ from torch.distributed.fsdp.wrap import enable_wrap, wrap
+
+ torch.distributed.init_process_group(
+ backend="nccl", init_method="tcp://localhost:9999", world_size=1, rank=0
+ )
+
+ wrapper_kwargs = dict(cpu_offload=CPUOffload(offload_params=True))
+
+ with enable_wrap(wrapper_cls=FullyShardedDataParallel, **wrapper_kwargs):
+ for layer_name, layer in model.layers.named_children():
+ wrapped_layer = wrap(layer)
+ setattr(model.layers, layer_name, wrapped_layer)
+ model = wrap(model)
+
+ return model
+
+
+@beartype
+def init_model_on_gpu_with_cpu_offloading(
+ model: torch.nn.Module, cuda_device_index: int = 0
+) -> torch.nn.Module:
+ """Initialize a PyTorch `Module` on a GPU with CPU offloading support.
+
+ :param model: The `Module` for which to enable GPU-based CPU offloading.
+ :param cuda_device_index: The index of the GPU to use.
+ :return: A `Module` with GPU-based CPU offloading enabled.
+ """
+ model = model.eval()
+ model_esm = enable_cpu_offloading(model.esm)
+ del model.esm
+ model.to(f"cuda:{cuda_device_index}")
+ model.esm = model_esm
+ return model
+
+
+@beartype
+def create_batched_sequence_datasest(
+ sequences: T.List[T.Tuple[str, str]], max_tokens_per_batch: int = 1024
+) -> T.Generator[T.Tuple[T.List[str], T.List[str]], None, None]:
+ """Create a dataset of batched sequences for increased computational throughput.
+
+ :param sequences: A list of `(header, seq)` tuples.
+ :param max_tokens_per_batch: The maximum number of tokens (i.e., residues) per batch.
+ :return: A `Generator` of batched headers and batched sequences, respectively.
+ """
+ batch_headers, batch_sequences, num_tokens = [], [], 0
+ for header, seq in sequences:
+ if (len(seq) + num_tokens > max_tokens_per_batch) and num_tokens > 0:
+ yield batch_headers, batch_sequences
+ batch_headers, batch_sequences, num_tokens = [], [], 0
+ batch_headers.append(header)
+ batch_sequences.append(seq)
+ num_tokens += len(seq)
+
+ yield batch_headers, batch_sequences
+
+
+def create_parser():
+ """Create an `ArgumentParser` for the script."""
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "-i",
+ "--fasta",
+ help="Path to input FASTA file",
+ type=Path,
+ required=True,
+ )
+ parser.add_argument(
+ "-o", "--pdb", help="Path to output PDB directory", type=Path, required=True
+ )
+ parser.add_argument(
+ "-m",
+ "--model-dir",
+ help="Parent path to Pretrained ESM data directory. ",
+ type=Path,
+ default=None,
+ )
+ parser.add_argument(
+ "--num-recycles",
+ type=int,
+ default=None,
+ help="Number of recycles to run. Defaults to number used in training (4).",
+ )
+ parser.add_argument(
+ "--max-tokens-per-batch",
+ type=int,
+ default=1024,
+ help="Maximum number of tokens per gpu forward-pass. This will group shorter sequences together "
+ "for batched prediction. Lowering this can help with out of memory issues, if these occur on "
+ "short sequences.",
+ )
+ parser.add_argument(
+ "--chunk-size",
+ type=int,
+ default=None,
+ help="Chunks axial attention computation to reduce memory usage from O(L^2) to O(L). "
+ "Equivalent to running a for loop over chunks of of each dimension. Lower values will "
+ "result in lower memory usage at the cost of speed. Recommended values: 128, 64, 32. "
+ "Default: None.",
+ )
+ parser.add_argument("--cpu-only", help="CPU only", action="store_true")
+ parser.add_argument("--cpu-offload", help="Enable CPU offloading", action="store_true")
+ parser.add_argument(
+ "--skip-existing", help="Skip predictions for existing structures", action="store_true"
+ )
+ parser.add_argument(
+ "--predict-only-unit1",
+ help="For multi-model datasets such as Binding MOAD, skip predictions for sequences not belonging to the `unit1` version of each complex",
+ action="store_true",
+ )
+ parser.add_argument(
+ "-d", "--cuda-device-index", help="Index of CUDA device to use", type=int, default=0
+ )
+ return parser
+
+
+@beartype
+def run(args: argparse.Namespace):
+ """Run the script with `argparse` arguments.
+
+ :param args: An `argparse` `Namespace` containing script arguments.
+ """
+ if not args.fasta.exists():
+ raise FileNotFoundError(args.fasta)
+
+ args.pdb.mkdir(exist_ok=True)
+
+ # Read fasta and sort sequences by length
+ logger.info(f"Reading sequences from {args.fasta}")
+ all_sequences = sorted(read_fasta(args.fasta), key=lambda header_seq: len(header_seq[1]))
+ all_sequences = [(header, seq) for header, seq in all_sequences if list(set(seq)) != ["-"]]
+ if args.skip_existing:
+ all_sequences = [
+ (header, seq)
+ for header, seq in all_sequences
+ if not (args.pdb / f"{header}.pdb").exists()
+ and (header.endswith("_unit1") or not args.predict_only_unit1)
+ ]
+ logger.info(f"Loaded {len(all_sequences)} sequences from {args.fasta}")
+
+ logger.info("Loading model")
+
+ # Use pre-downloaded ESM weights from model_pth.
+ if args.model_dir is not None:
+ # if pretrained model path is available
+ torch.hub.set_dir(args.model_dir)
+
+ model = esm.pretrained.esmfold_v1()
+
+ model = model.eval()
+ model.set_chunk_size(args.chunk_size)
+
+ if args.cpu_only:
+ model.esm.float() # convert to fp32 as ESM-2 in fp16 is not supported on CPU
+ model.cpu()
+ elif args.cpu_offload:
+ model = init_model_on_gpu_with_cpu_offloading(
+ model, cuda_device_index=args.cuda_device_index
+ )
+ else:
+ model.to(f"cuda:{args.cuda_device_index}")
+ logger.info("Starting Predictions")
+ batched_sequences = create_batched_sequence_datasest(all_sequences, args.max_tokens_per_batch)
+
+ num_completed = 0
+ num_sequences = len(all_sequences)
+ for headers, sequences in batched_sequences:
+ start = timer()
+ try:
+ output = model.infer(sequences, num_recycles=args.num_recycles)
+ except RuntimeError as e:
+ if e.args[0].startswith("CUDA out of memory"):
+ if len(sequences) > 1:
+ logger.info(
+ f"Failed (CUDA out of memory) to predict batch of size {len(sequences)}. "
+ "Try lowering `--max-tokens-per-batch`."
+ )
+ else:
+ logger.info(
+ f"Failed (CUDA out of memory) on sequence {headers[0]} of length {len(sequences[0])}."
+ )
+ continue
+ else:
+ logger.error(f"Failed to predict sequences {headers} due to unexpected error: {e}")
+ raise
+
+ try:
+ output = {key: value.cpu() for key, value in output.items()}
+ pdbs = model.output_to_pdb(output)
+ except Exception as e:
+ logger.info(
+ f"Failed to predict sequence {headers[0]} of length {len(sequences[0])} due to an unexpected error: {e}. "
+ "This is likely due to a sequence that is too short or contains only gaps."
+ )
+ continue
+ tottime = timer() - start
+ time_string = f"{tottime / len(headers):0.1f}s"
+ if len(sequences) > 1:
+ time_string = time_string + f" (amortized, batch size {len(sequences)})"
+ for header, seq, pdb_string, mean_plddt, ptm in zip(
+ headers, sequences, pdbs, output["mean_plddt"], output["ptm"]
+ ):
+ output_file = args.pdb / f"{header}.pdb"
+ output_file.write_text(pdb_string)
+ num_completed += 1
+ logger.info(
+ f"Predicted structure for {header} with length {len(seq)}, pLDDT {mean_plddt:0.1f}, "
+ f"pTM {ptm:0.3f} in {time_string}. "
+ f"{num_completed} / {num_sequences} completed."
+ )
+
+
+def main():
+ """Run the script with arguments."""
+ parser = create_parser()
+ args = parser.parse_args()
+ run(args)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/forks/FlowDock/flowdock/data/components/esmfold_sequence_preparation.py b/forks/FlowDock/flowdock/data/components/esmfold_sequence_preparation.py
new file mode 100644
index 00000000..3c79b3fc
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/esmfold_sequence_preparation.py
@@ -0,0 +1,93 @@
+import hydra
+import rootutils
+from beartype.typing import Dict, Literal
+from omegaconf import DictConfig
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock import register_custom_omegaconf_resolvers
+from flowdock.utils import RankedLogger
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+def read_fasta(filename: str) -> Dict[str, str]:
+ """Read sequences from a FASTA file and return them as a dictionary.
+
+ :param filename: Path to the input FASTA file.
+ :return: Dictionary containing sequences with identifiers as keys.
+ """
+ sequences = {}
+ current_id = ""
+ with open(filename) as file:
+ for line in file:
+ if line.startswith(">"):
+ current_id = line.strip()[1:]
+ sequences[current_id] = ""
+ else:
+ sequences[current_id] += line.strip()
+ return sequences
+
+
+def combine_sequences(
+ sequences: Dict[str, str], dataset: Literal["pdbbind", "moad", "dockgen", "pdbsidechain"]
+) -> Dict[str, str]:
+ """Combine sequences of the same protein complex by their PDB codes.
+
+ :param sequences: Dictionary containing sequences with identifiers as keys.
+ :param dataset: Name of the dataset for which to combine sequences.
+ :return: Dictionary containing combined sequences with PDB codes as keys.
+ """
+ combined_sequences = {}
+ for identifier, sequence in sequences.items():
+ if sequence and list(set(sequence)) not in [[""], ["-"]]:
+ if dataset == "pdbbind":
+ pdb_code = identifier.split("_")[0]
+ elif dataset == "moad":
+ pdb_code = "_".join(identifier.split("_")[0:2])
+ elif dataset == "dockgen":
+ pdb_code = identifier.split("_chain_")[0]
+ elif dataset == "pdbsidechain":
+ # NOTE: the van der Mers dataset selects only a single chain from each PDB complex structure
+ pdb_code = identifier.split("_")[0] + "_" + identifier.split("_")[-1]
+ else:
+ raise ValueError(f"Invalid dataset: {dataset}")
+ if pdb_code in combined_sequences:
+ combined_sequences[pdb_code] += ":" + sequence
+ else:
+ combined_sequences[pdb_code] = sequence
+ return combined_sequences
+
+
+def write_combined_fasta(combined_sequences: Dict[str, str], output_filename: str) -> None:
+ """Write combined sequences to an output FASTA file.
+
+ :param combined_sequences: Dictionary containing combined sequences with PDB codes as keys.
+ :param output_filename: Path to the output FASTA file.
+ """
+ with open(output_filename, "w") as file:
+ for pdb_code, sequence in combined_sequences.items():
+ file.write(f">{pdb_code}\n{sequence}\n")
+
+
+@hydra.main(
+ version_base="1.3",
+ config_path="../../../configs/data",
+ config_name="esmfold_sequence_preparation.yaml",
+)
+def main(cfg: DictConfig):
+ """Read protein chain sequences from a FASTA file, combine sequences of the same protein
+ complex, and write the combined sequences to an output FASTA file.
+
+ :param cfg: Configuration dictionary from the hydra YAML file.
+ """
+ sequences = read_fasta(cfg.input_fasta_file)
+ combined_sequences = combine_sequences(sequences, cfg.dataset)
+ write_combined_fasta(combined_sequences, cfg.output_fasta_file)
+
+ log.info("ESMFold FASTA file preparation complete.")
+
+
+if __name__ == "__main__":
+ register_custom_omegaconf_resolvers()
+ main()
diff --git a/forks/FlowDock/flowdock/data/components/moad_dataset.py b/forks/FlowDock/flowdock/data/components/moad_dataset.py
new file mode 100644
index 00000000..29486b84
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/moad_dataset.py
@@ -0,0 +1,914 @@
+# Adapted from: https://github.com/gcorso/DiffDock
+
+import ast
+import copy
+import glob
+import os
+import pickle # nosec
+import random
+from multiprocessing import Pool
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import rootutils
+import torch
+from beartype.typing import Optional
+from prody import confProDy
+from rdkit import Chem
+from scipy.spatial import distance
+from torch_geometric.data import Dataset
+from tqdm import tqdm
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock.data.components.mol_features import process_molecule
+from flowdock.utils import RankedLogger
+from flowdock.utils.data_utils import (
+ centralize_complex_graph,
+ combine_molecules,
+ merge_protein_and_ligands,
+ pdb_filepath_to_protein,
+ process_protein,
+)
+from flowdock.utils.model_utils import sample_inplace_to_torch
+from flowdock.utils.utils import fasta_to_dict, read_strings_from_txt
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+confProDy(verbosity="none")
+
+
+class BindingMOADDataset(Dataset):
+ """PyTorch Geometric Dataset for BindingMOAD dataset."""
+
+ def __init__(
+ self,
+ root,
+ dockgen_dir,
+ clusters_filepath,
+ transform=None,
+ cache_path=os.path.join("data", "cache"),
+ split="train",
+ limit_complexes=0,
+ num_workers=0,
+ max_lig_size=None,
+ min_multi_lig_distance=None,
+ remove_hs=False,
+ esm_embeddings_path=None,
+ dockgen_esm_embeddings_path=None,
+ esm_embeddings_sequences_path=None,
+ dockgen_esm_embeddings_sequences_path=None,
+ apo_protein_structure_dir=None,
+ dockgen_apo_protein_structure_dir=None,
+ require_ligand=False,
+ include_miscellaneous_atoms=False,
+ min_ligand_size=0,
+ multiplicity=1,
+ max_receptor_size=None,
+ remove_promiscuous_targets=None,
+ unroll_clusters=False,
+ remove_pdbbind=False,
+ enforce_timesplit=False,
+ no_randomness=False,
+ single_cluster_name=None,
+ total_dataset_size=None,
+ pdbbind_split_train=None,
+ pdbbind_split_val=None,
+ split_time=None,
+ min_protein_length: Optional[int] = 10,
+ max_protein_length: Optional[int] = 4000,
+ is_test_dataset=False,
+ a2h_assessment_csv_filepath=None,
+ dockgen_a2h_assessment_csv_filepath=None,
+ filter_using_a2h_assessment=False,
+ a2h_min_tmscore=None,
+ a2h_max_rmsd=None,
+ a2h_min_protein_length=None,
+ a2h_max_protein_length=None,
+ a2h_min_ligand_length=None,
+ a2h_max_ligand_length=None,
+ binding_affinity_values_dict=None,
+ n_lig_patches=32,
+ ):
+ """Initializes the BindingMOADDataset."""
+ if dockgen_a2h_assessment_csv_filepath is not None:
+ assert (
+ split != "train"
+ ), "DockGen `a2h` filtering is only for validation and test datasets."
+
+ super().__init__(root, transform)
+ self.moad_dir = root
+ self.dockgen_dir = dockgen_dir
+ self.include_miscellaneous_atoms = include_miscellaneous_atoms
+ self.max_lig_size = max_lig_size
+ self.min_multi_lig_distance = min_multi_lig_distance
+ self.split = split
+ self.limit_complexes = limit_complexes
+ self.num_workers = num_workers
+ self.remove_hs = remove_hs
+ self.require_ligand = require_ligand
+ self.esm_embeddings_path = esm_embeddings_path
+ self.dockgen_esm_embeddings_path = dockgen_esm_embeddings_path
+ self.esm_embeddings_sequences_path = esm_embeddings_sequences_path
+ self.dockgen_esm_embeddings_sequences_path = dockgen_esm_embeddings_sequences_path
+ self.apo_protein_structure_dir = apo_protein_structure_dir
+ self.dockgen_apo_protein_structure_dir = dockgen_apo_protein_structure_dir
+ self.multiplicity = multiplicity
+ self.no_randomness = no_randomness
+ self.total_dataset_size = total_dataset_size
+ self.min_protein_length = min_protein_length
+ self.max_protein_length = max_protein_length
+ self.is_test_dataset = is_test_dataset
+ self.binding_affinity_values_dict = binding_affinity_values_dict
+ self.n_lig_patches = n_lig_patches
+
+ self.prot_cache_path = os.path.join(
+ cache_path,
+ f"MOAD_limit{self.limit_complexes}_INDEX{self.split}"
+ + (
+ ""
+ if (
+ self.dockgen_esm_embeddings_path
+ if self.split != "train"
+ else self.esm_embeddings_path
+ )
+ is None
+ else "_esmEmbeddings"
+ )
+ + ("" if not self.include_miscellaneous_atoms else "_miscAtoms")
+ + ("" if min_protein_length is None else f"_minProteinLength{min_protein_length}")
+ + ("" if max_protein_length is None else f"_maxProteinLength{max_protein_length}"),
+ )
+
+ self.lig_cache_path = os.path.join(
+ cache_path,
+ f"MOAD_limit{self.limit_complexes}_INDEX{self.split}"
+ + f"_maxLigSize{self.max_lig_size}_H{int(not self.remove_hs)}"
+ + (
+ ""
+ if self.min_multi_lig_distance is None
+ else f"_minMultiLigDist{self.min_multi_lig_distance}"
+ ),
+ )
+
+ self.single_cluster_name = single_cluster_name
+ if split == "train":
+ split = "PDBBind"
+
+ with open(clusters_filepath, "rb") as f:
+ self.split_clusters = pickle.load(f)[split] # nosec
+
+ clusters_path = os.path.join(self.moad_dir, "new_cluster_to_ligands.pkl")
+ with open(clusters_path, "rb") as f:
+ self.cluster_to_ligands = pickle.load(f) # nosec
+
+ if not self.check_all_receptors():
+ os.makedirs(self.prot_cache_path, exist_ok=True)
+ self.preprocessing_receptors()
+
+ if not os.path.exists(os.path.join(self.lig_cache_path, "ligands.pkl")):
+ os.makedirs(self.lig_cache_path, exist_ok=True)
+ self.preprocessing_ligands()
+
+ log.info("Loading ligands from memory: ", os.path.join(self.lig_cache_path, "ligands.pkl"))
+ with open(os.path.join(self.lig_cache_path, "ligands.pkl"), "rb") as f:
+ self.ligands = pickle.load(f) # nosec
+
+ if require_ligand:
+ with open(os.path.join(self.lig_cache_path, "rdkit_ligands.pkl"), "rb") as f:
+ self.rdkit_ligands = pickle.load(f) # nosec
+ self.rdkit_ligands = {
+ ligs[0]["metadata"]["sample_ID"]: mol
+ for mol, ligs in zip(self.rdkit_ligands, self.ligands)
+ }
+
+ len_before = len(self.ligands)
+ if self.single_cluster_name is not None:
+ self.ligands = [
+ ligs
+ for ligs in self.ligands
+ if ligs[0]["metadata"]["sample_ID"]
+ in self.cluster_to_ligands[self.single_cluster_name]
+ ]
+ log.info(
+ f"Kept {len(self.ligands)} ligands in {self.single_cluster_name} out of {len_before}"
+ )
+
+ len_before = len(self.ligands)
+ self.ligands = {
+ ligs[0]["metadata"]["sample_ID"]: ligs
+ for ligs in self.ligands
+ if min_ligand_size == 0
+ or all(lig["features"]["sdf_coordinates"].shape[0] >= min_ligand_size for lig in ligs)
+ }
+ log.info(
+ f"Removed {len_before - len(self.ligands)} ligands below minimum size out of {len_before}"
+ )
+
+ receptors_names = {
+ (
+ ligs[0]["metadata"]["sample_ID"]
+ if self.split != "train"
+ else ligs[0]["metadata"]["sample_ID"][:6]
+ )
+ for ligs in self.ligands.values()
+ }
+ self.collect_receptors(receptors_names, max_receptor_size, remove_promiscuous_targets)
+
+ # filter ligands for which the receptor failed
+ tot_before = len(self.ligands)
+ self.ligands = {
+ k: v
+ for k, v in self.ligands.items()
+ if (k if self.split != "train" else k[:6]) in self.receptors
+ }
+ log.info(
+ f"Removed {tot_before - len(self.ligands)} ligands with no receptor out of {tot_before}"
+ )
+
+ if remove_pdbbind:
+ assert (
+ pdbbind_split_train is not None and pdbbind_split_val is not None
+ ), "PDBBind splits must be provided"
+ complexes_pdbbind = read_strings_from_txt(pdbbind_split_train) + read_strings_from_txt(
+ pdbbind_split_val
+ )
+ with open(os.path.join(self.moad_dir, "pdbbind_clusters_ecod.pkl"), "rb") as f:
+ pdbbind_to_cluster = pickle.load(f) # nosec
+ clusters_pdbbind = {pdbbind_to_cluster[c] for c in complexes_pdbbind}
+ self.split_clusters = [c for c in self.split_clusters if c not in clusters_pdbbind]
+ self.cluster_to_ligands = {
+ k: v for k, v in self.cluster_to_ligands.items() if k not in clusters_pdbbind
+ }
+ ligand_accepted = []
+ for c, ligands in self.cluster_to_ligands.items():
+ ligand_accepted += [(lig if self.split != "train" else lig[:6]) for lig in ligands]
+ ligand_accepted = set(ligand_accepted)
+ tot_before = len(self.ligands)
+ self.ligands = {k: v for k, v in self.ligands.items() if k in ligand_accepted}
+ log.info(
+ f"Removed {tot_before - len(self.ligands)} ligands in overlap with PDBBind out of {tot_before}"
+ )
+
+ if enforce_timesplit:
+ assert split_time is not None, "Time split must be provided"
+ with open(split_time) as f:
+ lines = f.readlines()
+ pdbids_from2019 = []
+ for i in range(6, len(lines), 4):
+ pdbids_from2019.append(lines[i][18:22])
+
+ pdbids_from2019 = set(pdbids_from2019)
+ len_before = len(self.ligands)
+ self.ligands = {
+ k: v for k, v in self.ligands.items() if k[:4].upper() not in pdbids_from2019
+ }
+ log.info(
+ f"Removed {len_before - len(self.ligands)} ligands from 2019 out of {len_before}"
+ )
+
+ if unroll_clusters:
+ rec_keys = {(k if self.split != "train" else k[:6]) for k in self.ligands.keys()}
+ self.cluster_to_ligands = {
+ k: [
+ k2
+ for k2 in self.ligands.keys()
+ if (k2 if self.split != "train" else k2[:6]) == k
+ ]
+ for k in rec_keys
+ }
+ self.split_clusters = list(rec_keys)
+ else:
+ for c in self.cluster_to_ligands.keys():
+ self.cluster_to_ligands[c] = [
+ v
+ for v in self.cluster_to_ligands[c]
+ if (v if self.split != "train" else v[:6]) in self.ligands
+ ]
+ self.split_clusters = [
+ c for c in self.split_clusters if len(self.cluster_to_ligands[c]) > 0
+ ]
+
+ # analyze and potentially filter the Binding MOAD dataset based on its apo-to-holo (a2h) structural assessment
+ if a2h_assessment_csv_filepath is not None and os.path.exists(a2h_assessment_csv_filepath):
+ a2h_assessment_df = pd.read_csv(a2h_assessment_csv_filepath)
+ a2h_assessment_df["ID"] = [
+ "_".join(Path(paths[0]).stem.split("_")[:2])
+ for paths in a2h_assessment_df["Filepath"].apply(ast.literal_eval).tolist()
+ ]
+ ligand_num_atoms = [
+ [int(num_atoms) for num_atoms in num_atoms_str.split(",")]
+ for num_atoms_str in a2h_assessment_df["Ligand_Num_Atoms"].tolist()
+ ]
+ a2h_assessment_df["Ligand_Total_Num_Atoms"] = np.array(
+ [np.array(num_atoms).sum() for num_atoms in ligand_num_atoms]
+ )
+ # import matplotlib.pyplot as plt
+ # import seaborn as sns
+ # plot_dir = Path(a2h_assessment_csv_filepath).parent / "plots"
+ # plot_dir.mkdir(exist_ok=True)
+ # plt.clf()
+ # sns.histplot(a2h_assessment_df["TM-score"])
+ # plt.title("Apo-To-Holo Protein TM-score")
+ # plt.savefig(plot_dir / "a2h_TM-score_hist.png")
+ # plt.clf()
+ # sns.histplot(a2h_assessment_df["RMSD"])
+ # plt.title("Apo-To-Holo Protein RMSD")
+ # plt.savefig(plot_dir / "a2h_RMSD_hist.png")
+ # plt.clf()
+ # plt.xlim(0, 1500)
+ # sns.histplot(a2h_assessment_df["Apo_Length"])
+ # plt.title("Apo Protein Length")
+ # plt.savefig(plot_dir / "apo_length_hist.png")
+ # plt.clf()
+ # plt.xlim(0, 500)
+ # sns.histplot(a2h_assessment_df["Ligand_Total_Num_Atoms"])
+ # plt.title("Ligand Total Number of Atoms")
+ # plt.savefig(plot_dir / "ligand_total_num_atoms_hist.png")
+ if filter_using_a2h_assessment and not is_test_dataset:
+ log.info(
+ f"Filtering the Binding MOAD {self.split} dataset based on its apo-to-holo (a2h) structural assessment"
+ )
+ a2h_assessment_df = a2h_assessment_df[
+ (a2h_assessment_df["TM-score"] >= a2h_min_tmscore)
+ & (a2h_assessment_df["RMSD"] <= a2h_max_rmsd)
+ & (a2h_assessment_df["Apo_Length"] >= a2h_min_protein_length)
+ & (a2h_assessment_df["Apo_Length"] <= a2h_max_protein_length)
+ & (a2h_assessment_df["Ligand_Total_Num_Atoms"] >= a2h_min_ligand_length)
+ & (a2h_assessment_df["Ligand_Total_Num_Atoms"] <= a2h_max_ligand_length)
+ ]
+ a2h_filtered_ids = {id: None for id in a2h_assessment_df["ID"].tolist()}
+ self.ligands = {
+ ligand: self.ligands[ligand]
+ for ligand in self.ligands
+ if ligand.lower() in a2h_filtered_ids
+ }
+ self.receptors = {
+ receptor: self.receptors[receptor]
+ for receptor in self.receptors
+ if receptor.lower() in a2h_filtered_ids
+ }
+ new_split_clusters = []
+ for cluster in self.split_clusters:
+ self.cluster_to_ligands[cluster] = [
+ ligand
+ for ligand in self.cluster_to_ligands[cluster]
+ if "_".join(ligand.lower().split("_")[:2]) in a2h_filtered_ids
+ ]
+ if len(self.cluster_to_ligands[cluster]) > 0:
+ new_split_clusters.append(cluster)
+ self.split_clusters = new_split_clusters
+
+ # analyze the DockGen dataset based on its apo-to-holo (a2h) structural assessment
+ if dockgen_a2h_assessment_csv_filepath is not None and os.path.exists(
+ dockgen_a2h_assessment_csv_filepath
+ ):
+ dockgen_a2h_assessment_df = pd.read_csv(dockgen_a2h_assessment_csv_filepath)
+ dockgen_a2h_assessment_df["ID"] = [
+ Path(paths[0]).stem.split("_holo_aligned")[0]
+ for paths in dockgen_a2h_assessment_df["Filepath"].apply(ast.literal_eval).tolist()
+ ]
+ dockgen_ligand_num_atoms = [
+ [num_atoms] for num_atoms in dockgen_a2h_assessment_df["Ligand_Num_Atoms"].tolist()
+ ]
+ dockgen_a2h_assessment_df["Ligand_Total_Num_Atoms"] = np.array(
+ [np.array(num_atoms).sum() for num_atoms in dockgen_ligand_num_atoms]
+ )
+ # import matplotlib.pyplot as plt
+ # import seaborn as sns
+ # dockgen_plot_dir = Path(dockgen_a2h_assessment_csv_filepath).parent / "plots"
+ # dockgen_plot_dir.mkdir(exist_ok=True)
+ # plt.clf()
+ # sns.histplot(dockgen_a2h_assessment_df["TM-score"])
+ # plt.title("Apo-To-Holo Protein TM-score")
+ # plt.savefig(dockgen_plot_dir / "a2h_TM-score_hist.png")
+ # plt.clf()
+ # sns.histplot(dockgen_a2h_assessment_df["RMSD"])
+ # plt.title("Apo-To-Holo Protein RMSD")
+ # plt.savefig(dockgen_plot_dir / "a2h_RMSD_hist.png")
+ # plt.clf()
+ # plt.xlim(0, 2000)
+ # sns.histplot(dockgen_a2h_assessment_df["Apo_Length"])
+ # plt.title("Apo Protein Length")
+ # plt.savefig(dockgen_plot_dir / "apo_length_hist.png")
+ # plt.clf()
+ # plt.xlim(0, 150)
+ # sns.histplot(dockgen_a2h_assessment_df["Ligand_Total_Num_Atoms"])
+ # plt.title("Ligand Total Number of Atoms")
+ # plt.savefig(dockgen_plot_dir / "ligand_total_num_atoms_hist.png")
+ if filter_using_a2h_assessment and not is_test_dataset:
+ log.info(
+ f"Filtering the DockGen {self.split} dataset based on its apo-to-holo (a2h) structural assessment"
+ )
+ dockgen_a2h_assessment_df = dockgen_a2h_assessment_df[
+ (dockgen_a2h_assessment_df["TM-score"] >= a2h_min_tmscore)
+ & (dockgen_a2h_assessment_df["RMSD"] <= a2h_max_rmsd)
+ & (dockgen_a2h_assessment_df["Apo_Length"] >= a2h_min_protein_length)
+ & (dockgen_a2h_assessment_df["Apo_Length"] <= a2h_max_protein_length)
+ & (
+ dockgen_a2h_assessment_df["Ligand_Total_Num_Atoms"]
+ >= a2h_min_ligand_length
+ )
+ & (
+ dockgen_a2h_assessment_df["Ligand_Total_Num_Atoms"]
+ <= a2h_max_ligand_length
+ )
+ ]
+ dockgen_a2h_filtered_ids = {
+ id: None for id in dockgen_a2h_assessment_df["ID"].tolist()
+ }
+ self.ligands = {
+ ligand: self.ligands[ligand]
+ for ligand in self.ligands
+ if ligand.replace(" ", "-") in dockgen_a2h_filtered_ids
+ }
+ self.receptors = {
+ receptor: self.receptors[receptor]
+ for receptor in self.receptors
+ if receptor.replace(" ", "-") in dockgen_a2h_filtered_ids
+ }
+ new_split_clusters = []
+ for cluster in self.split_clusters:
+ self.cluster_to_ligands[cluster] = [
+ ligand
+ for ligand in self.cluster_to_ligands[cluster]
+ if ligand.replace(" ", "-") in dockgen_a2h_filtered_ids
+ ]
+ if len(self.cluster_to_ligands[cluster]) > 0:
+ new_split_clusters.append(cluster)
+ self.split_clusters = new_split_clusters
+
+ list_names = [
+ name for cluster in self.split_clusters for name in self.cluster_to_ligands[cluster]
+ ]
+ log.info(
+ f"{len(list_names)} total complexes available from {len(self.split_clusters)} clusters after all {self.split} filtering"
+ )
+ with open(os.path.join(self.prot_cache_path, f"moad_{self.split}_names.txt"), "w") as f:
+ f.write("\n".join(list_names))
+
+ def len(self):
+ """Returns the number of complexes in the dataset."""
+ return (
+ len(self.split_clusters) * self.multiplicity
+ if self.total_dataset_size is None
+ else self.total_dataset_size
+ )
+
+ def get_by_name(self, ligand_name):
+ """Returns the complex graph for a given ligand name."""
+ lig_samples = copy.deepcopy(self.ligands[ligand_name])
+ protein_samples = copy.deepcopy(
+ self.receptors[(ligand_name if self.split != "train" else ligand_name[:6])]
+ )
+ complex_graph = merge_protein_and_ligands(
+ lig_samples,
+ protein_samples,
+ n_lig_patches=self.n_lig_patches,
+ )
+ if self.require_ligand:
+ complex_graph["metadata"]["mol"] = copy.deepcopy(self.rdkit_ligands[ligand_name])
+ return complex_graph
+
+ def get(self, idx):
+ """Returns the complex graph for a given index."""
+ if self.total_dataset_size is not None:
+ raise NotImplementedError(
+ "Random sampling not supported for total_dataset_size currently"
+ )
+ idx = random.randint(0, len(self.split_clusters) - 1) # nosec
+
+ idx = idx % len(self.split_clusters)
+ cluster = self.split_clusters[idx]
+
+ if self.no_randomness:
+ ligand_name = sorted(self.cluster_to_ligands[cluster])[0]
+ else:
+ ligand_name = random.choice(self.cluster_to_ligands[cluster]) # nosec
+
+ complex_graph = sample_inplace_to_torch(
+ self.get_by_name(ligand_name if self.split != "train" else ligand_name[:6])
+ )
+
+ if self.binding_affinity_values_dict is not None:
+ try:
+ # associate (super)ligands with their binding affinities;
+ # NOTE: this assumes that the number of available binding
+ # affinities for a given complex is equal to the number of
+ # ligands in the complex, which may not always be the case
+ ligand_code = ligand_name[:4].upper()
+ ligand_binding_affinities = list(
+ self.binding_affinity_values_dict[ligand_code].values()
+ )
+ assert (
+ len(ligand_binding_affinities) == complex_graph["metadata"]["num_molid"]
+ ), f"Mismatch between the number of ligands and available binding affinities for {ligand_name}."
+ complex_graph["features"]["affinity"] = torch.tensor(
+ ligand_binding_affinities, dtype=torch.float32
+ )
+ except Exception as e:
+ log.info(
+ f"Failed to get binding affinity for {ligand_name} due to: {e}. Substituting binding affinity value(s) with `NaN`."
+ )
+ complex_graph["features"]["affinity"] = torch.tensor(
+ [torch.nan for _ in range(complex_graph["metadata"]["num_molid"])],
+ dtype=torch.float32,
+ )
+ else:
+ complex_graph["features"]["affinity"] = torch.tensor(
+ [torch.nan for _ in range(complex_graph["metadata"]["num_molid"])],
+ dtype=torch.float32,
+ )
+
+ return centralize_complex_graph(complex_graph)
+
+ def get_all_complexes(self):
+ """Returns all the complexes in the dataset."""
+ complexes = {}
+ for cluster in self.split_clusters:
+ for ligand_name in self.cluster_to_ligands[cluster]:
+ complexes[ligand_name] = self.get_by_name(ligand_name)
+ return complexes
+
+ def preprocessing_receptors(self):
+ """Preprocesses the receptors and saves them to the cache path."""
+ log.info(
+ f"Processing receptors from [{self.split}] and saving them to [{self.prot_cache_path}]"
+ )
+
+ complex_names_all = sorted(
+ [item for cluster in self.split_clusters for item in self.cluster_to_ligands[cluster]]
+ )
+ if self.limit_complexes is not None and self.limit_complexes != 0:
+ complex_names_all = complex_names_all[: self.limit_complexes]
+
+ receptor_names_all = [
+ (item if self.split != "train" else item[:6]) for item in complex_names_all
+ ]
+ receptor_names_all = sorted(list(dict.fromkeys(receptor_names_all)))
+ log.info(f"Loading {len(receptor_names_all)} receptors.")
+
+ esm_embeddings_path = (
+ self.dockgen_esm_embeddings_path if self.split != "train" else self.esm_embeddings_path
+ )
+ esm_embeddings_sequences_path = (
+ self.dockgen_esm_embeddings_sequences_path
+ if self.split != "train"
+ else self.esm_embeddings_sequences_path
+ )
+
+ if esm_embeddings_path is not None:
+ log.info("Loading ESM embeddings")
+ sequences_to_embeddings = {}
+ receptor_names_all_dict = {name: None for name in receptor_names_all}
+ sequences_dict = fasta_to_dict(esm_embeddings_sequences_path)
+ for embedding_filepath in os.listdir(esm_embeddings_path):
+ key = Path(embedding_filepath).stem
+ if key.split("_chain_")[0].replace("-", " ") in receptor_names_all_dict:
+ embedding = torch.load(os.path.join(esm_embeddings_path, embedding_filepath))[
+ "representations"
+ ][33]
+ seq = sequences_dict[key]
+ id = int(key.split("_chain_")[1])
+ sequences_to_embeddings[seq + f":{id}"] = embedding
+ else:
+ sequences_to_embeddings = None
+
+ # running preprocessing in parallel on multiple workers and saving the progress every 1000 complexes
+ list_indices = list(range(len(receptor_names_all) // 1000 + 1))
+ random.shuffle(list_indices)
+ for i in list_indices:
+ if os.path.exists(os.path.join(self.prot_cache_path, f"receptors{i}.pkl")):
+ continue
+ receptor_names = receptor_names_all[1000 * i : 1000 * (i + 1)]
+ receptor_graphs = []
+ if self.num_workers > 1:
+ p = Pool(self.num_workers, maxtasksperchild=1)
+ p.__enter__()
+ with tqdm(
+ total=len(receptor_names),
+ desc=f"Loading receptors {i}/{len(receptor_names_all)//1000+1}",
+ ) as pbar:
+ map_fn = p.imap_unordered if self.num_workers > 1 else map
+ for t in map_fn(
+ self.get_receptor,
+ zip(receptor_names, [sequences_to_embeddings] * len(receptor_names)),
+ ):
+ if t is not None:
+ log.info(len(receptor_graphs))
+ receptor_graphs.append(t)
+ pbar.update()
+ if self.num_workers > 1:
+ p.__exit__(None, None, None)
+
+ log.info(f"Number of receptors: {len(receptor_graphs)}")
+ with open(os.path.join(self.prot_cache_path, f"receptors{i}.pkl"), "wb") as f:
+ pickle.dump((receptor_graphs), f)
+ return receptor_names_all
+
+ def check_all_receptors(self):
+ """Checks if all the receptors are preprocessed."""
+ complex_names_all = sorted(
+ [item for c in self.split_clusters for item in self.cluster_to_ligands[c]]
+ )
+ if self.limit_complexes is not None and self.limit_complexes != 0:
+ complex_names_all = complex_names_all[: self.limit_complexes]
+ receptor_names_all = [item[:6] for item in complex_names_all]
+ receptor_names_all = list(dict.fromkeys(receptor_names_all))
+ for i in range(len(receptor_names_all) // 1000 + 1):
+ if not os.path.exists(os.path.join(self.prot_cache_path, f"receptors{i}.pkl")):
+ return False
+ return True
+
+ def collect_receptors(
+ self, receptors_to_keep=None, max_receptor_size=None, remove_promiscuous_targets=None
+ ):
+ """Collects the receptors to keep and saves them to the receptors dictionary."""
+ complex_names_all = sorted(
+ [item for c in self.split_clusters for item in self.cluster_to_ligands[c]]
+ )
+ if self.limit_complexes is not None and self.limit_complexes != 0:
+ complex_names_all = complex_names_all[: self.limit_complexes]
+ receptor_names_all = [
+ (item if self.split != "train" else item[:6]) for item in complex_names_all
+ ]
+ receptor_names_all = sorted(list(dict.fromkeys(receptor_names_all)))
+
+ receptor_graphs_all = []
+ total_recovered = 0
+ log.info(f"Loading {len(receptor_names_all)} receptors to keep {len(receptors_to_keep)}.")
+ for i in range(len(receptor_names_all) // 1000 + 1):
+ log.info(f'Prot path: {os.path.join(self.prot_cache_path, f"receptors{i}.pkl")}')
+ with open(os.path.join(self.prot_cache_path, f"receptors{i}.pkl"), "rb") as f:
+ item = pickle.load(f) # nosec
+ total_recovered += len(item)
+ if receptors_to_keep is not None:
+ item = [t for t in item if t["metadata"]["sample_ID"] in receptors_to_keep]
+ receptor_graphs_all.extend(item)
+
+ cur_len = len(receptor_graphs_all)
+ log.info(
+ f"Kept {len(receptor_graphs_all)} receptors out of {len(receptor_names_all)} total and recovered {total_recovered}"
+ )
+
+ if max_receptor_size is not None:
+ receptor_graphs_all = [
+ rec
+ for rec in receptor_graphs_all
+ if rec["features"]["res_atom_positions"].shape[0] <= max_receptor_size
+ ]
+ log.info(
+ f"Kept {len(receptor_graphs_all)} receptors out of {cur_len} after filtering by size"
+ )
+ cur_len = len(receptor_graphs_all)
+
+ if remove_promiscuous_targets is not None:
+ promiscuous_targets = set()
+ for name in complex_names_all:
+ item = name.split("_")
+ if int(item[3]) > remove_promiscuous_targets:
+ promiscuous_targets.add(name[:6])
+ receptor_graphs_all = [
+ rec
+ for rec in receptor_graphs_all
+ if rec["metadata"]["sample_ID"] not in promiscuous_targets
+ ]
+ log.info(
+ f"Kept {len(receptor_graphs_all)} receptors out of {cur_len} after removing promiscuous targets"
+ )
+
+ self.receptors = {}
+ for r in receptor_graphs_all:
+ self.receptors[r["metadata"]["sample_ID"]] = r
+
+ def get_receptor(self, par):
+ """Returns the receptor graph for a given receptor name."""
+ name, sequences_to_embeddings = par
+ if self.split == "train":
+ holo_protein_filepath = os.path.join(
+ self.moad_dir, "pdb_protein", f"{name}_protein.pdb"
+ )
+ else:
+ holo_protein_filepath = os.path.join(
+ self.dockgen_dir,
+ name.replace(" ", "-"),
+ f"{name.replace(' ', '-')}_protein_processed.pdb",
+ )
+ if not os.path.exists(holo_protein_filepath):
+ log.error(f"Holo receptor not found for {name}: {holo_protein_filepath}")
+ return None
+
+ try:
+ holo_af_protein = pdb_filepath_to_protein(holo_protein_filepath)
+ holo_protein_sample = process_protein(
+ holo_af_protein,
+ sample_name=f"{name}_",
+ )
+ complex_graph = holo_protein_sample
+
+ except Exception as e:
+ log.error(f"Skipping holo {name} because of the error: {e}")
+ return None
+
+ if np.isnan(complex_graph["features"]["res_atom_positions"]).any():
+ log.error(
+ f"NaN in holo receptor pos for {name}. Skipping preprocessing for this example..."
+ )
+ return None
+
+ try:
+ apo_protein_structure_dir = (
+ self.dockgen_apo_protein_structure_dir
+ if self.split != "train"
+ else self.apo_protein_structure_dir
+ )
+ if apo_protein_structure_dir is not None:
+ apo_protein_filepath = os.path.join(
+ apo_protein_structure_dir,
+ f"{name.replace(' ', '-')}_holo_aligned_esmfold_protein.pdb",
+ )
+ apo_af_protein = pdb_filepath_to_protein(apo_protein_filepath)
+ apo_protein_sample = process_protein(
+ apo_af_protein,
+ sample_name=f"{name}_",
+ sequences_to_embeddings=sequences_to_embeddings,
+ )
+ for key in complex_graph.keys():
+ for subkey, value in apo_protein_sample[key].items():
+ complex_graph[key]["apo_" + subkey] = value
+ if not np.array_equal(
+ complex_graph["features"]["res_type"],
+ complex_graph["features"]["apo_res_type"],
+ ):
+ log.error(
+ f"Residue type mismatch between holo protein and apo protein for {name}. Skipping preprocessing for this example..."
+ )
+ return None
+ if np.isnan(complex_graph["features"]["apo_res_atom_positions"]).any():
+ log.error(
+ f"NaN in apo receptor pos for {name}. Skipping preprocessing for this example..."
+ )
+ return None
+
+ except Exception as e:
+ log.error(f"Skipping apo {name} because of the error: {e}")
+ return None
+
+ if (
+ self.min_protein_length is not None
+ and complex_graph["metadata"]["num_a"] < self.min_protein_length
+ and not self.is_test_dataset
+ ):
+ log.info(f"Skipping {name} because of its length {complex_graph['metadata']['num_a']}")
+ return None
+ if (
+ self.max_protein_length is not None
+ and complex_graph["metadata"]["num_a"] > self.max_protein_length
+ and not self.is_test_dataset
+ ):
+ log.info(f"Skipping {name} because of its length {complex_graph['metadata']['num_a']}")
+ return None
+
+ complex_graph["metadata"]["sample_ID"] = name
+ return complex_graph
+
+ def preprocessing_ligands(self):
+ """Preprocesses the ligands and saves them to the cache path."""
+ log.info(
+ f"Processing complexes from [{self.split}] and saving them to [{self.lig_cache_path}]"
+ )
+
+ complex_names_all = sorted(
+ [item for c in self.split_clusters for item in self.cluster_to_ligands[c]]
+ )
+ if self.limit_complexes is not None and self.limit_complexes != 0:
+ complex_names_all = complex_names_all[: self.limit_complexes]
+
+ complex_names_all = [
+ (item if self.split != "train" else item[:6]) for item in complex_names_all
+ ]
+ complex_names_all = sorted(list(dict.fromkeys(complex_names_all)))
+ log.info(f"Loading {len(complex_names_all)} (multi-)ligands.")
+
+ # running preprocessing in parallel on multiple workers and saving the progress every 1000 complexes
+ list_indices = list(range(len(complex_names_all) // 1000 + 1))
+ random.shuffle(list_indices)
+ for i in list_indices:
+ if os.path.exists(os.path.join(self.lig_cache_path, f"ligands{i}.pkl")):
+ continue
+ complex_names = complex_names_all[1000 * i : 1000 * (i + 1)]
+ ligand_graphs, rdkit_ligands = [], []
+ if self.num_workers > 1:
+ p = Pool(self.num_workers, maxtasksperchild=1)
+ p.__enter__()
+ with tqdm(
+ total=len(complex_names),
+ desc=f"Loading complexes {i}/{len(complex_names_all)//1000+1}",
+ ) as pbar:
+ map_fn = p.imap_unordered if self.num_workers > 1 else map
+ for t in map_fn(self.get_ligand, complex_names):
+ if t is not None:
+ ligand_graphs.append(t[0])
+ rdkit_ligands.append(t[1])
+ pbar.update()
+ if self.num_workers > 1:
+ p.__exit__(None, None, None)
+
+ with open(os.path.join(self.lig_cache_path, f"ligands{i}.pkl"), "wb") as f:
+ pickle.dump((ligand_graphs), f)
+ with open(os.path.join(self.lig_cache_path, f"rdkit_ligands{i}.pkl"), "wb") as f:
+ pickle.dump((rdkit_ligands), f)
+
+ ligand_graphs_all = []
+ for i in range(len(complex_names_all) // 1000 + 1):
+ with open(os.path.join(self.lig_cache_path, f"ligands{i}.pkl"), "rb") as f:
+ item = pickle.load(f) # nosec
+ ligand_graphs_all.extend(item)
+ with open(os.path.join(self.lig_cache_path, "ligands.pkl"), "wb") as f:
+ pickle.dump((ligand_graphs_all), f)
+
+ rdkit_ligands_all = []
+ for i in range(len(complex_names_all) // 1000 + 1):
+ with open(os.path.join(self.lig_cache_path, f"rdkit_ligands{i}.pkl"), "rb") as f:
+ item = pickle.load(f) # nosec
+ rdkit_ligands_all.extend(item)
+ with open(os.path.join(self.lig_cache_path, "rdkit_ligands.pkl"), "wb") as f:
+ pickle.dump((rdkit_ligands_all), f)
+
+ def get_ligand(self, name):
+ """Returns the ligand graph(s) for a given ligand name."""
+ if self.split == "train":
+ lig_paths = glob.glob(os.path.join(self.moad_dir, "pdb_superligand", f"{name}_*.pdb"))
+ else:
+ # NOTE: we refer to the DockGen data directory for all validation and test data
+ lig_paths = glob.glob(
+ os.path.join(
+ self.dockgen_dir,
+ name.replace(" ", "-"),
+ f"{name.replace(' ', '-')}_ligand*.pdb",
+ )
+ )
+
+ if not lig_paths:
+ log.error(f"No ligands found for {name}")
+ return None
+
+ try:
+ # NOTE: for `6nco_1_KQP_0` in the DockGen test dataset, we ignore sanitization to avoid an explicit valence error
+ is_unsanitizeable_ligand = self.split != "train" and name in [
+ "6wjy_2_U41_0",
+ "6nco_1_KQP_0",
+ ]
+ ligands = [
+ Chem.MolFromPDBFile(lig_path, sanitize=not is_unsanitizeable_ligand)
+ for lig_path in lig_paths
+ ]
+ ligands_pos_list = [ligand.GetConformer().GetPositions() for ligand in ligands]
+
+ lig = combine_molecules(ligands)
+ if not is_unsanitizeable_ligand:
+ Chem.SanitizeMol(lig)
+ except Exception as e:
+ log.error(
+ f"Failed to load, combine, or sanitize Binding MOAD ligand(s) for {name} due to: {e}. Skipping preprocessing for this example..."
+ )
+ return None
+
+ if self.min_multi_lig_distance is not None and len(ligands_pos_list) > 1:
+ # calculate pairwise distances between all ligands
+ for i in range(len(ligands_pos_list)):
+ for j in range(i + 1, len(ligands_pos_list)):
+ ligand_i_pos = ligands_pos_list[i]
+ ligand_j_pos = ligands_pos_list[j]
+ distances = distance.cdist(ligand_i_pos, ligand_j_pos, metric="euclidean")
+ min_dist = np.min(distances)
+ if min_dist < self.min_multi_lig_distance:
+ log.error(
+ f"At a minimum distance of {min_dist}, ligands {i + 1} and {j + 1} were too close to each other for {name}. Skipping preprocessing for this example"
+ )
+ return None
+
+ if self.max_lig_size is not None and lig.GetNumHeavyAtoms() > self.max_lig_size:
+ log.error(
+ f"(Multi-)ligand with {lig.GetNumHeavyAtoms()} heavy atoms is larger than max_lig_size {self.max_lig_size}. Skipping preprocessing for this example..."
+ )
+ return None
+
+ try:
+ lig_samples = [
+ process_molecule(
+ ligand,
+ ref_conf_xyz=np.array(ligand.GetConformer().GetPositions()),
+ return_as_dict=True,
+ )
+ for ligand in ligands
+ ]
+
+ except Exception as e:
+ log.error(f"Skipping {name} because of the error: {e}")
+ return None
+
+ for lig_sample in lig_samples:
+ lig_sample["metadata"]["sample_ID"] = name
+ return lig_samples, lig
diff --git a/forks/FlowDock/flowdock/data/components/mol_features.py b/forks/FlowDock/flowdock/data/components/mol_features.py
new file mode 100644
index 00000000..d6df3ecf
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/mol_features.py
@@ -0,0 +1,902 @@
+# Adapted from: https://github.com/zrqiao/NeuralPLexer
+
+import msgpack
+import msgpack_numpy as m
+import networkx
+import networkx as nx
+import numpy as np
+import rootutils
+import torch
+from beartype.typing import List, Optional
+from rdkit import Chem
+from rdkit.Chem import AllChem
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+# Periodic table metadata
+# Runtime accessing from mendeleev is surprisingly slow, tabulate here
+from flowdock.data.components.physical import PGROUP_IDS, PPERIOD_IDS
+from flowdock.utils import RankedLogger
+
+BONDORDER = {
+ Chem.rdchem.BondType.SINGLE: 1,
+ Chem.rdchem.BondType.AROMATIC: 1.5,
+ Chem.rdchem.BondType.DOUBLE: 2,
+ Chem.rdchem.BondType.TRIPLE: 3,
+}
+
+m.patch()
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+def one_of_k_encoding(x, allowable_set) -> List[bool]:
+ """
+ Maps inputs not in the allowable set to the last element.
+ Modified from https://github.com/XuhanLiu/NGFP.
+
+ :param x: input
+ :param allowable_set: allowable set
+ :return: one-hot encoding
+ """
+ if x not in allowable_set:
+ x = allowable_set[-1]
+ return list(map(lambda s: x == s, allowable_set))
+
+
+def get_atom_encoding(atom: Chem.rdchem.Atom) -> np.ndarray:
+ """Get atom encoding.
+
+ :param atom: Chem.rdchem.Atom atom object
+ :return: encoding_list
+ """
+ encoding_list = one_of_k_encoding(
+ PGROUP_IDS[atom.GetSymbol()], list(range(1, 19))
+ ) + one_of_k_encoding(PPERIOD_IDS[atom.GetSymbol()], list(range(1, 6)))
+ return np.array(encoding_list)
+
+
+def get_bond_encoding(bond: Chem.rdchem.Bond) -> np.ndarray:
+ """Get bond encoding.
+
+ :param bond: Chem.rdchem.Bond bond object
+ :return: np.array
+ """
+ bt = bond.GetBondType()
+ return np.array(
+ [
+ bt == Chem.rdchem.BondType.SINGLE,
+ bt == Chem.rdchem.BondType.DOUBLE,
+ bt == Chem.rdchem.BondType.TRIPLE,
+ bt == Chem.rdchem.BondType.AROMATIC,
+ ]
+ )
+
+
+def mol_to_graph(mol):
+ """Convert RDKit Mol to NetworkX graph
+ Adapted from https://github.com/deepchem/deepchem
+
+ Convert mol into a graph representation atoms are nodes, and bonds
+ are vertices stored as graph
+
+ Parameters
+ ----------
+ mol: rdkit Mol
+ The molecule to convert into a graph.
+
+ Returns
+ -------
+ graph: networkx.Graph
+ Contains atoms indices as nodes, edges as bonds.
+ """
+ G = nx.DiGraph()
+ num_atoms = mol.GetNumAtoms()
+ G.add_nodes_from(range(num_atoms))
+ if mol.GetNumBonds() == 0:
+ # assert num_atoms == 1
+ for i in range(mol.GetNumAtoms()):
+ G.add_edge(i, i, bond_idx=i)
+ return G
+ for i in range(mol.GetNumBonds()):
+ from_idx = mol.GetBondWithIdx(i).GetBeginAtomIdx()
+ to_idx = mol.GetBondWithIdx(i).GetEndAtomIdx()
+ G.add_edge(from_idx, to_idx, bond_idx=2 * i)
+ G.add_edge(to_idx, from_idx, bond_idx=2 * i + 1)
+ return G
+
+
+def compute_bond_pair_triangles(nx_graph):
+ """Compute bond pair triangles."""
+ # Line graph edge trackers
+ bab_dict, aaa_dict = dict(), dict()
+ in_atoms, mid_atoms, out_atoms = [], [], []
+ for atom_id in list(nx_graph):
+ neighbor_nodes_array = np.array(nx_graph[atom_id])
+ # Allow for repeated nodes iff num_atoms < 3
+ if len(nx_graph) == 1:
+ # Single-atom species
+ local_in = np.array([atom_id], dtype=int)
+ local_out = np.array([atom_id], dtype=int)
+ elif len(nx_graph) == 2:
+ # Allow for i=k to handle diatomic molecules
+ local_in = np.array(neighbor_nodes_array, dtype=int)
+ local_out = np.array(neighbor_nodes_array, dtype=int)
+ else:
+ # Masking out i=k for others
+ nna_rep = np.repeat(
+ neighbor_nodes_array[:, np.newaxis], len(neighbor_nodes_array), axis=1
+ )
+ mask = ~np.eye(nna_rep.shape[0], dtype=bool)
+ local_in = nna_rep[mask]
+ local_out = nna_rep.transpose()[mask]
+ local_mid = np.full_like(local_in, atom_id)
+ in_atoms.append(local_in)
+ mid_atoms.append(local_mid)
+ out_atoms.append(local_out)
+ in_atoms = np.concatenate(in_atoms)
+ mid_atoms = np.concatenate(mid_atoms)
+ out_atoms = np.concatenate(out_atoms)
+ in_bonds = np.array(
+ [nx_graph[in_atoms[i]][mid_atoms[i]]["bond_idx"] for i in range(len(in_atoms))]
+ )
+ out_bonds = np.array(
+ [nx_graph[mid_atoms[i]][out_atoms[i]]["bond_idx"] for i in range(len(out_atoms))]
+ )
+ for triangle_idx in range(len(in_atoms)):
+ bab_dict[
+ (in_bonds[triangle_idx], mid_atoms[triangle_idx], out_bonds[triangle_idx])
+ ] = triangle_idx
+ aaa_dict[
+ (in_atoms[triangle_idx], mid_atoms[triangle_idx], out_atoms[triangle_idx])
+ ] = triangle_idx
+ return (
+ np.stack([in_atoms, mid_atoms, out_atoms, in_bonds, out_bonds], axis=1),
+ bab_dict,
+ aaa_dict,
+ )
+
+
+def get_conformers_as_tensor(mol, n_conf, return_new_mol=False):
+ """Get conformers as tensor."""
+ conf_mol = Chem.AddHs(Chem.Mol(mol))
+ try:
+ AllChem.EmbedMultipleConfs(conf_mol, clearConfs=True, numConfs=n_conf, numThreads=0)
+ assert len(conf_mol.GetConformers()) == n_conf
+ except AssertionError:
+ AllChem.EmbedMultipleConfs(
+ conf_mol,
+ clearConfs=True,
+ numConfs=n_conf,
+ numThreads=0,
+ useRandomCoords=True,
+ )
+ assert len(conf_mol.GetConformers()) == n_conf
+ AllChem.MMFFOptimizeMoleculeConfs(conf_mol, mmffVariant="MMFF94", numThreads=0)
+ conf_mol = Chem.RemoveHs(conf_mol)
+ xyzs = np.array([c.GetPositions() for c in conf_mol.GetConformers()])
+ assert xyzs.shape[0] == n_conf
+ if return_new_mol:
+ return xyzs, conf_mol
+ return xyzs
+
+
+def is_potential_stereo_bond(bond: Chem.rdchem.Bond):
+ """Check if bond is a potential stereo bond."""
+ bond_order = BONDORDER.get(bond.GetBondType(), 1)
+ if bond_order > 1.4 and bond_order < 3.0:
+ return True
+ else:
+ return False
+
+
+def is_potential_stereo_center(atom: Chem.rdchem.Atom):
+ """
+ Note: Instead of checking whether a real chiral tag can be assigned to the query atom,
+ this function regards all neighbor atoms as unique groups.
+
+ Based on rules from https://www.rdkit.org/docs/RDKit_Book.html#brief-description-of-the-findpotentialstereo-algorithm
+ """
+ assert atom.GetSymbol() != "H"
+ if atom.GetDegree() >= 4:
+ return True
+ elif atom.GetTotalDegree() >= 4:
+ return True
+ elif atom.GetSymbol() in ["P", "As", "S", "Se"] and atom.GetDegree() >= 3:
+ return True
+ else:
+ return False
+
+
+def compute_stereo_encodings(
+ query_atom_id: int,
+ bonded_atom_id: int,
+ ref_bond_pos: int,
+ mol: Chem.rdchem.Mol,
+ nx_graph: networkx.DiGraph,
+ atom_idx_on_triangle: tuple,
+ ref_geom_xyz,
+ cutoff=0.05,
+):
+ """
+ Three types of stereochemistry annotations:
+ 1. In-clique index of the atom(s) the query atom is bonded to;
+ 2. On which side of the tangent plane (or unsure);
+ 3. If off-plane, above or below the plane (or unsure);
+ All stereogenic bonds must be explicitly labelled within the input graph.
+
+ :return: torch.Tensor with stereochemistry annotations
+ """
+ # Is a single atom or not
+ stereo_enc = [atom_idx_on_triangle[0] == atom_idx_on_triangle[1]]
+ # Is spinor type or not
+ stereo_enc += [atom_idx_on_triangle[0] == atom_idx_on_triangle[2]]
+ stereo_enc += [atom_idx_on_triangle[i] == query_atom_id for i in range(3)]
+ # Whether bonded to other atoms in the triplet
+ stereo_enc += [atom_idx_on_triangle[i] in nx_graph[query_atom_id].keys() for i in range(3)]
+ # stereo_enc += [
+ # mol.GetAtomWithIdx(bonded_atom_id).GetHybridization()
+ # == Chem.rdchem.HybridizationType.SP
+ # ]
+ stereo_enc += one_of_k_encoding(ref_bond_pos, [0, 1])
+ if ref_geom_xyz is None:
+ stereo_enc += [False] * 4
+ return np.array(stereo_enc)
+ in_bond_vec = ref_geom_xyz[atom_idx_on_triangle[1]] - ref_geom_xyz[atom_idx_on_triangle[0]]
+ in_vec = in_bond_vec / np.linalg.norm(in_bond_vec)
+ out_bond_vec = ref_geom_xyz[atom_idx_on_triangle[2]] - ref_geom_xyz[atom_idx_on_triangle[1]]
+ out_vec = out_bond_vec / np.linalg.norm(out_bond_vec)
+ z_vec = np.cross(in_vec, out_vec)
+
+ nx_graph[query_atom_id][bonded_atom_id]["bond_idx"] // 2
+ ref_bond_id = (
+ nx_graph[atom_idx_on_triangle[ref_bond_pos]][atom_idx_on_triangle[ref_bond_pos + 1]][
+ "bond_idx"
+ ]
+ // 2
+ )
+
+ # Resolving bond-centered stereochemistry
+ # if mol.GetBondWithIdx(ref_bond_id).GetStereo() in [
+ # Chem.rdchem.BondStereo.STEREOCIS,
+ # Chem.rdchem.BondStereo.STEREOTRANS,
+ # Chem.rdchem.BondStereo.STEREOE,
+ # Chem.rdchem.BondStereo.STEREOZ,
+ # ]:
+ if is_potential_stereo_bond(mol.GetBondWithIdx(ref_bond_id)):
+ query_bond_vec = ref_geom_xyz[query_atom_id] - ref_geom_xyz[bonded_atom_id]
+ query_bond_vec = query_bond_vec / np.linalg.norm(query_bond_vec)
+ if ref_bond_pos == 0:
+ ref_bond_vec = in_vec
+ if bonded_atom_id == atom_idx_on_triangle[0]:
+ query_bond_vec = -query_bond_vec
+ elif ref_bond_pos == 1:
+ if bonded_atom_id == atom_idx_on_triangle[1]:
+ query_bond_vec = -query_bond_vec
+ ref_bond_vec = out_vec
+ else:
+ raise ValueError
+ p_z_bond = np.dot(
+ np.cross(ref_bond_vec, query_bond_vec),
+ z_vec,
+ )
+ # print("bond p_z:", p_z_bond)
+ if np.abs(p_z_bond) > cutoff:
+ bond_stereo_enc = one_of_k_encoding(
+ p_z_bond > 0,
+ [True, False],
+ )
+ else:
+ # Cannot resolve E/Z geometry
+ bond_stereo_enc = [False, False]
+ else:
+ bond_stereo_enc = [False, False]
+ stereo_enc += bond_stereo_enc
+
+ # Resolving atom-centered stereochemistry
+ if bonded_atom_id == atom_idx_on_triangle[1] and is_potential_stereo_center(
+ mol.GetAtomWithIdx(bonded_atom_id)
+ ):
+ query_bonded_vec = ref_geom_xyz[query_atom_id] - ref_geom_xyz[bonded_atom_id]
+ query_bonded_vec = query_bonded_vec / np.linalg.norm(query_bonded_vec)
+ p_z_atom = np.dot(query_bonded_vec, z_vec)
+ # print("atom p_z:", p_z_atom)
+ if np.abs(p_z_atom) > cutoff:
+ atom_stereo_enc = one_of_k_encoding(
+ p_z_atom > 0,
+ [True, False],
+ )
+ else:
+ atom_stereo_enc = [False, False]
+ else:
+ atom_stereo_enc = [False, False]
+ stereo_enc += atom_stereo_enc
+
+ return np.array(stereo_enc)
+
+
+def compute_all_stereo_chemistry_encodings(
+ mol, nx_graph, atom_idx_on_triangles, aaa_dict, only_2d=False, ref_conf_xyz=None
+):
+ """Compute all stereochemistry encodings."""
+ if not only_2d:
+ if ref_conf_xyz is None:
+ ref_conf_xyz = get_conformers_as_tensor(mol, 1)[0]
+ else:
+ ref_conf_xyz = None
+ triangle_pairs_list, stereo_encodings_list = [], []
+ for triangle_atoms in atom_idx_on_triangles:
+ triangle_atoms = tuple(int(aidx) for aidx in triangle_atoms)
+ for query_atom in nx_graph[triangle_atoms[0]].keys():
+ if query_atom == triangle_atoms[1]:
+ continue
+ triangle_pairs_list.append(
+ [
+ aaa_dict[triangle_atoms],
+ aaa_dict[(query_atom, triangle_atoms[0], triangle_atoms[1])],
+ ]
+ )
+ stereo_encodings_list.append(
+ compute_stereo_encodings(
+ query_atom,
+ triangle_atoms[0],
+ 0,
+ mol,
+ nx_graph,
+ triangle_atoms,
+ ref_conf_xyz,
+ )
+ )
+ for query_atom in nx_graph[triangle_atoms[1]].keys():
+ if query_atom == triangle_atoms[0]:
+ continue
+ if query_atom == triangle_atoms[2]:
+ continue
+ # Outgoing bonds
+ triangle_pairs_list.append(
+ [
+ aaa_dict[triangle_atoms],
+ aaa_dict[(triangle_atoms[0], triangle_atoms[1], query_atom)],
+ ]
+ )
+ stereo_encodings_list.append(
+ compute_stereo_encodings(
+ query_atom,
+ triangle_atoms[1],
+ 0,
+ mol,
+ nx_graph,
+ triangle_atoms,
+ ref_conf_xyz,
+ )
+ )
+ # Incoming bonds
+ triangle_pairs_list.append(
+ [
+ aaa_dict[triangle_atoms],
+ aaa_dict[(query_atom, triangle_atoms[1], triangle_atoms[2])],
+ ]
+ )
+ stereo_encodings_list.append(
+ compute_stereo_encodings(
+ query_atom,
+ triangle_atoms[1],
+ 1,
+ mol,
+ nx_graph,
+ triangle_atoms,
+ ref_conf_xyz,
+ )
+ )
+ for query_atom in nx_graph[triangle_atoms[2]].keys():
+ if query_atom == triangle_atoms[1]:
+ continue
+ triangle_pairs_list.append(
+ [
+ aaa_dict[triangle_atoms],
+ aaa_dict[(triangle_atoms[1], triangle_atoms[2], query_atom)],
+ ]
+ )
+ stereo_encodings_list.append(
+ compute_stereo_encodings(
+ query_atom,
+ triangle_atoms[2],
+ 1,
+ mol,
+ nx_graph,
+ triangle_atoms,
+ ref_conf_xyz,
+ )
+ )
+
+ return np.array(triangle_pairs_list), stereo_encodings_list
+
+
+def tensorize_indexers(
+ bond_atom_ids,
+ triangle_ids,
+ triangle_pair_ids,
+ prop_init_ids=None,
+ prop_ids=None,
+ prop_pair_ids=None,
+ allow_dummy=False,
+):
+ """Tensorize indexers for bond-centered and atom-centered stereochemistry annotations."""
+ # Incidence matrices flattened - querying nodes from hyperedges
+ if len(bond_atom_ids) == 0:
+ if not allow_dummy:
+ raise ValueError("There must be at least one bond (ij), got 0")
+ bond_atom_ids = np.zeros((0, 2), dtype=np.int_)
+ if len(triangle_ids) == 0:
+ if not allow_dummy:
+ raise ValueError("There must be at least one triplet (ijk), got 0")
+ triangle_ids = np.zeros((0, 5), dtype=np.int_)
+ prop_init_ids = np.zeros((0, 3), dtype=np.int_)
+ # prop_ids = np.zeros((0, 2), dtype=np.int_)
+ if len(triangle_pair_ids) == 0:
+ triangle_pair_ids = np.zeros((0, 2), dtype=np.int_)
+ # prop_pair_ids = np.zeros((0, 3), dtype=np.int_)
+ indexer = {
+ "gather_idx_ij_i": np.array(bond_atom_ids[:, 0], dtype=np.int_),
+ "gather_idx_ij_j": np.array(bond_atom_ids[:, 1], dtype=np.int_),
+ "gather_idx_ijk_i": np.array(triangle_ids[:, 0], dtype=np.int_),
+ "gather_idx_ijk_j": np.array(triangle_ids[:, 1], dtype=np.int_),
+ "gather_idx_ijk_k": np.array(triangle_ids[:, 2], dtype=np.int_),
+ "gather_idx_ijk_ij": np.array(triangle_ids[:, 3], dtype=np.int_),
+ "gather_idx_ijk_jk": np.array(triangle_ids[:, 4], dtype=np.int_),
+ "gather_idx_ijkl_ijk": np.array(triangle_pair_ids[:, 0], dtype=np.int_),
+ "gather_idx_ijkl_jkl": np.array(triangle_pair_ids[:, 1], dtype=np.int_),
+ # "gather_idx_u0ijk_ijk": np.array(prop_init_ids[:, 0], dtype=np.int_),
+ # "gather_idx_u0ijk_u0": np.array(prop_init_ids[:, 1], dtype=np.int_),
+ }
+ return indexer
+
+
+def process_molecule(
+ mol: Chem.Mol,
+ only_2d: bool = False,
+ return_mol: bool = False,
+ ref_conf_xyz: Optional[np.ndarray] = None,
+ return_as_dict: bool = False,
+) -> tuple:
+ """Process (multi-fragment) molecule and return metadata, indexer and features.
+
+ :param mol: Chem.Mol molecule object
+ :param only_2d: whether to compute 2D stereochemistry only
+ :param return_mol: whether to return the molecule object
+ :param ref_conf_xyz: reference 3D coordinates
+ :param return_as_dict: whether to return as dictionary
+ :return: feature set dictionary
+ """
+ # NOTE: all bonds are directional
+ atom_encodings_list = [get_atom_encoding(atom) for atom in mol.GetAtoms()]
+ atomic_numbers = [atom.GetAtomicNum() for atom in mol.GetAtoms()]
+ nx_graph = mol_to_graph(mol)
+ bond_atom_ids = np.array(nx_graph.edges())
+ if mol.GetNumBonds() == 0:
+ bond_encodings_list = [np.zeros((4,)) for _ in range(len(bond_atom_ids))]
+ else:
+ bond_encodings_list = [
+ get_bond_encoding(mol.GetBondWithIdx(eid // 2)) for eid in range(len(bond_atom_ids))
+ ]
+ triangle_aaabb_ids, _, aaa_dict = compute_bond_pair_triangles(nx_graph)
+ atom_idx_on_triangles = triangle_aaabb_ids[:, :3]
+ triangle_pairs_list, stereo_encodings_list = compute_all_stereo_chemistry_encodings(
+ mol,
+ nx_graph,
+ atom_idx_on_triangles,
+ aaa_dict,
+ only_2d=only_2d,
+ ref_conf_xyz=ref_conf_xyz,
+ )
+ metadata = {
+ "num_molid": 1,
+ "num_ligand_atom": len(atom_encodings_list),
+ "num_i": len(atom_encodings_list),
+ "num_j": len(atom_encodings_list),
+ "num_k": len(atom_encodings_list),
+ "num_u": len(atom_encodings_list),
+ "num_ij": len(bond_atom_ids),
+ "num_jk": len(bond_atom_ids),
+ "num_ijk": len(triangle_aaabb_ids),
+ "num_jkl": len(triangle_aaabb_ids),
+ "num_ligand_clique": len(triangle_aaabb_ids),
+ "num_ijkl": len(triangle_pairs_list),
+ }
+ indexer = tensorize_indexers(bond_atom_ids, triangle_aaabb_ids, triangle_pairs_list)
+ # NOTE: the target `molid` is always `0`
+ indexer["gather_idx_i_molid"] = np.zeros((metadata["num_i"],), dtype=np.int_)
+ indexer["gather_idx_ijk_molid"] = np.zeros((metadata["num_ijk"],), dtype=np.int_)
+ if len(stereo_encodings_list) == 0:
+ stereo_encodings = np.zeros((0, 14), dtype=np.int_)
+ else:
+ stereo_encodings = np.stack(stereo_encodings_list, axis=0)
+ features = {
+ "atomic_numbers": np.array(atomic_numbers, dtype=int),
+ "atom_encodings": np.stack(atom_encodings_list, axis=0),
+ "bond_encodings": np.stack(bond_encodings_list, axis=0),
+ "stereo_chemistry_encodings": stereo_encodings,
+ "sdf_coordinates": ref_conf_xyz,
+ }
+ if return_mol:
+ if return_as_dict:
+ return {
+ "metadata": metadata,
+ "indexer": indexer,
+ "features": features,
+ "misc": {},
+ "mol": mol,
+ }
+ return mol, metadata, indexer, features
+ if return_as_dict:
+ return {
+ "metadata": metadata,
+ "indexer": indexer,
+ "features": features,
+ "misc": {},
+ }
+ return metadata, indexer, features
+
+
+def smiles2inputs(tokenizer, smiles, pad_length=128):
+ """Adapted from megamolbart."""
+
+ assert isinstance(smiles, str)
+ if pad_length:
+ assert pad_length >= len(smiles) + 2
+
+ tokens = tokenizer.tokenize([smiles], pad=True)
+
+ # Append to tokens and mask if appropriate
+ if pad_length:
+ for i in range(len(tokens["original_tokens"])):
+ num_i = len(tokens["original_tokens"][i])
+ n_pad = pad_length - len(tokens["original_tokens"][i])
+ tokens["original_tokens"][i] += [tokenizer.pad_token] * n_pad
+ tokens["masked_pad_masks"][i] += [1] * n_pad
+
+ token_ids = torch.tensor(tokenizer.convert_tokens_to_ids(tokens["original_tokens"]))
+ pad_mask = torch.tensor(tokens["masked_pad_masks"]).bool()
+ encode_input = {
+ "features": {"encoder_input": token_ids, "encoder_pad_mask": pad_mask},
+ "indexer": {"gather_idx_i_molid": np.zeros((num_i,), dtype=np.int_)},
+ "metadata": {"num_i": num_i, "num_molid": 1},
+ }
+
+ return encode_input
+
+
+def process_mol_file(
+ fname,
+ featurize=True,
+ tokenizer=None,
+ return_mol=False,
+ sanitize=True,
+ coord_feats=True,
+ pair_feats=False,
+ discard_coords=False,
+ **kwargs,
+):
+ """Process molecule file and return feature set."""
+ if fname.endswith("msgpack"):
+ with open(fname, "rb") as data_file:
+ byte_data = data_file.read()
+ data_loaded = msgpack.unpackb(byte_data)
+ return data_loaded
+ elif fname.endswith("sdf"):
+ fsuppl = Chem.SDMolSupplier(fname, sanitize=sanitize)
+ mol = next(fsuppl)
+ if not sanitize:
+ mol.UpdatePropertyCache(strict=False)
+ elif fname.endswith("mol2"):
+ mol = Chem.MolFromMol2File(fname, sanitize=sanitize)
+ elif fname.endswith("pdb"):
+ mol = Chem.MolFromPDBFile(fname, sanitize=sanitize)
+ if not sanitize:
+ mol.UpdatePropertyCache(strict=False)
+ else:
+ log.warning("No suffix found for ligand input, assuming SMILES input")
+ mol = Chem.MolFromSmiles(fname, sanitize=sanitize)
+ if sanitize:
+ mol = Chem.RemoveHs(mol, updateExplicitCount=True)
+ else:
+ mol = Chem.RemoveHs(mol, sanitize=False)
+ if not featurize:
+ return mol
+ if tokenizer is not None:
+ return smiles2inputs(tokenizer, Chem.rdmolfiles.MolToSmiles(mol))
+ if discard_coords:
+ conf_xyz = get_conformers_as_tensor(mol, 1)[0]
+ else:
+ conf_xyz = np.array(mol.GetConformer().GetPositions())
+ metadata, indexer, features = process_molecule(
+ mol, return_mol=False, ref_conf_xyz=conf_xyz, **kwargs
+ )
+ if coord_feats:
+ features["sdf_coordinates"] = conf_xyz
+ feature_set = {
+ "metadata": metadata,
+ "indexer": indexer,
+ "features": features,
+ "misc": {},
+ }
+ if pair_feats:
+ attach_pair_idx_and_encodings(feature_set, **kwargs)
+ if return_mol:
+ return feature_set, mol
+ return feature_set
+
+
+def attach_pair_idx_and_encodings(feature_set, max_n_frames=None, lazy_eval=False):
+ """Attach pair index and encodings to the feature set."""
+ if lazy_eval:
+ if "scatter_idx_u0ijk_Uijk" in feature_set["indexer"].keys():
+ return feature_set
+ num_triplets = feature_set["metadata"]["num_ijk"]
+ num_atoms = feature_set["metadata"]["num_i"]
+ num_frame_pairs = feature_set["metadata"]["num_ijkl"]
+ if max_n_frames is None:
+ max_n_frames = num_triplets
+ else:
+ max_n_frames = max(min(num_triplets, max_n_frames), 1)
+ key_frame_idx = np.random.choice(num_triplets, size=max_n_frames, replace=False)
+ key_atom_idx = feature_set["indexer"]["gather_idx_ijk_j"][key_frame_idx]
+ num_key_frames = key_frame_idx.shape[0]
+ # scatter_idx_u0ijk_Uijk = (
+ # feature_set["indexer"]["gather_idx_u0ijk_ijk"]
+ # + feature_set["indexer"]["gather_idx_u0ijk_u0"] * num_triplets
+ # )
+ gather_idx_UI_Uijk = (
+ np.broadcast_to(
+ np.arange(num_key_frames)[:, None],
+ (num_key_frames, num_key_frames),
+ ).flatten()
+ * num_triplets
+ + np.broadcast_to(
+ key_frame_idx[None, :],
+ (num_key_frames, num_key_frames),
+ ).flatten()
+ )
+ gather_idx_Uijkl_Uijk = (
+ np.broadcast_to(
+ np.arange(num_key_frames)[:, None], (num_key_frames, num_frame_pairs)
+ ).flatten()
+ * num_triplets
+ + np.broadcast_to(
+ feature_set["indexer"]["gather_idx_ijkl_ijk"][None, :],
+ (num_key_frames, num_frame_pairs),
+ ).flatten()
+ )
+ gather_idx_Uijkl_ujkl = (
+ np.broadcast_to(
+ np.arange(num_key_frames)[:, None], (num_key_frames, num_frame_pairs)
+ ).flatten()
+ * num_triplets
+ + np.broadcast_to(
+ feature_set["indexer"]["gather_idx_ijkl_jkl"][None, :],
+ (num_key_frames, num_frame_pairs),
+ ).flatten()
+ )
+ gather_idx_Uijkl_ijkl = np.broadcast_to(
+ np.arange(num_frame_pairs)[None, :], (num_key_frames, num_frame_pairs)
+ ).flatten()
+
+ adjacency_mat = np.zeros((num_atoms, num_atoms), dtype=np.int_)
+ adjacency_mat[
+ feature_set["indexer"]["gather_idx_ij_i"],
+ feature_set["indexer"]["gather_idx_ij_j"],
+ ] = 1
+ sum_pair_path_dist = [np.eye(num_atoms, dtype=np.int_)]
+ for path_length in range(3):
+ sum_pair_path_dist.append(np.matmul(sum_pair_path_dist[-1], adjacency_mat))
+ sum_pair_path_dist = np.stack(sum_pair_path_dist, axis=2)
+ atom_pair_feature_mat = np.zeros((num_atoms, num_atoms, 4), dtype=np.float_)
+ atom_pair_feature_mat[
+ feature_set["indexer"]["gather_idx_ij_i"],
+ feature_set["indexer"]["gather_idx_ij_j"],
+ ] = feature_set["features"]["bond_encodings"]
+ atom_pair_feature_mat = np.concatenate(
+ [atom_pair_feature_mat, (sum_pair_path_dist > 0).astype(np.float_)], axis=2
+ )
+ uv_adj_mat = np.sum(sum_pair_path_dist, axis=2) > 0
+ gather_idx_uv_u = np.broadcast_to(np.arange(num_atoms)[:, None], (num_atoms, num_atoms))[
+ uv_adj_mat
+ ]
+ gather_idx_uv_v = np.broadcast_to(np.arange(num_atoms)[None, :], (num_atoms, num_atoms))[
+ uv_adj_mat
+ ]
+
+ atom_frame_pair_feat_initial_ = np.concatenate(
+ [
+ atom_pair_feature_mat[key_atom_idx, :][:, feature_set["indexer"]["gather_idx_ijk_i"]],
+ atom_pair_feature_mat[key_atom_idx, :][:, feature_set["indexer"]["gather_idx_ijk_j"]],
+ atom_pair_feature_mat[key_atom_idx, :][:, feature_set["indexer"]["gather_idx_ijk_k"]],
+ ],
+ axis=2,
+ ).reshape((num_key_frames * num_triplets, atom_pair_feature_mat.shape[2] * 3))
+
+ # Generate on-the-fly to reduce disk usage
+ feature_set["indexer"].update(
+ {
+ "gather_idx_U_u": key_atom_idx,
+ "gather_idx_I_ijk": key_frame_idx,
+ "gather_idx_I_molid": np.zeros((num_key_frames,), dtype=np.int_),
+ "gather_idx_UI_Uijk": gather_idx_UI_Uijk,
+ "gather_idx_UI_u": np.broadcast_to(
+ key_atom_idx[:, None], (num_key_frames, num_key_frames)
+ ).flatten(),
+ "gather_idx_UI_U": np.broadcast_to(
+ np.arange(num_key_frames)[:, None], (num_key_frames, num_key_frames)
+ ).flatten(),
+ "gather_idx_UI_I": np.broadcast_to(
+ np.arange(num_key_frames)[None, :], (num_key_frames, num_key_frames)
+ ).flatten(),
+ # "scatter_idx_u0ijk_Uijk": scatter_idx_u0ijk_Uijk,
+ "gather_idx_Uijk_u": np.broadcast_to(
+ key_atom_idx[:, None], (num_key_frames, num_triplets)
+ ).flatten(),
+ "gather_idx_Uijk_ijk": np.broadcast_to(
+ np.arange(num_triplets)[None, :], (num_key_frames, num_triplets)
+ ).flatten(),
+ "gather_idx_Uijkl_Uijk": gather_idx_Uijkl_Uijk,
+ "gather_idx_Uijkl_ujkl": gather_idx_Uijkl_ujkl,
+ "gather_idx_Uijkl_ijkl": gather_idx_Uijkl_ijkl,
+ "gather_idx_uv_u": gather_idx_uv_u,
+ "gather_idx_uv_v": gather_idx_uv_v,
+ }
+ )
+ feature_set["features"].update(
+ {
+ "atom_pair_encodings": atom_pair_feature_mat[uv_adj_mat],
+ "atom_frame_pair_encodings": atom_frame_pair_feat_initial_,
+ }
+ )
+ feature_set["metadata"]["num_v"] = num_atoms
+ feature_set["metadata"]["num_I"] = num_key_frames
+ feature_set["metadata"]["num_U"] = num_key_frames
+ feature_set["metadata"]["num_Uijk"] = num_triplets * num_key_frames
+ feature_set["metadata"]["num_ujkl"] = num_triplets * num_key_frames
+ feature_set["metadata"]["num_Uijkl"] = num_frame_pairs * num_key_frames
+ feature_set["metadata"]["num_uv"] = gather_idx_uv_u.shape[0]
+ feature_set["metadata"]["num_UI"] = num_key_frames * num_key_frames
+ return feature_set
+
+
+def iterable_query(k1, k2, dc):
+ """Query a key from a list of dictionaries."""
+ return [d[k1][k2] for d in dc]
+
+
+def collate_idx_numpy(idx_ten_list, dst_sample_sizes):
+ """Collate NumPy index tensors."""
+ elewise_dst_offsets = np.repeat(
+ np.cumsum(np.array([0] + dst_sample_sizes[:-1], dtype=np.int_), axis=0),
+ np.array([idx_ten.shape[0] for idx_ten in idx_ten_list], dtype=np.int_),
+ )
+ col_idx_ten = np.concatenate(idx_ten_list, axis=0) + elewise_dst_offsets
+ return col_idx_ten
+
+
+def collate_idx_tensors(idx_ten_list, dst_sample_sizes):
+ """Collate PyTorch index tensors."""
+ device = idx_ten_list[0].device
+ elewise_dst_offsets = torch.repeat_interleave(
+ torch.cumsum(
+ torch.tensor([0] + dst_sample_sizes[:-1], dtype=torch.long, device=device),
+ dim=0,
+ ),
+ torch.tensor(
+ [idx_ten.size(0) for idx_ten in idx_ten_list],
+ dtype=torch.long,
+ device=device,
+ ),
+ )
+ col_idx_ten = torch.cat(idx_ten_list, dim=0).add_(elewise_dst_offsets)
+ return col_idx_ten
+
+
+def collate_samples(list_of_samples: list, exclude=[]):
+ """Collate samples into a batch."""
+ list_of_samples = [sample for sample in list_of_samples if sample is not None]
+ batch_metadata = dict()
+ for key in list_of_samples[0]["metadata"].keys():
+ batch_metadata[f"{key}_per_sample"] = iterable_query("metadata", key, list_of_samples)
+ if key.startswith("num_"):
+ if key.endswith("_per_sample"):
+ # Track the batch-level summary only
+ continue
+ batch_metadata[key] = sum(batch_metadata[f"{key}_per_sample"])
+
+ batch_indexer = dict()
+ for indexer_name in list_of_samples[0]["indexer"].keys():
+ indexer_name_parts = indexer_name.split("_")
+ batch_indexer[indexer_name] = collate_idx_tensors(
+ iterable_query("indexer", indexer_name, list_of_samples),
+ batch_metadata[f"num_{indexer_name_parts[-1]}_per_sample"],
+ )
+
+ batch_features = dict()
+ for feature_name in list_of_samples[0]["features"].keys():
+ if feature_name in exclude:
+ continue
+ batch_features[feature_name] = torch.cat(
+ iterable_query("features", feature_name, list_of_samples), dim=0
+ ).float()
+ ret_batch = {
+ "metadata": batch_metadata,
+ "indexer": batch_indexer,
+ "features": batch_features,
+ "batch_size": len(list_of_samples),
+ }
+ if "misc" in list_of_samples[0].keys():
+ ret_batch["misc"] = {
+ key: [x for sample in list_of_samples for x in sample["misc"][key]]
+ for key in list_of_samples[0]["misc"].keys()
+ }
+ if "labels" in list_of_samples[0].keys():
+ ret_batch["labels"] = dict()
+ for label_name in list_of_samples[0]["labels"].keys():
+ ret_batch["labels"][label_name] = torch.stack(
+ iterable_query("labels", label_name, list_of_samples), dim=0
+ ).float()
+ return ret_batch
+
+
+def collate_numpy_samples(list_of_samples: list):
+ """Collate NumPy samples into a batch."""
+ list_of_samples = [sample for sample in list_of_samples if sample is not None]
+ if len(list_of_samples) == 0:
+ return {
+ "metadata": {},
+ "indexer": {},
+ "features": {},
+ "labels": {},
+ "misc": {},
+ "batch_size": 0,
+ }
+ batch_metadata = dict()
+ for key in list_of_samples[0]["metadata"].keys():
+ batch_metadata[f"{key}_per_sample"] = iterable_query("metadata", key, list_of_samples)
+ if key.startswith("num_"):
+ if key.endswith("_per_sample"):
+ continue
+ batch_metadata[key] = sum(batch_metadata[f"{key}_per_sample"])
+
+ batch_indexer = dict()
+ for indexer_name in list_of_samples[0]["indexer"].keys():
+ i2 = indexer_name.split("_")[-1]
+ batch_indexer[indexer_name] = collate_idx_numpy(
+ iterable_query("indexer", indexer_name, list_of_samples),
+ batch_metadata[f"num_{i2}_per_sample"],
+ )
+
+ batch_features = dict()
+ for feature_name in list_of_samples[0]["features"].keys():
+ batch_features[feature_name] = np.concatenate(
+ iterable_query("features", feature_name, list_of_samples), axis=0
+ ).astype(np.float32)
+ batch_misc = {
+ key: [x for sample in list_of_samples for x in sample["misc"][key]]
+ for key in list_of_samples[0]["misc"].keys()
+ }
+ if "labels" not in list_of_samples[0].keys():
+ return {
+ "metadata": batch_metadata,
+ "indexer": batch_indexer,
+ "features": batch_features,
+ "misc": batch_misc,
+ "batch_size": len(list_of_samples),
+ }
+
+ batch_labels = dict()
+ for label_name in list_of_samples[0]["labels"].keys():
+ batch_labels[label_name] = np.stack(
+ iterable_query("labels", label_name, list_of_samples), axis=0
+ ).astype(np.float32)
+ return {
+ "metadata": batch_metadata,
+ "indexer": batch_indexer,
+ "features": batch_features,
+ "labels": batch_labels,
+ "misc": batch_misc,
+ "batch_size": len(list_of_samples),
+ }
diff --git a/forks/FlowDock/flowdock/data/components/parse_chi.py b/forks/FlowDock/flowdock/data/components/parse_chi.py
new file mode 100644
index 00000000..98d77430
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/parse_chi.py
@@ -0,0 +1,162 @@
+# From: https://github.com/gcorso/DiffDock
+# From Nick Polizzi
+
+import os
+from collections import defaultdict
+
+import numpy as np
+import prody as pr
+import rootutils
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock.data.components.constants import (
+ aa_idx2aa_short,
+ aa_long2short,
+ aa_short2aa_idx,
+ atom_order,
+ chi,
+)
+from flowdock.utils import RankedLogger
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+def get_dihedral_indices(resname, chi_num):
+ """Return the atom indices for the specified dihedral angle."""
+ if resname not in chi:
+ return np.array([np.nan] * 4)
+ if chi_num not in chi[resname]:
+ return np.array([np.nan] * 4)
+ return np.array([atom_order[resname].index(x) for x in chi[resname][chi_num]])
+
+
+dihedral_indices = defaultdict(list)
+for aa in atom_order.keys():
+ for i in range(1, 5):
+ inds = get_dihedral_indices(aa, i)
+ dihedral_indices[aa].append(inds)
+ dihedral_indices[aa] = np.array(dihedral_indices[aa])
+
+
+def vector_batch(a, b):
+ return a - b
+
+
+def unit_vector_batch(v):
+ return v / np.linalg.norm(v, axis=1, keepdims=True)
+
+
+def dihedral_angle_batch(p):
+ b0 = vector_batch(p[:, 0], p[:, 1])
+ b1 = vector_batch(p[:, 1], p[:, 2])
+ b2 = vector_batch(p[:, 2], p[:, 3])
+
+ n1 = np.cross(b0, b1)
+ n2 = np.cross(b1, b2)
+
+ m1 = np.cross(n1, b1 / np.linalg.norm(b1, axis=1, keepdims=True))
+
+ x = np.sum(n1 * n2, axis=1)
+ y = np.sum(m1 * n2, axis=1)
+
+ deg = np.degrees(np.arctan2(y, x))
+
+ deg[deg < 0] += 360
+
+ return deg
+
+
+def batch_compute_dihedral_angles(sidechains):
+ sidechains_np = np.array(sidechains)
+ dihedral_angles = dihedral_angle_batch(sidechains_np)
+ return dihedral_angles
+
+
+def get_coords(prody_pdb):
+ resindices = sorted(set(prody_pdb.ca.getResindices()))
+ coords = np.zeros((len(resindices), 14, 3))
+ for i, resind in enumerate(resindices):
+ sel = prody_pdb.select(f"resindex {resind}")
+ resname = sel.getResnames()[0]
+ for j, name in enumerate(
+ atom_order[aa_long2short[resname] if resname in aa_long2short else "X"]
+ ):
+ sel_resnum_name = sel.select(f"name {name}")
+ if sel_resnum_name is not None:
+ coords[i, j, :] = sel_resnum_name.getCoords()[0]
+ else:
+ coords[i, j, :] = [np.nan, np.nan, np.nan]
+ return coords
+
+
+def get_onehot_sequence(seq):
+ onehot = np.zeros((len(seq), 20))
+ for i, aa in enumerate(seq):
+ idx = aa_short2aa_idx[aa] if aa in aa_short2aa_idx else 7 # 7 is the index for GLY
+ onehot[i, idx] = 1
+ return onehot
+
+
+def get_dihedral_indices(onehot_sequence):
+ return np.array(
+ [dihedral_indices[aa_idx2aa_short[aa_idx]] for aa_idx in np.where(onehot_sequence)[1]]
+ )
+
+
+def _get_chi_angles(coords, indices):
+ X = coords
+ Y = indices.astype(int)
+ N = coords.shape[0]
+ mask = np.isnan(indices)
+ Y[mask] = 0
+ Z = X[np.arange(N)[:, None, None], Y, :]
+ Z[mask] = np.nan
+ chi_angles = batch_compute_dihedral_angles(Z.reshape(-1, 4, 3)).reshape(N, 4)
+ return chi_angles
+
+
+def get_chi_angles(coords, seq, return_onehot=False):
+ """Get chi angles of sidechains.
+
+ Parameters
+ ----------
+ prody_pdb : prody.AtomGroup
+ prody pdb object or selection
+ return_coords : bool, optional
+ return coordinates of prody_pdb in (N, 14, 3) array format, by default False
+ return_onehot : bool, optional
+ return one-hot sequence of prody_pdb, by default False
+
+ Returns
+ -------
+ numpy array of shape (N, 4)
+ Array contains chi angles of sidechains in row-order of residue indices in prody_pdb.
+ If a chi angle is not defined for a residue, due to missing atoms or GLY / ALA, it is set to np.nan.
+ """
+ onehot = get_onehot_sequence(seq)
+ dihedral_indices = get_dihedral_indices(onehot)
+ if return_onehot:
+ return _get_chi_angles(coords, dihedral_indices), onehot
+ return _get_chi_angles(coords, dihedral_indices)
+
+
+def test_get_chi_angles(print_chi_angles=False):
+ # need internet connection of '6w70.pdb' in working directory
+ pdb = pr.parsePDB("6w70")
+ prody_pdb = pdb.select("chain A")
+ chi_angles = get_chi_angles(prody_pdb)
+ assert chi_angles.shape == (prody_pdb.ca.numAtoms(), 4)
+ assert chi_angles[0, 0] < 56.0 and chi_angles[0, 0] > 55.0
+ log.info("test_get_chi_angles passed")
+ try:
+ os.remove("6w70.pdb.gz")
+ except Exception as e:
+ log.warning(f"Error removing file: {e}")
+ if print_chi_angles:
+ log.info(chi_angles)
+ return True
+
+
+if __name__ == "__main__":
+ test_get_chi_angles(print_chi_angles=True)
diff --git a/forks/FlowDock/flowdock/data/components/pdbbind_dataset.py b/forks/FlowDock/flowdock/data/components/pdbbind_dataset.py
new file mode 100644
index 00000000..80177294
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/pdbbind_dataset.py
@@ -0,0 +1,580 @@
+# Adapted from: https://github.com/gcorso/DiffDock
+
+import ast
+import binascii
+import copy
+import glob
+import os
+import pickle # nosec
+import random
+from collections import defaultdict
+from multiprocessing import Pool
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import rootutils
+import torch
+from beartype.typing import Optional
+from rdkit import Chem
+from rdkit.Chem import AddHs, MolFromSmiles, RemoveAllHs, RemoveHs
+from torch_geometric.data import Dataset
+from tqdm import tqdm
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock.data.components.mol_features import process_molecule
+from flowdock.data.components.process_mols import generate_conformer, read_molecule
+from flowdock.utils import RankedLogger
+from flowdock.utils.data_utils import (
+ centralize_complex_graph,
+ merge_protein_and_ligands,
+ pdb_filepath_to_protein,
+ process_protein,
+)
+from flowdock.utils.model_utils import sample_inplace_to_torch
+from flowdock.utils.utils import read_strings_from_txt
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+class PDBBindDataset(Dataset):
+ """A PyTorch Geometric Dataset for PDBBind dataset."""
+
+ def __init__(
+ self,
+ root,
+ transform=None,
+ cache_path=os.path.join("data", "cache"),
+ split_path="data" + os.sep,
+ limit_complexes=0,
+ num_workers=0,
+ max_lig_size=None,
+ remove_hs=False,
+ num_conformers=1,
+ esm_embeddings_path=None,
+ apo_protein_structure_dir=None,
+ require_ligand=False,
+ include_miscellaneous_atoms=False,
+ protein_path_list=None,
+ ligand_descriptions=None,
+ keep_local_structures=False,
+ protein_file="protein_processed",
+ ligand_file="ligand",
+ min_protein_length: Optional[int] = 10,
+ max_protein_length: Optional[int] = 4000,
+ is_test_dataset=False,
+ a2h_assessment_csv_filepath=None,
+ filter_using_a2h_assessment=False,
+ a2h_min_tmscore=None,
+ a2h_max_rmsd=None,
+ a2h_min_protein_length=None,
+ a2h_max_protein_length=None,
+ a2h_min_ligand_length=None,
+ a2h_max_ligand_length=None,
+ binding_affinity_values_dict=None,
+ n_lig_patches=32,
+ ):
+ """Initializes the dataset."""
+
+ super().__init__(root, transform)
+ self.pdbbind_dir = root
+ self.include_miscellaneous_atoms = include_miscellaneous_atoms
+ self.max_lig_size = max_lig_size
+ self.split_path = split_path
+ self.limit_complexes = limit_complexes
+ self.num_workers = num_workers
+ self.remove_hs = remove_hs
+ self.esm_embeddings_path = esm_embeddings_path
+ self.apo_protein_structure_dir = apo_protein_structure_dir
+ self.use_old_wrong_embedding_order = False
+ self.require_ligand = require_ligand
+ self.protein_path_list = protein_path_list
+ self.ligand_descriptions = ligand_descriptions
+ self.keep_local_structures = keep_local_structures
+ self.protein_file = protein_file
+ self.fixed_knn_radius_graph = True
+ self.ligand_file = ligand_file
+ self.min_protein_length = min_protein_length
+ self.max_protein_length = max_protein_length
+ self.is_test_dataset = is_test_dataset
+ self.binding_affinity_values_dict = binding_affinity_values_dict
+ self.n_lig_patches = n_lig_patches
+
+ split = os.path.splitext(os.path.basename(self.split_path))[0]
+ self.full_cache_path = os.path.join(
+ cache_path,
+ f"PDBBind_limit{self.limit_complexes}"
+ f"_INDEX{split}"
+ f"_maxLigSize{self.max_lig_size}_H{int(not self.remove_hs)}"
+ + ("" if self.esm_embeddings_path is None else "_esmEmbeddings")
+ + "_full"
+ + ("" if not keep_local_structures else "_keptLocalStruct")
+ + (
+ ""
+ if protein_path_list is None or ligand_descriptions is None
+ else str(binascii.crc32("".join(ligand_descriptions + protein_path_list).encode()))
+ )
+ + ("" if protein_file == "protein_processed" else "_" + protein_file)
+ + ("" if not self.include_miscellaneous_atoms else "_miscAtoms")
+ + ("" if self.use_old_wrong_embedding_order else "_chainOrd")
+ + ("" if min_protein_length is None else f"_minProteinLength{min_protein_length}")
+ + ("" if max_protein_length is None else f"_maxProteinLength{max_protein_length}"),
+ )
+ self.num_conformers = num_conformers
+
+ if not self.check_all_complexes():
+ os.makedirs(self.full_cache_path, exist_ok=True)
+ if protein_path_list is None or ligand_descriptions is None:
+ self.preprocessing()
+ else:
+ self.inference_preprocessing()
+
+ self.complex_graphs, self.rdkit_ligands = self.collect_all_complexes()
+
+ # analyze and potentially filter the PDBBind dataset based on its apo-to-holo (a2h) structural assessment
+ if a2h_assessment_csv_filepath is not None and os.path.exists(a2h_assessment_csv_filepath):
+ a2h_assessment_df = pd.read_csv(a2h_assessment_csv_filepath)
+ a2h_assessment_df["ID"] = [
+ Path(paths[0]).stem[:4]
+ for paths in a2h_assessment_df["Filepath"].apply(ast.literal_eval).tolist()
+ ]
+ ligand_num_atoms = [
+ [int(num_atoms) for num_atoms in num_atoms_str.split(",")]
+ for num_atoms_str in a2h_assessment_df["Ligand_Num_Atoms"].tolist()
+ ]
+ a2h_assessment_df["Ligand_Total_Num_Atoms"] = np.array(
+ [np.array(num_atoms).sum() for num_atoms in ligand_num_atoms]
+ )
+ # import matplotlib.pyplot as plt
+ # import seaborn as sns
+ # plot_dir = Path(a2h_assessment_csv_filepath).parent / "plots"
+ # plot_dir.mkdir(exist_ok=True)
+ # plt.clf()
+ # sns.histplot(a2h_assessment_df["TM-score"])
+ # plt.title("Apo-To-Holo Protein TM-score")
+ # plt.savefig(plot_dir / "a2h_TM-score_hist.png")
+ # plt.clf()
+ # sns.histplot(a2h_assessment_df["RMSD"])
+ # plt.title("Apo-To-Holo Protein RMSD")
+ # plt.savefig(plot_dir / "a2h_RMSD_hist.png")
+ # plt.clf()
+ # plt.xlim(0, 1500)
+ # sns.histplot(a2h_assessment_df["Apo_Length"])
+ # plt.title("Apo Protein Length")
+ # plt.savefig(plot_dir / "apo_length_hist.png")
+ # plt.clf()
+ # plt.xlim(0, 500)
+ # sns.histplot(a2h_assessment_df["Ligand_Total_Num_Atoms"])
+ # plt.title("Ligand Total Number of Atoms")
+ # plt.savefig(plot_dir / "ligand_total_num_atoms_hist.png")
+ if filter_using_a2h_assessment and not is_test_dataset:
+ log.info(
+ f"Filtering the PDBBind {split} dataset based on its apo-to-holo (a2h) structural assessment"
+ )
+ a2h_assessment_df = a2h_assessment_df[
+ (a2h_assessment_df["TM-score"] >= a2h_min_tmscore)
+ & (a2h_assessment_df["RMSD"] <= a2h_max_rmsd)
+ & (a2h_assessment_df["Apo_Length"] >= a2h_min_protein_length)
+ & (a2h_assessment_df["Apo_Length"] <= a2h_max_protein_length)
+ & (a2h_assessment_df["Ligand_Total_Num_Atoms"] >= a2h_min_ligand_length)
+ & (a2h_assessment_df["Ligand_Total_Num_Atoms"] <= a2h_max_ligand_length)
+ ]
+ a2h_filtered_ids = {id: None for id in a2h_assessment_df["ID"].tolist()}
+ new_complex_graphs, new_rdkit_ligands = [], []
+ for complex_id, complex_obj in enumerate(self.complex_graphs):
+ if complex_obj["metadata"]["sample_ID"].lower() in a2h_filtered_ids:
+ new_complex_graphs.append(complex_obj)
+ new_rdkit_ligands.append(self.rdkit_ligands[complex_id])
+ self.complex_graphs = new_complex_graphs
+ self.rdkit_ligands = new_rdkit_ligands
+
+ list_names = [complex_obj["metadata"]["sample_ID"] for complex_obj in self.complex_graphs]
+ log.info(
+ f"{len(list_names)} total complexes available from {self.full_cache_path} after all {split} filtering"
+ )
+ with open(
+ os.path.join(
+ self.full_cache_path,
+ f"pdbbind_{os.path.splitext(os.path.basename(self.split_path))[0][:3]}_names.txt",
+ ),
+ "w",
+ ) as f:
+ f.write("\n".join(list_names))
+
+ def len(self):
+ """Returns the number of complexes in the dataset."""
+ return len(self.complex_graphs)
+
+ def get(self, idx, default_ligand_ccd_id: str = "XXX"):
+ """Returns a HeteroData object for a given index."""
+ complex_graph = sample_inplace_to_torch(copy.deepcopy(self.complex_graphs[idx]))
+ if self.require_ligand:
+ complex_graph["metadata"]["mol"] = RemoveAllHs(copy.deepcopy(self.rdkit_ligands[idx]))
+ if self.binding_affinity_values_dict is not None:
+ try:
+ complex_graph["features"]["affinity"] = torch.tensor(
+ [
+ self.binding_affinity_values_dict[
+ complex_graph["metadata"]["sample_ID"].lower()
+ ][default_ligand_ccd_id]
+ ],
+ dtype=torch.float32,
+ )
+ except Exception as e:
+ log.info(
+ f"Binding affinity value not found for {complex_graph['metadata']['sample_ID']} due to: {e}"
+ )
+ complex_graph["features"]["affinity"] = torch.tensor(
+ [torch.nan], dtype=torch.float32
+ )
+ else:
+ complex_graph["features"]["affinity"] = torch.tensor([torch.nan], dtype=torch.float32)
+ return centralize_complex_graph(complex_graph)
+
+ def preprocessing(self):
+ """Preprocesses the complexes for training."""
+ log.info(
+ f"Processing complexes from [{self.split_path}] and saving them to [{self.full_cache_path}]"
+ )
+
+ complex_names_all = read_strings_from_txt(self.split_path)
+ if self.limit_complexes is not None and self.limit_complexes != 0:
+ complex_names_all = complex_names_all[: self.limit_complexes]
+ log.info(f"Loading {len(complex_names_all)} complexes.")
+
+ if self.esm_embeddings_path is not None:
+ log.info("Loading ESM embeddings")
+ chain_embeddings_dictlist = defaultdict(list)
+ chain_indices_dictlist = defaultdict(list)
+ for embedding_filepath in os.listdir(self.esm_embeddings_path):
+ key = Path(embedding_filepath).stem
+ key_name = key.split("_chain_")[0]
+ if key_name in complex_names_all:
+ embedding = torch.load(
+ os.path.join(self.esm_embeddings_path, embedding_filepath)
+ )["representations"][33]
+ chain_embeddings_dictlist[key_name].append(embedding)
+ chain_indices_dictlist[key_name].append(int(key.split("_chain_")[1]))
+ lm_embeddings_chains_all = []
+ for name in complex_names_all:
+ complex_chains_embeddings = chain_embeddings_dictlist[name]
+ complex_chains_indices = chain_indices_dictlist[name]
+ chain_reorder_idx = np.argsort(complex_chains_indices)
+ reordered_chains = {
+ idx: complex_chains_embeddings[i] for idx, i in enumerate(chain_reorder_idx)
+ }
+ lm_embeddings_chains_all.append(reordered_chains)
+ else:
+ lm_embeddings_chains_all = [None] * len(complex_names_all)
+
+ # running preprocessing in parallel on multiple workers and saving the progress every 1000 complexes
+ list_indices = list(range(len(complex_names_all) // 1000 + 1))
+ random.shuffle(list_indices)
+ for i in list_indices:
+ if os.path.exists(os.path.join(self.full_cache_path, f"heterographs{i}.pkl")):
+ continue
+ complex_names = complex_names_all[1000 * i : 1000 * (i + 1)]
+ lm_embeddings_chains = lm_embeddings_chains_all[1000 * i : 1000 * (i + 1)]
+ complex_graphs, rdkit_ligands = [], []
+ if self.num_workers > 1:
+ p = Pool(self.num_workers, maxtasksperchild=1)
+ p.__enter__()
+ with tqdm(
+ total=len(complex_names),
+ desc=f"Loading complexes {i}/{len(complex_names_all)//1000+1}",
+ ) as pbar:
+ map_fn = p.imap_unordered if self.num_workers > 1 else map
+ for t in map_fn(
+ self.get_complex,
+ zip(
+ complex_names,
+ lm_embeddings_chains,
+ [None] * len(complex_names),
+ [None] * len(complex_names),
+ ),
+ ):
+ if t is not None:
+ complex_graphs.extend(t[0])
+ rdkit_ligands.extend(t[1])
+ pbar.update()
+ if self.num_workers > 1:
+ p.__exit__(None, None, None)
+
+ with open(os.path.join(self.full_cache_path, f"heterographs{i}.pkl"), "wb") as f:
+ pickle.dump((complex_graphs), f)
+ with open(os.path.join(self.full_cache_path, f"rdkit_ligands{i}.pkl"), "wb") as f:
+ pickle.dump((rdkit_ligands), f)
+
+ def inference_preprocessing(self):
+ """Preprocesses the complexes for inference."""
+ ligands_list = []
+ log.info("Reading molecules and generating local structures with RDKit")
+ for ligand_description in tqdm(self.ligand_descriptions):
+ mol = MolFromSmiles(ligand_description) # check if it is a smiles or a path
+ if mol is not None:
+ mol = AddHs(mol)
+ generate_conformer(mol)
+ ligands_list.append(mol)
+ else:
+ mol = read_molecule(ligand_description, remove_hs=False, sanitize=True)
+ if not self.keep_local_structures:
+ mol.RemoveAllConformers()
+ mol = AddHs(mol)
+ generate_conformer(mol)
+ ligands_list.append(mol)
+
+ if self.esm_embeddings_path is not None:
+ log.info("Reading language model embeddings.")
+ lm_embeddings_chains_all = []
+ if not os.path.exists(self.esm_embeddings_path):
+ raise Exception("ESM embeddings path does not exist: ", self.esm_embeddings_path)
+ for protein_path in self.protein_path_list:
+ embeddings_paths = sorted(
+ glob.glob(
+ os.path.join(self.esm_embeddings_path, os.path.basename(protein_path))
+ + "*"
+ )
+ )
+ lm_embeddings_chains = []
+ for embeddings_path in embeddings_paths:
+ lm_embeddings_chains.append(torch.load(embeddings_path)["representations"][33])
+ lm_embeddings_chains_all.append(lm_embeddings_chains)
+ else:
+ lm_embeddings_chains_all = [None] * len(self.protein_path_list)
+
+ log.info("Generating graphs for ligands and proteins")
+ # running preprocessing in parallel on multiple workers and saving the progress every 1000 complexes
+ list_indices = list(range(len(self.protein_path_list) // 1000 + 1))
+ random.shuffle(list_indices)
+ for i in list_indices:
+ if os.path.exists(os.path.join(self.full_cache_path, f"heterographs{i}.pkl")):
+ continue
+ protein_paths_chunk = self.protein_path_list[1000 * i : 1000 * (i + 1)]
+ ligand_description_chunk = self.ligand_descriptions[1000 * i : 1000 * (i + 1)]
+ ligands_chunk = ligands_list[1000 * i : 1000 * (i + 1)]
+ lm_embeddings_chains = lm_embeddings_chains_all[1000 * i : 1000 * (i + 1)]
+ complex_graphs, rdkit_ligands = [], []
+ if self.num_workers > 1:
+ p = Pool(self.num_workers, maxtasksperchild=1)
+ p.__enter__()
+ with tqdm(
+ total=len(protein_paths_chunk),
+ desc=f"Loading complexes {i}/{len(protein_paths_chunk)//1000+1}",
+ ) as pbar:
+ map_fn = p.imap_unordered if self.num_workers > 1 else map
+ for t in map_fn(
+ self.get_complex,
+ zip(
+ protein_paths_chunk,
+ lm_embeddings_chains,
+ ligands_chunk,
+ ligand_description_chunk,
+ ),
+ ):
+ if t is not None:
+ complex_graphs.extend(t[0])
+ rdkit_ligands.extend(t[1])
+ pbar.update()
+ if self.num_workers > 1:
+ p.__exit__(None, None, None)
+
+ with open(os.path.join(self.full_cache_path, f"heterographs{i}.pkl"), "wb") as f:
+ pickle.dump((complex_graphs), f)
+ with open(os.path.join(self.full_cache_path, f"rdkit_ligands{i}.pkl"), "wb") as f:
+ pickle.dump((rdkit_ligands), f)
+
+ def check_all_complexes(self):
+ """Checks if all complexes are already in the cache."""
+ if os.path.exists(os.path.join(self.full_cache_path, "heterographs.pkl")):
+ return True
+
+ complex_names_all = read_strings_from_txt(self.split_path)
+ if self.limit_complexes is not None and self.limit_complexes != 0:
+ complex_names_all = complex_names_all[: self.limit_complexes]
+ for i in range(len(complex_names_all) // 1000 + 1):
+ if not os.path.exists(os.path.join(self.full_cache_path, f"heterographs{i}.pkl")):
+ return False
+ return True
+
+ def collect_all_complexes(self):
+ """Collects all complexes from the cache."""
+ log.info("Collecting all complexes from cache", self.full_cache_path)
+ if os.path.exists(os.path.join(self.full_cache_path, "heterographs.pkl")):
+ with open(os.path.join(self.full_cache_path, "heterographs.pkl"), "rb") as f:
+ complex_graphs = pickle.load(f) # nosec
+ if self.require_ligand:
+ with open(os.path.join(self.full_cache_path, "rdkit_ligands.pkl"), "rb") as f:
+ rdkit_ligands = pickle.load(f) # nosec
+ else:
+ rdkit_ligands = None
+ return complex_graphs, rdkit_ligands
+
+ complex_names_all = read_strings_from_txt(self.split_path)
+ if self.limit_complexes is not None and self.limit_complexes != 0:
+ complex_names_all = complex_names_all[: self.limit_complexes]
+ complex_graphs_all = []
+ for i in range(len(complex_names_all) // 1000 + 1):
+ with open(os.path.join(self.full_cache_path, f"heterographs{i}.pkl"), "rb") as f:
+ log.info(f"Loading heterographs{i}.pkl")
+ item = pickle.load(f) # nosec
+ complex_graphs_all.extend(item)
+
+ rdkit_ligands_all = []
+ for i in range(len(complex_names_all) // 1000 + 1):
+ with open(os.path.join(self.full_cache_path, f"rdkit_ligands{i}.pkl"), "rb") as f:
+ item = pickle.load(f) # nosec
+ rdkit_ligands_all.extend(item)
+
+ return complex_graphs_all, rdkit_ligands_all
+
+ def get_complex(self, par):
+ """Returns a list of HeteroData objects and a list of RDKit molecules for a given
+ complex."""
+ name, lm_embedding_chains, ligand, _ = par
+ if not os.path.exists(os.path.join(self.pdbbind_dir, name)) and ligand is None:
+ log.error(f"Data directory not found for {name}")
+ return [], []
+
+ try:
+ lig = read_mol(self.pdbbind_dir, name, suffix=self.ligand_file, remove_hs=False)
+ if self.max_lig_size is not None and lig.GetNumHeavyAtoms() > self.max_lig_size:
+ log.error(
+ f"Ligand with {lig.GetNumHeavyAtoms()} heavy atoms is larger than max_lig_size {self.max_lig_size}. Skipping preprocessing for this example..."
+ )
+ return [], []
+
+ if self.remove_hs:
+ lig = RemoveHs(lig)
+
+ lig_samples = [
+ process_molecule(
+ lig,
+ ref_conf_xyz=np.array(lig.GetConformer().GetPositions()),
+ return_as_dict=True,
+ )
+ ]
+
+ except Exception as e:
+ log.error(f"Skipping {name} because of error: {e}")
+ return [], []
+
+ try:
+ holo_protein_filepath = os.path.join(
+ self.pdbbind_dir, name, f"{name}_{self.protein_file}.pdb"
+ )
+ holo_af_protein = pdb_filepath_to_protein(holo_protein_filepath)
+ holo_protein_sample = process_protein(
+ holo_af_protein,
+ sample_name=f"{name}_",
+ )
+ complex_graph = merge_protein_and_ligands(
+ lig_samples,
+ holo_protein_sample,
+ n_lig_patches=self.n_lig_patches,
+ )
+
+ except Exception as e:
+ log.error(f"Skipping holo {name} because of error: {e}")
+ return [], []
+
+ if np.isnan(complex_graph["features"]["res_atom_positions"]).any():
+ log.error(
+ f"NaN in holo receptor pos for {name}. Skipping preprocessing for this example..."
+ )
+ return None
+
+ try:
+ if self.apo_protein_structure_dir is not None:
+ apo_protein_filepath = os.path.join(
+ self.apo_protein_structure_dir, f"{name}_holo_aligned_esmfold_protein.pdb"
+ )
+ apo_af_protein = pdb_filepath_to_protein(apo_protein_filepath)
+ apo_protein_sample = process_protein(
+ apo_af_protein,
+ sample_name=f"{name}_",
+ sequences_to_embeddings=lm_embedding_chains,
+ )
+ for key in complex_graph.keys():
+ for subkey, value in apo_protein_sample[key].items():
+ complex_graph[key]["apo_" + subkey] = value
+ if not np.array_equal(
+ complex_graph["features"]["res_type"],
+ complex_graph["features"]["apo_res_type"],
+ ):
+ log.error(
+ f"Residue type mismatch between holo protein and apo protein for {name}. Skipping preprocessing for this example..."
+ )
+ return None
+ if np.isnan(complex_graph["features"]["apo_res_atom_positions"]).any():
+ log.error(
+ f"NaN in apo receptor pos for {name}. Skipping preprocessing for this example..."
+ )
+ return None
+
+ except Exception as e:
+ log.error(f"Skipping apo {name} because of error: {e}")
+ return [], []
+
+ if (
+ self.min_protein_length is not None
+ and complex_graph["metadata"]["num_a"] < self.min_protein_length
+ and not self.is_test_dataset
+ ):
+ log.info(f"Skipping {name} because of its length {complex_graph['metadata']['num_a']}")
+ return [], []
+ if (
+ self.max_protein_length is not None
+ and complex_graph["metadata"]["num_a"] > self.max_protein_length
+ and not self.is_test_dataset
+ ):
+ log.info(f"Skipping {name} because of its length {complex_graph['metadata']['num_a']}")
+ return [], []
+
+ complex_graph["metadata"]["sample_ID"] = name
+ return [complex_graph], [lig]
+
+
+def read_mol(pdbbind_dir, name, suffix="ligand", remove_hs=False):
+ """Reads a ligand from the given directory and returns it as an RDKit molecule."""
+ lig = read_molecule(
+ os.path.join(pdbbind_dir, name, f"{name}_{suffix}.mol2"),
+ remove_hs=remove_hs,
+ sanitize=True,
+ )
+ if lig is None: # read sdf file if mol2 file cannot be sanitized
+ log.info(
+ "Reading the .mol2 file failed. We found a .sdf file instead and are trying to use that. Be aware that the .sdf files from PDBBind 2020 are missing chirality tags, although we will do our best to impute such information automatically using RDKit. Reference: https://www.blopig.com/blog/2021/09/watch-out-when-using-pdbbind."
+ )
+ lig = read_molecule(
+ os.path.join(pdbbind_dir, name, f"{name}_{suffix}.sdf"),
+ remove_hs=remove_hs,
+ sanitize=True,
+ )
+ Chem.rdmolops.AssignAtomChiralTagsFromStructure(lig)
+ return lig
+
+
+def read_mols(pdbbind_dir, name, remove_hs=False):
+ """Reads all ligands from the given directory and returns them as a list of RDKit molecules."""
+ ligs = []
+ for file in os.listdir(os.path.join(pdbbind_dir, name)):
+ if file.endswith(".mol2") and "rdkit" not in file:
+ lig = read_molecule(
+ os.path.join(pdbbind_dir, name, file), remove_hs=remove_hs, sanitize=True
+ )
+ if lig is None and os.path.exists(
+ os.path.join(pdbbind_dir, name, file[:-4] + ".sdf")
+ ): # read sdf file if mol2 file cannot be sanitized
+ log.info(
+ "Using the .mol2 file failed. We found a .sdf file instead and are trying to use that. Be aware that the .sdf files from PDBBind 2020 are missing chirality tags, although we will do our best to impute such information automatically using RDKit. Reference: https://www.blopig.com/blog/2021/09/watch-out-when-using-pdbbind."
+ )
+ lig = read_molecule(
+ os.path.join(pdbbind_dir, name, file[:-4] + ".sdf"),
+ remove_hs=remove_hs,
+ sanitize=True,
+ )
+ Chem.rdmolops.AssignAtomChiralTagsFromStructure(lig)
+ if lig is not None:
+ ligs.append(lig)
+ return ligs
diff --git a/forks/FlowDock/flowdock/data/components/pdbsidechain_dataset.py b/forks/FlowDock/flowdock/data/components/pdbsidechain_dataset.py
new file mode 100644
index 00000000..249412ce
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/pdbsidechain_dataset.py
@@ -0,0 +1,1028 @@
+# Adapted from: https://github.com/gcorso/DiffDock
+# Significant contribution from Ben Fry
+
+import ast
+import copy
+import os.path
+import pickle # nosec
+import random
+from multiprocessing import Pool
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import rootutils
+import torch
+from beartype.typing import Optional
+from rdkit import Chem
+from rdkit.Chem import MolFromSmiles
+from scipy.spatial.distance import pdist, squareform
+from torch_geometric.data import Dataset
+from tqdm import tqdm
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock.data.components import residue_constants
+from flowdock.data.components.constants import (
+ aa_long2short,
+ aa_to_cg_indices,
+ amino_acid_smiles,
+ cg_rdkit_indices,
+)
+from flowdock.data.components.mol_features import process_molecule
+from flowdock.data.components.process_mols import generate_conformer
+from flowdock.utils import RankedLogger
+from flowdock.utils.data_utils import (
+ atom37_to_atom14,
+ centralize_complex_graph,
+ combine_molecules,
+ convert_protein_pts_to_pdb,
+ get_mol_with_new_conformer_coords,
+ merge_protein_and_ligands,
+ min_max_normalize_array,
+ pdb_filepath_to_protein,
+ process_protein,
+)
+from flowdock.utils.model_utils import sample_inplace_to_torch
+from flowdock.utils.utils import fasta_to_dict, read_strings_from_txt
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+def compute_num_ca_neighbors(
+ coords, cg_coords, idx, is_valid_bb_node, max_dist=5, buffer_residue_num=7
+):
+ """Counts number of residues with heavy atoms within max_dist (Angstroms) of this sidechain
+ that are not residues within +/- buffer_residue_num in primary sequence.
+
+ From Ben's code
+ Note: Gabriele removed the chain_index
+ """
+
+ # Extract coordinates of all residues in the protein.
+ bb_coords = coords
+
+ # Compute the indices that we should not consider interactions.
+ excluded_neighbors = [
+ idx - x for x in reversed(range(0, buffer_residue_num + 1)) if (idx - x) >= 0
+ ]
+ excluded_neighbors.extend([idx + x for x in range(1, buffer_residue_num + 1)])
+
+ # Create indices of an N x M distance matrix where N is num BB nodes and M is num CG nodes.
+ e_idx = torch.stack(
+ [
+ torch.arange(bb_coords.shape[0])
+ .unsqueeze(-1)
+ .expand((-1, cg_coords.shape[0]))
+ .flatten(),
+ torch.arange(cg_coords.shape[0])
+ .unsqueeze(0)
+ .expand((bb_coords.shape[0], -1))
+ .flatten(),
+ ]
+ )
+
+ # Expand bb_coords and cg_coords into the same dimensionality.
+ bb_coords_exp = bb_coords[e_idx[0]]
+ cg_coords_exp = cg_coords[e_idx[1]].unsqueeze(1)
+
+ # Every row is distance of chemical group to each atom in backbone coordinate frame.
+ bb_exp_idces, _ = (torch.cdist(bb_coords_exp, cg_coords_exp).squeeze(-1) < max_dist).nonzero(
+ as_tuple=True
+ )
+ bb_idces_within_thresh = torch.unique(e_idx[0][bb_exp_idces])
+
+ # Only count residues that are not adjacent or origin in primary sequence and are valid backbone residues (fully resolved coordinate frame).
+ bb_idces_within_thresh = bb_idces_within_thresh[
+ ~torch.isin(bb_idces_within_thresh, torch.tensor(excluded_neighbors))
+ & is_valid_bb_node[bb_idces_within_thresh]
+ ]
+
+ return len(bb_idces_within_thresh)
+
+
+def identify_valid_vandermers(args):
+ """Constructs a tensor containing all the number of contacts for each residue that can be
+ sampled from for chemical groups.
+
+ By using every sidechain as a chemical group, we will load the actual chemical groups at
+ training time. These can be used to sample as probabilities once divided by the sum.
+ """
+ complex_graph, max_dist, buffer_residue_num = args
+
+ # Constructs a mask tracking whether index is a valid coordinate frame / residue label to train over.
+ # is_in_residue_vocabulary = torch.tensor([x in aa_short2long for x in data['seq']]).bool()
+ coords, seq = (
+ complex_graph["features"]["coords"],
+ complex_graph["features"]["seq"],
+ )
+ is_valid_bb_node = (
+ coords[:, :4].isnan().sum(dim=(1, 2)) == 0
+ ).bool() # * is_in_residue_vocabulary
+
+ valid_cg_idces = []
+ for idx, aa in enumerate(seq):
+ if aa not in aa_to_cg_indices:
+ valid_cg_idces.append(0)
+ else:
+ indices = aa_to_cg_indices[aa]
+ cg_coordinates = coords[idx][indices]
+
+ # remove chemical group residues that aren't fully resolved (i.e., contain NaNs or all-zero rows).
+ if (
+ torch.any(cg_coordinates.isnan()).item()
+ or torch.any(cg_coordinates.eq(0).all(dim=1)).item()
+ ):
+ valid_cg_idces.append(0)
+ continue
+
+ nbr_count = compute_num_ca_neighbors(
+ coords,
+ cg_coordinates,
+ idx,
+ is_valid_bb_node,
+ max_dist=max_dist,
+ buffer_residue_num=buffer_residue_num,
+ )
+ valid_cg_idces.append(nbr_count)
+
+ return complex_graph["metadata"]["sample_ID"], torch.tensor(valid_cg_idces)
+
+
+def fast_identify_valid_vandermers(coords, seq, max_dist=5, buffer_residue_num=7):
+ """Fast version of identify_valid_vandermers that only computes the number of neighbors for
+ each residue."""
+
+ offset = 10000 + max_dist
+ R = coords.shape[0]
+
+ coords = coords.numpy().reshape(-1, 3)
+ pdist_mat = squareform(pdist(coords))
+ pdist_mat = pdist_mat.reshape((R, 14, R, 14))
+ pdist_mat = np.nan_to_num(pdist_mat, nan=offset)
+ pdist_mat = np.min(pdist_mat, axis=(1, 3))
+
+ # compute pairwise distances
+ pdist_mat = pdist_mat + np.diag(np.ones(len(seq)) * offset)
+ for i in range(1, buffer_residue_num + 1):
+ pdist_mat += np.diag(np.ones(len(seq) - i) * offset, k=i) + np.diag(
+ np.ones(len(seq) - i) * offset, k=-i
+ )
+
+ # get number of residues that are within max_dist of each other
+ nbr_count = np.sum(pdist_mat < max_dist, axis=1)
+ return torch.tensor(nbr_count)
+
+
+def subgraph_mol(mol, atoms_to_keep):
+ """Subgraphs a molecule based on the atoms to keep."""
+ new_mol = Chem.RWMol(mol)
+ atoms_to_remove = [
+ atom.GetIdx() for atom in new_mol.GetAtoms() if atom.GetIdx() not in atoms_to_keep
+ ]
+ for atom_idx in reversed(sorted(atoms_to_remove)):
+ new_mol.RemoveAtom(atom_idx)
+ return new_mol.GetMol()
+
+
+def compute_cg_features(aa, aa_smile):
+ """Given an amino acid and a smiles string returns the stacked tensor of chemical group atom
+ encodings.
+
+ The order of the output tensor rows corresponds to the index the atoms appear in
+ aa_to_cg_indices from constants.
+ """
+
+ # Handle any residues that we don't have chemical groups for (ex: GLY if not using bb_cnh and bb_cco)
+ aa_short = aa_long2short[aa]
+ if aa_short not in aa_to_cg_indices:
+ return None
+
+ # Create rdkit molecule from smiles string.
+ mol = Chem.MolFromSmiles(aa_smile)
+
+ atoms_to_keep = [i for i, _ in cg_rdkit_indices[aa].items()]
+ mol = subgraph_mol(mol, atoms_to_keep)
+
+ return mol
+
+
+class PDBSidechainDataset(Dataset):
+ """A dataset for extracting sidechain "ligands" from PDB structures."""
+
+ def __init__(
+ self,
+ root,
+ transform=None,
+ cache_path=os.path.join("data", "cache"),
+ pdbsidechain_metadata_dir=os.path.join("data", "pdbsidechain"),
+ esm_embeddings_path=None,
+ esm_embeddings_sequences_path=None,
+ apo_protein_structure_dir=None,
+ split="train",
+ limit_complexes=0,
+ num_workers=0,
+ remove_hs=True,
+ multiplicity=1,
+ vandermers_max_dist=5,
+ vandermers_buffer_residue_num=7,
+ vandermers_min_contacts=5,
+ vandermers_max_surrogate_binding_affinity=15.0,
+ vandermers_second_ligand_max_closeness=10.0,
+ vandermers_extract_second_ligand=False,
+ merge_clusters=1,
+ vandermers_extraction=True,
+ add_random_ligand=False,
+ min_protein_length: Optional[int] = 10,
+ max_protein_length: Optional[int] = 4000,
+ is_test_dataset=False,
+ vandermers_use_prob_as_surrogate_binding_affinity=False,
+ a2h_assessment_csv_filepath=None,
+ filter_using_a2h_assessment=False,
+ a2h_min_tmscore=None,
+ a2h_max_rmsd=None,
+ a2h_min_protein_length=None,
+ a2h_max_protein_length=None,
+ postprocess_min_protein_length=None,
+ postprocess_max_protein_length=None,
+ n_lig_patches=32,
+ ):
+ """Initializes the dataset."""
+
+ super().__init__(root, transform)
+ assert remove_hs is True, "The argument `remove_hs` must be `True` for now"
+ self.root = root
+ self.split = split
+ self.pdbsidechain_metadata_dir = pdbsidechain_metadata_dir
+ self.esm_embeddings_path = esm_embeddings_path + (
+ f"_limit{limit_complexes}" if limit_complexes == 4 else ""
+ )
+ self.esm_embeddings_sequences_path = esm_embeddings_sequences_path
+ self.apo_protein_structure_dir = apo_protein_structure_dir + (
+ f"_limit{limit_complexes}" if limit_complexes == 4 else ""
+ )
+ self.limit_complexes = limit_complexes
+ self.multiplicity = multiplicity
+ self.num_workers = num_workers
+ self.vandermers_second_ligand_max_closeness = vandermers_second_ligand_max_closeness
+ self.vandermers_extract_second_ligand = vandermers_extract_second_ligand
+ self.merge_clusters = merge_clusters
+ self.vandermers_extraction = vandermers_extraction
+ self.add_random_ligand = add_random_ligand
+ self.min_protein_length = min_protein_length
+ self.max_protein_length = max_protein_length
+ self.is_test_dataset = is_test_dataset
+ self.postprocess_min_protein_length = postprocess_min_protein_length
+ self.postprocess_max_protein_length = postprocess_max_protein_length
+ self.vandermers_use_prob_as_surrogate_binding_affinity = (
+ vandermers_use_prob_as_surrogate_binding_affinity
+ )
+ self.n_lig_patches = n_lig_patches
+
+ if postprocess_min_protein_length is not None:
+ assert postprocess_min_protein_length >= 10, "`min_protein_length` must be >= 10"
+ log.info(
+ f"Postprocess-filtering out proteins with length < {postprocess_min_protein_length}"
+ )
+ if postprocess_max_protein_length is not None:
+ assert postprocess_max_protein_length <= 4000, "`max_protein_length` must be <= 4000"
+ log.info(
+ f"Postprocess-filtering out proteins with length > {postprocess_max_protein_length}"
+ )
+
+ if vandermers_extraction:
+ self.cg_node_feature_lookup_dict = {
+ aa_long2short[aa]: compute_cg_features(aa, aa_smile)
+ for aa, aa_smile in amino_acid_smiles.items()
+ }
+
+ self.cache_path = os.path.join(
+ cache_path,
+ f"PDB_limit{self.limit_complexes}_INDEX{self.split}"
+ + ("" if min_protein_length is None else f"_minProteinLength{min_protein_length}")
+ + ("" if max_protein_length is None else f"_maxProteinLength{max_protein_length}"),
+ )
+ self.read_split()
+
+ if not self.check_all_proteins():
+ os.makedirs(self.cache_path, exist_ok=True)
+ self.preprocess()
+
+ self.vandermers_max_dist = vandermers_max_dist
+ self.vandermers_buffer_residue_num = vandermers_buffer_residue_num
+ self.vandermers_min_contacts = vandermers_min_contacts
+ self.vandermers_max_surrogate_binding_affinity = vandermers_max_surrogate_binding_affinity
+ self.collect_proteins()
+
+ filtered_proteins = []
+ if vandermers_extraction:
+ for complex_graph in tqdm(self.protein_graphs):
+ if complex_graph["metadata"]["sample_ID"] in self.vandermers and torch.any(
+ self.vandermers[complex_graph["metadata"]["sample_ID"]] >= 10
+ ):
+ filtered_proteins.append(complex_graph)
+ log.info(
+ f"Computed vandermers and kept {len(filtered_proteins)} proteins out of {len(self.protein_graphs)}"
+ )
+ else:
+ filtered_proteins = self.protein_graphs
+
+ second_filter = []
+ for complex_graph in tqdm(filtered_proteins):
+ second_filter.append(complex_graph)
+ log.info(
+ f"Checked embeddings available and kept {len(second_filter)} proteins out of {len(filtered_proteins)}"
+ )
+
+ self.protein_graphs = second_filter
+
+ # filter clusters that have no protein graphs
+ self.split_clusters = list({g["features"]["cluster"] for g in self.protein_graphs})
+ self.cluster_to_complexes = {c: [] for c in self.split_clusters}
+ for p in self.protein_graphs:
+ self.cluster_to_complexes[p["features"]["cluster"]].append(p)
+ self.split_clusters = [
+ c for c in self.split_clusters if len(self.cluster_to_complexes[c]) > 0
+ ]
+ log.info(
+ f"Total elements in set: {len(self.split_clusters) * self.multiplicity // self.merge_clusters}",
+ )
+
+ # analyze and potentially filter the van der Mers dataset based on its apo-to-holo (a2h) structural assessment
+ if a2h_assessment_csv_filepath is not None and os.path.exists(a2h_assessment_csv_filepath):
+ a2h_assessment_df = pd.read_csv(a2h_assessment_csv_filepath)
+ a2h_assessment_df["ID"] = [
+ "_".join(Path(paths[0]).stem.split("_")[:2])
+ for paths in a2h_assessment_df["Filepath"].apply(ast.literal_eval).tolist()
+ ]
+ # import matplotlib.pyplot as plt
+ # import seaborn as sns
+ # plot_dir = Path(a2h_assessment_csv_filepath).parent / "plots"
+ # plot_dir.mkdir(exist_ok=True)
+ # plt.clf()
+ # sns.histplot(a2h_assessment_df["TM-score"])
+ # plt.title("Apo-To-Holo Protein TM-score")
+ # plt.savefig(plot_dir / "a2h_TM-score_hist.png")
+ # plt.clf()
+ # sns.histplot(a2h_assessment_df["RMSD"])
+ # plt.title("Apo-To-Holo Protein RMSD")
+ # plt.savefig(plot_dir / "a2h_RMSD_hist.png")
+ # plt.clf()
+ # plt.xlim(0, 1000)
+ # sns.histplot(a2h_assessment_df["Apo_Length"])
+ # plt.title("Apo Protein Length")
+ # plt.savefig(plot_dir / "apo_length_hist.png")
+ if filter_using_a2h_assessment and not is_test_dataset:
+ log.info(
+ f"Filtering the van der Mers {self.split} dataset based on its apo-to-holo (a2h) structural assessment"
+ )
+ a2h_assessment_df = a2h_assessment_df[
+ (a2h_assessment_df["TM-score"] >= a2h_min_tmscore)
+ & (a2h_assessment_df["RMSD"] <= a2h_max_rmsd)
+ & (a2h_assessment_df["Apo_Length"] >= a2h_min_protein_length)
+ & (a2h_assessment_df["Apo_Length"] <= a2h_max_protein_length)
+ ]
+ a2h_filtered_ids = {id: None for id in a2h_assessment_df["ID"].tolist()}
+ third_filter = [
+ p
+ for p in self.protein_graphs
+ if p["metadata"]["sample_ID"] in a2h_filtered_ids
+ ]
+ self.protein_graphs = third_filter
+
+ # filter clusters that have no protein graphs
+ self.split_clusters = list({g["features"]["cluster"] for g in self.protein_graphs})
+ self.cluster_to_complexes = {c: [] for c in self.split_clusters}
+ for p in self.protein_graphs:
+ self.cluster_to_complexes[p["features"]["cluster"]].append(p)
+ self.split_clusters = [
+ c for c in self.split_clusters if len(self.cluster_to_complexes[c]) > 0
+ ]
+ log.info(
+ f"Total elements in set after `a2h` filtering: {len(self.split_clusters) * self.multiplicity // self.merge_clusters}",
+ )
+
+ self.name_to_complex = {p["metadata"]["sample_ID"]: p for p in self.protein_graphs}
+ log.info(
+ f"{len(self.name_to_complex)} total protein chains available from {len(self.split_clusters) * self.multiplicity // self.merge_clusters} clusters after all {self.split} filtering"
+ )
+ with open(
+ os.path.join(
+ self.cache_path,
+ f"pdbsidechain_{self.split}_names.txt",
+ ),
+ "w",
+ ) as f:
+ f.write("\n".join(list(self.name_to_complex.keys())))
+
+ self.define_probabilities()
+
+ if self.add_random_ligand:
+ raise NotImplementedError("add_random_ligand is not implemented yet")
+ assert os.path.exists(
+ os.path.join(self.pdbsidechain_metadata_dir, "smiles_list.csv")
+ ), "smiles_list.csv not found"
+ # read csv with all smiles
+ with open(os.path.join(self.pdbsidechain_metadata_dir, "smiles_list.csv")) as f:
+ self.smiles_list = f.readlines()
+ self.smiles_list = [s.split(",")[0] for s in self.smiles_list]
+
+ def define_probabilities(self):
+ """Defines the probabilities for each residue to be selected as a sidechain."""
+ if not self.vandermers_extraction:
+ return
+
+ if self.vandermers_min_contacts is not None:
+ self.probabilities = torch.arange(1000) - self.vandermers_min_contacts + 1
+ self.probabilities[: self.vandermers_min_contacts] = 0
+ else:
+ assert os.path.exists(
+ os.path.join(self.pdbsidechain_metadata_dir, "pdbbind_counts.pkl")
+ ), "pdbbind_counts.pkl not found"
+ with open(
+ os.path.join(self.pdbsidechain_metadata_dir, "pdbbind_counts.pkl"), "rb"
+ ) as f:
+ pdbbind_counts = pickle.load(f) # nosec
+
+ pdb_counts = torch.ones(1000)
+ for contacts in self.vandermers.values():
+ pdb_counts.index_add_(0, contacts, torch.ones(contacts.shape))
+ log.info(f"pdbbind_counts[:30]: {pdbbind_counts[:30]}")
+ log.info(f"pdb_counts[:30]: {pdb_counts[:30]}")
+
+ self.probabilities = pdbbind_counts / pdb_counts
+ self.probabilities[:7] = 0
+
+ def len(self):
+ """Returns the number of proteins in the dataset."""
+ return len(self.split_clusters) * self.multiplicity // self.merge_clusters
+
+ def get(self, idx=None, protein=None, smiles=None):
+ """Returns a protein graph from the dataset."""
+ assert idx is not None or (
+ protein is not None and smiles is not None
+ ), "provide idx or protein or smile"
+
+ if protein is None or smiles is None:
+ idx = idx % len(self.split_clusters)
+ if self.merge_clusters > 1:
+ idx = idx * self.merge_clusters
+ idx = idx + random.randint(0, self.merge_clusters - 1) # nosec
+ idx = min(idx, len(self.split_clusters) - 1)
+ cluster = self.split_clusters[idx]
+ protein_samples = copy.deepcopy(
+ random.choice(self.cluster_to_complexes[cluster]) # nosec
+ )
+ else:
+ protein_samples = copy.deepcopy(self.name_to_complex[protein])
+
+ if self.vandermers_extraction:
+ # select sidechain to remove
+ vandermers_contacts = self.vandermers[protein_samples["metadata"]["sample_ID"]]
+ vandermers_probs = self.probabilities[vandermers_contacts].numpy()
+ original_vandermers_probs = vandermers_probs.copy()
+
+ if not np.any(vandermers_contacts.numpy() >= 10):
+ new_idx = random.randint(0, self.len()) # nosec
+ log.warning(
+ f"No vandermers with >= 10 contacts found. Retrying with new example from cluster with index {new_idx}."
+ )
+ return self.get(new_idx)
+
+ sidechain_idx = np.random.choice(
+ np.arange(len(vandermers_probs)), p=vandermers_probs / np.sum(vandermers_probs)
+ )
+
+ # remove part of the sequence
+ residues_to_keep = np.ones(len(protein_samples["features"]["seq"]), dtype=bool)
+ residues_to_keep[
+ max(0, sidechain_idx - self.vandermers_buffer_residue_num) : min(
+ sidechain_idx + self.vandermers_buffer_residue_num + 1,
+ len(protein_samples["features"]["seq"]),
+ )
+ ] = False
+
+ if self.vandermers_extract_second_ligand:
+ pos_idx = protein_samples["features"]["coords"][sidechain_idx, 1, :]
+ far_enough = (
+ torch.sum(
+ (protein_samples["features"]["coords"][..., 1, :] - pos_idx[None, :]) ** 2,
+ dim=-1,
+ )
+ > self.vandermers_second_ligand_max_closeness**2
+ )
+ vandermers_probs = vandermers_probs * far_enough.float().numpy()
+ vandermers_probs[
+ max(0, sidechain_idx - self.vandermers_buffer_residue_num) : min(
+ sidechain_idx + self.vandermers_buffer_residue_num + 1,
+ len(protein_samples["features"]["seq"]),
+ )
+ ] = 0
+ if np.all(vandermers_probs <= 0):
+ new_idx = random.randint(0, self.len()) # nosec
+ log.warning(
+ f"No second vandermer available. Retrying with new example from cluster with index {new_idx}."
+ )
+ return self.get(new_idx)
+ sc2_idx = np.random.choice(
+ np.arange(len(vandermers_probs)),
+ p=vandermers_probs / np.sum(vandermers_probs),
+ )
+
+ residues_to_keep[
+ max(0, sc2_idx - self.vandermers_buffer_residue_num) : min(
+ sc2_idx + self.vandermers_buffer_residue_num + 1,
+ len(protein_samples["features"]["seq"]),
+ )
+ ] = False
+
+ # create the sidechain ligand
+ sidechain_aa = protein_samples["features"]["seq"][sidechain_idx]
+ lig_list = [self.cg_node_feature_lookup_dict[sidechain_aa]]
+ ligand_pos_list = [
+ protein_samples["features"]["coords"][sidechain_idx][
+ protein_samples["features"]["mask"][sidechain_idx]
+ ]
+ ]
+ sidechain_indices = [sidechain_idx]
+ if self.vandermers_extract_second_ligand:
+ sidechain_aa2 = protein_samples["features"]["seq"][sc2_idx]
+ lig_list.append(self.cg_node_feature_lookup_dict[sidechain_aa2])
+ ligand_pos_list.append(
+ protein_samples["features"]["coords"][sc2_idx][
+ protein_samples["features"]["mask"][sc2_idx]
+ ]
+ )
+ sidechain_indices.append(sc2_idx)
+ try:
+ lig_samples = []
+ for lig, ligand_pos in zip(lig_list, ligand_pos_list):
+ lig_mol = get_mol_with_new_conformer_coords(
+ lig,
+ ligand_pos.numpy() if torch.is_tensor(ligand_pos) else ligand_pos,
+ )
+ lig_samples.append(
+ process_molecule(
+ lig_mol,
+ ref_conf_xyz=np.array(lig_mol.GetConformer().GetPositions()),
+ return_as_dict=True,
+ )
+ )
+ except Exception as e:
+ new_idx = random.randint(0, self.len()) # nosec
+ log.warning(
+ f"Vandermer could not be constructed for the cluster with index {idx} due to: {e}. Retrying with new example from cluster with index {new_idx}."
+ )
+ return self.get(new_idx)
+
+ # subset the protein features to keep only the residues that are not part of the selected sidechain
+ residues_to_keep = torch.from_numpy(residues_to_keep)
+ num_residues = residues_to_keep.shape[0]
+ for key in protein_samples:
+ for subkey in protein_samples[key].keys():
+ if (
+ isinstance(protein_samples[key][subkey], int)
+ and protein_samples[key][subkey] == num_residues
+ ):
+ protein_samples[key][subkey] = residues_to_keep.sum().item()
+ elif (
+ isinstance(protein_samples[key][subkey], np.ndarray)
+ and protein_samples[key][subkey].shape[0] == num_residues
+ ):
+ protein_samples[key][subkey] = protein_samples[key][subkey][
+ residues_to_keep
+ ]
+ elif (
+ isinstance(protein_samples[key][subkey], list)
+ and len(protein_samples[key][subkey]) > 0
+ and isinstance(protein_samples[key][subkey][0], tuple)
+ and len(protein_samples[key][subkey][0]) == 2
+ and isinstance(protein_samples[key][subkey][0][0], str)
+ and isinstance(protein_samples[key][subkey][0][1], str)
+ ):
+ new_tuples_list = []
+ for tuples_list in protein_samples[key][subkey]:
+ if len(tuples_list[1]) == num_residues:
+ new_tuple = (
+ tuples_list[0],
+ "".join(np.array(list(tuples_list[1]))[residues_to_keep]),
+ )
+ else:
+ new_tuple = tuples_list
+ new_tuples_list.append(new_tuple)
+ protein_samples[key][subkey] = new_tuples_list
+
+ for key in [
+ "coords",
+ "seq",
+ "mask",
+ "cluster",
+ "orig_seq",
+ "to_keep",
+ ]:
+ # clean up features prior to batching
+ if key in protein_samples["features"]:
+ del protein_samples["features"][key]
+ complex_graph = merge_protein_and_ligands(
+ lig_samples,
+ protein_samples,
+ n_lig_patches=self.n_lig_patches,
+ )
+ else:
+ lig_list = [None]
+ complex_graph = protein_samples
+
+ if self.add_random_ligand:
+ raise NotImplementedError("add_random_ligand is not implemented yet")
+ if smiles is not None:
+ mol = MolFromSmiles(smiles)
+ try:
+ generate_conformer(mol)
+ except Exception as e:
+ log.error(f"Failed to generate the given ligand, returning `None` due to: {e}")
+ return None
+ else:
+ success = False
+ while not success:
+ smiles = random.choice(self.smiles_list) # nosec
+ mol = MolFromSmiles(smiles)
+ try:
+ success = not generate_conformer(mol)
+ except Exception as e:
+ log.error(f"Changing ligand due to: {e}")
+
+ lig = get_mol_with_new_conformer_coords(
+ mol,
+ mol.GetConformer().GetPositions()
+ - np.mean(mol.GetConformer().GetPositions(), axis=0, keepdims=True),
+ )
+ lig_samples = [
+ process_molecule(
+ lig,
+ ref_conf_xyz=np.array(lig.GetConformer().GetPositions()),
+ return_as_dict=True,
+ )
+ ]
+ for key in [
+ "coords",
+ "seq",
+ "mask",
+ "cluster",
+ "orig_seq",
+ "to_keep",
+ ]:
+ # clean up features prior to batching
+ if key in protein_samples["features"]:
+ del protein_samples["features"][key]
+ complex_graph = merge_protein_and_ligands(
+ lig_samples,
+ protein_samples,
+ n_lig_patches=self.n_lig_patches,
+ )
+
+ complex_graph["metadata"]["mol"] = combine_molecules(lig_list) if all(lig_list) else None
+ complex_graph = sample_inplace_to_torch(complex_graph)
+
+ affinity = [torch.nan for _ in range(len(lig_list))]
+ if self.vandermers_extraction and self.vandermers_use_prob_as_surrogate_binding_affinity:
+ normalized_vandermers_probs = min_max_normalize_array(
+ original_vandermers_probs / np.sum(original_vandermers_probs)
+ )
+ affinity = [
+ normalized_vandermers_probs[sidechain_index]
+ * self.vandermers_max_surrogate_binding_affinity
+ for sidechain_index in sidechain_indices
+ ]
+ complex_graph["features"]["affinity"] = torch.tensor(affinity, dtype=torch.float32)
+
+ return centralize_complex_graph(complex_graph)
+
+ def read_split(self):
+ """Reads the split from the CSV file."""
+ # read CSV file
+ df = pd.read_csv(os.path.join(self.root, "list.csv"))
+ log.info("Loaded list CSV file")
+
+ # get clusters and filter by split
+ if self.split == "train":
+ val_clusters = set(
+ read_strings_from_txt(os.path.join(self.root, "valid_clusters.txt"))
+ )
+ test_clusters = set(
+ read_strings_from_txt(os.path.join(self.root, "test_clusters.txt"))
+ )
+ clusters = df["CLUSTER"].unique()
+ clusters = [
+ int(c) for c in clusters if c not in val_clusters and c not in test_clusters
+ ]
+ elif self.split == "val":
+ clusters = [
+ int(s)
+ for s in read_strings_from_txt(os.path.join(self.root, "valid_clusters.txt"))
+ ]
+ elif self.split == "test":
+ clusters = [
+ int(s) for s in read_strings_from_txt(os.path.join(self.root, "test_clusters.txt"))
+ ]
+ else:
+ raise ValueError("Split must be train, val or test")
+ log.info(f"{self.split} clusters: {len(clusters)}")
+ clusters = set(clusters)
+
+ self.chains_in_cluster = []
+ complexes_in_cluster = set()
+ for chain, cluster in zip(df["CHAINID"], df["CLUSTER"]):
+ if cluster not in clusters:
+ continue
+ # limit to one chain per complex
+ if chain[:4] not in complexes_in_cluster:
+ self.chains_in_cluster.append((chain, cluster))
+ complexes_in_cluster.add(chain[:4])
+ log.info(f"Filtered chains in cluster: {len(self.chains_in_cluster)}")
+
+ if self.limit_complexes > 0:
+ self.chains_in_cluster = self.chains_in_cluster[: self.limit_complexes]
+
+ def check_all_proteins(self):
+ """Checks if all proteins have been preprocessed."""
+ for i in range(len(self.chains_in_cluster) // 10000 + 1):
+ if not os.path.exists(os.path.join(self.cache_path, f"protein_graphs{i}.pkl")):
+ return False
+ return True
+
+ def collect_proteins(self):
+ """Collects all proteins from the cache."""
+ self.protein_graphs = []
+ self.vandermers = {}
+ total_recovered = 0
+ log.info(f"Loading {len(self.chains_in_cluster)} protein graphs.")
+ list_indices = list(range(len(self.chains_in_cluster) // 10000 + 1))
+ random.shuffle(list_indices)
+ for i in list_indices:
+ with open(os.path.join(self.cache_path, f"protein_graphs{i}.pkl"), "rb") as f:
+ log.info(f"Loading protein graphs {i}")
+ item = pickle.load(f) # nosec
+ if (
+ self.postprocess_min_protein_length is not None
+ and self.postprocess_max_protein_length is not None
+ ):
+ item = [
+ sample
+ for sample in item
+ if self.postprocess_min_protein_length
+ <= len(sample["features"]["res_type"])
+ <= self.postprocess_max_protein_length
+ ]
+ elif self.postprocess_min_protein_length is not None:
+ item = [
+ sample
+ for sample in item
+ if len(sample["features"]["res_type"])
+ >= self.postprocess_min_protein_length
+ ]
+ elif self.postprocess_max_protein_length is not None:
+ item = [
+ sample
+ for sample in item
+ if len(sample["features"]["res_type"])
+ <= self.postprocess_max_protein_length
+ ]
+ total_recovered += len(item)
+ self.protein_graphs.extend(item)
+
+ if not self.vandermers_extraction:
+ continue
+
+ if os.path.exists(
+ os.path.join(
+ self.cache_path,
+ f"vandermers{i}_{self.vandermers_max_dist}_{self.vandermers_buffer_residue_num}.pkl",
+ )
+ ):
+ with open(
+ os.path.join(
+ self.cache_path,
+ f"vandermers{i}_{self.vandermers_max_dist}_{self.vandermers_buffer_residue_num}.pkl",
+ ),
+ "rb",
+ ) as f:
+ vandermers = pickle.load(f) # nosec
+ self.vandermers.update(vandermers)
+ continue
+
+ vandermers = {}
+ if self.num_workers > 1:
+ p = Pool(self.num_workers, maxtasksperchild=1)
+ p.__enter__()
+ with tqdm(total=len(item), desc=f"Computing vandermers {i}") as pbar:
+ map_fn = p.imap_unordered if self.num_workers > 1 else map
+ arguments = zip(
+ item,
+ [self.vandermers_max_dist] * len(item),
+ [self.vandermers_buffer_residue_num] * len(item),
+ )
+ for t in map_fn(identify_valid_vandermers, arguments):
+ if t is not None:
+ vandermers[t[0]] = t[1]
+ pbar.update()
+ if self.num_workers > 1:
+ p.__exit__(None, None, None)
+
+ with open(
+ os.path.join(
+ self.cache_path,
+ f"vandermers{i}_{self.vandermers_max_dist}_{self.vandermers_buffer_residue_num}.pkl",
+ ),
+ "wb",
+ ) as f:
+ pickle.dump(vandermers, f)
+ self.vandermers.update(vandermers)
+
+ log.info(
+ f"Kept {len(self.protein_graphs)} proteins out of {len(self.chains_in_cluster)} total"
+ )
+
+ def preprocess(self):
+ """Preprocesses all proteins in the dataset."""
+ if self.esm_embeddings_path is not None:
+ log.info("Loading ESM embeddings")
+ sequences_to_embeddings = {}
+ receptor_names_all_dict = {chain[0]: None for chain in self.chains_in_cluster}
+ sequences_dict = fasta_to_dict(self.esm_embeddings_sequences_path)
+ for embedding_filepath in os.listdir(self.esm_embeddings_path):
+ key = Path(embedding_filepath).stem
+ chain_name = key.split("_chain_")[0]
+ chain_id = key.split("_chain_")[1]
+ full_chain_id = f"{chain_name}_{chain_id}"
+ if full_chain_id in receptor_names_all_dict:
+ embedding = torch.load(
+ os.path.join(self.esm_embeddings_path, embedding_filepath)
+ )["representations"][33]
+ seq = sequences_dict[key]
+ sequences_to_embeddings[seq + f":{chain_id}"] = embedding
+ else:
+ sequences_to_embeddings = None
+
+ # running preprocessing in parallel on multiple workers and saving the progress every 10000 proteins
+ list_indices = list(range(len(self.chains_in_cluster) // 10000 + 1))
+ random.shuffle(list_indices)
+ for i in list_indices:
+ if os.path.exists(os.path.join(self.cache_path, f"protein_graphs{i}.pkl")):
+ continue
+ chains_names = self.chains_in_cluster[10000 * i : 10000 * (i + 1)]
+ protein_graphs = []
+ if self.num_workers > 1:
+ p = Pool(self.num_workers, maxtasksperchild=1)
+ p.__enter__()
+ with tqdm(
+ total=len(chains_names),
+ desc=f"Loading protein batch {i}/{len(self.chains_in_cluster) // 10000 + 1}",
+ ) as pbar:
+ map_fn = p.imap_unordered if self.num_workers > 1 else map
+ for t in map_fn(
+ self.load_chain,
+ zip(chains_names, [sequences_to_embeddings] * len(chains_names)),
+ ):
+ if t is not None:
+ protein_graphs.append(t)
+ pbar.update()
+ if self.num_workers > 1:
+ p.__exit__(None, None, None)
+
+ with open(os.path.join(self.cache_path, f"protein_graphs{i}.pkl"), "wb") as f:
+ pickle.dump(protein_graphs, f)
+
+ log.info("Finished preprocessing and saving protein graphs")
+
+ def load_chain(self, c):
+ """Loads a protein graph from a chain name."""
+ chain, cluster = c[0]
+ sequences_to_embeddings = c[1]
+ chain_id = chain.split("_")[1]
+ pt_filepath = os.path.join(self.root, "pdb", chain[1:3], f"{chain}.pt")
+ holo_protein_filepath = os.path.join(self.root, "pdb", chain[1:3], f"{chain}.pdb")
+ if not os.path.exists(pt_filepath):
+ log.error(f"PyTorch file not found: {chain}")
+ return None
+ if not os.path.exists(holo_protein_filepath):
+ try:
+ convert_protein_pts_to_pdb([pt_filepath], holo_protein_filepath)
+ if not os.path.exists(holo_protein_filepath):
+ log.error(
+ f"Failed to convert PyTorch {pt_filepath} to PDB {holo_protein_filepath}"
+ )
+ return None
+
+ except Exception as e:
+ log.error(
+ f"Error in converting {pt_filepath} to {holo_protein_filepath} due to: {e}"
+ )
+ return None
+
+ try:
+ holo_af_protein = pdb_filepath_to_protein(holo_protein_filepath)
+ holo_protein_sample = process_protein(
+ holo_af_protein,
+ sample_name=f"{chain.split('_')[0]}_",
+ )
+ complex_graph = holo_protein_sample
+
+ except Exception as e:
+ log.error(f"Error in extracting holo receptor {chain}: {e}")
+ return None
+
+ if np.isnan(complex_graph["features"]["res_atom_positions"]).any():
+ log.error(
+ f"NaN in holo receptor pos for {chain}. Skipping preprocessing for this example..."
+ )
+ return None
+
+ try:
+ if self.apo_protein_structure_dir is not None:
+ apo_protein_filepath = os.path.join(
+ self.apo_protein_structure_dir, f"{chain}_holo_aligned_esmfold_protein.pdb"
+ )
+ apo_af_protein = pdb_filepath_to_protein(apo_protein_filepath)
+ apo_protein_sample = process_protein(
+ apo_af_protein,
+ sample_name=f"{chain.split('_')[0]}_",
+ sequences_to_embeddings=sequences_to_embeddings,
+ chain_id=chain_id,
+ )
+ for key in complex_graph.keys():
+ for subkey, value in apo_protein_sample[key].items():
+ complex_graph[key]["apo_" + subkey] = value
+ if not np.array_equal(
+ complex_graph["features"]["res_type"],
+ complex_graph["features"]["apo_res_type"],
+ ):
+ log.error(
+ f"Residue type mismatch between holo protein and apo protein for {chain}. Skipping preprocessing for this example..."
+ )
+ return None
+ if np.isnan(complex_graph["features"]["apo_res_atom_positions"]).any():
+ log.error(
+ f"NaN in apo receptor pos for {chain}. Skipping preprocessing for this example..."
+ )
+ return None
+
+ except Exception as e:
+ log.error(f"Skipping apo {chain} because of the error: {e}")
+ return None
+
+ if (
+ self.min_protein_length is not None
+ and complex_graph["metadata"]["num_a"] < self.min_protein_length
+ and not self.is_test_dataset
+ ):
+ log.info(
+ f"Skipping {chain} because of its length {complex_graph['metadata']['num_a']}"
+ )
+ return None
+ if (
+ self.max_protein_length is not None
+ and complex_graph["metadata"]["num_a"] > self.max_protein_length
+ and not self.is_test_dataset
+ ):
+ log.info(
+ f"Skipping {chain} because of its length {complex_graph['metadata']['num_a']}"
+ )
+ return None
+
+ # NOTE: can manually confirm that OpenFold's `atom14` representation perfectly matches RoseTTAFold2's `atom14` representation
+ # Also NOTE: RoseTTAFold2's `atom27` representation linked below is actually its `atom14` representation plus 13 hydrogen atoms per residue (starting at index 14)
+ # OpenFold reference code: https://github.com/aqlaboratory/openfold/blob/80c85b54e1a81d9a66df3f1b6c257ff97f10acd3/openfold/np/residue_constants.py#L601
+ # RoseTTAFold2 reference code: https://github.com/uw-ipd/RoseTTAFold2/blob/a26f29bb92d4946ef8bf225edba9240c9e647747/network/chemical.py#L17
+ atom14_coords, atom14_mask = atom37_to_atom14(
+ torch.from_numpy(complex_graph["features"]["res_type"]),
+ torch.from_numpy(complex_graph["features"]["res_atom_positions"]),
+ torch.from_numpy(complex_graph["features"]["res_atom_mask"]),
+ )
+ complex_seq = "".join(
+ [
+ residue_constants.restypes[res_idx]
+ for res_idx in complex_graph["features"]["res_type"]
+ ]
+ )
+
+ complex_graph["features"]["coords"] = atom14_coords
+ complex_graph["features"]["mask"] = atom14_mask.bool()
+ complex_graph["features"]["seq"] = complex_seq
+ complex_graph["features"]["cluster"] = cluster
+ complex_graph["metadata"]["sample_ID"] = chain
+
+ return complex_graph
+
+
+if __name__ == "__main__":
+ dataset = PDBSidechainDataset(
+ root=os.path.join("data", "pdbsidechain", "pdb_2021aug02"),
+ split="train",
+ multiplicity=1,
+ limit_complexes=150,
+ )
+ log.info(f"Length of dataset: {len(dataset)}")
+ log.info(f"First dataset example: {dataset[0]}")
+ for p in dataset:
+ log.info(f"Example: {p['metadata']['sample_ID']}")
+ break
diff --git a/forks/FlowDock/flowdock/data/components/physical.py b/forks/FlowDock/flowdock/data/components/physical.py
new file mode 100644
index 00000000..5e5dc2d7
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/physical.py
@@ -0,0 +1,84 @@
+# Adapted from: https://github.com/zrqiao/NeuralPLexer
+
+import os
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from mendeleev.fetch import fetch_table
+
+ptable = fetch_table("elements")
+PATOMIC_NUMBERS = {row["symbol"]: row["atomic_number"] for _, row in ptable.iterrows()}
+PGROUP_IDS = {row["symbol"]: row["group_id"] for _, row in ptable.iterrows()}
+PPERIOD_IDS = {row["symbol"]: row["period"] for _, row in ptable.iterrows()}
+VDW_RADII_MAPPING = {row["atomic_number"]: row["vdw_radius"] for _, row in ptable.iterrows()}
+UFF_VDW_RADII_MAPPING = {
+ row["atomic_number"]: row["vdw_radius_uff"] for _, row in ptable.iterrows()
+}
+
+
+def get_vdw_radii_array():
+ """Get van der Waals radii array."""
+ return np.array(
+ [VDW_RADII_MAPPING[i] if i in VDW_RADII_MAPPING.keys() else -1.0 for i in range(119)]
+ )
+
+
+def get_vdw_radii_array_uff():
+ """Get UFF van der Waals radii array."""
+ return np.array(
+ [
+ UFF_VDW_RADII_MAPPING[i] if i in UFF_VDW_RADII_MAPPING.keys() else -1.0
+ for i in range(119)
+ ]
+ )
+
+
+uff_params_df = (
+ pd.read_csv(
+ os.path.join(
+ Path(__file__).parent.absolute(),
+ "uff_parameters.csv",
+ ),
+ )
+ .drop_duplicates("Element")
+ .set_index("Element")
+)
+
+
+def calc_heavy_atom_LJ_clash_fraction(mol1, mol2, LJ_CUTOFF=100):
+ """Calculate the fraction of heavy atoms in mol1 that have a Lennard-Jones interaction energy
+ greater than LJ_CUTOFF with any atom in mol2."""
+ D_ij, x_ij = get_epsilon_sigma_uff(mol1, mol2)
+ coords1 = np.array(mol1.GetConformer().GetPositions())
+ coords2 = np.array(mol2.GetConformer().GetPositions())
+ distance_matrix = np.linalg.norm(
+ coords1[:, np.newaxis, :] - coords2[np.newaxis, :, :],
+ axis=2,
+ )
+ pairwise_LJ = D_ij * ((x_ij / distance_matrix) ** 12 - 2 * (x_ij / distance_matrix) ** 6)
+ return np.sum(np.sum(pairwise_LJ, axis=1) > LJ_CUTOFF) / pairwise_LJ.shape[0]
+
+
+def calc_heavy_atom_LJ_int_energy(mol1, mol2):
+ """Calculate the Lennard-Jones interaction energy between two molecules."""
+ D_ij, x_ij = get_epsilon_sigma_uff(mol1, mol2)
+ coords1 = np.array(mol1.GetConformer().GetPositions())
+ coords2 = np.array(mol2.GetConformer().GetPositions())
+ distance_matrix = np.linalg.norm(
+ coords1[:, np.newaxis, :] - coords2[np.newaxis, :, :],
+ axis=2,
+ )
+ pairwise_LJ = D_ij * ((x_ij / distance_matrix) ** 12 - 2 * (x_ij / distance_matrix) ** 6)
+ return np.sum(pairwise_LJ)
+
+
+def get_epsilon_sigma_uff(m1, m2):
+ """Get epsilon and sigma for Lorentz-Berthelot mixing rule."""
+ Di = [float(uff_params_df.loc[atom.GetSymbol(), "energy"]) for atom in m1.GetAtoms()]
+ Dj = [float(uff_params_df.loc[atom.GetSymbol(), "energy"]) for atom in m2.GetAtoms()]
+ xi = [float(uff_params_df.loc[atom.GetSymbol(), "distance"]) for atom in m1.GetAtoms()]
+ xj = [float(uff_params_df.loc[atom.GetSymbol(), "distance"]) for atom in m2.GetAtoms()]
+ x_ij = (np.array(xi)[:, np.newaxis] + np.array(xj)[np.newaxis, :]) / 2
+ D_ij = (np.array(Di)[:, np.newaxis] * np.array(Dj)[np.newaxis, :]) ** (1 / 2)
+ return D_ij, x_ij
diff --git a/forks/FlowDock/flowdock/data/components/process_mols.py b/forks/FlowDock/flowdock/data/components/process_mols.py
new file mode 100644
index 00000000..dcd238dd
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/process_mols.py
@@ -0,0 +1,76 @@
+# From: https://github.com/gcorso/DiffDock
+
+import warnings
+
+import rootutils
+from rdkit import Chem
+from rdkit.Chem import AllChem
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock.utils import RankedLogger
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+def generate_conformer(mol):
+ """Generate conformer for the ligand."""
+ ps = AllChem.ETKDGv2()
+ failures, id = 0, -1
+ while failures < 3 and id == -1:
+ if failures > 0:
+ log.warning(f"rdkit coords could not be generated. trying again {failures}.")
+ id = AllChem.EmbedMolecule(mol, ps)
+ failures += 1
+ if id == -1:
+ log.warning(
+ "rdkit coords could not be generated without using random coords. using random coords now."
+ )
+ ps.useRandomCoords = True
+ AllChem.EmbedMolecule(mol, ps)
+ AllChem.MMFFOptimizeMolecule(mol, confId=0)
+ return True
+ # else:
+ # AllChem.MMFFOptimizeMolecule(mol, confId=0)
+ return False
+
+
+def read_molecule(molecule_file, sanitize=False, calc_charges=False, remove_hs=False):
+ """Read molecule from file and return RDKit molecule object."""
+ if molecule_file.endswith(".mol2"):
+ mol = Chem.MolFromMol2File(molecule_file, sanitize=False, removeHs=False)
+ elif molecule_file.endswith(".sdf"):
+ supplier = Chem.SDMolSupplier(molecule_file, sanitize=False, removeHs=False)
+ mol = supplier[0]
+ elif molecule_file.endswith(".pdbqt"):
+ with open(molecule_file) as file:
+ pdbqt_data = file.readlines()
+ pdb_block = ""
+ for line in pdbqt_data:
+ pdb_block += f"{line[:66]}\n"
+ mol = Chem.MolFromPDBBlock(pdb_block, sanitize=False, removeHs=False)
+ elif molecule_file.endswith(".pdb"):
+ mol = Chem.MolFromPDBFile(molecule_file, sanitize=False, removeHs=False)
+ else:
+ return ValueError(
+ "Expect the format of the molecule_file to be "
+ "one of .mol2, .sdf, .pdbqt and .pdb, got {}".format(molecule_file)
+ )
+
+ try:
+ if sanitize or calc_charges:
+ Chem.SanitizeMol(mol)
+
+ if calc_charges:
+ # Compute Gasteiger charges on the molecule.
+ try:
+ AllChem.ComputeGasteigerCharges(mol)
+ except Exception as e:
+ warnings.warn("Unable to compute charges for the molecule.")
+
+ if remove_hs:
+ mol = Chem.RemoveHs(mol, sanitize=sanitize)
+ except Exception as e:
+ return None
+
+ return mol
diff --git a/forks/FlowDock/flowdock/data/components/residue_constants.py b/forks/FlowDock/flowdock/data/components/residue_constants.py
new file mode 100644
index 00000000..5dee23d9
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/residue_constants.py
@@ -0,0 +1,1107 @@
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Constants used in AlphaFold."""
+
+import collections
+import functools
+import os
+from typing import List, Mapping, Tuple
+
+import numpy as np
+import tree
+
+# Internal import (35fd).
+
+
+# Distance from one CA to next CA [trans configuration: omega = 180].
+ca_ca = 3.80209737096
+
+# Format: The list for each AA type contains chi1, chi2, chi3, chi4 in
+# this order (or a relevant subset from chi1 onwards). ALA and GLY don't have
+# chi angles so their chi angle lists are empty.
+chi_angles_atoms = {
+ "ALA": [],
+ # Chi5 in arginine is always 0 +- 5 degrees, so ignore it.
+ "ARG": [
+ ["N", "CA", "CB", "CG"],
+ ["CA", "CB", "CG", "CD"],
+ ["CB", "CG", "CD", "NE"],
+ ["CG", "CD", "NE", "CZ"],
+ ],
+ "ASN": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]],
+ "ASP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]],
+ "CYS": [["N", "CA", "CB", "SG"]],
+ "GLN": [
+ ["N", "CA", "CB", "CG"],
+ ["CA", "CB", "CG", "CD"],
+ ["CB", "CG", "CD", "OE1"],
+ ],
+ "GLU": [
+ ["N", "CA", "CB", "CG"],
+ ["CA", "CB", "CG", "CD"],
+ ["CB", "CG", "CD", "OE1"],
+ ],
+ "GLY": [],
+ "HIS": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "ND1"]],
+ "ILE": [["N", "CA", "CB", "CG1"], ["CA", "CB", "CG1", "CD1"]],
+ "LEU": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
+ "LYS": [
+ ["N", "CA", "CB", "CG"],
+ ["CA", "CB", "CG", "CD"],
+ ["CB", "CG", "CD", "CE"],
+ ["CG", "CD", "CE", "NZ"],
+ ],
+ "MET": [
+ ["N", "CA", "CB", "CG"],
+ ["CA", "CB", "CG", "SD"],
+ ["CB", "CG", "SD", "CE"],
+ ],
+ "PHE": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
+ "PRO": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"]],
+ "SER": [["N", "CA", "CB", "OG"]],
+ "THR": [["N", "CA", "CB", "OG1"]],
+ "TRP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
+ "TYR": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
+ "VAL": [["N", "CA", "CB", "CG1"]],
+}
+
+# If chi angles given in fixed-length array, this matrix determines how to mask
+# them for each AA type. The order is as per restype_order (see below).
+chi_angles_mask = [
+ [0.0, 0.0, 0.0, 0.0], # ALA
+ [1.0, 1.0, 1.0, 1.0], # ARG
+ [1.0, 1.0, 0.0, 0.0], # ASN
+ [1.0, 1.0, 0.0, 0.0], # ASP
+ [1.0, 0.0, 0.0, 0.0], # CYS
+ [1.0, 1.0, 1.0, 0.0], # GLN
+ [1.0, 1.0, 1.0, 0.0], # GLU
+ [0.0, 0.0, 0.0, 0.0], # GLY
+ [1.0, 1.0, 0.0, 0.0], # HIS
+ [1.0, 1.0, 0.0, 0.0], # ILE
+ [1.0, 1.0, 0.0, 0.0], # LEU
+ [1.0, 1.0, 1.0, 1.0], # LYS
+ [1.0, 1.0, 1.0, 0.0], # MET
+ [1.0, 1.0, 0.0, 0.0], # PHE
+ [1.0, 1.0, 0.0, 0.0], # PRO
+ [1.0, 0.0, 0.0, 0.0], # SER
+ [1.0, 0.0, 0.0, 0.0], # THR
+ [1.0, 1.0, 0.0, 0.0], # TRP
+ [1.0, 1.0, 0.0, 0.0], # TYR
+ [1.0, 0.0, 0.0, 0.0], # VAL
+]
+
+# The following chi angles are pi periodic: they can be rotated by a multiple
+# of pi without affecting the structure.
+chi_pi_periodic = [
+ [0.0, 0.0, 0.0, 0.0], # ALA
+ [0.0, 0.0, 0.0, 0.0], # ARG
+ [0.0, 0.0, 0.0, 0.0], # ASN
+ [0.0, 1.0, 0.0, 0.0], # ASP
+ [0.0, 0.0, 0.0, 0.0], # CYS
+ [0.0, 0.0, 0.0, 0.0], # GLN
+ [0.0, 0.0, 1.0, 0.0], # GLU
+ [0.0, 0.0, 0.0, 0.0], # GLY
+ [0.0, 0.0, 0.0, 0.0], # HIS
+ [0.0, 0.0, 0.0, 0.0], # ILE
+ [0.0, 0.0, 0.0, 0.0], # LEU
+ [0.0, 0.0, 0.0, 0.0], # LYS
+ [0.0, 0.0, 0.0, 0.0], # MET
+ [0.0, 1.0, 0.0, 0.0], # PHE
+ [0.0, 0.0, 0.0, 0.0], # PRO
+ [0.0, 0.0, 0.0, 0.0], # SER
+ [0.0, 0.0, 0.0, 0.0], # THR
+ [0.0, 0.0, 0.0, 0.0], # TRP
+ [0.0, 1.0, 0.0, 0.0], # TYR
+ [0.0, 0.0, 0.0, 0.0], # VAL
+ [0.0, 0.0, 0.0, 0.0], # UNK
+]
+
+# Atoms positions relative to the 8 rigid groups, defined by the pre-omega, phi,
+# psi and chi angles:
+# 0: 'backbone group',
+# 1: 'pre-omega-group', (empty)
+# 2: 'phi-group', (currently empty, because it defines only hydrogens)
+# 3: 'psi-group',
+# 4,5,6,7: 'chi1,2,3,4-group'
+# The atom positions are relative to the axis-end-atom of the corresponding
+# rotation axis. The x-axis is in direction of the rotation axis, and the y-axis
+# is defined such that the dihedral-angle-definiting atom (the last entry in
+# chi_angles_atoms above) is in the xy-plane (with a positive y-coordinate).
+# format: [atomname, group_idx, rel_position]
+rigid_group_atom_positions = {
+ "ALA": [
+ ["N", 0, (-0.525, 1.363, 0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.526, -0.000, -0.000)],
+ ["CB", 0, (-0.529, -0.774, -1.205)],
+ ["O", 3, (0.627, 1.062, 0.000)],
+ ],
+ "ARG": [
+ ["N", 0, (-0.524, 1.362, -0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.525, -0.000, -0.000)],
+ ["CB", 0, (-0.524, -0.778, -1.209)],
+ ["O", 3, (0.626, 1.062, 0.000)],
+ ["CG", 4, (0.616, 1.390, -0.000)],
+ ["CD", 5, (0.564, 1.414, 0.000)],
+ ["NE", 6, (0.539, 1.357, -0.000)],
+ ["NH1", 7, (0.206, 2.301, 0.000)],
+ ["NH2", 7, (2.078, 0.978, -0.000)],
+ ["CZ", 7, (0.758, 1.093, -0.000)],
+ ],
+ "ASN": [
+ ["N", 0, (-0.536, 1.357, 0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.526, -0.000, -0.000)],
+ ["CB", 0, (-0.531, -0.787, -1.200)],
+ ["O", 3, (0.625, 1.062, 0.000)],
+ ["CG", 4, (0.584, 1.399, 0.000)],
+ ["ND2", 5, (0.593, -1.188, 0.001)],
+ ["OD1", 5, (0.633, 1.059, 0.000)],
+ ],
+ "ASP": [
+ ["N", 0, (-0.525, 1.362, -0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.527, 0.000, -0.000)],
+ ["CB", 0, (-0.526, -0.778, -1.208)],
+ ["O", 3, (0.626, 1.062, -0.000)],
+ ["CG", 4, (0.593, 1.398, -0.000)],
+ ["OD1", 5, (0.610, 1.091, 0.000)],
+ ["OD2", 5, (0.592, -1.101, -0.003)],
+ ],
+ "CYS": [
+ ["N", 0, (-0.522, 1.362, -0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.524, 0.000, 0.000)],
+ ["CB", 0, (-0.519, -0.773, -1.212)],
+ ["O", 3, (0.625, 1.062, -0.000)],
+ ["SG", 4, (0.728, 1.653, 0.000)],
+ ],
+ "GLN": [
+ ["N", 0, (-0.526, 1.361, -0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.526, 0.000, 0.000)],
+ ["CB", 0, (-0.525, -0.779, -1.207)],
+ ["O", 3, (0.626, 1.062, -0.000)],
+ ["CG", 4, (0.615, 1.393, 0.000)],
+ ["CD", 5, (0.587, 1.399, -0.000)],
+ ["NE2", 6, (0.593, -1.189, -0.001)],
+ ["OE1", 6, (0.634, 1.060, 0.000)],
+ ],
+ "GLU": [
+ ["N", 0, (-0.528, 1.361, 0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.526, -0.000, -0.000)],
+ ["CB", 0, (-0.526, -0.781, -1.207)],
+ ["O", 3, (0.626, 1.062, 0.000)],
+ ["CG", 4, (0.615, 1.392, 0.000)],
+ ["CD", 5, (0.600, 1.397, 0.000)],
+ ["OE1", 6, (0.607, 1.095, -0.000)],
+ ["OE2", 6, (0.589, -1.104, -0.001)],
+ ],
+ "GLY": [
+ ["N", 0, (-0.572, 1.337, 0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.517, -0.000, -0.000)],
+ ["O", 3, (0.626, 1.062, -0.000)],
+ ],
+ "HIS": [
+ ["N", 0, (-0.527, 1.360, 0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.525, 0.000, 0.000)],
+ ["CB", 0, (-0.525, -0.778, -1.208)],
+ ["O", 3, (0.625, 1.063, 0.000)],
+ ["CG", 4, (0.600, 1.370, -0.000)],
+ ["CD2", 5, (0.889, -1.021, 0.003)],
+ ["ND1", 5, (0.744, 1.160, -0.000)],
+ ["CE1", 5, (2.030, 0.851, 0.002)],
+ ["NE2", 5, (2.145, -0.466, 0.004)],
+ ],
+ "ILE": [
+ ["N", 0, (-0.493, 1.373, -0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.527, -0.000, -0.000)],
+ ["CB", 0, (-0.536, -0.793, -1.213)],
+ ["O", 3, (0.627, 1.062, -0.000)],
+ ["CG1", 4, (0.534, 1.437, -0.000)],
+ ["CG2", 4, (0.540, -0.785, -1.199)],
+ ["CD1", 5, (0.619, 1.391, 0.000)],
+ ],
+ "LEU": [
+ ["N", 0, (-0.520, 1.363, 0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.525, -0.000, -0.000)],
+ ["CB", 0, (-0.522, -0.773, -1.214)],
+ ["O", 3, (0.625, 1.063, -0.000)],
+ ["CG", 4, (0.678, 1.371, 0.000)],
+ ["CD1", 5, (0.530, 1.430, -0.000)],
+ ["CD2", 5, (0.535, -0.774, 1.200)],
+ ],
+ "LYS": [
+ ["N", 0, (-0.526, 1.362, -0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.526, 0.000, 0.000)],
+ ["CB", 0, (-0.524, -0.778, -1.208)],
+ ["O", 3, (0.626, 1.062, -0.000)],
+ ["CG", 4, (0.619, 1.390, 0.000)],
+ ["CD", 5, (0.559, 1.417, 0.000)],
+ ["CE", 6, (0.560, 1.416, 0.000)],
+ ["NZ", 7, (0.554, 1.387, 0.000)],
+ ],
+ "MET": [
+ ["N", 0, (-0.521, 1.364, -0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.525, 0.000, 0.000)],
+ ["CB", 0, (-0.523, -0.776, -1.210)],
+ ["O", 3, (0.625, 1.062, -0.000)],
+ ["CG", 4, (0.613, 1.391, -0.000)],
+ ["SD", 5, (0.703, 1.695, 0.000)],
+ ["CE", 6, (0.320, 1.786, -0.000)],
+ ],
+ "PHE": [
+ ["N", 0, (-0.518, 1.363, 0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.524, 0.000, -0.000)],
+ ["CB", 0, (-0.525, -0.776, -1.212)],
+ ["O", 3, (0.626, 1.062, -0.000)],
+ ["CG", 4, (0.607, 1.377, 0.000)],
+ ["CD1", 5, (0.709, 1.195, -0.000)],
+ ["CD2", 5, (0.706, -1.196, 0.000)],
+ ["CE1", 5, (2.102, 1.198, -0.000)],
+ ["CE2", 5, (2.098, -1.201, -0.000)],
+ ["CZ", 5, (2.794, -0.003, -0.001)],
+ ],
+ "PRO": [
+ ["N", 0, (-0.566, 1.351, -0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.527, -0.000, 0.000)],
+ ["CB", 0, (-0.546, -0.611, -1.293)],
+ ["O", 3, (0.621, 1.066, 0.000)],
+ ["CG", 4, (0.382, 1.445, 0.0)],
+ # ['CD', 5, (0.427, 1.440, 0.0)],
+ ["CD", 5, (0.477, 1.424, 0.0)], # manually made angle 2 degrees larger
+ ],
+ "SER": [
+ ["N", 0, (-0.529, 1.360, -0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.525, -0.000, -0.000)],
+ ["CB", 0, (-0.518, -0.777, -1.211)],
+ ["O", 3, (0.626, 1.062, -0.000)],
+ ["OG", 4, (0.503, 1.325, 0.000)],
+ ],
+ "THR": [
+ ["N", 0, (-0.517, 1.364, 0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.526, 0.000, -0.000)],
+ ["CB", 0, (-0.516, -0.793, -1.215)],
+ ["O", 3, (0.626, 1.062, 0.000)],
+ ["CG2", 4, (0.550, -0.718, -1.228)],
+ ["OG1", 4, (0.472, 1.353, 0.000)],
+ ],
+ "TRP": [
+ ["N", 0, (-0.521, 1.363, 0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.525, -0.000, 0.000)],
+ ["CB", 0, (-0.523, -0.776, -1.212)],
+ ["O", 3, (0.627, 1.062, 0.000)],
+ ["CG", 4, (0.609, 1.370, -0.000)],
+ ["CD1", 5, (0.824, 1.091, 0.000)],
+ ["CD2", 5, (0.854, -1.148, -0.005)],
+ ["CE2", 5, (2.186, -0.678, -0.007)],
+ ["CE3", 5, (0.622, -2.530, -0.007)],
+ ["NE1", 5, (2.140, 0.690, -0.004)],
+ ["CH2", 5, (3.028, -2.890, -0.013)],
+ ["CZ2", 5, (3.283, -1.543, -0.011)],
+ ["CZ3", 5, (1.715, -3.389, -0.011)],
+ ],
+ "TYR": [
+ ["N", 0, (-0.522, 1.362, 0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.524, -0.000, -0.000)],
+ ["CB", 0, (-0.522, -0.776, -1.213)],
+ ["O", 3, (0.627, 1.062, -0.000)],
+ ["CG", 4, (0.607, 1.382, -0.000)],
+ ["CD1", 5, (0.716, 1.195, -0.000)],
+ ["CD2", 5, (0.713, -1.194, -0.001)],
+ ["CE1", 5, (2.107, 1.200, -0.002)],
+ ["CE2", 5, (2.104, -1.201, -0.003)],
+ ["OH", 5, (4.168, -0.002, -0.005)],
+ ["CZ", 5, (2.791, -0.001, -0.003)],
+ ],
+ "VAL": [
+ ["N", 0, (-0.494, 1.373, -0.000)],
+ ["CA", 0, (0.000, 0.000, 0.000)],
+ ["C", 0, (1.527, -0.000, -0.000)],
+ ["CB", 0, (-0.533, -0.795, -1.213)],
+ ["O", 3, (0.627, 1.062, -0.000)],
+ ["CG1", 4, (0.540, 1.429, -0.000)],
+ ["CG2", 4, (0.533, -0.776, 1.203)],
+ ],
+}
+
+# A list of atoms (excluding hydrogen) for each AA type. PDB naming convention.
+residue_atoms = {
+ "ALA": ["C", "CA", "CB", "N", "O"],
+ "ARG": ["C", "CA", "CB", "CG", "CD", "CZ", "N", "NE", "O", "NH1", "NH2"],
+ "ASP": ["C", "CA", "CB", "CG", "N", "O", "OD1", "OD2"],
+ "ASN": ["C", "CA", "CB", "CG", "N", "ND2", "O", "OD1"],
+ "CYS": ["C", "CA", "CB", "N", "O", "SG"],
+ "GLU": ["C", "CA", "CB", "CG", "CD", "N", "O", "OE1", "OE2"],
+ "GLN": ["C", "CA", "CB", "CG", "CD", "N", "NE2", "O", "OE1"],
+ "GLY": ["C", "CA", "N", "O"],
+ "HIS": ["C", "CA", "CB", "CG", "CD2", "CE1", "N", "ND1", "NE2", "O"],
+ "ILE": ["C", "CA", "CB", "CG1", "CG2", "CD1", "N", "O"],
+ "LEU": ["C", "CA", "CB", "CG", "CD1", "CD2", "N", "O"],
+ "LYS": ["C", "CA", "CB", "CG", "CD", "CE", "N", "NZ", "O"],
+ "MET": ["C", "CA", "CB", "CG", "CE", "N", "O", "SD"],
+ "PHE": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O"],
+ "PRO": ["C", "CA", "CB", "CG", "CD", "N", "O"],
+ "SER": ["C", "CA", "CB", "N", "O", "OG"],
+ "THR": ["C", "CA", "CB", "CG2", "N", "O", "OG1"],
+ "TRP": [
+ "C",
+ "CA",
+ "CB",
+ "CG",
+ "CD1",
+ "CD2",
+ "CE2",
+ "CE3",
+ "CZ2",
+ "CZ3",
+ "CH2",
+ "N",
+ "NE1",
+ "O",
+ ],
+ "TYR": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O", "OH"],
+ "VAL": ["C", "CA", "CB", "CG1", "CG2", "N", "O"],
+}
+
+# Naming swaps for ambiguous atom names.
+# Due to symmetries in the amino acids the naming of atoms is ambiguous in
+# 4 of the 20 amino acids.
+# (The LDDT paper lists 7 amino acids as ambiguous, but the naming ambiguities
+# in LEU, VAL and ARG can be resolved by using the 3d constellations of
+# the 'ambiguous' atoms and their neighbours)
+residue_atom_renaming_swaps = {
+ "ASP": {"OD1": "OD2"},
+ "GLU": {"OE1": "OE2"},
+ "PHE": {"CD1": "CD2", "CE1": "CE2"},
+ "TYR": {"CD1": "CD2", "CE1": "CE2"},
+}
+
+# Van der Waals radii [Angstroem] of the atoms (from Wikipedia)
+van_der_waals_radius = {
+ "C": 1.7,
+ "N": 1.55,
+ "O": 1.52,
+ "S": 1.8,
+}
+
+Bond = collections.namedtuple("Bond", ["atom1_name", "atom2_name", "length", "stddev"])
+BondAngle = collections.namedtuple(
+ "BondAngle", ["atom1_name", "atom2_name", "atom3name", "angle_rad", "stddev"]
+)
+
+
+@functools.lru_cache(maxsize=None)
+def load_stereo_chemical_props() -> (
+ Tuple[
+ Mapping[str, List[Bond]],
+ Mapping[str, List[Bond]],
+ Mapping[str, List[BondAngle]],
+ ]
+):
+ """Load stereo_chemical_props.txt into a nice structure.
+
+ Load literature values for bond lengths and bond angles and translate
+ bond angles into the length of the opposite edge of the triangle
+ ("residue_virtual_bonds").
+
+ Returns:
+ residue_bonds: Dict that maps resname -> list of Bond tuples.
+ residue_virtual_bonds: Dict that maps resname -> list of Bond tuples.
+ residue_bond_angles: Dict that maps resname -> list of BondAngle tuples.
+ """
+ stereo_chemical_props_path = os.path.join(
+ os.path.dirname(os.path.abspath(__file__)), "stereo_chemical_props.txt"
+ )
+ with open(stereo_chemical_props_path) as f:
+ stereo_chemical_props = f.read()
+ lines_iter = iter(stereo_chemical_props.splitlines())
+ # Load bond lengths.
+ residue_bonds = {}
+ next(lines_iter) # Skip header line.
+ for line in lines_iter:
+ if line.strip() == "-":
+ break
+ bond, resname, length, stddev = line.split()
+ atom1, atom2 = bond.split("-")
+ if resname not in residue_bonds:
+ residue_bonds[resname] = []
+ residue_bonds[resname].append(Bond(atom1, atom2, float(length), float(stddev)))
+ residue_bonds["UNK"] = []
+
+ # Load bond angles.
+ residue_bond_angles = {}
+ next(lines_iter) # Skip empty line.
+ next(lines_iter) # Skip header line.
+ for line in lines_iter:
+ if line.strip() == "-":
+ break
+ bond, resname, angle_degree, stddev_degree = line.split()
+ atom1, atom2, atom3 = bond.split("-")
+ if resname not in residue_bond_angles:
+ residue_bond_angles[resname] = []
+ residue_bond_angles[resname].append(
+ BondAngle(
+ atom1,
+ atom2,
+ atom3,
+ float(angle_degree) / 180.0 * np.pi,
+ float(stddev_degree) / 180.0 * np.pi,
+ )
+ )
+ residue_bond_angles["UNK"] = []
+
+ def make_bond_key(atom1_name, atom2_name):
+ """Unique key to lookup bonds."""
+ return "-".join(sorted([atom1_name, atom2_name]))
+
+ # Translate bond angles into distances ("virtual bonds").
+ residue_virtual_bonds = {}
+ for resname, bond_angles in residue_bond_angles.items():
+ # Create a fast lookup dict for bond lengths.
+ bond_cache = {}
+ for b in residue_bonds[resname]:
+ bond_cache[make_bond_key(b.atom1_name, b.atom2_name)] = b
+ residue_virtual_bonds[resname] = []
+ for ba in bond_angles:
+ bond1 = bond_cache[make_bond_key(ba.atom1_name, ba.atom2_name)]
+ bond2 = bond_cache[make_bond_key(ba.atom2_name, ba.atom3name)]
+
+ # Compute distance between atom1 and atom3 using the law of cosines
+ # c^2 = a^2 + b^2 - 2ab*cos(gamma).
+ gamma = ba.angle_rad
+ length = np.sqrt(
+ bond1.length**2
+ + bond2.length**2
+ - 2 * bond1.length * bond2.length * np.cos(gamma)
+ )
+
+ # Propagation of uncertainty assuming uncorrelated errors.
+ dl_outer = 0.5 / length
+ dl_dgamma = (2 * bond1.length * bond2.length * np.sin(gamma)) * dl_outer
+ dl_db1 = (2 * bond1.length - 2 * bond2.length * np.cos(gamma)) * dl_outer
+ dl_db2 = (2 * bond2.length - 2 * bond1.length * np.cos(gamma)) * dl_outer
+ stddev = np.sqrt(
+ (dl_dgamma * ba.stddev) ** 2
+ + (dl_db1 * bond1.stddev) ** 2
+ + (dl_db2 * bond2.stddev) ** 2
+ )
+ residue_virtual_bonds[resname].append(
+ Bond(ba.atom1_name, ba.atom3name, length, stddev)
+ )
+
+ return (residue_bonds, residue_virtual_bonds, residue_bond_angles)
+
+
+# Between-residue bond lengths for general bonds (first element) and for Proline
+# (second element).
+between_res_bond_length_c_n = [1.329, 1.341]
+between_res_bond_length_stddev_c_n = [0.014, 0.016]
+
+# Between-residue cos_angles.
+between_res_cos_angles_c_n_ca = [-0.5203, 0.0353] # degrees: 121.352 +- 2.315
+between_res_cos_angles_ca_c_n = [-0.4473, 0.0311] # degrees: 116.568 +- 1.995
+
+# This mapping is used when we need to store atom data in a format that requires
+# fixed atom data size for every residue (e.g. a numpy array).
+atom_types = [
+ "N",
+ "CA",
+ "C",
+ "CB",
+ "O",
+ "CG",
+ "CG1",
+ "CG2",
+ "OG",
+ "OG1",
+ "SG",
+ "CD",
+ "CD1",
+ "CD2",
+ "ND1",
+ "ND2",
+ "OD1",
+ "OD2",
+ "SD",
+ "CE",
+ "CE1",
+ "CE2",
+ "CE3",
+ "NE",
+ "NE1",
+ "NE2",
+ "OE1",
+ "OE2",
+ "CH2",
+ "NH1",
+ "NH2",
+ "OH",
+ "CZ",
+ "CZ2",
+ "CZ3",
+ "NZ",
+ "OXT",
+]
+atom_order = {atom_type: i for i, atom_type in enumerate(atom_types)}
+atom_type_num = len(atom_types) # := 37.
+
+atom_padding_length = 15 # Allow for terminal atoms
+element_id = {
+ "C": 6,
+ "N": 7,
+ "O": 8,
+ "S": 16,
+}
+
+# A compact atom encoding with 14 columns
+# pylint: disable=line-too-long
+# pylint: disable=bad-whitespace
+restype_name_to_atom14_names = {
+ "ALA": ["N", "CA", "C", "O", "CB", "", "", "", "", "", "", "", "", ""],
+ "ARG": [
+ "N",
+ "CA",
+ "C",
+ "O",
+ "CB",
+ "CG",
+ "CD",
+ "NE",
+ "CZ",
+ "NH1",
+ "NH2",
+ "",
+ "",
+ "",
+ ],
+ "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2", "", "", "", "", "", ""],
+ "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2", "", "", "", "", "", ""],
+ "CYS": ["N", "CA", "C", "O", "CB", "SG", "", "", "", "", "", "", "", ""],
+ "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2", "", "", "", "", ""],
+ "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2", "", "", "", "", ""],
+ "GLY": ["N", "CA", "C", "O", "", "", "", "", "", "", "", "", "", ""],
+ "HIS": [
+ "N",
+ "CA",
+ "C",
+ "O",
+ "CB",
+ "CG",
+ "ND1",
+ "CD2",
+ "CE1",
+ "NE2",
+ "",
+ "",
+ "",
+ "",
+ ],
+ "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1", "", "", "", "", "", ""],
+ "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "", "", "", "", "", ""],
+ "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ", "", "", "", "", ""],
+ "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE", "", "", "", "", "", ""],
+ "PHE": [
+ "N",
+ "CA",
+ "C",
+ "O",
+ "CB",
+ "CG",
+ "CD1",
+ "CD2",
+ "CE1",
+ "CE2",
+ "CZ",
+ "",
+ "",
+ "",
+ ],
+ "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD", "", "", "", "", "", "", ""],
+ "SER": ["N", "CA", "C", "O", "CB", "OG", "", "", "", "", "", "", "", ""],
+ "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2", "", "", "", "", "", "", ""],
+ "TRP": [
+ "N",
+ "CA",
+ "C",
+ "O",
+ "CB",
+ "CG",
+ "CD1",
+ "CD2",
+ "NE1",
+ "CE2",
+ "CE3",
+ "CZ2",
+ "CZ3",
+ "CH2",
+ ],
+ "TYR": [
+ "N",
+ "CA",
+ "C",
+ "O",
+ "CB",
+ "CG",
+ "CD1",
+ "CD2",
+ "CE1",
+ "CE2",
+ "CZ",
+ "OH",
+ "",
+ "",
+ ],
+ "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "", "", "", "", "", "", ""],
+ "UNK": ["", "", "", "", "", "", "", "", "", "", "", "", "", ""],
+}
+# pylint: enable=line-too-long
+# pylint: enable=bad-whitespace
+
+
+# This is the standard residue order when coding AA type as a number.
+# Reproduce it by taking 3-letter AA codes and sorting them alphabetically.
+restypes = [
+ "A",
+ "R",
+ "N",
+ "D",
+ "C",
+ "Q",
+ "E",
+ "G",
+ "H",
+ "I",
+ "L",
+ "K",
+ "M",
+ "F",
+ "P",
+ "S",
+ "T",
+ "W",
+ "Y",
+ "V",
+]
+restype_order = {restype: i for i, restype in enumerate(restypes)}
+restype_num = len(restypes) # := 20.
+unk_restype_index = restype_num # Catch-all index for unknown restypes.
+
+restypes_with_x = restypes + ["X"]
+restype_order_with_x = {restype: i for i, restype in enumerate(restypes_with_x)}
+
+
+def sequence_to_onehot(
+ sequence: str, mapping: Mapping[str, int], map_unknown_to_x: bool = False
+) -> np.ndarray:
+ """Maps the given sequence into a one-hot encoded matrix.
+
+ Args:
+ sequence: An amino acid sequence.
+ mapping: A dictionary mapping amino acids to integers.
+ map_unknown_to_x: If True, any amino acid that is not in the mapping will be
+ mapped to the unknown amino acid 'X'. If the mapping doesn't contain
+ amino acid 'X', an error will be thrown. If False, any amino acid not in
+ the mapping will throw an error.
+
+ Returns:
+ A numpy array of shape (seq_len, num_unique_aas) with one-hot encoding of
+ the sequence.
+
+ Raises:
+ ValueError: If the mapping doesn't contain values from 0 to
+ num_unique_aas - 1 without any gaps.
+ """
+ num_entries = max(mapping.values()) + 1
+
+ if sorted(set(mapping.values())) != list(range(num_entries)):
+ raise ValueError(
+ "The mapping must have values from 0 to num_unique_aas-1 "
+ "without any gaps. Got: %s" % sorted(mapping.values())
+ )
+
+ one_hot_arr = np.zeros((len(sequence), num_entries), dtype=np.int32)
+
+ for aa_index, aa_type in enumerate(sequence):
+ if map_unknown_to_x:
+ if aa_type.isalpha() and aa_type.isupper():
+ aa_id = mapping.get(aa_type, mapping["X"])
+ else:
+ raise ValueError(f"Invalid character in the sequence: {aa_type}")
+ else:
+ aa_id = mapping[aa_type]
+ one_hot_arr[aa_index, aa_id] = 1
+
+ return one_hot_arr
+
+
+restype_1to3 = {
+ "A": "ALA",
+ "R": "ARG",
+ "N": "ASN",
+ "D": "ASP",
+ "C": "CYS",
+ "Q": "GLN",
+ "E": "GLU",
+ "G": "GLY",
+ "H": "HIS",
+ "I": "ILE",
+ "L": "LEU",
+ "K": "LYS",
+ "M": "MET",
+ "F": "PHE",
+ "P": "PRO",
+ "S": "SER",
+ "T": "THR",
+ "W": "TRP",
+ "Y": "TYR",
+ "V": "VAL",
+}
+
+
+# NB: restype_3to1 differs from Bio.PDB.protein_letters_3to1 by being a simple
+# 1-to-1 mapping of 3 letter names to one letter names. The latter contains
+# many more, and less common, three letter names as keys and maps many of these
+# to the same one letter name (including 'X' and 'U' which we don't use here).
+restype_3to1 = {v: k for k, v in restype_1to3.items()}
+
+# Define a restype name for all unknown residues.
+unk_restype = "UNK"
+
+resnames = [restype_1to3[r] for r in restypes] + [unk_restype]
+resname_to_idx = {resname: i for i, resname in enumerate(resnames)}
+
+
+# The mapping here uses hhblits convention, so that B is mapped to D, J and O
+# are mapped to X, U is mapped to C, and Z is mapped to E. Other than that the
+# remaining 20 amino acids are kept in alphabetical order.
+# There are 2 non-amino acid codes, X (representing any amino acid) and
+# "-" representing a missing amino acid in an alignment. The id for these
+# codes is put at the end (20 and 21) so that they can easily be ignored if
+# desired.
+HHBLITS_AA_TO_ID = {
+ "A": 0,
+ "B": 2,
+ "C": 1,
+ "D": 2,
+ "E": 3,
+ "F": 4,
+ "G": 5,
+ "H": 6,
+ "I": 7,
+ "J": 20,
+ "K": 8,
+ "L": 9,
+ "M": 10,
+ "N": 11,
+ "O": 20,
+ "P": 12,
+ "Q": 13,
+ "R": 14,
+ "S": 15,
+ "T": 16,
+ "U": 1,
+ "V": 17,
+ "W": 18,
+ "X": 20,
+ "Y": 19,
+ "Z": 3,
+ "-": 21,
+}
+
+# Partial inversion of HHBLITS_AA_TO_ID.
+ID_TO_HHBLITS_AA = {
+ 0: "A",
+ 1: "C", # Also U.
+ 2: "D", # Also B.
+ 3: "E", # Also Z.
+ 4: "F",
+ 5: "G",
+ 6: "H",
+ 7: "I",
+ 8: "K",
+ 9: "L",
+ 10: "M",
+ 11: "N",
+ 12: "P",
+ 13: "Q",
+ 14: "R",
+ 15: "S",
+ 16: "T",
+ 17: "V",
+ 18: "W",
+ 19: "Y",
+ 20: "X", # Includes J and O.
+ 21: "-",
+}
+
+restypes_with_x_and_gap = restypes + ["X", "-"]
+MAP_HHBLITS_AATYPE_TO_OUR_AATYPE = tuple(
+ restypes_with_x_and_gap.index(ID_TO_HHBLITS_AA[i]) for i in range(len(restypes_with_x_and_gap))
+)
+
+
+def _make_standard_atom_mask() -> np.ndarray:
+ """Returns [num_res_types, num_atom_types] mask array."""
+ # +1 to account for unknown (all 0s).
+ mask = np.zeros([restype_num + 1, atom_type_num], dtype=np.int32)
+ for restype, restype_letter in enumerate(restypes):
+ restype_name = restype_1to3[restype_letter]
+ atom_names = residue_atoms[restype_name]
+ for atom_name in atom_names:
+ atom_type = atom_order[atom_name]
+ mask[restype, atom_type] = 1
+ return mask
+
+
+STANDARD_ATOM_MASK = _make_standard_atom_mask()
+
+
+# A one hot representation for the first and second atoms defining the axis
+# of rotation for each chi-angle in each residue.
+def chi_angle_atom(atom_index: int) -> np.ndarray:
+ """Define chi-angle rigid groups via one-hot representations."""
+ chi_angles_index = {}
+ one_hots = []
+
+ for k, v in chi_angles_atoms.items():
+ indices = [atom_types.index(s[atom_index]) for s in v]
+ indices.extend([-1] * (4 - len(indices)))
+ chi_angles_index[k] = indices
+
+ for r in restypes:
+ res3 = restype_1to3[r]
+ one_hot = np.eye(atom_type_num)[chi_angles_index[res3]]
+ one_hots.append(one_hot)
+
+ one_hots.append(np.zeros([4, atom_type_num])) # Add zeros for residue `X`.
+ one_hot = np.stack(one_hots, axis=0)
+ one_hot = np.transpose(one_hot, [0, 2, 1])
+
+ return one_hot
+
+
+chi_atom_1_one_hot = chi_angle_atom(1)
+chi_atom_2_one_hot = chi_angle_atom(2)
+
+# An array like chi_angles_atoms but using indices rather than names.
+chi_angles_atom_indices = [chi_angles_atoms[restype_1to3[r]] for r in restypes]
+chi_angles_atom_indices = tree.map_structure(
+ lambda atom_name: atom_order[atom_name], chi_angles_atom_indices
+)
+chi_angles_atom_indices = np.array(
+ [chi_atoms + ([[0, 0, 0, 0]] * (4 - len(chi_atoms))) for chi_atoms in chi_angles_atom_indices]
+)
+
+# Mapping from (res_name, atom_name) pairs to the atom's chi group index
+# and atom index within that group.
+chi_groups_for_atom = collections.defaultdict(list)
+for res_name, chi_angle_atoms_for_res in chi_angles_atoms.items():
+ for chi_group_i, chi_group in enumerate(chi_angle_atoms_for_res):
+ for atom_i, atom in enumerate(chi_group):
+ chi_groups_for_atom[(res_name, atom)].append((chi_group_i, atom_i))
+chi_groups_for_atom = dict(chi_groups_for_atom)
+
+
+def _make_rigid_transformation_4x4(ex, ey, translation):
+ """Create a rigid 4x4 transformation matrix from two axes and transl."""
+ # Normalize ex.
+ ex_normalized = ex / np.linalg.norm(ex)
+
+ # make ey perpendicular to ex
+ ey_normalized = ey - np.dot(ey, ex_normalized) * ex_normalized
+ ey_normalized /= np.linalg.norm(ey_normalized)
+
+ # compute ez as cross product
+ eznorm = np.cross(ex_normalized, ey_normalized)
+ m = np.stack([ex_normalized, ey_normalized, eznorm, translation]).transpose()
+ m = np.concatenate([m, [[0.0, 0.0, 0.0, 1.0]]], axis=0)
+ return m
+
+
+# create an array with (restype, atomtype) --> rigid_group_idx
+# and an array with (restype, atomtype, coord) for the atom positions
+# and compute affine transformation matrices (4,4) from one rigid group to the
+# previous group
+restype_atom37_to_rigid_group = np.zeros([21, 37], dtype=np.int32)
+restype_atom37_mask = np.zeros([21, 37], dtype=np.float32)
+restype_atom37_rigid_group_positions = np.zeros([21, 37, 3], dtype=np.float32)
+restype_atom14_to_rigid_group = np.zeros([21, 14], dtype=np.int32)
+restype_atom14_mask = np.zeros([21, 14], dtype=np.float32)
+restype_atom14_rigid_group_positions = np.zeros([21, 14, 3], dtype=np.float32)
+restype_rigid_group_default_frame = np.zeros([21, 8, 4, 4], dtype=np.float32)
+
+
+def _make_rigid_group_constants():
+ """Fill the arrays above."""
+ for restype, restype_letter in enumerate(restypes):
+ resname = restype_1to3[restype_letter]
+ for atomname, group_idx, atom_position in rigid_group_atom_positions[resname]:
+ atomtype = atom_order[atomname]
+ restype_atom37_to_rigid_group[restype, atomtype] = group_idx
+ restype_atom37_mask[restype, atomtype] = 1
+ restype_atom37_rigid_group_positions[restype, atomtype, :] = atom_position
+
+ atom14idx = restype_name_to_atom14_names[resname].index(atomname)
+ restype_atom14_to_rigid_group[restype, atom14idx] = group_idx
+ restype_atom14_mask[restype, atom14idx] = 1
+ restype_atom14_rigid_group_positions[restype, atom14idx, :] = atom_position
+
+ for restype, restype_letter in enumerate(restypes):
+ resname = restype_1to3[restype_letter]
+ atom_positions = {
+ name: np.array(pos) for name, _, pos in rigid_group_atom_positions[resname]
+ }
+
+ # backbone to backbone is the identity transform
+ restype_rigid_group_default_frame[restype, 0, :, :] = np.eye(4)
+
+ # pre-omega-frame to backbone (currently dummy identity matrix)
+ restype_rigid_group_default_frame[restype, 1, :, :] = np.eye(4)
+
+ # phi-frame to backbone
+ mat = _make_rigid_transformation_4x4(
+ ex=atom_positions["N"] - atom_positions["CA"],
+ ey=np.array([1.0, 0.0, 0.0]),
+ translation=atom_positions["N"],
+ )
+ restype_rigid_group_default_frame[restype, 2, :, :] = mat
+
+ # psi-frame to backbone
+ mat = _make_rigid_transformation_4x4(
+ ex=atom_positions["C"] - atom_positions["CA"],
+ ey=atom_positions["CA"] - atom_positions["N"],
+ translation=atom_positions["C"],
+ )
+ restype_rigid_group_default_frame[restype, 3, :, :] = mat
+
+ # chi1-frame to backbone
+ if chi_angles_mask[restype][0]:
+ base_atom_names = chi_angles_atoms[resname][0]
+ base_atom_positions = [atom_positions[name] for name in base_atom_names]
+ mat = _make_rigid_transformation_4x4(
+ ex=base_atom_positions[2] - base_atom_positions[1],
+ ey=base_atom_positions[0] - base_atom_positions[1],
+ translation=base_atom_positions[2],
+ )
+ restype_rigid_group_default_frame[restype, 4, :, :] = mat
+
+ # chi2-frame to chi1-frame
+ # chi3-frame to chi2-frame
+ # chi4-frame to chi3-frame
+ # luckily all rotation axes for the next frame start at (0,0,0) of the
+ # previous frame
+ for chi_idx in range(1, 4):
+ if chi_angles_mask[restype][chi_idx]:
+ axis_end_atom_name = chi_angles_atoms[resname][chi_idx][2]
+ axis_end_atom_position = atom_positions[axis_end_atom_name]
+ mat = _make_rigid_transformation_4x4(
+ ex=axis_end_atom_position,
+ ey=np.array([-1.0, 0.0, 0.0]),
+ translation=axis_end_atom_position,
+ )
+ restype_rigid_group_default_frame[restype, 4 + chi_idx, :, :] = mat
+
+
+_make_rigid_group_constants()
+
+
+def make_atom14_dists_bounds(overlap_tolerance=1.5, bond_length_tolerance_factor=15):
+ """Compute upper and lower bounds for bonds to assess violations."""
+ restype_atom14_bond_lower_bound = np.zeros([21, 14, 14], np.float32)
+ restype_atom14_bond_upper_bound = np.zeros([21, 14, 14], np.float32)
+ restype_atom14_bond_stddev = np.zeros([21, 14, 14], np.float32)
+ residue_bonds, residue_virtual_bonds, _ = load_stereo_chemical_props()
+ for restype, restype_letter in enumerate(restypes):
+ resname = restype_1to3[restype_letter]
+ atom_list = restype_name_to_atom14_names[resname]
+
+ # create lower and upper bounds for clashes
+ for atom1_idx, atom1_name in enumerate(atom_list):
+ if not atom1_name:
+ continue
+ atom1_radius = van_der_waals_radius[atom1_name[0]]
+ for atom2_idx, atom2_name in enumerate(atom_list):
+ if (not atom2_name) or atom1_idx == atom2_idx:
+ continue
+ atom2_radius = van_der_waals_radius[atom2_name[0]]
+ lower = atom1_radius + atom2_radius - overlap_tolerance
+ upper = 1e10
+ restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
+ restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
+ restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
+ restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper
+
+ # overwrite lower and upper bounds for bonds and angles
+ for b in residue_bonds[resname] + residue_virtual_bonds[resname]:
+ atom1_idx = atom_list.index(b.atom1_name)
+ atom2_idx = atom_list.index(b.atom2_name)
+ lower = b.length - bond_length_tolerance_factor * b.stddev
+ upper = b.length + bond_length_tolerance_factor * b.stddev
+ restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
+ restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
+ restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
+ restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper
+ restype_atom14_bond_stddev[restype, atom1_idx, atom2_idx] = b.stddev
+ restype_atom14_bond_stddev[restype, atom2_idx, atom1_idx] = b.stddev
+ return {
+ "lower_bound": restype_atom14_bond_lower_bound, # shape (21,14,14)
+ "upper_bound": restype_atom14_bond_upper_bound, # shape (21,14,14)
+ "stddev": restype_atom14_bond_stddev, # shape (21,14,14)
+ }
+
+
+def _make_restype_atom14_to_atom37():
+ """Map from atom14 to atom37 per residue type."""
+ restype_atom14_to_atom37 = [] # mapping (restype, atom14) --> atom37
+ for rt in restypes:
+ atom_names = restype_name_to_atom14_names[restype_1to3[rt]]
+ restype_atom14_to_atom37.append([(atom_order[name] if name else 0) for name in atom_names])
+ # Add dummy mapping for restype 'UNK'
+ restype_atom14_to_atom37.append([0] * 14)
+ restype_atom14_to_atom37 = np.array(restype_atom14_to_atom37, dtype=np.int32)
+ return restype_atom14_to_atom37
+
+
+RESTYPE_ATOM14_TO_ATOM37 = _make_restype_atom14_to_atom37()
+
+
+def _make_restype_atom14_mask():
+ """Mask of which atoms are present for which residue type in atom14."""
+ restype_atom14_mask = []
+
+ for rt in restypes:
+ atom_names = restype_name_to_atom14_names[restype_1to3[rt]]
+ restype_atom14_mask.append([(1.0 if name else 0.0) for name in atom_names])
+
+ restype_atom14_mask.append([0.0] * 14)
+ restype_atom14_mask = np.array(restype_atom14_mask, dtype=np.float32)
+ return restype_atom14_mask
+
+
+RESTYPE_ATOM14_MASK = _make_restype_atom14_mask()
diff --git a/forks/FlowDock/flowdock/data/components/uff_parameters.csv b/forks/FlowDock/flowdock/data/components/uff_parameters.csv
new file mode 100644
index 00000000..6e637af9
--- /dev/null
+++ b/forks/FlowDock/flowdock/data/components/uff_parameters.csv
@@ -0,0 +1,226 @@
+Atom,Element,Coordination,bond ri (Å),angle θ0 (deg),distance,energy,scale ζ,charge Zi,Source,,Vi,Uj,Xi,Hardness,Radius
+Ac6+3,Ac,6,1.983,90,3.478,0.033,12,3.9,UFF,,0,0,2.835,2.835,2.108
+Ag1+1,Ag,2,1.386,180,3.148,0.036,12,1.956,UFF,,0,0.2,4.436,3.134,1.622
+Ag1f1,Ag,2,1.22,180,3.148,0.036,12,1.956,UFF4MOF2,,,,,,
+Ag2f2,Ag,3,1.34,120,3.148,0.036,12,1.956,UFF4MOF2,,,,,,
+Ag3f2,Ag,4,1.48,109.47,3.148,0.036,12,1.956,UFF4MOF2,,,,,,
+Ag4f2,Ag,4,1.51,90,3.148,0.036,12,1.956,UFF4MOF2,,,,,,
+Al3,Al,4,1.244,109.47,4.499,0.505,11.278,1.792,UFF,,0,1.25,4.06,3.59,1.201
+Al3f2,Al,4,1.28,109.47,4.499,0.505,11.278,1.792,UFF4MOF2,,,,,,
+Al6+3,Al,6,1.22,90,4.499,0.505,11.278,1.792,UFF4MOF,,,,,,
+Am6+4,Am,6,1.66,90,3.381,0.014,12,3.9,UFF,,0,0,2.9895,3.0035,1.942
+Ar4+4,Ar,4,1.032,90,3.868,0.185,15.763,0.3,UFF,,0,1.25,9.465,6.355,2.108
+As3+3,As,4,1.211,92.1,4.23,0.309,13,2.864,UFF,,1.5,0.7,5.188,3.809,1.204
+At,At,1,1.545,180,4.75,0.284,15,2.24,UFF,,0,0.1,4.75,4.75,1.47
+Au1f1,Au,2,1.11,180,3.293,0.039,12,2.625,UFF4MOF2,,,,,,
+Au4+3,Au,4,1.262,90,3.293,0.039,12,2.625,UFF,,0,0.1,4.894,2.586,1.618
+B_2,B,3,0.828,120,4.083,0.18,12.052,1.755,UFF,,0,2,5.11,4.75,0.822
+B_3,B,4,0.838,109.47,4.083,0.18,12.052,1.755,UFF,,0,2,5.11,4.75,0.822
+Ba3f2,Ba,4,2.04,109.47,3.703,0.364,12,2.727,UFF4MOF2,,,,,,
+Ba6+2,Ba,6,2.277,90,3.703,0.364,12,2.727,UFF,,0,0.1,2.814,2.396,2.442
+Be3+2,Be,4,1.074,109.47,2.745,0.085,12,1.565,UFF,,0,2,4.877,4.443,1.24
+Bi3+3,Bi,4,1.512,90,4.37,0.518,13,2.47,UFF,,1,0.1,4.69,3.74,1.514
+Bk6+3,Bk,6,1.761,90,3.339,0.013,12,3.9,UFF,,0,0,3.1935,3.0355,1.9
+Br,Br,1,1.192,180,4.189,0.251,15,2.519,UFF,,0,0.7,7.79,4.425,1.141
+C_1,C,2,0.706,180,3.851,0.105,12.73,1.912,UFF,,0,2,5.343,5.063,0.759
+C_2,C,3,0.732,120,3.851,0.105,12.73,1.912,UFF,,0,2,5.343,5.063,0.759
+C_3,C,4,0.757,109.47,3.851,0.105,12.73,1.912,UFF,,2.119,2,5.343,5.063,0.759
+C_R,C,3,0.729,120,3.851,0.105,12.73,1.912,UFF,,0,2,5.343,5.063,0.759
+Ca1f1,Ca,4,1.56,180,3.399,0.238,12,2.141,UFF4MOF2,,,,,,
+Ca3f2,Ca,4,1.705,109.47,3.399,0.238,12,2.141,UFF4MOF2,,,,,,
+Ca6+2,Ca,6,1.761,90,3.399,0.238,12,2.141,UFF,,0,0.7,3.231,2.88,2
+Cd1f1,Cd,2,1.4,180,2.848,0.228,12,1.65,UFF4MOF2,,,,,,
+Cd3+2,Cd,4,1.403,109.47,2.848,0.228,12,1.65,UFF,,0,0.2,5.034,3.957,1.6
+Cd3f2,Cd,4,1.29,109.47,2.848,0.228,12,1.65,UFF4MOF2,,,,,,
+Cd4f2,Cd,4,1.46,90,2.848,0.228,12,1.65,UFF4MOF2,,,,,,
+Cd8f4,Cd,8,1.64,109.47,2.848,0.228,12,1.65,UFF4MOF2,,,,,,
+Ce6+3,Ce,6,1.841,90,3.556,0.013,12,3.3,UFF,,0,0.1,2.774,2.692,1.925
+Ce8f4,Ce,8,1.76,109.47,3.556,0.013,12,3.3,UFF4MOF2,,,,,,
+Cf6+3,Cf,6,1.75,90,3.313,0.013,12,3.9,UFF,,0,0,3.197,3.101,1.9
+Cl,Cl,1,1.044,180,3.947,0.227,14.866,2.348,UFF,,0,1.25,8.564,4.946,0.994
+Cm6+3,Cm,6,1.801,90,3.326,0.013,12,3.9,UFF,,0,0,2.8315,3.1895,1.9
+Co1f1,Co,2,1.28,180,2.872,0.014,12,1.308,UFF4MOF2,,,,,,
+Co2f2,Co,3,1.25,120,2.872,0.014,12,1.308,UFF4MOF2,,,,,,
+Co3+2,Co,4,1.24,109.47,2.872,0.014,12,1.308,UFF4MOF,,,,,,
+Co4+2,Co,4,1.16,90,2.872,0.014,12,1.308,UFF4MOF,,,,,,
+Co6+3,Co,6,1.241,90,2.872,0.014,12,2.43,UFF,,0,0.7,4.105,4.175,1.406
+Cr4+2,Cr,4,1.1,90,3.023,0.015,12,2.463,UFF4MOF,,,,,,
+Cr6+3,Cr,6,1.345,90,3.023,0.015,12,2.463,UFF,,0,0.7,3.415,3.865,1.402
+Cr6f3,Cr,6,1.28,90,3.023,0.015,12,2.463,UFF4MOF,,,,,,
+Cs,Cs,1,2.57,180,4.517,0.045,12,1.573,UFF,,0,0.1,2.183,1.711,2.984
+Cu1f1,Cu,2,1.24,180,3.495,0.005,12,2.43,UFF4MOF2,,,,,,
+Cu2f2,Cu,3,1.11,120,3.495,0.005,12,2.43,UFF4MOF2,,,,,,
+Cu3+1,Cu,4,1.302,109.47,3.495,0.005,12,1.756,UFF,,0,0.7,4.2,4.22,1.434
+Cu3f2,Cu,4,1.19,109.47,3.495,0.005,12,2.43,UFF4MOF2,,,,,,
+Cu4+2,Cu,4,1.28,90,3.495,0.005,12,2.43,UFF4MOF,,,,,,
+Dy6+3,Dy,6,1.71,90,3.428,0.007,12,3.3,UFF,,0,0.1,3.0555,2.8715,1.934
+Dy6f3,Dy,6,1.58,90,3.428,0.007,12,3.3,UFF4MOF2,,,,,,
+Dy8f4,Dy,8,1.7,109.47,3.428,0.007,12,3.3,UFF4MOF2,,,,,,
+Er6+3,Er,6,1.673,90,3.391,0.007,12,3.3,UFF,,0,0.1,3.1865,2.9145,1.915
+Er8f4,Er,8,1.64,109.47,3.391,0.007,12,3.3,UFF4MOF2,,,,,,
+Es6+3,Es,6,1.724,90,3.299,0.012,12,3.9,UFF,,0,0,3.333,3.089,1.9
+Eu6+3,Eu,6,1.771,90,3.493,0.008,12,3.3,UFF,,0,0.1,2.8785,2.7875,2.227
+Eu6f3,Eu,6,1.6,90,3.493,0.008,12,3.3,UFF4MOF2,,,,,,
+Eu8f4,Eu,8,1.74,109.47,3.493,0.008,12,3.3,UFF4MOF2,,,,,,
+F_,F,1,0.668,180,3.364,0.05,14.762,1.735,UFF,,0,2,10.874,7.474,0.706
+Fe1f1,Fe,2,1.33,180,2.912,0.013,12,2.43,UFF4MOF2,,,,,,
+Fe3+2,Fe,4,1.27,109.47,2.912,0.013,12,2.43,UFF,,0,0.7,3.76,4.14,1.393
+Fe4+2,Fe,4,1.1,90,2.912,0.013,12,2.43,UFF4MOF,,,,,,
+Fe6+2,Fe,6,1.335,90,2.912,0.013,12,2.43,UFF,,0,0.7,3.76,4.14,1.393
+Fe6+3,Fe,6,1.32,90,2.912,0.013,12,2.43,UFF4MOF,,,,,,
+Fm6+3,Fm,6,1.712,90,3.286,0.012,12,3.9,UFF,,0,0,3.4,3.1,1.9
+Fr,Fr,1,2.88,180,4.9,0.05,12,1.847,UFF,,0,0,2,2,2.3
+Ga3+3,Ga,4,1.26,109.47,4.383,0.415,11,1.821,UFF,,0,0.7,3.641,3.16,1.211
+Ga3f2,Ga,4,1.15,109.47,4.383,0.415,11,1.821,UFF4MOF2,,,,,,
+Ga6f3,Ga,6,1.48,90,4.383,0.415,11,1.821,UFF4MOF2,,,,,,
+Gd6+3,Gd,6,1.735,90,3.368,0.009,12,3.3,UFF,,0,0.1,3.1665,2.9745,1.968
+Gd6f3,Gd,6,1.55,90,3.368,0.009,12,3.3,UFF4MOF2,,,,,,
+Gd8f4,Gd,8,1.7,109.47,3.368,0.009,12,3.3,UFF4MOF2,,,,,,
+Ge3,Ge,4,1.197,109.47,4.28,0.379,12,2.789,UFF,,0.701,0.7,4.051,3.438,1.189
+H_,H,1,0.354,180,2.886,0.044,12,0.712,UFF,,0,0,4.528,6.9452,0.371
+H_b,H,2,0.46,83.5,2.886,0.044,12,0.712,UFF,,0,0,4.528,6.9452,0.371
+He4+4,He,4,0.849,90,2.362,0.056,15.24,0.098,UFF,,0,0,9.66,14.92,1.3
+Hf3+4,Hf,4,1.611,109.47,3.141,0.072,12,3.921,UFF,,0,0.1,3.7,3.4,1.759
+Hf8f4,Hf,8,1.46,109.47,3.141,0.072,12,3.921,UFF4MOF2,,,,,,
+Hg1+2,Hg,2,1.34,180,2.705,0.385,12,1.75,UFF,,0,0.1,6.27,4.16,1.6
+Hg3f2,Hg,4,1.248,109.47,2.705,0.385,12,1.75,UFF4MOF2,,,,,,
+Ho6+3,Ho,6,1.696,90,3.409,0.007,12,3.416,UFF,,0,0.1,3.127,2.891,1.925
+Ho8f4,Ho,8,1.7,109.47,3.409,0.007,12,3.416,UFF4MOF2,,,,,,
+I_,I,1,1.382,180,4.5,0.339,15,2.65,UFF,,0,0.2,6.822,3.762,1.333
+In3+3,In,4,1.459,109.47,4.463,0.599,11,2.07,UFF,,0,0.2,3.506,2.896,1.404
+In3f2,In,4,1.33,109.47,4.463,0.599,11,2.07,UFF4MOF2,,,,,,
+In6f3,In,6,1.53,90,4.463,0.599,11,2.07,UFF4MOF2,,,,,,
+In8f4,In,8,1.53,109.47,4.463,0.599,11,2.07,UFF4MOF2,,,,,,
+Ir6+3,Ir,6,1.371,90,2.84,0.073,12,3.731,UFF,,0,0.1,5,4,1.866
+K_,K,1,1.953,180,3.812,0.035,12,1.165,UFF,,0,0.7,2.421,1.92,2.586
+K_3f2,K,4,2.38,109.47,3.812,0.035,12,1.165,UFF4MOF2,,,,,,
+K_4f2,K,4,2.01,90,3.812,0.035,12,1.165,UFF4MOF2,,,,,,
+Kr4+4,Kr,4,1.147,90,4.141,0.22,16,0.452,UFF,,0,0.7,8.505,5.715,2.27
+La2f2,La,3,1.86,120,3.522,0.017,12,3.3,UFF4MOF2,,,,,,
+La3+3,La,4,1.943,109.47,3.522,0.017,12,3.3,UFF,,0,0.1,2.8355,2.7415,2.071
+La8f4,La,8,1.66,109.47,3.522,0.017,12,3.3,UFF4MOF2,,,,,,
+Li,Li,1,1.336,180,2.451,0.025,12,1.026,UFF,,0,2,3.006,2.386,1.557
+Li3f2,Li,4,1.28,109.47,2.451,0.025,12,1.026,UFF4MOF2,,,,,,
+Lu1f1,Lu,2,1.65,180,3.64,0.041,12,3.271,UFF4MOF2,,,,,,
+Lu6+3,Lu,6,1.671,90,3.64,0.041,12,3.271,UFF,,0,0.1,2.9629,2.4629,1.896
+Lu8f4,Lu,8,1.66,109.47,3.64,0.041,12,3.271,UFF4MOF2,,,,,,
+Lw6+3,Lw,6,1.698,90,3.236,0.011,12,3.9,UFF,,0,0,3.5,3.2,1.9
+Md6+3,Md,6,1.689,90,3.274,0.011,12,3.9,UFF,,0,0,3.47,3.11,1.9
+Mg3+2,Mg,4,1.421,109.47,3.021,0.111,12,1.787,UFF,,0,1.25,3.951,3.693,1.5
+Mg6f3,Mg,6,1.525,90,3.021,0.111,12,1.787,UFF4MOF2,,,,,,
+Mn1f1,Mn,2,1.38,180,2.961,0.013,12,2.43,UFF4MOF2,,,,,,
+Mn3f2,Mn,4,1.18,109.47,2.961,0.013,12,2.43,UFF4MOF2,,,,,,
+Mn4+2,Mn,4,1.26,90,2.961,0.013,12,2.43,UFF4MOF,,,,,,
+Mn6+2,Mn,6,1.382,90,2.961,0.013,12,2.43,UFF,,0,0.7,3.325,4.105,1.533
+Mn6+3,Mn,6,1.34,90,2.961,0.013,12,2.43,UFF4MOF,,,,,,
+Mn8f4,Mn,8,1.52,109.47,2.961,0.013,12,2.43,UFF4MOF2,,,,,,
+Mo3+6,Mo,4,1.484,109.47,3.052,0.056,12,3.4,UFF,,0,0.2,3.465,3.755,1.53
+Mo3f2,Mo,8,1.24,109.47,3.052,0.056,12,3.4,UFF4MOF2,,,,,,
+Mo4f2,Mo,4,1.4,90,3.052,0.056,12,3.4,UFF4MOF2,,,,,,
+Mo6+6,Mo,6,1.467,90,3.052,0.056,12,3.4,UFF,,0,0.2,3.465,3.755,1.53
+Mo8f4,Mo,8,1.28,109.47,3.052,0.056,12,3.4,UFF4MOF2,,,,,,
+N_1,N,2,0.656,180,3.66,0.069,13.407,2.544,UFF,,0,2,6.899,5.88,0.715
+N_2,N,3,0.685,111.2,3.66,0.069,13.407,2.544,UFF,,0,2,6.899,5.88,0.715
+N_3,N,4,0.7,106.7,3.66,0.069,13.407,2.544,UFF,,0.45,2,6.899,5.88,0.715
+N_R,N,3,0.699,120,3.66,0.069,13.407,2.544,UFF,,0,2,6.899,5.88,0.715
+Na,Na,1,1.539,180,2.983,0.03,12,1.081,UFF,,0,1.25,2.843,2.296,2.085
+Na3f2,Na,4,1.623,109.47,2.983,0.03,12,1.081,UFF4MOF2,,,,,,
+Na4f2,Na,4,1.79,90,2.983,0.03,12,1.081,UFF4MOF2,,,,,,
+Nb3+5,Nb,4,1.473,109.47,3.165,0.059,12,3.618,UFF,,0,0.2,3.55,3.38,1.603
+Nb8f4,Nb,8,1.37,109.47,3.165,0.059,12,3.618,UFF4MOF2,,,,,,
+Nd6+3,Nd,6,1.816,90,3.575,0.01,12,3.3,UFF,,0,0.1,2.8685,2.6205,2.007
+Nd8f4,Nd,8,1.78,109.47,3.575,0.01,12,3.3,UFF4MOF2,,,,,,
+Ne4+4,Ne,4,0.92,90,3.243,0.042,15.44,0.194,UFF,,0,2,11.04,10.55,1.768
+Ni4+2,Ni,4,1.164,90,2.834,0.015,12,2.43,UFF,,0,0.7,4.465,4.205,1.398
+No6+3,No,6,1.679,90,3.248,0.011,12,3.9,UFF,,0,0,3.475,3.175,1.9
+Np6+4,Np,6,1.666,90,3.424,0.019,12,3.9,UFF,,0,0,3.549,2.717,1.8
+O_1,O,2,0.639,180,3.5,0.06,14.085,2.3,UFF,,0,2,8.741,6.682,0.669
+O_2,O,3,0.634,120,3.5,0.06,14.085,2.3,UFF,,0,2,8.741,6.682,0.669
+O_2_z,O,3,0.528,120,3.5,0.06,14.085,2.3,UFF4MOF,,,,,,
+O_3,O,4,0.658,104.51,3.5,0.06,14.085,2.3,UFF,,0.018,2,8.741,6.682,0.669
+O_3_f,O,4,0.634,109.47,3.5,0.06,14.085,2.3,UFF4MOF,,,,,,
+O_3_z,O,4,0.528,146,3.5,0.06,14.085,2.3,UFF,,0.018,2,8.741,6.682,0.669
+O_R,O,3,0.68,110,3.5,0.06,14.085,2.3,UFF,,0,2,8.741,6.682,0.669
+Os6+6,Os,6,1.372,90,3.12,0.037,12,3.7,UFF,,0,0.1,5.14,3.63,1.7
+P_3+3,P,4,1.101,93.8,4.147,0.305,13.072,2.863,UFF,,2.4,1.25,5.463,4,1.102
+P_3+5,P,4,1.056,109.47,4.147,0.305,13.072,2.863,UFF,,2.4,1.25,5.463,4,1.102
+P_3+q,P,4,1.056,109.47,4.147,0.305,13.072,2.863,UFF,,2.4,1.25,5.463,4,1.102
+Pa6+4,Pa,6,1.711,90,3.424,0.022,12,3.9,UFF,,0,0,2.985,2.905,1.8
+Pb3,Pb,4,1.459,109.47,4.297,0.663,12,2.846,UFF,,0.1,0.1,3.9,3.53,1.444
+Pb4f2,Pb,4,1.67,90,4.297,0.663,12,2.846,UFF4MOF2,,,,,,
+Pd4+2,Pd,4,1.338,90,2.899,0.048,12,3.21,UFF,,0,0.2,4.32,4,1.544
+Pd6f3,Pd,6,1.19,90,2.899,0.048,12,3.21,UFF4MOF2,,,,,,
+Pm6+3,Pm,6,1.801,90,3.547,0.009,12,3.3,UFF,,0,0.1,2.881,2.673,2
+Po3+2,Po,4,1.5,90,4.709,0.325,14,2.33,UFF,,0.3,0.1,4.21,4.21,1.48
+Pr6+3,Pr,6,1.823,90,3.606,0.01,12,3.3,UFF,,0,0.1,2.858,2.564,2.007
+Pr8f4,Pr,8,1.83,109.47,3.606,0.01,12,3.3,UFF4MOF2,,,,,,
+Pt4+2,Pt,4,1.364,90,2.754,0.08,12,3.382,UFF,,0,0.1,4.79,4.43,1.557
+Pt4f2,Pt,4,1.125,90,2.754,0.08,12,3.382,UFF4MOF2,,,,,,
+Pu6+4,Pu,6,1.657,90,3.424,0.016,12,3.9,UFF,,0,0,3.243,2.819,1.84
+Ra6+2,Ra,6,2.512,90,3.677,0.404,12,2.92,UFF,,0,0,2.843,2.434,2.2
+Rb,Rb,1,2.26,180,4.114,0.04,12,1.592,UFF,,0,0.2,2.331,1.846,2.77
+Re3+7,Re,4,1.314,109.47,2.954,0.066,12,3.7,UFF,,0,0.1,3.96,3.92,1.6
+Re6+5,Re,6,1.372,90,2.954,0.066,12,3.7,UFF,,0,0.1,3.96,3.92,1.6
+Re6f3,Re,4,1.23,90,2.954,0.066,12,3.7,UFF4MOF2,,,,,,
+Rh4f2,Rh,4,1.22,90,2.929,0.053,12,3.5,UFF4MOF2,,,,,,
+Rh6+3,Rh,6,1.332,90,2.929,0.053,12,3.5,UFF,,0,0.2,3.975,4.005,1.509
+Rn4+4,Rn,4,1.42,90,4.765,0.248,16,0.583,UFF,,0,0.1,5.37,5.37,2.2
+Ru4f2,Ru,4,1.32,90,2.963,0.056,12,3.4,UFF4MOF2,,,,,,
+Ru6+2,Ru,6,1.478,90,2.963,0.056,12,3.4,UFF,,0,0.2,3.575,4.015,1.5
+S_2,S,3,0.854,120,4.035,0.274,13.969,2.703,UFF,,0,1.25,6.928,4.486,1.047
+S_3_f,S,3,0.854,109.47,4.035,0.274,13.969,2.703,UFF4MOF2,,,,,,
+S_3+2,S,4,1.064,92.1,4.035,0.274,13.969,2.703,UFF,,0.484,1.25,6.928,4.486,1.047
+S_3+4,S,4,1.049,103.2,4.035,0.274,13.969,2.703,UFF,,0.484,1.25,6.928,4.486,1.047
+S_3+6,S,4,1.027,109.47,4.035,0.274,13.969,2.703,UFF,,0.484,1.25,6.928,4.486,1.047
+S_R,S,4,1.077,92.2,4.035,0.274,13.969,2.703,UFF,,0,1.25,6.928,4.486,1.047
+Sb3+3,Sb,4,1.407,91.6,4.42,0.449,13,2.704,UFF,,1.1,0.2,4.899,3.342,1.404
+Sc3+3,Sc,4,1.513,109.47,3.295,0.019,12,2.592,UFF,,0,0.7,3.395,3.08,1.75
+Sc6+3,Sc,6,1.44,90,3.295,0.019,12,2.595,UFF4MOF,,,,,,
+Se3+2,Se,4,1.19,90.6,4.205,0.291,14,2.764,UFF,,0.335,0.7,6.428,4.131,1.224
+Si3,SI,4,1.117,109.47,4.295,0.402,12.175,2.323,UFF,,1.225,1.25,4.168,3.487,1.176
+Sm6+3,Sm,6,1.78,90,3.52,0.008,12,3.3,UFF,,0,0.1,2.9115,2.7195,1.978
+Sm8f4,Sm,8,1.78,109.47,3.52,0.008,12,3.3,UFF4MOF2,,,,,,
+Sn3,Sn,4,1.398,109.47,4.392,0.567,12,2.961,UFF,,0.199,0.2,3.987,3.124,1.354
+Sr6+2,Sr,6,2.052,90,3.641,0.235,12,2.449,UFF,,0,0.2,3.024,2.44,2.415
+Sr8f4,Sr,8,1.82,109.47,3.641,0.235,12,2.449,UFF4MOF2,,,,,,
+Ta3+5,Ta,4,1.511,109.47,3.17,0.081,12,4.075,UFF,,0,0.1,5.1,2.85,1.605
+Tb6+3,Tb,6,1.732,90,3.451,0.007,12,3.3,UFF,,0,0.1,3.018,2.834,1.954
+Tb8f4,Tb,8,1.64,109.47,3.451,0.007,12,3.3,UFF4MOF2,,,,,,
+Tc6+5,Tc,6,1.322,90,2.998,0.048,12,3.4,UFF,,0,0.2,3.29,3.99,1.5
+Te3+2,Te,4,1.386,90.25,4.47,0.398,14,2.882,UFF,,0.3,0.2,5.816,3.526,1.38
+Th6+4,Th,6,1.721,90,3.396,0.026,12,4.202,UFF,,0,0,3.175,2.905,2.018
+Ti3+4,Ti,4,1.412,109.47,3.175,0.017,12,2.659,UFF,,0,0.7,3.47,3.38,1.607
+Ti4+2,Ti,4,1.38,90,3.175,0.017,12,2.659,UFF4MOF,,,,,,
+Ti6+4,Ti,6,1.412,90,3.175,0.017,12,2.659,UFF,,0,0.7,3.47,3.38,1.607
+Ti8f4,Ti,8,1.46,109.47,3.175,0.017,12,2.659,UFF4MOF2,,,,,,
+Tl3+3,Tl,4,1.518,120,4.347,0.68,11,2.068,UFF,,0,0.1,3.2,2.9,1.53
+Tm6+3,Tm,6,1.66,90,3.374,0.006,12,3.3,UFF,,0,0.1,3.2514,2.9329,2
+Tm8f4,Tm,8,1.67,109.47,3.374,0.006,12,3.3,UFF4MOF2,,,,,,
+U_6+4,U,6,1.684,90,3.395,0.022,12,3.9,UFF,,0,0,3.341,2.853,1.713
+U_6f3,U,6,1.65,90,3.395,0.022,12,3.9,UFF4MOF2,,,,,,
+U_8f4,U,8,1.73,109.47,3.395,0.022,12,3.9,UFF4MOF2,,,,,,
+V_3+5,V,4,1.402,109.47,3.144,0.016,12,2.679,UFF,,0,0.7,3.65,3.41,1.47
+V_3f2,V,4,1.12,109.47,3.144,0.016,12,2.679,UFF4MOF2,,,,,,
+V_4+2,V,4,1.18,90,3.144,0.016,12,2.679,UFF4MOF,,,,,,
+V_6+3,V,6,1.3,90,3.144,0.016,12,2.679,UFF4MOF,,,,,,
+W_3+4,W,4,1.526,109.47,3.069,0.067,12,3.7,UFF,,0,0.1,4.63,3.31,1.538
+W_3+6,W,4,1.38,109.47,3.069,0.067,12,3.7,UFF,,0,0.1,4.63,3.31,1.538
+W_3f2,W,4,1.16,109.47,3.069,0.067,12,3.7,UFF4MOF2,,,,,,
+W_4f2,W,4,1.345,90,3.069,0.067,12,3.7,UFF4MOF2,,,,,,
+W_6+6,W,6,1.392,90,3.069,0.067,12,3.7,UFF,,0,0.1,4.63,3.31,1.538
+W_8f4,W,4,1.27,109.47,3.069,0.067,12,3.7,UFF4MOF2,,,,,,
+Xe4+4,Xe,4,1.267,90,4.404,0.332,12,0.556,UFF,,0,0.2,7.595,4.975,2.459
+Y_3+3,Y,4,1.698,109.47,3.345,0.072,12,3.257,UFF,,0,0.2,3.83,2.81,1.998
+Y_6f3,Y,6,1.6,90,3.345,0.072,12,3.257,UFF4MOF2,,,,,,
+Y_8f4,Y,8,1.68,109.47,3.345,0.072,12,3.257,UFF4MOF2,,,,,,
+Yb6+3,Yb,6,1.637,90,3.355,0.228,12,2.618,UFF,,0,0.1,3.2889,2.965,2.158
+Yb6f3,Yb,6,1.45,90,3.355,0.228,12,2.618,UFF4MOF2,,,,,,
+Yb8f4,Yb,8,1.62,109.47,3.355,0.228,12,2.618,UFF4MOF2,,,,,,
+Zn1f1,Zn,2,1.3,180,2.763,0.124,12,1.308,UFF4MOF2,,,,,,
+Zn2f2,Zn,3,1.3,120,2.763,0.124,12,1.308,UFF4MOF2,,,,,,
+Zn3+2,Zn,4,1.193,109.47,2.763,0.124,12,1.308,UFF,,0,0.7,5.106,4.285,1.4
+Zn3f2,Zn,4,1.24,109.47,2.763,0.124,12,1.308,UFF4MOF,,,,,,
+Zn4+2,Zn,4,1.34,90,2.763,0.124,12,1.308,UFF4MOF,,,,,,
+Zr3+4,Zr,4,1.564,109.47,3.124,0.069,12,3.667,UFF,,0,0.2,3.4,3.55,1.758
+Zr6f3,Zr,6,1.25,90,3.124,0.069,12,3.667,UFF4MOF2,,,,,,
+Zr8f4,Zr,8,1.68,109.47,3.124,0.069,12,3.667,UFF4MOF2,,,,,,
diff --git a/forks/FlowDock/flowdock/eval.py b/forks/FlowDock/flowdock/eval.py
new file mode 100644
index 00000000..34a375b2
--- /dev/null
+++ b/forks/FlowDock/flowdock/eval.py
@@ -0,0 +1,160 @@
+import os
+
+import hydra
+import lovely_tensors as lt
+import rootutils
+import torch
+from beartype.typing import Any, Dict, List, Tuple
+from lightning import LightningDataModule, LightningModule, Trainer
+from lightning.fabric.plugins.environments.cluster_environment import ClusterEnvironment
+from lightning.pytorch.loggers import Logger
+from lightning.pytorch.strategies.strategy import Strategy
+from omegaconf import DictConfig, open_dict
+
+lt.monkey_patch()
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+# ------------------------------------------------------------------------------------ #
+# the setup_root above is equivalent to:
+# - adding project root dir to PYTHONPATH
+# (so you don't need to force user to install project as a package)
+# (necessary before importing any local modules e.g. `from flowdock import utils`)
+# - setting up PROJECT_ROOT environment variable
+# (which is used as a base for paths in "configs/paths/default.yaml")
+# (this way all filepaths are the same no matter where you run the code)
+# - loading environment variables from ".env" in root dir
+#
+# you can remove it if you:
+# 1. either install project as a package or move entry files to project root dir
+# 2. set `root_dir` to "." in "configs/paths/default.yaml"
+#
+# more info: https://github.com/ashleve/rootutils
+# ------------------------------------------------------------------------------------ #
+
+from flowdock import register_custom_omegaconf_resolvers, resolve_omegaconf_variable
+from flowdock.utils import (
+ RankedLogger,
+ extras,
+ instantiate_loggers,
+ log_hyperparameters,
+ task_wrapper,
+)
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+@task_wrapper
+def evaluate(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+ """Evaluates given checkpoint on a datamodule testset.
+
+ This method is wrapped in optional @task_wrapper decorator, that controls the behavior during
+ failure. Useful for multiruns, saving info about the crash, etc.
+
+ :param cfg: DictConfig configuration composed by Hydra.
+ :return: Tuple[dict, dict] with metrics and dict with all instantiated objects.
+ """
+ assert cfg.ckpt_path, "Please provide a checkpoint path to evaluate!"
+ assert os.path.exists(cfg.ckpt_path), f"Checkpoint path {cfg.ckpt_path} does not exist!"
+
+ log.info(
+ f"Setting `float32_matmul_precision` to {cfg.model.cfg.task.float32_matmul_precision}."
+ )
+ torch.set_float32_matmul_precision(precision=cfg.model.cfg.task.float32_matmul_precision)
+
+ log.info(f"Instantiating datamodule <{cfg.data._target_}>")
+ datamodule: LightningDataModule = hydra.utils.instantiate(cfg.data, stage="test")
+
+ # Establish model input arguments
+ with open_dict(cfg):
+ if cfg.model.cfg.task.start_time == "auto":
+ cfg.model.cfg.task.start_time = 1.0
+ else:
+ cfg.model.cfg.task.start_time = float(cfg.model.cfg.task.start_time)
+
+ log.info(f"Instantiating model <{cfg.model._target_}>")
+ model: LightningModule = hydra.utils.instantiate(cfg.model)
+
+ log.info("Instantiating loggers...")
+ logger: List[Logger] = instantiate_loggers(cfg.get("logger"))
+
+ plugins = None
+ if "_target_" in cfg.environment:
+ log.info(f"Instantiating environment <{cfg.environment._target_}>")
+ plugins: ClusterEnvironment = hydra.utils.instantiate(cfg.environment)
+
+ strategy = getattr(cfg.trainer, "strategy", None)
+ if "_target_" in cfg.strategy:
+ log.info(f"Instantiating strategy <{cfg.strategy._target_}>")
+ strategy: Strategy = hydra.utils.instantiate(cfg.strategy)
+ if (
+ "mixed_precision" in strategy.__dict__
+ and getattr(strategy, "mixed_precision", None) is not None
+ ):
+ strategy.mixed_precision.param_dtype = (
+ resolve_omegaconf_variable(cfg.strategy.mixed_precision.param_dtype)
+ if getattr(cfg.strategy.mixed_precision, "param_dtype", None) is not None
+ else None
+ )
+ strategy.mixed_precision.reduce_dtype = (
+ resolve_omegaconf_variable(cfg.strategy.mixed_precision.reduce_dtype)
+ if getattr(cfg.strategy.mixed_precision, "reduce_dtype", None) is not None
+ else None
+ )
+ strategy.mixed_precision.buffer_dtype = (
+ resolve_omegaconf_variable(cfg.strategy.mixed_precision.buffer_dtype)
+ if getattr(cfg.strategy.mixed_precision, "buffer_dtype", None) is not None
+ else None
+ )
+
+ log.info(f"Instantiating trainer <{cfg.trainer._target_}>")
+ trainer: Trainer = (
+ hydra.utils.instantiate(
+ cfg.trainer,
+ logger=logger,
+ plugins=plugins,
+ strategy=strategy,
+ )
+ if strategy is not None
+ else hydra.utils.instantiate(
+ cfg.trainer,
+ logger=logger,
+ plugins=plugins,
+ )
+ )
+
+ object_dict = {
+ "cfg": cfg,
+ "datamodule": datamodule,
+ "model": model,
+ "logger": logger,
+ "trainer": trainer,
+ }
+
+ if logger:
+ log.info("Logging hyperparameters!")
+ log_hyperparameters(object_dict)
+
+ log.info("Starting testing!")
+ trainer.test(model=model, datamodule=datamodule, ckpt_path=cfg.ckpt_path)
+
+ metric_dict = trainer.callback_metrics
+
+ return metric_dict, object_dict
+
+
+@hydra.main(version_base="1.3", config_path="../configs", config_name="eval.yaml")
+def main(cfg: DictConfig) -> None:
+ """Main entry point for evaluation.
+
+ :param cfg: DictConfig configuration composed by Hydra.
+ """
+ # apply extra utilities
+ # (e.g. ask for tags if none are provided in cfg, print cfg tree, etc.)
+ extras(cfg)
+
+ evaluate(cfg)
+
+
+if __name__ == "__main__":
+ register_custom_omegaconf_resolvers()
+ main()
diff --git a/forks/FlowDock/flowdock/models/__init__.py b/forks/FlowDock/flowdock/models/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/forks/FlowDock/flowdock/models/components/__init__.py b/forks/FlowDock/flowdock/models/components/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/forks/FlowDock/flowdock/models/components/callbacks/ema.py b/forks/FlowDock/flowdock/models/components/callbacks/ema.py
new file mode 100644
index 00000000..b91eb6e5
--- /dev/null
+++ b/forks/FlowDock/flowdock/models/components/callbacks/ema.py
@@ -0,0 +1,452 @@
+# Adapted from https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/callbacks/ema.py
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import contextlib
+import copy
+import os
+import threading
+from pathlib import Path
+from typing import Any, Dict, Iterable, Optional, Union
+
+import lightning.pytorch as pl
+import torch
+from lightning.pytorch import Callback
+from lightning.pytorch.callbacks import ModelCheckpoint
+from lightning.pytorch.utilities.exceptions import MisconfigurationException
+from lightning.pytorch.utilities.rank_zero import rank_zero_info
+
+
+class EMA(Callback):
+ """Implements Exponential Moving Averaging (EMA).
+
+ When training a model, this callback will maintain moving averages of the trained parameters.
+ When evaluating, we use the moving averages copy of the trained parameters.
+ When saving, we save an additional set of parameters with the prefix `ema`.
+
+ Args:
+ decay: The exponential decay used when calculating the moving average. Has to be between 0-1.
+ validate_original_weights: Validate the original weights, as apposed to the EMA weights.
+ every_n_steps: Apply EMA every N steps.
+ cpu_offload: Offload weights to CPU.
+ """
+
+ def __init__(
+ self,
+ decay: float,
+ validate_original_weights: bool = False,
+ every_n_steps: int = 1,
+ cpu_offload: bool = False,
+ ):
+ if not (0 <= decay <= 1):
+ raise MisconfigurationException("EMA decay value must be between 0 and 1")
+ self.decay = decay
+ self.validate_original_weights = validate_original_weights
+ self.every_n_steps = every_n_steps
+ self.cpu_offload = cpu_offload
+
+ def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+ """Add the EMA optimizer to the trainer."""
+ device = pl_module.device if not self.cpu_offload else torch.device("cpu")
+ trainer.optimizers = [
+ EMAOptimizer(
+ optim,
+ device=device,
+ decay=self.decay,
+ every_n_steps=self.every_n_steps,
+ current_step=trainer.global_step,
+ )
+ for optim in trainer.optimizers
+ if not isinstance(optim, EMAOptimizer)
+ ]
+
+ def on_validation_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+ """Swap the model weights with the EMA weights."""
+ if self._should_validate_ema_weights(trainer):
+ self.swap_model_weights(trainer)
+
+ def on_validation_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+ """Swap the model weights back to the original weights."""
+ if self._should_validate_ema_weights(trainer):
+ self.swap_model_weights(trainer)
+
+ def on_test_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+ """Swap the model weights with the EMA weights."""
+ if self._should_validate_ema_weights(trainer):
+ self.swap_model_weights(trainer)
+
+ def on_test_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+ """Swap the model weights back to the original weights."""
+ if self._should_validate_ema_weights(trainer):
+ self.swap_model_weights(trainer)
+
+ def _should_validate_ema_weights(self, trainer: "pl.Trainer") -> bool:
+ """Check if the EMA weights should be validated."""
+ return not self.validate_original_weights and self._ema_initialized(trainer)
+
+ def _ema_initialized(self, trainer: "pl.Trainer") -> bool:
+ """Check if the EMA weights have been initialized."""
+ return any(isinstance(optimizer, EMAOptimizer) for optimizer in trainer.optimizers)
+
+ def swap_model_weights(self, trainer: "pl.Trainer", saving_ema_model: bool = False):
+ """Swaps the model weights with the EMA weights."""
+ for optimizer in trainer.optimizers:
+ assert isinstance(optimizer, EMAOptimizer)
+ optimizer.switch_main_parameter_weights(saving_ema_model)
+
+ @contextlib.contextmanager
+ def save_ema_model(self, trainer: "pl.Trainer"):
+ """Saves an EMA copy of the model + EMA optimizer states for resume."""
+ self.swap_model_weights(trainer, saving_ema_model=True)
+ try:
+ yield
+ finally:
+ self.swap_model_weights(trainer, saving_ema_model=False)
+
+ @contextlib.contextmanager
+ def save_original_optimizer_state(self, trainer: "pl.Trainer"):
+ """Save the original optimizer state."""
+ for optimizer in trainer.optimizers:
+ assert isinstance(optimizer, EMAOptimizer)
+ optimizer.save_original_optimizer_state = True
+ try:
+ yield
+ finally:
+ for optimizer in trainer.optimizers:
+ optimizer.save_original_optimizer_state = False
+
+ def on_load_checkpoint(
+ self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", checkpoint: Dict[str, Any]
+ ) -> None:
+ """Load the EMA state from the checkpoint if it exists."""
+ checkpoint_callback = trainer.checkpoint_callback
+
+ # use the connector as NeMo calls the connector directly in the exp_manager when restoring.
+ connector = trainer._checkpoint_connector
+ # Replace connector._ckpt_path with below to avoid calling into lightning's protected API
+ ckpt_path = trainer.ckpt_path
+
+ if (
+ ckpt_path
+ and checkpoint_callback is not None
+ and "EMA" in type(checkpoint_callback).__name__
+ ):
+ ext = checkpoint_callback.FILE_EXTENSION
+ if ckpt_path.endswith(f"-EMA{ext}"):
+ rank_zero_info(
+ "loading EMA based weights. "
+ "The callback will treat the loaded EMA weights as the main weights"
+ " and create a new EMA copy when training."
+ )
+ return
+ ema_path = ckpt_path.replace(ext, f"-EMA{ext}")
+ if os.path.exists(ema_path):
+ ema_state_dict = torch.load(ema_path, map_location=torch.device("cpu"))
+
+ checkpoint["optimizer_states"] = ema_state_dict["optimizer_states"]
+ del ema_state_dict
+ rank_zero_info("EMA state has been restored.")
+ else:
+ raise MisconfigurationException(
+ "Unable to find the associated EMA weights when re-loading, "
+ f"training will start with new EMA weights. Expected them to be at: {ema_path}",
+ )
+
+
+@torch.no_grad()
+def ema_update(ema_model_tuple, current_model_tuple, decay):
+ """Update the EMA model with the current model."""
+ torch._foreach_mul_(ema_model_tuple, decay)
+ torch._foreach_add_(
+ ema_model_tuple,
+ current_model_tuple,
+ alpha=(1.0 - decay),
+ )
+
+
+def run_ema_update_cpu(ema_model_tuple, current_model_tuple, decay, pre_sync_stream=None):
+ """Run EMA update on CPU."""
+ if pre_sync_stream is not None:
+ pre_sync_stream.synchronize()
+
+ ema_update(ema_model_tuple, current_model_tuple, decay)
+
+
+class EMAOptimizer(torch.optim.Optimizer):
+ r"""EMAOptimizer is a wrapper for torch.optim.Optimizer that computes Exponential Moving Average
+ of parameters registered in the optimizer.
+
+ EMA parameters are automatically updated after every step of the optimizer
+ with the following formula:
+
+ ema_weight = decay * ema_weight + (1 - decay) * training_weight
+
+ To access EMA parameters, use ``swap_ema_weights()`` context manager to
+ perform a temporary in-place swap of regular parameters with EMA
+ parameters.
+
+ Notes:
+ - EMAOptimizer is not compatible with APEX AMP O2.
+
+ Args:
+ optimizer (torch.optim.Optimizer): optimizer to wrap
+ device (torch.device): device for EMA parameters
+ decay (float): decay factor
+
+ Returns:
+ returns an instance of torch.optim.Optimizer that computes EMA of
+ parameters
+
+ Example:
+ model = Model().to(device)
+ opt = torch.optim.Adam(model.parameters())
+
+ opt = EMAOptimizer(opt, device, 0.9999)
+
+ for epoch in range(epochs):
+ training_loop(model, opt)
+
+ regular_eval_accuracy = evaluate(model)
+
+ with opt.swap_ema_weights():
+ ema_eval_accuracy = evaluate(model)
+ """
+
+ def __init__(
+ self,
+ optimizer: torch.optim.Optimizer,
+ device: torch.device,
+ decay: float = 0.9999,
+ every_n_steps: int = 1,
+ current_step: int = 0,
+ ):
+ self.optimizer = optimizer
+ self.decay = decay
+ self.device = device
+ self.current_step = current_step
+ self.every_n_steps = every_n_steps
+ self.save_original_optimizer_state = False
+
+ self.first_iteration = True
+ self.rebuild_ema_params = True
+ self.stream = None
+ self.thread = None
+
+ self.ema_params = ()
+ self.in_saving_ema_model_context = False
+
+ def all_parameters(self) -> Iterable[torch.Tensor]:
+ """Return an iterator over all parameters in the optimizer."""
+ return (param for group in self.param_groups for param in group["params"])
+
+ def step(self, closure=None, grad_scaler=None, **kwargs):
+ """Perform a single optimization step."""
+ self.join()
+
+ if self.first_iteration:
+ if any(p.is_cuda for p in self.all_parameters()):
+ self.stream = torch.cuda.Stream()
+
+ self.first_iteration = False
+
+ if self.rebuild_ema_params:
+ opt_params = list(self.all_parameters())
+
+ self.ema_params += tuple(
+ copy.deepcopy(param.data.detach()).to(self.device)
+ for param in opt_params[len(self.ema_params) :]
+ )
+ self.rebuild_ema_params = False
+
+ if (
+ getattr(self.optimizer, "_step_supports_amp_scaling", False)
+ and grad_scaler is not None
+ ):
+ loss = self.optimizer.step(closure=closure, grad_scaler=grad_scaler)
+ else:
+ loss = self.optimizer.step(closure)
+
+ if self._should_update_at_step():
+ self.update()
+ self.current_step += 1
+ return loss
+
+ def _should_update_at_step(self) -> bool:
+ """Check if the EMA parameters should be updated at the current step."""
+ return self.current_step % self.every_n_steps == 0
+
+ @torch.no_grad()
+ def update(self):
+ """Update the EMA parameters."""
+ if self.stream is not None:
+ self.stream.wait_stream(torch.cuda.current_stream())
+
+ with torch.cuda.stream(self.stream):
+ current_model_state = tuple(
+ param.data.to(self.device, non_blocking=True) for param in self.all_parameters()
+ )
+
+ if self.device.type == "cuda":
+ ema_update(self.ema_params, current_model_state, self.decay)
+
+ if self.device.type == "cpu":
+ self.thread = threading.Thread(
+ target=run_ema_update_cpu,
+ args=(
+ self.ema_params,
+ current_model_state,
+ self.decay,
+ self.stream,
+ ),
+ )
+ self.thread.start()
+
+ def swap_tensors(self, tensor1, tensor2):
+ """Swaps the tensors in-place."""
+ tmp = torch.empty_like(tensor1)
+ tmp.copy_(tensor1)
+ tensor1.copy_(tensor2)
+ tensor2.copy_(tmp)
+
+ def switch_main_parameter_weights(self, saving_ema_model: bool = False):
+ """Switches the main parameter weights with the EMA weights."""
+ self.join()
+ self.in_saving_ema_model_context = saving_ema_model
+ for param, ema_param in zip(self.all_parameters(), self.ema_params):
+ self.swap_tensors(param.data, ema_param)
+
+ @contextlib.contextmanager
+ def swap_ema_weights(self, enabled: bool = True):
+ r"""A context manager to in-place swap regular parameters with EMA parameters. It swaps back
+ to the original regular parameters on context manager exit.
+
+ Args:
+ enabled (bool): whether the swap should be performed
+ """
+
+ if enabled:
+ self.switch_main_parameter_weights()
+ try:
+ yield
+ finally:
+ if enabled:
+ self.switch_main_parameter_weights()
+
+ def __getattr__(self, name):
+ """Forward all other attribute calls to the optimizer."""
+ return getattr(self.optimizer, name)
+
+ def join(self):
+ """Wait for the update to complete."""
+ if self.stream is not None:
+ self.stream.synchronize()
+
+ if self.thread is not None:
+ self.thread.join()
+
+ def state_dict(self):
+ """Return the state dict for the optimizer."""
+ self.join()
+
+ if self.save_original_optimizer_state:
+ return self.optimizer.state_dict()
+
+ # if we are in the context of saving an EMA model, the EMA weights are in the modules' actual weights
+ ema_params = (
+ self.ema_params
+ if not self.in_saving_ema_model_context
+ else list(self.all_parameters())
+ )
+ state_dict = {
+ "opt": self.optimizer.state_dict(),
+ "ema": ema_params,
+ "current_step": self.current_step,
+ "decay": self.decay,
+ "every_n_steps": self.every_n_steps,
+ }
+ return state_dict
+
+ def load_state_dict(self, state_dict):
+ """Load the state dict for the optimizer."""
+ self.join()
+
+ self.optimizer.load_state_dict(state_dict["opt"])
+ self.ema_params = tuple(
+ param.to(self.device) for param in copy.deepcopy(state_dict["ema"])
+ )
+ self.current_step = state_dict["current_step"]
+ self.decay = state_dict["decay"]
+ self.every_n_steps = state_dict["every_n_steps"]
+ self.rebuild_ema_params = False
+
+ def add_param_group(self, param_group):
+ """Add a param group to the optimizer."""
+ self.optimizer.add_param_group(param_group)
+ self.rebuild_ema_params = True
+
+
+class EMAModelCheckpoint(ModelCheckpoint):
+ """EMA version of ModelCheckpoint that saves EMA checkpoints as well."""
+
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ def _ema_callback(self, trainer: "pl.Trainer") -> Optional[EMA]:
+ """Returns the EMA callback if it exists."""
+ ema_callback = None
+ for callback in trainer.callbacks:
+ if isinstance(callback, EMA):
+ ema_callback = callback
+ return ema_callback
+
+ def _save_checkpoint(self, trainer: "pl.Trainer", filepath: str) -> None:
+ """Saves the checkpoint file and the EMA checkpoint file if it exists."""
+ ema_callback = self._ema_callback(trainer)
+ if ema_callback is not None:
+ with ema_callback.save_original_optimizer_state(trainer):
+ super()._save_checkpoint(trainer, filepath)
+
+ # save EMA copy of the model as well.
+ with ema_callback.save_ema_model(trainer):
+ filepath = self._ema_format_filepath(filepath)
+ if self.verbose:
+ rank_zero_info(f"Saving EMA weights to separate checkpoint {filepath}")
+ super()._save_checkpoint(trainer, filepath)
+ else:
+ super()._save_checkpoint(trainer, filepath)
+
+ def _remove_checkpoint(self, trainer: "pl.Trainer", filepath: str) -> None:
+ """Removes the checkpoint file and the EMA checkpoint file if it exists."""
+ super()._remove_checkpoint(trainer, filepath)
+ ema_callback = self._ema_callback(trainer)
+ if ema_callback is not None:
+ # remove EMA copy of the state dict as well.
+ filepath = self._ema_format_filepath(filepath)
+ super()._remove_checkpoint(trainer, filepath)
+
+ def _ema_format_filepath(self, filepath: str) -> str:
+ """Appends '-EMA' to the filepath."""
+ return filepath.replace(self.FILE_EXTENSION, f"-EMA{self.FILE_EXTENSION}")
+
+ def _has_ema_ckpts(self, checkpoints: Iterable[Path]) -> bool:
+ """Checks if any of the checkpoints are EMA checkpoints."""
+ return any(self._is_ema_filepath(checkpoint_path) for checkpoint_path in checkpoints)
+
+ def _is_ema_filepath(self, filepath: Union[Path, str]) -> bool:
+ """Checks if the filepath is an EMA checkpoint."""
+ return str(filepath).endswith(f"-EMA{self.FILE_EXTENSION}")
+
+ @property
+ def _saved_checkpoint_paths(self) -> Iterable[Path]:
+ """Returns all the saved checkpoint paths in the directory."""
+ return Path(self.dirpath).rglob("*.ckpt")
diff --git a/forks/FlowDock/flowdock/models/components/cpm.py b/forks/FlowDock/flowdock/models/components/cpm.py
new file mode 100644
index 00000000..b66942d3
--- /dev/null
+++ b/forks/FlowDock/flowdock/models/components/cpm.py
@@ -0,0 +1,857 @@
+# Adapted from: https://github.com/zrqiao/NeuralPLexer
+
+import random
+
+import rootutils
+import torch
+from beartype.typing import Any, Dict, Optional, Tuple
+from omegaconf import DictConfig
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock.models.components.embedding import (
+ GaussianFourierEncoding1D,
+ RelativeGeometryEncoding,
+)
+from flowdock.models.components.hetero_graph import make_multi_relation_graph_batcher
+from flowdock.models.components.modules import (
+ BiDirectionalTriangleAttention,
+ TransformerLayer,
+)
+from flowdock.utils import RankedLogger
+from flowdock.utils.frame_utils import cartesian_to_internal, get_frame_matrix
+from flowdock.utils.model_utils import GELUMLP
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+STATE_DICT = Dict[str, Any]
+
+
+class ProtFormer(torch.nn.Module):
+ """Protein relational reasoning with downsampled edges."""
+
+ def __init__(
+ self,
+ dim: int,
+ pair_dim: int,
+ n_blocks: int = 4,
+ n_heads: int = 8,
+ dropout: float = 0.0,
+ ):
+ """Initialize the ProtFormer model."""
+ super().__init__()
+ self.dim = dim
+ self.pair_dim = pair_dim
+ self.n_heads = n_heads
+ self.n_blocks = n_blocks
+ self.time_encoding = GaussianFourierEncoding1D(16)
+ self.res_in_mlp = GELUMLP(dim + 32, dim, dropout=dropout)
+ self.chain_pos_encoding = GaussianFourierEncoding1D(self.pair_dim // 4)
+
+ self.rel_geom_enc = RelativeGeometryEncoding(15, self.pair_dim)
+ self.template_nenc = GELUMLP(64 + 37 * 3, self.dim, n_hidden_feats=128)
+ self.template_eenc = RelativeGeometryEncoding(15, self.pair_dim)
+ self.template_binding_site_enc = torch.nn.Linear(1, 64, bias=False)
+ self.pp_edge_embed = GELUMLP(
+ pair_dim + self.pair_dim // 4 * 2 + dim * 2,
+ self.pair_dim,
+ n_hidden_feats=dim,
+ dropout=dropout,
+ )
+
+ self.graph_stacks = torch.nn.ModuleList(
+ [
+ TransformerLayer(
+ dim,
+ n_heads,
+ head_dim=pair_dim // n_heads,
+ edge_channels=pair_dim,
+ edge_update=True,
+ dropout=dropout,
+ )
+ for _ in range(self.n_blocks)
+ ]
+ )
+ self.ABab_mha = TransformerLayer(
+ pair_dim,
+ n_heads,
+ bidirectional=True,
+ )
+
+ self.triangle_stacks = torch.nn.ModuleList(
+ [
+ BiDirectionalTriangleAttention(pair_dim, pair_dim // n_heads, n_heads)
+ for _ in range(self.n_blocks)
+ ]
+ )
+ self.graph_relations = [
+ (
+ "residue_to_residue",
+ "gather_idx_ab_a",
+ "gather_idx_ab_b",
+ "prot_res",
+ "prot_res",
+ ),
+ (
+ "sampled_residue_to_sampled_residue",
+ "gather_idx_AB_a",
+ "gather_idx_AB_b",
+ "prot_res",
+ "prot_res",
+ ),
+ ]
+
+ def compute_chain_pe(
+ self,
+ residue_index,
+ res_chain_index,
+ src_rid,
+ dst_rid,
+ ):
+ """Compute chain positional encoding for a pair of residues."""
+ chain_disp = residue_index[src_rid] - residue_index[dst_rid]
+ chain_rope = self.chain_pos_encoding(chain_disp.div(8).unsqueeze(-1)).div(
+ chain_disp.div(8).abs().add(1).unsqueeze(-1)
+ )
+ # Mask cross-chain entries
+ chain_mask = res_chain_index[src_rid] == res_chain_index[dst_rid]
+ chain_rope = chain_rope * chain_mask[..., None]
+ return chain_rope
+
+ def compute_chain_pair_pe(
+ self,
+ residue_index,
+ res_chain_index,
+ AB_broadcasted_rid,
+ ab_rid,
+ AB_broadcasted_cid,
+ ab_cid,
+ ):
+ """Compute chain positional encoding for a pair of residues."""
+ chain_disp_row = residue_index[AB_broadcasted_rid] - residue_index[ab_rid]
+ chain_disp_col = residue_index[AB_broadcasted_cid] - residue_index[ab_cid]
+ chain_disp = torch.stack([chain_disp_row, chain_disp_col], dim=-1)
+ chain_rope = self.chain_pos_encoding(chain_disp.div(8).unsqueeze(-1)).div(
+ chain_disp.div(8).abs().add(1).unsqueeze(-1)
+ )
+ # Mask cross-chain entries
+ chain_mask_row = res_chain_index[AB_broadcasted_rid] == res_chain_index[ab_rid]
+ chain_mask_col = res_chain_index[AB_broadcasted_cid] == res_chain_index[ab_cid]
+ chain_mask = torch.stack([chain_mask_row, chain_mask_col], dim=-1)
+ chain_rope = (chain_rope * chain_mask[..., None]).flatten(-2, -1)
+ return chain_rope
+
+ def eval_protein_template_encodings(self, batch, edge_idx, use_plddt=False):
+ """Evaluate template encodings for protein residues."""
+ with torch.no_grad():
+ template_bb_coords = batch["features"]["apo_res_atom_positions"][:, :3]
+ template_bb_frames = get_frame_matrix(
+ template_bb_coords[:, 0, :],
+ template_bb_coords[:, 1, :],
+ template_bb_coords[:, 2, :],
+ )
+ # Add template local representations & lddt
+ template_local_coords = cartesian_to_internal(
+ batch["features"]["apo_res_atom_positions"],
+ template_bb_frames.unsqueeze(1),
+ )
+ template_local_coords[~batch["features"]["apo_res_atom_mask"].bool()] = 0
+ # if use_plddt:
+ # template_plddt_enc = F.one_hot(
+ # torch.bucketize(
+ # batch["features"]["apo_pLDDT"],
+ # torch.linspace(0, 1, 65, device=template_bb_coords.device)[:-1],
+ # right=True,
+ # )
+ # - 1,
+ # num_classes=64,
+ # )
+ # else:
+ # template_plddt_enc = torch.zeros(
+ # template_local_coords.shape[0], 64, device=template_bb_coords.device
+ # )
+ if self.training:
+ use_sidechain_coords = random.randint(0, 1) # nosec
+ template_local_coords = template_local_coords * use_sidechain_coords
+ # use_plddt_input = random.randint(0, 1)
+ # template_plddt_enc = template_plddt_enc * use_plddt_input
+ if "binding_site_mask" in batch["features"].keys():
+ # Externally-specified binding residue list
+ binding_site_enc = self.template_binding_site_enc(
+ batch["features"]["binding_site_mask"][:, None].float()
+ )
+ else:
+ binding_site_enc = torch.zeros(
+ template_local_coords.shape[0], 64, device=template_bb_coords.device
+ )
+ template_nfeat = self.template_nenc(
+ torch.cat([template_local_coords.flatten(-2, -1), binding_site_enc], dim=-1)
+ )
+ template_efeat = self.template_eenc(template_bb_frames, edge_idx)
+ template_alignment_mask = batch["features"]["apo_res_alignment_mask"].float()
+ if self.training:
+ # template_alignment_mask = template_alignment_mask * use_template
+ nomasking_rate = random.randint(9, 10) / 10 # nosec
+ template_alignment_mask = template_alignment_mask * (
+ torch.rand_like(template_alignment_mask) < nomasking_rate
+ )
+ template_nfeat = template_nfeat * template_alignment_mask.unsqueeze(-1)
+ template_efeat = (
+ template_efeat
+ * template_alignment_mask[edge_idx[0]].unsqueeze(-1)
+ * template_alignment_mask[edge_idx[1]].unsqueeze(-1)
+ )
+ return template_nfeat, template_efeat
+
+ def forward(self, batch, **kwargs):
+ """Forward pass of the ProtFormer model."""
+ return self.forward_prot_sample(batch, **kwargs)
+
+ def forward_prot_sample(
+ self,
+ batch,
+ embed_coords=True,
+ in_attr_suffix="",
+ out_attr_suffix="",
+ use_template=False,
+ use_plddt=False,
+ **kwargs,
+ ):
+ """Forward pass of the ProtFormer model for a single protein sample."""
+ features = batch["features"]
+ indexer = batch["indexer"]
+ metadata = batch["metadata"]
+ device = features["res_type"].device
+
+ time_encoding = self.time_encoding(features["timestep_encoding_prot"])
+ if not embed_coords:
+ time_encoding = torch.zeros_like(time_encoding)
+
+ residue_rep = (
+ self.res_in_mlp(
+ torch.cat(
+ [
+ features["res_embedding_in"],
+ time_encoding,
+ ],
+ dim=1,
+ )
+ )
+ + features["res_embedding_in"]
+ )
+ batch_size = metadata["num_structid"]
+
+ # Prepare indexers
+ # Use max to ensure segmentation faults are 100% invoked
+ # in case there are any bad indices
+ max(metadata["num_a_per_sample"])
+ n_protein_patches = batch["metadata"]["n_prot_patches_per_sample"]
+
+ indexer["gather_idx_pid_b"] = indexer["gather_idx_pid_a"]
+ # Evaluate gather_idx_AB_a and gather_idx_AB_b
+ # Assign a to rows and b to columns
+ # Simple broadcasting for single-structure batches
+ indexer["gather_idx_AB_a"] = (
+ indexer["gather_idx_pid_a"]
+ .view(batch_size, n_protein_patches)[:, :, None]
+ .expand(-1, -1, n_protein_patches)
+ .contiguous()
+ .flatten()
+ )
+ indexer["gather_idx_AB_b"] = (
+ indexer["gather_idx_pid_b"]
+ .view(batch_size, n_protein_patches)[:, None, :]
+ .expand(-1, n_protein_patches, -1)
+ .contiguous()
+ .flatten()
+ )
+
+ # Handle all batch offsets here
+ graph_batcher = make_multi_relation_graph_batcher(self.graph_relations, indexer, metadata)
+ merged_edge_idx = graph_batcher.collate_idx_list(indexer)
+
+ input_protein_coords_padded = features["input_protein_coords"]
+ backbone_frames = get_frame_matrix(
+ input_protein_coords_padded[:, 0, :],
+ input_protein_coords_padded[:, 1, :],
+ input_protein_coords_padded[:, 2, :],
+ )
+ batch["features"]["backbone_frames"] = backbone_frames
+ # Adding geometrical info to pair representations
+
+ chain_pe = self.compute_chain_pe(
+ features["residue_index"],
+ features["res_chain_id"],
+ merged_edge_idx[0],
+ merged_edge_idx[1],
+ )
+ geometry_pe = self.rel_geom_enc(backbone_frames, merged_edge_idx)
+ if not embed_coords:
+ geometry_pe = torch.zeros_like(geometry_pe)
+ merged_edge_reps = self.pp_edge_embed(
+ torch.cat(
+ [
+ geometry_pe,
+ chain_pe,
+ residue_rep[merged_edge_idx[0]],
+ residue_rep[merged_edge_idx[1]],
+ ],
+ dim=-1,
+ )
+ )
+ if use_template and "apo_res_atom_positions" in features.keys():
+ (
+ template_res_encodings,
+ template_geom_encodings,
+ ) = self.eval_protein_template_encodings(batch, merged_edge_idx, use_plddt=use_plddt)
+ residue_rep = residue_rep + template_res_encodings
+ merged_edge_reps = merged_edge_reps + template_geom_encodings
+ edge_reps = graph_batcher.offload_edge_attr(merged_edge_reps)
+
+ node_reps = {"prot_res": residue_rep}
+
+ gather_idx_res_protpatch = indexer["gather_idx_a_pid"]
+ # Pointer: AB->AB, ab->AB
+ gather_idx_ab_AB = (
+ indexer["gather_idx_ab_structid"] * n_protein_patches**2
+ + (gather_idx_res_protpatch % n_protein_patches)[indexer["gather_idx_ab_a"]]
+ * n_protein_patches
+ + (gather_idx_res_protpatch % n_protein_patches)[indexer["gather_idx_ab_b"]]
+ )
+
+ # Intertwine graph iterations and triangle iterations
+ for block_id in range(self.n_blocks):
+ # Communicate between atomistic and patch resolutions
+ # Up-sampling for interface edge embeddings
+ rec_pair_rep = edge_reps["residue_to_residue"]
+ AB_grid_attr_flat = edge_reps["sampled_residue_to_sampled_residue"]
+ # Upper-left block: intra-window visual-attention
+ # Cross-attention between random and grid edges
+ rec_pair_rep, AB_grid_attr_flat = self.ABab_mha(
+ rec_pair_rep,
+ AB_grid_attr_flat,
+ (
+ torch.arange(metadata["num_ab"], device=device),
+ gather_idx_ab_AB,
+ ),
+ )
+ AB_grid_attr = AB_grid_attr_flat.view(
+ batch_size,
+ n_protein_patches,
+ n_protein_patches,
+ self.pair_dim,
+ )
+
+ # Inter-patch triangle attentions, refining intermolecular edges
+ _, AB_grid_attr = self.triangle_stacks[block_id](
+ AB_grid_attr,
+ AB_grid_attr,
+ AB_grid_attr.unsqueeze(-4),
+ )
+
+ # Transfer grid-formatted representations to edges
+ edge_reps["residue_to_residue"] = rec_pair_rep
+ edge_reps["sampled_residue_to_sampled_residue"] = AB_grid_attr.flatten(0, 2)
+ merged_node_reps = graph_batcher.collate_node_attr(node_reps)
+ merged_edge_reps = graph_batcher.collate_edge_attr(edge_reps)
+
+ # Graph transformer iteration
+ _, merged_node_reps, merged_edge_reps = self.graph_stacks[block_id](
+ merged_node_reps,
+ merged_node_reps,
+ merged_edge_idx,
+ merged_edge_reps,
+ )
+ node_reps = graph_batcher.offload_node_attr(merged_node_reps)
+ edge_reps = graph_batcher.offload_edge_attr(merged_edge_reps)
+
+ batch["features"][f"rec_res_attr{out_attr_suffix}"] = node_reps["prot_res"]
+ batch["features"][f"res_res_pair_attr{out_attr_suffix}"] = edge_reps["residue_to_residue"]
+ batch["features"][f"res_res_grid_attr_flat{out_attr_suffix}"] = edge_reps[
+ "sampled_residue_to_sampled_residue"
+ ]
+ batch["indexer"]["gather_idx_AB_a"] = indexer["gather_idx_AB_a"]
+ batch["indexer"]["gather_idx_AB_b"] = indexer["gather_idx_AB_b"]
+ batch["indexer"]["gather_idx_ab_AB"] = gather_idx_ab_AB
+ return batch
+
+
+class BindingFormer(ProtFormer):
+ """Edge inference on protein-ligand graphs."""
+
+ def __init__(
+ self,
+ dim: int,
+ pair_dim: int,
+ n_blocks: int = 4,
+ n_heads: int = 8,
+ n_ligand_patches: int = 16,
+ dropout: float = 0.0,
+ ):
+ """Initialize the BindingFormer model."""
+ super().__init__(
+ dim,
+ pair_dim,
+ n_blocks,
+ n_heads,
+ dropout,
+ )
+ self.dim = dim
+ self.n_heads = n_heads
+ self.n_blocks = n_blocks
+ self.n_ligand_patches = n_ligand_patches
+ self.pl_edge_embed = GELUMLP(dim * 2, self.pair_dim, n_hidden_feats=dim, dropout=dropout)
+ self.AaJ_mha = TransformerLayer(pair_dim, n_heads, bidirectional=True)
+
+ self.graph_relations = [
+ (
+ "residue_to_residue",
+ "gather_idx_ab_a",
+ "gather_idx_ab_b",
+ "prot_res",
+ "prot_res",
+ ),
+ (
+ "sampled_residue_to_sampled_residue",
+ "gather_idx_AB_a",
+ "gather_idx_AB_b",
+ "prot_res",
+ "prot_res",
+ ),
+ (
+ "sampled_residue_to_sampled_lig_triplet",
+ "gather_idx_AJ_a",
+ "gather_idx_AJ_J",
+ "prot_res",
+ "lig_trp",
+ ),
+ (
+ "sampled_lig_triplet_to_sampled_residue",
+ "gather_idx_AJ_J",
+ "gather_idx_AJ_a",
+ "lig_trp",
+ "prot_res",
+ ),
+ (
+ "residue_to_sampled_lig_triplet",
+ "gather_idx_aJ_a",
+ "gather_idx_aJ_J",
+ "prot_res",
+ "lig_trp",
+ ),
+ (
+ "sampled_lig_triplet_to_residue",
+ "gather_idx_aJ_J",
+ "gather_idx_aJ_a",
+ "lig_trp",
+ "prot_res",
+ ),
+ (
+ "sampled_lig_triplet_to_sampled_lig_triplet",
+ "gather_idx_IJ_I",
+ "gather_idx_IJ_J",
+ "lig_trp",
+ "lig_trp",
+ ),
+ ]
+
+ def forward(
+ self,
+ batch,
+ observed_block_contacts=None,
+ in_attr_suffix="",
+ out_attr_suffix="",
+ ):
+ """Forward pass of the BindingFormer model."""
+ features = batch["features"]
+ indexer = batch["indexer"]
+ metadata = batch["metadata"]
+ device = features["res_type"].device
+ # Synchronize with a language model
+ residue_rep = features[f"rec_res_attr{in_attr_suffix}"]
+ rec_pair_rep = features[f"res_res_pair_attr{in_attr_suffix}"]
+ # Inherit the last-layer pair representations from protein encoder
+ AB_grid_attr_flat = features[f"res_res_grid_attr_flat{in_attr_suffix}"]
+
+ # Prepare indexers
+ batch_size = metadata["num_structid"]
+ n_a_per_sample = max(metadata["num_a_per_sample"])
+ n_protein_patches = batch["metadata"]["n_prot_patches_per_sample"]
+
+ if not batch["misc"]["protein_only"]:
+ n_ligand_patches = max(metadata["num_I_per_sample"])
+ max(metadata["num_molid_per_sample"])
+ lig_frame_rep = features[f"lig_trp_attr{in_attr_suffix}"]
+ UI_grid_attr = features["lig_af_grid_attr_projected"]
+ IJ_grid_attr = (UI_grid_attr + UI_grid_attr.transpose(1, 2)) / 2
+
+ aJ_grid_attr = self.pl_edge_embed(
+ torch.cat(
+ [
+ residue_rep.view(batch_size, n_a_per_sample, self.dim)[:, :, None].expand(
+ -1, -1, n_ligand_patches, -1
+ ),
+ lig_frame_rep.view(batch_size, n_ligand_patches, self.dim)[
+ :, None, :
+ ].expand(-1, n_a_per_sample, -1, -1),
+ ],
+ dim=-1,
+ )
+ )
+ AJ_grid_attr = IJ_grid_attr.new_zeros(
+ batch_size, n_protein_patches, n_ligand_patches, self.pair_dim
+ )
+ gather_idx_I_I = torch.arange(
+ batch_size * n_ligand_patches, device=AJ_grid_attr.device
+ )
+ gather_idx_a_a = torch.arange(batch_size * n_a_per_sample, device=AJ_grid_attr.device)
+ # Note: off-diagonal (AJ) blocks are zero-initialized in the prior stack
+ indexer["gather_idx_IJ_I"] = (
+ gather_idx_I_I.view(batch_size, n_ligand_patches)[:, :, None]
+ .expand(-1, -1, n_ligand_patches)
+ .contiguous()
+ .flatten()
+ )
+ indexer["gather_idx_IJ_J"] = (
+ gather_idx_I_I.view(batch_size, n_ligand_patches)[:, None, :]
+ .expand(-1, n_ligand_patches, -1)
+ .contiguous()
+ .flatten()
+ )
+ indexer["gather_idx_AJ_a"] = (
+ indexer["gather_idx_pid_a"]
+ .view(batch_size, n_protein_patches)[:, :, None]
+ .expand(-1, -1, n_ligand_patches)
+ .contiguous()
+ .flatten()
+ )
+ indexer["gather_idx_AJ_J"] = (
+ gather_idx_I_I.view(batch_size, n_ligand_patches)[:, None, :]
+ .expand(-1, n_protein_patches, -1)
+ .contiguous()
+ .flatten()
+ )
+ indexer["gather_idx_aJ_a"] = (
+ gather_idx_a_a.view(batch_size, n_a_per_sample)[:, :, None]
+ .expand(-1, -1, n_ligand_patches)
+ .contiguous()
+ .flatten()
+ )
+ indexer["gather_idx_aJ_J"] = (
+ gather_idx_I_I.view(batch_size, n_ligand_patches)[:, None, :]
+ .expand(-1, n_a_per_sample, -1)
+ .contiguous()
+ .flatten()
+ )
+ batch["indexer"] = indexer
+
+ if observed_block_contacts is not None:
+ # Generative feedback from block one-hot sampling
+ # AJ_grid_attr = (
+ # AJ_grid_attr
+ # + observed_block_contacts.transpose(1, 2)
+ # .contiguous()
+ # .flatten(0, 1)[indexer["gather_idx_I_molid"]]
+ # .view(batch_size, n_ligand_patches, n_protein_patches, -1)
+ # .transpose(1, 2)
+ # .contiguous()
+ # )
+ AJ_grid_attr = AJ_grid_attr + observed_block_contacts
+
+ graph_batcher = make_multi_relation_graph_batcher(
+ self.graph_relations, indexer, metadata
+ )
+ merged_edge_idx = graph_batcher.collate_idx_list(indexer)
+ node_reps = {
+ "prot_res": residue_rep,
+ "lig_trp": lig_frame_rep,
+ }
+ edge_reps = {
+ "residue_to_residue": rec_pair_rep,
+ "sampled_residue_to_sampled_residue": AB_grid_attr_flat,
+ "sampled_lig_triplet_to_sampled_residue": AJ_grid_attr.flatten(0, 2),
+ "sampled_residue_to_sampled_lig_triplet": AJ_grid_attr.flatten(0, 2),
+ "sampled_lig_triplet_to_residue": aJ_grid_attr.flatten(0, 2),
+ "residue_to_sampled_lig_triplet": aJ_grid_attr.flatten(0, 2),
+ "sampled_lig_triplet_to_sampled_lig_triplet": IJ_grid_attr.flatten(0, 2),
+ }
+ edge_reps = graph_batcher.zero_pad_edge_attr(edge_reps, self.dim, device)
+ else:
+ graph_batcher = make_multi_relation_graph_batcher(
+ self.graph_relations[:2], indexer, metadata
+ )
+ merged_edge_idx = graph_batcher.collate_idx_list(indexer)
+
+ node_reps = {
+ "prot_res": residue_rep,
+ }
+ edge_reps = {
+ "residue_to_residue": rec_pair_rep,
+ "sampled_residue_to_sampled_residue": AB_grid_attr_flat,
+ }
+ edge_reps = graph_batcher.zero_pad_edge_attr(edge_reps, self.dim, device)
+
+ # Intertwine graph iterations and triangle iterations
+ gather_idx_res_protpatch = indexer["gather_idx_a_pid"]
+ for block_id in range(self.n_blocks):
+ # Communicate between atomistic and patch resolutions
+ # Up-sampling for interface edge embeddings
+ rec_pair_rep = edge_reps["residue_to_residue"]
+ AB_grid_attr_flat = edge_reps["sampled_residue_to_sampled_residue"]
+ AB_grid_attr = AB_grid_attr_flat.view(
+ batch_size,
+ n_protein_patches,
+ n_protein_patches,
+ self.pair_dim,
+ )
+
+ if not batch["misc"]["protein_only"]:
+ # Symmetrize off-diagonal blocks
+ AJ_grid_attr_flat_ = (
+ edge_reps["sampled_residue_to_sampled_lig_triplet"]
+ + edge_reps["sampled_lig_triplet_to_sampled_residue"]
+ ) / 2
+ AJ_grid_attr = AJ_grid_attr_flat_.contiguous().view(
+ batch_size, n_protein_patches, n_ligand_patches, -1
+ )
+ aJ_grid_attr_flat_ = (
+ edge_reps["residue_to_sampled_lig_triplet"]
+ + edge_reps["sampled_lig_triplet_to_residue"]
+ ) / 2
+ aJ_grid_attr = aJ_grid_attr_flat_.contiguous().view(
+ batch_size, n_a_per_sample, n_ligand_patches, -1
+ )
+ IJ_grid_attr = (
+ edge_reps["sampled_lig_triplet_to_sampled_lig_triplet"]
+ .contiguous()
+ .view(batch_size, n_ligand_patches, n_ligand_patches, -1)
+ )
+ AJ_grid_attr_temp_, aJ_grid_attr_temp_ = self.AaJ_mha(
+ AJ_grid_attr.flatten(0, 1),
+ aJ_grid_attr.flatten(0, 1),
+ (
+ gather_idx_res_protpatch,
+ torch.arange(gather_idx_res_protpatch.shape[0], device=device),
+ ),
+ )
+ AJ_grid_attr = AJ_grid_attr_temp_.contiguous().view(
+ batch_size, n_protein_patches, n_ligand_patches, -1
+ )
+ aJ_grid_attr = aJ_grid_attr_temp_.contiguous().view(
+ batch_size, n_a_per_sample, n_ligand_patches, -1
+ )
+ merged_grid_rep = torch.cat(
+ [
+ torch.cat([AB_grid_attr, AJ_grid_attr], dim=2),
+ torch.cat([AJ_grid_attr.transpose(1, 2), IJ_grid_attr], dim=2),
+ ],
+ dim=1,
+ )
+ else:
+ merged_grid_rep = AB_grid_attr
+
+ # Inter-patch triangle attentions
+ _, merged_grid_rep = self.triangle_stacks[block_id](
+ merged_grid_rep,
+ merged_grid_rep,
+ merged_grid_rep.unsqueeze(-4),
+ )
+
+ # Dis-assemble the grid representation
+ AB_grid_attr = merged_grid_rep[:, :n_protein_patches, :n_protein_patches]
+ # Transfer grid-formatted representations to edges
+ edge_reps["residue_to_residue"] = rec_pair_rep
+ edge_reps["sampled_residue_to_sampled_residue"] = AB_grid_attr.flatten(0, 2)
+
+ if not batch["misc"]["protein_only"]:
+ AJ_grid_attr = merged_grid_rep[
+ :, :n_protein_patches, n_protein_patches:
+ ].contiguous()
+ IJ_grid_attr = merged_grid_rep[
+ :, n_protein_patches:, n_protein_patches:
+ ].contiguous()
+
+ edge_reps["sampled_residue_to_sampled_lig_triplet"] = AJ_grid_attr.flatten(0, 2)
+ edge_reps["sampled_lig_triplet_to_sampled_residue"] = AJ_grid_attr.flatten(0, 2)
+ edge_reps["residue_to_sampled_lig_triplet"] = aJ_grid_attr.flatten(0, 2)
+ edge_reps["sampled_lig_triplet_to_residue"] = aJ_grid_attr.flatten(0, 2)
+ edge_reps["sampled_lig_triplet_to_sampled_lig_triplet"] = IJ_grid_attr.flatten(
+ 0, 2
+ )
+ merged_node_reps = graph_batcher.collate_node_attr(node_reps)
+ merged_edge_reps = graph_batcher.collate_edge_attr(edge_reps)
+
+ # Graph transformer iteration
+ _, merged_node_reps, merged_edge_reps = self.graph_stacks[block_id](
+ merged_node_reps,
+ merged_node_reps,
+ merged_edge_idx,
+ merged_edge_reps,
+ )
+ node_reps = graph_batcher.offload_node_attr(merged_node_reps)
+ edge_reps = graph_batcher.offload_edge_attr(merged_edge_reps)
+
+ batch["features"][f"rec_res_attr{out_attr_suffix}"] = node_reps["prot_res"]
+ batch["features"][f"res_res_pair_attr{out_attr_suffix}"] = edge_reps["residue_to_residue"]
+ batch["features"][f"res_res_grid_attr_flat{out_attr_suffix}"] = edge_reps[
+ "sampled_residue_to_sampled_residue"
+ ]
+ if not batch["misc"]["protein_only"]:
+ batch["features"][f"lig_trp_attr{out_attr_suffix}"] = node_reps["lig_trp"]
+ batch["features"][f"res_trp_grid_attr_flat{out_attr_suffix}"] = edge_reps[
+ "sampled_residue_to_sampled_lig_triplet"
+ ]
+ batch["features"][f"res_trp_pair_attr_flat{out_attr_suffix}"] = edge_reps[
+ "residue_to_sampled_lig_triplet"
+ ]
+ batch["features"][f"trp_trp_grid_attr_flat{out_attr_suffix}"] = edge_reps[
+ "sampled_lig_triplet_to_sampled_lig_triplet"
+ ]
+ batch["metadata"]["n_lig_patches_per_sample"] = n_ligand_patches
+ return batch
+
+
+def resolve_protein_encoder(
+ protein_model_cfg: DictConfig,
+ task_cfg: DictConfig,
+ state_dict: Optional[STATE_DICT] = None,
+) -> Tuple[torch.nn.Module, torch.nn.Module]:
+ """Instantiates a ProtFormer model for protein encoding.
+
+ :param protein_model_cfg: Protein model configuration.
+ :param task_cfg: Task configuration.
+ :param state_dict: Optional (potentially-pretrained) state dictionary.
+ :return: Protein encoder model and residue input projector.
+ """
+ node_dim = protein_model_cfg.residue_dim
+ model = ProtFormer(
+ node_dim,
+ protein_model_cfg.pair_dim,
+ n_heads=protein_model_cfg.n_heads,
+ n_blocks=protein_model_cfg.n_encoder_stacks,
+ dropout=task_cfg.dropout,
+ )
+ if protein_model_cfg.use_esm_embedding:
+ # protein sequence language model
+ res_in_projector = torch.nn.Linear(protein_model_cfg.plm_embed_dim, node_dim, bias=False)
+ else:
+ # one-hot amino acid types
+ res_in_projector = torch.nn.Linear(
+ protein_model_cfg.n_aa_types,
+ node_dim,
+ bias=False,
+ )
+ if protein_model_cfg.from_pretrained and state_dict is not None:
+ try:
+ # NOTE: we must avoid enforcing strict key matching
+ # due to the (new) weights `template_binding_site_enc.weight`
+ model.load_state_dict(
+ {
+ ".".join(k.split(".")[1:]): v
+ for k, v in state_dict.items()
+ if k.startswith("protein_encoder")
+ },
+ strict=False,
+ )
+ log.info("Successfully loaded pretrained protein encoder weights.")
+ except Exception as e:
+ log.warning(f"Skipping loading of pretrained protein encoder weights due to: {e}.")
+
+ try:
+ res_in_projector.load_state_dict(
+ {
+ ".".join(k.split(".")[1:]): v
+ for k, v in state_dict.items()
+ if k.startswith(
+ "plm_adapter"
+ if protein_model_cfg.use_esm_embedding
+ else "res_in_projector"
+ )
+ }
+ )
+ log.info("Successfully loaded pretrained protein input projector weights.")
+ except Exception as e:
+ log.warning(
+ f"Skipping loading of pretrained protein input projector weights due to: {e}."
+ )
+ return model, res_in_projector
+
+
+def resolve_pl_contact_stack(
+ protein_model_cfg: DictConfig,
+ ligand_model_cfg: DictConfig,
+ contact_cfg: DictConfig,
+ task_cfg: DictConfig,
+ state_dict: Optional[STATE_DICT] = None,
+) -> Tuple[torch.nn.Module, torch.nn.Module, torch.nn.Module, torch.nn.Module]:
+ """Instantiates a BindingFormer model for protein-ligand contact prediction.
+
+ :param protein_model_cfg: Protein model configuration.
+ :param ligand_model_cfg: Ligand model configuration.
+ :param contact_cfg: Contact prediction configuration.
+ :param task_cfg: Task configuration.
+ :param state_dict: Optional (potentially-pretrained) state dictionary.
+ :return: Protein-ligand contact prediction model, contact code embedding, distance bins, and
+ distogram head.
+ """
+ pl_contact_stack = BindingFormer(
+ protein_model_cfg.residue_dim,
+ protein_model_cfg.pair_dim,
+ n_heads=protein_model_cfg.n_heads,
+ n_blocks=contact_cfg.n_stacks,
+ n_ligand_patches=ligand_model_cfg.n_patches,
+ dropout=contact_cfg.dropout if contact_cfg.get("dropout") else task_cfg.dropout,
+ )
+ contact_code_embed = torch.nn.Embedding(2, protein_model_cfg.pair_dim)
+ # Distogram heads
+ dist_bins = torch.nn.Parameter(torch.linspace(2, 22, 32), requires_grad=False)
+ dgram_head = GELUMLP(
+ protein_model_cfg.pair_dim,
+ 32,
+ n_hidden_feats=protein_model_cfg.pair_dim,
+ zero_init=True,
+ )
+ if contact_cfg.from_pretrained and state_dict is not None:
+ try:
+ # NOTE: we must avoid enforcing strict key matching
+ # due to the (new) weights `template_binding_site_enc.weight`
+ pl_contact_stack.load_state_dict(
+ {
+ ".".join(k.split(".")[1:]): v
+ for k, v in state_dict.items()
+ if k.startswith("pl_contact_stack")
+ },
+ strict=False,
+ )
+ log.info("Successfully loaded pretrained protein-ligand contact prediction weights.")
+ except Exception as e:
+ log.warning(
+ f"Skipping loading of pretrained protein-ligand contact prediction weights due to: {e}."
+ )
+
+ try:
+ contact_code_embed.load_state_dict(
+ {
+ ".".join(k.split(".")[1:]): v
+ for k, v in state_dict.items()
+ if k.startswith("contact_code_embed")
+ }
+ )
+ log.info("Successfully loaded pretrained contact code embedding weights.")
+ except Exception as e:
+ log.warning(
+ f"Skipping loading of pretrained contact code embedding weights due to: {e}."
+ )
+
+ try:
+ dgram_head.load_state_dict(
+ {
+ ".".join(k.split(".")[1:]): v
+ for k, v in state_dict.items()
+ if k.startswith("dgram_head")
+ }
+ )
+ log.info("Successfully loaded pretrained distogram head weights.")
+ except Exception as e:
+ log.warning(f"Skipping loading of pretrained distogram head weights due to: {e}.")
+ return pl_contact_stack, contact_code_embed, dist_bins, dgram_head
diff --git a/forks/FlowDock/flowdock/models/components/embedding.py b/forks/FlowDock/flowdock/models/components/embedding.py
new file mode 100644
index 00000000..016ec3ca
--- /dev/null
+++ b/forks/FlowDock/flowdock/models/components/embedding.py
@@ -0,0 +1,105 @@
+# Adapted from: https://github.com/zrqiao/NeuralPLexer
+
+import math
+
+import rootutils
+import torch
+from beartype.typing import Tuple
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock.utils.frame_utils import RigidTransform
+
+
+class GaussianFourierEncoding1D(torch.nn.Module):
+ """Gaussian Fourier Encoding for 1D data."""
+
+ def __init__(
+ self,
+ n_basis: int,
+ eps: float = 1e-2,
+ ):
+ """Initialize Gaussian Fourier Encoding."""
+ super().__init__()
+ self.eps = eps
+ self.fourier_freqs = torch.nn.Parameter(torch.randn(n_basis) * math.pi)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ ):
+ """Forward pass of Gaussian Fourier Encoding."""
+ encodings = torch.cat(
+ [
+ torch.sin(self.fourier_freqs.mul(x)),
+ torch.cos(self.fourier_freqs.mul(x)),
+ ],
+ dim=-1,
+ )
+ return encodings
+
+
+class GaussianRBFEncoding1D(torch.nn.Module):
+ """Gaussian RBF Encoding for 1D data."""
+
+ def __init__(
+ self,
+ n_basis: int,
+ x_max: float,
+ sigma: float = 1.0,
+ ):
+ """Initialize Gaussian RBF Encoding."""
+ super().__init__()
+ self.sigma = sigma
+ self.rbf_centers = torch.nn.Parameter(
+ torch.linspace(0, x_max, n_basis), requires_grad=False
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ ):
+ """Forward pass of Gaussian RBF Encoding."""
+ encodings = torch.exp(-((x.unsqueeze(-1) - self.rbf_centers).div(self.sigma).square()))
+ return encodings
+
+
+class RelativeGeometryEncoding(torch.nn.Module):
+ "Compute radial basis functions and iterresidue/pseudoresidue orientations."
+
+ def __init__(self, n_basis: int, out_dim: int, d_max: float = 20.0):
+ """Initialize RelativeGeometryEncoding."""
+ super().__init__()
+ self.rbf_encoding = GaussianRBFEncoding1D(n_basis, d_max)
+ self.rel_geom_projector = torch.nn.Linear(n_basis + 15, out_dim, bias=False)
+
+ def forward(self, frames: RigidTransform, merged_edge_idx: Tuple[torch.Tensor, torch.Tensor]):
+ """Forward pass of RelativeGeometryEncoding."""
+ frame_t, frame_R = frames.t, frames.R
+ pair_dists = torch.norm(
+ frame_t[merged_edge_idx[0]] - frame_t[merged_edge_idx[1]],
+ dim=-1,
+ )
+ pair_directions_l = torch.matmul(
+ (frame_t[merged_edge_idx[1]] - frame_t[merged_edge_idx[0]]).unsqueeze(-2),
+ frame_R[merged_edge_idx[0]],
+ ).squeeze(-2) / pair_dists.square().add(1).sqrt().unsqueeze(-1)
+ pair_directions_r = torch.matmul(
+ (frame_t[merged_edge_idx[0]] - frame_t[merged_edge_idx[1]]).unsqueeze(-2),
+ frame_R[merged_edge_idx[1]],
+ ).squeeze(-2) / pair_dists.square().add(1).sqrt().unsqueeze(-1)
+ pair_orientations = torch.matmul(
+ frame_R.transpose(-2, -1).contiguous()[merged_edge_idx[0]],
+ frame_R[merged_edge_idx[1]],
+ )
+ return self.rel_geom_projector(
+ torch.cat(
+ [
+ self.rbf_encoding(pair_dists),
+ pair_directions_l,
+ pair_directions_r,
+ pair_orientations.flatten(-2, -1),
+ ],
+ dim=-1,
+ )
+ )
diff --git a/forks/FlowDock/flowdock/models/components/esdm.py b/forks/FlowDock/flowdock/models/components/esdm.py
new file mode 100644
index 00000000..ea14b006
--- /dev/null
+++ b/forks/FlowDock/flowdock/models/components/esdm.py
@@ -0,0 +1,884 @@
+# Adapted from: https://github.com/zrqiao/NeuralPLexer
+
+import rootutils
+import torch
+from beartype.typing import Any, Dict, Optional, Tuple
+from omegaconf import DictConfig
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock.models.components.embedding import (
+ GaussianFourierEncoding1D,
+ RelativeGeometryEncoding,
+)
+from flowdock.models.components.hetero_graph import make_multi_relation_graph_batcher
+from flowdock.models.components.modules import PointSetAttention
+from flowdock.utils import RankedLogger
+from flowdock.utils.frame_utils import RigidTransform, get_frame_matrix
+from flowdock.utils.model_utils import GELUMLP, AveragePooling, SumPooling, segment_mean
+
+STATE_DICT = Dict[str, Any]
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+class LocalUpdateUsingReferenceRotations(torch.nn.Module):
+ """Update local geometric representations using reference rotations."""
+
+ def __init__(
+ self,
+ fiber_dim: int,
+ extra_feat_dim: int = 0,
+ eps: float = 1e-4,
+ dropout: float = 0.0,
+ hidden_dim: Optional[int] = None,
+ zero_init: bool = False,
+ ):
+ """Initialize the LocalUpdateUsingReferenceRotations module."""
+ super().__init__()
+ self.dim = fiber_dim * 5 + extra_feat_dim
+ self.fiber_dim = fiber_dim
+ self.mlp = GELUMLP(
+ self.dim,
+ fiber_dim * 4,
+ dropout=dropout,
+ zero_init=zero_init,
+ n_hidden_feats=hidden_dim,
+ )
+ self.eps = eps
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ rotation_mats: torch.Tensor,
+ extra_feats=None,
+ ):
+ """Forward pass of the LocalUpdateUsingReferenceRotations module."""
+ # Vector norms are evaluated without applying rigid transform
+ vecx_local = torch.matmul(
+ x[:, 1:].transpose(-2, -1),
+ rotation_mats,
+ )
+ x1_local = torch.cat(
+ [
+ x[:, 0],
+ vecx_local.flatten(-2, -1),
+ x[:, 1:].square().sum(dim=-2).add(self.eps).sqrt(),
+ ],
+ dim=-1,
+ )
+ if extra_feats is not None:
+ x1_local = torch.cat([x1_local, extra_feats], dim=-1)
+ x1_local = self.mlp(x1_local).view(-1, 4, self.fiber_dim)
+ vecx1_out = torch.matmul(
+ rotation_mats,
+ x1_local[:, 1:],
+ )
+ x1_out = torch.cat([x1_local[:, :1], vecx1_out], dim=-2)
+ return x1_out
+
+
+class LocalUpdateUsingChannelWiseGating(torch.nn.Module):
+ """Update local geometric representations using channel-wise gating."""
+
+ def __init__(
+ self,
+ fiber_dim: int,
+ eps: float = 1e-4,
+ dropout: float = 0.0,
+ hidden_dim: Optional[int] = None,
+ zero_init: bool = False,
+ ):
+ """Initialize the LocalUpdateUsingChannelWiseGating module."""
+ super().__init__()
+ self.dim = fiber_dim * 2
+ self.fiber_dim = fiber_dim
+ self.mlp = GELUMLP(
+ self.dim,
+ self.dim,
+ dropout=dropout,
+ n_hidden_feats=hidden_dim,
+ zero_init=zero_init,
+ )
+ self.gate = torch.nn.Sigmoid()
+ self.lin_out = torch.nn.Linear(fiber_dim, fiber_dim, bias=False)
+ if zero_init:
+ self.lin_out.weight.data.fill_(0.0)
+ self.eps = eps
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ ):
+ """Forward pass of the LocalUpdateUsingChannelWiseGating module."""
+ x1 = torch.cat(
+ [
+ x[:, 0],
+ x[:, 1:].square().sum(dim=-2).add(self.eps).sqrt(),
+ ],
+ dim=-1,
+ )
+ x1 = self.mlp(x1)
+ # Gated nonlinear operation on l=1 representations
+ x1_scalar, x1_gatein = torch.split(x1, self.fiber_dim, dim=-1)
+ x1_gate = self.gate(x1_gatein).unsqueeze(-2)
+ vecx1_out = self.lin_out(x[:, 1:]).mul(x1_gate)
+ x1_out = torch.cat([x1_scalar.unsqueeze(-2), vecx1_out], dim=-2)
+ return x1_out
+
+
+class EquivariantTransformerBlock(torch.nn.Module):
+ """Equivariant Transformer Block module."""
+
+ def __init__(
+ self,
+ fiber_dim: int,
+ heads: int = 8,
+ point_dim: int = 4,
+ eps: float = 1e-4,
+ edge_dim: Optional[int] = None,
+ target_frames: bool = False,
+ edge_update: bool = False,
+ dropout: float = 0.0,
+ ):
+ """Initialize the EquivariantTransformerBlock module."""
+ super().__init__()
+ self.attn_conv = PointSetAttention(
+ fiber_dim,
+ heads=heads,
+ point_dim=point_dim,
+ edge_dim=edge_dim,
+ edge_update=edge_update,
+ )
+ self.fiber_dim = fiber_dim
+ self.target_frames = target_frames
+ self.eps = eps
+ self.edge_update = edge_update
+ if target_frames:
+ self.local_update = LocalUpdateUsingReferenceRotations(
+ fiber_dim, eps=eps, dropout=dropout, zero_init=True
+ )
+ else:
+ self.local_update = LocalUpdateUsingChannelWiseGating(
+ fiber_dim, eps=eps, dropout=dropout, zero_init=True
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ edge_index: torch.LongTensor,
+ t: torch.Tensor,
+ R: torch.Tensor = None,
+ x_edge: torch.Tensor = None,
+ ):
+ """Forward pass of the EquivariantTransformerBlock module."""
+ if self.edge_update:
+ xout, edge_out = self.attn_conv(x, x, edge_index, t, t, x_edge=x_edge)
+ x_edge = x_edge + edge_out
+ else:
+ xout = self.attn_conv(x, x, edge_index, t, t, x_edge=x_edge)
+ x = x + xout
+ if self.target_frames:
+ x = self.local_update(x, R) + x
+ else:
+ x = self.local_update(x) + x
+ return x, x_edge
+
+
+class EquivariantStructureDenoisingModule(torch.nn.Module):
+ """Equivariant Structure Denoising Module."""
+
+ def __init__(
+ self,
+ fiber_dim: int,
+ input_dim: int,
+ input_pair_dim: int,
+ hidden_dim: int = 1024,
+ n_stacks: int = 4,
+ n_heads: int = 8,
+ dropout: float = 0.0,
+ ):
+ """Initialize the EquivariantStructureDenoisingModule module."""
+ super().__init__()
+ self.input_dim = input_dim
+ self.input_pair_dim = input_pair_dim
+ self.fiber_dim = fiber_dim
+ self.protatm_padding_dim = 37
+ self.n_blocks = n_stacks
+ self.input_node_projector = torch.nn.Linear(input_dim, fiber_dim, bias=False)
+ self.input_node_vec_projector = torch.nn.Linear(input_dim, fiber_dim * 3, bias=False)
+ self.input_pair_projector = torch.nn.Linear(input_pair_dim, fiber_dim, bias=False)
+ # Inherit the residue representations
+ self.atm_embed = GELUMLP(input_dim + 32, fiber_dim)
+ self.ipa_modules = torch.nn.ModuleList(
+ [
+ EquivariantTransformerBlock(
+ fiber_dim,
+ heads=n_heads,
+ point_dim=fiber_dim // (n_heads * 2),
+ edge_dim=fiber_dim,
+ target_frames=True,
+ edge_update=True,
+ dropout=dropout,
+ )
+ for _ in range(n_stacks)
+ ]
+ )
+ self.res_adapters = torch.nn.ModuleList(
+ [
+ LocalUpdateUsingReferenceRotations(
+ fiber_dim,
+ extra_feat_dim=input_dim,
+ hidden_dim=hidden_dim,
+ dropout=dropout,
+ zero_init=True,
+ )
+ for _ in range(n_stacks)
+ ]
+ )
+ self.protatm_type_encoding = GELUMLP(self.protatm_padding_dim + input_dim, input_pair_dim)
+ self.time_encoding = GaussianFourierEncoding1D(16)
+ self.rel_geom_enc = RelativeGeometryEncoding(15, fiber_dim)
+ self.rel_geom_embed = GELUMLP(fiber_dim, fiber_dim, n_hidden_feats=fiber_dim)
+ # [displacement, scale]
+ self.out_drift_res = torch.nn.ModuleList(
+ [torch.nn.Linear(fiber_dim, 1, bias=False) for _ in range(n_stacks)]
+ )
+ # for i in range(n_stacks):
+ # self.out_drift_res[i].weight.data.fill_(0.0)
+ self.out_scale_res = torch.nn.ModuleList(
+ [GELUMLP(fiber_dim, 1, zero_init=True) for _ in range(n_stacks)]
+ )
+ self.out_drift_atm = torch.nn.ModuleList(
+ [torch.nn.Linear(fiber_dim, 1, bias=False) for _ in range(n_stacks)]
+ )
+ # for i in range(n_stacks):
+ # self.out_drift_atm[i].weight.data.fill_(0.0)
+ self.out_scale_atm = torch.nn.ModuleList(
+ [GELUMLP(fiber_dim, 1, zero_init=True) for _ in range(n_stacks)]
+ )
+
+ # Pre-tabulated edges
+ self.graph_relations = [
+ (
+ "residue_to_residue",
+ "gather_idx_ab_a",
+ "gather_idx_ab_b",
+ "prot_res",
+ "prot_res",
+ ),
+ (
+ "sampled_residue_to_sampled_residue",
+ "gather_idx_AB_a",
+ "gather_idx_AB_b",
+ "prot_res",
+ "prot_res",
+ ),
+ (
+ "prot_atm_to_prot_atm_graph",
+ "protatm_protatm_idx_src",
+ "protatm_protatm_idx_dst",
+ "prot_atm",
+ "prot_atm",
+ ),
+ (
+ "prot_atm_to_prot_atm_knn",
+ "knn_idx_protatm_protatm_src",
+ "knn_idx_protatm_protatm_dst",
+ "prot_atm",
+ "prot_atm",
+ ),
+ (
+ "prot_atm_to_residue",
+ "protatm_res_idx_protatm",
+ "protatm_res_idx_res",
+ "prot_atm",
+ "prot_res",
+ ),
+ (
+ "residue_to_prot_atm",
+ "protatm_res_idx_res",
+ "protatm_res_idx_protatm",
+ "prot_res",
+ "prot_atm",
+ ),
+ (
+ "sampled_lig_triplet_to_lig_atm",
+ "gather_idx_UI_I",
+ "gather_idx_UI_u",
+ "lig_trp",
+ "lig_atm",
+ ),
+ (
+ "lig_atm_to_sampled_lig_triplet",
+ "gather_idx_UI_u",
+ "gather_idx_UI_I",
+ "lig_atm",
+ "lig_trp",
+ ),
+ (
+ "lig_atm_to_lig_atm_graph",
+ "gather_idx_uv_u",
+ "gather_idx_uv_v",
+ "lig_atm",
+ "lig_atm",
+ ),
+ (
+ "sampled_residue_to_sampled_lig_triplet",
+ "gather_idx_AJ_a",
+ "gather_idx_AJ_J",
+ "prot_res",
+ "lig_trp",
+ ),
+ (
+ "sampled_lig_triplet_to_sampled_residue",
+ "gather_idx_AJ_J",
+ "gather_idx_AJ_a",
+ "lig_trp",
+ "prot_res",
+ ),
+ (
+ "residue_to_sampled_lig_triplet",
+ "gather_idx_aJ_a",
+ "gather_idx_aJ_J",
+ "prot_res",
+ "lig_trp",
+ ),
+ (
+ "sampled_lig_triplet_to_residue",
+ "gather_idx_aJ_J",
+ "gather_idx_aJ_a",
+ "lig_trp",
+ "prot_res",
+ ),
+ (
+ "sampled_lig_triplet_to_sampled_lig_triplet",
+ "gather_idx_IJ_I",
+ "gather_idx_IJ_J",
+ "lig_trp",
+ "lig_trp",
+ ),
+ (
+ "prot_atm_to_lig_atm_knn",
+ "knn_idx_protatm_ligatm_src",
+ "knn_idx_protatm_ligatm_dst",
+ "prot_atm",
+ "lig_atm",
+ ),
+ (
+ "lig_atm_to_prot_atm_knn",
+ "knn_idx_ligatm_protatm_src",
+ "knn_idx_ligatm_protatm_dst",
+ "lig_atm",
+ "prot_atm",
+ ),
+ (
+ "lig_atm_to_lig_atm_knn",
+ "knn_idx_ligatm_ligatm_src",
+ "knn_idx_ligatm_ligatm_dst",
+ "lig_atm",
+ "lig_atm",
+ ),
+ ]
+ self.graph_relations_no_ligand = self.graph_relations[:6]
+
+ def init_scalar_vec_rep(self, x, x_v=None, frame=None):
+ """Initialize scalar and vector representations."""
+ if frame is None:
+ # Zero-initialize the vector channels
+ vec_shape = (*x.shape[:-1], 3, x.shape[-1])
+ res = torch.cat([x.unsqueeze(-2), torch.zeros(vec_shape, device=x.device)], dim=-2)
+ else:
+ x_v = x_v.view(*x.shape[:-1], 3, x.shape[-1])
+ x_v_glob = torch.matmul(frame.R, x_v)
+ res = torch.cat([x.unsqueeze(-2), x_v_glob], dim=-2)
+ return res
+
+ def forward(
+ self,
+ batch,
+ frozen_lig=False,
+ frozen_prot=False,
+ **kwargs,
+ ):
+ """Forward pass of the EquivariantStructureDenoisingModule module."""
+ features = batch["features"]
+ indexer = batch["indexer"]
+ metadata = batch["metadata"]
+ metadata["num_structid"]
+ max(metadata["num_a_per_sample"])
+
+ prot_res_rep_in = features["rec_res_attr_decin"]
+ timestep_prot = features["timestep_encoding_prot"]
+ device = features["res_type"].device
+
+ # Protein all-atom representation initialization
+ protatm_padding_mask = features["res_atom_mask"]
+ protatm_atom37_onehot = torch.nn.functional.one_hot(
+ features["protatm_to_atom37_index"], num_classes=self.protatm_padding_dim
+ )
+ protatm_res_pair_encoding = self.protatm_type_encoding(
+ torch.cat(
+ [
+ prot_res_rep_in[indexer["protatm_res_idx_res"]],
+ protatm_atom37_onehot,
+ ],
+ dim=-1,
+ )
+ )
+ # Gathered AA features from individual graphs
+ prot_atm_rep_in = features["prot_atom_attr_projected"]
+ prot_atm_rep_int = self.atm_embed(
+ torch.cat(
+ [
+ prot_atm_rep_in,
+ self.time_encoding(timestep_prot)[indexer["protatm_res_idx_res"]],
+ ],
+ dim=-1,
+ )
+ )
+
+ prot_atm_coords_padded = features["input_protein_coords"]
+ prot_atm_coords_flat = prot_atm_coords_padded[protatm_padding_mask]
+
+ # Embed the rigid body node representations
+ backbone_frames = get_frame_matrix(
+ prot_atm_coords_padded[:, 0],
+ prot_atm_coords_padded[:, 1],
+ prot_atm_coords_padded[:, 2],
+ )
+ prot_res_rep = self.init_scalar_vec_rep(
+ self.input_node_projector(prot_res_rep_in),
+ x_v=self.input_node_vec_projector(prot_res_rep_in),
+ frame=backbone_frames,
+ )
+ prot_atm_rep = self.init_scalar_vec_rep(prot_atm_rep_int)
+ # gather AA features from individual graphs
+ node_reps = {
+ "prot_res": prot_res_rep,
+ "prot_atm": prot_atm_rep,
+ }
+ # Embed pair representations
+ edge_reps = {
+ "residue_to_residue": features["res_res_pair_attr_decin"],
+ "prot_atm_to_prot_atm_graph": features["prot_atom_pair_attr_projected"],
+ "prot_atm_to_prot_atm_knn": features["knn_feat_protatm_protatm"],
+ "prot_atm_to_residue": protatm_res_pair_encoding,
+ "residue_to_prot_atm": protatm_res_pair_encoding,
+ "sampled_residue_to_sampled_residue": features["res_res_grid_attr_flat_decin"],
+ }
+
+ if not batch["misc"]["protein_only"]:
+ max(metadata["num_i_per_sample"])
+ timestep_lig = features["timestep_encoding_lig"]
+ lig_atm_rep_in = features["lig_atom_attr_projected"]
+ lig_frame_rep_in = features["lig_trp_attr_decin"]
+ # Ligand atom embedding. Two timescales
+ lig_atm_rep_int = self.atm_embed(
+ torch.cat(
+ [lig_atm_rep_in, self.time_encoding(timestep_lig)],
+ dim=-1,
+ )
+ )
+ lig_atm_rep = self.init_scalar_vec_rep(lig_atm_rep_int)
+
+ # Prepare ligand atom - sidechain atom indexers
+ # Initialize coordinate features
+ lig_atm_coords = features["input_ligand_coords"].clone()
+
+ lig_frame_atm_idx = (
+ indexer["gather_idx_ijk_i"][indexer["gather_idx_I_ijk"]],
+ indexer["gather_idx_ijk_j"][indexer["gather_idx_I_ijk"]],
+ indexer["gather_idx_ijk_k"][indexer["gather_idx_I_ijk"]],
+ )
+ ligand_trp_frames = get_frame_matrix(
+ lig_atm_coords[lig_frame_atm_idx[0]],
+ lig_atm_coords[lig_frame_atm_idx[1]],
+ lig_atm_coords[lig_frame_atm_idx[2]],
+ )
+ lig_frame_rep = self.init_scalar_vec_rep(
+ self.input_node_projector(lig_frame_rep_in),
+ x_v=self.input_node_vec_projector(lig_frame_rep_in),
+ frame=ligand_trp_frames,
+ )
+ node_reps.update(
+ {
+ "lig_atm": lig_atm_rep,
+ "lig_trp": lig_frame_rep,
+ }
+ )
+ edge_reps.update(
+ {
+ "lig_atm_to_lig_atm_graph": features["lig_atom_pair_attr_projected"],
+ "sampled_lig_triplet_to_lig_atm": features["lig_af_pair_attr_projected"],
+ "lig_atm_to_sampled_lig_triplet": features["lig_af_pair_attr_projected"],
+ "sampled_residue_to_sampled_lig_triplet": features[
+ "res_trp_grid_attr_flat_decin"
+ ],
+ "sampled_lig_triplet_to_sampled_residue": features[
+ "res_trp_grid_attr_flat_decin"
+ ],
+ "residue_to_sampled_lig_triplet": features["res_trp_pair_attr_flat_decin"],
+ "sampled_lig_triplet_to_residue": features["res_trp_pair_attr_flat_decin"],
+ "sampled_lig_triplet_to_sampled_lig_triplet": features[
+ "trp_trp_grid_attr_flat_decin"
+ ],
+ "prot_atm_to_lig_atm_knn": features["knn_feat_protatm_ligatm"],
+ "lig_atm_to_prot_atm_knn": features["knn_feat_ligatm_protatm"],
+ "lig_atm_to_lig_atm_knn": features["knn_feat_ligatm_ligatm"],
+ }
+ )
+
+ # Message passing
+ protatm_res_idx_res = indexer["protatm_res_idx_res"]
+ if batch["misc"]["protein_only"]:
+ graph_relations = self.graph_relations_no_ligand
+ else:
+ graph_relations = self.graph_relations
+ graph_batcher = make_multi_relation_graph_batcher(graph_relations, indexer, metadata)
+ merged_edge_idx = graph_batcher.collate_idx_list(indexer)
+ merged_node_reps = graph_batcher.collate_node_attr(node_reps)
+ merged_edge_reps = graph_batcher.collate_edge_attr(
+ graph_batcher.zero_pad_edge_attr(edge_reps, self.input_pair_dim, device)
+ )
+ merged_edge_reps = self.input_pair_projector(merged_edge_reps)
+ assert merged_edge_idx[0].shape[0] == merged_edge_reps.shape[0]
+ assert merged_edge_idx[1].shape[0] == merged_edge_reps.shape[0]
+
+ dummy_prot_atm_frames = RigidTransform(prot_atm_coords_flat, R=None)
+ if not batch["misc"]["protein_only"]:
+ dummy_lig_atm_frames = RigidTransform(lig_atm_coords, R=None)
+ merged_node_t = graph_batcher.collate_node_attr(
+ {
+ "prot_res": backbone_frames.t,
+ "prot_atm": dummy_prot_atm_frames.t,
+ "lig_atm": dummy_lig_atm_frames.t,
+ "lig_trp": ligand_trp_frames.t,
+ }
+ )
+ merged_node_R = graph_batcher.collate_node_attr(
+ {
+ "prot_res": backbone_frames.R,
+ "prot_atm": dummy_prot_atm_frames.R,
+ "lig_atm": dummy_lig_atm_frames.R,
+ "lig_trp": ligand_trp_frames.R,
+ }
+ )
+ else:
+ merged_node_t = graph_batcher.collate_node_attr(
+ {
+ "prot_res": backbone_frames.t,
+ "prot_atm": dummy_prot_atm_frames.t,
+ }
+ )
+ merged_node_R = graph_batcher.collate_node_attr(
+ {
+ "prot_res": backbone_frames.R,
+ "prot_atm": dummy_prot_atm_frames.R,
+ }
+ )
+ merged_node_frames = RigidTransform(merged_node_t, merged_node_R)
+ merged_edge_reps = merged_edge_reps + (
+ self.rel_geom_embed(
+ self.rel_geom_enc(merged_node_frames, merged_edge_idx) + merged_edge_reps
+ )
+ )
+
+ # No need to reassign embeddings but need to update point coordinates & frames
+ for block_id in range(self.n_blocks):
+ dummy_prot_atm_frames = RigidTransform(prot_atm_coords_flat, R=None)
+ if not batch["misc"]["protein_only"]:
+ dummy_lig_atm_frames = RigidTransform(lig_atm_coords, R=None)
+ merged_node_t = graph_batcher.collate_node_attr(
+ {
+ "prot_res": backbone_frames.t,
+ "prot_atm": dummy_prot_atm_frames.t,
+ "lig_atm": dummy_lig_atm_frames.t,
+ "lig_trp": ligand_trp_frames.t,
+ }
+ )
+ merged_node_R = graph_batcher.collate_node_attr(
+ {
+ "prot_res": backbone_frames.R,
+ "prot_atm": dummy_prot_atm_frames.R,
+ "lig_atm": dummy_lig_atm_frames.R,
+ "lig_trp": ligand_trp_frames.R,
+ }
+ )
+ else:
+ merged_node_t = graph_batcher.collate_node_attr(
+ {
+ "prot_res": backbone_frames.t,
+ "prot_atm": dummy_prot_atm_frames.t,
+ }
+ )
+ merged_node_R = graph_batcher.collate_node_attr(
+ {
+ "prot_res": backbone_frames.R,
+ "prot_atm": dummy_prot_atm_frames.R,
+ }
+ )
+ # PredictDrift iteration
+ merged_node_reps, merged_edge_reps = self.ipa_modules[block_id](
+ merged_node_reps,
+ merged_edge_idx,
+ t=merged_node_t,
+ R=merged_node_R,
+ x_edge=merged_edge_reps,
+ )
+ offloaded_node_reps = graph_batcher.offload_node_attr(merged_node_reps)
+ if "lig_trp" in offloaded_node_reps.keys():
+ lig_frame_rep = offloaded_node_reps["lig_trp"]
+ offloaded_node_reps["lig_trp"] = lig_frame_rep + self.res_adapters[block_id](
+ lig_frame_rep, ligand_trp_frames.R, extra_feats=lig_frame_rep_in
+ )
+ prot_res_rep = offloaded_node_reps["prot_res"]
+ offloaded_node_reps["prot_res"] = prot_res_rep + self.res_adapters[block_id](
+ prot_res_rep, backbone_frames.R, extra_feats=prot_res_rep_in
+ )
+ merged_node_reps = graph_batcher.collate_node_attr(offloaded_node_reps)
+ # Displacement vectors in the global coordinate system
+ if not batch["misc"]["protein_only"]:
+ drift_trp = (
+ self.out_drift_res[block_id](offloaded_node_reps["lig_trp"][:, 1:]).squeeze(-1)
+ * torch.sigmoid(
+ self.out_scale_res[block_id](offloaded_node_reps["lig_trp"][:, 0])
+ )
+ * 10
+ )
+ drift_trp_gathered = segment_mean(
+ drift_trp,
+ indexer["gather_idx_I_molid"],
+ metadata["num_molid"],
+ )[indexer["gather_idx_i_molid"]]
+ drift_atm = self.out_drift_atm[block_id](
+ offloaded_node_reps["lig_atm"][:, 1:]
+ ).squeeze(-1) * torch.sigmoid(
+ self.out_scale_atm[block_id](offloaded_node_reps["lig_atm"][:, 0])
+ )
+ if not frozen_lig:
+ lig_atm_coords = lig_atm_coords + drift_atm + drift_trp_gathered
+ ligand_trp_frames = get_frame_matrix(
+ lig_atm_coords[lig_frame_atm_idx[0]],
+ lig_atm_coords[lig_frame_atm_idx[1]],
+ lig_atm_coords[lig_frame_atm_idx[2]],
+ )
+
+ drift_bb = (
+ self.out_drift_res[block_id](offloaded_node_reps["prot_res"][:, 1:]).squeeze(-1)
+ * torch.sigmoid(
+ self.out_scale_res[block_id](offloaded_node_reps["prot_res"][:, 0])
+ )
+ * 10
+ )
+ drift_bb_gathered = drift_bb[protatm_res_idx_res]
+ drift_prot_atm_int = self.out_drift_atm[block_id](
+ offloaded_node_reps["prot_atm"][:, 1:]
+ ).squeeze(-1) * torch.sigmoid(
+ self.out_scale_atm[block_id](offloaded_node_reps["prot_atm"][:, 0])
+ )
+ if not frozen_prot:
+ prot_atm_coords_flat = (
+ prot_atm_coords_flat + drift_prot_atm_int + drift_bb_gathered
+ )
+
+ prot_atm_coords_padded = torch.zeros_like(features["input_protein_coords"])
+ prot_atm_coords_padded[protatm_padding_mask] = prot_atm_coords_flat
+ backbone_frames = get_frame_matrix(
+ prot_atm_coords_padded[:, 0],
+ prot_atm_coords_padded[:, 1],
+ prot_atm_coords_padded[:, 2],
+ )
+
+ ret = {
+ "final_embedding_prot_atom": offloaded_node_reps["prot_atm"],
+ "final_embedding_prot_res": offloaded_node_reps["prot_res"],
+ "final_coords_prot_atom": prot_atm_coords_flat,
+ "final_coords_prot_atom_padded": prot_atm_coords_padded,
+ }
+ if not batch["misc"]["protein_only"]:
+ ret["final_embedding_lig_atom"] = offloaded_node_reps["lig_atm"]
+ ret["final_coords_lig_atom"] = lig_atm_coords
+ else:
+ ret["final_embedding_lig_atom"] = None
+ ret["final_coords_lig_atom"] = None
+ return ret
+
+
+def resolve_score_head(
+ protein_model_cfg: DictConfig,
+ score_cfg: DictConfig,
+ task_cfg: DictConfig,
+ state_dict: Optional[STATE_DICT] = None,
+) -> torch.nn.Module:
+ """Instantiates an EquivariantStructureDenoisingModule model for protein-ligand complex
+ structure denoising.
+
+ :param protein_model_cfg: Protein model configuration.
+ :param score_cfg: Score configuration.
+ :param task_cfg: Task configuration.
+ :param state_dict: Optional (potentially-pretrained) state dictionary.
+ :return: EquivariantStructureDenoisingModule model.
+ """
+ model = EquivariantStructureDenoisingModule(
+ score_cfg.fiber_dim,
+ input_dim=protein_model_cfg.residue_dim,
+ input_pair_dim=protein_model_cfg.pair_dim,
+ hidden_dim=score_cfg.hidden_dim,
+ n_stacks=score_cfg.n_stacks,
+ n_heads=protein_model_cfg.n_heads,
+ dropout=task_cfg.dropout,
+ )
+ if score_cfg.from_pretrained and state_dict is not None:
+ try:
+ model.load_state_dict(
+ {
+ ".".join(k.split(".")[1:]): v
+ for k, v in state_dict.items()
+ if k.startswith("score_head")
+ }
+ )
+ log.info("Successfully loaded pretrained score weights.")
+ except Exception as e:
+ log.warning(f"Skipping loading of pretrained score weights due to: {e}.")
+ return model
+
+
+def resolve_confidence_head(
+ protein_model_cfg: DictConfig,
+ confidence_cfg: DictConfig,
+ task_cfg: DictConfig,
+ state_dict: Optional[STATE_DICT] = None,
+) -> Tuple[torch.nn.Module, torch.nn.Module]:
+ """Instantiates an EquivariantStructureDenoisingModule model for confidence prediction.
+
+ :param protein_model_cfg: Protein model configuration.
+ :param confidence_cfg: Confidence configuration.
+ :param task_cfg: Task configuration.
+ :param state_dict: Optional (potentially-pretrained) state dictionary.
+ :return: EquivariantStructureDenoisingModule model and plDDT gram head weights.
+ """
+ confidence_head = EquivariantStructureDenoisingModule(
+ confidence_cfg.fiber_dim,
+ input_dim=protein_model_cfg.residue_dim,
+ input_pair_dim=protein_model_cfg.pair_dim,
+ hidden_dim=confidence_cfg.hidden_dim,
+ n_stacks=confidence_cfg.n_stacks,
+ n_heads=protein_model_cfg.n_heads,
+ dropout=task_cfg.dropout,
+ )
+ plddt_gram_head = GELUMLP(
+ protein_model_cfg.pair_dim,
+ 8,
+ n_hidden_feats=protein_model_cfg.pair_dim,
+ zero_init=True,
+ )
+ if confidence_cfg.from_pretrained and state_dict is not None:
+ try:
+ confidence_head.load_state_dict(
+ {
+ ".".join(k.split(".")[1:]): v
+ for k, v in state_dict.items()
+ if k.startswith("confidence_head")
+ }
+ )
+ log.info("Successfully loaded pretrained confidence weights.")
+ except Exception as e:
+ log.warning(f"Skipping loading of pretrained confidence weights due to: {e}.")
+
+ try:
+ plddt_gram_head.load_state_dict(
+ {
+ ".".join(k.split(".")[1:]): v
+ for k, v in state_dict.items()
+ if k.startswith("plddt_gram_head")
+ }
+ )
+ log.info("Successfully loaded pretrained pLDDT gram head weights.")
+ except Exception as e:
+ log.warning(f"Skipping loading of pretrained pLDDT gram head weights due to: {e}.")
+ return confidence_head, plddt_gram_head
+
+
+def resolve_affinity_head(
+ ligand_model_cfg: DictConfig,
+ affinity_cfg: DictConfig,
+ task_cfg: DictConfig,
+ learnable_pooling: bool = True,
+ state_dict: Optional[STATE_DICT] = None,
+) -> Tuple[torch.nn.Module, torch.nn.Module, torch.nn.Module]:
+ """Instantiates an EquivariantStructureDenoisingModule model for affinity prediction.
+
+ :param ligand_model_cfg: Ligand model configuration.
+ :param affinity_cfg: Affinity configuration.
+ :param task_cfg: Task configuration.
+ :param learnable_pooling: Whether to use learnable ligand pooling modules.
+ :param state_dict: Optional (potentially-pretrained) state dictionary.
+ :return: EquivariantStructureDenoisingModule model as well as a ligand pooling module and
+ projection head.
+ """
+ affinity_head = EquivariantStructureDenoisingModule(
+ affinity_cfg.fiber_dim,
+ input_dim=ligand_model_cfg.node_channels,
+ input_pair_dim=ligand_model_cfg.pair_channels,
+ hidden_dim=affinity_cfg.hidden_dim,
+ n_stacks=affinity_cfg.n_stacks,
+ n_heads=ligand_model_cfg.n_heads,
+ dropout=affinity_cfg.dropout if affinity_cfg.get("dropout") else task_cfg.dropout,
+ )
+ if affinity_cfg.ligand_pooling in ["sum", "add", "summation", "addition"]:
+ ligand_pooling = SumPooling(learnable=learnable_pooling, hidden_dim=affinity_cfg.fiber_dim)
+ elif affinity_cfg.ligand_pooling in ["mean", "avg", "average"]:
+ ligand_pooling = AveragePooling(
+ learnable=learnable_pooling, hidden_dim=affinity_cfg.fiber_dim
+ )
+ else:
+ raise NotImplementedError(
+ f"Unsupported ligand pooling method: {affinity_cfg.ligand_pooling}"
+ )
+ affinity_proj_head = GELUMLP(
+ affinity_cfg.fiber_dim,
+ 1,
+ n_hidden_feats=affinity_cfg.fiber_dim,
+ zero_init=True,
+ )
+ if affinity_cfg.from_pretrained and state_dict is not None:
+ try:
+ affinity_head.load_state_dict(
+ {
+ ".".join(k.split(".")[1:]): v
+ for k, v in state_dict.items()
+ if k.startswith("affinity_head")
+ }
+ )
+ log.info("Successfully loaded pretrained affinity head weights.")
+ except Exception as e:
+ log.warning(f"Skipping loading of pretrained affinity head weights due to: {e}.")
+
+ if learnable_pooling:
+ try:
+ ligand_pooling.load_state_dict(
+ {
+ ".".join(k.split(".")[1:]): v
+ for k, v in state_dict.items()
+ if k.startswith("ligand_pooling")
+ }
+ )
+ log.info("Successfully loaded pretrained ligand pooling weights.")
+ except Exception as e:
+ log.warning(f"Skipping loading of pretrained ligand pooling weights due to: {e}.")
+
+ try:
+ affinity_proj_head.load_state_dict(
+ {
+ ".".join(k.split(".")[1:]): v
+ for k, v in state_dict.items()
+ if k.startswith("affinity_proj_head")
+ }
+ )
+ log.info("Successfully loaded pretrained affinity projection head weights.")
+ except Exception as e:
+ log.warning(
+ f"Skipping loading of pretrained affinity projection head weights due to: {e}."
+ )
+ return affinity_head, ligand_pooling, affinity_proj_head
diff --git a/forks/FlowDock/flowdock/models/components/flowdock.py b/forks/FlowDock/flowdock/models/components/flowdock.py
new file mode 100644
index 00000000..3870e528
--- /dev/null
+++ b/forks/FlowDock/flowdock/models/components/flowdock.py
@@ -0,0 +1,2026 @@
+import os
+import random
+from functools import partial
+
+import rootutils
+import torch
+import torch.nn.functional as F
+import tqdm
+from beartype.typing import Any, Callable, Dict, Literal, Optional, Tuple, Union
+from omegaconf import DictConfig
+from pytorch3d.ops import corresponding_points_alignment
+
+from flowdock.utils.inspect_ode_samplers import clamp_tensor
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock.data.components.mol_features import collate_numpy_samples
+from flowdock.data.components.physical import (
+ get_vdw_radii_array,
+ get_vdw_radii_array_uff,
+)
+from flowdock.models.components.cpm import (
+ resolve_pl_contact_stack,
+ resolve_protein_encoder,
+)
+from flowdock.models.components.esdm import (
+ resolve_affinity_head,
+ resolve_confidence_head,
+ resolve_score_head,
+)
+from flowdock.models.components.losses import (
+ compute_fape_from_atom37,
+ compute_lddt_ca,
+ compute_lddt_pli,
+ compute_TMscore_lbound,
+ compute_TMscore_raw,
+)
+from flowdock.models.components.mht_encoder import (
+ resolve_ligand_encoder,
+ resolve_relational_reasoning_module,
+)
+from flowdock.models.components.noise import (
+ sample_complex_harmonic_prior,
+ sample_esmfold_prior,
+ sample_gaussian_prior,
+ sample_ligand_harmonic_prior,
+)
+from flowdock.models.components.transforms import (
+ DefaultPLCoordinateConverter,
+ LatentCoordinateConverter,
+)
+from flowdock.utils import RankedLogger
+from flowdock.utils.data_utils import (
+ erase_holo_coordinates,
+ get_standard_aa_features,
+ prepare_batch,
+)
+from flowdock.utils.frame_utils import apply_similarity_transform
+from flowdock.utils.model_utils import (
+ distogram_to_gaussian_contact_logits,
+ eval_true_contact_maps,
+ inplace_to_device,
+ inplace_to_torch,
+ sample_res_rowmask_from_contacts,
+ sample_reslig_contact_matrix,
+ segment_argmin,
+ segment_mean,
+ segment_sum,
+ topk_edge_mask_from_logits,
+)
+
+MODEL_BATCH = Dict[str, Any]
+NANOMETERS_TO_ANGSTROM = 10.0
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+class FlowDock(torch.nn.Module):
+ """A geometric conditional flow matching model for protein-ligand docking."""
+
+ def __init__(
+ self,
+ cfg: DictConfig,
+ ) -> None:
+ """Initialize a `FlowDock` module.
+
+ :param cfg: A model configuration dictionary.
+ """
+ super().__init__()
+ self.cfg = cfg
+ self.ligand_cfg = cfg.mol_encoder
+ self.protein_cfg = cfg.protein_encoder
+ self.relational_reasoning_cfg = cfg.relational_reasoning
+ self.contact_cfg = cfg.contact_predictor
+ self.score_cfg = cfg.score_head
+ self.confidence_cfg = cfg.confidence
+ self.affinity_cfg = cfg.affinity
+ self.global_cfg = cfg.task
+ self.protatm_padding_dim = self.protein_cfg.atom_padding_dim # := 37
+ self.max_n_edges = self.global_cfg.edge_crop_size
+ self.latent_model = cfg.latent_model
+ self.prior_type = cfg.prior_type
+
+ # VDW radius mapping, in Angstrom
+ self.atnum2vdw = torch.nn.Parameter(
+ torch.tensor(get_vdw_radii_array() / 100.0),
+ requires_grad=False,
+ )
+ self.atnum2vdw_uff = torch.nn.Parameter(
+ torch.tensor(get_vdw_radii_array_uff() / 100.0),
+ requires_grad=False,
+ )
+
+ # graph hyperparameters
+ self.BINDING_SITE_CUTOFF = 6.0
+ self.INTERNAL_TARGET_VARIANCE_SCALE = self.global_cfg.internal_max_sigma
+ self.GLOBAL_TARGET_VARIANCE_SCALE = self.global_cfg.global_max_sigma
+ self.CONTACT_SCALE = 5.0 # fixed hyperparameter
+
+ (
+ standard_aa_template_featset,
+ standard_aa_graph_featset,
+ ) = get_standard_aa_features()
+ self.standard_aa_template_featset = inplace_to_torch(standard_aa_template_featset)
+ self.standard_aa_molgraph_featset = inplace_to_torch(
+ collate_numpy_samples(standard_aa_graph_featset)
+ )
+
+ # load pretrained weights as desired
+ from_pretrained = (
+ self.ligand_cfg.from_pretrained
+ or self.protein_cfg.from_pretrained
+ or self.relational_reasoning_cfg.from_pretrained
+ or self.contact_cfg.from_pretrained
+ or self.score_cfg.from_pretrained
+ or self.confidence_cfg.from_pretrained
+ or self.affinity_cfg.from_pretrained
+ )
+ if from_pretrained:
+ assert cfg.mol_encoder.checkpoint_file is not None and os.path.exists(
+ cfg.mol_encoder.checkpoint_file
+ ), "Pretrained model weights not found."
+ pretrained_state_dict = (
+ torch.load(cfg.mol_encoder.checkpoint_file)["state_dict"] if from_pretrained else None
+ )
+
+ # ligand encoder
+ self.ligand_encoder = resolve_ligand_encoder(
+ self.ligand_cfg, self.global_cfg, state_dict=pretrained_state_dict
+ )
+ self.lig_masking_rate = self.global_cfg.max_masking_rate
+
+ # protein structure encoder
+ self.protein_encoder, res_in_projector = resolve_protein_encoder(
+ self.protein_cfg, self.global_cfg, state_dict=pretrained_state_dict
+ )
+ # protein sequence encoder
+ if self.protein_cfg.use_esm_embedding:
+ # protein sequence language model
+ self.plm_adapter = res_in_projector
+ else:
+ # one-hot amino acid types
+ self.res_in_projector = res_in_projector
+
+ # relational reasoning module
+ (
+ self.molgraph_single_projector,
+ self.molgraph_pair_projector,
+ self.covalent_embed,
+ ) = resolve_relational_reasoning_module(
+ self.protein_cfg,
+ self.ligand_cfg,
+ self.relational_reasoning_cfg,
+ state_dict=pretrained_state_dict,
+ )
+
+ # contact prediction module
+ (
+ self.pl_contact_stack,
+ self.contact_code_embed,
+ self.dist_bins,
+ self.dgram_head,
+ ) = resolve_pl_contact_stack(
+ self.protein_cfg,
+ self.ligand_cfg,
+ self.contact_cfg,
+ self.global_cfg,
+ state_dict=pretrained_state_dict,
+ )
+
+ # structure denoising module
+ self.score_head = resolve_score_head(
+ self.protein_cfg, self.score_cfg, self.global_cfg, state_dict=pretrained_state_dict
+ )
+
+ # confidence prediction module
+ if self.confidence_cfg.enabled:
+ self.confidence_head, self.plddt_gram_head = resolve_confidence_head(
+ self.protein_cfg,
+ self.confidence_cfg,
+ self.global_cfg,
+ state_dict=pretrained_state_dict,
+ )
+
+ # affinity prediction module
+ if self.affinity_cfg.enabled:
+ (
+ self.affinity_head,
+ self.ligand_pooling,
+ self.affinity_proj_head,
+ ) = resolve_affinity_head(
+ self.ligand_cfg,
+ self.affinity_cfg,
+ self.global_cfg,
+ learnable_pooling=True,
+ state_dict=pretrained_state_dict,
+ )
+
+ self.freeze_pretraining_params()
+
+ def freeze_pretraining_params(self):
+ """Freeze pretraining parameters."""
+ if self.global_cfg.freeze_mol_encoder:
+ log.info("Freezing ligand encoder parameters.")
+ self.ligand_encoder.eval()
+ for p in self.ligand_encoder.parameters():
+ p.requires_grad = False
+ if self.global_cfg.freeze_protein_encoder:
+ log.info("Freezing protein encoder parameters.")
+ for module in [
+ (
+ self.plm_adapter
+ if self.protein_cfg.use_esm_embedding
+ else self.res_in_projector
+ ),
+ self.protein_encoder,
+ ]:
+ module.eval()
+ for p in module.parameters():
+ p.requires_grad = False
+ if self.global_cfg.freeze_relational_reasoning:
+ log.info("Freezing relational reasoning module parameters.")
+ for module in [
+ self.molgraph_single_projector,
+ self.molgraph_pair_projector,
+ self.covalent_embed,
+ ]:
+ module.eval()
+ for p in module.parameters():
+ p.requires_grad = False
+ if self.global_cfg.freeze_contact_predictor:
+ log.info("Freezing contact prediction module parameters.")
+ for module in [
+ self.pl_contact_stack,
+ self.contact_code_embed,
+ self.dgram_head,
+ ]:
+ module.eval()
+ for p in module.parameters():
+ p.requires_grad = False
+ if self.global_cfg.freeze_score_head:
+ log.info("Freezing structure denoising module parameters.")
+ self.score_head.eval()
+ for p in self.score_head.parameters():
+ p.requires_grad = False
+ if self.confidence_cfg.enabled and self.global_cfg.freeze_confidence:
+ log.info("Freezing confidence prediction module parameters.")
+ for module in [self.confidence_head, self.plddt_gram_head]:
+ module.eval()
+ for p in module.parameters():
+ p.requires_grad = False
+ if self.affinity_cfg.enabled and self.global_cfg.freeze_affinity:
+ log.info("Freezing affinity prediction module parameters.")
+ self.affinity_head.eval()
+ for p in self.affinity_head.parameters():
+ p.requires_grad = False
+
+ @staticmethod
+ def assign_timestep_encodings(batch: MODEL_BATCH, t_normalized: Union[float, torch.Tensor]):
+ """Assign timestep encodings to the batch.
+
+ :param batch: A batch dictionary.
+ :param t_normalized: The normalized timestep.
+ """
+ # NOTE: `t_normalized` must be in the range `[0, 1]`
+ features = batch["features"]
+ indexer = batch["indexer"]
+ device = features["res_type"].device
+ if not isinstance(t_normalized, torch.Tensor):
+ t_normalized = torch.full(
+ (batch["metadata"]["num_structid"], 1),
+ t_normalized,
+ device=device,
+ )
+ if t_normalized.shape != (batch["metadata"]["num_structid"], 1):
+ assert (
+ t_normalized.numel() == 1
+ ), f"To properly shape-coerce time step tensor of shape {t_normalized.shape}, the input tensor must contain a single value."
+ t_normalized = torch.full(
+ (batch["metadata"]["num_structid"], 1),
+ t_normalized.item(),
+ device=device,
+ )
+ t_prot = t_normalized[indexer["gather_idx_a_structid"]]
+ batch["features"]["timestep_encoding_prot"] = t_prot
+
+ if not batch["misc"]["protein_only"]:
+ batch["features"]["timestep_encoding_lig"] = t_normalized[
+ indexer["gather_idx_i_structid"]
+ ]
+
+ def resolve_latent_converter(self, *args):
+ if self.latent_model == "default":
+ return DefaultPLCoordinateConverter(self.global_cfg, *args)
+ else:
+ raise NotImplementedError
+
+ def forward_interp(
+ self,
+ batch: MODEL_BATCH,
+ x_int_0: torch.Tensor,
+ t: torch.Tensor,
+ latent_converter: LatentCoordinateConverter,
+ umeyama_correction: bool = True,
+ erase_data: bool = False,
+ ) -> torch.Tensor:
+ """Interpolate latent internal coordinates.
+
+ Note that this function adds small amounts of Gaussian noise
+ to the ground-truth protein and ligand coordinates, to discourage
+ the model from overfitting to experimental noise in the training data.
+ Reference: https://www.science.org/doi/10.1126/science.add2187
+
+ :param batch: A batch dictionary.
+ :param x_int_0: Dimension-less (ground-truth) internal coordinates.
+ :param t: The current normalized timestep.
+ :param latent_converter: The latent coordinate converter.
+ :param umeyama_correction: Whether to apply the Umeyama correction.
+ :param erase_data: Whether to erase data.
+ :return: Interpolated latent internal coordinates.
+ """
+ (
+ ca_lat,
+ apo_ca_lat,
+ cother_lat,
+ apo_cother_lat,
+ ca_lat_centroid_coords,
+ apo_ca_lat_centroid_coords,
+ lig_lat,
+ ) = torch.split(
+ x_int_0,
+ [
+ latent_converter._n_res_per_sample,
+ latent_converter._n_res_per_sample,
+ latent_converter._n_cother_per_sample,
+ latent_converter._n_cother_per_sample,
+ latent_converter._n_molid_per_sample,
+ latent_converter._n_molid_per_sample,
+ latent_converter._n_ligha_per_sample,
+ ],
+ dim=1,
+ )
+ x_int_0_ = torch.cat(
+ [
+ ca_lat,
+ cother_lat,
+ lig_lat,
+ ],
+ dim=1,
+ )
+ try:
+ assert self.global_cfg.single_protein_batch, "Only single protein batch is supported."
+ if self.prior_type == "gaussian":
+ noisy_x_int_0, noisy_x_int_1 = sample_gaussian_prior(
+ x_int_0_, latent_converter, sigma=1.0, x0_sigma=1e-4
+ )
+ elif self.prior_type == "harmonic":
+ noisy_x_int_0, noisy_x_int_1 = sample_complex_harmonic_prior(
+ x_int_0_, latent_converter, batch, x0_sigma=1e-4
+ )
+ elif self.prior_type == "esmfold":
+ # NOTE: the following unnormalization step assumes that `self.latent_model == "default"`
+ apo_lig_lat_ = sample_ligand_harmonic_prior(
+ lig_lat, apo_ca_lat * latent_converter.ca_scale, batch
+ )
+ x_int_1_ = torch.cat(
+ [
+ apo_ca_lat, # NOTE: already normalized
+ apo_cother_lat, # NOTE: already normalized
+ apo_lig_lat_ / latent_converter.other_scale,
+ ],
+ dim=1,
+ )
+ noisy_x_int_0, noisy_x_int_1 = sample_esmfold_prior(
+ x_int_0_, x_int_1_, sigma=1e-4, x0_sigma=1e-4
+ )
+ else:
+ raise NotImplementedError(f"Unsupported prior type: {self.prior_type}")
+ except Exception as e:
+ log.error(
+ f"Failed to converge within `{self.prior_type}` noise function of `forward_interp()` due to: {e}."
+ )
+ raise e
+
+ # NOTE: by this point, both `noisy_x_int_0` and `noisy_x_int_1` are normalized
+ if umeyama_correction and not erase_data:
+ try:
+ # align the complex structure based solely the optimal Ca atom alignment
+ # NOTE: we do not perform such alignments during the initial (`t=1`) sampling timestep
+ noisy_x_ca_int_1 = noisy_x_int_1.split(
+ [
+ latent_converter._n_res_per_sample,
+ latent_converter._n_cother_per_sample,
+ latent_converter._n_ligha_per_sample,
+ ],
+ dim=1,
+ )[0]
+ noisy_x_ca_int_0 = noisy_x_int_0.split(
+ [
+ latent_converter._n_res_per_sample,
+ latent_converter._n_cother_per_sample,
+ latent_converter._n_ligha_per_sample,
+ ],
+ dim=1,
+ )[0]
+ similarity_transform = corresponding_points_alignment(
+ X=noisy_x_ca_int_1, Y=noisy_x_ca_int_0, estimate_scale=False
+ )
+ noisy_x_int_1 = apply_similarity_transform(noisy_x_int_1, *similarity_transform)
+ except Exception as e:
+ log.warning(
+ f"Failed optimal noise alignment within `forward_interp()` due to: {e}. Skipping optimal noise alignment..."
+ )
+ raise e
+
+ # interpolate between target and prior distributions
+ noisy_x_int_t = (1 - t) * noisy_x_int_0 + t * noisy_x_int_1
+
+ # recalculate (and renormalize) centroid coordinates
+ (
+ x_int_t_ca_lat,
+ x_int_t_cother_lat,
+ x_int_t_lig_lat,
+ ) = torch.split(
+ noisy_x_int_t,
+ [
+ latent_converter._n_res_per_sample,
+ latent_converter._n_cother_per_sample,
+ latent_converter._n_ligha_per_sample,
+ ],
+ dim=1,
+ )
+ (
+ x_int_1_ca_lat,
+ x_int_1_cother_lat,
+ _,
+ ) = torch.split(
+ noisy_x_int_1,
+ [
+ latent_converter._n_res_per_sample,
+ latent_converter._n_cother_per_sample,
+ latent_converter._n_ligha_per_sample,
+ ],
+ dim=1,
+ )
+ x_int = torch.cat(
+ [
+ x_int_t_ca_lat,
+ x_int_1_ca_lat,
+ x_int_t_cother_lat,
+ x_int_1_cother_lat,
+ ca_lat_centroid_coords,
+ apo_ca_lat_centroid_coords,
+ x_int_t_lig_lat,
+ ],
+ dim=1,
+ )
+ return x_int
+
+ def forward_interp_plcomplex_latinp(
+ self,
+ batch: MODEL_BATCH,
+ t: torch.Tensor,
+ latent_converter: LatentCoordinateConverter,
+ umeyama_correction: bool = True,
+ erase_data: bool = False,
+ ) -> MODEL_BATCH:
+ """Noise-interpolate protein-ligand complex latent internal coordinates.
+
+ :param batch: A batch dictionary.
+ :param t: The current normalized timestep.
+ :param latent_converter: The latent coordinate converter.
+ :param umeyama_correction: Whether to apply the Umeyama correction.
+ :param erase_data: Whether to erase data.
+ :return: Batch dictionary with interpolated latent internal coordinates.
+ """
+ # Dimension-less internal coordinates
+ # [B, N, 3]
+ x_int = latent_converter.to_latent(batch)
+ if erase_data:
+ x_int = erase_holo_coordinates(batch, x_int, latent_converter)
+ x_int_t = self.forward_interp(
+ batch,
+ x_int,
+ t,
+ latent_converter,
+ umeyama_correction=umeyama_correction,
+ erase_data=erase_data,
+ )
+ return latent_converter.assign_to_batch(batch, x_int_t)
+
+ def prepare_protein_patch_indexers(
+ self, batch: MODEL_BATCH, randomize_anchors: bool = False
+ ) -> MODEL_BATCH:
+ """Prepare protein patch indexers for the batch.
+
+ :param batch: A batch dictionary.
+ :param randomize_anchors: Whether to randomize the anchors.
+ :return: Batch dictionary with protein patch indexers.
+ """
+ features = batch["features"]
+ metadata = batch["metadata"]
+ indexer = batch["indexer"]
+ batch_size = metadata["num_structid"]
+ device = features["res_type"].device
+
+ # Prepare indexers
+ # Use max to ensure segmentation faults are 100% invoked
+ # in case there are any bad indices
+ max(metadata["num_a_per_sample"])
+ n_a_per_sample = max(metadata["num_a_per_sample"])
+ assert (
+ n_a_per_sample * batch_size == metadata["num_a"]
+ ), "Invalid (batched) number of residues"
+ n_protein_patches = min(self.protein_cfg.n_patches, n_a_per_sample)
+ batch["metadata"]["n_prot_patches_per_sample"] = n_protein_patches
+
+ # Uniform segmentation
+ res_idx_in_batch = torch.arange(metadata["num_a"], device=device)
+ batch["indexer"]["gather_idx_a_pid"] = (
+ res_idx_in_batch // n_a_per_sample
+ ) * n_protein_patches + (
+ ((res_idx_in_batch % n_a_per_sample) * n_protein_patches) // n_a_per_sample
+ )
+
+ if randomize_anchors:
+ # Random down-sampling, assigning residues to the patch grid
+ # This maps grid row/column idx to sampled residue idx
+ batch["indexer"]["gather_idx_pid_a"] = segment_argmin(
+ batch["features"]["res_type"].new_zeros(n_a_per_sample * batch_size),
+ indexer["gather_idx_a_pid"],
+ n_protein_patches * batch_size,
+ randomize=True,
+ )
+ else:
+ batch["indexer"]["gather_idx_pid_a"] = segment_mean(
+ res_idx_in_batch,
+ indexer["gather_idx_a_pid"],
+ n_protein_patches * batch_size,
+ ).long()
+
+ return batch
+
+ def prepare_protein_backbone_indexers(
+ self, batch: MODEL_BATCH, **kwargs: Dict[str, Any]
+ ) -> MODEL_BATCH:
+ """Prepare protein backbone indexers for the batch.
+
+ :param batch: A batch dictionary.
+ :param kwargs: Additional keyword arguments.
+ :return: Batch dictionary with protein backbone indexers.
+ """
+ features = batch["features"]
+ indexer = batch["indexer"]
+ metadata = batch["metadata"]
+ device = features["res_type"].device
+
+ protatm_coords_padded = features["input_protein_coords"]
+ batch_size = metadata["num_structid"]
+
+ assert self.global_cfg.single_protein_batch, "Only single protein batch is supported."
+ num_res_per_struct = max(metadata["num_a_per_sample"])
+ # Check that the samples are clones of the same complex
+ assert (
+ batch_size * num_res_per_struct == protatm_coords_padded.shape[0]
+ ), "Invalid number of residues."
+
+ input_prot_coords_folded = protatm_coords_padded.unflatten(
+ 0, (batch_size, num_res_per_struct)
+ )
+ single_struct_chain_id = indexer["gather_idx_a_chainid"][:num_res_per_struct]
+ single_struct_res_id = features["residue_index"][:num_res_per_struct]
+ ca_ca_dist = (
+ input_prot_coords_folded[:, :, None, 1] - input_prot_coords_folded[:, None, :, 1]
+ ).norm(dim=-1)
+ ca_ca_knn_mask = topk_edge_mask_from_logits(
+ -ca_ca_dist / self.CONTACT_SCALE,
+ self.protein_cfg.max_residue_degree,
+ randomize=True,
+ )
+ chain_mask = single_struct_chain_id[None, :, None] == single_struct_chain_id[None, None, :]
+ sequence_dist = single_struct_res_id[None, :, None] - single_struct_res_id[None, None, :]
+ sequence_proximity_mask = (torch.abs(sequence_dist) <= 4) & chain_mask
+ prot_res_res_edge_mask = ca_ca_knn_mask | sequence_proximity_mask
+
+ dense_row_idx_3D = (
+ torch.arange(batch_size * num_res_per_struct, device=device)
+ .view(batch_size, num_res_per_struct)[:, :, None]
+ .expand(-1, -1, num_res_per_struct)
+ ).contiguous()
+ dense_col_idx_3D = dense_row_idx_3D.transpose(1, 2).contiguous()
+ batch["metadata"]["num_prot_res"] = metadata["num_a"]
+ batch["indexer"]["gather_idx_ab_a"] = dense_row_idx_3D[prot_res_res_edge_mask]
+ batch["indexer"]["gather_idx_ab_b"] = dense_col_idx_3D[prot_res_res_edge_mask]
+ batch["indexer"]["gather_idx_ab_structid"] = indexer["gather_idx_a_structid"][
+ indexer["gather_idx_ab_a"]
+ ]
+ batch["metadata"]["num_ab"] = batch["indexer"]["gather_idx_ab_a"].shape[0]
+
+ if self.global_cfg.constrained_inpainting:
+ # Diversified spherical cropping scheme
+ assert self.global_cfg.single_protein_batch, "Only single protein batch is supported."
+ batch_size = metadata["num_structid"]
+ # Assert single ligand samples
+ assert batch_size == metadata["num_molid"], "Invalid number of ligands."
+ ligand_coords = batch["features"]["sdf_coordinates"].reshape(batch_size, -1, 3)
+ ligand_centroids = torch.mean(ligand_coords, dim=1)
+ if kwargs["training"]:
+ # 3A perturbations around the ligand centroid
+ perturbed_centroids = ligand_centroids + torch.rand_like(ligand_centroids) * 1.73
+ site_radius = torch.amax(
+ torch.norm(ligand_coords - perturbed_centroids[:, None, :], dim=-1),
+ dim=1,
+ )
+ perturbed_site_radius = (
+ site_radius + (0.5 + torch.rand_like(site_radius)) * self.BINDING_SITE_CUTOFF
+ )
+ else:
+ perturbed_centroids = ligand_centroids
+ site_radius = torch.amax(
+ torch.norm(ligand_coords - perturbed_centroids[:, None, :], dim=-1),
+ dim=1,
+ )
+ perturbed_site_radius = site_radius + self.BINDING_SITE_CUTOFF
+ centroid_ca_dist = (
+ batch["features"]["res_atom_positions"][:, 1].contiguous().view(batch_size, -1, 3)
+ - perturbed_centroids[:, None, :]
+ ).norm(dim=-1)
+ binding_site_mask = (centroid_ca_dist < perturbed_site_radius[:, None]).flatten(0, 1)
+ batch["features"]["binding_site_mask"] = binding_site_mask
+ batch["features"]["template_alignment_mask"] = (~binding_site_mask) & batch[
+ "features"
+ ]["template_alignment_mask"].bool()
+
+ return batch
+
+ def initialize_protein_embeddings(self, batch: MODEL_BATCH):
+ """Initialize protein embeddings.
+
+ :param batch: A batch dictionary.
+ """
+ features = batch["features"]
+
+ # Protein residue and residue-pair embeddings
+ if "res_embedding_in" not in features:
+ if self.protein_cfg.use_esm_embedding:
+ assert (
+ self.global_cfg.single_protein_batch
+ ), "Only single protein batch is supported."
+ features["res_embedding_in"] = self.plm_adapter(
+ batch["features"]["apo_lm_embeddings"]
+ if "apo_lm_embeddings" in batch["features"]
+ else batch["features"]["lm_embeddings"]
+ )
+ assert (
+ features["res_embedding_in"].shape[0] == features["res_atom_types"].shape[0]
+ ), "Invalid number of residues."
+ else:
+ features["res_embedding_in"] = self.res_in_projector(
+ F.one_hot(
+ features["res_type"].long(),
+ num_classes=self.protein_cfg.n_aa_types,
+ ).float()
+ )
+
+ def initialize_protatm_indexer_and_embeddings(self, batch: MODEL_BATCH):
+ """Assign coordinate-independent edges and protein features from PIFormer.
+
+ :param batch: A batch dictionary.
+ """
+ features = batch["features"]
+ device = features["res_type"].device
+ assert self.global_cfg.single_protein_batch, "Only single protein batch is supported."
+ self.standard_aa_molgraph_featset = inplace_to_device(
+ self.standard_aa_molgraph_featset, device
+ )
+ self.standard_aa_template_featset = inplace_to_device(
+ self.standard_aa_template_featset, device
+ )
+ self.standard_aa_molgraph_featset = self.ligand_encoder(self.standard_aa_molgraph_featset)
+ with torch.no_grad():
+ assert (
+ self.standard_aa_molgraph_featset["metadata"]["num_i"] == 167
+ ), "Invalid number of atoms."
+ template_atom_idx_in_batch_padded = torch.full(
+ (20, 37), fill_value=-1, dtype=torch.long, device=device
+ )
+ template_atom37_mask = self.standard_aa_template_featset["features"][
+ "res_atom_mask"
+ ].bool()
+ template_atom_idx_in_batch_padded[template_atom37_mask] = torch.arange(
+ 167, device=device
+ )
+ atom_idx_in_batch_to_restype_idx = (
+ torch.arange(20, device=device)[:, None]
+ .expand(-1, 37)
+ .contiguous()[template_atom37_mask]
+ )
+ atom_idx_in_batch_to_atom37_idx = (
+ torch.arange(37, device=device)[None, :]
+ .expand(20, -1)
+ .contiguous()[template_atom37_mask]
+ )
+ template_padded_edge_mask_per_aa = torch.zeros(
+ (20, 37, 37), dtype=torch.bool, device=device
+ )
+ template_aa_graph_indexer = self.standard_aa_molgraph_featset["indexer"]
+ template_padded_edge_mask_per_aa[
+ atom_idx_in_batch_to_restype_idx[template_aa_graph_indexer["gather_idx_uv_u"]],
+ atom_idx_in_batch_to_atom37_idx[template_aa_graph_indexer["gather_idx_uv_u"]],
+ atom_idx_in_batch_to_atom37_idx[template_aa_graph_indexer["gather_idx_uv_v"]],
+ ] = True
+
+ # Gather adjacency matrix to the input protein
+ features = batch["features"]
+ metadata = batch["metadata"]
+ # Prepare intra-residue protein atom - protein atom indexers
+ n_res_first = max(metadata["num_a_per_sample"])
+ batch["features"]["res_atom_mask"] = features["res_atom_mask"].bool()
+ protatm_padding_mask = batch["features"]["res_atom_mask"][:n_res_first]
+ n_protatm_first = int(protatm_padding_mask.sum())
+ protatm_res_idx_res_first = (
+ torch.arange(n_res_first, device=device)[:, None]
+ .expand(-1, 37)
+ .contiguous()[protatm_padding_mask]
+ )
+ protatm_to_atom37_idx_first = (
+ torch.arange(37, device=device)[None, :]
+ .expand(n_res_first, -1)
+ .contiguous()[protatm_padding_mask]
+ )
+ same_residue_mask = (
+ protatm_res_idx_res_first[:, None] == protatm_res_idx_res_first[None, :]
+ ).contiguous()
+ aa_graph_edge_mask = torch.zeros(
+ (n_protatm_first, n_protatm_first), dtype=torch.bool, device=device
+ )
+ src_idx_sameres = (
+ torch.arange(n_protatm_first, device=device)[:, None]
+ .expand(-1, n_protatm_first)
+ .contiguous()[same_residue_mask]
+ )
+ dst_idx_sameres = (
+ torch.arange(n_protatm_first, device=device)[None, :]
+ .expand(n_protatm_first, -1)
+ .contiguous()[same_residue_mask]
+ )
+ aa_graph_edge_mask[
+ src_idx_sameres, dst_idx_sameres
+ ] = template_padded_edge_mask_per_aa[
+ features["res_type"].long()[protatm_res_idx_res_first[src_idx_sameres]],
+ protatm_to_atom37_idx_first[src_idx_sameres],
+ protatm_to_atom37_idx_first[dst_idx_sameres],
+ ]
+ src_idx_first = (
+ torch.arange(n_protatm_first, device=device)[:, None]
+ .expand(-1, n_protatm_first)
+ .contiguous()[aa_graph_edge_mask]
+ )
+ dst_idx_first = (
+ torch.arange(n_protatm_first, device=device)[None, :]
+ .expand(n_protatm_first, -1)
+ .contiguous()[aa_graph_edge_mask]
+ )
+ batch_size = metadata["num_structid"]
+ src_idx = (
+ (
+ src_idx_first[None, :].expand(batch_size, -1)
+ + torch.arange(batch_size, device=device)[:, None] * n_protatm_first
+ )
+ .contiguous()
+ .flatten()
+ )
+ dst_idx = (
+ (
+ dst_idx_first[None, :].expand(batch_size, -1)
+ + torch.arange(batch_size, device=device)[:, None] * n_protatm_first
+ )
+ .contiguous()
+ .flatten()
+ )
+ batch["metadata"]["num_protatm_per_sample"] = n_protatm_first
+ batch["indexer"]["protatm_protatm_idx_src"] = src_idx
+ batch["indexer"]["protatm_protatm_idx_dst"] = dst_idx
+ batch["metadata"]["num_prot_atm"] = n_protatm_first * batch_size
+ batch["indexer"]["protatm_res_idx_res"] = (
+ (
+ protatm_res_idx_res_first[None, :].expand(batch_size, -1)
+ + torch.arange(batch_size, device=device)[:, None] * n_res_first
+ )
+ .contiguous()
+ .flatten()
+ )
+ batch["indexer"]["protatm_res_idx_protatm"] = torch.arange(
+ batch["metadata"]["num_prot_atm"], device=device
+ )
+ # Gather graph features to the protein feature set
+ template_padded_node_feat_per_aa = torch.zeros(
+ (20, 37, self.protein_cfg.residue_dim), device=device
+ )
+ template_padded_node_feat_per_aa[template_atom37_mask] = self.molgraph_single_projector(
+ self.standard_aa_molgraph_featset["features"]["lig_atom_attr"]
+ )
+ protatm_padding_mask = batch["features"]["res_atom_mask"]
+ protatm_to_atom37_idx = (
+ protatm_to_atom37_idx_first[None, :].expand(batch_size, -1).contiguous().flatten(0, 1)
+ )
+ batch["features"]["protatm_to_atom37_index"] = protatm_to_atom37_idx
+ batch["features"]["protatm_to_atomic_number"] = features["res_atom_types"].long()[
+ protatm_padding_mask
+ ]
+ batch["features"]["prot_atom_attr_projected"] = (
+ template_padded_node_feat_per_aa[
+ features["res_type"].long()[protatm_res_idx_res_first],
+ protatm_to_atom37_idx_first,
+ ][None, :]
+ .expand(batch_size, -1, -1)
+ .contiguous()
+ .flatten(0, 1)
+ )
+ template_padded_edge_feat_per_aa = torch.zeros(
+ (20, 37, 37, self.protein_cfg.pair_dim), device=device
+ )
+ template_padded_edge_feat_per_aa[
+ atom_idx_in_batch_to_restype_idx[template_aa_graph_indexer["gather_idx_uv_u"]],
+ atom_idx_in_batch_to_atom37_idx[template_aa_graph_indexer["gather_idx_uv_u"]],
+ atom_idx_in_batch_to_atom37_idx[template_aa_graph_indexer["gather_idx_uv_v"]],
+ ] = self.molgraph_pair_projector(
+ self.standard_aa_molgraph_featset["features"]["lig_atom_pair_attr"]
+ )
+
+ batch["features"]["prot_atom_pair_attr_projected"] = (
+ template_padded_edge_feat_per_aa[
+ features["res_type"].long()[protatm_res_idx_res_first[src_idx_first]],
+ protatm_to_atom37_idx_first[src_idx_first],
+ protatm_to_atom37_idx_first[dst_idx_first],
+ ][None, :, :]
+ .expand(batch_size, -1, -1)
+ .contiguous()
+ .flatten(0, 1)
+ )
+ return batch
+
+ def initialize_ligand_embeddings(self, batch: MODEL_BATCH, **kwargs: Dict[str, Any]):
+ """Initialize ligand embeddings.
+
+ :param batch: A batch dictionary.
+ :param kwargs: Additional keyword arguments.
+ """
+ metadata = batch["metadata"]
+ batch["features"]
+ indexer = batch["indexer"]
+ batch_size = metadata["num_structid"]
+
+ # Ligand atom, frame and pair embeddings
+ if kwargs["training"]:
+ masking_rate = random.uniform(0, self.lig_masking_rate) # nosec
+ else:
+ masking_rate = 0
+ batch = self.ligand_encoder(batch, masking_rate=masking_rate)
+ batch["features"]["lig_atom_attr_projected"] = self.molgraph_single_projector(
+ batch["features"]["lig_atom_attr"]
+ )
+ # Downsampled ligand frames
+ batch["features"]["lig_trp_attr_projected"] = self.molgraph_single_projector(
+ batch["features"]["lig_trp_attr"]
+ )
+ batch["features"]["lig_atom_pair_attr_projected"] = self.molgraph_pair_projector(
+ batch["features"]["lig_atom_pair_attr"]
+ )
+ lig_af_pair_attr_flat_ = self.molgraph_pair_projector(
+ batch["features"]["lig_af_pair_attr"]
+ )
+ batch["features"]["lig_af_pair_attr_projected"] = lig_af_pair_attr_flat_
+
+ if self.global_cfg.single_protein_batch:
+ lig_af_pair_attr = lig_af_pair_attr_flat_.new_zeros(
+ batch_size,
+ max(metadata["num_U_per_sample"]),
+ max(metadata["num_I_per_sample"]),
+ self.protein_cfg.pair_dim,
+ )
+ n_U_first = max(metadata["num_U_per_sample"])
+ n_I_first = max(metadata["num_I_per_sample"])
+ lig_af_pair_attr[
+ indexer["gather_idx_UI_U"] // n_U_first,
+ indexer["gather_idx_UI_U"] % n_U_first,
+ indexer["gather_idx_UI_I"] % n_I_first,
+ ] = lig_af_pair_attr_flat_
+
+ batch["features"]["lig_af_grid_attr_projected"] = lig_af_pair_attr
+ else:
+ raise NotImplementedError("Only single protein batch is supported.")
+
+ def run_encoder_stack(
+ self,
+ batch: MODEL_BATCH,
+ **kwargs: Dict[str, Any],
+ ) -> MODEL_BATCH:
+ """Run the encoder stack.
+
+ :param lit_module: A LightningModule instance.
+ :param batch: A batch dictionary.
+ :param training: Whether the model is in training mode.
+ :param kwargs: Additional keyword arguments.
+ :return: A batch dictionary.
+ """
+ with torch.no_grad():
+ batch = self.prepare_protein_patch_indexers(
+ batch, randomize_anchors=kwargs["training"]
+ )
+ self.prepare_protein_backbone_indexers(batch, **kwargs)
+ self.initialize_protein_embeddings(batch)
+ self.initialize_protatm_indexer_and_embeddings(batch)
+
+ batch = self.protein_encoder(
+ batch,
+ in_attr_suffix="",
+ out_attr_suffix="_projected",
+ **kwargs,
+ )
+ if batch["misc"]["protein_only"]:
+ return batch
+
+ # NOTE: here, we are assuming a static ligand graph
+ if "lig_atom_attr" not in batch["features"]:
+ self.initialize_ligand_embeddings(batch, **kwargs)
+ return batch
+
+ def run_contact_map_stack(
+ self,
+ batch: MODEL_BATCH,
+ iter_id: Union[int, str],
+ observed_block_contacts: Optional[torch.Tensor] = None,
+ **kwargs: Dict[str, Any],
+ ) -> MODEL_BATCH:
+ """Run the contact map stack.
+
+ :param batch: A batch dictionary.
+ :param iter_id: The current iteration ID.
+ :param observed_block_contacts: Optional observed block contacts.
+ :param kwargs: Additional keyword arguments.
+ :return: A batch dictionary.
+ """
+ features = batch["features"]
+ device = features["res_type"].device
+ if observed_block_contacts is not None:
+ # Merge into 8AA blocks and gather to patches
+ patch8_idx = (
+ torch.arange(
+ observed_block_contacts.shape[1],
+ device=device,
+ )
+ // 8
+ )
+ merged_contacts_reswise = (
+ segment_sum(
+ observed_block_contacts.transpose(0, 1).contiguous(),
+ patch8_idx,
+ max(patch8_idx) + 1,
+ )
+ .bool()[patch8_idx]
+ .transpose(0, 1)
+ .contiguous()
+ )
+ merged_contacts_gathered = (
+ merged_contacts_reswise.flatten(0, 1)[batch["indexer"]["gather_idx_pid_a"]]
+ .contiguous()
+ .view(
+ observed_block_contacts.shape[0],
+ -1,
+ observed_block_contacts.shape[2],
+ )
+ )
+ block_contact_embedding = self.contact_code_embed(merged_contacts_gathered.long())
+ else:
+ block_contact_embedding = None
+ batch = self.pl_contact_stack(
+ batch,
+ in_attr_suffix="_projected",
+ out_attr_suffix=f"_out_{iter_id}",
+ observed_block_contacts=block_contact_embedding,
+ )
+
+ if batch["misc"]["protein_only"]:
+ return batch
+
+ metadata = batch["metadata"]
+ batch_size = metadata["num_structid"]
+ n_a_per_sample = max(metadata["num_a_per_sample"])
+ n_I_per_sample = metadata["n_lig_patches_per_sample"]
+ res_lig_pair_attr = batch["features"][f"res_trp_pair_attr_flat_out_{iter_id}"]
+ raw_dgram_logits = self.dgram_head(res_lig_pair_attr).view(
+ batch_size, n_a_per_sample, n_I_per_sample, 32
+ )
+ batch["outputs"][f"res_lig_distogram_out_{iter_id}"] = F.log_softmax(
+ raw_dgram_logits, dim=-1
+ )
+ return batch
+
+ def infer_geometry_prior(
+ self,
+ batch: MODEL_BATCH,
+ cached_block_contacts: Optional[torch.Tensor] = None,
+ binding_site_mask: Optional[torch.Tensor] = None,
+ logit_clamp_value: Optional[torch.Tensor] = None,
+ **kwargs: Dict[str, Any],
+ ):
+ """Infer a geometry prior.
+
+ :param batch: A batch dictionary.
+ :param cached_block_contacts: Cached block contacts.
+ :param binding_site_mask: Binding site mask.
+ :param logit_clamp_value: Logit clamp value.
+ """
+ # Parse self.task_cfg.block_contact_decoding_scheme
+ assert (
+ self.global_cfg.block_contact_decoding_scheme == "beam"
+ ), "Only beam search is supported."
+ n_lig_frames = max(batch["metadata"]["num_I_per_sample"])
+ # Autoregressive block-contact sampling
+ if cached_block_contacts is None:
+ # Start from the prior distribution
+ sampled_block_contacts = None
+ last_distogram = batch["outputs"]["res_lig_distogram_out_0"]
+ for iter_id in tqdm.tqdm(range(n_lig_frames), desc="Block contact sampling"):
+ last_contact_map = distogram_to_gaussian_contact_logits(
+ last_distogram, self.dist_bins, self.CONTACT_SCALE
+ )
+ sampled_block_contacts = sample_reslig_contact_matrix(
+ batch, last_contact_map, last=sampled_block_contacts
+ ).detach()
+ self.run_contact_map_stack(
+ batch, iter_id, observed_block_contacts=sampled_block_contacts
+ )
+ last_distogram = batch["outputs"][f"res_lig_distogram_out_{iter_id}"]
+ batch["outputs"]["sampled_block_contacts_last"] = sampled_block_contacts
+ # Check that all ligands are assigned to one protein chain segment
+ num_assigned_per_lig = segment_sum(
+ torch.sum(sampled_block_contacts, dim=1).contiguous().flatten(0, 1),
+ batch["indexer"]["gather_idx_I_molid"],
+ batch["metadata"]["num_molid"],
+ )
+ assert torch.all(num_assigned_per_lig >= 1)
+ else:
+ sampled_block_contacts = cached_block_contacts
+
+ # Use the cached contacts and only sample once
+ self.run_contact_map_stack(
+ batch, n_lig_frames, observed_block_contacts=sampled_block_contacts
+ )
+ last_distogram = batch["outputs"][f"res_lig_distogram_out_{n_lig_frames}"]
+ res_lig_contact_logit_pred = distogram_to_gaussian_contact_logits(
+ last_distogram, self.dist_bins, self.CONTACT_SCALE
+ )
+
+ if binding_site_mask is not None:
+ res_lig_contact_logit_pred = res_lig_contact_logit_pred - (
+ ~binding_site_mask[:, :, None] * 1e9
+ )
+ if not kwargs["training"] and logit_clamp_value is not None:
+ res_lig_contact_logit_pred = (
+ res_lig_contact_logit_pred - (res_lig_contact_logit_pred < logit_clamp_value) * 1e9
+ )
+ batch["outputs"]["geometry_prior_L"] = res_lig_contact_logit_pred.flatten()
+
+ def init_randexp_kNN_edges_and_covmask(
+ self, batch: MODEL_BATCH, detect_covalent: bool = False
+ ):
+ """Initialize random expansion kNN edges and covalent mask.
+
+ :param batch: A batch dictionary.
+ :param detect_covalent: Whether to detect covalent bonds.
+ """
+ device = batch["features"]["res_type"].device
+ batch_size = batch["metadata"]["num_structid"]
+ protatm_padding_mask = batch["features"]["res_atom_mask"]
+ prot_atm_coords_padded = batch["features"]["input_protein_coords"]
+ protatm_coords = prot_atm_coords_padded[protatm_padding_mask].contiguous()
+ n_protatm_per_sample = batch["metadata"]["num_protatm_per_sample"]
+ protatm_coords = protatm_coords.view(batch_size, n_protatm_per_sample, 3)
+ if not batch["misc"]["protein_only"]:
+ n_ligatm_per_sample = max(batch["metadata"]["num_i_per_sample"])
+ ligatm_coords = batch["features"]["input_ligand_coords"]
+ ligatm_coords = ligatm_coords.view(batch_size, n_ligatm_per_sample, 3)
+ atm_coords = torch.cat([protatm_coords, ligatm_coords], dim=1)
+ else:
+ atm_coords = protatm_coords
+ distance_mat = torch.norm(atm_coords[:, :, None] - atm_coords[:, None, :], dim=-1)
+ distance_mat[distance_mat == 0] = 1e9
+ knn_edge_mask = topk_edge_mask_from_logits(
+ -distance_mat / self.CONTACT_SCALE,
+ self.cfg.score_head.max_atom_degree,
+ randomize=True,
+ )
+ if (not batch["misc"]["protein_only"]) and detect_covalent:
+ prot_atomic_numbers = batch["features"]["protatm_to_atomic_number"].view(
+ batch_size, n_protatm_per_sample
+ )
+ lig_atomic_numbers = (
+ batch["features"]["atomic_numbers"].long().view(batch_size, n_ligatm_per_sample)
+ )
+ atom_vdw = self.atnum2vdw[torch.cat([prot_atomic_numbers, lig_atomic_numbers], dim=1)]
+ average_vdw = (atom_vdw[:, :, None] + atom_vdw[:, None, :]) / 2
+ intermol_iscov_mask = distance_mat < average_vdw * 1.3
+ intermol_iscov_mask[:, :n_protatm_per_sample, :n_protatm_per_sample] = False
+ gather_idx_i_molid = batch["indexer"]["gather_idx_i_molid"].view(
+ batch_size, n_ligatm_per_sample
+ )
+ lig_samemol_mask = gather_idx_i_molid[:, :, None] == gather_idx_i_molid[:, None, :]
+ intermol_iscov_mask[
+ :, n_protatm_per_sample:, n_protatm_per_sample:
+ ] = intermol_iscov_mask[:, n_protatm_per_sample:, n_protatm_per_sample:] & (
+ ~lig_samemol_mask
+ )
+ knn_edge_mask = knn_edge_mask | intermol_iscov_mask
+ else:
+ intermol_iscov_mask = torch.zeros_like(distance_mat, dtype=torch.bool)
+ p_idx = torch.arange(batch_size * n_protatm_per_sample, device=device).view(
+ batch_size, n_protatm_per_sample
+ )
+ pp_edge_mask = knn_edge_mask[:, :n_protatm_per_sample, :n_protatm_per_sample]
+ batch["indexer"]["knn_idx_protatm_protatm_src"] = (
+ p_idx[:, :, None].expand(-1, -1, n_protatm_per_sample).contiguous()[pp_edge_mask]
+ )
+ batch["indexer"]["knn_idx_protatm_protatm_dst"] = (
+ p_idx[:, None, :].expand(-1, n_protatm_per_sample, -1).contiguous()[pp_edge_mask]
+ )
+ batch["features"]["knn_feat_protatm_protatm"] = self.covalent_embed(
+ intermol_iscov_mask[:, :n_protatm_per_sample, :n_protatm_per_sample][
+ pp_edge_mask
+ ].long()
+ )
+ if not batch["misc"]["protein_only"]:
+ l_idx = torch.arange(batch_size * n_ligatm_per_sample, device=device).view(
+ batch_size, n_ligatm_per_sample
+ )
+ pl_edge_mask = knn_edge_mask[:, :n_protatm_per_sample, n_protatm_per_sample:]
+ batch["indexer"]["knn_idx_protatm_ligatm_src"] = (
+ p_idx[:, :, None].expand(-1, -1, n_ligatm_per_sample).contiguous()[pl_edge_mask]
+ )
+ batch["indexer"]["knn_idx_protatm_ligatm_dst"] = (
+ l_idx[:, None, :].expand(-1, n_protatm_per_sample, -1).contiguous()[pl_edge_mask]
+ )
+ batch["features"]["knn_feat_protatm_ligatm"] = self.covalent_embed(
+ intermol_iscov_mask[:, :n_protatm_per_sample, n_protatm_per_sample:][
+ pl_edge_mask
+ ].long()
+ )
+ lp_edge_mask = knn_edge_mask[:, n_protatm_per_sample:, :n_protatm_per_sample]
+ batch["indexer"]["knn_idx_ligatm_protatm_src"] = (
+ l_idx[:, :, None].expand(-1, -1, n_protatm_per_sample).contiguous()[lp_edge_mask]
+ )
+ batch["indexer"]["knn_idx_ligatm_protatm_dst"] = (
+ p_idx[:, None, :].expand(-1, n_ligatm_per_sample, -1).contiguous()[lp_edge_mask]
+ )
+ batch["features"]["knn_feat_ligatm_protatm"] = self.covalent_embed(
+ intermol_iscov_mask[:, n_protatm_per_sample:, :n_protatm_per_sample][
+ lp_edge_mask
+ ].long()
+ )
+ ll_edge_mask = knn_edge_mask[:, n_protatm_per_sample:, n_protatm_per_sample:]
+ batch["indexer"]["knn_idx_ligatm_ligatm_src"] = (
+ l_idx[:, :, None].expand(-1, -1, n_ligatm_per_sample).contiguous()[ll_edge_mask]
+ )
+ batch["indexer"]["knn_idx_ligatm_ligatm_dst"] = (
+ l_idx[:, None, :].expand(-1, n_ligatm_per_sample, -1).contiguous()[ll_edge_mask]
+ )
+ batch["features"]["knn_feat_ligatm_ligatm"] = self.covalent_embed(
+ intermol_iscov_mask[:, n_protatm_per_sample:, n_protatm_per_sample:][
+ ll_edge_mask
+ ].long()
+ )
+
+ def init_esdm_inputs(
+ self, batch: MODEL_BATCH, embedding_iter_id: Union[int, str]
+ ) -> MODEL_BATCH:
+ """Initialize the inputs for the ESDM.
+
+ :param batch: A batch dictionary.
+ :param embedding_iter_id: The embedding iteration ID.
+ :param kwargs: Additional keyword arguments.
+ :return: A batch dictionary.
+ """
+ with torch.no_grad():
+ self.init_randexp_kNN_edges_and_covmask(
+ batch,
+ detect_covalent=self.global_cfg.detect_covalent,
+ )
+ batch["features"]["rec_res_attr_decin"] = batch["features"][
+ f"rec_res_attr_out_{embedding_iter_id}"
+ ]
+ batch["features"]["res_res_pair_attr_decin"] = batch["features"][
+ f"res_res_pair_attr_out_{embedding_iter_id}"
+ ]
+ batch["features"]["res_res_grid_attr_flat_decin"] = batch["features"][
+ f"res_res_grid_attr_flat_out_{embedding_iter_id}"
+ ]
+ if batch["misc"]["protein_only"]:
+ return batch
+ batch["features"]["lig_trp_attr_decin"] = batch["features"][
+ f"lig_trp_attr_out_{embedding_iter_id}"
+ ]
+ # Use protein-ligand edges from the contact predictor
+ batch["features"]["res_trp_grid_attr_flat_decin"] = batch["features"][
+ f"res_trp_grid_attr_flat_out_{embedding_iter_id}"
+ ]
+ batch["features"]["res_trp_pair_attr_flat_decin"] = batch["features"][
+ f"res_trp_pair_attr_flat_out_{embedding_iter_id}"
+ ]
+ batch["features"]["trp_trp_grid_attr_flat_decin"] = batch["features"][
+ f"trp_trp_grid_attr_flat_out_{embedding_iter_id}"
+ ]
+ return batch
+
+ def run_score_head(
+ self,
+ batch: MODEL_BATCH,
+ embedding_iter_id: Union[int, str],
+ frozen_lig: Optional[bool] = None,
+ frozen_prot: Optional[bool] = None,
+ **kwargs: Dict[str, Any],
+ ) -> MODEL_BATCH:
+ """Run the score head.
+
+ :param batch: A batch dictionary.
+ :param embedding_iter_id: The embedding iteration ID.
+ :param frozen_lig: Whether to freeze the ligand backbone.
+ :param frozen_prot: Whether to freeze the protein backbone.
+ :param kwargs: Additional keyword arguments.
+ :return: A batch dictionary with the score head output.
+ """
+ batch = self.init_esdm_inputs(batch, embedding_iter_id)
+ return self.score_head(
+ batch,
+ frozen_lig=frozen_lig
+ if frozen_lig is not None
+ else self.global_cfg.frozen_ligand_backbone,
+ frozen_prot=frozen_prot
+ if frozen_prot is not None
+ else self.global_cfg.frozen_protein_backbone,
+ **kwargs,
+ )
+
+ def run_confidence_head(self, batch: MODEL_BATCH, **kwargs: Dict[str, Any]) -> MODEL_BATCH:
+ """Run the confidence head.
+
+ :param batch: A batch dictionary.
+ :param kwargs: Additional keyword arguments.
+ :return: A batch dictionary with the confidence head output.
+ """
+ return self.confidence_head(batch, frozen_lig=False, frozen_prot=False, **kwargs)
+
+ def run_affinity_head(self, batch: MODEL_BATCH, **kwargs: Dict[str, Any]) -> MODEL_BATCH:
+ """Run the affinity head.
+
+ :param batch: A batch dictionary.
+ :param kwargs: Additional keyword arguments.
+ :return: Affinity head output.
+ """
+ aff_out = self.affinity_head(batch, frozen_lig=False, frozen_prot=False, **kwargs)
+ aff_pooled = self.ligand_pooling(
+ aff_out["final_embedding_lig_atom"][:, 0],
+ batch["indexer"]["gather_idx_i_molid"],
+ batch["metadata"]["num_molid"],
+ )
+ return self.affinity_proj_head(aff_pooled).squeeze(1)
+
+ def run_auxiliary_estimation(
+ self,
+ batch: MODEL_BATCH,
+ struct: MODEL_BATCH,
+ return_avg_stats: bool = False,
+ **kwargs: Dict[str, Any],
+ ) -> Union[MODEL_BATCH, Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]]:
+ """Run auxiliary estimations.
+
+ :param batch: A batch dictionary.
+ :param struct: A batch dictionary.
+ :param return_avg_stats: Whether to return average statistics.
+ :param kwargs: Additional keyword arguments.
+ :return: A batch dictionary or a tuple of average (optional) statistics.
+ """
+ batch_size = batch["metadata"]["num_structid"]
+ batch["features"]["input_protein_coords"] = struct["receptor_padded"].clone()
+ if struct["ligands"] is not None:
+ batch["features"]["input_ligand_coords"] = struct["ligands"].clone()
+ else:
+ batch["features"]["input_ligand_coords"] = None
+ self.assign_timestep_encodings(batch, 0.0)
+ batch = self.run_encoder_stack(batch, use_template=False, use_plddt=False, **kwargs)
+ self.run_contact_map_stack(batch, iter_id="auxiliary")
+ batch = self.init_esdm_inputs(batch, "auxiliary")
+ if self.affinity_cfg.enabled:
+ batch["outputs"]["affinity_logits"] = self.run_affinity_head(batch)
+ if not self.confidence_cfg.enabled:
+ return batch
+ conf_out = self.run_confidence_head(batch)
+ conf_rep = (
+ conf_out["final_embedding_prot_res"][:, 0]
+ .contiguous()
+ .view(batch_size, -1, self.cfg.confidence.fiber_dim)
+ )
+ if struct["ligands"] is not None:
+ conf_rep_lig = (
+ conf_out["final_embedding_lig_atom"][:, 0]
+ .contiguous()
+ .view(batch_size, -1, self.cfg.confidence.fiber_dim)
+ )
+ conf_rep = torch.cat([conf_rep, conf_rep_lig], dim=1)
+ plddt_logits = F.log_softmax(self.plddt_gram_head(conf_rep), dim=-1)
+ batch["outputs"]["plddt_logits"] = plddt_logits
+ plddt_gram = torch.exp(plddt_logits)
+ batch["outputs"]["plddt"] = torch.cumsum(plddt_gram[:, :, :4], dim=-1).mean(dim=-1)
+
+ if return_avg_stats:
+ plddt_avg = (batch["outputs"]["plddt"].view(batch_size, -1).mean(dim=1)).detach()
+ if struct["ligands"] is not None:
+ plddt_avg_lig = (
+ batch["outputs"]["plddt"]
+ .view(batch_size, -1)[:, batch["metadata"]["num_a_per_sample"][0] :]
+ .mean(dim=1)
+ .detach()
+ )
+ plddt_avg_ligs = segment_mean(
+ batch["outputs"]["plddt"]
+ .view(batch_size, -1)[:, batch["metadata"]["num_a_per_sample"][0] :]
+ .reshape(-1),
+ batch["indexer"]["gather_idx_i_molid"],
+ batch["metadata"]["num_molid"],
+ ).detach()
+ else:
+ plddt_avg_lig = None
+ plddt_avg_ligs = None
+ return plddt_avg, plddt_avg_lig, plddt_avg_ligs
+
+ return batch
+
+ @staticmethod
+ def reverse_interp_ode_step(
+ x0_hat: torch.Tensor, xt: torch.Tensor, t: torch.Tensor, s: torch.Tensor
+ ) -> torch.Tensor:
+ """Reverse process sampling using an Euler ODE solver.
+
+ :param x0_hat: The denoised state.
+ :param xt: The intermediate (noisy) state.
+ :param t: The current timestep.
+ :param s: The next timestep after `t`.
+ :return: The interpolated state.
+ """
+ step_size = t - s
+ return xt + step_size * x0_hat
+
+ @staticmethod
+ def reverse_interp_vdode_step(
+ x0_hat: torch.Tensor,
+ xt: torch.Tensor,
+ t: torch.Tensor,
+ s: torch.Tensor,
+ eta: float = 1.0,
+ ) -> torch.Tensor:
+ """Reverse process sampling using an Euler Variance Diminishing (VD) ODE solver.
+
+ Note that the LHS and RHS time step scaling factors are clamped to the range
+ `[1e-6, 1 - 1e-6]` by default.
+
+ :param x0_hat: The denoised state.
+ :param xt: The intermediate (noisy) state.
+ :param t: The current timestep.
+ :param s: The next timestep after `t`.
+ :param eta: The variance diminishing factor to employ.
+ :return: The interpolated state.
+ """
+ return clamp_tensor(1 - ((s / t) * eta)) * x0_hat + clamp_tensor((s / t) * eta) * xt
+
+ def reverse_interp_plcomplex_latinp(
+ self,
+ batch: MODEL_BATCH,
+ t: torch.Tensor,
+ s: torch.Tensor,
+ latent_converter: LatentCoordinateConverter,
+ score_converter: LatentCoordinateConverter,
+ sampler_step_fn: Callable,
+ umeyama_correction: bool = True,
+ use_template: bool = False,
+ ) -> MODEL_BATCH:
+ """Reverse-interpolate protein-ligand complex latent internal coordinates.
+
+ :param batch: A batch dictionary.
+ :param t: The current timestep.
+ :param s: The next timestep after `t`.
+ :param latent_converter: The latent coordinate converter.
+ :param score_converter: The latent score converter.
+ :param sampler_step_fn: The sampling step function to use.
+ :param umeyama_correction: Whether to apply the Umeyama correction.
+ :param use_template: Whether to use a given protein structure template.
+ :return: A batch dictionary.
+ """
+ batch_size = batch["metadata"]["num_structid"]
+
+ # derive dimension-less internal coordinates
+ x_int_t = latent_converter.to_latent(batch)
+ (
+ x_int_t_ca_lat,
+ x_int_t_apo_ca_lat,
+ x_int_t_cother_lat,
+ x_int_t_apo_cother_lat,
+ x_int_t_ca_lat_centroid_coords,
+ x_int_t_apo_ca_lat_centroid_coords,
+ x_int_t_lig_lat,
+ ) = torch.split(
+ x_int_t,
+ [
+ latent_converter._n_res_per_sample,
+ latent_converter._n_res_per_sample,
+ latent_converter._n_cother_per_sample,
+ latent_converter._n_cother_per_sample,
+ latent_converter._n_molid_per_sample,
+ latent_converter._n_molid_per_sample,
+ latent_converter._n_ligha_per_sample,
+ ],
+ dim=1,
+ )
+ x_int_t_ = torch.cat(
+ [
+ x_int_t_ca_lat,
+ x_int_t_cother_lat,
+ x_int_t_lig_lat,
+ ],
+ dim=1,
+ )
+ self.assign_timestep_encodings(batch, t)
+
+ batch = self.forward(
+ batch,
+ iter_id=score_converter.iter_id,
+ contact_prediction=True,
+ score=True,
+ observed_block_contacts=score_converter.sampled_block_contacts,
+ use_template=use_template,
+ training=False,
+ )
+ if umeyama_correction:
+ _last_pred_ca_trace = (
+ batch["outputs"]["denoised_prediction"]["final_coords_prot_atom_padded"][:, 1]
+ .view(batch_size, -1, 3)
+ .detach()
+ )
+ if score_converter._last_pred_ca_trace is not None:
+ similarity_transform = corresponding_points_alignment(
+ _last_pred_ca_trace,
+ score_converter._last_pred_ca_trace,
+ estimate_scale=False,
+ )
+ _last_pred_ca_trace = apply_similarity_transform(
+ _last_pred_ca_trace, *similarity_transform
+ )
+ protatm_padding_mask = batch["features"]["res_atom_mask"]
+ pred_protatm_coords = (
+ batch["outputs"]["denoised_prediction"]["final_coords_prot_atom_padded"][
+ protatm_padding_mask
+ ]
+ .contiguous()
+ .view(batch_size, -1, 3)
+ )
+ aligned_pred_protatm_coords = (
+ apply_similarity_transform(pred_protatm_coords, *similarity_transform)
+ .contiguous()
+ .flatten(0, 1)
+ )
+ batch["outputs"]["denoised_prediction"]["final_coords_prot_atom_padded"][
+ protatm_padding_mask
+ ] = aligned_pred_protatm_coords
+ batch["outputs"]["denoised_prediction"][
+ "final_coords_prot_atom"
+ ] = aligned_pred_protatm_coords
+ if not batch["misc"]["protein_only"]:
+ pred_ligatm_coords = batch["outputs"]["denoised_prediction"][
+ "final_coords_lig_atom"
+ ].view(batch_size, -1, 3)
+ aligned_pred_ligatm_coords = (
+ apply_similarity_transform(pred_ligatm_coords, *similarity_transform)
+ .contiguous()
+ .flatten(0, 1)
+ )
+ batch["outputs"]["denoised_prediction"][
+ "final_coords_lig_atom"
+ ] = aligned_pred_ligatm_coords
+ score_converter._last_pred_ca_trace = _last_pred_ca_trace
+
+ # interpolate in the latent (reduced) coordinates space
+ x_int_hat_t = score_converter.to_latent(batch)
+ (
+ x_int_hat_t_ca_lat,
+ _,
+ x_int_hat_t_cother_lat,
+ _,
+ _,
+ _,
+ x_int_hat_t_lig_lat,
+ ) = torch.split(
+ x_int_hat_t,
+ [
+ latent_converter._n_res_per_sample,
+ latent_converter._n_res_per_sample,
+ latent_converter._n_cother_per_sample,
+ latent_converter._n_cother_per_sample,
+ latent_converter._n_molid_per_sample,
+ latent_converter._n_molid_per_sample,
+ latent_converter._n_ligha_per_sample,
+ ],
+ dim=1,
+ )
+ x_int_hat_t_ = torch.cat(
+ [
+ x_int_hat_t_ca_lat,
+ x_int_hat_t_cother_lat,
+ x_int_hat_t_lig_lat,
+ ],
+ dim=1,
+ )
+
+ # only interpolate using Ca atom, non-Ca atom, and ligand atom coordinates
+ x_int_tm = sampler_step_fn(x_int_hat_t_, x_int_t_, t, s)
+
+ # reassemble outputs
+ (
+ x_int_tm_ca_lat,
+ x_int_tm_cother_lat,
+ x_int_tm_lig_lat,
+ ) = torch.split(
+ x_int_tm,
+ [
+ latent_converter._n_res_per_sample,
+ latent_converter._n_cother_per_sample,
+ latent_converter._n_ligha_per_sample,
+ ],
+ dim=1,
+ )
+ x_int_tm_ = torch.cat(
+ [
+ x_int_tm_ca_lat,
+ x_int_t_apo_ca_lat,
+ x_int_tm_cother_lat,
+ x_int_t_apo_cother_lat,
+ x_int_t_ca_lat_centroid_coords,
+ x_int_t_apo_ca_lat_centroid_coords,
+ x_int_tm_lig_lat,
+ ],
+ dim=1,
+ )
+ return latent_converter.assign_to_batch(batch, x_int_tm_)
+
+ def sample_pl_complex_structures(
+ self,
+ batch: MODEL_BATCH,
+ num_steps: int = 100,
+ return_summary_stats: int = False,
+ return_all_states: bool = False,
+ sampler: Literal["ODE", "VDODE"] = "VDODE",
+ sampler_eta: float = 1.0,
+ umeyama_correction: bool = True,
+ start_time: float = 1.0,
+ exact_prior: bool = False,
+ eval_input_protein: bool = False,
+ align_to_ground_truth: bool = True,
+ use_template: Optional[bool] = None,
+ **kwargs,
+ ) -> MODEL_BATCH:
+ """Sample protein-ligand complex structures.
+
+ :param batch: A batch dictionary.
+ :param num_steps: The number of steps.
+ :param return_summary_stats: Whether to return summary statistics.
+ :param return_all_states: Whether to return all states along with sampling metrics.
+ :param sampler: The reverse process sampler to use.
+ :param sampler_eta: The variance diminishing factor to employ for the `VDODE` sampler,
+ which offers a trade-off between exploration (1.0) and exploitation (> 1.0).
+ :param umeyama_correction: Apply optimal alignment between the denoised structure and
+ previous step outputs.
+ :param start_time: The start time.
+ :param exact_prior: Whether to use the exact prior.
+ :param eval_input_protein: Whether to evaluate the input protein structure.
+ :param align_to_ground_truth: Whether to align to the ground truth.
+ :param use_template: Whether to use a given protein structure template.
+ :param kwargs: Additional keyword arguments.
+ :return: A batch dictionary.
+ """
+ assert num_steps > 0, "Invalid number of steps."
+ assert 0.0 <= start_time <= 1.0, "Invalid start time."
+ if use_template is None:
+ use_template = self.global_cfg.use_template
+
+ features = batch["features"]
+ indexer = batch["indexer"]
+ metadata = batch["metadata"]
+ res_atom_mask = batch["features"]["res_atom_mask"].bool()
+ device = features["res_type"].device
+ batch_size = metadata["num_structid"]
+
+ if "num_molid" in batch["metadata"].keys() and batch["metadata"]["num_molid"] > 0:
+ batch["misc"]["protein_only"] = False
+ else:
+ batch["misc"]["protein_only"] = True
+
+ forward_lat_converter = self.resolve_latent_converter(
+ [
+ ("features", "res_atom_positions"),
+ ("features", "input_protein_coords"),
+ ],
+ [("features", "sdf_coordinates"), ("features", "input_ligand_coords")],
+ )
+ reverse_lat_converter = self.resolve_latent_converter(
+ [
+ ("features", "input_protein_coords"),
+ ("features", "input_protein_coords"),
+ ],
+ [
+ ("features", "input_ligand_coords"),
+ ("features", "input_ligand_coords"),
+ ],
+ )
+ reverse_score_converter = self.resolve_latent_converter(
+ [
+ (
+ "outputs",
+ "denoised_prediction",
+ "final_coords_prot_atom_padded",
+ ),
+ None,
+ ],
+ [
+ (
+ "outputs",
+ "denoised_prediction",
+ "final_coords_lig_atom",
+ ),
+ None,
+ ],
+ )
+
+ with torch.no_grad():
+ if not batch["misc"]["protein_only"]:
+ # Autoregressive block contact map prior
+ if exact_prior:
+ batch = self.prepare_protein_patch_indexers(batch)
+ _, contact_logit_matrix = eval_true_contact_maps(
+ batch, self.CONTACT_SCALE, **kwargs
+ )
+ else:
+ batch = self.forward_interp_plcomplex_latinp(
+ batch,
+ start_time,
+ forward_lat_converter,
+ erase_data=(start_time >= 1.0),
+ )
+ self.assign_timestep_encodings(batch, start_time)
+ # Sample the categorical contact encodings under the hood
+ batch = self.forward(
+ batch,
+ contact_prediction=True,
+ infer_geometry_prior=True,
+ use_template=use_template,
+ training=False,
+ )
+ # Sample initial ligand coordinates from the geometry prior
+ contact_logit_matrix = batch["outputs"]["geometry_prior_L"]
+
+ sampled_lig_res_anchor_mask = sample_res_rowmask_from_contacts(
+ batch, contact_logit_matrix, self.global_cfg.single_protein_batch
+ )
+ num_cont_to_sample = max(metadata["num_I_per_sample"])
+ sampled_block_contacts = None
+ for _ in range(num_cont_to_sample):
+ sampled_block_contacts = sample_reslig_contact_matrix(
+ batch, contact_logit_matrix, last=sampled_block_contacts
+ )
+ forward_lat_converter.lig_res_anchor_mask = sampled_lig_res_anchor_mask
+ reverse_lat_converter.lig_res_anchor_mask = sampled_lig_res_anchor_mask
+ reverse_score_converter.lig_res_anchor_mask = sampled_lig_res_anchor_mask
+ reverse_score_converter.iter_id = num_cont_to_sample
+ reverse_score_converter.sampled_block_contacts = sampled_block_contacts
+ else:
+ reverse_score_converter.iter_id = 0
+ reverse_score_converter.sampled_block_contacts = None
+
+ if sampler == "ODE":
+ sampler_step_fn = self.reverse_interp_ode_step
+ elif sampler == "VDODE":
+ sampler_step_fn = partial(self.reverse_interp_vdode_step, eta=sampler_eta)
+ else:
+ raise NotImplementedError(f"Reverse process sampler {sampler} not implemented.")
+
+ with torch.no_grad():
+ # NOTE: Here, we assume the predicted contacts are robust to a resampling of geometric noise
+ batch = self.forward_interp_plcomplex_latinp(
+ batch,
+ start_time,
+ forward_lat_converter,
+ erase_data=(start_time >= 1.0),
+ )
+
+ if return_all_states:
+ all_frames = [
+ {
+ "ligands": batch["features"]["input_ligand_coords"].cpu(),
+ "receptor": batch["features"]["input_protein_coords"][res_atom_mask].cpu(),
+ "receptor_padded": batch["features"]["input_protein_coords"].cpu(),
+ }
+ ]
+ if eval_input_protein:
+ protein_fape_input, _ = compute_fape_from_atom37(
+ batch,
+ device,
+ batch["features"]["input_protein_coords"],
+ batch["features"]["res_atom_positions"],
+ )
+ tm_lbound_input = compute_TMscore_lbound(
+ batch,
+ batch["features"]["input_protein_coords"],
+ batch["features"]["res_atom_positions"],
+ )
+ tm_lbound_mirrored_input = compute_TMscore_lbound(
+ batch,
+ -batch["features"]["input_protein_coords"],
+ batch["features"]["res_atom_positions"],
+ )
+ tm_aligned_ca_input = compute_TMscore_raw(
+ batch,
+ batch["features"]["input_protein_coords"][:, 1],
+ batch["features"]["res_atom_positions"][:, 1],
+ )
+ lddt_ca_input = compute_lddt_ca(
+ batch,
+ batch["features"]["input_protein_coords"],
+ batch["features"]["res_atom_positions"],
+ )
+ input_ret = {
+ "FAPE_protein_input": protein_fape_input,
+ "TM_aligned_ca_input": tm_aligned_ca_input,
+ "TM_lbound_input": tm_lbound_input,
+ "TM_lbound_mirrored_input": tm_lbound_mirrored_input,
+ "lDDT-Ca_input": lddt_ca_input,
+ }
+ # NOTE: We follow https://arxiv.org/pdf/2402.04845.pdf for all symbolic conventions
+ schedule = torch.linspace(start_time, 0, num_steps + 1, device=device)
+ for t, s in tqdm.tqdm(
+ zip(schedule[:-1], schedule[1:]), desc=f"Structure generation using {sampler}"
+ ):
+ batch = self.reverse_interp_plcomplex_latinp(
+ batch,
+ t[None, None],
+ s[None, None],
+ reverse_lat_converter,
+ reverse_score_converter,
+ sampler_step_fn,
+ umeyama_correction=umeyama_correction,
+ use_template=use_template,
+ )
+ if return_all_states:
+ # all_frames.append(
+ # {
+ # "ligands": batch["features"]["input_ligand_coords"],
+ # "receptor": batch["features"]["input_protein_coords"][
+ # res_atom_mask
+ # ],
+ # "receptor_padded": batch["features"][
+ # "input_protein_coords"
+ # ],
+ # }
+ # )
+ all_frames.append(
+ {
+ "ligands": batch["outputs"]["denoised_prediction"][
+ "final_coords_lig_atom"
+ ].cpu(),
+ "receptor": batch["outputs"]["denoised_prediction"][
+ "final_coords_prot_atom"
+ ].cpu(),
+ "receptor_padded": batch["outputs"]["denoised_prediction"][
+ "final_coords_prot_atom_padded"
+ ].cpu(),
+ }
+ )
+
+ mean_x1 = batch["outputs"]["denoised_prediction"]["final_coords_lig_atom"]
+ mean_x2_padded = batch["outputs"]["denoised_prediction"][
+ "final_coords_prot_atom_padded"
+ ]
+ protatm_padding_mask = batch["features"]["res_atom_mask"]
+ mean_x2 = mean_x2_padded[protatm_padding_mask]
+ if align_to_ground_truth:
+ similarity_transform = corresponding_points_alignment(
+ mean_x2_padded[:, 1].view(batch_size, -1, 3),
+ batch["features"]["res_atom_positions"][:, 1].view(batch_size, -1, 3),
+ estimate_scale=False,
+ )
+ mean_x2 = (
+ apply_similarity_transform(
+ mean_x2.view(batch_size, -1, 3), *similarity_transform
+ )
+ .contiguous()
+ .flatten(0, 1)
+ )
+ mean_x2_padded[protatm_padding_mask] = mean_x2
+ if mean_x1 is not None:
+ mean_x1 = (
+ apply_similarity_transform(
+ mean_x1.view(batch_size, -1, 3), *similarity_transform
+ )
+ .contiguous()
+ .flatten(0, 1)
+ )
+
+ if return_all_states:
+ all_frames.append(
+ {
+ "ligands": mean_x1.cpu(),
+ "receptor": mean_x2.cpu(),
+ "receptor_padded": mean_x2_padded.cpu(),
+ }
+ )
+ protein_fape, _ = compute_fape_from_atom37(
+ batch,
+ device,
+ mean_x2_padded,
+ batch["features"]["res_atom_positions"],
+ )
+ tm_lbound = compute_TMscore_lbound(
+ batch,
+ mean_x2_padded,
+ batch["features"]["res_atom_positions"],
+ )
+ tm_lbound_mirrored = compute_TMscore_lbound(
+ batch,
+ -mean_x2_padded,
+ batch["features"]["res_atom_positions"],
+ )
+ tm_aligned_ca = compute_TMscore_raw(
+ batch,
+ mean_x2_padded[:, 1],
+ batch["features"]["res_atom_positions"][:, 1],
+ )
+ lddt_ca = compute_lddt_ca(
+ batch,
+ mean_x2_padded,
+ batch["features"]["res_atom_positions"],
+ )
+ ret = {
+ "FAPE_protein": protein_fape,
+ "TM_aligned_ca": tm_aligned_ca,
+ "TM_lbound": tm_lbound,
+ "TM_lbound_mirrored": tm_lbound_mirrored,
+ "lDDT-Ca": lddt_ca,
+ }
+ if eval_input_protein:
+ ret.update(input_ret)
+ if mean_x1 is not None:
+ n_I_per_sample = max(metadata["num_I_per_sample"])
+ lig_frame_atm_idx = torch.stack(
+ [
+ indexer["gather_idx_ijk_i"][indexer["gather_idx_I_ijk"]][:n_I_per_sample],
+ indexer["gather_idx_ijk_j"][indexer["gather_idx_I_ijk"]][:n_I_per_sample],
+ indexer["gather_idx_ijk_k"][indexer["gather_idx_I_ijk"]][:n_I_per_sample],
+ ],
+ dim=0,
+ )
+ _, lig_fape, _ = compute_fape_from_atom37(
+ batch,
+ device,
+ mean_x2_padded,
+ batch["features"]["res_atom_positions"],
+ pred_lig_coords=mean_x1,
+ target_lig_coords=batch["features"]["sdf_coordinates"],
+ lig_frame_atm_idx=lig_frame_atm_idx,
+ split_pl_views=True,
+ )
+ coords_pred_prot = mean_x2_padded[res_atom_mask].view(
+ metadata["num_structid"], -1, 3
+ )
+ coords_ref_prot = batch["features"]["res_atom_positions"][res_atom_mask].view(
+ metadata["num_structid"], -1, 3
+ )
+ coords_pred_lig = mean_x1.view(metadata["num_structid"], -1, 3)
+ coords_ref_lig = batch["features"]["sdf_coordinates"].view(
+ metadata["num_structid"], -1, 3
+ )
+ lig_rmsd = segment_mean(
+ (
+ (coords_pred_lig - coords_pred_prot.mean(dim=1, keepdim=True))
+ - (coords_ref_lig - coords_ref_prot.mean(dim=1, keepdim=True))
+ )
+ .square()
+ .sum(dim=-1)
+ .flatten(0, 1),
+ indexer["gather_idx_i_molid"],
+ metadata["num_molid"],
+ ).sqrt()
+ lig_centroid_distance = (
+ segment_mean(
+ (coords_pred_lig - coords_pred_prot.mean(dim=1, keepdim=True)).flatten(
+ 0, 1
+ ),
+ indexer["gather_idx_i_molid"],
+ metadata["num_molid"],
+ )
+ - segment_mean(
+ (coords_ref_lig - coords_ref_prot.mean(dim=1, keepdim=True)).flatten(0, 1),
+ indexer["gather_idx_i_molid"],
+ metadata["num_molid"],
+ )
+ ).norm(dim=-1)
+ lig_hit_score_1A = (lig_rmsd < 1.0).float()
+ lig_hit_score_2A = (lig_rmsd < 2.0).float()
+ lig_hit_score_4A = (lig_rmsd < 4.0).float()
+ lddt_pli = compute_lddt_pli(
+ batch,
+ mean_x2_padded,
+ batch["features"]["res_atom_positions"],
+ mean_x1,
+ batch["features"]["sdf_coordinates"],
+ )
+ ret.update(
+ {
+ "ligand_RMSD": lig_rmsd,
+ "ligand_centroid_distance": lig_centroid_distance,
+ "lDDT-pli": lddt_pli,
+ "FAPE_ligview": lig_fape,
+ "ligand_hit_score_1A": lig_hit_score_1A,
+ "ligand_hit_score_2A": lig_hit_score_2A,
+ "ligand_hit_score_4A": lig_hit_score_4A,
+ }
+ )
+
+ if return_summary_stats:
+ return ret
+
+ if return_all_states:
+ ret.update({"all_frames": all_frames})
+
+ ret.update(
+ {
+ "ligands": mean_x1,
+ "receptor": mean_x2,
+ "receptor_padded": mean_x2_padded,
+ }
+ )
+ return ret
+
+ def forward(
+ self,
+ batch: MODEL_BATCH,
+ iter_id: Union[int, str] = 0,
+ observed_block_contacts: Optional[torch.Tensor] = None,
+ contact_prediction: bool = True,
+ infer_geometry_prior: bool = False,
+ score: bool = False,
+ use_template: bool = False,
+ **kwargs: Dict[str, Any],
+ ) -> Union[MODEL_BATCH, torch.Tensor]:
+ """Perform a forward pass through the model.
+
+ :param batch: A batch dictionary.
+ :param training: Whether the model is in training mode.
+ :param iter_id: The current iteration ID.
+ :param observed_block_contacts: Observed block contacts.
+ :param contact_prediction: Whether to predict contacts.
+ :param infer_geometry_prior: Whether to predict using a geometry prior.
+ :param score: Whether to predict a denoised complex structure.
+ :param use_template: Whether to use a template protein structure.
+ :param kwargs: Additional keyword arguments.
+ :return: Batch dictionary with outputs or ligand binding affinity.
+ """
+ prepare_batch(batch)
+
+ batch = self.run_encoder_stack(
+ batch,
+ use_template=use_template,
+ use_plddt=self.global_cfg.use_plddt,
+ **kwargs,
+ )
+
+ if contact_prediction:
+ self.run_contact_map_stack(
+ batch,
+ iter_id,
+ observed_block_contacts=observed_block_contacts,
+ **kwargs,
+ )
+
+ if infer_geometry_prior:
+ assert (
+ batch["misc"]["protein_only"] is False
+ ), "Only protein-ligand complexes are supported for a geometry prior."
+ self.infer_geometry_prior(batch, **kwargs)
+
+ if score:
+ batch["outputs"]["denoised_prediction"] = self.run_score_head(
+ batch, embedding_iter_id=iter_id, **kwargs
+ )
+
+ return batch
+
+
+if __name__ == "__main__":
+ _ = FlowDock()
diff --git a/forks/FlowDock/flowdock/models/components/hetero_graph.py b/forks/FlowDock/flowdock/models/components/hetero_graph.py
new file mode 100644
index 00000000..5d5cab15
--- /dev/null
+++ b/forks/FlowDock/flowdock/models/components/hetero_graph.py
@@ -0,0 +1,166 @@
+# Adapted from: https://github.com/zrqiao/NeuralPLexer
+
+from dataclasses import dataclass
+
+import torch
+from beartype.typing import Dict, List, Tuple
+
+
+@dataclass(frozen=True)
+class Relation:
+ edge_type: str
+ edge_rev_name: str
+ edge_frd_name: str
+ src_node_type: str
+ dst_node_type: str
+ num_edges: int
+
+
+class MultiRelationGraphBatcher:
+ """Collate sub-graphs of different node/edge types into a single instance.
+
+ Returned multi-relation edge indices are stored in LongTensor of shape [2, N_edges].
+ """
+
+ def __init__(
+ self,
+ relation_forms: List[Relation],
+ graph_metadata: Dict[str, int],
+ ):
+ """Initialize the batcher."""
+ self._relation_forms = relation_forms
+ self._make_offset_dict(graph_metadata)
+
+ def _make_offset_dict(self, graph_metadata):
+ """Create offset dictionaries for node and edge types."""
+ self._node_chunk_sizes = {}
+ self._edge_chunk_sizes = {}
+ self._offsets_lower = {}
+ self._offsets_upper = {}
+ all_node_types = set()
+ for relation in self._relation_forms:
+ assert (
+ f"num_{relation.src_node_type}" in graph_metadata.keys()
+ ), f"Missing metadata: num_{relation.src_node_type}"
+ assert (
+ f"num_{relation.dst_node_type}" in graph_metadata.keys()
+ ), f"Missing metadata: num_{relation.src_node_type}"
+ all_node_types.add(relation.src_node_type)
+ all_node_types.add(relation.dst_node_type)
+ offset = 0
+ # Fix node type ordering
+ self.all_node_types = list(all_node_types)
+ for node_type in self.all_node_types:
+ self._offsets_lower[node_type] = offset
+ self._node_chunk_sizes[node_type] = graph_metadata[f"num_{node_type}"]
+ new_offset = offset + self._node_chunk_sizes[node_type]
+ self._offsets_upper[node_type] = new_offset
+ offset = new_offset
+
+ def collate_single_relation_graphs(self, indexer, node_attr_dict, edge_attr_dict):
+ """Collate sub-graphs of different node/edge types into a single instance."""
+ return {
+ "node_attr": self.collate_node_attr(node_attr_dict),
+ "edge_attr": self.collate_edge_attr(edge_attr_dict),
+ "edge_index": self.collate_idx_list(indexer),
+ }
+
+ def collate_idx_list(
+ self,
+ indexer: Dict[str, torch.Tensor],
+ ) -> torch.Tensor:
+ """Collate edge indices for all relations."""
+ ret_eidxs_rev, ret_eidxs_frd = [], []
+ for relation in self._relation_forms:
+ assert relation.edge_rev_name in indexer.keys()
+ assert relation.edge_frd_name in indexer.keys()
+ assert indexer[relation.edge_rev_name].dim() == 1
+ assert indexer[relation.edge_frd_name].dim() == 1
+ assert torch.all(
+ indexer[relation.edge_rev_name] < self._node_chunk_sizes[relation.src_node_type]
+ ), f"Node index on edge exceeding boundary: {relation.edge_type}, {self._node_chunk_sizes[relation.src_node_type]}, {self._node_chunk_sizes[relation.dst_node_type]}, {max(indexer[relation.edge_rev_name])}, {max(indexer[relation.edge_frd_name])}"
+ assert torch.all(
+ indexer[relation.edge_frd_name] < self._node_chunk_sizes[relation.dst_node_type]
+ ), f"Node index on edge exceeding boundary: {relation.edge_type}, {self._node_chunk_sizes[relation.src_node_type]}, {self._node_chunk_sizes[relation.dst_node_type]}, {max(indexer[relation.edge_rev_name])}, {max(indexer[relation.edge_frd_name])}"
+ ret_eidxs_rev.append(
+ indexer[relation.edge_rev_name] + self._offsets_lower[relation.src_node_type]
+ )
+ ret_eidxs_frd.append(
+ indexer[relation.edge_frd_name] + self._offsets_lower[relation.dst_node_type]
+ )
+ ret_eidxs_rev = torch.cat(ret_eidxs_rev, dim=0)
+ ret_eidxs_frd = torch.cat(ret_eidxs_frd, dim=0)
+ return torch.stack([ret_eidxs_rev, ret_eidxs_frd], dim=0)
+
+ def collate_node_attr(self, node_attr_dict: Dict[str, torch.Tensor]):
+ """Collate node attributes for all node types."""
+ for node_type in self.all_node_types:
+ assert (
+ node_attr_dict[node_type].shape[0] == self._node_chunk_sizes[node_type]
+ ), f"Node count mismatch: {node_type}, {node_attr_dict[node_type].shape[0]}, {self._node_chunk_sizes[node_type]}"
+ return torch.cat([node_attr_dict[node_type] for node_type in self.all_node_types], dim=0)
+
+ def collate_edge_attr(self, edge_attr_dict: Dict[str, torch.Tensor]):
+ """Collate edge attributes for all relations."""
+ # for relation in self._relation_forms:
+ # print(relation.edge_type, edge_attr_dict[relation.edge_type].shape)
+ return torch.cat(
+ [edge_attr_dict[relation.edge_type] for relation in self._relation_forms],
+ dim=0,
+ )
+
+ def zero_pad_edge_attr(
+ self,
+ edge_attr_dict: Dict[str, torch.Tensor],
+ embedding_dim: int,
+ device: torch.device,
+ ):
+ """Zero pad edge attributes for all relations."""
+ for relation in self._relation_forms:
+ if edge_attr_dict[relation.edge_type] is None:
+ edge_attr_dict[relation.edge_type] = torch.zeros(
+ (relation.num_edges, embedding_dim),
+ device=device,
+ )
+ return edge_attr_dict
+
+ def offload_node_attr(self, cat_node_attr: torch.Tensor):
+ """Offload node attributes for all node types."""
+ node_chunk_sizes = [self._node_chunk_sizes[node_type] for node_type in self.all_node_types]
+ node_attr_split = torch.split(cat_node_attr, node_chunk_sizes)
+ return {
+ self.all_node_types[i]: node_attr_split[i] for i in range(len(self.all_node_types))
+ }
+
+ def offload_edge_attr(self, cat_edge_attr: torch.Tensor):
+ """Offload edge attributes for all relations."""
+ edge_chunk_sizes = [relation.num_edges for relation in self._relation_forms]
+ edge_attr_split = torch.split(cat_edge_attr, edge_chunk_sizes)
+ return {
+ self._relation_forms[i].edge_type: edge_attr_split[i]
+ for i in range(len(self._relation_forms))
+ }
+
+
+def make_multi_relation_graph_batcher(
+ list_of_relations: List[Tuple[str, str, str, str, str]],
+ indexer,
+ metadata,
+):
+ """Make a multi-relation graph batcher."""
+ # Use one instantiation of the indexer to compute chunk sizes
+ relation_forms = [
+ Relation(
+ edge_type=rl_tuple[0],
+ edge_rev_name=rl_tuple[1],
+ edge_frd_name=rl_tuple[2],
+ src_node_type=rl_tuple[3],
+ dst_node_type=rl_tuple[4],
+ num_edges=indexer[rl_tuple[1]].shape[0],
+ )
+ for rl_tuple in list_of_relations
+ ]
+ return MultiRelationGraphBatcher(
+ relation_forms,
+ metadata,
+ )
diff --git a/forks/FlowDock/flowdock/models/components/losses.py b/forks/FlowDock/flowdock/models/components/losses.py
new file mode 100644
index 00000000..d9d1731b
--- /dev/null
+++ b/forks/FlowDock/flowdock/models/components/losses.py
@@ -0,0 +1,1439 @@
+import random
+
+import rootutils
+import torch
+import torch.nn.functional as F
+from beartype.typing import Any, Dict, Literal, Optional, Tuple, Union
+from lightning import LightningModule
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock.utils.frame_utils import cartesian_to_internal, get_frame_matrix
+from flowdock.utils.metric_utils import compute_per_atom_lddt
+from flowdock.utils.model_utils import (
+ distance_to_gaussian_contact_logits,
+ distogram_to_gaussian_contact_logits,
+ eval_true_contact_maps,
+ sample_res_rowmask_from_contacts,
+ sample_reslig_contact_matrix,
+ segment_mean,
+)
+
+MODEL_BATCH = Dict[str, Any]
+MODEL_STAGE = Literal["train", "val", "test", "predict"]
+LOSS_MODES = Literal[
+ "structure_prediction",
+ "auxiliary_estimation",
+ "auxiliary_estimation_without_structure_prediction",
+]
+
+
+def compute_contact_prediction_losses(
+ pred_distograms: torch.Tensor,
+ ref_dist_mat: torch.Tensor,
+ dist_bins: torch.Tensor,
+ contact_scale: float,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Compute the contact prediction losses for a given batch.
+
+ :param pred_distograms: The predicted distograms.
+ :param ref_dist_mat: The reference distance matrix.
+ :param dist_bins: The distance bins.
+ :param contact_scale: The contact scale.
+ :return: The distogram and forward KL losses.
+ """
+ # True onehot distance and distogram loss
+ distance_bin_idx = torch.bucketize(ref_dist_mat, dist_bins[:-1], right=True)
+ distogram_loss = F.cross_entropy(pred_distograms.flatten(0, -2), distance_bin_idx.flatten())
+ # Evaluate contact logits via log(\sum_j p_j \exp(-\alpha*r_j^2))
+ ref_contact_logits = distance_to_gaussian_contact_logits(ref_dist_mat, contact_scale)
+ pred_contact_logits = distogram_to_gaussian_contact_logits(
+ pred_distograms,
+ dist_bins,
+ contact_scale,
+ )
+ forward_kl_loss = F.kl_div(
+ F.log_softmax(
+ pred_contact_logits.flatten(-2, -1),
+ dim=-1,
+ ),
+ F.log_softmax(
+ ref_contact_logits.flatten(-2, -1),
+ dim=-1,
+ ),
+ log_target=True,
+ reduction="batchmean",
+ )
+ return distogram_loss, forward_kl_loss
+
+
+def compute_protein_distogram_loss(
+ batch: MODEL_BATCH,
+ target_coords: torch.Tensor,
+ dist_bins: torch.Tensor,
+ dgram_head: torch.nn.Module,
+ entry: str = "res_res_grid_attr_flat",
+) -> torch.Tensor:
+ """Compute the protein distogram loss for a given batch.
+
+ :param batch: A batch dictionary.
+ :param target_coords: The target coordinates.
+ :param dist_bins: The distance bins.
+ :param dgram_head: The distogram head to use for loss calculation.
+ :param entry: The entry to use.
+ :return: The distogram loss.
+ """
+ n_protein_patches = batch["metadata"]["n_prot_patches_per_sample"]
+ sampled_grid_features = batch["features"][entry]
+ sampled_ca_coords = target_coords[batch["indexer"]["gather_idx_pid_a"]].view(
+ batch["metadata"]["num_structid"], n_protein_patches, 3
+ )
+ sampled_ca_dist = torch.norm(
+ sampled_ca_coords[:, :, None] - sampled_ca_coords[:, None, :], dim=-1
+ )
+ # Using AF2 parameters
+ distance_bin_idx = torch.bucketize(sampled_ca_dist, dist_bins[:-1], right=True)
+ distogram_loss = F.cross_entropy(dgram_head(sampled_grid_features), distance_bin_idx.flatten())
+ return distogram_loss
+
+
+def compute_fape_from_atom37(
+ batch: MODEL_BATCH,
+ device: Union[str, torch.device],
+ pred_prot_coords: torch.Tensor, # [N_res, 37, 3]
+ target_prot_coords: torch.Tensor, # [N_res, 37, 3]
+ pred_lig_coords: Optional[torch.Tensor] = None, # [N_atom, 3]
+ target_lig_coords: Optional[torch.Tensor] = None, # [N_atom, 3]
+ lig_frame_atm_idx: Optional[torch.Tensor] = None, # [3, N_atom]
+ split_pl_views: bool = False,
+ cap_size: int = 8000,
+) -> Union[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
+ """Compute the Frame Aligned Point Error (FAPE) loss from `atom37` coordinates.
+
+ :param batch: A batch dictionary.
+ :param device: The device to use.
+ :param pred_prot_coords: The predicted protein coordinates.
+ :param target_prot_coords: The target protein coordinates.
+ :param pred_lig_coords: The predicted ligand coordinates.
+ :param target_lig_coords: The target ligand coordinates.
+ :param lig_frame_atm_idx: The ligand frame atom indices.
+ :param split_pl_views: Whether to split the protein-ligand views.
+ :param cap_size: The capped size.
+ :return: The FAPE loss.
+ """
+ features = batch["features"]
+ batch_size = batch["metadata"]["num_structid"]
+ with torch.no_grad():
+ atom_mask = (
+ features["res_atom_mask"].bool().view(batch["metadata"]["num_structid"], -1, 37)
+ ).clone()
+ atom_mask[:, :, [6, 7, 12, 13, 16, 17, 20, 21, 26, 27, 29, 30]] = False
+ pred_prot_coords = pred_prot_coords.view(batch_size, -1, 37, 3)
+ target_prot_coords = target_prot_coords.view(batch_size, -1, 37, 3)
+ pred_bb_frames = get_frame_matrix(
+ pred_prot_coords[:, :, 0, :],
+ pred_prot_coords[:, :, 1, :],
+ pred_prot_coords[:, :, 2, :],
+ )
+ # pred_bb_frames.R = pred_bb_frames.R.detach()
+ target_bb_frames = get_frame_matrix(
+ target_prot_coords[:, :, 0, :],
+ target_prot_coords[:, :, 1, :],
+ target_prot_coords[:, :, 2, :],
+ )
+ pred_prot_coords_flat = pred_prot_coords[atom_mask].view(batch_size, -1, 3)
+ target_prot_coords_flat = target_prot_coords[atom_mask].view(batch_size, -1, 3)
+ if pred_lig_coords is not None:
+ assert target_prot_coords is not None, "Target protein coordinates must be provided."
+ assert lig_frame_atm_idx is not None, "Ligand frame atom indices must be provided."
+ pred_lig_coords = pred_lig_coords.view(batch_size, -1, 3)
+ target_lig_coords = target_lig_coords.view(batch_size, -1, 3)
+ pred_coords = torch.cat([pred_prot_coords_flat, pred_lig_coords], dim=1)
+ target_coords = torch.cat([target_prot_coords_flat, target_lig_coords], dim=1)
+ pred_lig_frames = get_frame_matrix(
+ pred_lig_coords[:, lig_frame_atm_idx[0]],
+ pred_lig_coords[:, lig_frame_atm_idx[1]],
+ pred_lig_coords[:, lig_frame_atm_idx[2]],
+ )
+ pred_frames = pred_bb_frames.concatenate(pred_lig_frames, dim=1)
+ target_lig_frames = get_frame_matrix(
+ target_lig_coords[:, lig_frame_atm_idx[0]],
+ target_lig_coords[:, lig_frame_atm_idx[1]],
+ target_lig_coords[:, lig_frame_atm_idx[2]],
+ )
+ target_frames = target_bb_frames.concatenate(target_lig_frames, dim=1)
+ else:
+ pred_coords = pred_prot_coords_flat
+ target_coords = target_prot_coords_flat
+ pred_frames = pred_bb_frames
+ target_frames = target_bb_frames
+ # Columns-frames, rows-points
+ # [B, 1, N, 3] - [B, F, 1, 3]
+ sampling_rate = cap_size / (batch_size * target_coords.shape[1])
+ sampling_mask = torch.rand(target_coords.shape[1], device=device) < sampling_rate
+ aligned_pred_points = cartesian_to_internal(
+ pred_coords[:, sampling_mask].unsqueeze(1), pred_frames.unsqueeze(2)
+ )
+ with torch.no_grad():
+ aligned_target_points = cartesian_to_internal(
+ target_coords[:, sampling_mask].unsqueeze(1), target_frames.unsqueeze(2)
+ )
+ pair_dist_aligned = (
+ torch.square(aligned_pred_points - aligned_target_points)
+ .sum(-1)
+ .add(1e-4)
+ .sqrt()
+ .sub(1e-2)
+ )
+ cropped_pair_dists = torch.clamp(pair_dist_aligned, max=10)
+ normalized_pair_dists = (
+ pair_dist_aligned / aligned_target_points.square().sum(-1).add(1e-4).sqrt()
+ )
+ if split_pl_views:
+ fape_protframe = cropped_pair_dists[:, : target_bb_frames.t.shape[1]].mean((1, 2)) / 10
+ fape_ligframe = cropped_pair_dists[:, target_bb_frames.t.shape[1] :].mean((1, 2)) / 10
+ return fape_protframe, fape_ligframe, normalized_pair_dists.mean((1, 2))
+ return cropped_pair_dists.mean((1, 2)) / 10, normalized_pair_dists.mean((1, 2))
+
+
+def compute_aa_distance_geometry_loss(
+ batch: MODEL_BATCH, pred_coords: torch.Tensor, target_coords: torch.Tensor
+) -> torch.Tensor:
+ """Compute the amino acid distance geometry loss for a given batch.
+
+ :param batch: A batch dictionary.
+ :param pred_coords: The predicted coordinates.
+ :param target_coords: The target coordinates.
+ :return: The distance geometry loss.
+ """
+ batch_size = batch["metadata"]["num_structid"]
+ features = batch["features"]
+ atom_mask = features["res_atom_mask"].bool()
+ # Add backbone atoms from previous residue
+ atom_mask = atom_mask.view(batch_size, -1, 37)
+ atom_mask = torch.cat([atom_mask[:, 1:], atom_mask[:, :-1, 0:3]], dim=2).flatten(0, 1)
+ pred_coords = pred_coords.view(batch_size, -1, 37, 3)
+ pred_coords = torch.cat([pred_coords[:, 1:], pred_coords[:, :-1, 0:3]], dim=2).flatten(0, 1)
+ target_coords = target_coords.view(batch_size, -1, 37, 3)
+ target_coords = torch.cat([target_coords[:, 1:], target_coords[:, :-1, 0:3]], dim=2).flatten(
+ 0, 1
+ )
+ local_pair_dist_target = (
+ (target_coords[:, None, :] - target_coords[:, :, None]).square().sum(-1).add(1e-4).sqrt()
+ )
+ local_pair_dist_pred = (
+ (pred_coords[:, None, :] - pred_coords[:, :, None]).square().sum(-1).add(1e-4).sqrt()
+ )
+ local_pair_mask = (
+ atom_mask[:, None, :] & atom_mask[:, :, None] & (local_pair_dist_target < 3.0)
+ )
+ ret = (local_pair_dist_target - local_pair_dist_pred).abs()[local_pair_mask]
+ return ret.view(batch["metadata"]["num_structid"], -1).mean(dim=1)
+
+
+def compute_sm_distance_geometry_loss(
+ batch: MODEL_BATCH, pred_coords: torch.Tensor, target_coords: torch.Tensor
+) -> torch.Tensor:
+ """Compute the small molecule distance geometry loss for a given batch.
+
+ :param batch: A batch dictionary.
+ :param pred_coords: The predicted coordinates.
+ :param target_coords: The target coordinates.
+ :return: The distance geometry loss.
+ """
+ batch_size = batch["metadata"]["num_structid"]
+ pred_coords = pred_coords.view(batch_size, -1, 3)
+ target_coords = target_coords.view(batch_size, -1, 3)
+ pair_dist_target = (
+ (target_coords[:, None, :] - target_coords[:, :, None]).square().sum(-1).add(1e-4).sqrt()
+ )
+ pair_dist_pred = (
+ (pred_coords[:, None, :] - pred_coords[:, :, None]).square().sum(-1).add(1e-4).sqrt()
+ )
+ local_pair_mask = pair_dist_target < 3.0
+ ret = (pair_dist_target - pair_dist_pred).abs()[local_pair_mask]
+ return ret.view(batch_size, -1).mean(dim=1)
+
+
+def compute_drmsd_and_clashloss(
+ batch: MODEL_BATCH,
+ device: Union[str, torch.device],
+ pred_prot_coords: torch.Tensor,
+ target_prot_coords: torch.Tensor,
+ atnum2vdw_uff: torch.nn.Parameter,
+ cap_size: int = 4000,
+ pred_lig_coords: Optional[torch.Tensor] = None,
+ target_lig_coords: Optional[torch.Tensor] = None,
+ ligatm_types: Optional[torch.Tensor] = None,
+ binding_site: bool = False,
+ pl_interface: bool = False,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ """Compute the differentiable root-mean-square deviation (dRMSD) and optional clash loss for a
+ given batch.
+
+ :param batch: A batch dictionary.
+ :param device: The device to use.
+ :param pred_prot_coords: The predicted protein coordinates.
+ :param target_prot_coords: The target protein coordinates.
+ :param atnum2vdw_uff: The atomic number to UFF VDW parameters mapping `Parameter`.
+ :param cap_size: The capped size.
+ :param pred_lig_coords: The predicted ligand coordinates.
+ :param target_lig_coords: The target ligand coordinates.
+ :param ligatm_types: The ligand atom types.
+ :param binding_site: Whether to compute the binding site.
+ :param pl_interface: Whether to compute the protein-ligand interface.
+ :return: The dRMSD and optional clash loss.
+ """
+ features = batch["features"]
+ with torch.no_grad():
+ if not binding_site:
+ atom_mask = features["res_atom_mask"].bool().clone()
+ else:
+ atom_mask = (
+ features["res_atom_mask"].bool() & features["binding_site_mask_clean"][:, None]
+ )
+ if pl_interface:
+ # Removing ambiguous atoms
+ atom_mask[:, [6, 7, 12, 13, 16, 17, 20, 21, 26, 27, 29, 30]] = False
+
+ batch_size = batch["metadata"]["num_structid"]
+ pred_prot_coords = pred_prot_coords[atom_mask].view(batch_size, -1, 3)
+ if pred_lig_coords is not None:
+ assert target_prot_coords is not None, "Target protein coordinates must be provided."
+ assert ligatm_types is not None, "Ligand atom types must be provided."
+ pred_lig_coords = pred_lig_coords.view(batch_size, -1, 3)
+ pred_coords = torch.cat([pred_prot_coords, pred_lig_coords], dim=1)
+ else:
+ pred_coords = pred_prot_coords
+ sampling_rate = cap_size / pred_coords.shape[1]
+ sampling_mask = torch.rand(pred_coords.shape[1], device=device) < sampling_rate
+ pred_coords = pred_coords[:, sampling_mask]
+ if pl_interface:
+ pred_dist = (
+ torch.square(pred_coords[:, :, None] - pred_lig_coords[:, None, :])
+ .sum(-1)
+ .add(1e-4)
+ .sqrt()
+ .sub(1e-2)
+ )
+ else:
+ pred_dist = (
+ torch.square(pred_coords[:, :, None] - pred_coords[:, None, :])
+ .sum(-1)
+ .add(1e-4)
+ .sqrt()
+ .sub(1e-2)
+ )
+ with torch.no_grad():
+ target_prot_coords = target_prot_coords[atom_mask].view(batch_size, -1, 3)
+ if pred_lig_coords is not None:
+ target_lig_coords = target_lig_coords.view(batch_size, -1, 3)
+ target_coords = torch.cat([target_prot_coords, target_lig_coords], dim=1)
+ else:
+ target_coords = target_prot_coords
+ target_coords = target_coords[:, sampling_mask]
+ if pl_interface:
+ target_dist = (
+ torch.square(target_coords[:, :, None] - target_lig_coords[:, None, :])
+ .sum(-1)
+ .add(1e-4)
+ .sqrt()
+ .sub(1e-2)
+ )
+ else:
+ target_dist = (
+ torch.square(target_coords[:, :, None] - target_coords[:, None, :])
+ .sum(-1)
+ .add(1e-4)
+ .sqrt()
+ .sub(1e-2)
+ )
+ # In Angstrom, using UFF params to compute clash loss
+ protatm_types = features["res_atom_types"].long()[atom_mask]
+ protatm_vdw = atnum2vdw_uff[protatm_types].view(batch_size, -1)
+ if pred_lig_coords is not None:
+ ligatm_vdw = atnum2vdw_uff[ligatm_types].view(batch_size, -1)
+ atm_vdw = torch.cat([protatm_vdw, ligatm_vdw], dim=1)
+ else:
+ atm_vdw = protatm_vdw
+ atm_vdw = atm_vdw[:, sampling_mask]
+ average_vdw = (atm_vdw[:, :, None] + atm_vdw[:, None, :]) / 2
+ # Use conservative cutoffs to avoid mis-penalization
+
+ dist_errors = (pred_dist - target_dist).square()
+ drmsd = dist_errors.add(1e-2).sqrt().sub(1e-1).mean(dim=(1, 2))
+ if pl_interface:
+ return drmsd, None
+
+ covalent_like = target_dist < (average_vdw * 1.2)
+ # Alphafold supplementary Eq. 46, modified
+ clash_pairwise = torch.clamp(average_vdw * 1.1 - pred_dist.add(1e-6), min=0.0)
+ clash_loss = clash_pairwise.mul(~covalent_like).sum(dim=2).mean(dim=1)
+ return drmsd, clash_loss
+
+
+def compute_template_weighted_centroid_drmsd(
+ batch: MODEL_BATCH,
+ pred_prot_coords: torch.Tensor,
+) -> torch.Tensor:
+ """Compute the template-weighted centroid dRMSD for a given batch.
+
+ :param batch: A batch dictionary.
+ :param pred_prot_coords: The predicted protein coordinates.
+ :return: The dRMSD.
+ """
+ batch_size = batch["metadata"]["num_structid"]
+
+ pred_cent_coords = (
+ pred_prot_coords.mul(batch["features"]["res_atom_mask"].bool()[:, :, None])
+ .sum(dim=1)
+ .div(batch["features"]["res_atom_mask"].bool().sum(dim=1)[:, None] + 1e-9)
+ ).view(batch_size, -1, 3)
+ pred_dist = (
+ torch.square(pred_cent_coords[:, :, None] - pred_cent_coords[:, None, :])
+ .sum(-1)
+ .add(1e-4)
+ .sqrt()
+ .sub(1e-2)
+ )
+ with torch.no_grad():
+ target_cent_coords = (
+ batch["features"]["res_atom_positions"]
+ .mul(batch["features"]["res_atom_mask"].bool()[:, :, None])
+ .sum(dim=1)
+ .div(batch["features"]["res_atom_mask"].bool().sum(dim=1)[:, None] + 1e-9)
+ ).view(batch_size, -1, 3)
+ template_cent_coords = (
+ batch["features"]["apo_res_atom_positions"]
+ .mul(batch["features"]["apo_res_atom_mask"].bool()[:, :, None])
+ .sum(dim=1)
+ .div(batch["features"]["apo_res_atom_mask"].bool().sum(dim=1)[:, None] + 1e-9)
+ ).view(batch_size, -1, 3)
+ target_dist = (
+ torch.square(target_cent_coords[:, :, None] - target_cent_coords[:, None, :])
+ .sum(-1)
+ .add(1e-4)
+ .sqrt()
+ .sub(1e-2)
+ )
+ template_dist = (
+ torch.square(template_cent_coords[:, :, None] - template_cent_coords[:, None, :])
+ .sum(-1)
+ .add(1e-4)
+ .sqrt()
+ .sub(1e-2)
+ )
+ template_alignment_mask = (
+ batch["features"]["apo_res_alignment_mask"].bool().view(batch_size, -1)
+ )
+ motion_mask = (
+ ((target_dist - template_dist).abs() > 2.0)
+ * template_alignment_mask[:, None, :]
+ * template_alignment_mask[:, :, None]
+ )
+
+ dist_errors = (pred_dist - target_dist).square()
+ drmsd = (dist_errors.add(1e-4).sqrt().sub(1e-2).mul(motion_mask).sum(dim=(1, 2))) / (
+ motion_mask.long().sum(dim=(1, 2)) + 1
+ )
+ return drmsd
+
+
+def compute_TMscore_lbound(
+ batch: MODEL_BATCH, pred_coords: torch.Tensor, target_coords: torch.Tensor
+) -> torch.Tensor:
+ """Compute the TM-score lower bound for a given batch.
+
+ :param batch: A batch dictionary.
+ :param pred_coords: The predicted coordinates.
+ :param target_coords: The target coordinates.
+ :return: The TM-score lower bound.
+ """
+ features = batch["features"]
+ atom_mask = features["res_atom_mask"].bool().view(batch["metadata"]["num_structid"], -1, 37)
+ pred_coords = pred_coords.view(batch["metadata"]["num_structid"], -1, 37, 3)
+ target_coords = target_coords.view(batch["metadata"]["num_structid"], -1, 37, 3)
+ pred_bb_frames = get_frame_matrix(
+ pred_coords[:, :, 0, :],
+ pred_coords[:, :, 1, :],
+ pred_coords[:, :, 2, :],
+ strict=True,
+ )
+ target_bb_frames = get_frame_matrix(
+ target_coords[:, :, 0, :],
+ target_coords[:, :, 1, :],
+ target_coords[:, :, 2, :],
+ strict=True,
+ )
+ pred_coords_flat = pred_coords[atom_mask].view(batch["metadata"]["num_structid"], -1, 3)
+ target_coords_flat = target_coords[atom_mask].view(batch["metadata"]["num_structid"], -1, 3)
+ # Columns-frames, rows-points
+ # [B, 1, N, 3] - [B, F, 1, 3]
+ aligned_pred_points = cartesian_to_internal(
+ pred_coords_flat.unsqueeze(1), pred_bb_frames.unsqueeze(2)
+ )
+ with torch.no_grad():
+ aligned_target_points = cartesian_to_internal(
+ target_coords_flat.unsqueeze(1), target_bb_frames.unsqueeze(2)
+ )
+ pair_dist_aligned = (aligned_pred_points - aligned_target_points).norm(dim=-1)
+ tm_normalizer = 1.24 * (max(target_coords.shape[1], 19) - 15) ** (1 / 3) - 1.8
+ per_frame_tm = torch.mean(1 / (1 + (pair_dist_aligned / tm_normalizer) ** 2), dim=2)
+ return torch.amax(per_frame_tm, dim=1)
+
+
+def compute_TMscore_raw(
+ batch: MODEL_BATCH, pred_coords: torch.Tensor, target_coords: torch.Tensor
+) -> torch.Tensor:
+ """Compute the raw TM-score for a given batch.
+
+ :param batch: A batch dictionary.
+ :param pred_coords: The predicted coordinates.
+ :param target_coords: The target coordinates.
+ :return: The raw TM-score.
+ """
+ pred_coords = pred_coords.view(batch["metadata"]["num_structid"], -1, 3)
+ target_coords = target_coords.view(batch["metadata"]["num_structid"], -1, 3)
+ pair_dist_aligned = (pred_coords - target_coords).norm(dim=-1)
+ tm_normalizer = 1.24 * (max(target_coords.shape[1], 19) - 15) ** (1 / 3) - 1.8
+ per_struct_tm = torch.mean(1 / (1 + (pair_dist_aligned / tm_normalizer) ** 2), dim=1)
+ return per_struct_tm
+
+
+def compute_lddt_ca(
+ batch: MODEL_BATCH, pred_coords: torch.Tensor, target_coords: torch.Tensor
+) -> torch.Tensor:
+ """Compute the local distance difference test (lDDT) for C-alpha atoms for a given batch.
+
+ :param batch: A batch dictionary.
+ :param pred_coords: The predicted coordinates.
+ :param target_coords: The target coordinates.
+ :return: The lDDT for C-alpha atoms.
+ """
+ pred_coords = pred_coords.view(batch["metadata"]["num_structid"], -1, 37, 3)
+ target_coords = target_coords.view(batch["metadata"]["num_structid"], -1, 37, 3)
+ pred_ca_flat = pred_coords[:, :, 1]
+ target_ca_flat = target_coords[:, :, 1]
+ target_dist = (target_ca_flat[:, :, None] - target_ca_flat[:, None, :]).norm(dim=-1)
+ pred_dist = (pred_ca_flat[:, :, None] - pred_ca_flat[:, None, :]).norm(dim=-1)
+ conserved_mask = target_dist < 15.0
+ lddt = 0
+ for threshold in [0.5, 1, 2, 4]:
+ below_threshold = (pred_dist - target_dist).abs() < threshold
+ lddt = lddt + below_threshold.mul(conserved_mask).sum((1, 2)) / conserved_mask.sum((1, 2))
+ return lddt / 4
+
+
+def compute_lddt_pli(
+ batch: MODEL_BATCH,
+ pred_prot_coords: torch.Tensor,
+ target_prot_coords: torch.Tensor,
+ pred_lig_coords: torch.Tensor,
+ target_lig_coords: torch.Tensor,
+) -> torch.Tensor:
+ """Compute the local distance difference test (lDDT) for protein-ligand interface atoms for a
+ given batch.
+
+ :param batch: A batch dictionary.
+ :param pred_prot_coords: The predicted protein coordinates.
+ :param target_prot_coords: The target protein coordinates.
+ :param pred_lig_coords: The predicted ligand coordinates.
+ :param target_lig_coords: The target ligand coordinates.
+ :return: The lDDT for protein-ligand interface atoms.
+ """
+ features = batch["features"]
+ batch_size = batch["metadata"]["num_structid"]
+ atom_mask = features["res_atom_mask"].bool()
+ pred_prot_coords = pred_prot_coords[atom_mask].view(batch_size, -1, 3)
+ target_prot_coords = target_prot_coords[atom_mask].view(batch_size, -1, 3)
+ pred_lig_coords = pred_lig_coords.view(batch_size, -1, 3)
+ target_lig_coords = target_lig_coords.view(batch_size, -1, 3)
+ target_dist = (target_prot_coords[:, :, None] - target_lig_coords[:, None, :]).norm(dim=-1)
+ pred_dist = (pred_prot_coords[:, :, None] - pred_lig_coords[:, None, :]).norm(dim=-1)
+ conserved_mask = target_dist < 6.0
+ lddt = 0
+ for threshold in [0.5, 1, 2, 4]:
+ below_threshold = (pred_dist - target_dist).abs() < threshold
+ lddt = lddt + below_threshold.mul(conserved_mask).sum((1, 2)) / conserved_mask.sum((1, 2))
+ return lddt / 4
+
+
+def eval_structure_prediction_losses(
+ lit_module: LightningModule,
+ batch: MODEL_BATCH,
+ batch_idx: int,
+ device: Union[str, torch.device],
+ stage: MODEL_STAGE,
+ t_1: float = 1.0,
+) -> MODEL_BATCH:
+ """Evaluate the structure prediction losses for a given batch.
+
+ :param lit_module: The LightningModule object to reference.
+ :param batch: A batch dictionary.
+ :param batch_idx: The batch index.
+ :param device: The device to use.
+ :param stage: The stage of the training.
+ :param t_1: The final timestep in the range [0, 1].
+ :return: Batch dictionary with losses.
+ """
+ assert 0 <= t_1 <= 1, "`t_1` must be in the range `[0, 1]`."
+ batch_size = batch["metadata"]["num_structid"]
+ max(batch["metadata"]["num_a_per_sample"])
+
+ if "num_molid" in batch["metadata"].keys() and batch["metadata"]["num_molid"] > 0:
+ batch["misc"]["protein_only"] = False
+ else:
+ batch["misc"]["protein_only"] = True
+
+ if "augmented_coordinates" in batch["features"].keys():
+ batch["features"]["sdf_coordinates"] = batch["features"]["augmented_coordinates"]
+ is_native_sample = 0
+ else:
+ is_native_sample = 1
+
+ # Sample the timestep for each structure
+ t = torch.rand((batch_size, 1), device=device)
+
+ prior_training = int(random.randint(0, 10) == 1) # nosec
+ if prior_training == 1:
+ t = torch.full_like(t, t_1)
+
+ if lit_module.training and lit_module.hparams.cfg.task.use_template:
+ use_template = bool(random.randint(0, 1)) # nosec
+ else:
+ use_template = lit_module.hparams.cfg.task.use_template
+
+ lit_module.net.assign_timestep_encodings(batch, t)
+ features = batch["features"]
+ indexer = batch["indexer"]
+ metadata = batch["metadata"]
+
+ loss = 0
+ forward_lat_converter = lit_module.net.resolve_latent_converter(
+ [
+ ("features", "res_atom_positions"),
+ ("features", "input_protein_coords"),
+ ],
+ [("features", "sdf_coordinates"), ("features", "input_ligand_coords")],
+ )
+ batch = lit_module.net.prepare_protein_patch_indexers(batch)
+ if not batch["misc"]["protein_only"]:
+ max(metadata["num_i_per_sample"])
+
+ # Evaluate the contact map
+ ref_dist_mat, contact_logit_matrix = eval_true_contact_maps(
+ batch, lit_module.net.CONTACT_SCALE
+ )
+ num_cont_to_sample = max(metadata["num_I_per_sample"])
+ sampled_block_contacts = [
+ None,
+ ]
+ # Onehot contact code sampling
+ with torch.no_grad():
+ for _ in range(num_cont_to_sample):
+ sampled_block_contacts.append(
+ sample_reslig_contact_matrix(
+ batch, contact_logit_matrix, last=sampled_block_contacts[-1]
+ )
+ )
+ forward_lat_converter.lig_res_anchor_mask = sample_res_rowmask_from_contacts(
+ batch,
+ contact_logit_matrix,
+ lit_module.hparams.cfg.task.single_protein_batch,
+ )
+ with torch.no_grad():
+ batch = lit_module.net.forward_interp_plcomplex_latinp(
+ batch, t[:, :, None], forward_lat_converter
+ )
+ if prior_training == 1:
+ iter_id = random.randint(0, num_cont_to_sample) # nosec
+ else:
+ iter_id = num_cont_to_sample
+ batch = lit_module.forward(
+ batch, contact_prediction=False, score=False, use_template=use_template
+ )
+ batch = lit_module.net.run_contact_map_stack(
+ batch,
+ iter_id=iter_id,
+ observed_block_contacts=sampled_block_contacts[iter_id],
+ )
+ pred_distogram = batch["outputs"][f"res_lig_distogram_out_{iter_id}"]
+ (
+ pl_distogram_loss,
+ pl_contact_loss_forward,
+ ) = compute_contact_prediction_losses(
+ pred_distogram, ref_dist_mat, lit_module.net.dist_bins, lit_module.net.CONTACT_SCALE
+ )
+ cont_loss = 0
+ cont_loss = (
+ cont_loss
+ + pl_distogram_loss
+ * lit_module.hparams.cfg.task.contact_loss_weight
+ * is_native_sample
+ )
+ cont_loss = (
+ cont_loss
+ + pl_contact_loss_forward
+ * lit_module.hparams.cfg.task.contact_loss_weight
+ * is_native_sample
+ )
+ lit_module.log(
+ f"{stage}_contact/contact_loss_distogram",
+ pl_distogram_loss.detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}_contact/contact_loss_forwardKL",
+ pl_contact_loss_forward.detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ if lit_module.hparams.cfg.task.freeze_contact_predictor:
+ # Keep the contact prediction parameters in the computational graph but with zero gradients
+ cont_loss *= 0.0
+ else:
+ with torch.no_grad():
+ batch = lit_module.net.forward_interp_plcomplex_latinp(
+ batch, t[:, :, None], forward_lat_converter
+ )
+ iter_id = 0
+ batch = lit_module.forward(
+ batch,
+ iter_id=0,
+ contact_prediction=True,
+ score=False,
+ use_template=use_template,
+ )
+ protein_distogram_loss = compute_protein_distogram_loss(
+ batch,
+ batch["features"]["res_atom_positions"][:, 1],
+ lit_module.net.dist_bins,
+ lit_module.net.dgram_head,
+ entry=f"res_res_grid_attr_flat_out_{iter_id}",
+ )
+ lit_module.log(
+ f"{stage}_contact/prot_distogram_loss",
+ protein_distogram_loss.detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ if lit_module.hparams.cfg.task.freeze_contact_predictor:
+ # Keep the distogram prediction parameters in the computational graph but with zero gradients
+ protein_distogram_loss *= 0.0
+
+ # NOTE: we keep the loss weighting time-independent since `sigma=1` for all prior distributions (where relevant)
+ lambda_weighting = t.new_ones(batch_size)
+
+ # Run score head and evaluate structure prediction losses
+ res_atom_mask = features["res_atom_mask"].bool()
+
+ scores = lit_module.net.run_score_head(batch, embedding_iter_id=iter_id)
+
+ if lit_module.training:
+ # # Sigmoid scaling
+ # violation_loss_ratio = 1 / (
+ # 1
+ # + math.exp(10 - 12 * lit_module.current_epoch / lit_module.trainer.max_epochs)
+ # )
+ # violation_loss_ratio = (lit_module.current_epoch / lit_module.trainer.max_epochs)
+ violation_loss_ratio = 1.0
+ else:
+ violation_loss_ratio = 1.0
+
+ if not batch["misc"]["protein_only"]:
+ if "binding_site_mask_clean" not in batch["features"]:
+ with torch.no_grad():
+ min_lig_res_dist_clean = (
+ (
+ batch["features"]["res_atom_positions"][:, 1].view(batch_size, -1, 3)[
+ :, :, None
+ ]
+ - batch["features"]["sdf_coordinates"].view(batch_size, -1, 3)[:, None, :]
+ )
+ .norm(dim=-1)
+ .amin(dim=2)
+ ).flatten(0, 1)
+ binding_site_mask_clean = (
+ min_lig_res_dist_clean < lit_module.net.BINDING_SITE_CUTOFF
+ )
+ batch["features"]["binding_site_mask_clean"] = binding_site_mask_clean
+ coords_pred_prot = scores["final_coords_prot_atom_padded"][res_atom_mask].view(
+ batch_size, -1, 3
+ )
+ coords_ref_prot = batch["features"]["res_atom_positions"][res_atom_mask].view(
+ batch_size, -1, 3
+ )
+ coords_pred_bs_prot = scores["final_coords_prot_atom_padded"][
+ res_atom_mask & batch["features"]["binding_site_mask_clean"][:, None]
+ ].view(batch_size, -1, 3)
+ coords_ref_bs_prot = batch["features"]["res_atom_positions"][
+ res_atom_mask & batch["features"]["binding_site_mask_clean"][:, None]
+ ].view(batch_size, -1, 3)
+ coords_pred_lig = scores["final_coords_lig_atom"].view(batch_size, -1, 3)
+ coords_ref_lig = batch["features"]["sdf_coordinates"].view(batch_size, -1, 3)
+ coords_pred = torch.cat([coords_pred_prot, coords_pred_lig], dim=1)
+ coords_ref = torch.cat([coords_ref_prot, coords_ref_lig], dim=1)
+ coords_pred_bs = torch.cat([coords_pred_bs_prot, coords_pred_lig], dim=1)
+ coords_ref_bs = torch.cat([coords_ref_bs_prot, coords_ref_lig], dim=1)
+ n_I_per_sample = max(metadata["num_I_per_sample"])
+ lig_frame_atm_idx = torch.stack(
+ [
+ indexer["gather_idx_ijk_i"][indexer["gather_idx_I_ijk"]][:n_I_per_sample],
+ indexer["gather_idx_ijk_j"][indexer["gather_idx_I_ijk"]][:n_I_per_sample],
+ indexer["gather_idx_ijk_k"][indexer["gather_idx_I_ijk"]][:n_I_per_sample],
+ ],
+ dim=0,
+ )
+ (
+ global_fape_pview,
+ global_fape_lview,
+ normalized_fape,
+ ) = compute_fape_from_atom37(
+ batch,
+ device,
+ scores["final_coords_prot_atom_padded"],
+ batch["features"]["res_atom_positions"],
+ pred_lig_coords=scores["final_coords_lig_atom"],
+ target_lig_coords=batch["features"]["sdf_coordinates"],
+ lig_frame_atm_idx=lig_frame_atm_idx,
+ split_pl_views=True,
+ )
+ aa_distgeom_error = compute_aa_distance_geometry_loss(
+ batch,
+ scores["final_coords_prot_atom_padded"],
+ batch["features"]["res_atom_positions"],
+ )
+ lig_distgeom_error = compute_sm_distance_geometry_loss(
+ batch,
+ scores["final_coords_lig_atom"],
+ batch["features"]["sdf_coordinates"],
+ )
+ glob_drmsd, _ = compute_drmsd_and_clashloss(
+ batch,
+ device,
+ scores["final_coords_prot_atom_padded"],
+ batch["features"]["res_atom_positions"],
+ lit_module.net.atnum2vdw_uff,
+ pred_lig_coords=scores["final_coords_lig_atom"],
+ target_lig_coords=batch["features"]["sdf_coordinates"],
+ ligatm_types=batch["features"]["atomic_numbers"].long(),
+ )
+ bs_drmsd, clash_error = compute_drmsd_and_clashloss(
+ batch,
+ device,
+ scores["final_coords_prot_atom_padded"],
+ batch["features"]["res_atom_positions"],
+ lit_module.net.atnum2vdw_uff,
+ pred_lig_coords=scores["final_coords_lig_atom"],
+ target_lig_coords=batch["features"]["sdf_coordinates"],
+ ligatm_types=batch["features"]["atomic_numbers"].long(),
+ binding_site=True,
+ )
+ pli_drmsd, _ = compute_drmsd_and_clashloss(
+ batch,
+ device,
+ scores["final_coords_prot_atom_padded"],
+ batch["features"]["res_atom_positions"],
+ lit_module.net.atnum2vdw_uff,
+ pred_lig_coords=scores["final_coords_lig_atom"],
+ target_lig_coords=batch["features"]["sdf_coordinates"],
+ ligatm_types=batch["features"]["atomic_numbers"].long(),
+ pl_interface=True,
+ )
+ distgeom_loss = (
+ aa_distgeom_error.mul(lambda_weighting) * max(metadata["num_a_per_sample"])
+ + lig_distgeom_error.mul(lambda_weighting) * max(metadata["num_i_per_sample"])
+ ).mean() / max(metadata["num_a_per_sample"])
+
+ fape_loss = (
+ (
+ global_fape_pview
+ + global_fape_lview
+ * (
+ lit_module.hparams.cfg.task.ligand_score_loss_weight
+ / lit_module.hparams.cfg.task.global_score_loss_weight
+ )
+ + normalized_fape
+ )
+ .mul(lambda_weighting)
+ .mean()
+ )
+
+ if not lit_module.hparams.cfg.task.freeze_score_head:
+ loss = (
+ loss
+ + fape_loss
+ * lit_module.hparams.cfg.task.global_score_loss_weight
+ * is_native_sample
+ )
+ loss = (
+ loss
+ + glob_drmsd.mul(lambda_weighting).mean()
+ * lit_module.hparams.cfg.task.drmsd_loss_weight
+ )
+ if use_template:
+ twe_drmsd = compute_template_weighted_centroid_drmsd(
+ batch, scores["final_coords_prot_atom_padded"]
+ )
+ if not lit_module.hparams.cfg.task.freeze_score_head:
+ loss = (
+ loss
+ + twe_drmsd.mul(lambda_weighting).mean()
+ * lit_module.hparams.cfg.task.drmsd_loss_weight
+ )
+ lit_module.log(
+ f"{stage}/drmsd_loss_weighted",
+ twe_drmsd.mul(lambda_weighting).mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/drmsd_weighted",
+ twe_drmsd.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ if not lit_module.hparams.cfg.task.freeze_score_head:
+ loss = (
+ loss
+ + bs_drmsd.mul(lambda_weighting).mean()
+ * lit_module.hparams.cfg.task.drmsd_loss_weight
+ )
+ loss = (
+ loss
+ + pli_drmsd.mul(lambda_weighting).mean()
+ * lit_module.hparams.cfg.task.drmsd_loss_weight
+ )
+
+ loss = (
+ loss
+ + distgeom_loss
+ * lit_module.hparams.cfg.task.local_distgeom_loss_weight
+ * violation_loss_ratio
+ )
+ loss = (
+ loss
+ + clash_error.mul(lambda_weighting).mean()
+ * lit_module.hparams.cfg.task.clash_loss_weight
+ * violation_loss_ratio
+ )
+ if not lit_module.hparams.cfg.task.freeze_contact_predictor:
+ loss = (0.1 + 0.9 * prior_training) * cont_loss + (1 - prior_training * 0.99) * loss
+ loss = (
+ loss + protein_distogram_loss * lit_module.hparams.cfg.task.distogram_loss_weight
+ )
+ with torch.no_grad():
+ tm_lbound = compute_TMscore_lbound(
+ batch,
+ scores["final_coords_prot_atom_padded"],
+ batch["features"]["res_atom_positions"],
+ )
+ lig_rmsd = segment_mean(
+ (
+ (coords_pred_lig - coords_pred_prot.mean(dim=1, keepdim=True))
+ - (coords_ref_lig - coords_ref_prot.mean(dim=1, keepdim=True))
+ )
+ .square()
+ .sum(dim=-1)
+ .flatten(0, 1),
+ indexer["gather_idx_i_molid"],
+ metadata["num_molid"],
+ ).sqrt()
+ lit_module.log(
+ f"{stage}/tm_lbound",
+ tm_lbound.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/ligand_rmsd_ubound",
+ lig_rmsd.mean().detach(),
+ on_epoch=True,
+ batch_size=lig_rmsd.shape[0],
+ )
+ # L1 score matching loss
+ dsm_loss_global = (
+ (
+ (coords_pred - coords_pred_prot.mean(dim=1, keepdim=True))
+ - (coords_ref - coords_ref_prot.mean(dim=1, keepdim=True))
+ )
+ .square()
+ .sum(dim=-1)
+ .add(1e-2)
+ .sqrt()
+ .sub(1e-1)
+ .mean(dim=1)
+ .mul(lambda_weighting)
+ )
+ dsm_loss_site = (
+ (
+ (coords_pred_bs - coords_pred_bs_prot.mean(dim=1, keepdim=True))
+ - (coords_ref_bs - coords_ref_bs_prot.mean(dim=1, keepdim=True))
+ )
+ .square()
+ .sum(dim=-1)
+ .add(1e-2)
+ .sqrt()
+ .sub(1e-1)
+ .mean(dim=1)
+ .mul(lambda_weighting)
+ )
+ dsm_loss_ligand = (
+ (
+ (coords_pred_lig - coords_pred.mean(dim=1, keepdim=True))
+ - (coords_ref_lig - coords_ref.mean(dim=1, keepdim=True))
+ )
+ .square()
+ .sum(dim=-1)
+ .add(1e-2)
+ .sqrt()
+ .sub(1e-1)
+ .mean(dim=1)
+ .mul(lambda_weighting)
+ )
+ lit_module.log(
+ f"{stage}/denoising_loss_global",
+ dsm_loss_global.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/denoising_loss_site",
+ dsm_loss_site.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/denoising_loss_ligand",
+ dsm_loss_ligand.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/drmsd_loss_global",
+ glob_drmsd.mul(lambda_weighting).mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/drmsd_loss_site",
+ bs_drmsd.mul(lambda_weighting).mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/drmsd_loss_pli",
+ pli_drmsd.mul(lambda_weighting).mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/drmsd_global",
+ glob_drmsd.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/drmsd_site",
+ bs_drmsd.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/drmsd_pli",
+ pli_drmsd.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/fape_global_proteinview",
+ global_fape_pview.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/fape_global_ligandview",
+ global_fape_lview.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/fape_normalized",
+ normalized_fape.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/fape_loss",
+ fape_loss.detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}_violation/aa_distgeom_error",
+ aa_distgeom_error.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}_violation/lig_distgeom_error",
+ lig_distgeom_error.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}_violation/clash_error",
+ clash_error.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}_violation/clash_loss",
+ clash_error.mul(lambda_weighting).mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}_violation/distgeom_loss",
+ distgeom_loss.detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ else:
+ coords_pred = scores["final_coords_prot_atom_padded"][res_atom_mask].view(
+ batch_size, -1, 3
+ )
+ coords_ref = batch["features"]["res_atom_positions"][res_atom_mask].view(batch_size, -1, 3)
+ global_fape_pview, normalized_fape = compute_fape_from_atom37(
+ batch,
+ device,
+ scores["final_coords_prot_atom_padded"],
+ batch["features"]["res_atom_positions"],
+ )
+ aa_distgeom_error = compute_aa_distance_geometry_loss(
+ batch,
+ scores["final_coords_prot_atom_padded"],
+ batch["features"]["res_atom_positions"],
+ )
+ glob_drmsd, clash_error = compute_drmsd_and_clashloss(
+ batch,
+ device,
+ scores["final_coords_prot_atom_padded"],
+ batch["features"]["res_atom_positions"],
+ lit_module.net.atnum2vdw_uff,
+ )
+ distgeom_loss = aa_distgeom_error.mul(lambda_weighting).mean()
+ fape_loss = (global_fape_pview + normalized_fape).mul(lambda_weighting).mean()
+
+ global_fape_pview.detach()
+ if not lit_module.hparams.cfg.task.freeze_score_head:
+ loss = (
+ loss
+ + distgeom_loss
+ * lit_module.hparams.cfg.task.local_distgeom_loss_weight
+ * violation_loss_ratio
+ )
+ loss = loss + fape_loss * lit_module.hparams.cfg.task.global_score_loss_weight
+ loss = (
+ loss
+ + glob_drmsd.mul(lambda_weighting).mean()
+ * lit_module.hparams.cfg.task.drmsd_loss_weight
+ )
+ if use_template:
+ twe_drmsd = compute_template_weighted_centroid_drmsd(
+ batch, scores["final_coords_prot_atom_padded"]
+ )
+ if not lit_module.hparams.cfg.task.freeze_score_head:
+ loss = (
+ loss
+ + twe_drmsd.mul(lambda_weighting).mean()
+ * lit_module.hparams.cfg.task.drmsd_loss_weight
+ )
+ lit_module.log(
+ f"{stage}/drmsd_loss_weighted",
+ twe_drmsd.mul(lambda_weighting).mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/drmsd_weighted",
+ twe_drmsd.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ if not lit_module.hparams.cfg.task.freeze_score_head:
+ loss = (
+ loss
+ + clash_error.mul(lambda_weighting).mean()
+ * lit_module.hparams.cfg.task.clash_loss_weight
+ * violation_loss_ratio
+ )
+ if not lit_module.hparams.cfg.task.freeze_contact_predictor:
+ loss = (
+ loss + protein_distogram_loss * lit_module.hparams.cfg.task.distogram_loss_weight
+ )
+
+ with torch.no_grad():
+ dsm_loss_global = (
+ (
+ (coords_pred - coords_pred.mean(dim=1, keepdim=True))
+ - (coords_ref - coords_ref.mean(dim=1, keepdim=True))
+ )
+ .square()
+ .sum(dim=-1)
+ .add(1e-2)
+ .sqrt()
+ .sub(1e-1)
+ .mean(dim=1)
+ .mul(lambda_weighting)
+ )
+ lit_module.log(
+ f"{stage}/denoising_loss_global",
+ dsm_loss_global.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ tm_lbound = compute_TMscore_lbound(
+ batch,
+ scores["final_coords_prot_atom_padded"],
+ batch["features"]["res_atom_positions"],
+ )
+ lit_module.log(
+ f"{stage}/tm_lbound",
+ tm_lbound.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/drmsd_loss_global",
+ glob_drmsd.mul(lambda_weighting).mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/drmsd_global",
+ glob_drmsd.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/fape_global_proteinview",
+ global_fape_pview.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/fape_normalized",
+ normalized_fape.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}/fape_loss",
+ fape_loss.detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}_violation/aa_distgeom_error",
+ aa_distgeom_error.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}_violation/clash_error",
+ clash_error.mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}_violation/clash_loss",
+ clash_error.mul(lambda_weighting).mean().detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}_violation/distgeom_loss",
+ distgeom_loss.detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ if torch.is_tensor(loss) and not torch.isnan(loss):
+ lit_module.log(
+ f"{stage}/loss",
+ loss.detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ sync_dist=(stage != "train"),
+ )
+ batch["outputs"]["loss"] = loss
+ if not torch.is_tensor(batch["outputs"]["loss"]) and batch["outputs"]["loss"] == 0:
+ batch["outputs"]["loss"] = None
+ return batch
+
+
+def eval_auxiliary_estimation_losses(
+ lit_module: LightningModule,
+ batch: MODEL_BATCH,
+ stage: MODEL_STAGE,
+ loss_mode: LOSS_MODES,
+ **kwargs: Dict[str, Any],
+) -> MODEL_BATCH:
+ """Evaluate the auxiliary estimation losses for a given batch.
+
+ :param lit_module: The LightningModule object to reference.
+ :param batch: A batch dictionary.
+ :param stage: The stage of the training.
+ :param loss_mode: The loss mode to use.
+ :param kwargs: Additional keyword arguments.
+ :return: Batch dictionary with losses.
+ """
+ use_template = bool(random.randint(0, 1)) # nosec
+ if use_template:
+ # Enable higher ligand diversity when using backbone template
+ start_time = 1.0
+ else:
+ start_time = random.randint(1, 5) / 5 # nosec
+ with torch.no_grad():
+ if loss_mode == "auxiliary_estimation_without_structure_prediction":
+ # Sample the structure without using the structure prediction head
+ # i.e., Provide the holo (ground-truth) protein and ligand structures for affinity estimation
+ output_struct = {
+ "receptor": batch["features"]["res_atom_positions"].flatten(0, 1),
+ "receptor_padded": batch["features"]["res_atom_positions"],
+ "ligands": batch["features"]["sdf_coordinates"],
+ }
+ else:
+ output_struct = lit_module.net.sample_pl_complex_structures(
+ batch,
+ sampler="VDODE",
+ sampler_eta=1.0,
+ num_steps=int(5 / start_time),
+ start_time=start_time,
+ exact_prior=True,
+ use_template=use_template,
+ cutoff=20.0, # Hot logits
+ )
+ batch_size = batch["metadata"]["num_structid"]
+ batch = lit_module.net.run_auxiliary_estimation(batch, output_struct, **kwargs)
+ if lit_module.hparams.cfg.confidence.enabled:
+ with torch.no_grad():
+ # Receptor centroids
+ ref_coords = (
+ (
+ batch["features"]["res_atom_positions"]
+ .mul(batch["features"]["res_atom_mask"].bool()[:, :, None])
+ .sum(dim=1)
+ .div(batch["features"]["res_atom_mask"].bool().sum(dim=1)[:, None] + 1e-9)
+ )
+ .contiguous()
+ .view(batch_size, -1, 3)
+ )
+ pred_coords = (
+ (
+ output_struct["receptor_padded"]
+ .mul(batch["features"]["res_atom_mask"].bool()[:, :, None])
+ .sum(dim=1)
+ .div(batch["features"]["res_atom_mask"].bool().sum(dim=1)[:, None] + 1e-9)
+ )
+ .contiguous()
+ .view(batch_size, -1, 3)
+ )
+ # The number of effective protein atoms used in plddt calculation
+ n_protatm_per_sample = pred_coords.shape[1]
+ if output_struct["ligands"] is not None:
+ ref_lig_coords = (
+ batch["features"]["sdf_coordinates"].contiguous().view(batch_size, -1, 3)
+ )
+ ref_coords = torch.cat([ref_coords, ref_lig_coords], dim=1)
+ pred_lig_coords = output_struct["ligands"].contiguous().view(batch_size, -1, 3)
+ pred_coords = torch.cat([pred_coords, pred_lig_coords], dim=1)
+ per_atom_lddt, per_atom_lddt_gram = compute_per_atom_lddt(
+ batch, pred_coords, ref_coords
+ )
+
+ plddt_dev = (per_atom_lddt - batch["outputs"]["plddt"]).abs().mean()
+ confidence_loss = (
+ F.cross_entropy(
+ batch["outputs"]["plddt_logits"].flatten(0, 1),
+ per_atom_lddt_gram.flatten(0, 1),
+ reduction="none",
+ )
+ .contiguous()
+ .view(batch_size, -1)
+ )
+ conf_loss = confidence_loss.mean()
+ if output_struct["ligands"] is not None:
+ plddt_dev_lig = (
+ (
+ per_atom_lddt.view(batch_size, -1)[:, n_protatm_per_sample:]
+ - batch["outputs"]["plddt"].view(batch_size, -1)[:, n_protatm_per_sample:]
+ )
+ .abs()
+ .mean()
+ )
+ conf_loss_lig = confidence_loss[:, n_protatm_per_sample:].mean()
+ conf_loss = conf_loss + conf_loss_lig # + plddt_dev_lig * 0.1
+ lit_module.log(
+ f"{stage}_confidence/plddt_dev_lig",
+ plddt_dev_lig.detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}_confidence/plddt_dev",
+ plddt_dev.detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ lit_module.log(
+ f"{stage}_confidence/loss",
+ conf_loss.detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ sync_dist=(stage != "train"),
+ )
+ if lit_module.hparams.cfg.task.freeze_confidence:
+ # Keep the confidence prediction parameters in the computational graph but with zero gradients
+ conf_loss *= 0
+ else:
+ conf_loss = 0
+ if lit_module.hparams.cfg.affinity.enabled:
+ num_molid_per_sample = batch["metadata"]["num_molid"] // batch_size
+ gather_idx_molid_structid = torch.arange(
+ batch_size, device=batch["outputs"]["affinity_logits"].device
+ ).repeat_interleave(num_molid_per_sample)
+
+ # Calculate affinity loss as the mean squared error between the predicted affinity logits and the ground-truth affinity values
+ affinity_logits = batch["outputs"]["affinity_logits"]
+ # Substitute missing ground-truth affinity values with the affinity head's (detached) predicted logits to indicate no learning signal for these examples
+ affinity = torch.where(
+ batch["features"]["affinity"].isnan(),
+ affinity_logits.detach(),
+ batch["features"]["affinity"],
+ )
+ aff_loss = segment_mean(
+ # Find the (batched) mean squared error over all ligand chains in the same complex, then calculate the mean of each batch
+ (affinity_logits - affinity).square(),
+ gather_idx_molid_structid,
+ batch_size,
+ ).mean()
+ lit_module.log(
+ f"{stage}_affinity/loss",
+ aff_loss.detach(),
+ on_epoch=True,
+ batch_size=batch_size,
+ sync_dist=(stage != "train"),
+ )
+ if lit_module.hparams.cfg.task.freeze_affinity:
+ # Keep the affinity prediction parameters in the computational graph but with zero gradients
+ aff_loss *= 0
+ else:
+ aff_loss = 0
+ plddt_loss = conf_loss * lit_module.hparams.cfg.task.plddt_loss_weight
+ affinity_loss = aff_loss * lit_module.hparams.cfg.task.affinity_loss_weight
+ batch["outputs"]["loss"] = plddt_loss + affinity_loss
+ if not torch.is_tensor(batch["outputs"]["loss"]) and batch["outputs"]["loss"] == 0:
+ batch["outputs"]["loss"] = None
+ return batch
diff --git a/forks/FlowDock/flowdock/models/components/mht_encoder.py b/forks/FlowDock/flowdock/models/components/mht_encoder.py
new file mode 100644
index 00000000..7dca7953
--- /dev/null
+++ b/forks/FlowDock/flowdock/models/components/mht_encoder.py
@@ -0,0 +1,364 @@
+# Adapted from: https://github.com/zrqiao/NeuralPLexer
+
+import rootutils
+import torch
+from beartype.typing import Any, Dict, Optional, Tuple
+from omegaconf import DictConfig
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock.models.components.hetero_graph import make_multi_relation_graph_batcher
+from flowdock.models.components.modules import TransformerLayer
+from flowdock.utils import RankedLogger
+from flowdock.utils.model_utils import GELUMLP, segment_softmax, segment_sum
+
+MODEL_BATCH = Dict[str, Any]
+STATE_DICT = Dict[str, Any]
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+class PathConvStack(torch.nn.Module):
+ """Path integral convolution stack for ligand encoding."""
+
+ def __init__(
+ self,
+ pair_channels: int,
+ n_heads: int = 8,
+ max_pi_length: int = 8,
+ dropout: float = 0.0,
+ ):
+ """Initialize PathConvStack model."""
+ super().__init__()
+ self.pair_channels = pair_channels
+ self.max_pi_length = max_pi_length
+ self.n_heads = n_heads
+
+ self.prop_value_layer = torch.nn.Linear(pair_channels, n_heads, bias=False)
+ self.triangle_pair_kernel_layer = torch.nn.Linear(pair_channels, n_heads, bias=False)
+ self.prop_update_mlp = GELUMLP(
+ n_heads * (max_pi_length + 1), pair_channels, dropout=dropout
+ )
+
+ def forward(
+ self,
+ prop_attr: torch.Tensor,
+ stereo_attr: torch.Tensor,
+ indexer: Dict[str, torch.LongTensor],
+ metadata: Dict[str, Any],
+ ) -> torch.Tensor:
+ """Forward pass for PathConvStack model.
+
+ :param prop_attr: Atom-frame pair attributes.
+ :param stereo_attr: Stereochemistry attributes.
+ :param indexer: A dictionary of indices.
+ :param metadata: A dictionary of metadata.
+ :return: Updated atom-frame pair attributes.
+ """
+ triangle_pair_kernel = self.triangle_pair_kernel_layer(stereo_attr)
+ # Segment-wise softmax, normalized by outgoing triangles
+ triangle_pair_alpha = segment_softmax(
+ triangle_pair_kernel, indexer["gather_idx_ijkl_jkl"], metadata["num_ijk"]
+ ) # .div(self.max_pi_length)
+ # Uijk,ijkl->ujkl pair representation update
+ kernel = triangle_pair_alpha[indexer["gather_idx_Uijkl_ijkl"]]
+ out_prop_attr = [self.prop_value_layer(prop_attr)]
+ for _ in range(self.max_pi_length):
+ gathered_prop_attr = out_prop_attr[-1][indexer["gather_idx_Uijkl_Uijk"]]
+ out_prop_attr.append(
+ segment_sum(
+ kernel.mul(gathered_prop_attr),
+ indexer["gather_idx_Uijkl_ujkl"],
+ metadata["num_Uijk"],
+ )
+ )
+ new_prop_attr = torch.cat(out_prop_attr, dim=-1)
+ new_prop_attr = self.prop_update_mlp(new_prop_attr) + prop_attr
+ return new_prop_attr
+
+
+class PIFormer(torch.nn.Module):
+ """PIFormer model for ligand encoding."""
+
+ def __init__(
+ self,
+ node_channels: int,
+ pair_channels: int,
+ n_atom_encodings: int,
+ n_bond_encodings: int,
+ n_atom_pos_encodings: int,
+ n_stereo_encodings: int,
+ heads: int,
+ head_dim: int,
+ max_path_length: int = 4,
+ n_transformer_stacks=4,
+ hidden_dim: Optional[int] = None,
+ dropout: float = 0.0,
+ ):
+ """Initialize PIFormer model."""
+ super().__init__()
+ self.node_channels = node_channels
+ self.pair_channels = pair_channels
+ self.max_pi_length = max_path_length
+ self.n_transformer_stacks = n_transformer_stacks
+ self.n_atom_encodings = n_atom_encodings
+ self.n_bond_encodings = n_bond_encodings
+ self.n_atom_pair_encodings = n_bond_encodings + 4
+ self.n_atom_pos_encodings = n_atom_pos_encodings
+
+ self.input_atom_layer = torch.nn.Linear(n_atom_encodings, node_channels)
+ self.input_pair_layer = torch.nn.Linear(self.n_atom_pair_encodings, pair_channels)
+ self.input_stereo_layer = torch.nn.Linear(n_stereo_encodings, pair_channels)
+ self.input_prop_layer = GELUMLP(
+ self.n_atom_pair_encodings * 3,
+ pair_channels,
+ )
+ self.path_integral_stacks = torch.nn.ModuleList(
+ [
+ PathConvStack(
+ pair_channels,
+ max_pi_length=max_path_length,
+ dropout=dropout,
+ )
+ for _ in range(n_transformer_stacks)
+ ]
+ )
+ self.graph_transformer_stacks = torch.nn.ModuleList(
+ [
+ TransformerLayer(
+ node_channels,
+ heads,
+ head_dim=head_dim,
+ edge_channels=pair_channels,
+ hidden_dim=hidden_dim,
+ dropout=dropout,
+ edge_update=True,
+ )
+ for _ in range(n_transformer_stacks)
+ ]
+ )
+
+ def forward(self, batch: MODEL_BATCH, masking_rate: float = 0.0) -> MODEL_BATCH:
+ """Forward pass for PIFormer model.
+
+ :param batch: A batch dictionary.
+ :param masking_rate: Masking rate.
+ :return: A batch dictionary.
+ """
+ features = batch["features"]
+ indexer = batch["indexer"]
+ metadata = batch["metadata"]
+ features["atom_encodings"] = features["atom_encodings"]
+ atom_attr = features["atom_encodings"]
+ atom_pair_attr = features["atom_pair_encodings"]
+ af_pair_attr = features["atom_frame_pair_encodings"]
+ stereo_enc = features["stereo_chemistry_encodings"]
+ batch["features"]["lig_atom_token"] = atom_attr.detach().clone()
+ batch["features"]["lig_pair_token"] = atom_pair_attr.detach().clone()
+
+ atom_mask = torch.rand(atom_attr.shape[0], device=atom_attr.device) > masking_rate
+ stereo_mask = torch.rand(stereo_enc.shape[0], device=stereo_enc.device) > masking_rate
+ atom_pair_mask = (
+ torch.rand(atom_pair_attr.shape[0], device=atom_pair_attr.device) > masking_rate
+ )
+ af_pair_mask = (
+ torch.rand(af_pair_attr.shape[0], device=atom_pair_attr.device) > masking_rate
+ )
+ atom_attr = atom_attr * atom_mask[:, None]
+ stereo_enc = stereo_enc * stereo_mask[:, None]
+ atom_pair_attr = atom_pair_attr * atom_pair_mask[:, None]
+ af_pair_attr = af_pair_attr * af_pair_mask[:, None]
+
+ # Embedding blocks
+ metadata["num_atom"] = metadata["num_u"]
+ metadata["num_frame"] = metadata["num_ijk"]
+ atom_attr = self.input_atom_layer(atom_attr)
+ atom_pair_attr = self.input_pair_layer(atom_pair_attr)
+ triangle_attr = atom_attr.new_zeros(metadata["num_frame"], self.node_channels)
+ # Initialize atom-frame pair attributes. Reusing uv indices
+ prop_attr = self.input_prop_layer(af_pair_attr)
+ stereo_attr = self.input_stereo_layer(stereo_enc)
+
+ graph_relations = [
+ ("atom_to_atom", "gather_idx_uv_u", "gather_idx_uv_v", "atom", "atom"),
+ (
+ "atom_to_frame",
+ "gather_idx_Uijk_u",
+ "gather_idx_Uijk_ijk",
+ "atom",
+ "frame",
+ ),
+ (
+ "frame_to_atom",
+ "gather_idx_Uijk_ijk",
+ "gather_idx_Uijk_u",
+ "frame",
+ "atom",
+ ),
+ (
+ "frame_to_frame",
+ "gather_idx_ijkl_ijk",
+ "gather_idx_ijkl_jkl",
+ "frame",
+ "frame",
+ ),
+ ]
+
+ graph_batcher = make_multi_relation_graph_batcher(graph_relations, indexer, metadata)
+ merged_edge_idx = graph_batcher.collate_idx_list(indexer)
+ node_reps = {"atom": atom_attr, "frame": triangle_attr}
+ edge_reps = {
+ "atom_to_atom": atom_pair_attr,
+ "atom_to_frame": prop_attr,
+ "frame_to_atom": prop_attr,
+ "frame_to_frame": stereo_attr,
+ }
+
+ # Graph path integral recursion
+ for block_id in range(self.n_transformer_stacks):
+ merged_node_attr = graph_batcher.collate_node_attr(node_reps)
+ merged_edge_attr = graph_batcher.collate_edge_attr(edge_reps)
+ _, merged_node_attr, merged_edge_attr = self.graph_transformer_stacks[block_id](
+ merged_node_attr,
+ merged_node_attr,
+ merged_edge_idx,
+ merged_edge_attr,
+ )
+ node_reps = graph_batcher.offload_node_attr(merged_node_attr)
+ edge_reps = graph_batcher.offload_edge_attr(merged_edge_attr)
+ prop_attr = edge_reps["atom_to_frame"]
+ stereo_attr = edge_reps["frame_to_frame"]
+ prop_attr = prop_attr + self.path_integral_stacks[block_id](
+ prop_attr,
+ stereo_attr,
+ indexer,
+ metadata,
+ )
+ edge_reps["atom_to_frame"] = prop_attr
+
+ node_reps["sampled_frame"] = node_reps["frame"][indexer["gather_idx_I_ijk"]]
+
+ batch["metadata"]["num_lig_atm"] = metadata["num_u"]
+ batch["metadata"]["num_lig_trp"] = metadata["num_I"]
+
+ batch["features"]["lig_atom_attr"] = node_reps["atom"]
+ # Downsampled ligand frames
+ batch["features"]["lig_trp_attr"] = node_reps["sampled_frame"]
+ batch["features"]["lig_atom_pair_attr"] = edge_reps["atom_to_atom"]
+ batch["features"]["lig_prop_attr"] = edge_reps["atom_to_frame"]
+ edge_reps["sampled_atom_to_sampled_frame"] = edge_reps["atom_to_frame"][
+ indexer["gather_idx_UI_Uijk"]
+ ]
+ batch["features"]["lig_af_pair_attr"] = edge_reps["sampled_atom_to_sampled_frame"]
+ return batch
+
+
+def resolve_ligand_encoder(
+ ligand_model_cfg: DictConfig,
+ task_cfg: DictConfig,
+ state_dict: Optional[STATE_DICT] = None,
+) -> torch.nn.Module:
+ """Instantiates a PIFormer model for ligand encoding.
+
+ :param ligand_model_cfg: Ligand model configuration.
+ :param task_cfg: Task configuration.
+ :param state_dict: Optional (potentially-pretrained) state dictionary.
+ :return: Ligand encoder model.
+ """
+ model = PIFormer(
+ ligand_model_cfg.node_channels,
+ ligand_model_cfg.pair_channels,
+ ligand_model_cfg.n_atom_encodings,
+ ligand_model_cfg.n_bond_encodings,
+ ligand_model_cfg.n_atom_pos_encodings,
+ ligand_model_cfg.n_stereo_encodings,
+ ligand_model_cfg.n_attention_heads,
+ ligand_model_cfg.attention_head_dim,
+ hidden_dim=ligand_model_cfg.hidden_dim,
+ max_path_length=ligand_model_cfg.max_path_integral_length,
+ n_transformer_stacks=ligand_model_cfg.n_transformer_stacks,
+ dropout=task_cfg.dropout,
+ )
+ if ligand_model_cfg.from_pretrained and state_dict is not None:
+ try:
+ model.load_state_dict(
+ {
+ ".".join(k.split(".")[1:]): v
+ for k, v in state_dict.items()
+ if k.startswith("ligand_encoder")
+ }
+ )
+ log.info(
+ "Successfully loaded pretrained ligand Molecular Heat Transformer (MHT) weights."
+ )
+ except Exception as e:
+ log.warning(
+ f"Skipping loading of pretrained ligand Molecular Heat Transformer (MHT) weights due to: {e}."
+ )
+ return model
+
+
+def resolve_relational_reasoning_module(
+ protein_model_cfg: DictConfig,
+ ligand_model_cfg: DictConfig,
+ relational_reasoning_cfg: DictConfig,
+ state_dict: Optional[STATE_DICT] = None,
+) -> Tuple[torch.nn.Module, torch.nn.Module, torch.nn.Module]:
+ """Instantiates relational reasoning module for ligand encoding.
+
+ :param protein_model_cfg: Protein model configuration.
+ :param ligand_model_cfg: Ligand model configuration.
+ :param relational_reasoning_cfg: Relational reasoning configuration.
+ :param state_dict: Optional (potentially-pretrained) state dictionary.
+ :return: Relational reasoning modules for ligand encoding.
+ """
+ molgraph_single_projector = torch.nn.Linear(
+ ligand_model_cfg.node_channels, protein_model_cfg.residue_dim, bias=False
+ )
+ molgraph_pair_projector = torch.nn.Linear(
+ ligand_model_cfg.pair_channels, protein_model_cfg.pair_dim, bias=False
+ )
+ covalent_embed = torch.nn.Embedding(2, protein_model_cfg.pair_dim)
+ if relational_reasoning_cfg.from_pretrained and state_dict is not None:
+ try:
+ molgraph_single_projector.load_state_dict(
+ {
+ ".".join(k.split(".")[1:]): v
+ for k, v in state_dict.items()
+ if k.startswith("molgraph_single_projector")
+ }
+ )
+ log.info("Successfully loaded pretrained ligand graph single projector weights.")
+ except Exception as e:
+ log.warning(
+ f"Skipping loading of pretrained ligand graph single projector weights due to: {e}."
+ )
+
+ try:
+ molgraph_pair_projector.load_state_dict(
+ {
+ ".".join(k.split(".")[1:]): v
+ for k, v in state_dict.items()
+ if k.startswith("molgraph_pair_projector")
+ }
+ )
+ log.info("Successfully loaded pretrained ligand graph pair projector weights.")
+ except Exception as e:
+ log.warning(
+ f"Skipping loading of pretrained ligand graph pair projector weights due to: {e}."
+ )
+
+ try:
+ covalent_embed.load_state_dict(
+ {
+ ".".join(k.split(".")[1:]): v
+ for k, v in state_dict.items()
+ if k.startswith("covalent_embed")
+ }
+ )
+ log.info("Successfully loaded pretrained ligand covalent embedding weights.")
+ except Exception as e:
+ log.warning(
+ f"Skipping loading of pretrained ligand covalent embedding weights due to: {e}."
+ )
+ return molgraph_single_projector, molgraph_pair_projector, covalent_embed
diff --git a/forks/FlowDock/flowdock/models/components/modules.py b/forks/FlowDock/flowdock/models/components/modules.py
new file mode 100644
index 00000000..7ee190fd
--- /dev/null
+++ b/forks/FlowDock/flowdock/models/components/modules.py
@@ -0,0 +1,423 @@
+import math
+
+import rootutils
+import torch
+import torch.nn.functional as F
+from beartype.typing import Optional, Tuple, Union
+from openfold.model.primitives import Attention
+from openfold.utils.tensor_utils import permute_final_dims
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock.utils.model_utils import GELUMLP, segment_softmax
+
+
+class MultiHeadAttentionConv(torch.nn.Module):
+ """Native Pytorch implementation."""
+
+ def __init__(
+ self,
+ dim: Union[int, Tuple[int, int]],
+ head_dim: int,
+ edge_dim: int = None,
+ n_heads: int = 1,
+ dropout: float = 0.0,
+ edge_lin: bool = True,
+ **kwargs,
+ ):
+ """Multi-Head Attention Convolution layer."""
+ super().__init__()
+ self.dim = dim
+ self.head_dim = head_dim
+ self.n_heads = n_heads
+ self.dropout = dropout
+ self.edge_dim = edge_dim
+ self.edge_lin = edge_lin
+ self._alpha = None
+
+ if isinstance(dim, int):
+ dim = (dim, dim)
+
+ self.lin_key = torch.nn.Linear(dim[0], n_heads * head_dim, bias=False)
+ self.lin_query = torch.nn.Linear(dim[1], n_heads * head_dim, bias=False)
+ self.lin_value = torch.nn.Linear(dim[0], n_heads * head_dim, bias=False)
+ if edge_lin is True:
+ self.lin_edge = torch.nn.Linear(edge_dim, n_heads, bias=False)
+ else:
+ self.lin_edge = self.register_parameter("lin_edge", None)
+
+ self.reset_parameters()
+
+ def reset_parameters(self):
+ """Reset the parameters of the layer."""
+ self.lin_key.reset_parameters()
+ self.lin_query.reset_parameters()
+ self.lin_value.reset_parameters()
+ if self.edge_lin:
+ self.lin_edge.reset_parameters()
+
+ def forward(
+ self,
+ x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
+ edge_index: torch.Tensor,
+ edge_attr: torch.Tensor = None,
+ return_attention_weights=None,
+ ) -> Union[torch.Tensor, Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]:
+ """Forward pass of the Multi-Head Attention Convolution layer.
+
+ :param x (torch.Tensor or Tuple[torch.Tensor, torch.Tensor]): The input features.
+ :param edge_index (torch.Tensor): The edge index tensor.
+ :param edge_attr (torch.Tensor, optional): The edge attribute tensor.
+ :param return_attention_weights (bool, optional): If set to `True`,
+ will additionally return the tuple
+ `(edge_index, attention_weights)`, holding the computed
+ attention weights for each edge. Default is `None`.
+ :return: The output features or the tuple
+ `(output_features, (edge_index, attention_weights))`.
+ """
+
+ H, C = self.n_heads, self.head_dim
+
+ if isinstance(x, torch.Tensor):
+ x = (x, x)
+
+ query = self.lin_query(x[1]).view(*x[1].shape[:-1], H, C)
+ key = self.lin_key(x[0]).view(*x[0].shape[:-1], H, C)
+ value = self.lin_value(x[0]).view(*x[0].shape[:-1], H, C)
+
+ attended_values = self.message(key, query, value, edge_attr, edge_index)
+ out = self.aggregate(attended_values, edge_index[1], query.shape[0])
+
+ alpha = self._alpha
+ self._alpha = None
+
+ out = out.contiguous().view(*out.shape[:-2], H * C)
+
+ if isinstance(return_attention_weights, bool):
+ assert alpha is not None
+ return out, (edge_index, alpha)
+ else:
+ return out
+
+ def message(
+ self,
+ key: torch.Tensor,
+ query: torch.Tensor,
+ value: torch.Tensor,
+ edge_attr: torch.Tensor,
+ index: torch.Tensor,
+ ) -> torch.Tensor:
+ """Add the relative positional encodings to attention scores.
+
+ :param key (torch.Tensor): The key tensor.
+ :param query (torch.Tensor): The query tensor.
+ :param value (torch.Tensor): The value tensor.
+ :param edge_attr (torch.Tensor): The edge attribute tensor.
+ :param index (torch.Tensor): The edge index tensor.
+ :return: The output tensor.
+ """
+ edge_bias = 0
+ if self.lin_edge is not None:
+ assert edge_attr is not None
+ edge_bias = self.lin_edge(edge_attr)
+
+ _alpha_z = (query[index[1]] * key[index[0]]).sum(dim=-1) / math.sqrt(
+ self.head_dim
+ ) + edge_bias
+ self._alpha = _alpha_z
+ alpha = segment_softmax(_alpha_z, index[1], query.shape[0])
+ alpha = F.dropout(alpha, p=self.dropout, training=self.training)
+
+ out = value[index[0]]
+ out *= alpha.unsqueeze(-1)
+ return out
+
+ def aggregate(
+ self, src: torch.Tensor, dst_idx: torch.Tensor, dst_size: torch.Size
+ ) -> torch.Tensor:
+ """Aggregate the source tensor to the destination tensor.
+
+ :param src (torch.Tensor): The source tensor.
+ :param dst_idx (torch.Tensor): The destination index tensor.
+ :param dst_size (torch.Size): The destination size tensor.
+ :return: The output tensor.
+ """
+ out = torch.zeros(
+ dst_size,
+ *src.shape[1:],
+ dtype=src.dtype,
+ device=src.device,
+ ).index_add_(0, dst_idx, src)
+ return out
+
+
+class TransformerLayer(torch.nn.Module):
+ """A single layer of a transformer model."""
+
+ def __init__(
+ self,
+ node_dim: int,
+ n_heads: int,
+ head_dim: Optional[int] = None,
+ hidden_dim: Optional[int] = None,
+ bidirectional: bool = False,
+ edge_channels: Optional[int] = None,
+ dropout: float = 0.0,
+ edge_update: bool = False,
+ ):
+ """Initialize the transformer layer."""
+ super().__init__()
+ edge_lin = edge_channels is not None
+ self.edge_update = edge_update
+ if head_dim is None:
+ head_dim = node_dim // n_heads
+ self.conv = MultiHeadAttentionConv(
+ node_dim,
+ head_dim,
+ edge_dim=edge_channels,
+ n_heads=n_heads,
+ edge_lin=edge_lin,
+ dropout=dropout,
+ )
+ self.bidirectional = bidirectional
+ self.projector = torch.nn.Linear(head_dim * n_heads, node_dim, bias=False)
+ self.norm = torch.nn.LayerNorm(node_dim)
+ self.mlp = GELUMLP(
+ node_dim,
+ node_dim,
+ n_hidden_feats=hidden_dim,
+ dropout=dropout,
+ zero_init=True,
+ )
+ if edge_update:
+ self.mlpe = GELUMLP(
+ n_heads + edge_channels, edge_channels, dropout=dropout, zero_init=True
+ )
+
+ def forward(
+ self,
+ x_s: torch.Tensor,
+ x_a: torch.Tensor,
+ edge_index: torch.Tensor,
+ edge_attr: Optional[torch.Tensor] = None,
+ ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+ """Forward pass through the transformer layer.
+
+ :param x_s (torch.Tensor): The source node features.
+ :param x_a (torch.Tensor): The target node features.
+ :param edge_index (torch.Tensor): The edge index tensor. :param edge_attr (torch.Tensor,
+ optional): The edge attribute tensor.
+ :return: The output source and target node features.
+ """
+ out_a, (edge_index, alpha) = self.conv(
+ (x_s, x_a),
+ edge_index,
+ edge_attr,
+ return_attention_weights=True,
+ )
+ x_a = x_a + self.projector(out_a)
+ x_a = self.mlp(self.norm(x_a)) + x_a
+ if self.bidirectional:
+ out_s = self.conv((x_a, x_s), (edge_index[1], edge_index[0]), edge_attr)
+ x_s = x_s + self.projector(out_s)
+ x_s = self.mlp(self.norm(x_s)) + x_s
+ if self.edge_update:
+ edge_attr = edge_attr + self.mlpe(torch.cat([alpha, edge_attr], dim=-1))
+ return x_s, x_a, edge_attr
+ else:
+ return x_s, x_a
+
+
+class PointSetAttention(torch.nn.Module):
+ """PointSetAttention module."""
+
+ def __init__(
+ self,
+ fiber_dim: int,
+ heads: int = 8,
+ point_dim: int = 4,
+ edge_dim: Optional[int] = None,
+ edge_update: bool = False,
+ dropout: float = 0.0,
+ ):
+ """Initialize the PointSetAttention module."""
+ super().__init__()
+ self.fiber_dim = fiber_dim
+ self.edge_dim = edge_dim
+ self.heads = heads
+ self.point_dim = point_dim
+ self.dropout = dropout
+ self.edge_update = edge_update
+ self.distance_scaling = 10 # 1 nm
+
+ # num attention contributions
+ num_attn_logits = 2
+
+ self.lin_query = torch.nn.Linear(fiber_dim, point_dim * heads, bias=False)
+ self.lin_key = torch.nn.Linear(fiber_dim, point_dim * heads, bias=False)
+ self.lin_value = torch.nn.Linear(fiber_dim, point_dim * heads, bias=False)
+ if edge_dim is not None:
+ self.lin_edge = torch.nn.Linear(edge_dim, heads, bias=False)
+ if edge_update:
+ self.edge_update_mlp = GELUMLP(heads + edge_dim, edge_dim)
+
+ # qkv projection for scalar attention (normal)
+ self.scalar_attn_logits_scale = (num_attn_logits * point_dim) ** -0.5
+
+ # qkv projection for point attention (coordinate and orientation aware)
+ point_weight_init_value = torch.log(torch.exp(torch.full((heads,), 1.0)) - 1.0)
+ self.point_weights = torch.nn.Parameter(point_weight_init_value)
+
+ self.point_attn_logits_scale = ((num_attn_logits * point_dim) * (9 / 2)) ** -0.5
+ point_weight_init_value = torch.log(torch.exp(torch.full((heads,), 1.0)) - 1.0)
+ self.point_weights = torch.nn.Parameter(point_weight_init_value)
+
+ # combine out - point dim * 4
+ self.to_out = torch.nn.Linear(heads * point_dim, fiber_dim, bias=False)
+
+ def forward(
+ self,
+ x_k: torch.Tensor,
+ x_q: torch.Tensor,
+ edge_index: torch.LongTensor,
+ point_centers_k: torch.Tensor,
+ point_centers_q: torch.Tensor,
+ x_edge: torch.Tensor = None,
+ ):
+ """Forward pass of the PointSetAttention module."""
+ H, P = self.heads, self.point_dim
+
+ q = self.lin_query(x_q)
+ k = self.lin_key(x_k)
+ v = self.lin_value(x_k)
+
+ scalar_q = q[..., 0, :].view(-1, H, P)
+ scalar_k = k[..., 0, :].view(-1, H, P)
+ scalar_v = v[..., 0, :].view(-1, H, P)
+
+ point_q_local = q[..., 1:, :].view(-1, 3, H, P)
+ point_k_local = k[..., 1:, :].view(-1, 3, H, P)
+ point_v_local = v[..., 1:, :].view(-1, 3, H, P)
+
+ point_q = point_q_local + point_centers_q[..., None, None] / self.distance_scaling
+ point_k = point_k_local + point_centers_k[..., None, None] / self.distance_scaling
+ point_v = point_v_local + point_centers_k[..., None, None] / self.distance_scaling
+
+ if self.edge_dim is not None:
+ edge_bias = self.lin_edge(x_edge)
+ else:
+ edge_bias = 0
+
+ attn_logits, attentions = self.compute_attention(
+ scalar_k, scalar_q, point_k, point_q, edge_bias, edge_index
+ )
+ res_scalar = self.aggregate(
+ attentions[:, :, None] * scalar_v[edge_index[0]],
+ edge_index[1],
+ scalar_q.shape[0],
+ )
+ res_points = self.aggregate(
+ attentions[:, None, :, None] * point_v[edge_index[0]],
+ edge_index[1],
+ point_q.shape[0],
+ )
+ res_points_local = res_points - point_centers_q[..., None, None] / self.distance_scaling
+
+ # [N, H, P], [N, 3, H, P] -> [N, 4, C]
+ res = torch.cat([res_scalar.unsqueeze(-3), res_points_local], dim=-3).flatten(-2, -1)
+ out = self.to_out(res) # [N, 4, C]
+ if self.edge_update:
+ edge_out = self.edge_update_mlp(torch.cat([attn_logits, x_edge], dim=-1))
+ return out, edge_out
+ return out
+
+ def compute_attention(self, scalar_k, scalar_q, point_k, point_q, edge_bias, index):
+ """Compute the attention scores."""
+ scalar_q = scalar_q[index[1]]
+ scalar_k = scalar_k[index[0]]
+ point_q = point_q[index[1]]
+ point_k = point_k[index[0]]
+
+ scalar_logits = (scalar_q * scalar_k).sum(dim=-1) * self.scalar_attn_logits_scale
+ point_weights = F.softplus(self.point_weights).unsqueeze(0)
+ point_logits = (
+ torch.square(point_q - point_k).sum(dim=(-3, -1)) * self.point_attn_logits_scale
+ )
+
+ logits = scalar_logits - 1 / 2 * point_logits * point_weights + edge_bias
+ alpha = segment_softmax(logits, index[1], scalar_q.shape[0])
+ alpha = F.dropout(alpha, p=self.dropout, training=self.training)
+ return logits, alpha
+
+ @staticmethod
+ def aggregate(src, dst_idx, dst_size):
+ """Aggregate the source tensor to the destination tensor."""
+ out = torch.zeros(
+ dst_size,
+ *src.shape[1:],
+ dtype=src.dtype,
+ device=src.device,
+ ).index_add_(0, dst_idx, src)
+ return out
+
+
+class BiDirectionalTriangleAttention(torch.nn.Module):
+ """
+ Adapted from https://github.com/aqlaboratory/openfold
+ supports rectangular pair representation tensors
+ """
+
+ def __init__(self, c_in: int, c_hidden: int, no_heads: int, inf: float = 1e9):
+ """Initialize the Bi-Directional Triangle Attention layer."""
+ super().__init__()
+
+ self.c_in = c_in
+ self.c_hidden = c_hidden
+ self.no_heads = no_heads
+ self.inf = inf
+
+ self.linear = torch.nn.Linear(c_in, self.no_heads, bias=False)
+
+ self.mha_1 = Attention(self.c_in, self.c_in, self.c_in, self.c_hidden, self.no_heads)
+ self.mha_2 = Attention(self.c_in, self.c_in, self.c_in, self.c_hidden, self.no_heads)
+ self.layer_norm = torch.nn.LayerNorm(self.c_in)
+
+ def forward(
+ self,
+ x1: torch.Tensor,
+ x2: torch.Tensor,
+ x_pair: torch.Tensor,
+ mask: Optional[torch.Tensor] = None,
+ use_lma: bool = False,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Forward pass of the Bi-Directional Triangle Attention layer."""
+ if mask is None:
+ # [*, I, J, K]
+ mask = x_pair.new_ones(
+ x_pair.shape[:-1],
+ )
+
+ # [*, I, J, C_in]
+ x1 = self.layer_norm(x1)
+ # [*, I, K, C_in]
+ x2 = self.layer_norm(x2)
+
+ # [*, I, 1, J, K]
+ mask_bias = (self.inf * (mask - 1))[..., :, None, :, :]
+
+ # [*, I, H, J, K]
+ triangle_bias = permute_final_dims(self.linear(x_pair), [0, 3, 1, 2])
+
+ biases_J2I = [mask_bias, triangle_bias]
+
+ x1_out = self.mha_1(q_x=x1, kv_x=x2, biases=biases_J2I, use_lma=use_lma)
+ x1 = x1 + x1_out
+
+ # transpose the triangle bias for I->J attention.
+ mask_bias_T_ = mask_bias.transpose(-2, -1).contiguous()
+ triangle_bias_T_ = triangle_bias.transpose(-2, -1).contiguous()
+ biases_I2J = [mask_bias_T_, triangle_bias_T_]
+ x2_out = self.mha_2(q_x=x2, kv_x=x1, biases=biases_I2J, use_lma=use_lma)
+ x2 = x2 + x2_out
+
+ return x1, x2
diff --git a/forks/FlowDock/flowdock/models/components/noise.py b/forks/FlowDock/flowdock/models/components/noise.py
new file mode 100644
index 00000000..bc92212c
--- /dev/null
+++ b/forks/FlowDock/flowdock/models/components/noise.py
@@ -0,0 +1,443 @@
+import numpy as np
+import rootutils
+import torch
+from beartype.typing import Any, Dict, List, Optional, Tuple, Union
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock.models.components.transforms import LatentCoordinateConverter
+from flowdock.utils import RankedLogger
+from flowdock.utils.model_utils import segment_mean
+
+MODEL_BATCH = Dict[str, Any]
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+class DiffusionSDE:
+ """Diffusion SDE class.
+
+ Adapted from: https://github.com/HannesStark/FlowSite
+ """
+
+ def __init__(self, sigma: torch.Tensor, tau_factor: float = 5.0):
+ """Initialize the Diffusion SDE class."""
+ self.lamb = 1 / sigma**2
+ self.tau_factor = tau_factor
+
+ def var(self, t: torch.Tensor) -> torch.Tensor:
+ """Calculate the variance of the diffusion SDE."""
+ return (1 - torch.exp(-self.lamb * t)) / self.lamb
+
+ def max_t(self) -> float:
+ """Calculate the maximum time of the diffusion SDE."""
+ return self.tau_factor / self.lamb
+
+ def mu_factor(self, t: torch.Tensor) -> torch.Tensor:
+ """Calculate the mu factor of the diffusion SDE."""
+ return torch.exp(-self.lamb * t / 2)
+
+
+class HarmonicSDE:
+ """Harmonic SDE class.
+
+ Adapted from: https://github.com/HannesStark/FlowSite
+ """
+
+ def __init__(self, J: Optional[torch.Tensor] = None, diagonalize: bool = True):
+ """Initialize the Harmonic SDE class."""
+ self.l_index = 1
+ self.use_cuda = False
+ if not diagonalize:
+ return
+ if J is not None:
+ self.D, self.P = np.linalg.eigh(J)
+ self.N = self.D.size
+
+ @staticmethod
+ def diagonalize(
+ N,
+ ptr: torch.Tensor,
+ edges: Optional[List[Tuple[int, int]]] = None,
+ antiedges: Optional[List[Tuple[int, int]]] = None,
+ a=1,
+ b=0.3,
+ lamb: Optional[torch.Tensor] = None,
+ device: Optional[Union[str, torch.device]] = None,
+ ):
+ """Diagonalize using the Harmonic SDE."""
+ device = device or ptr.device
+ J = torch.zeros((N, N), device=device)
+ if edges is None:
+ for i, j in zip(np.arange(N - 1), np.arange(1, N)):
+ J[i, i] += a
+ J[j, j] += a
+ J[i, j] = J[j, i] = -a
+ else:
+ for i, j in edges:
+ J[i, i] += a
+ J[j, j] += a
+ J[i, j] = J[j, i] = -a
+ if antiedges is not None:
+ for i, j in antiedges:
+ J[i, i] -= b
+ J[j, j] -= b
+ J[i, j] = J[j, i] = b
+ if edges is not None:
+ J += torch.diag(lamb)
+
+ Ds, Ps = [], []
+ for start, end in zip(ptr[:-1], ptr[1:]):
+ D, P = torch.linalg.eigh(J[start:end, start:end])
+ D_ = D
+ if edges is None:
+ D_inv = 1 / D
+ D_inv[0] = 0
+ D_ = D_inv
+ Ds.append(D_)
+ Ps.append(P)
+ return torch.cat(Ds), torch.block_diag(*Ps)
+
+ def eigens(self, t):
+ """Calculate the eigenvalues of `sigma_t` using the Harmonic SDE."""
+ np_ = torch if self.use_cuda else np
+ D = 1 / self.D * (1 - np_.exp(-t * self.D))
+ t = torch.tensor(t, device="cuda").float() if self.use_cuda else t
+ return np_.where(D != 0, D, t)
+
+ def conditional(self, mask, x2):
+ """Calculate the conditional distribution using the Harmonic SDE."""
+ J_11 = self.J[~mask][:, ~mask]
+ J_12 = self.J[~mask][:, mask]
+ h = -J_12 @ x2
+ mu = np.linalg.inv(J_11) @ h
+ D, P = np.linalg.eigh(J_11)
+ z = np.random.randn(*mu.shape)
+ return (P / D**0.5) @ z + mu
+
+ def A(self, t, invT=False):
+ """Calculate the matrix `A` using the Harmonic SDE."""
+ D = self.eigens(t)
+ A = self.P * (D**0.5)
+ if not invT:
+ return A
+ AinvT = self.P / (D**0.5)
+ return A, AinvT
+
+ def Sigma_inv(self, t):
+ """Calculate the inverse of the covariance matrix `Sigma` using the Harmonic SDE."""
+ D = 1 / self.eigens(t)
+ return (self.P * D) @ self.P.T
+
+ def Sigma(self, t):
+ """Calculate the covariance matrix `Sigma` using the Harmonic SDE."""
+ D = self.eigens(t)
+ return (self.P * D) @ self.P.T
+
+ @property
+ def J(self):
+ """Return the matrix `J`."""
+ return (self.P * self.D) @ self.P.T
+
+ def rmsd(self, t):
+ """Calculate the root mean square deviation using the Harmonic SDE."""
+ l_index = self.l_index
+ D = 1 / self.D * (1 - np.exp(-t * self.D))
+ return np.sqrt(3 * D[l_index:].mean())
+
+ def sample(self, t, x=None, score=False, k=None, center=True, adj=False):
+ """Sample from the Harmonic SDE."""
+ l_index = self.l_index
+ np_ = torch if self.use_cuda else np
+ if x is None:
+ if self.use_cuda:
+ x = torch.zeros((self.N, 3), device="cuda").float()
+ else:
+ x = np.zeros((self.N, 3))
+ if t == 0:
+ return x
+ z = (
+ np.random.randn(self.N, 3)
+ if not self.use_cuda
+ else torch.randn(self.N, 3, device="cuda").float()
+ )
+ D = self.eigens(t)
+ xx = self.P.T @ x
+ if center:
+ z[0] = 0
+ xx[0] = 0
+ if k:
+ z[k + l_index :] = 0
+ xx[k + l_index :] = 0
+
+ out = np_.exp(-t * self.D / 2)[:, None] * xx + np_.sqrt(D)[:, None] * z
+
+ if score:
+ score = -(1 / np_.sqrt(D))[:, None] * z
+ if adj:
+ score = score + self.D[:, None] * out
+ return self.P @ out, self.P @ score
+ return self.P @ out
+
+ def score_norm(self, t, k=None, adj=False):
+ """Calculate the score norm using the Harmonic SDE."""
+ if k == 0:
+ return 0
+ l_index = self.l_index
+ np_ = torch if self.use_cuda else np
+ k = k or self.N - 1
+ D = 1 / self.eigens(t)
+ if adj:
+ D = D * np_.exp(-self.D * t)
+ return (D[l_index : k + l_index].sum() / self.N) ** 0.5
+
+ def inject(self, t, modes):
+ """Inject noise along the given modes using the Harmonic SDE."""
+ z = (
+ np.random.randn(self.N, 3)
+ if not self.use_cuda
+ else torch.randn(self.N, 3, device="cuda").float()
+ )
+ z[~modes] = 0
+ A = self.A(t, invT=False)
+ return A @ z
+
+ def score(self, x0, xt, t):
+ """Calculate the score of the diffusion kernel using the Harmonic SDE."""
+ Sigma_inv = self.Sigma_inv(t)
+ mu_t = (self.P * np.exp(-t * self.D / 2)) @ (self.P.T @ x0)
+ return Sigma_inv @ (mu_t - xt)
+
+ def project(self, X, k, center=False):
+ """Project onto the first `k` nonzero modes using the Harmonic SDE."""
+ l_index = self.l_index
+ D = self.P.T @ X
+ D[k + l_index :] = 0
+ if center:
+ D[0] = 0
+ return self.P @ D
+
+ def unproject(self, X, mask, k, return_Pinv=False):
+ """Find the vector along the first k nonzero modes whose mask is closest to `X`"""
+ l_index = self.l_index
+ PP = self.P[mask, : k + l_index]
+ Pinv = np.linalg.pinv(PP)
+ out = self.P[:, : k + l_index] @ Pinv @ X
+ if return_Pinv:
+ return out, Pinv
+ return out
+
+ def energy(self, X):
+ """Calculate the energy using the Harmonic SDE."""
+ l_index = self.l_index
+ return (self.D[:, None] * (self.P.T @ X) ** 2).sum(-1)[l_index:] / 2
+
+ @property
+ def free_energy(self):
+ """Calculate the free energy using the Harmonic SDE."""
+ l_index = self.l_index
+ return 3 * np.log(self.D[l_index:]).sum() / 2
+
+ def KL_H(self, t):
+ """Calculate the Kullback-Leibler divergence using the Harmonic SDE."""
+ l_index = self.l_index
+ D = self.D[l_index:]
+ return -3 * 0.5 * (np.log(1 - np.exp(-D * t)) + np.exp(-D * t)).sum(0)
+
+
+def sample_gaussian_prior(
+ x0: torch.Tensor,
+ latent_converter: LatentCoordinateConverter,
+ sigma: float,
+ x0_sigma: float = 1e-4,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Sample noise from a Gaussian prior distribution.
+
+ :param x0: ground-truth tensor
+ :param latent_converter: The latent coordinate converter
+ :param sigma: standard deviation of the Gaussian noise
+ :param x0_sigma: standard deviation of the Gaussian noise for the ground-truth tensor
+ :return: tuple of ground-truth and predicted tensors with additive Gaussian prior noise
+ """
+ prior = torch.randn_like(x0)
+ x_int_0 = x0 + prior * x0_sigma # add small Gaussian noise to the ground-truth tensor
+ (
+ x1_ca_lat,
+ x1_cother_lat,
+ x1_lig_lat,
+ ) = torch.split(
+ prior * sigma,
+ [
+ latent_converter._n_res_per_sample,
+ latent_converter._n_cother_per_sample,
+ latent_converter._n_ligha_per_sample,
+ ],
+ dim=1,
+ )
+ x_int_1 = torch.cat(
+ [
+ x1_ca_lat,
+ x1_cother_lat,
+ x1_lig_lat,
+ ],
+ dim=1,
+ )
+ return x_int_0, x_int_1
+
+
+def sample_protein_harmonic_prior(
+ protein_ca_x0: torch.Tensor,
+ protein_cother_x0: torch.Tensor,
+ batch: MODEL_BATCH,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Sample protein noise from a harmonic prior distribution.
+ Adapted from: https://github.com/bjing2016/alphaflow
+
+ Note that this function represents non-Ca atoms as Gaussian noise
+ centered around each harmonically-noised Ca atom.
+
+ :param protein_ca_x0: ground-truth protein Ca-atom tensor
+ :param protein_cother_x0: ground-truth protein other-atom tensor
+ :param batch: A batch dictionary
+ :return: tuple of harmonic protein Ca atom noise and Gaussian protein other atom noise
+ """
+ indexer = batch["indexer"]
+ metadata = batch["metadata"]
+ protein_bid = indexer["gather_idx_a_structid"]
+ protein_num_nodes = protein_ca_x0.size(0) * protein_ca_x0.size(1)
+ ptr = torch.cumsum(torch.bincount(protein_bid), dim=0)
+ ptr = torch.cat((torch.tensor([0], device=protein_bid.device), ptr))
+ try:
+ D_inv, P = HarmonicSDE.diagonalize(
+ protein_num_nodes,
+ ptr,
+ a=3 / (3.8**2),
+ )
+ except Exception as e:
+ log.error(
+ f"Failed to call HarmonicSDE.diagonalize() for protein(s) {metadata['sample_ID_per_sample']} due to: {e}"
+ )
+ raise e
+ noise = torch.randn((protein_num_nodes, 3), device=protein_ca_x0.device)
+ harmonic_ca_noise = P @ (torch.sqrt(D_inv)[:, None] * noise)
+ gaussian_cother_noise = (
+ torch.randn_like(protein_cother_x0.flatten(0, 1))
+ + harmonic_ca_noise[indexer["gather_idx_a_cotherid"]]
+ )
+ return (
+ harmonic_ca_noise.view(protein_ca_x0.size()).contiguous(),
+ gaussian_cother_noise.view(protein_cother_x0.size()).contiguous(),
+ )
+
+
+def sample_ligand_harmonic_prior(
+ lig_x0: torch.Tensor, protein_ca_x0: torch.Tensor, batch: MODEL_BATCH, sigma: float = 1.0
+) -> torch.Tensor:
+ """
+ Sample ligand noise from a harmonic prior distribution.
+ Adapted from: https://github.com/HannesStark/FlowSite
+
+ :param lig_x0: ground-truth ligand tensor
+ :param protein_x0: ground-truth protein Ca-atom tensor
+ :param batch: A batch dictionary
+ :param sigma: standard deviation of the harmonic noise
+ :return: tensor of harmonic noise
+ """
+ indexer = batch["indexer"]
+ metadata = batch["metadata"]
+ lig_num_nodes = lig_x0.size(0) * lig_x0.size(1)
+ num_molid_per_sample = max(metadata["num_molid_per_sample"])
+ # NOTE: here, we distinguish each ligand chain in a complex for harmonic chain sampling
+ lig_bid = indexer["gather_idx_i_molid"]
+ protein_sigma = (
+ segment_mean(
+ torch.square(protein_ca_x0).flatten(0, 1),
+ indexer["gather_idx_a_structid"],
+ metadata["num_structid"],
+ ).mean(dim=-1)
+ ** 0.5
+ ).repeat_interleave(num_molid_per_sample)
+ sde = DiffusionSDE(protein_sigma * sigma)
+ edges = torch.stack(
+ (
+ indexer["gather_idx_ij_i"],
+ indexer["gather_idx_ij_j"],
+ )
+ )
+ edges = edges[:, edges[0] < edges[1]] # de-duplicate edges
+ ptr = torch.cumsum(torch.bincount(lig_bid), dim=0)
+ ptr = torch.cat((torch.tensor([0], device=lig_bid.device), ptr))
+ try:
+ D, P = HarmonicSDE.diagonalize(
+ lig_num_nodes,
+ ptr,
+ edges=edges.T,
+ lamb=sde.lamb[lig_bid],
+ )
+ except Exception as e:
+ log.error(
+ f"Failed to call HarmonicSDE.diagonalize() for ligand(s) {metadata['sample_ID_per_sample']} due to: {e}"
+ )
+ raise e
+ noise = torch.randn((lig_num_nodes, 3), device=lig_x0.device)
+ prior = P @ (noise / torch.sqrt(D)[:, None])
+ return prior.view(lig_x0.size()).contiguous()
+
+
+def sample_complex_harmonic_prior(
+ x0: torch.Tensor,
+ latent_converter: LatentCoordinateConverter,
+ batch: MODEL_BATCH,
+ x0_sigma: float = 1e-4,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Sample protein-ligand complex noise from a harmonic prior distribution.
+ From: https://github.com/bjing2016/alphaflow
+
+ :param x0: ground-truth tensor
+ :param latent_converter: The latent coordinate converter
+ :param batch: A batch dictionary
+ :param x0_sigma: standard deviation of the Gaussian noise for the ground-truth tensor
+ :return: tuple of ground-truth and predicted tensors with additive Gaussian and harmonic prior noise, respectively
+ """
+ ca_lat, cother_lat, lig_lat = x0.split(
+ [
+ latent_converter._n_res_per_sample,
+ latent_converter._n_cother_per_sample,
+ latent_converter._n_ligha_per_sample,
+ ],
+ dim=1,
+ )
+ harmonic_ca_lat, gaussian_cother_lat = sample_protein_harmonic_prior(
+ ca_lat,
+ cother_lat,
+ batch,
+ )
+ harmonic_lig_lat = sample_ligand_harmonic_prior(lig_lat, harmonic_ca_lat, batch)
+ x1 = torch.cat(
+ [
+ # NOTE: the following normalization steps assume that `self.latent_model == "default"`
+ harmonic_ca_lat / latent_converter.ca_scale,
+ gaussian_cother_lat / latent_converter.other_scale,
+ harmonic_lig_lat / latent_converter.other_scale,
+ ],
+ dim=1,
+ )
+ gaussian_prior = torch.randn_like(x0)
+ return x0 + gaussian_prior * x0_sigma, x1
+
+
+def sample_esmfold_prior(
+ x0: torch.Tensor, x1: torch.Tensor, sigma: float, x0_sigma: float = 1e-4
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Sample noise from an ESMFold prior distribution.
+
+ :param x0: ground-truth tensor
+ :param x1: predicted tensor
+ :param sigma: standard deviation of the ESMFold prior's additive Gaussian noise
+ :param x0_sigma: standard deviation of the Gaussian noise for the ground-truth tensor
+ :return: tuple of ground-truth and predicted tensors with additive Gaussian prior noise
+ """
+ prior_noise = torch.randn_like(x0)
+ return x0 + prior_noise * x0_sigma, x1 + prior_noise * sigma
diff --git a/forks/FlowDock/flowdock/models/components/transforms.py b/forks/FlowDock/flowdock/models/components/transforms.py
new file mode 100644
index 00000000..843f7d96
--- /dev/null
+++ b/forks/FlowDock/flowdock/models/components/transforms.py
@@ -0,0 +1,241 @@
+# Adapted from: https://github.com/zrqiao/NeuralPLexer
+
+import rootutils
+import torch
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock.utils.model_utils import segment_mean
+
+
+class LatentCoordinateConverter:
+ """Transform the batched feature dict to latent coordinate arrays."""
+
+ def __init__(self, config, prot_atom37_namemap, lig_namemap):
+ """Initialize the converter."""
+ super().__init__()
+ self.config = config
+ self.prot_namemap = prot_atom37_namemap
+ self.lig_namemap = lig_namemap
+ self.cached_noise = None
+ self._last_pred_ca_trace = None
+
+ @staticmethod
+ def nested_get(dic, keys):
+ """Get the value in the nested dictionary."""
+ for key in keys:
+ dic = dic[key]
+ return dic
+
+ @staticmethod
+ def nested_set(dic, keys, value):
+ """Set the value in the nested dictionary."""
+ for key in keys[:-1]:
+ dic = dic.setdefault(key, {})
+ dic[keys[-1]] = value
+
+ def to_latent(self, batch):
+ """Convert the batched feature dict to latent coordinates."""
+ return None
+
+ def assign_to_batch(self, batch, x_int):
+ """Assign the latent coordinates to the batched feature dict."""
+ return None
+
+
+class DefaultPLCoordinateConverter(LatentCoordinateConverter):
+ """Minimal conversion, using internal coords for sidechains and global coords for others."""
+
+ def __init__(self, config, prot_atom37_namemap, lig_namemap):
+ """Initialize the converter."""
+ super().__init__(config, prot_atom37_namemap, lig_namemap)
+ # Scale parameters in Angstrom
+ self.ca_scale = config.global_max_sigma
+ self.other_scale = config.internal_max_sigma
+
+ def to_latent(self, batch: dict):
+ """Convert the batched feature dict to latent coordinates."""
+ indexer = batch["indexer"]
+ metadata = batch["metadata"]
+ self._batch_size = metadata["num_structid"]
+ atom37_mask = batch["features"]["res_atom_mask"].bool()
+ self._cother_mask = atom37_mask.clone()
+ self._cother_mask[:, 1] = False
+ atom37_coords = self.nested_get(batch, self.prot_namemap[0])
+ try:
+ apo_available = True
+ apo_atom37_coords = self.nested_get(
+ batch, self.prot_namemap[0][:-1] + ("apo_" + self.prot_namemap[0][-1],)
+ )
+ except KeyError:
+ apo_available = False
+ apo_atom37_coords = torch.zeros_like(atom37_coords)
+ ca_atom_centroid_coords = segment_mean(
+ # NOTE: in contrast to NeuralPLexer, we center all coordinates at the origin using the Ca atom centroids
+ atom37_coords[:, 1],
+ indexer["gather_idx_a_structid"],
+ self._batch_size,
+ )
+ if apo_available:
+ apo_ca_atom_centroid_coords = segment_mean(
+ apo_atom37_coords[:, 1],
+ indexer["gather_idx_a_structid"],
+ self._batch_size,
+ )
+ else:
+ apo_ca_atom_centroid_coords = torch.zeros_like(ca_atom_centroid_coords)
+ ca_coords_glob = (
+ (atom37_coords[:, 1] - ca_atom_centroid_coords[indexer["gather_idx_a_structid"]])
+ .contiguous()
+ .view(self._batch_size, -1, 3)
+ )
+ if apo_available:
+ apo_ca_coords_glob = (
+ (
+ apo_atom37_coords[:, 1]
+ - apo_ca_atom_centroid_coords[indexer["gather_idx_a_structid"]]
+ )
+ .contiguous()
+ .view(self._batch_size, -1, 3)
+ )
+ else:
+ apo_ca_coords_glob = torch.zeros_like(ca_coords_glob)
+ cother_coords_int = (
+ (atom37_coords - ca_atom_centroid_coords[indexer["gather_idx_a_structid"], None])[
+ self._cother_mask
+ ]
+ .contiguous()
+ .view(self._batch_size, -1, 3)
+ )
+ if apo_available:
+ apo_cother_coords_int = (
+ (
+ apo_atom37_coords
+ - apo_ca_atom_centroid_coords[indexer["gather_idx_a_structid"], None]
+ )[self._cother_mask]
+ .contiguous()
+ .view(self._batch_size, -1, 3)
+ )
+ else:
+ apo_cother_coords_int = torch.zeros_like(cother_coords_int)
+ self._n_res_per_sample = ca_coords_glob.shape[1]
+ self._n_cother_per_sample = cother_coords_int.shape[1]
+ if batch["misc"]["protein_only"]:
+ self._n_ligha_per_sample = 0
+ x_int = torch.cat(
+ [
+ ca_coords_glob / self.ca_scale,
+ apo_ca_coords_glob / self.ca_scale,
+ cother_coords_int / self.other_scale,
+ apo_cother_coords_int / self.other_scale,
+ ],
+ dim=1,
+ )
+ return x_int
+ lig_ha_coords = self.nested_get(batch, self.lig_namemap[0])
+ lig_ha_coords_int = (
+ lig_ha_coords - ca_atom_centroid_coords[indexer["gather_idx_i_structid"]]
+ )
+ lig_ha_coords_int = lig_ha_coords_int.contiguous().view(self._batch_size, -1, 3)
+ ca_atom_centroid_coords = ca_atom_centroid_coords.contiguous().view(
+ self._batch_size, -1, 3
+ )
+ apo_ca_atom_centroid_coords = apo_ca_atom_centroid_coords.contiguous().view(
+ self._batch_size, -1, 3
+ )
+ x_int = torch.cat(
+ [
+ ca_coords_glob / self.ca_scale,
+ apo_ca_coords_glob / self.ca_scale,
+ cother_coords_int / self.other_scale,
+ apo_cother_coords_int / self.other_scale,
+ ca_atom_centroid_coords / self.ca_scale,
+ apo_ca_atom_centroid_coords / self.ca_scale,
+ lig_ha_coords_int / self.other_scale,
+ ],
+ dim=1,
+ )
+ # NOTE: since we use the Ca atom centroids for centralization, we have only one molid per sample
+ self._n_molid_per_sample = ca_atom_centroid_coords.shape[1]
+ self._n_ligha_per_sample = lig_ha_coords_int.shape[1]
+ return x_int
+
+ def assign_to_batch(self, batch: dict, x_lat: torch.Tensor):
+ """Assign the latent coordinates to the batched feature dict."""
+ indexer = batch["indexer"]
+ new_atom37_coords = x_lat.new_zeros(self._batch_size * self._n_res_per_sample, 37, 3)
+ apo_new_atom37_coords = x_lat.new_zeros(self._batch_size * self._n_res_per_sample, 37, 3)
+ if batch["misc"]["protein_only"]:
+ ca_lat, apo_ca_lat, cother_lat, apo_cother_lat = torch.split(
+ x_lat,
+ [
+ self._n_res_per_sample,
+ self._n_res_per_sample,
+ self._n_cother_per_sample,
+ self._n_cother_per_sample,
+ ],
+ dim=1,
+ )
+ else:
+ (
+ ca_lat,
+ apo_ca_lat,
+ cother_lat,
+ apo_cother_lat,
+ ca_cent_lat,
+ _,
+ lig_lat,
+ ) = torch.split(
+ x_lat,
+ [
+ self._n_res_per_sample,
+ self._n_res_per_sample,
+ self._n_cother_per_sample,
+ self._n_cother_per_sample,
+ self._n_molid_per_sample,
+ self._n_molid_per_sample,
+ self._n_ligha_per_sample,
+ ],
+ dim=1,
+ )
+ new_ca_glob = (ca_lat * self.ca_scale).contiguous().flatten(0, 1)
+ apo_new_ca_glob = (apo_ca_lat * self.ca_scale).contiguous().flatten(0, 1)
+ new_atom37_coords[self._cother_mask] = (
+ (cother_lat * self.other_scale).contiguous().flatten(0, 1)
+ )
+ apo_new_atom37_coords[self._cother_mask] = (
+ (apo_cother_lat * self.other_scale).contiguous().flatten(0, 1)
+ )
+ new_atom37_coords = new_atom37_coords
+ apo_new_atom37_coords = apo_new_atom37_coords
+ new_atom37_coords[~self._cother_mask] = 0
+ apo_new_atom37_coords[~self._cother_mask] = 0
+ new_atom37_coords[:, 1] = new_ca_glob
+ apo_new_atom37_coords[:, 1] = apo_new_ca_glob
+ self.nested_set(batch, self.prot_namemap[1], new_atom37_coords)
+ self.nested_set(
+ batch,
+ self.prot_namemap[1][:-1] + ("apo_" + self.prot_namemap[1][-1],),
+ apo_new_atom37_coords,
+ )
+ if batch["misc"]["protein_only"]:
+ self.nested_set(batch, self.lig_namemap[1], None)
+ self.empty_cache()
+ return batch
+ new_ligha_coords_int = (lig_lat * self.other_scale).contiguous().flatten(0, 1)
+ new_ligha_coords_cent = (ca_cent_lat * self.ca_scale).contiguous().flatten(0, 1)
+ new_ligha_coords = (
+ new_ligha_coords_int + new_ligha_coords_cent[indexer["gather_idx_i_structid"]]
+ )
+ self.nested_set(batch, self.lig_namemap[1], new_ligha_coords)
+ self.empty_cache()
+ return batch
+
+ def empty_cache(self):
+ """Empty the cached variables."""
+ self._batch_size = None
+ self._cother_mask = None
+ self._n_res_per_sample = None
+ self._n_cother_per_sample = None
+ self._n_ligha_per_sample = None
+ self._n_molid_per_sample = None
diff --git a/forks/FlowDock/flowdock/models/flowdock_fm_module.py b/forks/FlowDock/flowdock/models/flowdock_fm_module.py
new file mode 100644
index 00000000..4ae9d200
--- /dev/null
+++ b/forks/FlowDock/flowdock/models/flowdock_fm_module.py
@@ -0,0 +1,936 @@
+import os
+
+import esm
+import numpy as np
+import rootutils
+import torch
+from beartype.typing import Any, Dict, Literal, Optional, Union
+from lightning import LightningModule
+from omegaconf import DictConfig
+from torchmetrics.functional.regression import (
+ mean_absolute_error,
+ mean_squared_error,
+ pearson_corrcoef,
+ spearman_corrcoef,
+)
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+from flowdock.models.components.losses import (
+ eval_auxiliary_estimation_losses,
+ eval_structure_prediction_losses,
+)
+from flowdock.utils import RankedLogger
+from flowdock.utils.data_utils import pdb_filepath_to_protein, prepare_batch
+from flowdock.utils.model_utils import extract_esm_embeddings
+from flowdock.utils.sampling_utils import multi_pose_sampling
+from flowdock.utils.visualization_utils import (
+ construct_prot_lig_pairs,
+ write_prot_lig_pairs_to_pdb_file,
+)
+
+MODEL_BATCH = Dict[str, Any]
+MODEL_STAGE = Literal["train", "val", "test", "predict"]
+LOSS_MODES_LIST = [
+ "structure_prediction",
+ "auxiliary_estimation",
+ "auxiliary_estimation_without_structure_prediction",
+]
+LOSS_MODES = Literal[
+ "structure_prediction",
+ "auxiliary_estimation",
+ "auxiliary_estimation_without_structure_prediction",
+]
+AUX_ESTIMATION_STAGES = ["train", "val", "test"]
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+class FlowDockFMLitModule(LightningModule):
+ """A `LightningModule` for geometric flow matching (FM) with FlowDock.
+
+ A `LightningModule` implements 8 key methods:
+
+ ```python
+ def __init__(self):
+ # Define initialization code here.
+
+ def setup(self, stage):
+ # Things to setup before each stage, 'fit', 'validate', 'test', 'predict'.
+ # This hook is called on every process when using DDP.
+
+ def training_step(self, batch, batch_idx):
+ # The complete training step.
+
+ def validation_step(self, batch, batch_idx):
+ # The complete validation step.
+
+ def test_step(self, batch, batch_idx):
+ # The complete test step.
+
+ def predict_step(self, batch, batch_idx):
+ # The complete predict step.
+
+ def configure_optimizers(self):
+ # Define and configure optimizers and LR schedulers.
+ ```
+
+ Docs:
+ https://lightning.ai/docs/pytorch/latest/common/lightning_module.html
+ """
+
+ def __init__(
+ self,
+ net: torch.nn.Module,
+ optimizer: torch.optim.Optimizer,
+ scheduler: torch.optim.lr_scheduler,
+ compile: bool,
+ cfg: DictConfig,
+ **kwargs: Dict[str, Any],
+ ):
+ """Initialize a `FlowDockFMLitModule`.
+
+ :param net: The model to train.
+ :param optimizer: The optimizer to use for training.
+ :param scheduler: The learning rate scheduler to use for training.
+ :param compile: Whether to compile the model before training.
+ :param cfg: The model configuration.
+ :param kwargs: Additional keyword arguments.
+ """
+ super().__init__()
+
+ # the model along with its hyperparameters
+ self.net = net(cfg)
+
+ # this line allows to access init params with 'self.hparams' attribute
+ # also ensures init params will be stored in ckpt
+ self.save_hyperparameters(logger=False, ignore=["net"])
+
+ # for validating input arguments
+ if self.hparams.cfg.task.loss_mode not in LOSS_MODES_LIST:
+ raise ValueError(
+ f"Invalid loss mode: {self.hparams.cfg.task.loss_mode}. Must be one of {LOSS_MODES}."
+ )
+
+ # for inspecting the model's outputs during validation and testing
+ (
+ self.training_step_outputs,
+ self.validation_step_outputs,
+ self.test_step_outputs,
+ self.predict_step_outputs,
+ ) = (
+ [],
+ [],
+ [],
+ [],
+ )
+
+ def forward(
+ self,
+ batch: MODEL_BATCH,
+ iter_id: Union[int, str] = 0,
+ observed_block_contacts: Optional[torch.Tensor] = None,
+ contact_prediction: bool = True,
+ infer_geometry_prior: bool = False,
+ score: bool = False,
+ affinity: bool = True,
+ use_template: bool = False,
+ **kwargs: Dict[str, Any],
+ ) -> MODEL_BATCH:
+ """Perform a forward pass through the model.
+
+ :param batch: A batch dictionary.
+ :param iter_id: The current iteration ID.
+ :param observed_block_contacts: Observed block contacts.
+ :param contact_prediction: Whether to predict contacts.
+ :param infer_geometry_prior: Whether to predict using a geometry prior.
+ :param score: Whether to predict a denoised complex structure.
+ :param affinity: Whether to predict ligand binding affinity.
+ :param use_template: Whether to use a template protein structure.
+ :param kwargs: Additional keyword arguments.
+ :return: Batch dictionary with outputs.
+ """
+ return self.net(
+ batch,
+ iter_id=iter_id,
+ observed_block_contacts=observed_block_contacts,
+ contact_prediction=contact_prediction,
+ infer_geometry_prior=infer_geometry_prior,
+ score=score,
+ affinity=affinity,
+ use_template=use_template,
+ training=self.training,
+ **kwargs,
+ )
+
+ def model_step(
+ self,
+ batch: MODEL_BATCH,
+ batch_idx: int,
+ stage: MODEL_STAGE,
+ loss_mode: Optional[LOSS_MODES] = None,
+ ) -> MODEL_BATCH:
+ """Perform a single model step on a batch of data.
+
+ :param batch: A batch dictionary.
+ :param batch_idx: The index of the current batch.
+ :param stage: The current model stage (i.e., `train`, `val`, `test`, or `predict`).
+ :param loss_mode: The loss mode to use for training.
+ :return: Batch dictionary with losses.
+ """
+ prepare_batch(batch)
+ predicting_aux_outputs = (
+ self.hparams.cfg.confidence.enabled or self.hparams.cfg.affinity.enabled
+ )
+ is_aux_loss_stage = stage in AUX_ESTIMATION_STAGES
+ is_aux_batch = batch_idx % self.hparams.cfg.task.aux_batch_freq == 0
+ struct_pred_loss_mode_requested = (
+ loss_mode is not None and loss_mode == "structure_prediction"
+ )
+ should_eval_aux_loss = (
+ predicting_aux_outputs
+ and is_aux_loss_stage
+ and is_aux_batch
+ and not struct_pred_loss_mode_requested
+ and (
+ not self.hparams.cfg.task.freeze_confidence
+ or (
+ not self.hparams.cfg.task.freeze_affinity
+ and batch["features"]["affinity"].any().item()
+ )
+ )
+ )
+ eval_aux_loss_mode_requested = (
+ predicting_aux_outputs
+ and loss_mode is not None
+ and "auxiliary_estimation" in loss_mode
+ )
+ if should_eval_aux_loss or eval_aux_loss_mode_requested:
+ return eval_auxiliary_estimation_losses(
+ self, batch, stage, loss_mode, training=self.training
+ )
+ loss_fn = eval_structure_prediction_losses
+ return loss_fn(self, batch, batch_idx, self.device, stage, t_1=1.0)
+
+ def on_train_start(self):
+ """Lightning hook that is called when training begins."""
+ pass
+
+ def training_step(self, batch: MODEL_BATCH, batch_idx: int) -> torch.Tensor:
+ """Perform a single training step on a batch of data from the training set.
+
+ :param batch: A batch dictionary.
+ :param batch_idx: The index of the current batch.
+ :return: A tensor of losses between model predictions and targets.
+ """
+ if self.hparams.cfg.task.overfitting_example_name is not None and not all(
+ name == self.hparams.cfg.task.overfitting_example_name
+ for name in batch["metadata"]["sample_ID_per_sample"]
+ ):
+ return None
+
+ try:
+ batch = self.model_step(batch, batch_idx, "train")
+ except Exception as e:
+ log.error(
+ f"Failed to perform training step for batch index {batch_idx} due to: {e}. Skipping example."
+ )
+ return None
+
+ if self.hparams.cfg.affinity.enabled and "affinity_logits" in batch["outputs"]:
+ training_outputs = {
+ "affinity_logits": batch["outputs"]["affinity_logits"],
+ "affinity": batch["features"]["affinity"],
+ }
+ self.training_step_outputs.append(training_outputs)
+
+ # return loss or backpropagation will fail
+ return batch["outputs"]["loss"]
+
+ def on_train_epoch_end(self):
+ """Lightning hook that is called when a training epoch ends."""
+ if self.hparams.cfg.affinity.enabled and any(
+ "affinity_logits" in output for output in self.training_step_outputs
+ ):
+ affinity_logits = torch.cat(
+ [
+ output["affinity_logits"]
+ for output in self.training_step_outputs
+ if "affinity_logits" in output
+ ]
+ )
+ affinity = torch.cat(
+ [
+ output["affinity"]
+ for output in self.training_step_outputs
+ if "affinity_logits" in output
+ ]
+ )
+ affinity_logits = affinity_logits[~affinity.isnan()]
+ affinity = affinity[~affinity.isnan()]
+ if affinity.numel() > 1:
+ # NOTE: there must be at least two affinity batches to properly score the affinity predictions
+ aff_rmse = torch.sqrt(mean_squared_error(affinity_logits, affinity))
+ aff_mae = mean_absolute_error(affinity_logits, affinity)
+ aff_pearson = pearson_corrcoef(affinity_logits, affinity)
+ aff_spearman = spearman_corrcoef(affinity_logits, affinity)
+ self.log(
+ "train_affinity/RMSE",
+ aff_rmse.detach(),
+ on_epoch=True,
+ batch_size=len(affinity),
+ sync_dist=False,
+ )
+ self.log(
+ "train_affinity/MAE",
+ aff_mae.detach(),
+ on_epoch=True,
+ batch_size=len(affinity),
+ sync_dist=False,
+ )
+ self.log(
+ "train_affinity/Pearson",
+ aff_pearson.detach(),
+ on_epoch=True,
+ batch_size=len(affinity),
+ sync_dist=False,
+ )
+ self.log(
+ "train_affinity/Spearman",
+ aff_spearman.detach(),
+ on_epoch=True,
+ batch_size=len(affinity),
+ sync_dist=False,
+ )
+ self.training_step_outputs.clear() # free memory
+
+ def on_validation_start(self):
+ """Lightning hook that is called when validation begins."""
+ # create a directory to store model outputs from each validation epoch
+ os.makedirs(
+ os.path.join(self.trainer.default_root_dir, "validation_epoch_outputs"), exist_ok=True
+ )
+
+ def validation_step(self, batch: MODEL_BATCH, batch_idx: int, dataloader_idx: int = 0):
+ """Perform a single validation step on a batch of data from the validation set.
+
+ :param batch: A batch dictionary.
+ :param batch_idx: The index of the current batch.
+ :param dataloader_idx: The index of the current dataloader.
+ """
+ if self.hparams.cfg.task.overfitting_example_name is not None and not all(
+ name == self.hparams.cfg.task.overfitting_example_name
+ for name in batch["metadata"]["sample_ID_per_sample"]
+ ):
+ return None
+
+ try:
+ prepare_batch(batch)
+ sampling_stats = self.net.sample_pl_complex_structures(
+ batch,
+ sampler="VDODE",
+ sampler_eta=1.0,
+ num_steps=10,
+ start_time=1.0,
+ exact_prior=False,
+ return_all_states=True,
+ eval_input_protein=True,
+ )
+ all_frames = sampling_stats["all_frames"]
+ del sampling_stats["all_frames"]
+ for metric_name in sampling_stats.keys():
+ log_stat = sampling_stats[metric_name].mean().detach()
+ batch_size = sampling_stats[metric_name].shape[0]
+ self.log(
+ f"val_sampling/{metric_name}",
+ log_stat,
+ on_step=True,
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ sampling_stats = self.net.sample_pl_complex_structures(
+ batch,
+ sampler="VDODE",
+ sampler_eta=1.0,
+ num_steps=10,
+ start_time=1.0,
+ exact_prior=False,
+ use_template=False,
+ )
+ for metric_name in sampling_stats.keys():
+ log_stat = sampling_stats[metric_name].mean().detach()
+ batch_size = sampling_stats[metric_name].shape[0]
+ self.log(
+ f"val_sampling_notemplate/{metric_name}",
+ log_stat,
+ on_step=True,
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ sampling_stats = self.net.sample_pl_complex_structures(
+ batch,
+ sampler="VDODE",
+ sampler_eta=1.0,
+ num_steps=10,
+ start_time=1.0,
+ return_summary_stats=True,
+ exact_prior=True,
+ )
+ for metric_name in sampling_stats.keys():
+ log_stat = sampling_stats[metric_name].mean().detach()
+ batch_size = sampling_stats[metric_name].shape[0]
+ self.log(
+ f"val_sampling_trueprior/{metric_name}",
+ log_stat,
+ on_step=True,
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ batch = self.model_step(batch, batch_idx, "val")
+ except Exception as e:
+ log.error(
+ f"Failed to perform validation step for batch index {batch_idx} of dataloader {dataloader_idx} due to: {e}. Skipping example."
+ )
+ return None
+
+ # store model outputs for inspection
+ validation_outputs = {}
+ if self.hparams.cfg.task.visualize_generated_samples:
+ validation_outputs = {
+ "name": batch["metadata"]["sample_ID_per_sample"],
+ "batch_size": batch["metadata"]["num_structid"],
+ "aatype": batch["features"]["res_type"].long().cpu().numpy(),
+ "res_atom_mask": batch["features"]["res_atom_mask"].cpu().numpy(),
+ "protein_coordinates_list": [
+ frame["receptor_padded"].cpu().numpy() for frame in all_frames
+ ],
+ "ligand_coordinates_list": [
+ frame["ligands"].cpu().numpy() for frame in all_frames
+ ],
+ "ligand_mol": batch["metadata"]["mol_per_sample"],
+ "protein_batch_indexer": batch["indexer"]["gather_idx_a_structid"].cpu().numpy(),
+ "ligand_batch_indexer": batch["indexer"]["gather_idx_i_structid"].cpu().numpy(),
+ "gt_protein_coordinates": batch["features"]["res_atom_positions"].cpu().numpy(),
+ "gt_ligand_coordinates": batch["features"]["sdf_coordinates"].cpu().numpy(),
+ "dataloader_idx": dataloader_idx,
+ }
+ if self.hparams.cfg.affinity.enabled and "affinity_logits" in batch["outputs"]:
+ validation_outputs.update(
+ {
+ "affinity_logits": batch["outputs"]["affinity_logits"],
+ "affinity": batch["features"]["affinity"],
+ "dataloader_idx": dataloader_idx,
+ }
+ )
+ if validation_outputs:
+ self.validation_step_outputs.append(validation_outputs)
+
+ def on_validation_epoch_end(self):
+ "Lightning hook that is called when a validation epoch ends."
+ if self.hparams.cfg.task.visualize_generated_samples:
+ for i, outputs in enumerate(self.validation_step_outputs):
+ for batch_index in range(outputs["batch_size"]):
+ prot_lig_pairs = construct_prot_lig_pairs(outputs, batch_index)
+ write_prot_lig_pairs_to_pdb_file(
+ prot_lig_pairs,
+ os.path.join(
+ self.trainer.default_root_dir,
+ "validation_epoch_outputs",
+ f"{outputs['name'][batch_index]}_validation_epoch_{self.current_epoch}_global_step_{self.global_step}_output_{i}_batch_{batch_index}_dataloader_{outputs['dataloader_idx']}.pdb",
+ ),
+ )
+ if self.hparams.cfg.affinity.enabled and any(
+ "affinity_logits" in output for output in self.validation_step_outputs
+ ):
+ affinity_logits = torch.cat(
+ [
+ output["affinity_logits"]
+ for output in self.validation_step_outputs
+ if "affinity_logits" in output
+ ]
+ )
+ affinity = torch.cat(
+ [
+ output["affinity"]
+ for output in self.validation_step_outputs
+ if "affinity_logits" in output
+ ]
+ )
+ affinity_logits = affinity_logits[~affinity.isnan()]
+ affinity = affinity[~affinity.isnan()]
+ if affinity.numel() > 1:
+ # NOTE: there must be at least two affinity batches to properly score the affinity predictions
+ aff_rmse = torch.sqrt(mean_squared_error(affinity_logits, affinity))
+ aff_mae = mean_absolute_error(affinity_logits, affinity)
+ aff_pearson = pearson_corrcoef(affinity_logits, affinity)
+ aff_spearman = spearman_corrcoef(affinity_logits, affinity)
+ self.log(
+ "val_affinity/RMSE",
+ aff_rmse.detach(),
+ on_epoch=True,
+ batch_size=len(affinity),
+ sync_dist=True,
+ )
+ self.log(
+ "val_affinity/MAE",
+ aff_mae.detach(),
+ on_epoch=True,
+ batch_size=len(affinity),
+ sync_dist=True,
+ )
+ self.log(
+ "val_affinity/Pearson",
+ aff_pearson.detach(),
+ on_epoch=True,
+ batch_size=len(affinity),
+ sync_dist=True,
+ )
+ self.log(
+ "val_affinity/Spearman",
+ aff_spearman.detach(),
+ on_epoch=True,
+ batch_size=len(affinity),
+ sync_dist=True,
+ )
+ self.validation_step_outputs.clear() # free memory
+
+ def on_test_start(self):
+ """Lightning hook that is called when testing begins."""
+ # create a directory to store model outputs from each test epoch
+ os.makedirs(
+ os.path.join(self.trainer.default_root_dir, "test_epoch_outputs"), exist_ok=True
+ )
+
+ def test_step(self, batch: MODEL_BATCH, batch_idx: int, dataloader_idx: int = 0):
+ """Perform a single test step on a batch of data from the test set.
+
+ :param batch: A batch dictionary.
+ :param batch_idx: The index of the current batch.
+ :param dataloader_idx: The index of the current dataloader.
+ """
+ if self.hparams.cfg.task.overfitting_example_name is not None and not all(
+ name == self.hparams.cfg.task.overfitting_example_name
+ for name in batch["metadata"]["sample_ID_per_sample"]
+ ):
+ return None
+
+ try:
+ prepare_batch(batch)
+ if self.hparams.cfg.task.eval_structure_prediction:
+ sampling_stats = self.net.sample_pl_complex_structures(
+ batch,
+ sampler=self.hparams.cfg.task.sampler,
+ sampler_eta=self.hparams.cfg.task.sampler_eta,
+ num_steps=self.hparams.cfg.task.num_steps,
+ start_time=self.hparams.cfg.task.start_time,
+ exact_prior=False,
+ return_all_states=True,
+ eval_input_protein=True,
+ )
+ all_frames = sampling_stats["all_frames"]
+ del sampling_stats["all_frames"]
+ for metric_name in sampling_stats.keys():
+ log_stat = sampling_stats[metric_name].mean().detach()
+ batch_size = sampling_stats[metric_name].shape[0]
+ self.log(
+ f"test_sampling/{metric_name}",
+ log_stat,
+ on_step=True,
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ sampling_stats = self.net.sample_pl_complex_structures(
+ batch,
+ sampler=self.hparams.cfg.task.sampler,
+ sampler_eta=self.hparams.cfg.task.sampler_eta,
+ num_steps=self.hparams.cfg.task.num_steps,
+ start_time=self.hparams.cfg.task.start_time,
+ exact_prior=False,
+ use_template=False,
+ )
+ for metric_name in sampling_stats.keys():
+ log_stat = sampling_stats[metric_name].mean().detach()
+ batch_size = sampling_stats[metric_name].shape[0]
+ self.log(
+ f"test_sampling_notemplate/{metric_name}",
+ log_stat,
+ on_step=True,
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ sampling_stats = self.net.sample_pl_complex_structures(
+ batch,
+ sampler=self.hparams.cfg.task.sampler,
+ sampler_eta=self.hparams.cfg.task.sampler_eta,
+ num_steps=self.hparams.cfg.task.num_steps,
+ start_time=self.hparams.cfg.task.start_time,
+ return_summary_stats=True,
+ exact_prior=True,
+ )
+ for metric_name in sampling_stats.keys():
+ log_stat = sampling_stats[metric_name].mean().detach()
+ batch_size = sampling_stats[metric_name].shape[0]
+ self.log(
+ f"test_sampling_trueprior/{metric_name}",
+ log_stat,
+ on_step=True,
+ on_epoch=True,
+ batch_size=batch_size,
+ )
+ batch = self.model_step(
+ batch, batch_idx, "test", loss_mode=self.hparams.cfg.task.loss_mode
+ )
+ except Exception as e:
+ log.error(
+ f"Failed to perform test step for {batch['metadata']['sample_ID_per_sample']} with batch index {batch_idx} of dataloader {dataloader_idx} due to: {e}."
+ )
+ raise e
+
+ # store model outputs for inspection
+ test_outputs = {}
+ if (
+ self.hparams.cfg.task.visualize_generated_samples
+ and self.hparams.cfg.task.eval_structure_prediction
+ ):
+ test_outputs.update(
+ {
+ "name": batch["metadata"]["sample_ID_per_sample"],
+ "batch_size": batch["metadata"]["num_structid"],
+ "aatype": batch["features"]["res_type"].long().cpu().numpy(),
+ "res_atom_mask": batch["features"]["res_atom_mask"].cpu().numpy(),
+ "protein_coordinates_list": [
+ frame["receptor_padded"].cpu().numpy() for frame in all_frames
+ ],
+ "ligand_coordinates_list": [
+ frame["ligands"].cpu().numpy() for frame in all_frames
+ ],
+ "ligand_mol": batch["metadata"]["mol_per_sample"],
+ "protein_batch_indexer": batch["indexer"]["gather_idx_a_structid"]
+ .cpu()
+ .numpy(),
+ "ligand_batch_indexer": batch["indexer"]["gather_idx_i_structid"]
+ .cpu()
+ .numpy(),
+ "gt_protein_coordinates": batch["features"]["res_atom_positions"]
+ .cpu()
+ .numpy(),
+ "gt_ligand_coordinates": batch["features"]["sdf_coordinates"].cpu().numpy(),
+ "dataloader_idx": dataloader_idx,
+ }
+ )
+ if self.hparams.cfg.affinity.enabled and "affinity_logits" in batch["outputs"]:
+ test_outputs.update(
+ {
+ "affinity_logits": batch["outputs"]["affinity_logits"],
+ "affinity": batch["features"]["affinity"],
+ "dataloader_idx": dataloader_idx,
+ }
+ )
+ if test_outputs:
+ self.test_step_outputs.append(test_outputs)
+
+ def on_test_epoch_end(self):
+ """Lightning hook that is called when a test epoch ends."""
+ if (
+ self.hparams.cfg.task.visualize_generated_samples
+ and self.hparams.cfg.task.eval_structure_prediction
+ ):
+ for i, outputs in enumerate(self.test_step_outputs):
+ for batch_index in range(outputs["batch_size"]):
+ prot_lig_pairs = construct_prot_lig_pairs(outputs, batch_index)
+ write_prot_lig_pairs_to_pdb_file(
+ prot_lig_pairs,
+ os.path.join(
+ self.trainer.default_root_dir,
+ "test_epoch_outputs",
+ f"{outputs['name'][batch_index]}_test_epoch_{self.current_epoch}_global_step_{self.global_step}_output_{i}_batch_{batch_index}_dataloader_{outputs['dataloader_idx']}.pdb",
+ ),
+ )
+ if self.hparams.cfg.affinity.enabled and any(
+ "affinity_logits" in output for output in self.test_step_outputs
+ ):
+ affinity_logits = torch.cat(
+ [
+ output["affinity_logits"]
+ for output in self.test_step_outputs
+ if "affinity_logits" in output
+ ]
+ )
+ affinity = torch.cat(
+ [
+ output["affinity"]
+ for output in self.test_step_outputs
+ if "affinity_logits" in output
+ ]
+ )
+ affinity_logits = affinity_logits[~affinity.isnan()]
+ affinity = affinity[~affinity.isnan()]
+ if affinity.numel() > 1:
+ # NOTE: there must be at least two affinity batches to properly score the affinity predictions
+ aff_rmse = torch.sqrt(mean_squared_error(affinity_logits, affinity))
+ aff_mae = mean_absolute_error(affinity_logits, affinity)
+ aff_pearson = pearson_corrcoef(affinity_logits, affinity)
+ aff_spearman = spearman_corrcoef(affinity_logits, affinity)
+ self.log(
+ "test_affinity/RMSE",
+ aff_rmse.detach(),
+ on_epoch=True,
+ batch_size=len(affinity),
+ sync_dist=True,
+ )
+ self.log(
+ "test_affinity/MAE",
+ aff_mae.detach(),
+ on_epoch=True,
+ batch_size=len(affinity),
+ sync_dist=True,
+ )
+ self.log(
+ "test_affinity/Pearson",
+ aff_pearson.detach(),
+ on_epoch=True,
+ batch_size=len(affinity),
+ sync_dist=True,
+ )
+ self.log(
+ "test_affinity/Spearman",
+ aff_spearman.detach(),
+ on_epoch=True,
+ batch_size=len(affinity),
+ sync_dist=True,
+ )
+ self.test_step_outputs.clear() # free memory
+
+ def on_predict_start(self):
+ """Lightning hook that is called when testing begins."""
+ # create a directory to store model outputs from each predict epoch
+ os.makedirs(
+ os.path.join(self.trainer.default_root_dir, "predict_epoch_outputs"), exist_ok=True
+ )
+
+ log.info("Loading pretrained ESM model...")
+ esm_model, self.esm_alphabet = esm.pretrained.load_model_and_alphabet_hub(
+ self.hparams.cfg.model.cfg.protein_encoder.esm_version
+ )
+ self.esm_model = esm_model.eval().float()
+ self.esm_batch_converter = self.esm_alphabet.get_batch_converter()
+ self.esm_model.cpu()
+
+ skip_loading_esmfold_weights = (
+ # skip loading ESMFold weights if the template protein structure for a single complex input is provided
+ self.hparams.cfg.task.csv_path is None
+ and self.hparams.cfg.task.input_template is not None
+ and os.path.exists(self.hparams.cfg.task.input_template)
+ )
+ if not skip_loading_esmfold_weights:
+ log.info("Loading pretrained ESMFold model...")
+ esmfold_model = esm.pretrained.esmfold_v1()
+ self.esmfold_model = esmfold_model.eval().float()
+ self.esmfold_model.set_chunk_size(self.hparams.cfg.esmfold_chunk_size)
+ self.esmfold_model.cpu()
+
+ def predict_step(self, batch: MODEL_BATCH, batch_idx: int, dataloader_idx: int = 0):
+ """Perform a single predict step on a batch of data from the predict set.
+
+ :param batch: A batch dictionary.
+ :param batch_idx: The index of the current batch.
+ :param dataloader_idx: The index of the current dataloader.
+ """
+ rec_path = batch["rec_path"][0]
+ ligand_paths = list(
+ path[0] for path in batch["lig_paths"]
+ ) # unpack a list of (batched) single-element string tuples
+ sample_id = batch["sample_id"][0] if "sample_id" in batch else "sample"
+ input_template = batch["input_template"][0] if "input_template" in batch else None
+
+ # generate ESM embeddings for the protein
+ protein = pdb_filepath_to_protein(rec_path)
+ sequences = [
+ "".join(np.array(list(chain_seq))[chain_mask])
+ for (_, chain_seq, chain_mask) in protein.letter_sequences
+ ]
+ esm_embeddings = extract_esm_embeddings(
+ self.esm_model,
+ self.esm_alphabet,
+ self.esm_batch_converter,
+ sequences,
+ device="cpu",
+ esm_repr_layer=self.hparams.cfg.model.cfg.protein_encoder.esm_repr_layer,
+ )
+ sequences_to_embeddings = {
+ f"{seq}:{i}": esm_embeddings[i].cpu().numpy() for i, seq in enumerate(sequences)
+ }
+
+ # generate initial ESMFold-predicted structure for the protein if a template is not provided
+ apo_rec_path = None
+ if input_template and os.path.exists(input_template):
+ apo_protein = pdb_filepath_to_protein(input_template)
+ apo_chain_seq_masked = "".join(
+ "".join(np.array(list(chain_seq))[chain_mask])
+ for (_, chain_seq, chain_mask) in apo_protein.letter_sequences
+ )
+ chain_seq_masked = "".join(
+ "".join(np.array(list(chain_seq))[chain_mask])
+ for (_, chain_seq, chain_mask) in protein.letter_sequences
+ )
+ if apo_chain_seq_masked != chain_seq_masked:
+ log.error(
+ f"Provided template protein structure {input_template} does not match the input protein sequence within {rec_path}. Skipping example {sample_id} at batch index {batch_idx} of dataloader {dataloader_idx}."
+ )
+ return None
+ log.info(f"Starting from provided template protein structure: {input_template}")
+ apo_rec_path = input_template
+ if apo_rec_path is None and self.hparams.cfg.prior_type == "esmfold":
+ esmfold_sequence = ":".join(sequences)
+ apo_rec_path = rec_path.replace(".pdb", "_apo.pdb")
+ with torch.no_grad():
+ esmfold_pdb_output = self.esmfold_model.infer_pdb(esmfold_sequence)
+ with open(apo_rec_path, "w") as f:
+ f.write(esmfold_pdb_output)
+
+ _, _, _, _, _, all_frames, batch_all, b_factors, plddt_rankings = multi_pose_sampling(
+ rec_path,
+ ligand_paths,
+ self.hparams.cfg,
+ self,
+ self.hparams.cfg.out_path,
+ separate_pdb=self.hparams.cfg.separate_pdb,
+ apo_receptor_path=apo_rec_path,
+ sample_id=sample_id,
+ protein=protein,
+ sequences_to_embeddings=sequences_to_embeddings,
+ return_all_states=self.hparams.cfg.task.visualize_generated_samples,
+ auxiliary_estimation_only=self.hparams.cfg.task.auxiliary_estimation_only,
+ )
+ # store model outputs for inspection
+ if self.hparams.cfg.task.visualize_generated_samples:
+ predict_outputs = {
+ "name": batch_all["metadata"]["sample_ID_per_sample"],
+ "batch_size": batch_all["metadata"]["num_structid"],
+ "aatype": batch_all["features"]["res_type"].long().cpu().numpy(),
+ "res_atom_mask": batch_all["features"]["res_atom_mask"].cpu().numpy(),
+ "protein_coordinates_list": [
+ frame["receptor_padded"].cpu().numpy() for frame in all_frames
+ ],
+ "ligand_coordinates_list": [
+ frame["ligands"].cpu().numpy() for frame in all_frames
+ ],
+ "ligand_mol": batch_all["metadata"]["mol_per_sample"],
+ "protein_batch_indexer": batch_all["indexer"]["gather_idx_a_structid"]
+ .cpu()
+ .numpy(),
+ "ligand_batch_indexer": batch_all["indexer"]["gather_idx_i_structid"]
+ .cpu()
+ .numpy(),
+ "b_factors": b_factors,
+ "plddt_rankings": plddt_rankings,
+ }
+ self.predict_step_outputs.append(predict_outputs)
+
+ def on_predict_epoch_end(self):
+ """Lightning hook that is called when a predict epoch ends."""
+ if self.hparams.cfg.task.visualize_generated_samples:
+ for i, outputs in enumerate(self.predict_step_outputs):
+ for batch_index in range(outputs["batch_size"]):
+ prot_lig_pairs = construct_prot_lig_pairs(outputs, batch_index)
+ ranking = (
+ outputs["plddt_rankings"][batch_index]
+ if "plddt_rankings" in outputs
+ else None
+ )
+ write_prot_lig_pairs_to_pdb_file(
+ prot_lig_pairs,
+ os.path.join(
+ self.hparams.cfg.out_path,
+ "predict_epoch_outputs",
+ f"{outputs['name'][batch_index]}{f'_rank{ranking + 1}' if ranking is not None else ''}_predict_epoch_{self.current_epoch}_global_step_{self.global_step}_output_{i}_batch_{batch_index}.pdb",
+ ),
+ )
+ self.predict_step_outputs.clear() # free memory
+
+ def on_after_backward(self):
+ """Skip updates in case of unstable gradients.
+
+ Reference: https://github.com/Lightning-AI/lightning/issues/4956
+ """
+ valid_gradients = True
+ for _, param in self.named_parameters():
+ if param.grad is not None:
+ valid_gradients = not (
+ torch.isnan(param.grad).any() or torch.isinf(param.grad).any()
+ )
+ if not valid_gradients:
+ break
+ if not valid_gradients:
+ log.warning(
+ "Detected `inf` or `nan` values in gradients. Not updating model parameters."
+ )
+ self.zero_grad()
+
+ def optimizer_step(
+ self,
+ epoch,
+ batch_idx,
+ optimizer,
+ optimizer_closure,
+ ):
+ """Override the optimizer step to dynamically update the learning rate.
+
+ :param epoch: The current epoch.
+ :param batch_idx: The index of the current batch.
+ :param optimizer: The optimizer to use for training.
+ :param optimizer_closure: The optimizer closure.
+ """
+ # update params
+ optimizer = optimizer.optimizer
+ optimizer.step(closure=optimizer_closure)
+
+ # warm up learning rate
+ if self.trainer.global_step < 1000:
+ lr_scale = min(1.0, float(self.trainer.global_step + 1) / 1000.0)
+ for pg in optimizer.param_groups:
+ # NOTE: `self.hparams.optimizer.keywords["lr"]` refers to the optimizer's initial learning rate
+ pg["lr"] = lr_scale * self.hparams.optimizer.keywords["lr"]
+
+ def setup(self, stage: str):
+ """Lightning hook that is called at the beginning of fit (train + validate), validate,
+ test, or predict.
+
+ This is a good hook when you need to build models dynamically or adjust something about
+ them. This hook is called on every process when using DDP.
+
+ :param stage: Either `"fit"`, `"validate"`, `"test"`, or `"predict"`.
+ """
+ if self.hparams.compile and stage == "fit":
+ self.net = torch.compile(self.net)
+
+ def configure_optimizers(self) -> Dict[str, Any]:
+ """Choose what optimizers and learning-rate schedulers to use in your optimization.
+ Normally you'd need one. But in the case of GANs or similar you might have multiple.
+
+ Examples:
+ https://lightning.ai/docs/pytorch/latest/common/lightning_module.html#configure-optimizers
+
+ :return: A dict containing the configured optimizers and learning-rate schedulers to be used for training.
+ """
+ try:
+ optimizer = self.hparams.optimizer(params=self.trainer.model.parameters())
+ except TypeError:
+ # NOTE: strategies such as DeepSpeed require `params` to instead be specified as `model_params`
+ optimizer = self.hparams.optimizer(model_params=self.trainer.model.parameters())
+ if self.hparams.scheduler is not None:
+ scheduler = self.hparams.scheduler(optimizer=optimizer)
+ return {
+ "optimizer": optimizer,
+ "lr_scheduler": {
+ "scheduler": scheduler,
+ "monitor": "val/loss",
+ "interval": "epoch",
+ "frequency": 1,
+ },
+ }
+ return {"optimizer": optimizer}
+
+
+if __name__ == "__main__":
+ _ = FlowDockFMLitModule(None, None, None, None)
diff --git a/forks/FlowDock/flowdock/sample.py b/forks/FlowDock/flowdock/sample.py
new file mode 100644
index 00000000..d43b380e
--- /dev/null
+++ b/forks/FlowDock/flowdock/sample.py
@@ -0,0 +1,279 @@
+import os
+
+import hydra
+import lovely_tensors as lt
+import pandas as pd
+import rootutils
+import torch
+from beartype.typing import Any, Dict, List, Tuple
+from lightning import LightningModule, Trainer
+from lightning.fabric.plugins.environments.cluster_environment import ClusterEnvironment
+from lightning.pytorch.loggers import Logger
+from lightning.pytorch.strategies.strategy import Strategy
+from omegaconf import DictConfig, open_dict
+from torch.utils.data import DataLoader
+
+lt.monkey_patch()
+
+rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+# ------------------------------------------------------------------------------------ #
+# the setup_root above is equivalent to:
+# - adding project root dir to PYTHONPATH
+# (so you don't need to force user to install project as a package)
+# (necessary before importing any local modules e.g. `from flowdock import utils`)
+# - setting up PROJECT_ROOT environment variable
+# (which is used as a base for paths in "configs/paths/default.yaml")
+# (this way all filepaths are the same no matter where you run the code)
+# - loading environment variables from ".env" in root dir
+#
+# you can remove it if you:
+# 1. either install project as a package or move entry files to project root dir
+# 2. set `root_dir` to "." in "configs/paths/default.yaml"
+#
+# more info: https://github.com/ashleve/rootutils
+# ------------------------------------------------------------------------------------ #
+
+from flowdock import register_custom_omegaconf_resolvers, resolve_omegaconf_variable
+from flowdock.utils import (
+ RankedLogger,
+ extras,
+ instantiate_loggers,
+ log_hyperparameters,
+ task_wrapper,
+)
+from flowdock.utils.data_utils import (
+ create_full_pdb_with_zero_coordinates,
+ create_temp_ligand_frag_files,
+)
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+AVAILABLE_SAMPLING_TASKS = ["batched_structure_sampling"]
+
+
+class SamplingDataset(torch.utils.data.Dataset):
+ """Dataset for sampling."""
+
+ def __init__(self, cfg: DictConfig):
+ """Initializes the SamplingDataset."""
+ if cfg.sampling_task == "batched_structure_sampling":
+ if cfg.csv_path is not None:
+ # handle variable CSV inputs
+ df_rows = []
+ self.df = pd.read_csv(cfg.csv_path)
+ for _, row in self.df.iterrows():
+ sample_id = row.id
+ input_receptor = row.input_receptor
+ input_ligand = row.input_ligand
+ input_template = row.input_template
+ assert input_receptor is not None, "Receptor path is required for sampling."
+ if input_ligand is not None:
+ if input_ligand.endswith(".sdf"):
+ ligand_paths = create_temp_ligand_frag_files(input_ligand)
+ else:
+ ligand_paths = list(input_ligand.split("|"))
+ else:
+ ligand_paths = None # handle `null` ligand input
+ if not input_receptor.endswith(".pdb"):
+ log.warning(
+ "Assuming the provided receptor input is a protein sequence. Creating a dummy PDB file."
+ )
+ create_full_pdb_with_zero_coordinates(
+ input_receptor, os.path.join(cfg.out_path, f"input_{sample_id}.pdb")
+ )
+ input_receptor = os.path.join(cfg.out_path, f"input_{sample_id}.pdb")
+ df_row = {
+ "sample_id": sample_id,
+ "rec_path": input_receptor,
+ "lig_paths": ligand_paths,
+ }
+ if input_template is not None:
+ df_row["input_template"] = input_template
+ df_rows.append(df_row)
+ self.df = pd.DataFrame(df_rows)
+ else:
+ sample_id = cfg.sample_id
+ input_receptor = cfg.input_receptor
+ input_ligand = cfg.input_ligand
+ if input_ligand is not None:
+ if input_ligand.endswith(".sdf"):
+ ligand_paths = create_temp_ligand_frag_files(input_ligand)
+ else:
+ ligand_paths = list(input_ligand.split("|"))
+ else:
+ ligand_paths = None # handle `null` ligand input
+ if not input_receptor.endswith(".pdb"):
+ log.warning(
+ "Assuming the provided receptor input is a protein sequence. Creating a dummy PDB file."
+ )
+ create_full_pdb_with_zero_coordinates(
+ input_receptor, os.path.join(cfg.out_path, "input.pdb")
+ )
+ input_receptor = os.path.join(cfg.out_path, "input.pdb")
+ self.df = pd.DataFrame(
+ [
+ {
+ "sample_id": sample_id,
+ "rec_path": input_receptor,
+ "lig_paths": ligand_paths,
+ }
+ ]
+ )
+ if cfg.input_template is not None:
+ self.df["input_template"] = cfg.input_template
+ else:
+ raise NotImplementedError(f"Sampling task {cfg.sampling_task} is not implemented.")
+
+ def __len__(self):
+ """Returns the length of the dataset."""
+ return len(self.df)
+
+ def __getitem__(self, idx: int) -> Tuple[str, str]:
+ """Returns the input receptor and input ligand."""
+ return self.df.iloc[idx].to_dict()
+
+
+@task_wrapper
+def sample(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+ """Samples using given checkpoint on a datamodule predictset.
+
+ This method is wrapped in optional @task_wrapper decorator, that controls the behavior during
+ failure. Useful for multiruns, saving info about the crash, etc.
+
+ :param cfg: DictConfig configuration composed by Hydra.
+ :return: Tuple[dict, dict] with metrics and dict with all instantiated objects.
+ """
+ assert cfg.ckpt_path, "Please provide a checkpoint path with which to sample!"
+ assert os.path.exists(cfg.ckpt_path), f"Checkpoint path {cfg.ckpt_path} does not exist!"
+ assert (
+ cfg.sampling_task in AVAILABLE_SAMPLING_TASKS
+ ), f"Sampling task {cfg.sampling_task} is not one of the following available tasks: {AVAILABLE_SAMPLING_TASKS}."
+ assert (cfg.input_receptor is not None and cfg.input_ligand is not None) or (
+ cfg.csv_path is not None and os.path.exists(cfg.csv_path)
+ ), "Please provide either an input receptor and ligand or a CSV file with receptor and ligand sequences/filepaths."
+
+ log.info(
+ f"Setting `float32_matmul_precision` to {cfg.model.cfg.task.float32_matmul_precision}."
+ )
+ torch.set_float32_matmul_precision(precision=cfg.model.cfg.task.float32_matmul_precision)
+
+ # Establish model input arguments
+ with open_dict(cfg):
+ # NOTE: Structure trajectories will not be visualized when performing auxiliary estimation only
+ cfg.model.cfg.prior_type = cfg.prior_type
+ cfg.model.cfg.task.detect_covalent = cfg.detect_covalent
+ cfg.model.cfg.task.use_template = cfg.use_template
+ cfg.model.cfg.task.csv_path = cfg.csv_path
+ cfg.model.cfg.task.input_receptor = cfg.input_receptor
+ cfg.model.cfg.task.input_ligand = cfg.input_ligand
+ cfg.model.cfg.task.input_template = cfg.input_template
+ cfg.model.cfg.task.visualize_generated_samples = (
+ cfg.visualize_sample_trajectories and not cfg.auxiliary_estimation_only
+ )
+ cfg.model.cfg.task.auxiliary_estimation_only = cfg.auxiliary_estimation_only
+ if cfg.latent_model is not None:
+ with open_dict(cfg):
+ cfg.model.cfg.latent_model = cfg.latent_model
+ with open_dict(cfg):
+ if cfg.start_time == "auto":
+ cfg.start_time = 1.0
+ else:
+ cfg.start_time = float(cfg.start_time)
+
+ log.info("Converting sampling inputs into a
- 
+
+
+
+Where are downloaded weights stored?
+
+By default, weights are automatically downloaded and stored in
+
## ⚡ Try it online
@@ -42,6 +61,13 @@ We provide a [web server](https://lab.chaidiscovery.com) so you can test the Cha
How can MSAs be provided to Chai-1?
++ +Chai-1 supports MSAs provided as an `aligned.pqt` file. This file format is similar to an `a3m` file, but has additional columns that provide metadata like the source database and sequence pairing keys. We provide code to convert `a3m` files to `aligned.pqt` files. For more information on how to provide MSAs to Chai-1, see [this documentation](examples/msas/README.md). + +
+
+## Using experimental restraints
+Chai-1 uniquely offers the ability to fold complexes with user-specified "restraints" as inputs. These restraints specify inter-chain contacts at various resolutions that are used to guide Chai-1 in folding the complex. See [restraints documentation](examples/restraints/README.md) for details.
+
+
+ 
+