diff --git a/.gitignore b/.gitignore
index 481e7811..0d8860ee 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,4 +1,5 @@
.vscode/
+.idea/
__pycache__/
*.egg-info
build
@@ -8,3 +9,4 @@ dist
data
openfold/resources/
tests/test_data/
+cutlass/
diff --git a/README.md b/README.md
index 59bcb846..7b72ba22 100644
--- a/README.md
+++ b/README.md
@@ -7,13 +7,31 @@ _Figure: Comparison of OpenFold and AlphaFold2 predictions to the experimental s
A faithful but trainable PyTorch reproduction of DeepMind's
[AlphaFold 2](https://github.com/deepmind/alphafold).
+## Contents
+
+- [OpenFold](#openfold)
+ - [Contents](#contents)
+ - [Features](#features)
+ - [Installation (Linux)](#installation-linux)
+ - [Download Alignment Databases](#download-alignment-databases)
+ - [Inference](#inference)
+ - [Monomer inference](#monomer-inference)
+ - [Multimer Inference](#multimer-inference)
+ - [Soloseq Inference](#soloseq-inference)
+ - [Training](#training)
+ - [Testing](#testing)
+ - [Building and Using the Docker Container](#building-and-using-the-docker-container)
+ - [Copyright Notice](#copyright-notice)
+ - [Contributing](#contributing)
+ - [Citing this Work](#citing-this-work)
+
## Features
OpenFold carefully reproduces (almost) all of the features of the original open
-source inference code (v2.0.1). The sole exception is model ensembling, which
-fared poorly in DeepMind's own ablation testing and is being phased out in future
-DeepMind experiments. It is omitted here for the sake of reducing clutter. In
-cases where the *Nature* paper differs from the source, we always defer to the
+source monomer (v2.0.1) and multimer (v2.3.2) inference code. The sole exception is
+model ensembling, which fared poorly in DeepMind's own ablation testing and is being
+phased out in future DeepMind experiments. It is omitted here for the sake of reducing
+clutter. In cases where the *Nature* paper differs from the source, we always defer to the
latter.
OpenFold is trainable in full precision, half precision, or `bfloat16` with or without DeepSpeed,
@@ -63,7 +81,7 @@ To install:
For some systems, it may help to append the Conda environment library path to `$LD_LIBRARY_PATH`. The `install_third_party_dependencies.sh` script does this once, but you may need this for each bash instance.
-## Usage
+## Download Alignment Databases
If you intend to generate your own alignments, e.g. for inference, you have two
choices for downloading protein databases, depending on whether you want to use
@@ -112,7 +130,16 @@ DeepMind's pretrained parameters, you will only be able to make changes that
do not affect the shapes of model parameters. For an example of initializing
the model, consult `run_pretrained_openfold.py`.
-### Inference
+## Inference
+
+OpenFold now supports three inference modes:
+- [Monomer Inference](#monomer-inference): OpenFold reproduction of AlphaFold2. Inference available with either DeepMind's pretrained parameters or OpenFold trained parameters.
+- [Multimer Inference](#multimer-inference): OpenFold reproduction of AlphaFold-Multimer. Inference available with DeepMind's pre-trained parameters.
+- [Single Sequence Inference (SoloSeq)](#soloseq-inference): Language Model based structure prediction, using [ESM-1b](https://github.com/facebookresearch/esm) embeddings.
+
+More instructions for each inference mode are provided below:
+
+### Monomer inference
To run inference on a sequence or multiple sequences using a set of DeepMind's
pretrained parameters, first download the OpenFold weights e.g.:
@@ -131,14 +158,14 @@ python3 run_pretrained_openfold.py \
--mgnify_database_path data/mgnify/mgy_clusters_2018_12.fa \
--pdb70_database_path data/pdb70/pdb70 \
--uniclust30_database_path data/uniclust30/uniclust30_2018_08/uniclust30_2018_08 \
- --output_dir ./ \
--bfd_database_path data/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
- --model_device "cuda:0" \
--jackhmmer_binary_path lib/conda/envs/openfold_venv/bin/jackhmmer \
--hhblits_binary_path lib/conda/envs/openfold_venv/bin/hhblits \
--hhsearch_binary_path lib/conda/envs/openfold_venv/bin/hhsearch \
--kalign_binary_path lib/conda/envs/openfold_venv/bin/kalign \
--config_preset "model_1_ptm" \
+ --model_device "cuda:0" \
+ --output_dir ./ \
--openfold_checkpoint_path openfold/resources/openfold_params/finetuning_ptm_2.pt
```
@@ -176,13 +203,6 @@ To enable it, add `--trace_model` to the inference command.
To get a speedup during inference, enable [FlashAttention](https://github.com/HazyResearch/flash-attention)
in the config. Note that it appears to work best for sequences with < 1000 residues.
-Input FASTA files containing multiple sequences are treated as complexes. In
-this case, the inference script runs AlphaFold-Gap, a hack proposed
-[here](https://twitter.com/minkbaek/status/1417538291709071362?lang=en), using
-the specified stock AlphaFold/OpenFold parameters (NOT AlphaFold-Multimer). To
-run inference with AlphaFold-Multimer, use the (experimental) `multimer` branch
-instead.
-
To minimize memory usage during inference on long sequences, consider the
following changes:
@@ -221,7 +241,78 @@ efficent AlphaFold-Multimer more than double the time. Use the
at once. The `run_pretrained_openfold.py` script can enable this config option with the
`--long_sequence_inference` command line option
-#### SoloSeq Inference
+Input FASTA files containing multiple sequences are treated as complexes. In
+this case, the inference script runs AlphaFold-Gap, a hack proposed
+[here](https://twitter.com/minkbaek/status/1417538291709071362?lang=en), using
+the specified stock AlphaFold/OpenFold parameters (NOT AlphaFold-Multimer).
+
+### Multimer Inference
+
+To run inference on a complex or multiple complexes using a set of DeepMind's pretrained parameters, run e.g.:
+
+```bash
+python3 run_pretrained_openfold.py \
+ fasta_dir \
+ data/pdb_mmcif/mmcif_files/ \
+ --uniref90_database_path data/uniref90/uniref90.fasta \
+ --mgnify_database_path data/mgnify/mgy_clusters_2022_05.fa \
+ --pdb_seqres_database_path data/pdb_seqres/pdb_seqres.txt \
+ --uniref30_database_path data/uniref30/UniRef30_2021_03 \
+ --uniprot_database_path data/uniprot/uniprot.fasta \
+ --bfd_database_path data/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
+ --jackhmmer_binary_path lib/conda/envs/openfold_venv/bin/jackhmmer \
+ --hhblits_binary_path lib/conda/envs/openfold_venv/bin/hhblits \
+ --hmmsearch_binary_path lib/conda/envs/openfold_venv/bin/hmmsearch \
+ --hmmbuild_binary_path lib/conda/envs/openfold_venv/bin/hmmbuild \
+ --kalign_binary_path lib/conda/envs/openfold_venv/bin/kalign \
+ --config_preset "model_1_multimer_v3" \
+ --model_device "cuda:0" \
+ --output_dir ./
+```
+
+As with monomer inference, if you've already computed alignments for the query, you can use
+the `--use_precomputed_alignments` option. Note that template searching in the multimer pipeline
+uses HMMSearch with the PDB SeqRes database, replacing HHSearch and PDB70 used in the monomer pipeline.
+
+**Upgrade from an existing OpenFold installation**
+
+The above command requires several upgrades to existing openfold installations.
+
+1. Re-download the alphafold parameters to get the latest
+AlphaFold-Multimer v3 weights:
+
+ ```bash
+ bash scripts/download_alphafold_params.sh openfold/resources
+ ```
+
+2. Download the [UniProt](https://www.uniprot.org/uniprotkb/)
+and [PDB SeqRes](https://www.rcsb.org/) databases:
+
+ ```bash
+ bash scripts/download_uniprot.sh data/
+ ```
+
+ The PDB SeqRes and PDB databases must be from the same date to avoid potential
+ errors during template searching. Remove the existing `data/pdb_mmcif` directory
+ and download both databases:
+
+ ```bash
+ bash scripts/download_pdb_mmcif.sh data/
+ bash scripts/download_pdb_seqres.sh data/
+ ```
+
+3. Additionally, AlphaFold-Multimer uses upgraded versions of the [MGnify](https://www.ebi.ac.uk/metagenomics)
+and [UniRef30](https://uniclust.mmseqs.com/) (previously UniClust30) databases. To download the upgraded databases, run:
+
+ ```bash
+ bash scripts/download_uniref30.sh data/
+ bash scripts/download_mgnify.sh data/
+ ```
+ Multimer inference can also run with the older database versions if desired.
+
+
+### Soloseq Inference
+
To run inference for a sequence using the SoloSeq single-sequence model, you can either precompute ESM-1b embeddings in bulk, or you can generate them during inference.
For generating ESM-1b embeddings in bulk, use the provided script: `scripts/precompute_embeddings.py`. The script takes a directory of FASTA files (one sequence per file) and generates ESM-1b embeddings in the same format and directory structure as required by SoloSeq. Following is an example command to use the script:
@@ -260,7 +351,7 @@ python3 run_pretrained_openfold.py \
--output_dir ./ \
--model_device "cuda:0" \
--config_preset "seq_model_esm1b_ptm" \
- --openfold_checkpoint_path openfold/resources/openfold_params/seq_model_esm1b_ptm.pt \
+ --openfold_checkpoint_path openfold/resources/openfold_soloseq_params/seq_model_esm1b_ptm.pt \
--uniref90_database_path data/uniref90/uniref90.fasta \
--pdb70_database_path data/pdb70/pdb70 \
--jackhmmer_binary_path lib/conda/envs/openfold_venv/bin/jackhmmer \
@@ -274,7 +365,7 @@ SoloSeq allows you to use the same flags and optimizations as the MSA-based Open
**NOTE:** Due to the nature of the ESM-1b embeddings, the sequence length for inference using the SoloSeq model is limited to 1022 residues. Sequences longer than that will be truncated.
-### Training
+## Training
To train the model, you will first need to precompute protein alignments.
@@ -412,9 +503,9 @@ environment. These run components of AlphaFold and OpenFold side by side and
ensure that output activations are adequately similar. For most modules, we
target a maximum pointwise difference of `1e-4`.
-## Building and using the docker container
+## Building and Using the Docker Container
-### Building the docker image
+**Building the Docker Image**
Openfold can be built as a docker container using the included dockerfile. To build it, run the following command from the root of this repository:
@@ -422,7 +513,7 @@ Openfold can be built as a docker container using the included dockerfile. To bu
docker build -t openfold .
```
-### Running the docker container
+**Running the Docker Container**
The built container contains both `run_pretrained_openfold.py` and `train_openfold.py` as well as all necessary software dependencies. It does not contain the model parameters, sequence, or structural databases. These should be downloaded to the host machine following the instructions in the Usage section above.
@@ -462,7 +553,7 @@ python3 /opt/openfold/run_pretrained_openfold.py \
--openfold_checkpoint_path /database/openfold_params/finetuning_ptm_2.pt
```
-## Copyright notice
+## Copyright Notice
While AlphaFold's and, by extension, OpenFold's source code is licensed under
the permissive Apache Licence, Version 2.0, DeepMind's pretrained parameters
@@ -475,7 +566,7 @@ replaces the original, more restrictive CC BY-NC 4.0 license as of January 2022.
If you encounter problems using OpenFold, feel free to create an issue! We also
welcome pull requests from the community.
-## Citing this work
+## Citing this Work
Please cite our paper:
@@ -504,4 +595,4 @@ If you use OpenProteinSet, please also cite:
primaryClass={q-bio.BM}
}
```
-Any work that cites OpenFold should also cite AlphaFold.
+Any work that cites OpenFold should also cite [AlphaFold](https://www.nature.com/articles/s41586-021-03819-2) and [AlphaFold-Multimer](https://www.biorxiv.org/content/10.1101/2021.10.04.463034v1) if applicable.
diff --git a/environment.yml b/environment.yml
index edf7a4f2..7b73b25c 100644
--- a/environment.yml
+++ b/environment.yml
@@ -14,6 +14,7 @@ dependencies:
- pytorch-lightning==1.5.10
- biopython==1.79
- numpy==1.21
+ - pandas==2.0
- PyYAML==5.4.1
- requests
- scipy==1.7
diff --git a/notebooks/OpenFold.ipynb b/notebooks/OpenFold.ipynb
old mode 100755
new mode 100644
index 013d79bf..8a844652
--- a/notebooks/OpenFold.ipynb
+++ b/notebooks/OpenFold.ipynb
@@ -3,8 +3,7 @@
{
"cell_type": "markdown",
"metadata": {
- "id": "view-in-github",
- "colab_type": "text"
+ "id": "view-in-github"
},
"source": [
"![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\")
"
@@ -52,25 +51,44 @@
"cell_type": "code",
"execution_count": null,
"metadata": {
- "cellView": "form",
"id": "rowN0bVYLe9n"
},
"outputs": [],
"source": [
"#@markdown ### Enter the amino acid sequence to fold ⬇️\n",
- "sequence = 'MAAHKGAEHHHKAAEHHEQAAKHHHAAAEHHEKGEHEQAAHHADTAYAHHKHAEEHAAQAAKHDAEHHAPKPH' #@param {type:\"string\"}\n",
+ "#@markdown For multiple sequences, separate sequences with a colon `:`\n",
+ "input_sequence = 'MKLKQVADKLEEVASKLYHNANELARVAKLLGER:MKLKQVADKLEEVASKLYHNANELARVAKLLGER: MKLKQVADKLEEVASKLYHNANELARVAKLLGER:MKLKQVADKLEEVASKLYHNANELARVAKLLGER' #@param {type:\"string\"}\n",
"\n",
"#@markdown ### Configure the model ⬇️\n",
"\n",
- "weight_set = 'OpenFold' #@param [\"OpenFold\", \"AlphaFold\"]\n",
+ "weight_set = 'AlphaFold' #@param [\"OpenFold\", \"AlphaFold\"]\n",
+ "model_mode = 'multimer' #@param [\"monomer\", \"multimer\"]\n",
"relax_prediction = True #@param {type:\"boolean\"}\n",
"\n",
+ "\n",
"# Remove all whitespaces, tabs and end lines; upper-case\n",
- "sequence = sequence.translate(str.maketrans('', '', ' \\n\\t')).upper()\n",
+ "input_sequence = input_sequence.translate(str.maketrans('', '', ' \\n\\t')).upper()\n",
"aatypes = set('ACDEFGHIKLMNPQRSTVWY') # 20 standard aatypes\n",
- "if not set(sequence).issubset(aatypes):\n",
- " raise Exception(f'Input sequence contains non-amino acid letters: {set(sequence) - aatypes}. OpenFold only supports 20 standard amino acids as inputs.')\n",
- "\n",
+ "allowed_chars = aatypes.union({':'})\n",
+ "if not set(input_sequence).issubset(allowed_chars):\n",
+ " raise Exception(f'Input sequence contains non-amino acid letters: {set(input_sequence) - allowed_chars}. OpenFold only supports 20 standard amino acids as inputs.')\n",
+ "\n",
+ "if ':' in input_sequence and weight_set != 'AlphaFold':\n",
+ " raise ValueError('Input sequence is a multimer, must select Alphafold weight set')\n",
+ "\n",
+ "import enum\n",
+ "@enum.unique\n",
+ "class ModelType(enum.Enum):\n",
+ " MONOMER = 0\n",
+ " MULTIMER = 1\n",
+ "\n",
+ "model_type_dict = {\n",
+ " 'monomer': ModelType.MONOMER,\n",
+ " 'multimer': ModelType.MULTIMER,\n",
+ "}\n",
+ "\n",
+ "model_type = model_type_dict[model_mode]\n",
+ "print(f'Length of input sequence : {len(input_sequence.replace(\":\", \"\"))}')\n",
"#@markdown After making your selections, execute this cell by pressing the\n",
"#@markdown *Play* button on the left."
]
@@ -79,17 +97,16 @@
"cell_type": "code",
"execution_count": null,
"metadata": {
- "cellView": "form",
"id": "woIxeCPygt7K"
},
"outputs": [],
"source": [
"#@title Install third-party software\n",
- "#@markdown Please execute this cell by pressing the *Play* button on \n",
+ "#@markdown Please execute this cell by pressing the *Play* button on\n",
"#@markdown the left.\n",
"\n",
"\n",
- "#@markdown **Note**: This installs the software on the Colab \n",
+ "#@markdown **Note**: This installs the software on the Colab\n",
"#@markdown notebook in the cloud and not on your computer.\n",
"\n",
"import os, time\n",
@@ -103,10 +120,8 @@
"os.system(\"wget -qnc https://github.com/conda-forge/miniforge/releases/latest/download/Mambaforge-Linux-x86_64.sh\")\n",
"os.system(\"bash Mambaforge-Linux-x86_64.sh -bfp /usr/local\")\n",
"os.system(\"mamba config --set auto_update_conda false\")\n",
- "os.system(f\"mamba install -y -c conda-forge -c bioconda kalign2=2.04 hhsuite=3.3.0 openmm=7.7.0 python={python_version} pdbfixer\")\n",
- "\n",
- "\n",
- "os.system(\"pip install -q \\\"torch<2\\\" biopython ml_collections py3Dmol modelcif\")\n",
+ "os.system(f\"mamba install -y -c conda-forge -c bioconda kalign2=2.04 hhsuite=3.3.0 openmm=7.7.0 python={python_version} pdbfixer biopython=1.79\")\n",
+ "os.system(\"pip install -q torch ml_collections py3Dmol modelcif\")\n",
"\n",
"try:\n",
" with io.capture_output() as captured:\n",
@@ -119,12 +134,12 @@
" %shell wget -q -P /content \\\n",
" https://git.scicore.unibas.ch/schwede/openstructure/-/raw/7102c63615b64735c4941278d92b554ec94415f8/modules/mol/alg/src/stereo_chemical_props.txt\n",
"\n",
- " %shell mkdir -p /content/openfold/openfold/resourcees\n",
- " \n",
- " commit = \"099769d2ecfd01a8baa8d950030df454a042c910\"\n",
+ " %shell mkdir -p /content/openfold/openfold/resources\n",
+ "\n",
+ " commit = \"e2e19f16676b1a409f9ba3a6f69b11ee7f5887c2\"\n",
" os.system(f\"pip install -q git+https://github.com/aqlaboratory/openfold.git@{commit}\")\n",
- " \n",
- " %shell cp -f /content/stereo_chemical_props.txt /usr/local/lib/python3.10/site-packages/openfold/resources/\n",
+ "\n",
+ " os.system(f\"cp -f -p /content/stereo_chemical_props.txt /usr/local/lib/python{python_version}/site-packages/openfold/resources/\")\n",
"\n",
"except subprocess.CalledProcessError as captured:\n",
" print(captured)"
@@ -134,18 +149,17 @@
"cell_type": "code",
"execution_count": null,
"metadata": {
- "cellView": "form",
"id": "VzJ5iMjTtoZw"
},
"outputs": [],
"source": [
- "#@title Download model weights \n",
- "#@markdown Please execute this cell by pressing the *Play* button on \n",
+ "#@title Download model weights\n",
+ "#@markdown Please execute this cell by pressing the *Play* button on\n",
"#@markdown the left.\n",
"\n",
"# Define constants\n",
"GIT_REPO='https://github.com/aqlaboratory/openfold'\n",
- "ALPHAFOLD_PARAM_SOURCE_URL = 'https://storage.googleapis.com/alphafold/alphafold_params_2022-01-19.tar'\n",
+ "ALPHAFOLD_PARAM_SOURCE_URL = 'https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar'\n",
"OPENFOLD_PARAMS_DIR = './openfold/openfold/resources/openfold_params'\n",
"ALPHAFOLD_PARAMS_DIR = './openfold/openfold/resources/params'\n",
"ALPHAFOLD_PARAMS_PATH = os.path.join(\n",
@@ -184,17 +198,17 @@
"cell_type": "code",
"execution_count": null,
"metadata": {
- "cellView": "form",
"id": "_FpxxMo-mvcP"
},
"outputs": [],
"source": [
"#@title Import Python packages\n",
- "#@markdown Please execute this cell by pressing the *Play* button on \n",
+ "#@markdown Please execute this cell by pressing the *Play* button on\n",
"#@markdown the left.\n",
"\n",
"import unittest.mock\n",
"import sys\n",
+ "from typing import Dict, Sequence\n",
"\n",
"sys.path.insert(0, f'/usr/local/lib/python{python_version}/dist-packages/')\n",
"sys.path.insert(0, f'/usr/local/lib/python{python_version}/site-packages/')\n",
@@ -234,22 +248,12 @@
" return \"UTF-8\"\n",
"locale.getpreferredencoding = getpreferredencoding\n",
"\n",
- "# A filthy hack to avoid slow Linear layer initialization\n",
- "import openfold.model.primitives\n",
- "\n",
- "def __default_linear_init__(self, *args, **kwargs):\n",
- " return torch.nn.Linear.__init__(\n",
- " self, \n",
- " *args[:2], \n",
- " **{k:v for k,v in kwargs.items() if k == \"bias\"}\n",
- " )\n",
- "\n",
- "openfold.model.primitives.Linear.__init__ = __default_linear_init__\n",
- "\n",
"from openfold import config\n",
"from openfold.data import feature_pipeline\n",
"from openfold.data import parsers\n",
"from openfold.data import data_pipeline\n",
+ "from openfold.data import msa_pairing\n",
+ "from openfold.data import feature_processing_multimer\n",
"from openfold.data.tools import jackhmmer\n",
"from openfold.model import model\n",
"from openfold.np import protein\n",
@@ -276,22 +280,16 @@
},
{
"cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "2tTeTTsLKPjB"
- },
- "outputs": [],
"source": [
"#@title Search against genetic databases\n",
"\n",
"#@markdown Once this cell has been executed, you will see\n",
- "#@markdown statistics about the multiple sequence alignment \n",
- "#@markdown (MSA) that will be used by OpenFold. In particular, \n",
- "#@markdown you’ll see how well each residue is covered by similar \n",
+ "#@markdown statistics about the multiple sequence alignment\n",
+ "#@markdown (MSA) that will be used by OpenFold. In particular,\n",
+ "#@markdown you’ll see how well each residue is covered by similar\n",
"#@markdown sequences in the MSA.\n",
"\n",
- "# --- Find the closest source ---\n",
+ "# --- Find the closest source --\n",
"test_url_pattern = 'https://storage.googleapis.com/alphafold-colab{:s}/latest/uniref90_2021_03.fasta.1'\n",
"ex = futures.ThreadPoolExecutor(3)\n",
"def fetch(source):\n",
@@ -304,114 +302,156 @@
" ex.shutdown()\n",
" break\n",
"\n",
- "# --- Search against genetic databases ---\n",
- "with open('target.fasta', 'wt') as f:\n",
- " f.write(f'>query\\n{sequence}')\n",
- "\n",
"# Run the search against chunks of genetic databases (since the genetic\n",
"# databases don't fit in Colab ramdisk).\n",
"\n",
"jackhmmer_binary_path = '/usr/bin/jackhmmer'\n",
- "dbs = []\n",
"\n",
- "num_jackhmmer_chunks = {'uniref90': 59, 'smallbfd': 17, 'mgnify': 71}\n",
- "total_jackhmmer_chunks = sum(num_jackhmmer_chunks.values())\n",
+ "# --- Parse multiple sequences, if there are any ---\n",
+ "def split_multiple_sequences(sequence):\n",
+ " seqs = sequence.split(':')\n",
+ " sorted_seqs = sorted(seqs, key=lambda s: len(s))\n",
+ "\n",
+ " # TODO: Handle the homomer case when writing fasta sequences\n",
+ " fasta_path_tuples = []\n",
+ " for idx, seq in enumerate(set(sorted_seqs)):\n",
+ " fasta_path = f'target_{idx+1}.fasta'\n",
+ " with open(fasta_path, 'wt') as f:\n",
+ " f.write(f'>query\\n{seq}\\n')\n",
+ " fasta_path_tuples.append((seq, fasta_path))\n",
+ " fasta_path_by_seq = dict(fasta_path_tuples)\n",
+ "\n",
+ " return sorted_seqs, fasta_path_by_seq\n",
+ "\n",
+ "sequences, fasta_path_by_sequence = split_multiple_sequences(input_sequence)\n",
+ "db_results_by_sequence = {seq: {} for seq in fasta_path_by_sequence.keys()}\n",
+ "\n",
+ "DB_ROOT_PATH = f'https://storage.googleapis.com/alphafold-colab{source}/latest/'\n",
+ "db_configs = {}\n",
+ "db_configs['smallbfd'] = {\n",
+ " 'database_path': f'{DB_ROOT_PATH}uniref90_2021_03.fasta',\n",
+ " 'z_value': 65984053,\n",
+ " 'num_jackhmmer_chunks': 17,\n",
+ "}\n",
+ "db_configs['mgnify'] = {\n",
+ " 'database_path': f'{DB_ROOT_PATH}mgy_clusters_2022_05.fasta',\n",
+ " 'z_value': 304820129,\n",
+ " 'num_jackhmmer_chunks': 120,\n",
+ "}\n",
+ "db_configs['uniref90'] = {\n",
+ " 'database_path': f'{DB_ROOT_PATH}uniref90_2022_01.fasta',\n",
+ " 'z_value': 144113457,\n",
+ " 'num_jackhmmer_chunks': 62,\n",
+ "}\n",
+ "\n",
+ "# Search UniProt and construct the all_seq features only for heteromers, not homomers.\n",
+ "if model_type == ModelType.MULTIMER and len(set(sequences)) > 1:\n",
+ " db_configs['uniprot'] = {\n",
+ " 'database_path': f'{DB_ROOT_PATH}uniprot_2021_04.fasta',\n",
+ " 'z_value': 225013025 + 565928,\n",
+ " 'num_jackhmmer_chunks': 101,\n",
+ " }\n",
+ "\n",
+ "total_jackhmmer_chunks = sum([d['num_jackhmmer_chunks'] for d in db_configs.values()])\n",
"with tqdm.notebook.tqdm(total=total_jackhmmer_chunks, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
" def jackhmmer_chunk_callback(i):\n",
" pbar.update(n=1)\n",
"\n",
- " pbar.set_description('Searching uniref90')\n",
- " jackhmmer_uniref90_runner = jackhmmer.Jackhmmer(\n",
- " binary_path=jackhmmer_binary_path,\n",
- " database_path=f'https://storage.googleapis.com/alphafold-colab{source}/latest/uniref90_2021_03.fasta',\n",
- " get_tblout=True,\n",
- " num_streamed_chunks=num_jackhmmer_chunks['uniref90'],\n",
- " streaming_callback=jackhmmer_chunk_callback,\n",
- " z_value=135301051)\n",
- " dbs.append(('uniref90', jackhmmer_uniref90_runner.query('target.fasta')))\n",
- "\n",
- " pbar.set_description('Searching smallbfd')\n",
- " jackhmmer_smallbfd_runner = jackhmmer.Jackhmmer(\n",
- " binary_path=jackhmmer_binary_path,\n",
- " database_path=f'https://storage.googleapis.com/alphafold-colab{source}/latest/bfd-first_non_consensus_sequences.fasta',\n",
- " get_tblout=True,\n",
- " num_streamed_chunks=num_jackhmmer_chunks['smallbfd'],\n",
- " streaming_callback=jackhmmer_chunk_callback,\n",
- " z_value=65984053)\n",
- " dbs.append(('smallbfd', jackhmmer_smallbfd_runner.query('target.fasta')))\n",
- "\n",
- " pbar.set_description('Searching mgnify')\n",
- " jackhmmer_mgnify_runner = jackhmmer.Jackhmmer(\n",
- " binary_path=jackhmmer_binary_path,\n",
- " database_path=f'https://storage.googleapis.com/alphafold-colab{source}/latest/mgy_clusters_2019_05.fasta',\n",
- " get_tblout=True,\n",
- " num_streamed_chunks=num_jackhmmer_chunks['mgnify'],\n",
- " streaming_callback=jackhmmer_chunk_callback,\n",
- " z_value=304820129)\n",
- " dbs.append(('mgnify', jackhmmer_mgnify_runner.query('target.fasta')))\n",
+ " for db_name, db_config in db_configs.items():\n",
+ " pbar.set_description(f'Searching {db_name}')\n",
+ " jackhmmer_runner = jackhmmer.Jackhmmer(\n",
+ " binary_path=jackhmmer_binary_path,\n",
+ " database_path=db_config['database_path'],\n",
+ " get_tblout=True,\n",
+ " num_streamed_chunks=db_config['num_jackhmmer_chunks'],\n",
+ " streaming_callback=jackhmmer_chunk_callback,\n",
+ " z_value=db_config['z_value'])\n",
+ "\n",
+ " db_results = jackhmmer_runner.query_multiple(fasta_path_by_sequence.values())\n",
+ " for seq, result in zip(fasta_path_by_sequence.keys(), db_results):\n",
+ " db_results_by_sequence[seq][db_name] = result\n",
"\n",
"\n",
"# --- Extract the MSAs and visualize ---\n",
"# Extract the MSAs from the Stockholm files.\n",
"# NB: deduplication happens later in data_pipeline.make_msa_features.\n",
"\n",
- "mgnify_max_hits = 501\n",
- "\n",
- "msas = []\n",
- "deletion_matrices = []\n",
- "full_msa = []\n",
- "for db_name, db_results in dbs:\n",
- " unsorted_results = []\n",
- " for i, result in enumerate(db_results):\n",
- " msa, deletion_matrix, target_names = parsers.parse_stockholm(result['sto'])\n",
- " e_values_dict = parsers.parse_e_values_from_tblout(result['tbl'])\n",
- " e_values = [e_values_dict[t.split('/')[0]] for t in target_names]\n",
- " zipped_results = zip(msa, deletion_matrix, target_names, e_values)\n",
- " if i != 0:\n",
- " # Only take query from the first chunk\n",
- " zipped_results = [x for x in zipped_results if x[2] != 'query']\n",
- " unsorted_results.extend(zipped_results)\n",
- " sorted_by_evalue = sorted(unsorted_results, key=lambda x: x[3])\n",
- " db_msas, db_deletion_matrices, _, _ = zip(*sorted_by_evalue)\n",
- " if db_msas:\n",
- " if db_name == 'mgnify':\n",
- " db_msas = db_msas[:mgnify_max_hits]\n",
- " db_deletion_matrices = db_deletion_matrices[:mgnify_max_hits]\n",
- " full_msa.extend(db_msas)\n",
- " msas.append(db_msas)\n",
- " deletion_matrices.append(db_deletion_matrices)\n",
- " msa_size = len(set(db_msas))\n",
- " print(f'{msa_size} Sequences Found in {db_name}')\n",
- "\n",
- "deduped_full_msa = list(dict.fromkeys(full_msa))\n",
- "total_msa_size = len(deduped_full_msa)\n",
- "print(f'\\n{total_msa_size} Sequences Found in Total\\n')\n",
- "\n",
- "aa_map = {restype: i for i, restype in enumerate('ABCDEFGHIJKLMNOPQRSTUVWXYZ-')}\n",
- "msa_arr = np.array([[aa_map[aa] for aa in seq] for seq in deduped_full_msa])\n",
- "num_alignments, num_res = msa_arr.shape\n",
+ "MAX_HITS_BY_DB = {\n",
+ " 'uniref90': 10000,\n",
+ " 'smallbfd': 5000,\n",
+ " 'mgnify': 501,\n",
+ " 'uniprot': 50000,\n",
+ "}\n",
+ "\n",
+ "msas_by_seq_by_db = {seq: {} for seq in sequences}\n",
+ "full_msa_by_seq = {seq: [] for seq in sequences}\n",
+ "\n",
+ "for seq, sequence_result in db_results_by_sequence.items():\n",
+ " print(f'parsing_results_for_sequence {seq}')\n",
+ " for db_name, db_results in sequence_result.items():\n",
+ " unsorted_results = []\n",
+ " for i, result in enumerate(db_results):\n",
+ " msa_obj = parsers.parse_stockholm(result['sto'])\n",
+ " e_values_dict = parsers.parse_e_values_from_tblout(result['tbl'])\n",
+ " target_names = msa_obj.descriptions\n",
+ " e_values = [e_values_dict[t.split('/')[0]] for t in target_names]\n",
+ " zipped_results = zip(msa_obj.sequences, msa_obj.deletion_matrix, target_names, e_values)\n",
+ " if i != 0:\n",
+ " # Only take query from the first chunk\n",
+ " zipped_results = [x for x in zipped_results if x[2] != 'query']\n",
+ " unsorted_results.extend(zipped_results)\n",
+ " sorted_by_evalue = sorted(unsorted_results, key=lambda x: x[3])\n",
+ " msas, del_matrix, targets, _ = zip(*sorted_by_evalue)\n",
+ " db_msas = parsers.Msa(msas, del_matrix, targets)\n",
+ " if db_msas:\n",
+ " if db_name in MAX_HITS_BY_DB:\n",
+ " db_msas.truncate(MAX_HITS_BY_DB[db_name])\n",
+ " msas_by_seq_by_db[seq][db_name] = db_msas\n",
+ " full_msa_by_seq[seq].extend(db_msas.sequences)\n",
+ " msa_size = len(set(db_msas.sequences))\n",
+ " print(f'{msa_size} Sequences Found in {db_name}')\n",
+ "\n",
"\n",
"fig = plt.figure(figsize=(12, 3))\n",
+ "max_num_alignments = 0\n",
+ "\n",
+ "for seq_idx, seq in enumerate(set(sequences)):\n",
+ " full_msas = full_msa_by_seq[seq]\n",
+ " deduped_full_msa = list(dict.fromkeys(full_msas))\n",
+ " total_msa_size = len(deduped_full_msa)\n",
+ " print(f'\\n{total_msa_size} Sequences Found in Total\\n')\n",
+ "\n",
+ " aa_map = {restype: i for i, restype in enumerate('ABCDEFGHIJKLMNOPQRSTUVWXYZ-')}\n",
+ " msa_arr = np.array([[aa_map[aa] for aa in seq] for seq in deduped_full_msa])\n",
+ " num_alignments, num_res = msa_arr.shape\n",
+ " plt.plot(np.sum(msa_arr != aa_map['-'], axis=0), label=f'Chain {seq_idx}')\n",
+ " max_num_alignments = max(num_alignments, max_num_alignments)\n",
+ "\n",
+ "\n",
"plt.title('Per-Residue Count of Non-Gap Amino Acids in the MSA')\n",
- "plt.plot(np.sum(msa_arr != aa_map['-'], axis=0), color='black')\n",
"plt.ylabel('Non-Gap Count')\n",
- "plt.yticks(range(0, num_alignments + 1, max(1, int(num_alignments / 3))))\n",
+ "plt.yticks(range(0, max_num_alignments + 1, max(1, int(max_num_alignments / 3))))\n",
+ "plt.legend()\n",
"plt.show()"
- ]
+ ],
+ "metadata": {
+ "id": "o7BqQN_gfYtq"
+ },
+ "execution_count": null,
+ "outputs": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
- "cellView": "form",
"id": "XUo6foMQxwS2"
},
"outputs": [],
"source": [
"#@title Run OpenFold and download prediction\n",
"\n",
- "#@markdown Once this cell has been executed, a zip-archive with \n",
- "#@markdown the obtained prediction will be automatically downloaded \n",
+ "#@markdown Once this cell has been executed, a zip-archive with\n",
+ "#@markdown the obtained prediction will be automatically downloaded\n",
"#@markdown to your computer.\n",
"\n",
"# Color bands for visualizing plddt\n",
@@ -423,13 +463,22 @@
"]\n",
"\n",
"# --- Run the model ---\n",
- "model_names = [ \n",
- " 'finetuning_3.pt', \n",
- " 'finetuning_4.pt', \n",
- " 'finetuning_5.pt', \n",
- " 'finetuning_ptm_2.pt',\n",
- " 'finetuning_no_templ_ptm_1.pt'\n",
- "]\n",
+ "if model_type == ModelType.MONOMER:\n",
+ " model_names = [\n",
+ " 'finetuning_3.pt',\n",
+ " 'finetuning_4.pt',\n",
+ " 'finetuning_5.pt',\n",
+ " 'finetuning_ptm_2.pt',\n",
+ " 'finetuning_no_templ_ptm_1.pt'\n",
+ " ]\n",
+ "elif model_type == ModelType.MULTIMER:\n",
+ " model_names = [\n",
+ " 'model_1_multimer_v3',\n",
+ " 'model_2_multimer_v3',\n",
+ " 'model_3_multimer_v3',\n",
+ " 'model_4_multimer_v3',\n",
+ " 'model_5_multimer_v3',\n",
+ " ]\n",
"\n",
"def _placeholder_template_feats(num_templates_, num_res_):\n",
" return {\n",
@@ -440,26 +489,72 @@
" 'template_sum_probs': np.zeros((num_templates_, 1), dtype=np.float32),\n",
" }\n",
"\n",
+ "\n",
+ "def make_features(\n",
+ " sequences: Sequence[str],\n",
+ " msas_by_seq_by_db: Dict[str, Dict[str, parsers.Msa]],\n",
+ " model_type: ModelType):\n",
+ " num_templates = 1 # Placeholder for generating fake template features\n",
+ " feature_dict = {}\n",
+ "\n",
+ " for idx, seq in enumerate(sequences, start=1):\n",
+ " _chain_id = f'chain_{idx}'\n",
+ " num_res = len(seq)\n",
+ "\n",
+ " feats = data_pipeline.make_sequence_features(seq, _chain_id, num_res)\n",
+ " msas_without_uniprot = [msas_by_seq_by_db[seq][db] for db in db_configs.keys() if db != 'uniprot']\n",
+ " msa_feats = data_pipeline.make_msa_features(msas_without_uniprot)\n",
+ " feats.update(msa_feats)\n",
+ " feats.update(_placeholder_template_feats(num_templates, num_res))\n",
+ "\n",
+ " if model_type == ModelType.MONOMER:\n",
+ " feature_dict[seq] = feats\n",
+ " if model_type == ModelType.MULTIMER:\n",
+ " # Perform extra pair processing steps for heteromers\n",
+ " if len(set(sequences)) > 1:\n",
+ " uniprot_msa = msas_by_seq_by_db[seq]['uniprot']\n",
+ " uniprot_msa_features = data_pipeline.make_msa_features([uniprot_msa])\n",
+ " valid_feat_names = msa_pairing.MSA_FEATURES + (\n",
+ " 'msa_species_identifiers',\n",
+ " )\n",
+ " pair_feats = {\n",
+ " f'{k}_all_seq': v for k, v in uniprot_msa_features.items()\n",
+ " if k in valid_feat_names\n",
+ " }\n",
+ " feats.update(pair_feats)\n",
+ "\n",
+ " feats = data_pipeline.convert_monomer_features(feats, _chain_id)\n",
+ " feature_dict[_chain_id] = feats\n",
+ "\n",
+ " if model_type == ModelType.MONOMER:\n",
+ " np_example = feature_dict[sequences[0]]\n",
+ " elif model_type == ModelType.MULTIMER:\n",
+ " all_chain_feats = data_pipeline.add_assembly_features(feature_dict)\n",
+ " features = feature_processing_multimer.pair_and_merge(all_chain_features=all_chain_feats)\n",
+ " np_example = data_pipeline.pad_msa(features, 512)\n",
+ "\n",
+ " return np_example\n",
+ "\n",
+ "\n",
"output_dir = 'prediction'\n",
"os.makedirs(output_dir, exist_ok=True)\n",
"\n",
"plddts = {}\n",
"pae_outputs = {}\n",
+ "weighted_ptms = {}\n",
"unrelaxed_proteins = {}\n",
"\n",
- "with tqdm.notebook.tqdm(total=len(model_names) + 1, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
- " for i, model_name in list(enumerate(model_names)):\n",
+ "with tqdm.notebook.tqdm(total=len(model_names), bar_format=TQDM_BAR_FORMAT) as pbar:\n",
+ " for i, model_name in enumerate(model_names, start = 1):\n",
" pbar.set_description(f'Running {model_name}')\n",
- " num_templates = 1 # dummy number --- is ignored\n",
- " num_res = len(sequence)\n",
- " \n",
- " feature_dict = {}\n",
- " feature_dict.update(data_pipeline.make_sequence_features(sequence, 'test', num_res))\n",
- " feature_dict.update(data_pipeline.make_msa_features(msas, deletion_matrices=deletion_matrices))\n",
- " feature_dict.update(_placeholder_template_feats(num_templates, num_res))\n",
+ "\n",
+ " feature_dict = make_features(sequences, msas_by_seq_by_db, model_type)\n",
"\n",
" if(weight_set == \"AlphaFold\"):\n",
- " config_preset = f\"model_{i}\"\n",
+ " if model_type == ModelType.MONOMER:\n",
+ " config_preset = f\"model_{i}\"\n",
+ " elif model_type == ModelType.MULTIMER:\n",
+ " config_preset = f'model_{i}_multimer_v3'\n",
" else:\n",
" if(\"_no_templ_\" in model_name):\n",
" config_preset = \"model_3\"\n",
@@ -469,6 +564,11 @@
" config_preset += \"_ptm\"\n",
"\n",
" cfg = config.model_config(config_preset)\n",
+ "\n",
+ " # Force the model to only use 3 recycling updates\n",
+ " cfg.data.common.max_recycling_iters = 3\n",
+ " cfg.model.recycle_early_stop_tolerance = -1\n",
+ "\n",
" openfold_model = model.AlphaFold(cfg)\n",
" openfold_model = openfold_model.eval()\n",
" if(weight_set == \"AlphaFold\"):\n",
@@ -490,7 +590,9 @@
"\n",
" pipeline = feature_pipeline.FeaturePipeline(cfg.data)\n",
" processed_feature_dict = pipeline.process_features(\n",
- " feature_dict, mode='predict'\n",
+ " feature_dict,\n",
+ " mode='predict',\n",
+ " is_multimer = (model_type == ModelType.MULTIMER),\n",
" )\n",
"\n",
" processed_feature_dict = tensor_tree_map(\n",
@@ -510,21 +612,32 @@
"\n",
" mean_plddt = prediction_result['plddt'].mean()\n",
"\n",
- " if 'predicted_aligned_error' in prediction_result:\n",
- " pae_outputs[model_name] = (\n",
- " prediction_result['predicted_aligned_error'],\n",
- " prediction_result['max_predicted_aligned_error']\n",
- " )\n",
- " else:\n",
- " # Get the pLDDT confidence metrics. Do not put pTM models here as they\n",
- " # should never get selected.\n",
+ " if model_type == ModelType.MONOMER:\n",
+ " if 'predicted_aligned_error' in prediction_result:\n",
+ " pae_outputs[model_name] = (\n",
+ " prediction_result['predicted_aligned_error'],\n",
+ " prediction_result['max_predicted_aligned_error']\n",
+ " )\n",
+ " else:\n",
+ " # Get the pLDDT confidence metrics. Do not put pTM models here as they\n",
+ " # should never get selected.\n",
+ " plddts[model_name] = prediction_result['plddt']\n",
+ " elif model_type == ModelType.MULTIMER:\n",
+ " # Multimer models are sorted by pTM+ipTM.\n",
" plddts[model_name] = prediction_result['plddt']\n",
+ " pae_outputs[model_name] = (prediction_result['predicted_aligned_error'],\n",
+ " prediction_result['max_predicted_aligned_error'])\n",
+ "\n",
+ " weighted_ptms[model_name] = prediction_result['weighted_ptm_score']\n",
"\n",
" # Set the b-factors to the per-residue plddt.\n",
" final_atom_mask = prediction_result['final_atom_mask']\n",
" b_factors = prediction_result['plddt'][:, None] * final_atom_mask\n",
" unrelaxed_protein = protein.from_prediction(\n",
- " processed_feature_dict, prediction_result, b_factors=b_factors\n",
+ " processed_feature_dict,\n",
+ " prediction_result,\n",
+ " remove_leading_feature_dimension=False,\n",
+ " b_factors=b_factors,\n",
" )\n",
" unrelaxed_proteins[model_name] = unrelaxed_protein\n",
"\n",
@@ -535,7 +648,10 @@
" pbar.update(n=1)\n",
"\n",
" # Find the best model according to the mean pLDDT.\n",
- " best_model_name = max(plddts.keys(), key=lambda x: plddts[x].mean())\n",
+ " if model_type == ModelType.MONOMER:\n",
+ " best_model_name = max(plddts.keys(), key=lambda x: plddts[x].mean())\n",
+ " elif model_type == ModelType.MULTIMER:\n",
+ " best_model_name = max(weighted_ptms.keys(), key=lambda x: weighted_ptms[x])\n",
" best_pdb = protein.to_pdb(unrelaxed_proteins[best_model_name])\n",
"\n",
" # --- AMBER relax the best model ---\n",
@@ -547,7 +663,7 @@
" stiffness=10.0,\n",
" exclude_residues=[],\n",
" max_outer_iterations=20,\n",
- " use_gpu=False,\n",
+ " use_gpu=True,\n",
" )\n",
" relaxed_pdb, _, _ = amber_relaxer.process(\n",
" prot=unrelaxed_proteins[best_model_name]\n",
@@ -598,6 +714,15 @@
" plt.title('Model Confidence', fontsize=20, pad=20)\n",
" return plt\n",
"\n",
+ "# Show the structure coloured by chain if the multimer model has been used.\n",
+ "if model_type == ModelType.MULTIMER:\n",
+ " multichain_view = py3Dmol.view(width=800, height=600)\n",
+ " multichain_view.addModelsAsFrames(to_visualize_pdb)\n",
+ " multichain_style = {'cartoon': {'colorscheme': 'chain'}}\n",
+ " multichain_view.setStyle({'model': -1}, multichain_style)\n",
+ " multichain_view.zoomTo()\n",
+ " multichain_view.show()\n",
+ "\n",
"# Color the structure by per-residue pLDDT\n",
"color_map = {i: bands[2] for i, bands in enumerate(PLDDT_BANDS)}\n",
"view = py3Dmol.view(width=800, height=600)\n",
@@ -643,6 +768,15 @@
" pae, max_pae = list(pae_outputs.values())[0]\n",
" plt.imshow(pae, vmin=0., vmax=max_pae, cmap='Greens_r')\n",
" plt.colorbar(fraction=0.046, pad=0.04)\n",
+ "\n",
+ " # Display lines at chain boundaries.\n",
+ " best_unrelaxed_prot = unrelaxed_proteins[best_model_name]\n",
+ " total_num_res = best_unrelaxed_prot.residue_index.shape[-1]\n",
+ " chain_ids = best_unrelaxed_prot.chain_index\n",
+ " for chain_boundary in np.nonzero(chain_ids[:-1] - chain_ids[1:]):\n",
+ " if chain_boundary.size:\n",
+ " plt.plot([0, total_num_res], [chain_boundary, chain_boundary], color='red')\n",
+ " plt.plot([chain_boundary, chain_boundary], [0, total_num_res], color='red')\n",
" plt.title('Predicted Aligned Error')\n",
" plt.xlabel('Scored residue')\n",
" plt.ylabel('Aligned residue')\n",
@@ -680,7 +814,7 @@
"source": [
"### Interpreting the prediction\n",
"\n",
- "Please see the [AlphaFold methods paper](https://www.nature.com/articles/s41586-021-03819-2) and the [AlphaFold predictions of the human proteome paper](https://www.nature.com/articles/s41586-021-03828-1), as well as [DeepMind's FAQ](https://alphafold.ebi.ac.uk/faq) on how to interpret AlphaFold/OpenFold predictions."
+ "Please see the [AlphaFold methods paper](https://www.nature.com/articles/s41586-021-03819-2) and the [AlphaFold predictions of the human proteome paper](https://www.nature.com/articles/s41586-021-03828-1), as well as [DeepMind's FAQ](https://alphafold.ebi.ac.uk/faq) on how to interpret AlphaFold/OpenFold predictions. More information about the predictions of the AlphaFold Multimer model can be found in the [Alphafold Multimer paper](https://www.biorxiv.org/content/10.1101/2022.03.11.484043v3.full.pdf)."
]
},
{
@@ -718,7 +852,7 @@
" * Colab CPU runtime is not supported. Try changing your runtime via _Runtime_ > _Change runtime type_ > _Hardware accelerator_ > _GPU_.\n",
" * The type of GPU allocated to your Colab varies. See the [Colab FAQ](https://research.google.com/colaboratory/faq.html) for more details.\n",
" * If you receive “Cannot connect to GPU backend”, you can try again later to see if Colab allocates you a GPU.\n",
- " * [Colab Pro](https://colab.research.google.com/signup) offers priority access to GPUs. \n",
+ " * [Colab Pro](https://colab.research.google.com/signup) offers priority access to GPUs.\n",
"* Does this tool install anything on my computer?\n",
" * No, everything happens in the cloud on Google Colab.\n",
" * At the end of the Colab execution a zip-archive with the obtained prediction will be automatically downloaded to your computer.\n",
@@ -766,13 +900,10 @@
}
],
"metadata": {
- "accelerator": "GPU",
"colab": {
- "collapsed_sections": [],
- "name": "OpenFold.ipynb",
"provenance": [],
"gpuType": "T4",
- "include_colab_link": true
+ "toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
@@ -780,8 +911,9 @@
},
"language_info": {
"name": "python"
- }
+ },
+ "accelerator": "GPU"
},
"nbformat": 4,
"nbformat_minor": 0
-}
+}
\ No newline at end of file
diff --git a/notebooks/environment.yml b/notebooks/environment.yml
index 22266554..8734d79b 100644
--- a/notebooks/environment.yml
+++ b/notebooks/environment.yml
@@ -3,15 +3,15 @@ channels:
- conda-forge
- bioconda
dependencies:
- - conda-forge::openmm=7.5.1
- - conda-forge::pdbfixer
+ - openmm=7.7
+ - pdbfixer
+ - ml-collections
+ - PyYAML==5.4.1
+ - requests
+ - typing-extensions
- bioconda::hmmer==3.3.2
- bioconda::hhsuite==3.3.0
- bioconda::kalign2==2.04
- pip:
- biopython==1.79
- dm-tree==0.1.6
- - ml-collections==0.1.0
- - PyYAML==5.4.1
- - requests==2.26.0
- - typing-extensions==3.10.0.2
diff --git a/openfold/__init__.py b/openfold/__init__.py
index 3a3faee7..86cac764 100644
--- a/openfold/__init__.py
+++ b/openfold/__init__.py
@@ -1,5 +1,6 @@
from . import model
from . import utils
+from . import data
from . import np
from . import resources
diff --git a/openfold/config.py b/openfold/config.py
index dfd231b2..b6259ed0 100644
--- a/openfold/config.py
+++ b/openfold/config.py
@@ -1,3 +1,4 @@
+import re
import copy
import importlib
import ml_collections as mlc
@@ -16,7 +17,7 @@ def string_to_setting(s):
path = s.split('.')
setting = config
for p in path:
- setting = setting[p]
+ setting = setting.get(p)
return setting
@@ -161,44 +162,70 @@ def model_config(
c.model.template.enabled = False
c.model.heads.tm.enabled = True
c.loss.tm.weight = 0.1
- # SINGLE SEQUENCE EMBEDDING PRESETS
- elif name == "seqemb_initial_training":
- c.data.train.max_msa_clusters = 1
- c.data.eval.max_msa_clusters = 1
- c.data.train.block_delete_msa = False
- c.data.train.max_distillation_msa_clusters = 1
- elif name == "seqemb_finetuning":
- c.data.train.max_msa_clusters = 1
- c.data.eval.max_msa_clusters = 1
- c.data.train.block_delete_msa = False
- c.data.train.max_distillation_msa_clusters = 1
- c.data.train.crop_size = 384
- c.loss.violation.weight = 1.
- c.loss.experimentally_resolved.weight = 0.01
- elif name == "seq_model_esm1b":
- c.data.common.use_templates = True
- c.data.common.use_template_torsion_angles = True
- c.model.template.enabled = True
- c.data.predict.max_msa_clusters = 1
- elif name == "seq_model_esm1b_ptm":
- c.data.common.use_templates = True
- c.data.common.use_template_torsion_angles = True
- c.model.template.enabled = True
- c.data.predict.max_msa_clusters = 1
- c.model.heads.tm.enabled = True
- c.loss.tm.weight = 0.1
+ elif name.startswith("seq"): # SINGLE SEQUENCE EMBEDDING PRESETS
+ c.update(seq_mode_config.copy_and_resolve_references())
+ if name == "seqemb_initial_training":
+ c.data.train.max_msa_clusters = 1
+ c.data.eval.max_msa_clusters = 1
+ c.data.train.block_delete_msa = False
+ c.data.train.max_distillation_msa_clusters = 1
+ elif name == "seqemb_finetuning":
+ c.data.train.max_msa_clusters = 1
+ c.data.eval.max_msa_clusters = 1
+ c.data.train.block_delete_msa = False
+ c.data.train.max_distillation_msa_clusters = 1
+ c.data.train.crop_size = 384
+ c.loss.violation.weight = 1.
+ c.loss.experimentally_resolved.weight = 0.01
+ elif name == "seq_model_esm1b":
+ c.data.common.use_templates = True
+ c.data.common.use_template_torsion_angles = True
+ c.model.template.enabled = True
+ c.data.predict.max_msa_clusters = 1
+ elif name == "seq_model_esm1b_ptm":
+ c.data.common.use_templates = True
+ c.data.common.use_template_torsion_angles = True
+ c.model.template.enabled = True
+ c.data.predict.max_msa_clusters = 1
+ c.model.heads.tm.enabled = True
+ c.loss.tm.weight = 0.1
+ elif "multimer" in name: # MULTIMER PRESETS
+ c.update(multimer_config_update.copy_and_resolve_references())
+
+ # Not used in multimer
+ del c.model.template.template_pointwise_attention
+ del c.loss.fape.backbone
+
+ # TODO: Change max_msa_clusters and max_extra_msa to multimer feats within model
+ if re.fullmatch("^model_[1-5]_multimer(_v2)?$", name):
+ #c.model.input_embedder.num_msa = 252
+ #c.model.extra_msa.extra_msa_embedder.num_extra_msa = 1152
+ c.data.train.crop_size = 384
+
+ c.data.train.max_msa_clusters = 252
+ c.data.eval.max_msa_clusters = 252
+ c.data.predict.max_msa_clusters = 252
+
+ c.data.train.max_extra_msa = 1152
+ c.data.eval.max_extra_msa = 1152
+ c.data.predict.max_extra_msa = 1152
+
+ c.model.evoformer_stack.fuse_projection_weights = False
+ c.model.extra_msa.extra_msa_stack.fuse_projection_weights = False
+ c.model.template.template_pair_stack.fuse_projection_weights = False
+ elif name == 'model_4_multimer_v3':
+ #c.model.extra_msa.extra_msa_embedder.num_extra_msa = 1152
+ c.data.train.max_extra_msa = 1152
+ c.data.eval.max_extra_msa = 1152
+ c.data.predict.max_extra_msa = 1152
+ elif name == 'model_5_multimer_v3':
+ #c.model.extra_msa.extra_msa_embedder.num_extra_msa = 1152
+ c.data.train.max_extra_msa = 1152
+ c.data.eval.max_extra_msa = 1152
+ c.data.predict.max_extra_msa = 1152
else:
raise ValueError("Invalid model name")
- if name.startswith("seq"):
- # Tell the data pipeline that we will use sequence embeddings instead of MSAs.
- c.data.seqemb_mode.enabled = True
- c.globals.seqemb_mode_enabled = True
- # In seqemb mode, we turn off the ExtraMSAStack and Evoformer's column attention.
- c.model.extra_msa.enabled = False
- c.model.evoformer_stack.no_column_attention = True
- c.update(seq_mode_config.copy_and_resolve_references())
-
if long_sequence_inference:
assert(not train)
c.globals.offload_inference = True
@@ -380,6 +407,8 @@ def model_config(
"max_templates": 4,
"crop": False,
"crop_size": None,
+ "spatial_crop_prob": None,
+ "interface_threshold": None,
"supervised": False,
"uniform_recycling": False,
},
@@ -394,6 +423,8 @@ def model_config(
"max_templates": 4,
"crop": False,
"crop_size": None,
+ "spatial_crop_prob": None,
+ "interface_threshold": None,
"supervised": True,
"uniform_recycling": False,
},
@@ -409,6 +440,8 @@ def model_config(
"shuffle_top_k_prefiltered": 20,
"crop": True,
"crop_size": 256,
+ "spatial_crop_prob": 0.,
+ "interface_threshold": None,
"supervised": True,
"clamp_prob": 0.9,
"max_distillation_msa_clusters": 1000,
@@ -426,7 +459,6 @@ def model_config(
},
# Recurring FieldReferences that can be changed globally here
"globals": {
- "seqemb_mode_enabled": False, # Global flag for enabling seq emb mode
"blocks_per_ckpt": blocks_per_ckpt,
"chunk_size": chunk_size,
# Use DeepSpeed memory-efficient attention kernel. Mutually
@@ -446,6 +478,8 @@ def model_config(
"c_e": c_e,
"c_s": c_s,
"eps": eps,
+ "is_multimer": False,
+ "seqemb_mode_enabled": False, # Global flag for enabling seq emb mode
},
"model": {
"_mask_trans": False,
@@ -470,7 +504,7 @@ def model_config(
"max_bin": 50.75,
"no_bins": 39,
},
- "template_angle_embedder": {
+ "template_single_embedder": {
# DISCREPANCY: c_in is supposed to be 51.
"c_in": 57,
"c_out": c_m,
@@ -489,6 +523,8 @@ def model_config(
"no_heads": 4,
"pair_transition_n": 2,
"dropout_rate": 0.25,
+ "tri_mul_first": False,
+ "fuse_projection_weights": False,
"blocks_per_ckpt": blocks_per_ckpt,
"tune_chunk_size": tune_chunk_size,
"inf": 1e9,
@@ -537,6 +573,8 @@ def model_config(
"transition_n": 4,
"msa_dropout": 0.15,
"pair_dropout": 0.25,
+ "opm_first": False,
+ "fuse_projection_weights": False,
"clear_cache_between_blocks": False,
"tune_chunk_size": tune_chunk_size,
"inf": 1e9,
@@ -560,6 +598,8 @@ def model_config(
"msa_dropout": 0.15,
"pair_dropout": 0.25,
"no_column_attention": False,
+ "opm_first": False,
+ "fuse_projection_weights": False,
"blocks_per_ckpt": blocks_per_ckpt,
"clear_cache_between_blocks": False,
"tune_chunk_size": tune_chunk_size,
@@ -607,6 +647,12 @@ def model_config(
"c_out": 37,
},
},
+ # A negative value indicates that no early stopping will occur, i.e.
+ # the model will always run `max_recycling_iters` number of recycling
+ # iterations. A positive value will enable early stopping if the
+ # difference in pairwise distances is less than the tolerance between
+ # recycling steps.
+ "recycle_early_stop_tolerance": -1.
},
"relax": {
"max_iterations": 0, # no max
@@ -652,6 +698,7 @@ def model_config(
"weight": 0.01,
},
"masked_msa": {
+ "num_classes": 23,
"eps": eps, # 1e-8,
"weight": 2.0,
},
@@ -664,6 +711,7 @@ def model_config(
"violation": {
"violation_tolerance_factor": 12.0,
"clash_overlap_tolerance": 1.5,
+ "average_clashes": False,
"eps": eps, # 1e-6,
"weight": 0.0,
},
@@ -676,12 +724,199 @@ def model_config(
"weight": 0.,
"enabled": tm_enabled,
},
+ "chain_center_of_mass": {
+ "clamp_distance": -4.0,
+ "weight": 0.,
+ "eps": eps,
+ "enabled": False,
+ },
"eps": eps,
},
"ema": {"decay": 0.999},
}
)
+multimer_config_update = mlc.ConfigDict({
+ "globals": {
+ "is_multimer": True
+ },
+ "data": {
+ "common": {
+ "feat": {
+ "aatype": [NUM_RES],
+ "all_atom_mask": [NUM_RES, None],
+ "all_atom_positions": [NUM_RES, None, None],
+ # "all_chains_entity_ids": [], # TODO: Resolve missing features, remove processed msa feats
+ # "all_crops_all_chains_mask": [],
+ # "all_crops_all_chains_positions": [],
+ # "all_crops_all_chains_residue_ids": [],
+ "assembly_num_chains": [],
+ "asym_id": [NUM_RES],
+ "atom14_atom_exists": [NUM_RES, None],
+ "atom37_atom_exists": [NUM_RES, None],
+ "bert_mask": [NUM_MSA_SEQ, NUM_RES],
+ "cluster_bias_mask": [NUM_MSA_SEQ],
+ "cluster_profile": [NUM_MSA_SEQ, NUM_RES, None],
+ "cluster_deletion_mean": [NUM_MSA_SEQ, NUM_RES],
+ "deletion_matrix": [NUM_MSA_SEQ, NUM_RES],
+ "deletion_mean": [NUM_RES],
+ "entity_id": [NUM_RES],
+ "entity_mask": [NUM_RES],
+ "extra_deletion_matrix": [NUM_EXTRA_SEQ, NUM_RES],
+ "extra_msa": [NUM_EXTRA_SEQ, NUM_RES],
+ "extra_msa_mask": [NUM_EXTRA_SEQ, NUM_RES],
+ # "mem_peak": [],
+ "msa": [NUM_MSA_SEQ, NUM_RES],
+ "msa_feat": [NUM_MSA_SEQ, NUM_RES, None],
+ "msa_mask": [NUM_MSA_SEQ, NUM_RES],
+ "msa_profile": [NUM_RES, None],
+ "num_alignments": [],
+ "num_templates": [],
+ # "queue_size": [],
+ "residue_index": [NUM_RES],
+ "residx_atom14_to_atom37": [NUM_RES, None],
+ "residx_atom37_to_atom14": [NUM_RES, None],
+ "resolution": [],
+ "seq_length": [],
+ "seq_mask": [NUM_RES],
+ "sym_id": [NUM_RES],
+ "target_feat": [NUM_RES, None],
+ "template_aatype": [NUM_TEMPLATES, NUM_RES],
+ "template_all_atom_mask": [NUM_TEMPLATES, NUM_RES, None],
+ "template_all_atom_positions": [
+ NUM_TEMPLATES, NUM_RES, None, None,
+ ],
+ "true_msa": [NUM_MSA_SEQ, NUM_RES]
+ },
+ "max_recycling_iters": 20, # For training, value is 3
+ "unsupervised_features": [
+ "aatype",
+ "residue_index",
+ "msa",
+ "num_alignments",
+ "seq_length",
+ "between_segment_residues",
+ "deletion_matrix",
+ "no_recycling_iters",
+ # Additional multimer features
+ "msa_mask",
+ "seq_mask",
+ "asym_id",
+ "entity_id",
+ "sym_id",
+ ]
+ },
+ "supervised": {
+ "clamp_prob": 1.
+ },
+ # TODO: Change max_msa_clusters and max_extra_msa to multimer feats within model:
+ # c.model.input_embedder.num_msa = 508
+ # c.model.extra_msa.extra_msa_embedder.num_extra_msa = 2048
+ "predict": {
+ "max_msa_clusters": 508,
+ "max_extra_msa": 2048
+ },
+ "eval": {
+ "max_msa_clusters": 508,
+ "max_extra_msa": 2048
+ },
+ "train": {
+ "max_msa_clusters": 508,
+ "max_extra_msa": 2048,
+ "block_delete_msa" : False,
+ "crop_size": 640,
+ "spatial_crop_prob": 0.5,
+ "interface_threshold": 10.,
+ "clamp_prob": 1.,
+ },
+ },
+ "model": {
+ "input_embedder": {
+ "tf_dim": 21,
+ #"num_msa": 508,
+ "max_relative_chain": 2,
+ "max_relative_idx": 32,
+ "use_chain_relative": True
+ },
+ "template": {
+ "template_single_embedder": {
+ "c_in": 34,
+ "c_out": c_m
+ },
+ "template_pair_embedder": {
+ "c_in": c_z,
+ "c_out": c_t,
+ "c_dgram": 39,
+ "c_aatype": 22
+ },
+ "template_pair_stack": {
+ "tri_mul_first": True,
+ "fuse_projection_weights": True
+ },
+ "c_t": c_t,
+ "c_z": c_z,
+ "use_unit_vector": True
+ },
+ "extra_msa": {
+ # "extra_msa_embedder": {
+ # "num_extra_msa": 2048
+ # },
+ "extra_msa_stack": {
+ "opm_first": True,
+ "fuse_projection_weights": True
+ }
+ },
+ "evoformer_stack": {
+ "opm_first": True,
+ "fuse_projection_weights": True
+ },
+ "structure_module": {
+ "trans_scale_factor": 20
+ },
+ "heads": {
+ "tm": {
+ "ptm_weight": 0.2,
+ "iptm_weight": 0.8,
+ "enabled": True
+ },
+ "masked_msa": {
+ "c_out": 22
+ },
+ },
+ "recycle_early_stop_tolerance": 0.5 # For training, value is -1.
+ },
+ "loss": {
+ "fape": {
+ "intra_chain_backbone": {
+ "clamp_distance": 10.0,
+ "loss_unit_distance": 10.0,
+ "weight": 0.5
+ },
+ "interface_backbone": {
+ "clamp_distance": 30.0,
+ "loss_unit_distance": 20.0,
+ "weight": 0.5
+ }
+ },
+ "masked_msa": {
+ "num_classes": 22
+ },
+ "violation": {
+ "average_clashes": True,
+ "weight": 0.03 # Not finetuning
+ },
+ "tm": {
+ "weight": 0.1,
+ "enabled": True
+ },
+ "chain_center_of_mass": {
+ "weight": 0.05,
+ "enabled": True
+ }
+ }
+})
+
+
seq_mode_config = mlc.ConfigDict({
"data": {
"common": {
@@ -700,12 +935,18 @@ def model_config(
"seqemb_mode_enabled": True,
},
"model": {
- "preembedding_embedder": { # Used in sequence embedding mode
+ "preembedding_embedder": { # Used in sequence embedding mode
"tf_dim": 22,
"preembedding_dim": preemb_dim_size,
"c_z": c_z,
"c_m": c_m,
"relpos_k": 32,
},
+ "extra_msa": {
+ "enabled": False # Disable Extra MSA Stack
+ },
+ "evoformer_stack": {
+ "no_column_attention": True # Turn off Evoformer's column attention
+ },
}
-})
\ No newline at end of file
+})
diff --git a/openfold/data/data_modules.py b/openfold/data/data_modules.py
index caafddbb..62f48e4b 100644
--- a/openfold/data/data_modules.py
+++ b/openfold/data/data_modules.py
@@ -4,42 +4,45 @@
import logging
import os
import pickle
-from typing import Optional, Sequence, List, Any
+from typing import Optional, Sequence, Any, Union
import ml_collections as mlc
import pytorch_lightning as pl
import torch
from torch.utils.data import RandomSampler
-
+from openfold.np.residue_constants import restypes
from openfold.data import (
data_pipeline,
feature_pipeline,
mmcif_parsing,
templates,
)
-from openfold.utils.tensor_utils import tensor_tree_map, dict_multimap
+from openfold.utils.tensor_utils import dict_multimap
+from openfold.utils.tensor_utils import (
+ tensor_tree_map,
+)
class OpenFoldSingleDataset(torch.utils.data.Dataset):
def __init__(self,
- data_dir: str,
- alignment_dir: str,
- template_mmcif_dir: str,
- max_template_date: str,
- config: mlc.ConfigDict,
- chain_data_cache_path: Optional[str] = None,
- kalign_binary_path: str = '/usr/bin/kalign',
- max_template_hits: int = 4,
- obsolete_pdbs_file_path: Optional[str] = None,
- template_release_dates_cache_path: Optional[str] = None,
- shuffle_top_k_prefiltered: Optional[int] = None,
- treat_pdb_as_distillation: bool = True,
- filter_path: Optional[str] = None,
- mode: str = "train",
- alignment_index: Optional[Any] = None,
- _output_raw: bool = False,
- _structure_index: Optional[Any] = None,
- ):
+ data_dir: str,
+ alignment_dir: str,
+ template_mmcif_dir: str,
+ max_template_date: str,
+ config: mlc.ConfigDict,
+ chain_data_cache_path: Optional[str] = None,
+ kalign_binary_path: str = '/usr/bin/kalign',
+ max_template_hits: int = 4,
+ obsolete_pdbs_file_path: Optional[str] = None,
+ template_release_dates_cache_path: Optional[str] = None,
+ shuffle_top_k_prefiltered: Optional[int] = None,
+ treat_pdb_as_distillation: bool = True,
+ filter_path: Optional[str] = None,
+ mode: str = "train",
+ alignment_index: Optional[Any] = None,
+ _output_raw: bool = False,
+ _structure_index: Optional[Any] = None,
+ ):
"""
Args:
data_dir:
@@ -101,21 +104,21 @@ def __init__(self,
self.supported_exts = [".cif", ".core", ".pdb"]
valid_modes = ["train", "eval", "predict"]
- if(mode not in valid_modes):
+ if mode not in valid_modes:
raise ValueError(f'mode must be one of {valid_modes}')
- if(template_release_dates_cache_path is None):
+ if template_release_dates_cache_path is None:
logging.warning(
"Template release dates cache does not exist. Remember to run "
"scripts/generate_mmcif_cache.py before running OpenFold"
)
- if(alignment_index is not None):
+ if alignment_index is not None:
self._chain_ids = list(alignment_index.keys())
else:
self._chain_ids = list(os.listdir(alignment_dir))
- if(filter_path is not None):
+ if filter_path is not None:
with open(filter_path, "r") as f:
chains_to_include = set([l.strip() for l in f.readlines()])
@@ -145,12 +148,15 @@ def __init__(self,
len(missing),
missing_examples,
chain_data_cache_path)
-
+
self._chain_id_to_idx_dict = {
chain: i for i, chain in enumerate(self._chain_ids)
}
- template_featurizer = templates.TemplateHitFeaturizer(
+ # If it's running template search for a monomer, then use hhsearch
+ # as demonstrated in AlphaFold's run_alphafold.py code
+ # https://github.com/deepmind/alphafold/blob/6c4d833fbd1c6b8e7c9a21dae5d4ada2ce777e10/run_alphafold.py#L462C1-L477
+ template_featurizer = templates.HhsearchHitFeaturizer(
mmcif_dir=template_mmcif_dir,
max_template_date=max_template_date,
max_hits=max_template_hits,
@@ -164,8 +170,8 @@ def __init__(self,
template_featurizer=template_featurizer,
)
- if(not self._output_raw):
- self.feature_pipeline = feature_pipeline.FeaturePipeline(config)
+ if not self._output_raw:
+ self.feature_pipeline = feature_pipeline.FeaturePipeline(config)
def _parse_mmcif(self, path, file_id, chain_id, alignment_dir, alignment_index):
with open(path, 'r') as f:
@@ -177,7 +183,7 @@ def _parse_mmcif(self, path, file_id, chain_id, alignment_dir, alignment_index):
# Crash if an error is encountered. Any parsing errors should have
# been dealt with at the alignment stage.
- if(mmcif_object.mmcif_object is None):
+ if mmcif_object.mmcif_object is None:
raise list(mmcif_object.errors.values())[0]
mmcif_object = mmcif_object.mmcif_object
@@ -203,48 +209,47 @@ def __getitem__(self, idx):
alignment_dir = os.path.join(self.alignment_dir, name)
alignment_index = None
- if(self.alignment_index is not None):
+ if self.alignment_index is not None:
alignment_dir = self.alignment_dir
alignment_index = self.alignment_index[name]
- if(self.mode == 'train' or self.mode == 'eval'):
+ if self.mode == 'train' or self.mode == 'eval':
spl = name.rsplit('_', 1)
- if(len(spl) == 2):
+ if len(spl) == 2:
file_id, chain_id = spl
else:
file_id, = spl
chain_id = None
path = os.path.join(self.data_dir, file_id)
- structure_index_entry = None
- if(self._structure_index is not None):
+ if self._structure_index is not None:
structure_index_entry = self._structure_index[name]
- assert(len(structure_index_entry["files"]) == 1)
+ assert (len(structure_index_entry["files"]) == 1)
filename, _, _ = structure_index_entry["files"][0]
ext = os.path.splitext(filename)[1]
else:
ext = None
for e in self.supported_exts:
- if(os.path.exists(path + e)):
+ if os.path.exists(path + e):
ext = e
break
- if(ext is None):
+ if ext is None:
raise ValueError("Invalid file type")
path += ext
- if(ext == ".cif"):
+ if ext == ".cif":
data = self._parse_mmcif(
path, file_id, chain_id, alignment_dir, alignment_index,
)
- elif(ext == ".core"):
+ elif ext == ".core":
data = self.data_pipeline.process_core(
path, alignment_dir, alignment_index,
seqemb_mode=self.config.seqemb_mode.enabled,
)
- elif(ext == ".pdb"):
+ elif ext == ".pdb":
structure_index = None
- if(self._structure_index is not None):
+ if self._structure_index is not None:
structure_index = self._structure_index[name]
data = self.data_pipeline.process_pdb(
pdb_path=path,
@@ -256,7 +261,7 @@ def __getitem__(self, idx):
seqemb_mode=self.config.seqemb_mode.enabled,
)
else:
- raise ValueError("Extension branch missing")
+ raise ValueError("Extension branch missing")
else:
path = os.path.join(name, name + ".fasta")
data = self.data_pipeline.process_fasta(
@@ -266,11 +271,11 @@ def __getitem__(self, idx):
seqemb_mode=self.config.seqemb_mode.enabled,
)
- if(self._output_raw):
+ if self._output_raw:
return data
feats = self.feature_pipeline.process_features(
- data, self.mode
+ data, self.mode
)
feats["batch_idx"] = torch.tensor(
@@ -281,51 +286,250 @@ def __getitem__(self, idx):
return feats
def __len__(self):
- return len(self._chain_ids)
+ return len(self._chain_ids)
-def deterministic_train_filter(
- chain_data_cache_entry: Any,
- max_resolution: float = 9.,
- max_single_aa_prop: float = 0.8,
-) -> bool:
- # Hard filters
- resolution = chain_data_cache_entry.get("resolution", None)
- if(resolution is not None and resolution > max_resolution):
- return False
+class OpenFoldSingleMultimerDataset(torch.utils.data.Dataset):
+ def __init__(self,
+ data_dir: str,
+ alignment_dir: str,
+ template_mmcif_dir: str,
+ max_template_date: str,
+ config: mlc.ConfigDict,
+ mmcif_data_cache_path: Optional[str] = None,
+ kalign_binary_path: str = '/usr/bin/kalign',
+ max_template_hits: int = 4,
+ obsolete_pdbs_file_path: Optional[str] = None,
+ template_release_dates_cache_path: Optional[str] = None,
+ shuffle_top_k_prefiltered: Optional[int] = None,
+ treat_pdb_as_distillation: bool = True,
+ filter_path: Optional[str] = None,
+ mode: str = "train",
+ alignment_index: Optional[Any] = None,
+ _output_raw: bool = False,
+ _structure_index: Optional[Any] = None,
+ ):
+ """
+ This class check each individual PDB ID and return its chain(s) features/ground truth
+ Args:
+ data_dir:
+ A path to a directory containing mmCIF files (in train
+ mode) or FASTA files (in inference mode).
+ alignment_dir:
+ A path to a directory containing only data in the format
+ output by an AlignmentRunner
+ (defined in openfold.features.alignment_runner).
+ I.e. a directory of directories named {PDB_ID}_{CHAIN_ID}
+ or simply {PDB_ID}, each containing .a3m, .sto, and .hhr
+ files.
+ template_mmcif_dir:
+ Path to a directory containing template mmCIF files.
+ config:
+ A dataset config object. See openfold.config
+ mmcif_data_cache_path:
+ Path to cache of all mmcifs files generated by
+ scripts/generate_mmcif_cache.py It should be a json file which records
+ what PDB ID contains which chain(s)
+ kalign_binary_path:
+ Path to kalign binary.
+ max_template_hits:
+ An upper bound on how many templates are considered. During
+ training, the templates ultimately used are subsampled
+ from this total quantity.
+ template_release_dates_cache_path:
+ Path to the output of scripts/generate_mmcif_cache.
+ obsolete_pdbs_file_path:
+ Path to the file containing replacements for obsolete PDBs.
+ shuffle_top_k_prefiltered:
+ Whether to uniformly shuffle the top k template hits before
+ parsing max_template_hits of them. Can be used to
+ approximate DeepMind's training-time template subsampling
+ scheme much more performantly.
+ treat_pdb_as_distillation:
+ Whether to assume that .pdb files in the data_dir are from
+ the self-distillation set (and should be subjected to
+ special distillation set preprocessing steps).
+ mode:
+ "train", "val", or "predict"
+ """
+ super(OpenFoldSingleMultimerDataset, self).__init__()
+ self.data_dir = data_dir
+ self.mmcif_data_cache_path = mmcif_data_cache_path
- seq = chain_data_cache_entry["seq"]
- counts = {}
- for aa in seq:
- counts.setdefault(aa, 0)
- counts[aa] += 1
- largest_aa_count = max(counts.values())
- largest_single_aa_prop = largest_aa_count / len(seq)
- if(largest_single_aa_prop > max_single_aa_prop):
- return False
+ if self.mmcif_data_cache_path is not None:
+ with open(self.mmcif_data_cache_path, "r") as infile:
+ self.mmcif_data_cache = json.load(infile)
+ assert isinstance(self.mmcif_data_cache, dict)
- return True
+ self.alignment_dir = alignment_dir
+ self.config = config
+ self.treat_pdb_as_distillation = treat_pdb_as_distillation
+ self.mode = mode
+ self.alignment_index = alignment_index
+ self._output_raw = _output_raw
+ self._structure_index = _structure_index
+ self.supported_exts = [".cif", ".core", ".pdb"]
-def get_stochastic_train_filter_prob(
- chain_data_cache_entry: Any,
-) -> List[float]:
- # Stochastic filters
- probabilities = []
-
- cluster_size = chain_data_cache_entry.get("cluster_size", None)
- if(cluster_size is not None and cluster_size > 0):
- probabilities.append(1 / cluster_size)
-
- chain_length = len(chain_data_cache_entry["seq"])
- probabilities.append((1 / 512) * (max(min(chain_length, 512), 256)))
+ valid_modes = ["train", "eval", "predict"]
+ if mode not in valid_modes:
+ raise ValueError(f'mode must be one of {valid_modes}')
- # Risk of underflow here?
- out = 1
- for p in probabilities:
- out *= p
+ if template_release_dates_cache_path is None:
+ logging.warning(
+ "Template release dates cache does not exist. Remember to run "
+ "scripts/generate_mmcif_cache.py before running OpenFold"
+ )
- return out
+ if self.mmcif_data_cache_path is not None:
+ self._mmcifs = list(self.mmcif_data_cache.keys())
+ elif self.alignment_index is not None:
+ self._mmcifs = [i.split("_")[0] for i in list(alignment_index.keys())]
+ elif self.alignment_dir is not None:
+ self._mmcifs = [i.split("_")[0] for i in os.listdir(self.alignment_dir)]
+ else:
+ raise ValueError("You must provide at least one of the mmcif_data_cache or alignment_dir")
+
+ if filter_path is not None:
+ with open(filter_path, "r") as f:
+ mmcifs_to_include = set([l.strip() for l in f.readlines()])
+
+ self._mmcifs = [
+ m for m in self._mmcifs if m in mmcifs_to_include
+ ]
+
+ self._mmcif_id_to_idx_dict = {
+ mmcif: i for i, mmcif in enumerate(self._mmcifs)
+ }
+
+ template_featurizer = templates.HmmsearchHitFeaturizer(
+ mmcif_dir=template_mmcif_dir,
+ max_template_date=max_template_date,
+ max_hits=max_template_hits,
+ kalign_binary_path=kalign_binary_path,
+ release_dates_path=template_release_dates_cache_path,
+ obsolete_pdbs_path=obsolete_pdbs_file_path,
+ _shuffle_top_k_prefiltered=shuffle_top_k_prefiltered,
+ )
+
+ data_processor = data_pipeline.DataPipeline(
+ template_featurizer=template_featurizer,
+ )
+ self.data_pipeline = data_pipeline.DataPipelineMultimer(
+ monomer_data_pipeline=data_processor
+ )
+ self.feature_pipeline = feature_pipeline.FeaturePipeline(config)
+
+ def _parse_mmcif(self, path, file_id, alignment_dir, alignment_index):
+ with open(path, 'r') as f:
+ mmcif_string = f.read()
+
+ mmcif_object = mmcif_parsing.parse(
+ file_id=file_id, mmcif_string=mmcif_string
+ )
+
+ # Crash if an error is encountered. Any parsing errors should have
+ # been dealt with at the alignment stage.
+ if mmcif_object.mmcif_object is None:
+ raise list(mmcif_object.errors.values())[0]
+
+ mmcif_object = mmcif_object.mmcif_object
+
+ data = self.data_pipeline.process_mmcif(
+ mmcif=mmcif_object,
+ alignment_dir=alignment_dir,
+ alignment_index=alignment_index
+ )
+
+ return data
+
+ def mmcif_id_to_idx(self, mmcif_id):
+ return self._mmcif_id_to_idx_dict[mmcif_id]
+
+ def idx_to_mmcif_id(self, idx):
+ return self._mmcifs[idx]
+
+ def __getitem__(self, idx):
+ mmcif_id = self.idx_to_mmcif_id(idx)
+
+ alignment_index = None
+ if self.alignment_index is not None:
+ alignment_index = {k: v for k, v in self.alignment_index.items()
+ if f'{mmcif_id}_' in k}
+
+ if self.mode == 'train' or self.mode == 'eval':
+ path = os.path.join(self.data_dir, f"{mmcif_id}")
+ ext = None
+ for e in self.supported_exts:
+ if os.path.exists(path + e):
+ ext = e
+ break
+
+ if ext is None:
+ raise ValueError("Invalid file type")
+
+ # TODO: Add pdb and core exts to data_pipeline for multimer
+ path += ext
+ if ext == ".cif":
+ data = self._parse_mmcif(
+ path, mmcif_id, self.alignment_dir, alignment_index,
+ )
+ else:
+ raise ValueError("Extension branch missing")
+ else:
+ path = os.path.join(self.data_dir, f"{mmcif_id}.fasta")
+ data = self.data_pipeline.process_fasta(
+ fasta_path=path,
+ alignment_dir=self.alignment_dir,
+ alignment_index=alignment_index
+ )
+
+ if self._output_raw:
+ return data
+
+ # process all_chain_features
+ data = self.feature_pipeline.process_features(data,
+ mode=self.mode,
+ is_multimer=True)
+
+ # if it's inference mode, only need all_chain_features
+ data["batch_idx"] = torch.tensor(
+ [idx for _ in range(data["aatype"].shape[-1])],
+ dtype=torch.int64,
+ device=data["aatype"].device)
+
+ return data
+
+ def __len__(self):
+ return len(self._mmcifs)
+
+
+def resolution_filter(resolution: int, max_resolution: float) -> bool:
+ """Check that the resolution is <= max_resolution permitted"""
+ return resolution is not None and resolution <= max_resolution
+
+
+def aa_count_filter(seqs: list, max_single_aa_prop: float) -> bool:
+ """Check if any single amino acid accounts for more than max_single_aa_prop percent of the sequence(s)"""
+ counts = {}
+ for seq in seqs:
+ for aa in seq:
+ counts.setdefault(aa, 0)
+ if aa not in restypes:
+ return False
+ else:
+ counts[aa] += 1
+
+ total_len = sum([len(i) for i in seqs])
+ largest_aa_count = max(counts.values())
+ largest_single_aa_prop = largest_aa_count / total_len
+ return largest_single_aa_prop <= max_single_aa_prop
+
+
+def all_seq_len_filter(seqs: list, minimum_number_of_residues: int) -> bool:
+ """Check if the total combined sequence lengths are >= minimum_numer_of_residues"""
+ total_len = sum([len(i) for i in seqs])
+ return total_len >= minimum_number_of_residues
class OpenFoldDataset(torch.utils.data.Dataset):
@@ -335,68 +539,106 @@ class OpenFoldDataset(torch.utils.data.Dataset):
length of an OpenFoldFilteredDataset is arbitrary. Samples are selected
and filtered once at initialization.
"""
+
def __init__(self,
- datasets: Sequence[OpenFoldSingleDataset],
- probabilities: Sequence[float],
- epoch_len: int,
- generator: torch.Generator = None,
- _roll_at_init: bool = True,
- ):
+ datasets: Union[Sequence[OpenFoldSingleDataset], Sequence[OpenFoldSingleMultimerDataset]],
+ probabilities: Sequence[float],
+ epoch_len: int,
+ generator: torch.Generator = None,
+ _roll_at_init: bool = True,
+ ):
self.datasets = datasets
self.probabilities = probabilities
self.epoch_len = epoch_len
self.generator = generator
- def looped_shuffled_dataset_idx(dataset_len):
- while True:
- # Uniformly shuffle each dataset's indices
- weights = [1. for _ in range(dataset_len)]
- shuf = torch.multinomial(
- torch.tensor(weights),
- num_samples=dataset_len,
- replacement=False,
- generator=self.generator,
- )
- for idx in shuf:
- yield idx
-
- def looped_samples(dataset_idx):
- max_cache_len = int(epoch_len * probabilities[dataset_idx])
- dataset = self.datasets[dataset_idx]
- idx_iter = looped_shuffled_dataset_idx(len(dataset))
- chain_data_cache = dataset.chain_data_cache
- while True:
- weights = []
- idx = []
- for _ in range(max_cache_len):
- candidate_idx = next(idx_iter)
- chain_id = dataset.idx_to_chain_id(candidate_idx)
- chain_data_cache_entry = chain_data_cache[chain_id]
- if(not deterministic_train_filter(chain_data_cache_entry)):
- continue
-
- p = get_stochastic_train_filter_prob(
- chain_data_cache_entry,
- )
- weights.append([1. - p, p])
- idx.append(candidate_idx)
-
- samples = torch.multinomial(
- torch.tensor(weights),
- num_samples=1,
- generator=self.generator,
- )
- samples = samples.squeeze()
+ self._samples = [self.looped_samples(i) for i in range(len(self.datasets))]
+ if _roll_at_init:
+ self.reroll()
- cache = [i for i, s in zip(idx, samples) if s]
+ @staticmethod
+ def deterministic_train_filter(
+ cache_entry: Any,
+ max_resolution: float = 9.,
+ max_single_aa_prop: float = 0.8,
+ *args, **kwargs
+ ) -> bool:
+ # Hard filters
+ resolution = cache_entry.get("resolution", None)
+ seqs = [cache_entry["seq"]]
+
+ return all([resolution_filter(resolution=resolution,
+ max_resolution=max_resolution),
+ aa_count_filter(seqs=seqs,
+ max_single_aa_prop=max_single_aa_prop)])
+
+ @staticmethod
+ def get_stochastic_train_filter_prob(
+ cache_entry: Any,
+ *args, **kwargs
+ ) -> float:
+ # Stochastic filters
+ probabilities = []
+
+ cluster_size = cache_entry.get("cluster_size", None)
+ if cluster_size is not None and cluster_size > 0:
+ probabilities.append(1 / cluster_size)
+
+ chain_length = len(cache_entry["seq"])
+ probabilities.append((1 / 512) * (max(min(chain_length, 512), 256)))
+
+ # Risk of underflow here?
+ out = 1
+ for p in probabilities:
+ out *= p
+
+ return out
+
+ def looped_shuffled_dataset_idx(self, dataset_len):
+ while True:
+ # Uniformly shuffle each dataset's indices
+ weights = [1. for _ in range(dataset_len)]
+ shuf = torch.multinomial(
+ torch.tensor(weights),
+ num_samples=dataset_len,
+ replacement=False,
+ generator=self.generator,
+ )
+ for idx in shuf:
+ yield idx
+
+ def looped_samples(self, dataset_idx):
+ max_cache_len = int(self.epoch_len * self.probabilities[dataset_idx])
+ dataset = self.datasets[dataset_idx]
+ idx_iter = self.looped_shuffled_dataset_idx(len(dataset))
+ chain_data_cache = dataset.chain_data_cache
+ while True:
+ weights = []
+ idx = []
+ for _ in range(max_cache_len):
+ candidate_idx = next(idx_iter)
+ chain_id = dataset.idx_to_chain_id(candidate_idx)
+ chain_data_cache_entry = chain_data_cache[chain_id]
+ if not self.deterministic_train_filter(chain_data_cache_entry):
+ continue
+
+ p = self.get_stochastic_train_filter_prob(
+ chain_data_cache_entry,
+ )
+ weights.append([1. - p, p])
+ idx.append(candidate_idx)
- for datapoint_idx in cache:
- yield datapoint_idx
+ samples = torch.multinomial(
+ torch.tensor(weights),
+ num_samples=1,
+ generator=self.generator,
+ )
+ samples = samples.squeeze()
- self._samples = [looped_samples(i) for i in range(len(self.datasets))]
+ cache = [i for i, s in zip(idx, samples) if s]
- if(_roll_at_init):
- self.reroll()
+ for datapoint_idx in cache:
+ yield datapoint_idx
def __getitem__(self, idx):
dataset_idx, datapoint_idx = self.datapoints[idx]
@@ -412,7 +654,6 @@ def reroll(self):
replacement=True,
generator=self.generator,
)
-
self.datapoints = []
for dataset_idx in dataset_choices:
samples = self._samples[dataset_idx]
@@ -420,10 +661,102 @@ def reroll(self):
self.datapoints.append((dataset_idx, datapoint_idx))
+class OpenFoldMultimerDataset(OpenFoldDataset):
+ """
+ Create a torch Dataset object for multimer training and
+ add filtering steps described in AlphaFold Multimer's paper:
+ https://www.biorxiv.org/content/10.1101/2021.10.04.463034v2.full.pdf Supplementary section 7.1
+ """
+
+ def __init__(self,
+ datasets: Sequence[OpenFoldSingleMultimerDataset],
+ probabilities: Sequence[float],
+ epoch_len: int,
+ generator: torch.Generator = None,
+ _roll_at_init: bool = True
+ ):
+ super(OpenFoldMultimerDataset, self).__init__(datasets=datasets,
+ probabilities=probabilities,
+ epoch_len=epoch_len,
+ generator=generator,
+ _roll_at_init=_roll_at_init)
+
+ @staticmethod
+ def deterministic_train_filter(
+ cache_entry: Any,
+ is_distillation: bool,
+ max_resolution: float = 9.,
+ max_single_aa_prop: float = 0.8,
+ minimum_number_of_residues: int = 200,
+ *args, **kwargs
+ ) -> bool:
+ """
+ Implement multimer training filtering criteria described in
+ https://www.biorxiv.org/content/10.1101/2021.10.04.463034v2.full.pdf Supplementary section 7.1
+ """
+ resolution = cache_entry.get("resolution", None)
+ seqs = cache_entry["seqs"]
+
+ return all([resolution_filter(resolution=resolution,
+ max_resolution=max_resolution),
+ aa_count_filter(seqs=seqs,
+ max_single_aa_prop=max_single_aa_prop),
+ (not is_distillation or all_seq_len_filter(seqs=seqs,
+ minimum_number_of_residues=minimum_number_of_residues))])
+
+ @staticmethod
+ def get_stochastic_train_filter_prob(
+ cache_entry: Any,
+ *args, **kwargs
+ ) -> list:
+ # Stochastic filters
+ cluster_sizes = cache_entry.get("cluster_sizes")
+ if cluster_sizes is not None:
+ return [1 / c if c > 0 else 1 for c in cluster_sizes]
+
+ num_chains = len(cache_entry["chain_ids"])
+ return [1.] * num_chains
+
+ def looped_samples(self, dataset_idx):
+ max_cache_len = int(self.epoch_len * self.probabilities[dataset_idx])
+ dataset = self.datasets[dataset_idx]
+ is_distillation = dataset.treat_pdb_as_distillation
+ idx_iter = self.looped_shuffled_dataset_idx(len(dataset))
+ mmcif_data_cache = dataset.mmcif_data_cache
+ while True:
+ weights = []
+ idx = []
+ for _ in range(max_cache_len):
+ candidate_idx = next(idx_iter)
+ mmcif_id = dataset.idx_to_mmcif_id(candidate_idx)
+ mmcif_data_cache_entry = mmcif_data_cache[mmcif_id]
+ if not self.deterministic_train_filter(cache_entry=mmcif_data_cache_entry,
+ is_distillation=is_distillation):
+ continue
+
+ chain_probs = self.get_stochastic_train_filter_prob(
+ mmcif_data_cache_entry,
+ )
+ weights.extend([[1. - p, p] for p in chain_probs])
+ idx.extend([candidate_idx] * len(chain_probs))
+
+ samples = torch.multinomial(
+ torch.tensor(weights),
+ num_samples=1,
+ generator=self.generator,
+ )
+ samples = samples.squeeze()
+
+ cache = [i for i, s in zip(idx, samples) if s]
+
+ for datapoint_idx in cache:
+ yield datapoint_idx
+
+
class OpenFoldBatchCollator:
def __call__(self, prots):
stack_fn = partial(torch.stack, dim=0)
- return dict_multimap(stack_fn, prots)
+ return dict_multimap(stack_fn, prots)
class OpenFoldDataLoader(torch.utils.data.DataLoader):
@@ -439,8 +772,8 @@ def _prep_batch_properties_probs(self):
stage_cfg = self.config[self.stage]
max_iters = self.config.common.max_recycling_iters
-
- if(stage_cfg.uniform_recycling):
+
+ if stage_cfg.uniform_recycling:
recycling_probs = [
1. / (max_iters + 1) for _ in range(max_iters + 1)
]
@@ -449,15 +782,15 @@ def _prep_batch_properties_probs(self):
0. for _ in range(max_iters + 1)
]
recycling_probs[-1] = 1.
-
+
keyed_probs.append(
("no_recycling_iters", recycling_probs)
)
keys, probs = zip(*keyed_probs)
max_len = max([len(p) for p in probs])
- padding = [[0.] * (max_len - len(p)) for p in probs]
-
+ padding = [[0.] * (max_len - len(p)) for p in probs]
+
self.prop_keys = keys
self.prop_probs_tensor = torch.tensor(
[p + pad for p, pad in zip(probs, padding)],
@@ -465,9 +798,10 @@ def _prep_batch_properties_probs(self):
)
def _add_batch_properties(self, batch):
+ gt_features = batch.pop('gt_features', None)
samples = torch.multinomial(
self.prop_probs_tensor,
- num_samples=1, # 1 per row
+ num_samples=1, # 1 per row
replacement=True,
generator=self.generator
)
@@ -479,8 +813,8 @@ def _add_batch_properties(self, batch):
for i, key in enumerate(self.prop_keys):
sample = int(samples[i][0])
sample_tensor = torch.tensor(
- sample,
- device=aatype.device,
+ sample,
+ device=aatype.device,
requires_grad=False
)
orig_shape = sample_tensor.shape
@@ -492,11 +826,12 @@ def _add_batch_properties(self, batch):
)
batch[key] = sample_tensor
- if(key == "no_recycling_iters"):
- no_recycling = sample
-
+ if key == "no_recycling_iters":
+ no_recycling = sample
+
resample_recycling = lambda t: t[..., :no_recycling + 1]
batch = tensor_tree_map(resample_recycling, batch)
+ batch['gt_features'] = gt_features
return batch
@@ -512,31 +847,31 @@ def _batch_prop_gen(iterator):
class OpenFoldDataModule(pl.LightningDataModule):
def __init__(self,
- config: mlc.ConfigDict,
- template_mmcif_dir: str,
- max_template_date: str,
- train_data_dir: Optional[str] = None,
- train_alignment_dir: Optional[str] = None,
- train_chain_data_cache_path: Optional[str] = None,
- distillation_data_dir: Optional[str] = None,
- distillation_alignment_dir: Optional[str] = None,
- distillation_chain_data_cache_path: Optional[str] = None,
- val_data_dir: Optional[str] = None,
- val_alignment_dir: Optional[str] = None,
- predict_data_dir: Optional[str] = None,
- predict_alignment_dir: Optional[str] = None,
- kalign_binary_path: str = '/usr/bin/kalign',
- train_filter_path: Optional[str] = None,
- distillation_filter_path: Optional[str] = None,
- obsolete_pdbs_file_path: Optional[str] = None,
- template_release_dates_cache_path: Optional[str] = None,
- batch_seed: Optional[int] = None,
- train_epoch_len: int = 50000,
- _distillation_structure_index_path: Optional[str] = None,
- alignment_index_path: Optional[str] = None,
- distillation_alignment_index_path: Optional[str] = None,
- **kwargs
- ):
+ config: mlc.ConfigDict,
+ template_mmcif_dir: str,
+ max_template_date: str,
+ train_data_dir: Optional[str] = None,
+ train_alignment_dir: Optional[str] = None,
+ train_chain_data_cache_path: Optional[str] = None,
+ distillation_data_dir: Optional[str] = None,
+ distillation_alignment_dir: Optional[str] = None,
+ distillation_chain_data_cache_path: Optional[str] = None,
+ val_data_dir: Optional[str] = None,
+ val_alignment_dir: Optional[str] = None,
+ predict_data_dir: Optional[str] = None,
+ predict_alignment_dir: Optional[str] = None,
+ kalign_binary_path: str = '/usr/bin/kalign',
+ train_filter_path: Optional[str] = None,
+ distillation_filter_path: Optional[str] = None,
+ obsolete_pdbs_file_path: Optional[str] = None,
+ template_release_dates_cache_path: Optional[str] = None,
+ batch_seed: Optional[int] = None,
+ train_epoch_len: int = 50000,
+ _distillation_structure_index_path: Optional[str] = None,
+ alignment_index_path: Optional[str] = None,
+ distillation_alignment_index_path: Optional[str] = None,
+ **kwargs
+ ):
super(OpenFoldDataModule, self).__init__()
self.config = config
@@ -564,7 +899,7 @@ def __init__(self,
self.batch_seed = batch_seed
self.train_epoch_len = train_epoch_len
- if(self.train_data_dir is None and self.predict_data_dir is None):
+ if self.train_data_dir is None and self.predict_data_dir is None:
raise ValueError(
'At least one of train_data_dir or predict_data_dir must be '
'specified'
@@ -572,65 +907,61 @@ def __init__(self,
self.training_mode = self.train_data_dir is not None
- if(self.training_mode and train_alignment_dir is None):
+ if self.training_mode and train_alignment_dir is None:
raise ValueError(
'In training mode, train_alignment_dir must be specified'
)
- elif(not self.training_mode and predict_alignment_dir is None):
+ elif not self.training_mode and predict_alignment_dir is None:
raise ValueError(
'In inference mode, predict_alignment_dir must be specified'
- )
- elif(val_data_dir is not None and val_alignment_dir is None):
+ )
+ elif val_data_dir is not None and val_alignment_dir is None:
raise ValueError(
'If val_data_dir is specified, val_alignment_dir must '
'be specified as well'
- )
+ )
# An ad-hoc measure for our particular filesystem restrictions
self._distillation_structure_index = None
- if(_distillation_structure_index_path is not None):
+ if _distillation_structure_index_path is not None:
with open(_distillation_structure_index_path, "r") as fp:
self._distillation_structure_index = json.load(fp)
-
+
self.alignment_index = None
- if(alignment_index_path is not None):
+ if alignment_index_path is not None:
with open(alignment_index_path, "r") as fp:
self.alignment_index = json.load(fp)
self.distillation_alignment_index = None
- if(distillation_alignment_index_path is not None):
+ if distillation_alignment_index_path is not None:
with open(distillation_alignment_index_path, "r") as fp:
self.distillation_alignment_index = json.load(fp)
def setup(self):
# Most of the arguments are the same for the three datasets
dataset_gen = partial(OpenFoldSingleDataset,
- template_mmcif_dir=self.template_mmcif_dir,
- max_template_date=self.max_template_date,
- config=self.config,
- kalign_binary_path=self.kalign_binary_path,
- template_release_dates_cache_path=
- self.template_release_dates_cache_path,
- obsolete_pdbs_file_path=
- self.obsolete_pdbs_file_path,
- )
-
- if(self.training_mode):
+ template_mmcif_dir=self.template_mmcif_dir,
+ max_template_date=self.max_template_date,
+ config=self.config,
+ kalign_binary_path=self.kalign_binary_path,
+ template_release_dates_cache_path=self.template_release_dates_cache_path,
+ obsolete_pdbs_file_path=self.obsolete_pdbs_file_path)
+
+ if self.training_mode:
train_dataset = dataset_gen(
data_dir=self.train_data_dir,
chain_data_cache_path=self.train_chain_data_cache_path,
alignment_dir=self.train_alignment_dir,
filter_path=self.train_filter_path,
max_template_hits=self.config.train.max_template_hits,
- shuffle_top_k_prefiltered=
- self.config.train.shuffle_top_k_prefiltered,
+ shuffle_top_k_prefiltered=self.config.train.shuffle_top_k_prefiltered,
treat_pdb_as_distillation=False,
mode="train",
alignment_index=self.alignment_index,
)
distillation_dataset = None
- if(self.distillation_data_dir is not None):
+ if self.distillation_data_dir is not None:
distillation_dataset = dataset_gen(
data_dir=self.distillation_data_dir,
chain_data_cache_path=self.distillation_chain_data_cache_path,
@@ -644,8 +975,8 @@ def setup(self):
)
d_prob = self.config.train.distillation_prob
-
- if(distillation_dataset is not None):
+
+ if distillation_dataset is not None:
datasets = [train_dataset, distillation_dataset]
d_prob = self.config.train.distillation_prob
probabilities = [1. - d_prob, d_prob]
@@ -654,10 +985,10 @@ def setup(self):
probabilities = [1.]
generator = None
- if(self.batch_seed is not None):
+ if self.batch_seed is not None:
generator = torch.Generator()
generator = generator.manual_seed(self.batch_seed + 1)
-
+
self.train_dataset = OpenFoldDataset(
datasets=datasets,
probabilities=probabilities,
@@ -665,8 +996,8 @@ def setup(self):
generator=generator,
_roll_at_init=False,
)
-
- if(self.val_data_dir is not None):
+
+ if self.val_data_dir is not None:
self.eval_dataset = dataset_gen(
data_dir=self.val_data_dir,
alignment_dir=self.val_alignment_dir,
@@ -676,7 +1007,7 @@ def setup(self):
)
else:
self.eval_dataset = None
- else:
+ else:
self.predict_dataset = dataset_gen(
data_dir=self.predict_data_dir,
alignment_dir=self.predict_alignment_dir,
@@ -687,18 +1018,17 @@ def setup(self):
def _gen_dataloader(self, stage):
generator = None
- if(self.batch_seed is not None):
+ if self.batch_seed is not None:
generator = torch.Generator()
generator = generator.manual_seed(self.batch_seed)
- dataset = None
- if(stage == "train"):
+ if stage == "train":
dataset = self.train_dataset
# Filter the dataset, if necessary
dataset.reroll()
- elif(stage == "eval"):
+ elif stage == "eval":
dataset = self.eval_dataset
- elif(stage == "predict"):
+ elif stage == "predict":
dataset = self.predict_dataset
else:
raise ValueError("Invalid stage")
@@ -718,15 +1048,121 @@ def _gen_dataloader(self, stage):
return dl
def train_dataloader(self):
- return self._gen_dataloader("train")
+ return self._gen_dataloader("train")
def val_dataloader(self):
- if(self.eval_dataset is not None):
+ if self.eval_dataset is not None:
return self._gen_dataloader("eval")
return None
def predict_dataloader(self):
- return self._gen_dataloader("predict")
+ return self._gen_dataloader("predict")
+
+
+class OpenFoldMultimerDataModule(OpenFoldDataModule):
+ """
+ Create a datamodule specifically for multimer training
+
+ Compared to OpenFoldDataModule, OpenFoldMultimerDataModule
+ requires mmcif_data_cache_path which is the product of
+ scripts/generate_mmcif_cache.py mmcif_data_cache_path should be
+ a file that record what chain(s) each mmcif file has
+ """
+
+ def __init__(self, config: mlc.ConfigDict,
+ template_mmcif_dir: str, max_template_date: str,
+ train_data_dir: Optional[str] = None,
+ train_mmcif_data_cache_path: Optional[str] = None,
+ val_mmcif_data_cache_path: Optional[str] = None,
+ **kwargs):
+ super(OpenFoldMultimerDataModule, self).__init__(config,
+ template_mmcif_dir,
+ max_template_date,
+ train_data_dir,
+ **kwargs)
+
+ self.train_mmcif_data_cache_path = train_mmcif_data_cache_path
+ self.training_mode = self.train_data_dir is not None
+ self.val_mmcif_data_cache_path = val_mmcif_data_cache_path
+
+ def setup(self):
+ # Most of the arguments are the same for the three datasets
+ dataset_gen = partial(OpenFoldSingleMultimerDataset,
+ template_mmcif_dir=self.template_mmcif_dir,
+ max_template_date=self.max_template_date,
+ config=self.config,
+ kalign_binary_path=self.kalign_binary_path,
+ template_release_dates_cache_path=self.template_release_dates_cache_path,
+ obsolete_pdbs_file_path=self.obsolete_pdbs_file_path)
+
+ if self.training_mode:
+ train_dataset = dataset_gen(
+ data_dir=self.train_data_dir,
+ mmcif_data_cache_path=self.train_mmcif_data_cache_path,
+ alignment_dir=self.train_alignment_dir,
+ filter_path=self.train_filter_path,
+ max_template_hits=self.config.train.max_template_hits,
+ shuffle_top_k_prefiltered=self.config.train.shuffle_top_k_prefiltered,
+ treat_pdb_as_distillation=False,
+ mode="train",
+ alignment_index=self.alignment_index,
+ )
+
+ distillation_dataset = None
+ if self.distillation_data_dir is not None:
+ distillation_dataset = dataset_gen(
+ data_dir=self.distillation_data_dir,
+ alignment_dir=self.distillation_alignment_dir,
+ filter_path=self.distillation_filter_path,
+ max_template_hits=self.config.train.max_template_hits,
+ treat_pdb_as_distillation=True,
+ mode="train",
+ alignment_index=self.distillation_alignment_index,
+ _structure_index=self._distillation_structure_index,
+ )
+
+ d_prob = self.config.train.distillation_prob
+
+ if distillation_dataset is not None:
+ datasets = [train_dataset, distillation_dataset]
+ d_prob = self.config.train.distillation_prob
+ probabilities = [1. - d_prob, d_prob]
+ else:
+ datasets = [train_dataset]
+ probabilities = [1.]
+
+ generator = None
+ if self.batch_seed is not None:
+ generator = torch.Generator()
+ generator = generator.manual_seed(self.batch_seed + 1)
+
+ self.train_dataset = OpenFoldMultimerDataset(
+ datasets=datasets,
+ probabilities=probabilities,
+ epoch_len=self.train_epoch_len,
+ generator=generator,
+ _roll_at_init=True,
+ )
+
+ if self.val_data_dir is not None:
+ self.eval_dataset = dataset_gen(
+ data_dir=self.val_data_dir,
+ alignment_dir=self.val_alignment_dir,
+ mmcif_data_cache_path=self.val_mmcif_data_cache_path,
+ filter_path=None,
+ max_template_hits=self.config.eval.max_template_hits,
+ mode="eval",
+ )
+ else:
+ self.eval_dataset = None
+ else:
+ self.predict_dataset = dataset_gen(
+ data_dir=self.predict_data_dir,
+ alignment_dir=self.predict_alignment_dir,
+ filter_path=None,
+ max_template_hits=self.config.predict.max_template_hits,
+ mode="predict",
+ )
class DummyDataset(torch.utils.data.Dataset):
diff --git a/openfold/data/data_pipeline.py b/openfold/data/data_pipeline.py
index 2ce9dfed..ce8494d0 100644
--- a/openfold/data/data_pipeline.py
+++ b/openfold/data/data_pipeline.py
@@ -14,38 +14,28 @@
# limitations under the License.
import os
-import datetime
+import copy
+import collections
+import contextlib
+import dataclasses
from multiprocessing import cpu_count
-from typing import Mapping, Optional, Sequence, Any
-
+import tempfile
+from typing import Mapping, Optional, Sequence, Any, MutableMapping, Union
import numpy as np
import torch
-
-from openfold.data import templates, parsers, mmcif_parsing
-from openfold.data.templates import get_custom_template_features
-from openfold.data.tools import jackhmmer, hhblits, hhsearch
-from openfold.data.tools.utils import to_date
+from openfold.data import templates, parsers, mmcif_parsing, msa_identifiers, msa_pairing, feature_processing_multimer
+from openfold.data.templates import get_custom_template_features, empty_template_feats
+from openfold.data.tools import jackhmmer, hhblits, hhsearch, hmmsearch
from openfold.np import residue_constants, protein
-
-FeatureDict = Mapping[str, np.ndarray]
-
-def empty_template_feats(n_res) -> FeatureDict:
- return {
- "template_aatype": np.zeros((0, n_res)).astype(np.int64),
- "template_all_atom_positions":
- np.zeros((0, n_res, 37, 3)).astype(np.float32),
- "template_sum_probs": np.zeros((0, 1)).astype(np.float32),
- "template_all_atom_mask": np.zeros((0, n_res, 37)).astype(np.float32),
- }
+FeatureDict = MutableMapping[str, np.ndarray]
+TemplateSearcher = Union[hhsearch.HHSearch, hmmsearch.Hmmsearch]
def make_template_features(
input_sequence: str,
hits: Sequence[Any],
template_featurizer: Any,
- query_pdb_code: Optional[str] = None,
- query_release_date: Optional[str] = None,
) -> FeatureDict:
hits_cat = sum(hits.values(), [])
if(len(hits_cat) == 0 or template_featurizer is None):
@@ -53,17 +43,10 @@ def make_template_features(
else:
templates_result = template_featurizer.get_templates(
query_sequence=input_sequence,
- query_pdb_code=query_pdb_code,
- query_release_date=query_release_date,
hits=hits_cat,
)
template_features = templates_result.features
- # The template featurizer doesn't format empty template features
- # properly. This is a quick fix.
- if(template_features["template_aatype"].shape[0] == 0):
- template_features = empty_template_feats(len(input_sequence))
-
return template_features
@@ -86,7 +69,7 @@ def unify_template_features(
assert(new_shape[1] == n_res)
new_shape[1] = sum(seq_lens)
new_array = np.zeros(new_shape, dtype=v.dtype)
-
+
if(k == "template_aatype"):
new_array[..., residue_constants.HHBLITS_AA_TO_ID['-']] = 1
@@ -172,13 +155,13 @@ def make_mmcif_features(
def _aatype_to_str_sequence(aatype):
return ''.join([
- residue_constants.restypes_with_x[aatype[i]]
+ residue_constants.restypes_with_x[aatype[i]]
for i in range(len(aatype))
])
def make_protein_features(
- protein_object: protein.Protein,
+ protein_object: protein.Protein,
description: str,
_is_distillation: bool = False,
) -> FeatureDict:
@@ -225,32 +208,35 @@ def make_pdb_features(
return pdb_feats
-def make_msa_features(
- msas: Sequence[Sequence[str]],
- deletion_matrices: Sequence[parsers.DeletionMatrix],
-) -> FeatureDict:
+def make_msa_features(msas: Sequence[parsers.Msa]) -> FeatureDict:
"""Constructs a feature dict of MSA features."""
if not msas:
raise ValueError("At least one MSA must be provided.")
int_msa = []
deletion_matrix = []
+ species_ids = []
seen_sequences = set()
for msa_index, msa in enumerate(msas):
if not msa:
raise ValueError(
f"MSA {msa_index} must contain at least one sequence."
)
- for sequence_index, sequence in enumerate(msa):
+ for sequence_index, sequence in enumerate(msa.sequences):
if sequence in seen_sequences:
continue
seen_sequences.add(sequence)
int_msa.append(
[residue_constants.HHBLITS_AA_TO_ID[res] for res in sequence]
)
- deletion_matrix.append(deletion_matrices[msa_index][sequence_index])
- num_res = len(msas[0][0])
+ deletion_matrix.append(msa.deletion_matrix[sequence_index])
+ identifiers = msa_identifiers.get_identifiers(
+ msa.descriptions[sequence_index]
+ )
+ species_ids.append(identifiers.species_id.encode('utf-8'))
+
+ num_res = len(msas[0].sequences[0])
num_alignments = len(int_msa)
features = {}
features["deletion_matrix_int"] = np.array(deletion_matrix, dtype=np.int32)
@@ -258,19 +244,41 @@ def make_msa_features(
features["num_alignments"] = np.array(
[num_alignments] * num_res, dtype=np.int32
)
+ features["msa_species_identifiers"] = np.array(species_ids, dtype=np.object_)
return features
-# Generate 1-sequence MSA features having only the input sequence
-def make_dummy_msa_feats(input_sequence):
- msas = [[input_sequence]]
- deletion_matrices = [[[0 for _ in input_sequence]]]
- msa_features = make_msa_features(
- msas=msas,
- deletion_matrices=deletion_matrices,
- )
+def run_msa_tool(
+ msa_runner,
+ fasta_path: str,
+ msa_out_path: str,
+ msa_format: str,
+ max_sto_sequences: Optional[int] = None,
+) -> Mapping[str, Any]:
+ """Runs an MSA tool, checking if output already exists first."""
+ if(msa_format == "sto" and max_sto_sequences is not None):
+ result = msa_runner.query(fasta_path, max_sto_sequences)[0]
+ else:
+ result = msa_runner.query(fasta_path)[0]
+
+ assert msa_out_path.split('.')[-1] == msa_format
+ with open(msa_out_path, "w") as f:
+ f.write(result[msa_format])
- return msa_features
+ return result
+
+
+def make_dummy_msa_obj(input_sequence) -> parsers.Msa:
+ deletion_matrix = [[0 for _ in input_sequence]]
+ return parsers.Msa(sequences=[input_sequence],
+ deletion_matrix=deletion_matrix,
+ descriptions=['dummy'])
+
+
+# Generate 1-sequence MSA features having only the input sequence
+def make_dummy_msa_feats(input_sequence) -> FeatureDict:
+ msa_data_obj = make_dummy_msa_obj(input_sequence)
+ return make_msa_features([msa_data_obj])
def make_sequence_features_with_custom_template(
@@ -290,10 +298,11 @@ def make_sequence_features_with_custom_template(
num_res=num_res,
)
- msa_data = [[sequence]]
- deletion_matrix = [[[0 for _ in sequence]]]
+ msa_data = [sequence]
+ deletion_matrix = [[0 for _ in sequence]]
+ msa_data_obj = parsers.Msa(sequences=msa_data, deletion_matrix=deletion_matrix, descriptions=None)
- msa_features = make_msa_features(msa_data, deletion_matrix)
+ msa_features = make_msa_features([msa_data_obj])
template_features = get_custom_template_features(
mmcif_path=mmcif_path,
query_sequence=sequence,
@@ -308,22 +317,25 @@ def make_sequence_features_with_custom_template(
**template_features.features
}
+
class AlignmentRunner:
"""Runs alignment tools and saves the results"""
def __init__(
self,
jackhmmer_binary_path: Optional[str] = None,
hhblits_binary_path: Optional[str] = None,
- hhsearch_binary_path: Optional[str] = None,
uniref90_database_path: Optional[str] = None,
mgnify_database_path: Optional[str] = None,
bfd_database_path: Optional[str] = None,
+ uniref30_database_path: Optional[str] = None,
uniclust30_database_path: Optional[str] = None,
- pdb70_database_path: Optional[str] = None,
+ uniprot_database_path: Optional[str] = None,
+ template_searcher: Optional[TemplateSearcher] = None,
use_small_bfd: Optional[bool] = None,
no_cpus: Optional[int] = None,
uniref_max_hits: int = 10000,
mgnify_max_hits: int = 5000,
+ uniprot_max_hits: int = 50000,
):
"""
Args:
@@ -331,8 +343,6 @@ def __init__(
Path to jackhmmer binary
hhblits_binary_path:
Path to hhblits binary
- hhsearch_binary_path:
- Path to hhsearch binary
uniref90_database_path:
Path to uniref90 database. If provided, jackhmmer_binary_path
must also be provided
@@ -341,16 +351,17 @@ def __init__(
must also be provided
bfd_database_path:
Path to BFD database. Depending on the value of use_small_bfd,
- one of hhblits_binary_path or jackhmmer_binary_path must be
+ one of hhblits_binary_path or jackhmmer_binary_path must be
provided.
+ uniref30_database_path:
+ Path to uniref30. Searched alongside BFD if use_small_bfd is
+ false.
uniclust30_database_path:
- Path to uniclust30. Searched alongside BFD if use_small_bfd is
+ Path to uniclust30. Searched alongside BFD if use_small_bfd is
false.
- pdb70_database_path:
- Path to pdb70 database.
use_small_bfd:
- Whether to search the BFD database alone with jackhmmer or
- in conjunction with uniclust30 with hhblits.
+ Whether to search the BFD database alone with jackhmmer or
+ in conjunction with uniref30/uniclust30 with hhblits.
no_cpus:
The number of CPUs available for alignment. By default, all
CPUs are used.
@@ -358,6 +369,8 @@ def __init__(
Max number of uniref hits
mgnify_max_hits:
Max number of mgnify hits
+ uniprot_max_hits:
+ Max number of uniprot hits
"""
db_map = {
"jackhmmer": {
@@ -366,6 +379,7 @@ def __init__(
uniref90_database_path,
mgnify_database_path,
bfd_database_path if use_small_bfd else None,
+ uniprot_database_path,
],
},
"hhblits": {
@@ -374,12 +388,6 @@ def __init__(
bfd_database_path if not use_small_bfd else None,
],
},
- "hhsearch": {
- "binary": hhsearch_binary_path,
- "dbs": [
- pdb70_database_path,
- ],
- },
}
for name, dic in db_map.items():
@@ -389,22 +397,16 @@ def __init__(
f"{name} DBs provided but {name} binary is None"
)
- if(not all([x is None for x in db_map["hhsearch"]["dbs"]])
- and uniref90_database_path is None):
- raise ValueError(
- """uniref90_database_path must be specified in order to perform
- template search"""
- )
-
self.uniref_max_hits = uniref_max_hits
self.mgnify_max_hits = mgnify_max_hits
+ self.uniprot_max_hits = uniprot_max_hits
self.use_small_bfd = use_small_bfd
if(no_cpus is None):
no_cpus = cpu_count()
self.jackhmmer_uniref90_runner = None
- if(jackhmmer_binary_path is not None and
+ if(jackhmmer_binary_path is not None and
uniref90_database_path is not None
):
self.jackhmmer_uniref90_runner = jackhmmer.Jackhmmer(
@@ -412,9 +414,9 @@ def __init__(
database_path=uniref90_database_path,
n_cpu=no_cpus,
)
-
+
self.jackhmmer_small_bfd_runner = None
- self.hhblits_bfd_uniclust_runner = None
+ self.hhblits_bfd_unirefclust_runner = None
if(bfd_database_path is not None):
if use_small_bfd:
self.jackhmmer_small_bfd_runner = jackhmmer.Jackhmmer(
@@ -424,9 +426,11 @@ def __init__(
)
else:
dbs = [bfd_database_path]
- if(uniclust30_database_path is not None):
+ if(uniref30_database_path is not None):
+ dbs.append(uniref30_database_path)
+ if (uniclust30_database_path is not None):
dbs.append(uniclust30_database_path)
- self.hhblits_bfd_uniclust_runner = hhblits.HHBlits(
+ self.hhblits_bfd_unirefclust_runner = hhblits.HHBlits(
binary_path=hhblits_binary_path,
databases=dbs,
n_cpu=no_cpus,
@@ -440,14 +444,23 @@ def __init__(
n_cpu=no_cpus,
)
- self.hhsearch_pdb70_runner = None
- if(pdb70_database_path is not None):
- self.hhsearch_pdb70_runner = hhsearch.HHSearch(
- binary_path=hhsearch_binary_path,
- databases=[pdb70_database_path],
- n_cpu=no_cpus,
+ self.jackhmmer_uniprot_runner = None
+ if(uniprot_database_path is not None):
+ self.jackhmmer_uniprot_runner = jackhmmer.Jackhmmer(
+ binary_path=jackhmmer_binary_path,
+ database_path=uniprot_database_path,
+ n_cpu=no_cpus
)
+ if(template_searcher is not None and
+ self.jackhmmer_uniref90_runner is None
+ ):
+ raise ValueError(
+ "Uniref90 runner must be specified to run template search"
+ )
+
+ self.template_searcher = template_searcher
+
def run(
self,
fasta_path: str,
@@ -455,52 +468,226 @@ def run(
):
"""Runs alignment tools on a sequence"""
if(self.jackhmmer_uniref90_runner is not None):
- jackhmmer_uniref90_result = self.jackhmmer_uniref90_runner.query(
- fasta_path
- )[0]
- uniref90_msa_as_a3m = parsers.convert_stockholm_to_a3m(
- jackhmmer_uniref90_result["sto"],
- max_sequences=self.uniref_max_hits
+ uniref90_out_path = os.path.join(output_dir, "uniref90_hits.sto")
+
+ jackhmmer_uniref90_result = run_msa_tool(
+ msa_runner=self.jackhmmer_uniref90_runner,
+ fasta_path=fasta_path,
+ msa_out_path=uniref90_out_path,
+ msa_format='sto',
+ max_sto_sequences=self.uniref_max_hits,
)
- uniref90_out_path = os.path.join(output_dir, "uniref90_hits.a3m")
- with open(uniref90_out_path, "w") as f:
- f.write(uniref90_msa_as_a3m)
- if(self.hhsearch_pdb70_runner is not None):
- hhsearch_result = self.hhsearch_pdb70_runner.query(
- uniref90_msa_as_a3m
- )
- pdb70_out_path = os.path.join(output_dir, "pdb70_hits.hhr")
- with open(pdb70_out_path, "w") as f:
- f.write(hhsearch_result)
+ template_msa = jackhmmer_uniref90_result["sto"]
+ template_msa = parsers.deduplicate_stockholm_msa(template_msa)
+ template_msa = parsers.remove_empty_columns_from_stockholm_msa(
+ template_msa
+ )
+
+ if(self.template_searcher is not None):
+ if(self.template_searcher.input_format == "sto"):
+ pdb_templates_result = self.template_searcher.query(
+ template_msa,
+ output_dir=output_dir
+ )
+ elif(self.template_searcher.input_format == "a3m"):
+ uniref90_msa_as_a3m = parsers.convert_stockholm_to_a3m(
+ template_msa
+ )
+ pdb_templates_result = self.template_searcher.query(
+ uniref90_msa_as_a3m,
+ output_dir=output_dir
+ )
+ else:
+ fmt = self.template_searcher.input_format
+ raise ValueError(
+ f"Unrecognized template input format: {fmt}"
+ )
if(self.jackhmmer_mgnify_runner is not None):
- jackhmmer_mgnify_result = self.jackhmmer_mgnify_runner.query(
- fasta_path
- )[0]
- mgnify_msa_as_a3m = parsers.convert_stockholm_to_a3m(
- jackhmmer_mgnify_result["sto"],
- max_sequences=self.mgnify_max_hits
+ mgnify_out_path = os.path.join(output_dir, "mgnify_hits.sto")
+ jackhmmer_mgnify_result = run_msa_tool(
+ msa_runner=self.jackhmmer_mgnify_runner,
+ fasta_path=fasta_path,
+ msa_out_path=mgnify_out_path,
+ msa_format='sto',
+ max_sto_sequences=self.mgnify_max_hits
)
- mgnify_out_path = os.path.join(output_dir, "mgnify_hits.a3m")
- with open(mgnify_out_path, "w") as f:
- f.write(mgnify_msa_as_a3m)
if(self.use_small_bfd and self.jackhmmer_small_bfd_runner is not None):
- jackhmmer_small_bfd_result = self.jackhmmer_small_bfd_runner.query(
- fasta_path
- )[0]
bfd_out_path = os.path.join(output_dir, "small_bfd_hits.sto")
- with open(bfd_out_path, "w") as f:
- f.write(jackhmmer_small_bfd_result["sto"])
- elif(self.hhblits_bfd_uniclust_runner is not None):
- hhblits_bfd_uniclust_result = (
- self.hhblits_bfd_uniclust_runner.query(fasta_path)
+ jackhmmer_small_bfd_result = run_msa_tool(
+ msa_runner=self.jackhmmer_small_bfd_runner,
+ fasta_path=fasta_path,
+ msa_out_path=bfd_out_path,
+ msa_format="sto",
+ )
+ elif(self.hhblits_bfd_unirefclust_runner is not None):
+ uni_name = "uni"
+ for db_name in self.hhblits_bfd_unirefclust_runner.databases:
+ if "uniref" in db_name.lower():
+ uni_name = f"{uni_name}ref"
+ elif "uniclust" in db_name.lower():
+ uni_name = f"{uni_name}clust"
+
+ bfd_out_path = os.path.join(output_dir, f"bfd_{uni_name}_hits.a3m")
+ hhblits_bfd_unirefclust_result = run_msa_tool(
+ msa_runner=self.hhblits_bfd_unirefclust_runner,
+ fasta_path=fasta_path,
+ msa_out_path=bfd_out_path,
+ msa_format="a3m",
+ )
+
+ if(self.jackhmmer_uniprot_runner is not None):
+ uniprot_out_path = os.path.join(output_dir, 'uniprot_hits.sto')
+ result = run_msa_tool(
+ self.jackhmmer_uniprot_runner,
+ fasta_path=fasta_path,
+ msa_out_path=uniprot_out_path,
+ msa_format='sto',
+ max_sto_sequences=self.uniprot_max_hits,
)
- if output_dir is not None:
- bfd_out_path = os.path.join(output_dir, "bfd_uniclust_hits.a3m")
- with open(bfd_out_path, "w") as f:
- f.write(hhblits_bfd_uniclust_result["a3m"])
+
+
+@dataclasses.dataclass(frozen=True)
+class _FastaChain:
+ sequence: str
+ description: str
+
+
+def _make_chain_id_map(
+ sequences: Sequence[str],
+ descriptions: Sequence[str],
+) -> Mapping[str, _FastaChain]:
+ """Makes a mapping from PDB-format chain ID to sequence and description."""
+ if len(sequences) != len(descriptions):
+ raise ValueError('sequences and descriptions must have equal length. '
+ f'Got {len(sequences)} != {len(descriptions)}.')
+ if len(sequences) > protein.PDB_MAX_CHAINS:
+ raise ValueError('Cannot process more chains than the PDB format supports. '
+ f'Got {len(sequences)} chains.')
+ chain_id_map = {}
+ for chain_id, sequence, description in zip(
+ protein.PDB_CHAIN_IDS, sequences, descriptions
+ ):
+ chain_id_map[chain_id] = _FastaChain(
+ sequence=sequence, description=description
+ )
+ return chain_id_map
+
+
+@contextlib.contextmanager
+def temp_fasta_file(fasta_str: str):
+ with tempfile.NamedTemporaryFile('w', suffix='.fasta') as fasta_file:
+ fasta_file.write(fasta_str)
+ fasta_file.seek(0)
+ yield fasta_file.name
+
+
+def convert_monomer_features(
+ monomer_features: FeatureDict,
+ chain_id: str
+) -> FeatureDict:
+ """Reshapes and modifies monomer features for multimer models."""
+ converted = {}
+ converted['auth_chain_id'] = np.asarray(chain_id, dtype=np.object_)
+ unnecessary_leading_dim_feats = {
+ 'sequence', 'domain_name', 'num_alignments', 'seq_length'
+ }
+ for feature_name, feature in monomer_features.items():
+ if feature_name in unnecessary_leading_dim_feats:
+ # asarray ensures it's a np.ndarray.
+ feature = np.asarray(feature[0], dtype=feature.dtype)
+ elif feature_name == 'aatype':
+ # The multimer model performs the one-hot operation itself.
+ feature = np.argmax(feature, axis=-1).astype(np.int32)
+ elif feature_name == 'template_aatype':
+ feature = np.argmax(feature, axis=-1).astype(np.int32)
+ new_order_list = residue_constants.MAP_HHBLITS_AATYPE_TO_OUR_AATYPE
+ feature = np.take(new_order_list, feature.astype(np.int32), axis=0)
+ elif feature_name == 'template_all_atom_masks':
+ feature_name = 'template_all_atom_mask'
+ converted[feature_name] = feature
+ return converted
+
+
+def int_id_to_str_id(num: int) -> str:
+ """Encodes a number as a string, using reverse spreadsheet style naming.
+
+ Args:
+ num: A positive integer.
+
+ Returns:
+ A string that encodes the positive integer using reverse spreadsheet style,
+ naming e.g. 1 = A, 2 = B, ..., 27 = AA, 28 = BA, 29 = CA, ... This is the
+ usual way to encode chain IDs in mmCIF files.
+ """
+ if num <= 0:
+ raise ValueError(f'Only positive integers allowed, got {num}.')
+
+ num = num - 1 # 1-based indexing.
+ output = []
+ while num >= 0:
+ output.append(chr(num % 26 + ord('A')))
+ num = num // 26 - 1
+ return ''.join(output)
+
+
+def add_assembly_features(
+ all_chain_features: MutableMapping[str, FeatureDict],
+) -> MutableMapping[str, FeatureDict]:
+ """Add features to distinguish between chains.
+
+ Args:
+ all_chain_features: A dictionary which maps chain_id to a dictionary of
+ features for each chain.
+
+ Returns:
+ all_chain_features: A dictionary which maps strings of the form
+ `_` to the corresponding chain features. E.g. two
+ chains from a homodimer would have keys A_1 and A_2. Two chains from a
+ heterodimer would have keys A_1 and B_1.
+ """
+ # Group the chains by sequence
+ seq_to_entity_id = {}
+ grouped_chains = collections.defaultdict(list)
+ for chain_id, chain_features in all_chain_features.items():
+ seq = str(chain_features['sequence'])
+ if seq not in seq_to_entity_id:
+ seq_to_entity_id[seq] = len(seq_to_entity_id) + 1
+ grouped_chains[seq_to_entity_id[seq]].append(chain_features)
+
+ new_all_chain_features = {}
+ chain_id = 1
+ for entity_id, group_chain_features in grouped_chains.items():
+ for sym_id, chain_features in enumerate(group_chain_features, start=1):
+ new_all_chain_features[
+ f'{int_id_to_str_id(entity_id)}_{sym_id}'] = chain_features
+ seq_length = chain_features['seq_length']
+ chain_features['asym_id'] = (
+ chain_id * np.ones(seq_length)
+ ).astype(np.int64)
+ chain_features['sym_id'] = (
+ sym_id * np.ones(seq_length)
+ ).astype(np.int64)
+ chain_features['entity_id'] = (
+ entity_id * np.ones(seq_length)
+ ).astype(np.int64)
+ chain_id += 1
+
+ return new_all_chain_features
+
+
+def pad_msa(np_example, min_num_seq):
+ np_example = dict(np_example)
+ num_seq = np_example['msa'].shape[0]
+ if num_seq < min_num_seq:
+ for feat in ('msa', 'deletion_matrix', 'bert_mask', 'msa_mask'):
+ np_example[feat] = np.pad(
+ np_example[feat], ((0, min_num_seq - num_seq), (0, 0)))
+ np_example['cluster_bias_mask'] = np.pad(
+ np_example['cluster_bias_mask'], ((0, min_num_seq - num_seq),))
+ return np_example
class DataPipeline:
@@ -514,10 +701,10 @@ def __init__(
def _parse_msa_data(
self,
alignment_dir: str,
- alignment_index: Optional[Any] = None,
+ alignment_index: Optional[Any] = None
) -> Mapping[str, Any]:
- msa_data = {}
- if(alignment_index is not None):
+ msa_data = {}
+ if alignment_index is not None:
fp = open(os.path.join(alignment_dir, alignment_index["db"]), "rb")
def read_msa(start, size):
@@ -526,49 +713,47 @@ def read_msa(start, size):
return msa
for (name, start, size) in alignment_index["files"]:
- ext = os.path.splitext(name)[-1]
+ filename, ext = os.path.splitext(name)
- if(ext == ".a3m"):
- msa, deletion_matrix = parsers.parse_a3m(
- read_msa(start, size)
- )
- data = {"msa": msa, "deletion_matrix": deletion_matrix}
- elif(ext == ".sto"):
- msa, deletion_matrix, _ = parsers.parse_stockholm(
+ if ext == ".a3m":
+ msa = parsers.parse_a3m(
read_msa(start, size)
)
- data = {"msa": msa, "deletion_matrix": deletion_matrix}
+ # The "hmm_output" exception is a crude way to exclude
+ # multimer template hits.
+ # Multimer "uniprot_hits" processed separately.
+ elif ext == ".sto" and filename not in ["uniprot_hits", "hmm_output"]:
+ msa = parsers.parse_stockholm(read_msa(start, size))
else:
continue
-
- msa_data[name] = data
-
+
+ msa_data[name] = msa
+
fp.close()
- else:
+ else:
for f in os.listdir(alignment_dir):
path = os.path.join(alignment_dir, f)
- ext = os.path.splitext(f)[-1]
+ filename, ext = os.path.splitext(f)
- if(ext == ".a3m"):
+ if ext == ".a3m":
with open(path, "r") as fp:
- msa, deletion_matrix = parsers.parse_a3m(fp.read())
- data = {"msa": msa, "deletion_matrix": deletion_matrix}
- elif(ext == ".sto"):
+ msa = parsers.parse_a3m(fp.read())
+ elif ext == ".sto" and filename not in ["uniprot_hits", "hmm_output"]:
with open(path, "r") as fp:
- msa, deletion_matrix, _ = parsers.parse_stockholm(
+ msa = parsers.parse_stockholm(
fp.read()
)
- data = {"msa": msa, "deletion_matrix": deletion_matrix}
else:
continue
-
- msa_data[f] = data
+
+ msa_data[f] = msa
return msa_data
- def _parse_template_hits(
+ def _parse_template_hit_files(
self,
alignment_dir: str,
+ input_sequence: str,
alignment_index: Optional[Any] = None
) -> Mapping[str, Any]:
all_hits = {}
@@ -585,6 +770,12 @@ def read_template(start, size):
if(ext == ".hhr"):
hits = parsers.parse_hhr(read_template(start, size))
all_hits[name] = hits
+ elif(name == "hmmsearch_output.sto"):
+ hits = parsers.parse_hmmsearch_sto(
+ read_template(start, size),
+ input_sequence,
+ )
+ all_hits[name] = hits
fp.close()
else:
@@ -596,13 +787,20 @@ def read_template(start, size):
with open(path, "r") as fp:
hits = parsers.parse_hhr(fp.read())
all_hits[f] = hits
+ elif(f == "hmm_output.sto"):
+ with open(path, "r") as fp:
+ hits = parsers.parse_hmmsearch_sto(
+ fp.read(),
+ input_sequence,
+ )
+ all_hits[f] = hits
return all_hits
def _get_msas(self,
alignment_dir: str,
input_sequence: Optional[str] = None,
- alignment_index: Optional[str] = None,
+ alignment_index: Optional[Any] = None,
):
msa_data = self._parse_msa_data(alignment_dir, alignment_index)
if(len(msa_data) == 0):
@@ -613,29 +811,23 @@ def _get_msas(self,
must be provided.
"""
)
- msa_data["dummy"] = {
- "msa": [input_sequence],
- "deletion_matrix": [[0 for _ in input_sequence]],
- }
- msas, deletion_matrices = zip(*[
- (v["msa"], v["deletion_matrix"]) for v in msa_data.values()
- ])
+ msa_data["dummy"] = make_dummy_msa_obj(input_sequence)
- return msas, deletion_matrices
+ return list(msa_data.values())
def _process_msa_feats(
self,
alignment_dir: str,
input_sequence: Optional[str] = None,
- alignment_index: Optional[str] = None
+ alignment_index: Optional[Any] = None
) -> Mapping[str, Any]:
- msas, deletion_matrices = self._get_msas(
+
+ msas = self._get_msas(
alignment_dir, input_sequence, alignment_index
)
msa_features = make_msa_features(
- msas=msas,
- deletion_matrices=deletion_matrices,
+ msas=msas
)
return msa_features
@@ -660,10 +852,10 @@ def process_fasta(
self,
fasta_path: str,
alignment_dir: str,
- alignment_index: Optional[str] = None,
+ alignment_index: Optional[Any] = None,
seqemb_mode: bool = False,
) -> FeatureDict:
- """Assembles features for a single sequence in a FASTA file"""
+ """Assembles features for a single sequence in a FASTA file"""
with open(fasta_path) as f:
fasta_str = f.read()
input_seqs, input_descs = parsers.parse_fasta(fasta_str)
@@ -675,7 +867,12 @@ def process_fasta(
input_description = input_descs[0]
num_res = len(input_sequence)
- hits = self._parse_template_hits(alignment_dir, alignment_index)
+ hits = self._parse_template_hit_files(
+ alignment_dir=alignment_dir,
+ input_sequence=input_sequence,
+ alignment_index=alignment_index,
+ )
+
template_features = make_template_features(
input_sequence,
hits,
@@ -700,7 +897,7 @@ def process_fasta(
**sequence_features,
**msa_features,
**template_features,
- **sequence_embedding_features,
+ **sequence_embedding_features
}
def process_mmcif(
@@ -708,7 +905,7 @@ def process_mmcif(
mmcif: mmcif_parsing.MmcifObject, # parsing is expensive, so no path
alignment_dir: str,
chain_id: Optional[str] = None,
- alignment_index: Optional[str] = None,
+ alignment_index: Optional[Any] = None,
seqemb_mode: bool = False,
) -> FeatureDict:
"""
@@ -727,12 +924,15 @@ def process_mmcif(
mmcif_feats = make_mmcif_features(mmcif, chain_id)
input_sequence = mmcif.chain_to_seqres[chain_id]
- hits = self._parse_template_hits(alignment_dir, alignment_index)
+ hits = self._parse_template_hit_files(
+ alignment_dir=alignment_dir,
+ input_sequence=input_sequence,
+ alignment_index=alignment_index)
+
template_features = make_template_features(
input_sequence,
hits,
- self.template_featurizer,
- query_release_date=to_date(mmcif.header["release_date"])
+ self.template_featurizer
)
sequence_embedding_features = {}
@@ -752,7 +952,7 @@ def process_pdb(
is_distillation: bool = True,
chain_id: Optional[str] = None,
_structure_index: Optional[str] = None,
- alignment_index: Optional[str] = None,
+ alignment_index: Optional[Any] = None,
seqemb_mode: bool = False,
) -> FeatureDict:
"""
@@ -772,15 +972,20 @@ def process_pdb(
pdb_str = f.read()
protein_object = protein.from_pdb_string(pdb_str, chain_id)
- input_sequence = _aatype_to_str_sequence(protein_object.aatype)
+ input_sequence = _aatype_to_str_sequence(protein_object.aatype)
description = os.path.splitext(os.path.basename(pdb_path))[0].upper()
pdb_feats = make_pdb_features(
- protein_object,
- description,
+ protein_object,
+ description,
is_distillation=is_distillation
)
- hits = self._parse_template_hits(alignment_dir, alignment_index)
+ hits = self._parse_template_hit_files(
+ alignment_dir=alignment_dir,
+ input_sequence=input_sequence,
+ alignment_index=alignment_index,
+ )
+
template_features = make_template_features(
input_sequence,
hits,
@@ -801,7 +1006,7 @@ def process_core(
self,
core_path: str,
alignment_dir: str,
- alignment_index: Optional[str] = None,
+ alignment_index: Optional[Any] = None,
seqemb_mode: bool = False,
) -> FeatureDict:
"""
@@ -811,11 +1016,16 @@ def process_core(
core_str = f.read()
protein_object = protein.from_proteinnet_string(core_str)
- input_sequence = _aatype_to_str_sequence(protein_object.aatype)
+ input_sequence = _aatype_to_str_sequence(protein_object.aatype)
description = os.path.splitext(os.path.basename(core_path))[0].upper()
core_feats = make_protein_features(protein_object, description)
-
- hits = self._parse_template_hits(alignment_dir, alignment_index)
+
+ hits = self._parse_template_hit_files(
+ alignment_dir=alignment_dir,
+ input_sequence=input_sequence,
+ alignment_index=alignment_index,
+ )
+
template_features = make_template_features(
input_sequence,
hits,
@@ -833,10 +1043,10 @@ def process_core(
return {**core_feats, **template_features, **msa_features, **sequence_embedding_features}
def process_multiseq_fasta(self,
- fasta_path: str,
- super_alignment_dir: str,
- ri_gap: int = 200,
- ) -> FeatureDict:
+ fasta_path: str,
+ super_alignment_dir: str,
+ ri_gap: int = 200,
+ ) -> FeatureDict:
"""
Assembles features for a multi-sequence FASTA. Uses Minkyung Baek's
hack from Twitter (a.k.a. AlphaFold-Gap).
@@ -845,7 +1055,7 @@ def process_multiseq_fasta(self,
fasta_str = f.read()
input_seqs, input_descs = parsers.parse_fasta(fasta_str)
-
+
# No whitespace allowed
input_descs = [i.split()[0] for i in input_descs]
@@ -872,14 +1082,15 @@ def process_multiseq_fasta(self,
alignment_dir = os.path.join(
super_alignment_dir, desc
)
- msas, deletion_mats = self._get_msas(
+ msas = self._get_msas(
alignment_dir, seq, None
)
- msa_list.append(msas)
- deletion_mat_list.append(deletion_mats)
+ msa_list.append([m.sequences for m in msas])
+ deletion_mat_list.append([m.deletion_matrix for m in msas])
final_msa = []
final_deletion_mat = []
+ final_msa_obj = []
msa_it = enumerate(zip(msa_list, deletion_mat_list))
for i, (msas, deletion_mats) in msa_it:
prec, post = sum(seq_lens[:i]), sum(seq_lens[i + 1:])
@@ -887,18 +1098,19 @@ def process_multiseq_fasta(self,
[prec * '-' + seq + post * '-' for seq in msa] for msa in msas
]
deletion_mats = [
- [prec * [0] + dml + post * [0] for dml in deletion_mat]
+ [prec * [0] + dml + post * [0] for dml in deletion_mat]
for deletion_mat in deletion_mats
]
- assert(len(msas[0][-1]) == len(input_sequence))
+ assert (len(msas[0][-1]) == len(input_sequence))
final_msa.extend(msas)
final_deletion_mat.extend(deletion_mats)
+ final_msa_obj.extend([parsers.Msa(sequences=msas[k], deletion_matrix=deletion_mats[k], descriptions=None)
+ for k in range(len(msas))])
msa_features = make_msa_features(
- msas=final_msa,
- deletion_matrices=final_deletion_mat,
+ msas=final_msa_obj
)
template_feature_list = []
@@ -906,7 +1118,10 @@ def process_multiseq_fasta(self,
alignment_dir = os.path.join(
super_alignment_dir, desc
)
- hits = self._parse_template_hits(alignment_dir, alignment_index=None)
+ hits = self._parse_template_hit_files(alignment_dir=alignment_dir,
+ input_sequence=seq,
+ alignment_index=None)
+
template_features = make_template_features(
seq,
hits,
@@ -918,6 +1133,228 @@ def process_multiseq_fasta(self,
return {
**sequence_features,
- **msa_features,
+ **msa_features,
**template_features,
}
+
+
+class DataPipelineMultimer:
+ """Runs the alignment tools and assembles the input features."""
+
+ def __init__(self,
+ monomer_data_pipeline: DataPipeline,
+ ):
+ """Initializes the data pipeline.
+ Args:
+ monomer_data_pipeline: An instance of pipeline.DataPipeline - that runs
+ the data pipeline for the monomer AlphaFold system.
+ jackhmmer_binary_path: Location of the jackhmmer binary.
+ uniprot_database_path: Location of the unclustered uniprot sequences, that
+ will be searched with jackhmmer and used for MSA pairing.
+ max_uniprot_hits: The maximum number of hits to return from uniprot.
+ use_precomputed_msas: Whether to use pre-existing MSAs; see run_alphafold.
+ """
+ self._monomer_data_pipeline = monomer_data_pipeline
+
+ def _process_single_chain(
+ self,
+ chain_id: str,
+ sequence: str,
+ description: str,
+ chain_alignment_dir: str,
+ chain_alignment_index: Optional[Any],
+ is_homomer_or_monomer: bool
+ ) -> FeatureDict:
+ """Runs the monomer pipeline on a single chain."""
+ chain_fasta_str = f'>{chain_id}\n{sequence}\n'
+
+ if chain_alignment_index is None and not os.path.exists(chain_alignment_dir):
+ raise ValueError(f"Alignments for {chain_id} not found...")
+
+ with temp_fasta_file(chain_fasta_str) as chain_fasta_path:
+ chain_features = self._monomer_data_pipeline.process_fasta(
+ fasta_path=chain_fasta_path,
+ alignment_dir=chain_alignment_dir,
+ alignment_index=chain_alignment_index
+ )
+
+ # We only construct the pairing features if there are 2 or more unique
+ # sequences.
+ if not is_homomer_or_monomer:
+ all_seq_msa_features = self._all_seq_msa_features(
+ chain_alignment_dir,
+ chain_alignment_index
+ )
+ chain_features.update(all_seq_msa_features)
+ return chain_features
+
+ @staticmethod
+ def _all_seq_msa_features(alignment_dir, alignment_index):
+ """Get MSA features for unclustered uniprot, for pairing."""
+ if alignment_index is not None:
+ fp = open(os.path.join(alignment_dir, alignment_index["db"]), "rb")
+
+ def read_msa(start, size):
+ fp.seek(start)
+ msa = fp.read(size).decode("utf-8")
+ return msa
+
+ start, size = next(iter((start, size) for name, start, size in alignment_index["files"]
+ if name == 'uniprot_hits.sto'))
+
+ msa = parsers.parse_stockholm(read_msa(start, size))
+ fp.close()
+ else:
+ uniprot_msa_path = os.path.join(alignment_dir, "uniprot_hits.sto")
+ if not os.path.exists(uniprot_msa_path):
+ chain_id = os.path.basename(os.path.normpath(alignment_dir))
+ raise ValueError(f"Missing 'uniprot_hits.sto' for {chain_id}. "
+ f"This is required for Multimer MSA pairing.")
+
+ with open(uniprot_msa_path, "r") as fp:
+ uniprot_msa_string = fp.read()
+ msa = parsers.parse_stockholm(uniprot_msa_string)
+
+ all_seq_features = make_msa_features([msa])
+ valid_feats = msa_pairing.MSA_FEATURES + (
+ 'msa_species_identifiers',
+ )
+ feats = {
+ f'{k}_all_seq': v for k, v in all_seq_features.items()
+ if k in valid_feats
+ }
+ return feats
+
+ def process_fasta(self,
+ fasta_path: str,
+ alignment_dir: str,
+ alignment_index: Optional[Any] = None
+ ) -> FeatureDict:
+ """Creates features."""
+ with open(fasta_path) as f:
+ input_fasta_str = f.read()
+
+ input_seqs, input_descs = parsers.parse_fasta(input_fasta_str)
+
+ all_chain_features = {}
+ sequence_features = {}
+ is_homomer_or_monomer = len(set(input_seqs)) == 1
+ for desc, seq in zip(input_descs, input_seqs):
+ if seq in sequence_features:
+ all_chain_features[desc] = copy.deepcopy(
+ sequence_features[seq]
+ )
+ continue
+
+ if alignment_index is not None:
+ chain_alignment_index = alignment_index.get(desc)
+ chain_alignment_dir = alignment_dir
+ else:
+ chain_alignment_index = None
+ chain_alignment_dir = os.path.join(alignment_dir, desc)
+
+ chain_features = self._process_single_chain(
+ chain_id=desc,
+ sequence=seq,
+ description=desc,
+ chain_alignment_dir=chain_alignment_dir,
+ chain_alignment_index=chain_alignment_index,
+ is_homomer_or_monomer=is_homomer_or_monomer
+ )
+
+ chain_features = convert_monomer_features(
+ chain_features,
+ chain_id=desc
+ )
+ all_chain_features[desc] = chain_features
+ sequence_features[seq] = chain_features
+
+ all_chain_features = add_assembly_features(all_chain_features)
+
+ np_example = feature_processing_multimer.pair_and_merge(
+ all_chain_features=all_chain_features,
+ )
+
+ # Pad MSA to avoid zero-sized extra_msa.
+ np_example = pad_msa(np_example, 512)
+
+ return np_example
+
+ def get_mmcif_features(
+ self, mmcif_object: mmcif_parsing.MmcifObject, chain_id: str
+ ) -> FeatureDict:
+ mmcif_feats = {}
+
+ all_atom_positions, all_atom_mask = mmcif_parsing.get_atom_coords(
+ mmcif_object=mmcif_object, chain_id=chain_id
+ )
+ mmcif_feats["all_atom_positions"] = all_atom_positions
+ mmcif_feats["all_atom_mask"] = all_atom_mask
+
+ mmcif_feats["resolution"] = np.array(
+ mmcif_object.header["resolution"], dtype=np.float32
+ )
+
+ mmcif_feats["release_date"] = np.array(
+ [mmcif_object.header["release_date"].encode("utf-8")], dtype=np.object_
+ )
+
+ mmcif_feats["is_distillation"] = np.array(0., dtype=np.float32)
+
+ return mmcif_feats
+
+ def process_mmcif(
+ self,
+ mmcif: mmcif_parsing.MmcifObject, # parsing is expensive, so no path
+ alignment_dir: str,
+ alignment_index: Optional[Any] = None,
+ ) -> FeatureDict:
+
+ all_chain_features = {}
+ sequence_features = {}
+ is_homomer_or_monomer = len(set(list(mmcif.chain_to_seqres.values()))) == 1
+ for chain_id, seq in mmcif.chain_to_seqres.items():
+ desc= "_".join([mmcif.file_id, chain_id])
+
+ if seq in sequence_features:
+ all_chain_features[desc] = copy.deepcopy(
+ sequence_features[seq]
+ )
+ continue
+
+ if alignment_index is not None:
+ chain_alignment_index = alignment_index.get(desc)
+ chain_alignment_dir = alignment_dir
+ else:
+ chain_alignment_index = None
+ chain_alignment_dir = os.path.join(alignment_dir, desc)
+
+ chain_features = self._process_single_chain(
+ chain_id=desc,
+ sequence=seq,
+ description=desc,
+ chain_alignment_dir=chain_alignment_dir,
+ chain_alignment_index=chain_alignment_index,
+ is_homomer_or_monomer=is_homomer_or_monomer
+ )
+
+ chain_features = convert_monomer_features(
+ chain_features,
+ chain_id=desc
+ )
+
+ mmcif_feats = self.get_mmcif_features(mmcif, chain_id)
+ chain_features.update(mmcif_feats)
+ all_chain_features[desc] = chain_features
+ sequence_features[seq] = chain_features
+
+ all_chain_features = add_assembly_features(all_chain_features)
+
+ np_example = feature_processing_multimer.pair_and_merge(
+ all_chain_features=all_chain_features,
+ )
+
+ # Pad MSA to avoid zero-sized extra_msa.
+ np_example = pad_msa(np_example, 512)
+
+ return np_example
diff --git a/openfold/data/data_transforms.py b/openfold/data/data_transforms.py
index c2e20be1..bd306a11 100755
--- a/openfold/data/data_transforms.py
+++ b/openfold/data/data_transforms.py
@@ -23,6 +23,9 @@
from openfold.config import NUM_RES, NUM_EXTRA_SEQ, NUM_TEMPLATES, NUM_MSA_SEQ
from openfold.np import residue_constants as rc
from openfold.utils.rigid_utils import Rotation, Rigid
+from openfold.utils.geometry.rigid_matrix_vector import Rigid3Array
+from openfold.utils.geometry.rotation_matrix import Rot3Array
+from openfold.utils.geometry.vector import Vec3Array
from openfold.utils.tensor_utils import (
tree_map,
tensor_tree_map,
@@ -86,18 +89,17 @@ def make_all_atom_aatype(protein):
def fix_templates_aatype(protein):
# Map one-hot to indices
num_templates = protein["template_aatype"].shape[0]
- if(num_templates > 0):
- protein["template_aatype"] = torch.argmax(
- protein["template_aatype"], dim=-1
- )
- # Map hhsearch-aatype to our aatype.
- new_order_list = rc.MAP_HHBLITS_AATYPE_TO_OUR_AATYPE
- new_order = torch.tensor(
- new_order_list, dtype=torch.int64, device=protein["aatype"].device,
- ).expand(num_templates, -1)
- protein["template_aatype"] = torch.gather(
- new_order, 1, index=protein["template_aatype"]
- )
+ protein["template_aatype"] = torch.argmax(
+ protein["template_aatype"], dim=-1
+ )
+ # Map hhsearch-aatype to our aatype.
+ new_order_list = rc.MAP_HHBLITS_AATYPE_TO_OUR_AATYPE
+ new_order = torch.tensor(
+ new_order_list, dtype=torch.int64, device=protein["template_aatype"].device,
+ ).expand(num_templates, -1)
+ protein["template_aatype"] = torch.gather(
+ new_order, 1, index=protein["template_aatype"]
+ )
return protein
@@ -447,13 +449,15 @@ def make_hhblits_profile(protein):
@curry1
-def make_masked_msa(protein, config, replace_fraction):
+def make_masked_msa(protein, config, replace_fraction, seed):
"""Create data for BERT on raw MSA."""
+ device = protein["msa"].device
+
# Add a random amino acid uniformly.
random_aa = torch.tensor(
[0.05] * 20 + [0.0, 0.0],
dtype=torch.float32,
- device=protein["aatype"].device
+ device=device
)
categorical_probs = (
@@ -473,11 +477,18 @@ def make_masked_msa(protein, config, replace_fraction):
assert mask_prob >= 0.0
categorical_probs = torch.nn.functional.pad(
- categorical_probs, pad_shapes, value=mask_prob
+ categorical_probs, pad_shapes, value=mask_prob,
)
sh = protein["msa"].shape
- mask_position = torch.rand(sh) < replace_fraction
+
+ g = None
+ if seed is not None:
+ g = torch.Generator(device=protein["msa"].device)
+ g.manual_seed(seed)
+
+ sample = torch.rand(sh, device=device, generator=g)
+ mask_position = sample < replace_fraction
bert_msa = shaped_categorical(categorical_probs)
bert_msa = torch.where(mask_position, bert_msa, protein["msa"])
@@ -670,7 +681,7 @@ def make_atom14_masks(protein):
def make_atom14_masks_np(batch):
batch = tree_map(
lambda n: torch.tensor(n, device="cpu"),
- batch,
+ batch,
np.ndarray
)
out = make_atom14_masks(batch)
@@ -736,7 +747,7 @@ def make_atom14_positions(protein):
for index, correspondence in enumerate(correspondences):
renaming_matrix[index, correspondence] = 1.0
all_matrices[resname] = renaming_matrix
-
+
renaming_matrices = torch.stack(
[all_matrices[restype] for restype in restype_3]
)
@@ -782,10 +793,14 @@ def make_atom14_positions(protein):
def atom37_to_frames(protein, eps=1e-8):
+ is_multimer = "asym_id" in protein
aatype = protein["aatype"]
all_atom_positions = protein["all_atom_positions"]
all_atom_mask = protein["all_atom_mask"]
+ if is_multimer:
+ all_atom_positions = Vec3Array.from_array(all_atom_positions)
+
batch_dims = len(aatype.shape[:-1])
restype_rigidgroup_base_atom_names = np.full([21, 8, 3], "", dtype=object)
@@ -832,19 +847,37 @@ def atom37_to_frames(protein, eps=1e-8):
no_batch_dims=batch_dims,
)
- base_atom_pos = batched_gather(
- all_atom_positions,
- residx_rigidgroup_base_atom37_idx,
- dim=-2,
- no_batch_dims=len(all_atom_positions.shape[:-2]),
- )
+ if is_multimer:
+ base_atom_pos = [batched_gather(
+ pos,
+ residx_rigidgroup_base_atom37_idx,
+ dim=-1,
+ no_batch_dims=len(all_atom_positions.shape[:-1]),
+ ) for pos in all_atom_positions]
+ base_atom_pos = Vec3Array.from_array(torch.stack(base_atom_pos, dim=-1))
+ else:
+ base_atom_pos = batched_gather(
+ all_atom_positions,
+ residx_rigidgroup_base_atom37_idx,
+ dim=-2,
+ no_batch_dims=len(all_atom_positions.shape[:-2]),
+ )
- gt_frames = Rigid.from_3_points(
- p_neg_x_axis=base_atom_pos[..., 0, :],
- origin=base_atom_pos[..., 1, :],
- p_xy_plane=base_atom_pos[..., 2, :],
- eps=eps,
- )
+ if is_multimer:
+ point_on_neg_x_axis = base_atom_pos[:, :, 0]
+ origin = base_atom_pos[:, :, 1]
+ point_on_xy_plane = base_atom_pos[:, :, 2]
+ gt_rotation = Rot3Array.from_two_vectors(
+ origin - point_on_neg_x_axis, point_on_xy_plane - origin)
+
+ gt_frames = Rigid3Array(gt_rotation, origin)
+ else:
+ gt_frames = Rigid.from_3_points(
+ p_neg_x_axis=base_atom_pos[..., 0, :],
+ origin=base_atom_pos[..., 1, :],
+ p_xy_plane=base_atom_pos[..., 2, :],
+ eps=eps,
+ )
group_exists = batched_gather(
restype_rigidgroup_mask,
@@ -865,9 +898,13 @@ def atom37_to_frames(protein, eps=1e-8):
rots = torch.tile(rots, (*((1,) * batch_dims), 8, 1, 1))
rots[..., 0, 0, 0] = -1
rots[..., 0, 2, 2] = -1
- rots = Rotation(rot_mats=rots)
- gt_frames = gt_frames.compose(Rigid(rots, None))
+ if is_multimer:
+ gt_frames = gt_frames.compose_rotation(
+ Rot3Array.from_array(rots))
+ else:
+ rots = Rotation(rot_mats=rots)
+ gt_frames = gt_frames.compose(Rigid(rots, None))
restype_rigidgroup_is_ambiguous = all_atom_mask.new_zeros(
*((1,) * batch_dims), 21, 8
@@ -901,12 +938,18 @@ def atom37_to_frames(protein, eps=1e-8):
no_batch_dims=batch_dims,
)
- residx_rigidgroup_ambiguity_rot = Rotation(
- rot_mats=residx_rigidgroup_ambiguity_rot
- )
- alt_gt_frames = gt_frames.compose(
- Rigid(residx_rigidgroup_ambiguity_rot, None)
- )
+ if is_multimer:
+ ambiguity_rot = Rot3Array.from_array(residx_rigidgroup_ambiguity_rot)
+
+ # Create the alternative ground truth frames.
+ alt_gt_frames = gt_frames.compose_rotation(ambiguity_rot)
+ else:
+ residx_rigidgroup_ambiguity_rot = Rotation(
+ rot_mats=residx_rigidgroup_ambiguity_rot
+ )
+ alt_gt_frames = gt_frames.compose(
+ Rigid(residx_rigidgroup_ambiguity_rot, None)
+ )
gt_frames_tensor = gt_frames.to_tensor_4x4()
alt_gt_frames_tensor = alt_gt_frames.to_tensor_4x4()
diff --git a/openfold/data/data_transforms_multimer.py b/openfold/data/data_transforms_multimer.py
new file mode 100644
index 00000000..fcdfb3d7
--- /dev/null
+++ b/openfold/data/data_transforms_multimer.py
@@ -0,0 +1,493 @@
+from typing import Sequence
+
+import torch
+
+from openfold.config import NUM_RES
+from openfold.data.data_transforms import curry1
+from openfold.np import residue_constants as rc
+from openfold.utils.tensor_utils import masked_mean
+
+
+def gumbel_noise(
+ shape: Sequence[int],
+ device: torch.device,
+ eps=1e-6,
+ generator=None,
+) -> torch.Tensor:
+ """Generate Gumbel Noise of given Shape.
+
+ This generates samples from Gumbel(0, 1).
+
+ Args:
+ shape: Shape of noise to return.
+
+ Returns:
+ Gumbel noise of given shape.
+ """
+ uniform_noise = torch.rand(
+ shape, dtype=torch.float32, device=device, generator=generator
+ )
+ gumbel = -torch.log(-torch.log(uniform_noise + eps) + eps)
+ return gumbel
+
+
+def gumbel_max_sample(logits: torch.Tensor, generator=None) -> torch.Tensor:
+ """Samples from a probability distribution given by 'logits'.
+
+ This uses Gumbel-max trick to implement the sampling in an efficient manner.
+
+ Args:
+ logits: Logarithm of probabilities to sample from, probabilities can be
+ unnormalized.
+
+ Returns:
+ Sample from logprobs in one-hot form.
+ """
+ z = gumbel_noise(logits.shape, device=logits.device, generator=generator)
+ return torch.nn.functional.one_hot(
+ torch.argmax(logits + z, dim=-1),
+ logits.shape[-1],
+ )
+
+
+def gumbel_argsort_sample_idx(
+ logits: torch.Tensor,
+ generator=None
+) -> torch.Tensor:
+ """Samples with replacement from a distribution given by 'logits'.
+
+ This uses Gumbel trick to implement the sampling an efficient manner. For a
+ distribution over k items this samples k times without replacement, so this
+ is effectively sampling a random permutation with probabilities over the
+ permutations derived from the logprobs.
+
+ Args:
+ logits: Logarithm of probabilities to sample from, probabilities can be
+ unnormalized.
+
+ Returns:
+ Sample from logprobs in one-hot form.
+ """
+ z = gumbel_noise(logits.shape, device=logits.device, generator=generator)
+ return torch.argsort(logits + z, dim=-1, descending=True)
+
+
+@curry1
+def make_masked_msa(batch, config, replace_fraction, seed, eps=1e-6):
+ """Create data for BERT on raw MSA."""
+ # Add a random amino acid uniformly.
+ random_aa = torch.Tensor(
+ [0.05] * 20 + [0., 0.],
+ device=batch['msa'].device
+ )
+
+ categorical_probs = (
+ config.uniform_prob * random_aa +
+ config.profile_prob * batch['msa_profile'] +
+ config.same_prob * torch.nn.functional.one_hot(batch['msa'], 22)
+ )
+
+ # Put all remaining probability on [MASK] which is a new column.
+ mask_prob = 1. - config.profile_prob - config.same_prob - config.uniform_prob
+
+ categorical_probs = torch.nn.functional.pad(
+ categorical_probs, [0,1], value=mask_prob
+ )
+
+ sh = batch['msa'].shape
+ mask_position = torch.rand(sh, device=batch['msa'].device) < replace_fraction
+ mask_position *= batch['msa_mask'].to(mask_position.dtype)
+
+ logits = torch.log(categorical_probs + eps)
+
+ g = None
+ if seed is not None:
+ g = torch.Generator(device=batch["msa"].device)
+ g.manual_seed(seed)
+
+ bert_msa = gumbel_max_sample(logits, generator=g)
+
+ bert_msa = torch.where(
+ mask_position,
+ torch.argmax(bert_msa, dim=-1),
+ batch['msa']
+ )
+ bert_msa *= batch['msa_mask'].to(bert_msa.dtype)
+
+ # Mix real and masked MSA.
+ if 'bert_mask' in batch:
+ batch['bert_mask'] *= mask_position.to(torch.float32)
+ else:
+ batch['bert_mask'] = mask_position.to(torch.float32)
+ batch['true_msa'] = batch['msa']
+ batch['msa'] = bert_msa
+
+ return batch
+
+
+@curry1
+def nearest_neighbor_clusters(batch, gap_agreement_weight=0.):
+ """Assign each extra MSA sequence to its nearest neighbor in sampled MSA."""
+ device = batch["msa_mask"].device
+
+ # Determine how much weight we assign to each agreement. In theory, we could
+ # use a full blosum matrix here, but right now let's just down-weight gap
+ # agreement because it could be spurious.
+ # Never put weight on agreeing on BERT mask.
+
+ weights = torch.Tensor(
+ [1.] * 21 + [gap_agreement_weight] + [0.],
+ device=device,
+ )
+
+ msa_mask = batch['msa_mask']
+ msa_one_hot = torch.nn.functional.one_hot(batch['msa'], 23)
+
+ extra_mask = batch['extra_msa_mask']
+ extra_one_hot = torch.nn.functional.one_hot(batch['extra_msa'], 23)
+
+ msa_one_hot_masked = msa_mask[:, :, None] * msa_one_hot
+ extra_one_hot_masked = extra_mask[:, :, None] * extra_one_hot
+
+ agreement = torch.einsum(
+ 'mrc, nrc->nm',
+ extra_one_hot_masked,
+ weights * msa_one_hot_masked
+ )
+
+ cluster_assignment = torch.nn.functional.softmax(1e3 * agreement, dim=0)
+ cluster_assignment *= torch.einsum('mr, nr->mn', msa_mask, extra_mask)
+
+ cluster_count = torch.sum(cluster_assignment, dim=-1)
+ cluster_count += 1. # We always include the sequence itself.
+
+ msa_sum = torch.einsum('nm, mrc->nrc', cluster_assignment, extra_one_hot_masked)
+ msa_sum += msa_one_hot_masked
+
+ cluster_profile = msa_sum / cluster_count[:, None, None]
+
+ extra_deletion_matrix = batch['extra_deletion_matrix']
+ deletion_matrix = batch['deletion_matrix']
+
+ del_sum = torch.einsum(
+ 'nm, mc->nc',
+ cluster_assignment,
+ extra_mask * extra_deletion_matrix
+ )
+ del_sum += deletion_matrix # Original sequence.
+ cluster_deletion_mean = del_sum / cluster_count[:, None]
+
+ batch['cluster_profile'] = cluster_profile
+ batch['cluster_deletion_mean'] = cluster_deletion_mean
+
+ return batch
+
+
+def create_target_feat(batch):
+ """Create the target features"""
+ batch["target_feat"] = torch.nn.functional.one_hot(
+ batch["aatype"], 21
+ ).to(torch.float32)
+ return batch
+
+
+def create_msa_feat(batch):
+ """Create and concatenate MSA features."""
+ device = batch["msa"]
+ msa_1hot = torch.nn.functional.one_hot(batch['msa'], 23)
+ deletion_matrix = batch['deletion_matrix']
+ has_deletion = torch.clamp(deletion_matrix, min=0., max=1.)[..., None]
+ pi = torch.acos(torch.zeros(1, device=deletion_matrix.device)) * 2
+ deletion_value = (torch.atan(deletion_matrix / 3.) * (2. / pi))[..., None]
+
+ deletion_mean_value = (
+ torch.atan(
+ batch['cluster_deletion_mean'] / 3.) *
+ (2. / pi)
+ )[..., None]
+
+ msa_feat = torch.cat(
+ [
+ msa_1hot,
+ has_deletion,
+ deletion_value,
+ batch['cluster_profile'],
+ deletion_mean_value
+ ],
+ dim=-1,
+ )
+
+ batch["msa_feat"] = msa_feat
+
+ return batch
+
+
+def build_extra_msa_feat(batch):
+ """Expand extra_msa into 1hot and concat with other extra msa features.
+
+ We do this as late as possible as the one_hot extra msa can be very large.
+
+ Args:
+ batch: a dictionary with the following keys:
+ * 'extra_msa': [num_seq, num_res] MSA that wasn't selected as a cluster
+ centre. Note - This isn't one-hotted.
+ * 'extra_deletion_matrix': [num_seq, num_res] Number of deletions at given
+ position.
+ num_extra_msa: Number of extra msa to use.
+
+ Returns:
+ Concatenated tensor of extra MSA features.
+ """
+ # 23 = 20 amino acids + 'X' for unknown + gap + bert mask
+ extra_msa = batch['extra_msa']
+ deletion_matrix = batch['extra_deletion_matrix']
+ msa_1hot = torch.nn.functional.one_hot(extra_msa, 23)
+ has_deletion = torch.clamp(deletion_matrix, min=0., max=1.)[..., None]
+ pi = torch.acos(torch.zeros(1, device=deletion_matrix.device)) * 2
+ deletion_value = (
+ (torch.atan(deletion_matrix / 3.) * (2. / pi))[..., None]
+ )
+ extra_msa_mask = batch['extra_msa_mask']
+ catted = torch.cat([msa_1hot, has_deletion, deletion_value], dim=-1)
+
+ return catted
+
+
+@curry1
+def sample_msa(batch, max_seq, max_extra_msa_seq, seed, inf=1e6):
+ """Sample MSA randomly, remaining sequences are stored as `extra_*`.
+
+ Args:
+ batch: batch to sample msa from.
+ max_seq: number of sequences to sample.
+ Returns:
+ Protein with sampled msa.
+ """
+ g = None
+ if seed is not None:
+ g = torch.Generator(device=batch["msa"].device)
+ g.manual_seed(seed)
+
+ # Sample uniformly among sequences with at least one non-masked position.
+ logits = (torch.clamp(torch.sum(batch['msa_mask'], dim=-1), 0., 1.) - 1.) * inf
+ # The cluster_bias_mask can be used to preserve the first row (target
+ # sequence) for each chain, for example.
+ if 'cluster_bias_mask' not in batch:
+ cluster_bias_mask = torch.nn.functional.pad(
+ batch['msa'].new_zeros(batch['msa'].shape[0] - 1),
+ (1, 0),
+ value=1.
+ )
+ else:
+ cluster_bias_mask = batch['cluster_bias_mask']
+
+ logits += cluster_bias_mask * inf
+ index_order = gumbel_argsort_sample_idx(logits, generator=g)
+ sel_idx = index_order[:max_seq]
+ extra_idx = index_order[max_seq:][:max_extra_msa_seq]
+
+ for k in ['msa', 'deletion_matrix', 'msa_mask', 'bert_mask']:
+ if k in batch:
+ batch['extra_' + k] = batch[k][extra_idx]
+ batch[k] = batch[k][sel_idx]
+
+ return batch
+
+
+def make_msa_profile(batch):
+ """Compute the MSA profile."""
+
+ # Compute the profile for every residue (over all MSA sequences).
+ batch["msa_profile"] = masked_mean(
+ batch['msa_mask'][..., None],
+ torch.nn.functional.one_hot(batch['msa'], 22),
+ dim=-3,
+ )
+
+ return batch
+
+
+def randint(lower, upper, generator, device):
+ return int(torch.randint(
+ lower,
+ upper + 1,
+ (1,),
+ device=device,
+ generator=generator,
+ )[0])
+
+
+def get_interface_residues(positions, atom_mask, asym_id, interface_threshold):
+ coord_diff = positions[..., None, :, :] - positions[..., None, :, :, :]
+ pairwise_dists = torch.sqrt(torch.sum(coord_diff ** 2, dim=-1))
+
+ diff_chain_mask = (asym_id[..., None, :] != asym_id[..., :, None]).float()
+ pair_mask = atom_mask[..., None, :] * atom_mask[..., None, :, :]
+ mask = (diff_chain_mask[..., None] * pair_mask).bool()
+
+ min_dist_per_res, _ = torch.where(mask, pairwise_dists, torch.inf).min(dim=-1)
+
+ valid_interfaces = torch.sum((min_dist_per_res < interface_threshold).float(), dim=-1)
+ interface_residues_idxs = torch.nonzero(valid_interfaces, as_tuple=True)[0]
+
+ return interface_residues_idxs
+
+
+def get_spatial_crop_idx(protein, crop_size, interface_threshold, generator):
+ positions = protein["all_atom_positions"]
+ atom_mask = protein["all_atom_mask"]
+ asym_id = protein["asym_id"]
+
+ interface_residues = get_interface_residues(positions=positions,
+ atom_mask=atom_mask,
+ asym_id=asym_id,
+ interface_threshold=interface_threshold)
+
+ if not torch.any(interface_residues):
+ return get_contiguous_crop_idx(protein, crop_size, generator)
+
+ target_res_idx = randint(lower=0,
+ upper=interface_residues.shape[-1] - 1,
+ generator=generator,
+ device=positions.device)
+
+ target_res = interface_residues[target_res_idx]
+
+ ca_idx = rc.atom_order["CA"]
+ ca_positions = positions[..., ca_idx, :]
+ ca_mask = atom_mask[..., ca_idx].bool()
+
+ coord_diff = ca_positions[..., None, :] - ca_positions[..., None, :, :]
+ ca_pairwise_dists = torch.sqrt(torch.sum(coord_diff ** 2, dim=-1))
+
+ to_target_distances = ca_pairwise_dists[target_res]
+ break_tie = (
+ torch.arange(
+ 0, to_target_distances.shape[-1], device=positions.device
+ ).float()
+ * 1e-3
+ )
+ to_target_distances = torch.where(ca_mask, to_target_distances, torch.inf) + break_tie
+
+ ret = torch.argsort(to_target_distances)[:crop_size]
+ return ret.sort().values
+
+
+def get_contiguous_crop_idx(protein, crop_size, generator):
+ unique_asym_ids, chain_idxs, chain_lens = protein["asym_id"].unique(dim=-1,
+ return_inverse=True,
+ return_counts=True)
+
+ shuffle_idx = torch.randperm(chain_lens.shape[-1], device=chain_lens.device, generator=generator)
+
+ _, idx_sorted = torch.sort(chain_idxs, stable=True)
+ cum_sum = chain_lens.cumsum(dim=0)
+ cum_sum = torch.cat((torch.tensor([0]), cum_sum[:-1]), dim=0)
+ asym_offsets = idx_sorted[cum_sum]
+
+ num_budget = crop_size
+ num_remaining = int(protein["seq_length"])
+
+ crop_idxs = []
+ for idx in shuffle_idx:
+ chain_len = int(chain_lens[idx])
+ num_remaining -= chain_len
+
+ crop_size_max = min(num_budget, chain_len)
+ crop_size_min = min(chain_len, max(0, num_budget - num_remaining))
+ chain_crop_size = randint(lower=crop_size_min,
+ upper=crop_size_max,
+ generator=generator,
+ device=chain_lens.device)
+
+ num_budget -= chain_crop_size
+
+ chain_start = randint(lower=0,
+ upper=chain_len - chain_crop_size,
+ generator=generator,
+ device=chain_lens.device)
+
+ asym_offset = asym_offsets[idx]
+ crop_idxs.append(
+ torch.arange(asym_offset + chain_start, asym_offset + chain_start + chain_crop_size)
+ )
+
+ return torch.concat(crop_idxs).sort().values
+
+
+@curry1
+def random_crop_to_size(
+ protein,
+ crop_size,
+ max_templates,
+ shape_schema,
+ spatial_crop_prob,
+ interface_threshold,
+ subsample_templates=False,
+ seed=None,
+):
+ """Crop randomly to `crop_size`, or keep as is if shorter than that."""
+ # We want each ensemble to be cropped the same way
+ g = None
+ if seed is not None:
+ g = torch.Generator(device=protein["seq_length"].device)
+ g.manual_seed(seed)
+
+ use_spatial_crop = torch.rand((1,),
+ device=protein["seq_length"].device,
+ generator=g) < spatial_crop_prob
+
+ num_res = protein["aatype"].shape[0]
+ if num_res <= crop_size:
+ crop_idxs = torch.arange(num_res)
+ elif use_spatial_crop:
+ crop_idxs = get_spatial_crop_idx(protein, crop_size, interface_threshold, g)
+ else:
+ crop_idxs = get_contiguous_crop_idx(protein, crop_size, g)
+
+ if "template_mask" in protein:
+ num_templates = protein["template_mask"].shape[-1]
+ else:
+ num_templates = 0
+
+ # No need to subsample templates if there aren't any
+ subsample_templates = subsample_templates and num_templates
+
+ if subsample_templates:
+ templates_crop_start = randint(lower=0,
+ upper=num_templates,
+ generator=g,
+ device=protein["seq_length"].device)
+ templates_select_indices = torch.randperm(
+ num_templates, device=protein["seq_length"].device, generator=g
+ )
+ else:
+ templates_crop_start = 0
+
+ num_res_crop_size = min(int(protein["seq_length"]), crop_size)
+ num_templates_crop_size = min(
+ num_templates - templates_crop_start, max_templates
+ )
+
+ for k, v in protein.items():
+ if k not in shape_schema or (
+ "template" not in k and NUM_RES not in shape_schema[k]
+ ):
+ continue
+
+ # randomly permute the templates before cropping them.
+ if k.startswith("template") and subsample_templates:
+ v = v[templates_select_indices]
+
+ for i, (dim_size, dim) in enumerate(zip(shape_schema[k], v.shape)):
+ is_num_res = dim_size == NUM_RES
+ if i == 0 and k.startswith("template"):
+ v = v[slice(templates_crop_start, templates_crop_start + num_templates_crop_size)]
+ elif is_num_res:
+ v = torch.index_select(v, i, crop_idxs)
+
+ protein[k] = v
+
+ protein["seq_length"] = protein["seq_length"].new_tensor(num_res_crop_size)
+
+ return protein
diff --git a/openfold/data/feature_pipeline.py b/openfold/data/feature_pipeline.py
index ab378438..a8682441 100644
--- a/openfold/data/feature_pipeline.py
+++ b/openfold/data/feature_pipeline.py
@@ -20,7 +20,7 @@
import numpy as np
import torch
-from openfold.data import input_pipeline
+from openfold.data import input_pipeline, input_pipeline_multimer
FeatureDict = Mapping[str, np.ndarray]
@@ -80,11 +80,14 @@ def np_example_to_features(
np_example: FeatureDict,
config: ml_collections.ConfigDict,
mode: str,
+ is_multimer: bool = False
):
np_example = dict(np_example)
- num_res = int(np_example["seq_length"][0])
- cfg, feature_names = make_data_config(config, mode=mode, num_res=num_res)
+ seq_length = np_example["seq_length"]
+ num_res = int(seq_length[0]) if seq_length.ndim != 0 else int(seq_length)
+ cfg, feature_names = make_data_config(config, mode=mode, num_res=num_res)
+
if "deletion_matrix_int" in np_example:
np_example["deletion_matrix"] = np_example.pop(
"deletion_matrix_int"
@@ -93,12 +96,20 @@ def np_example_to_features(
tensor_dict = np_to_tensor_dict(
np_example=np_example, features=feature_names
)
+
with torch.no_grad():
- features = input_pipeline.process_tensors_from_config(
- tensor_dict,
- cfg.common,
- cfg[mode],
- )
+ if is_multimer:
+ features = input_pipeline_multimer.process_tensors_from_config(
+ tensor_dict,
+ cfg.common,
+ cfg[mode],
+ )
+ else:
+ features = input_pipeline.process_tensors_from_config(
+ tensor_dict,
+ cfg.common,
+ cfg[mode],
+ )
if mode == "train":
p = torch.rand(1).item()
@@ -128,10 +139,15 @@ def __init__(
def process_features(
self,
raw_features: FeatureDict,
- mode: str = "train",
+ mode: str = "train",
+ is_multimer: bool = False,
) -> FeatureDict:
+ # if(is_multimer and mode != "predict"):
+ # raise ValueError("Multimer mode is not currently trainable")
+
return np_example_to_features(
np_example=raw_features,
config=self.config,
mode=mode,
+ is_multimer=is_multimer,
)
diff --git a/openfold/data/feature_processing_multimer.py b/openfold/data/feature_processing_multimer.py
new file mode 100644
index 00000000..518babf2
--- /dev/null
+++ b/openfold/data/feature_processing_multimer.py
@@ -0,0 +1,245 @@
+# Copyright 2021 DeepMind Technologies Limited
+# Copyright 2022 AlQuraishi Laboratory
+#
+# 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.
+
+"""Feature processing logic for multimer data pipeline."""
+
+from typing import Iterable, MutableMapping, List, Mapping
+
+from openfold.data import msa_pairing
+from openfold.np import residue_constants
+import numpy as np
+
+
+# TODO: Move this into the config
+REQUIRED_FEATURES = frozenset({
+ 'aatype', 'all_atom_mask', 'all_atom_positions', 'all_chains_entity_ids',
+ 'all_crops_all_chains_mask', 'all_crops_all_chains_positions',
+ 'all_crops_all_chains_residue_ids', 'assembly_num_chains', 'asym_id',
+ 'bert_mask', 'cluster_bias_mask', 'deletion_matrix', 'deletion_mean',
+ 'entity_id', 'entity_mask', 'mem_peak', 'msa', 'msa_mask', 'num_alignments',
+ 'num_templates', 'queue_size', 'residue_index', 'resolution',
+ 'seq_length', 'seq_mask', 'sym_id', 'template_aatype',
+ 'template_all_atom_mask', 'template_all_atom_positions'
+})
+
+MAX_TEMPLATES = 4
+MSA_CROP_SIZE = 2048
+
+
+def _is_homomer_or_monomer(chains: Iterable[Mapping[str, np.ndarray]]) -> bool:
+ """Checks if a list of chains represents a homomer/monomer example."""
+ # Note that an entity_id of 0 indicates padding.
+ num_unique_chains = len(np.unique(np.concatenate(
+ [np.unique(chain['entity_id'][chain['entity_id'] > 0]) for
+ chain in chains])))
+ return num_unique_chains == 1
+
+
+def pair_and_merge(
+ all_chain_features: MutableMapping[str, Mapping[str, np.ndarray]],
+) -> Mapping[str, np.ndarray]:
+ """Runs processing on features to augment, pair and merge.
+
+ Args:
+ all_chain_features: A MutableMap of dictionaries of features for each chain.
+
+ Returns:
+ A dictionary of features.
+ """
+
+ process_unmerged_features(all_chain_features)
+
+ np_chains_list = list(all_chain_features.values())
+
+ pair_msa_sequences = not _is_homomer_or_monomer(np_chains_list)
+
+ if pair_msa_sequences:
+ np_chains_list = msa_pairing.create_paired_features(
+ chains=np_chains_list
+ )
+ np_chains_list = msa_pairing.deduplicate_unpaired_sequences(np_chains_list)
+ np_chains_list = crop_chains(
+ np_chains_list,
+ msa_crop_size=MSA_CROP_SIZE,
+ pair_msa_sequences=pair_msa_sequences,
+ max_templates=MAX_TEMPLATES
+ )
+ np_example = msa_pairing.merge_chain_features(
+ np_chains_list=np_chains_list, pair_msa_sequences=pair_msa_sequences,
+ max_templates=MAX_TEMPLATES
+ )
+ np_example = process_final(np_example)
+ return np_example
+
+
+def crop_chains(
+ chains_list: List[Mapping[str, np.ndarray]],
+ msa_crop_size: int,
+ pair_msa_sequences: bool,
+ max_templates: int
+) -> List[Mapping[str, np.ndarray]]:
+ """Crops the MSAs for a set of chains.
+
+ Args:
+ chains_list: A list of chains to be cropped.
+ msa_crop_size: The total number of sequences to crop from the MSA.
+ pair_msa_sequences: Whether we are operating in sequence-pairing mode.
+ max_templates: The maximum templates to use per chain.
+
+ Returns:
+ The chains cropped.
+ """
+
+ # Apply the cropping.
+ cropped_chains = []
+ for chain in chains_list:
+ cropped_chain = _crop_single_chain(
+ chain,
+ msa_crop_size=msa_crop_size,
+ pair_msa_sequences=pair_msa_sequences,
+ max_templates=max_templates)
+ cropped_chains.append(cropped_chain)
+
+ return cropped_chains
+
+
+def _crop_single_chain(chain: Mapping[str, np.ndarray],
+ msa_crop_size: int,
+ pair_msa_sequences: bool,
+ max_templates: int) -> Mapping[str, np.ndarray]:
+ """Crops msa sequences to `msa_crop_size`."""
+ msa_size = chain['num_alignments']
+
+ if pair_msa_sequences:
+ msa_size_all_seq = chain['num_alignments_all_seq']
+ msa_crop_size_all_seq = np.minimum(msa_size_all_seq, msa_crop_size // 2)
+
+ # We reduce the number of un-paired sequences, by the number of times a
+ # sequence from this chain's MSA is included in the paired MSA. This keeps
+ # the MSA size for each chain roughly constant.
+ msa_all_seq = chain['msa_all_seq'][:msa_crop_size_all_seq, :]
+ num_non_gapped_pairs = np.sum(
+ np.any(msa_all_seq != msa_pairing.MSA_GAP_IDX, axis=1))
+ num_non_gapped_pairs = np.minimum(num_non_gapped_pairs,
+ msa_crop_size_all_seq)
+
+ # Restrict the unpaired crop size so that paired+unpaired sequences do not
+ # exceed msa_seqs_per_chain for each chain.
+ max_msa_crop_size = np.maximum(msa_crop_size - num_non_gapped_pairs, 0)
+ msa_crop_size = np.minimum(msa_size, max_msa_crop_size)
+ else:
+ msa_crop_size = np.minimum(msa_size, msa_crop_size)
+
+ include_templates = 'template_aatype' in chain and max_templates
+ if include_templates:
+ num_templates = chain['template_aatype'].shape[0]
+ templates_crop_size = np.minimum(num_templates, max_templates)
+
+ for k in chain:
+ k_split = k.split('_all_seq')[0]
+ if k_split in msa_pairing.TEMPLATE_FEATURES:
+ chain[k] = chain[k][:templates_crop_size, :]
+ elif k_split in msa_pairing.MSA_FEATURES:
+ if '_all_seq' in k and pair_msa_sequences:
+ chain[k] = chain[k][:msa_crop_size_all_seq, :]
+ else:
+ chain[k] = chain[k][:msa_crop_size, :]
+
+ chain['num_alignments'] = np.asarray(msa_crop_size, dtype=np.int32)
+ if include_templates:
+ chain['num_templates'] = np.asarray(templates_crop_size, dtype=np.int32)
+ if pair_msa_sequences:
+ chain['num_alignments_all_seq'] = np.asarray(
+ msa_crop_size_all_seq, dtype=np.int32)
+ return chain
+
+
+def process_final(
+ np_example: Mapping[str, np.ndarray]
+) -> Mapping[str, np.ndarray]:
+ """Final processing steps in data pipeline, after merging and pairing."""
+ np_example = _correct_msa_restypes(np_example)
+ np_example = _make_seq_mask(np_example)
+ np_example = _make_msa_mask(np_example)
+ np_example = _filter_features(np_example)
+
+ return np_example
+
+
+def _correct_msa_restypes(np_example):
+ """Correct MSA restype to have the same order as residue_constants."""
+ new_order_list = residue_constants.MAP_HHBLITS_AATYPE_TO_OUR_AATYPE
+ np_example['msa'] = np.take(new_order_list, np_example['msa'], axis=0)
+ np_example['msa'] = np_example['msa'].astype(np.int32)
+ return np_example
+
+
+def _make_seq_mask(np_example):
+ np_example['seq_mask'] = (np_example['entity_id'] > 0).astype(np.float32)
+ return np_example
+
+
+def _make_msa_mask(np_example):
+ """Mask features are all ones, but will later be zero-padded."""
+
+ np_example['msa_mask'] = np.ones_like(np_example['msa'], dtype=np.float32)
+
+ seq_mask = (np_example['entity_id'] > 0).astype(np.float32)
+ np_example['msa_mask'] *= seq_mask[None]
+
+ return np_example
+
+
+def _filter_features(
+ np_example: Mapping[str, np.ndarray]
+) -> Mapping[str, np.ndarray]:
+ """Filters features of example to only those requested."""
+ return {k: v for (k, v) in np_example.items() if k in REQUIRED_FEATURES}
+
+
+def process_unmerged_features(
+ all_chain_features: MutableMapping[str, Mapping[str, np.ndarray]]
+):
+ """Postprocessing stage for per-chain features before merging."""
+ num_chains = len(all_chain_features)
+ for chain_features in all_chain_features.values():
+ # Convert deletion matrices to float.
+ chain_features['deletion_matrix'] = np.asarray(
+ chain_features.pop('deletion_matrix_int'), dtype=np.float32
+ )
+ if 'deletion_matrix_int_all_seq' in chain_features:
+ chain_features['deletion_matrix_all_seq'] = np.asarray(
+ chain_features.pop('deletion_matrix_int_all_seq'), dtype=np.float32
+ )
+
+ chain_features['deletion_mean'] = np.mean(
+ chain_features['deletion_matrix'], axis=0
+ )
+
+ if 'all_atom_positions' not in chain_features:
+ # Add all_atom_mask and dummy all_atom_positions based on aatype.
+ all_atom_mask = residue_constants.STANDARD_ATOM_MASK[
+ chain_features['aatype']]
+ chain_features['all_atom_mask'] = all_atom_mask.astype(dtype=np.float32)
+ chain_features['all_atom_positions'] = np.zeros(
+ list(all_atom_mask.shape) + [3])
+
+ # Add assembly_num_chains.
+ chain_features['assembly_num_chains'] = np.asarray(num_chains)
+
+ # Add entity_mask.
+ for chain_features in all_chain_features.values():
+ chain_features['entity_mask'] = (
+ chain_features['entity_id'] != 0).astype(np.int32)
diff --git a/openfold/data/input_pipeline.py b/openfold/data/input_pipeline.py
index dee7aa0b..5dc3df6c 100644
--- a/openfold/data/input_pipeline.py
+++ b/openfold/data/input_pipeline.py
@@ -107,7 +107,8 @@ def ensembled_transform_fns(common_cfg, mode_cfg, ensemble_seed):
# the masked locations and secret corrupted locations.
transforms.append(
data_transforms.make_masked_msa(
- common_cfg.masked_msa, mode_cfg.masked_msa_replace_fraction
+ common_cfg.masked_msa, mode_cfg.masked_msa_replace_fraction,
+ seed=(msa_seed + 1) if msa_seed else None,
)
)
diff --git a/openfold/data/input_pipeline_multimer.py b/openfold/data/input_pipeline_multimer.py
new file mode 100644
index 00000000..835b4aad
--- /dev/null
+++ b/openfold/data/input_pipeline_multimer.py
@@ -0,0 +1,180 @@
+# Copyright 2021 AlQuraishi Laboratory
+# 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.
+
+import random
+import torch
+
+from openfold.data import (
+ data_transforms,
+ data_transforms_multimer,
+)
+
+
+def groundtruth_transforms_fns():
+ transforms = [data_transforms.make_atom14_masks,
+ data_transforms.make_atom14_positions,
+ data_transforms.atom37_to_frames,
+ data_transforms.atom37_to_torsion_angles(""),
+ data_transforms.make_pseudo_beta(""),
+ data_transforms.get_backbone_frames,
+ data_transforms.get_chi_angles]
+ return transforms
+
+
+def nonensembled_transform_fns():
+ """Input pipeline data transformers that are not ensembled."""
+ transforms = [
+ data_transforms.cast_to_64bit_ints,
+ data_transforms_multimer.make_msa_profile,
+ data_transforms_multimer.create_target_feat,
+ data_transforms.make_atom14_masks
+ ]
+
+ return transforms
+
+
+def ensembled_transform_fns(common_cfg, mode_cfg, ensemble_seed):
+ """Input pipeline data transformers that can be ensembled and averaged."""
+ transforms = []
+
+ pad_msa_clusters = mode_cfg.max_msa_clusters
+ max_msa_clusters = pad_msa_clusters
+ max_extra_msa = mode_cfg.max_extra_msa
+
+ msa_seed = None
+ if(not common_cfg.resample_msa_in_recycling):
+ msa_seed = ensemble_seed
+
+ transforms.append(
+ data_transforms_multimer.sample_msa(
+ max_msa_clusters,
+ max_extra_msa,
+ seed=msa_seed,
+ )
+ )
+
+ if "masked_msa" in common_cfg:
+ # Masked MSA should come *before* MSA clustering so that
+ # the clustering and full MSA profile do not leak information about
+ # the masked locations and secret corrupted locations.
+ transforms.append(
+ data_transforms_multimer.make_masked_msa(
+ common_cfg.masked_msa,
+ mode_cfg.masked_msa_replace_fraction,
+ seed=(msa_seed + 1) if msa_seed else None,
+ )
+ )
+
+ transforms.append(data_transforms_multimer.nearest_neighbor_clusters())
+ transforms.append(data_transforms_multimer.create_msa_feat)
+
+ crop_feats = dict(common_cfg.feat)
+
+ if mode_cfg.fixed_size:
+ transforms.append(data_transforms.select_feat(list(crop_feats)))
+
+ if mode_cfg.crop:
+ transforms.append(
+ data_transforms_multimer.random_crop_to_size(
+ crop_size=mode_cfg.crop_size,
+ max_templates=mode_cfg.max_templates,
+ shape_schema=crop_feats,
+ spatial_crop_prob=mode_cfg.spatial_crop_prob,
+ interface_threshold=mode_cfg.interface_threshold,
+ subsample_templates=mode_cfg.subsample_templates,
+ seed=ensemble_seed + 1,
+ )
+ )
+ transforms.append(
+ data_transforms.make_fixed_size(
+ shape_schema=crop_feats,
+ msa_cluster_size=pad_msa_clusters,
+ extra_msa_size=mode_cfg.max_extra_msa,
+ num_res=mode_cfg.crop_size,
+ num_templates=mode_cfg.max_templates,
+ )
+ )
+ else:
+ transforms.append(
+ data_transforms.crop_templates(mode_cfg.max_templates)
+ )
+
+ return transforms
+
+
+def prepare_ground_truth_features(tensors):
+ """Prepare ground truth features that are only needed for loss calculation during training"""
+
+ gt_features = ['all_atom_mask', 'all_atom_positions', 'asym_id', 'sym_id', 'entity_id']
+ gt_tensors = {k: v for k, v in tensors.items() if k in gt_features}
+ gt_tensors['aatype'] = tensors['aatype'].to(torch.long)
+ gt_tensors = compose(groundtruth_transforms_fns())(gt_tensors)
+ return gt_tensors
+
+
+def process_tensors_from_config(tensors, common_cfg, mode_cfg):
+ """Based on the config, apply filters and transformations to the data."""
+
+ process_gt_feats = mode_cfg.supervised
+ gt_tensors = {}
+ if process_gt_feats:
+ gt_tensors = prepare_ground_truth_features(tensors)
+
+ ensemble_seed = random.randint(0, torch.iinfo(torch.int32).max)
+ tensors['aatype'] = tensors['aatype'].to(torch.long)
+ nonensembled = nonensembled_transform_fns()
+ tensors = compose(nonensembled)(tensors)
+ if("no_recycling_iters" in tensors):
+ num_recycling = int(tensors["no_recycling_iters"])
+ else:
+ num_recycling = common_cfg.max_recycling_iters
+
+ def wrap_ensemble_fn(data, i):
+ """Function to be mapped over the ensemble dimension."""
+ d = data.copy()
+ fns = ensembled_transform_fns(
+ common_cfg,
+ mode_cfg,
+ ensemble_seed,
+ )
+ fn = compose(fns)
+ d["ensemble_index"] = i
+ return fn(d)
+
+ tensors = map_fn(
+ lambda x: wrap_ensemble_fn(tensors, x), torch.arange(num_recycling + 1)
+ )
+
+ if process_gt_feats:
+ tensors['gt_features'] = gt_tensors
+
+ return tensors
+
+@data_transforms.curry1
+def compose(x, fs):
+ for f in fs:
+ x = f(x)
+ return x
+
+
+def map_fn(fun, x):
+ ensembles = [fun(elem) for elem in x]
+ features = ensembles[0].keys()
+ ensembled_dict = {}
+ for feat in features:
+ ensembled_dict[feat] = torch.stack(
+ [dict_i[feat] for dict_i in ensembles], dim=-1
+ )
+ return ensembled_dict
diff --git a/openfold/data/mmcif_parsing.py b/openfold/data/mmcif_parsing.py
index 0212f2fd..1bc5ef86 100644
--- a/openfold/data/mmcif_parsing.py
+++ b/openfold/data/mmcif_parsing.py
@@ -16,6 +16,7 @@
"""Parses the mmCIF file format."""
import collections
import dataclasses
+import functools
import io
import json
import logging
@@ -173,6 +174,7 @@ def mmcif_loop_to_dict(
return {entry[index]: entry for entry in entries}
+@functools.lru_cache(16, typed=False)
def parse(
*, file_id: str, mmcif_string: str, catch_all_errors: bool = True
) -> ParsingResult:
@@ -346,7 +348,7 @@ def _get_header(parsed_info: MmCIFDict) -> PdbHeader:
raw_resolution = parsed_info[res_key][0]
header["resolution"] = float(raw_resolution)
except ValueError:
- logging.info(
+ logging.debug(
"Invalid resolution format: %s", parsed_info[res_key]
)
@@ -474,6 +476,20 @@ def get_atom_coords(
pos[residue_constants.atom_order["SD"]] = [x, y, z]
mask[residue_constants.atom_order["SD"]] = 1.0
+ # Fix naming errors in arginine residues where NH2 is incorrectly
+ # assigned to be closer to CD than NH1
+ cd = residue_constants.atom_order['CD']
+ nh1 = residue_constants.atom_order['NH1']
+ nh2 = residue_constants.atom_order['NH2']
+ if(
+ res.get_resname() == 'ARG' and
+ all(mask[atom_index] for atom_index in (cd, nh1, nh2)) and
+ (np.linalg.norm(pos[nh1] - pos[cd]) >
+ np.linalg.norm(pos[nh2] - pos[cd]))
+ ):
+ pos[nh1], pos[nh2] = pos[nh2].copy(), pos[nh1].copy()
+ mask[nh1], mask[nh2] = mask[nh2].copy(), mask[nh1].copy()
+
all_atom_positions[res_index] = pos
all_atom_mask[res_index] = mask
diff --git a/openfold/data/msa_identifiers.py b/openfold/data/msa_identifiers.py
new file mode 100644
index 00000000..07dbfa47
--- /dev/null
+++ b/openfold/data/msa_identifiers.py
@@ -0,0 +1,91 @@
+# 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.
+
+"""Utilities for extracting identifiers from MSA sequence descriptions."""
+
+import dataclasses
+import re
+from typing import Optional
+
+
+# Sequences coming from UniProtKB database come in the
+# `db|UniqueIdentifier|EntryName` format, e.g. `tr|A0A146SKV9|A0A146SKV9_FUNHE`
+# or `sp|P0C2L1|A3X1_LOXLA` (for TREMBL/Swiss-Prot respectively).
+_UNIPROT_PATTERN = re.compile(
+ r"""
+ ^
+ # UniProtKB/TrEMBL or UniProtKB/Swiss-Prot
+ (?:tr|sp)
+ \|
+ # A primary accession number of the UniProtKB entry.
+ (?P[A-Za-z0-9]{6,10})
+ # Occasionally there is a _0 or _1 isoform suffix, which we ignore.
+ (?:_\d)?
+ \|
+ # TREMBL repeats the accession ID here. Swiss-Prot has a mnemonic
+ # protein ID code.
+ (?:[A-Za-z0-9]+)
+ _
+ # A mnemonic species identification code.
+ (?P([A-Za-z0-9]){1,5})
+ # Small BFD uses a final value after an underscore, which we ignore.
+ (?:_\d+)?
+ $
+ """,
+ re.VERBOSE)
+
+
+@dataclasses.dataclass(frozen=True)
+class Identifiers:
+ species_id: str = ''
+
+
+def _parse_sequence_identifier(msa_sequence_identifier: str) -> Identifiers:
+ """Gets accession id and species from an msa sequence identifier.
+
+ The sequence identifier has the format specified by
+ _UNIPROT_TREMBL_ENTRY_NAME_PATTERN or _UNIPROT_SWISSPROT_ENTRY_NAME_PATTERN.
+ An example of a sequence identifier: `tr|A0A146SKV9|A0A146SKV9_FUNHE`
+
+ Args:
+ msa_sequence_identifier: a sequence identifier.
+
+ Returns:
+ An `Identifiers` instance with a uniprot_accession_id and species_id. These
+ can be empty in the case where no identifier was found.
+ """
+ matches = re.search(_UNIPROT_PATTERN, msa_sequence_identifier.strip())
+ if matches:
+ return Identifiers(
+ species_id=matches.group('SpeciesIdentifier')
+ )
+ return Identifiers()
+
+
+def _extract_sequence_identifier(description: str) -> Optional[str]:
+ """Extracts sequence identifier from description. Returns None if no match."""
+ split_description = description.split()
+ if split_description:
+ return split_description[0].partition('/')[0]
+ else:
+ return None
+
+
+def get_identifiers(description: str) -> Identifiers:
+ """Computes extra MSA features from the description."""
+ sequence_identifier = _extract_sequence_identifier(description)
+ if sequence_identifier is None:
+ return Identifiers()
+ else:
+ return _parse_sequence_identifier(sequence_identifier)
diff --git a/openfold/data/msa_pairing.py b/openfold/data/msa_pairing.py
new file mode 100644
index 00000000..9b82acaf
--- /dev/null
+++ b/openfold/data/msa_pairing.py
@@ -0,0 +1,482 @@
+# 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.
+
+"""Pairing logic for multimer data pipeline."""
+
+import collections
+import functools
+import string
+from typing import Any, Dict, Iterable, List, Sequence, Mapping
+
+import numpy as np
+import pandas as pd
+import scipy.linalg
+
+from openfold.np import residue_constants
+
+
+# TODO: This stuff should probably also be in a config
+MSA_GAP_IDX = residue_constants.restypes_with_x_and_gap.index('-')
+SEQUENCE_GAP_CUTOFF = 0.5
+SEQUENCE_SIMILARITY_CUTOFF = 0.9
+
+MSA_PAD_VALUES = {'msa_all_seq': MSA_GAP_IDX,
+ 'msa_mask_all_seq': 1,
+ 'deletion_matrix_all_seq': 0,
+ 'deletion_matrix_int_all_seq': 0,
+ 'msa': MSA_GAP_IDX,
+ 'msa_mask': 1,
+ 'deletion_matrix': 0,
+ 'deletion_matrix_int': 0}
+
+MSA_FEATURES = ('msa', 'msa_mask', 'deletion_matrix', 'deletion_matrix_int')
+SEQ_FEATURES = ('residue_index', 'aatype', 'all_atom_positions',
+ 'all_atom_mask', 'seq_mask', 'between_segment_residues',
+ 'has_alt_locations', 'has_hetatoms', 'asym_id', 'entity_id',
+ 'sym_id', 'entity_mask', 'deletion_mean',
+ 'prediction_atom_mask',
+ 'literature_positions', 'atom_indices_to_group_indices',
+ 'rigid_group_default_frame')
+TEMPLATE_FEATURES = ('template_aatype', 'template_all_atom_positions',
+ 'template_all_atom_mask')
+CHAIN_FEATURES = ('num_alignments', 'seq_length')
+
+
+def create_paired_features(
+ chains: Iterable[Mapping[str, np.ndarray]],
+) -> List[Mapping[str, np.ndarray]]:
+ """Returns the original chains with paired NUM_SEQ features.
+
+ Args:
+ chains: A list of feature dictionaries for each chain.
+
+ Returns:
+ A list of feature dictionaries with sequence features including only
+ rows to be paired.
+ """
+ chains = list(chains)
+ chain_keys = chains[0].keys()
+
+ if len(chains) < 2:
+ return chains
+ else:
+ updated_chains = []
+ paired_chains_to_paired_row_indices = pair_sequences(chains)
+ paired_rows = reorder_paired_rows(
+ paired_chains_to_paired_row_indices)
+
+ for chain_num, chain in enumerate(chains):
+ new_chain = {k: v for k, v in chain.items() if '_all_seq' not in k}
+ for feature_name in chain_keys:
+ if feature_name.endswith('_all_seq'):
+ feats_padded = pad_features(chain[feature_name], feature_name)
+ new_chain[feature_name] = feats_padded[paired_rows[:, chain_num]]
+ new_chain['num_alignments_all_seq'] = np.asarray(
+ len(paired_rows[:, chain_num]))
+ updated_chains.append(new_chain)
+ return updated_chains
+
+
+def pad_features(feature: np.ndarray, feature_name: str) -> np.ndarray:
+ """Add a 'padding' row at the end of the features list.
+
+ The padding row will be selected as a 'paired' row in the case of partial
+ alignment - for the chain that doesn't have paired alignment.
+
+ Args:
+ feature: The feature to be padded.
+ feature_name: The name of the feature to be padded.
+
+ Returns:
+ The feature with an additional padding row.
+ """
+ assert feature.dtype != np.dtype(np.string_)
+ if feature_name in ('msa_all_seq', 'msa_mask_all_seq',
+ 'deletion_matrix_all_seq', 'deletion_matrix_int_all_seq'):
+ num_res = feature.shape[1]
+ padding = MSA_PAD_VALUES[feature_name] * np.ones([1, num_res],
+ feature.dtype)
+ elif feature_name == 'msa_species_identifiers_all_seq':
+ padding = [b'']
+ else:
+ return feature
+ feats_padded = np.concatenate([feature, padding], axis=0)
+ return feats_padded
+
+
+def _make_msa_df(chain_features: Mapping[str, np.ndarray]) -> pd.DataFrame:
+ """Makes dataframe with msa features needed for msa pairing."""
+ chain_msa = chain_features['msa_all_seq']
+ query_seq = chain_msa[0]
+ per_seq_similarity = np.sum(
+ query_seq[None] == chain_msa, axis=-1) / float(len(query_seq))
+ per_seq_gap = np.sum(chain_msa == 21, axis=-1) / float(len(query_seq))
+ msa_df = pd.DataFrame({
+ 'msa_species_identifiers':
+ chain_features['msa_species_identifiers_all_seq'],
+ 'msa_row':
+ np.arange(len(
+ chain_features['msa_species_identifiers_all_seq'])),
+ 'msa_similarity': per_seq_similarity,
+ 'gap': per_seq_gap
+ })
+ return msa_df
+
+
+def _create_species_dict(msa_df: pd.DataFrame) -> Dict[bytes, pd.DataFrame]:
+ """Creates mapping from species to msa dataframe of that species."""
+ species_lookup = {}
+ for species, species_df in msa_df.groupby('msa_species_identifiers'):
+ species_lookup[species] = species_df
+ return species_lookup
+
+
+def _match_rows_by_sequence_similarity(this_species_msa_dfs: List[pd.DataFrame]
+ ) -> List[List[int]]:
+ """Finds MSA sequence pairings across chains based on sequence similarity.
+
+ Each chain's MSA sequences are first sorted by their sequence similarity to
+ their respective target sequence. The sequences are then paired, starting
+ from the sequences most similar to their target sequence.
+
+ Args:
+ this_species_msa_dfs: a list of dataframes containing MSA features for
+ sequences for a specific species.
+
+ Returns:
+ A list of lists, each containing M indices corresponding to paired MSA rows,
+ where M is the number of chains.
+ """
+ all_paired_msa_rows = []
+
+ num_seqs = [len(species_df) for species_df in this_species_msa_dfs
+ if species_df is not None]
+ take_num_seqs = np.min(num_seqs)
+
+ sort_by_similarity = (
+ lambda x: x.sort_values('msa_similarity', axis=0, ascending=False))
+
+ for species_df in this_species_msa_dfs:
+ if species_df is not None:
+ species_df_sorted = sort_by_similarity(species_df)
+ msa_rows = species_df_sorted.msa_row.iloc[:take_num_seqs].values
+ else:
+ msa_rows = [-1] * take_num_seqs # take the last 'padding' row
+ all_paired_msa_rows.append(msa_rows)
+ all_paired_msa_rows = list(np.array(all_paired_msa_rows).transpose())
+ return all_paired_msa_rows
+
+
+def pair_sequences(
+ examples: List[Mapping[str, np.ndarray]],
+) -> Dict[int, np.ndarray]:
+ """Returns indices for paired MSA sequences across chains."""
+
+ num_examples = len(examples)
+
+ all_chain_species_dict = []
+ common_species = set()
+ for chain_features in examples:
+ msa_df = _make_msa_df(chain_features)
+ species_dict = _create_species_dict(msa_df)
+ all_chain_species_dict.append(species_dict)
+ common_species.update(set(species_dict))
+
+ common_species = sorted(common_species)
+ common_species.remove(b'') # Remove target sequence species.
+
+ all_paired_msa_rows = [np.zeros(len(examples), int)]
+ all_paired_msa_rows_dict = {k: [] for k in range(num_examples)}
+ all_paired_msa_rows_dict[num_examples] = [np.zeros(len(examples), int)]
+
+ for species in common_species:
+ if not species:
+ continue
+ this_species_msa_dfs = []
+ species_dfs_present = 0
+ for species_dict in all_chain_species_dict:
+ if species in species_dict:
+ this_species_msa_dfs.append(species_dict[species])
+ species_dfs_present += 1
+ else:
+ this_species_msa_dfs.append(None)
+
+ # Skip species that are present in only one chain.
+ if species_dfs_present <= 1:
+ continue
+
+ if np.any(
+ np.array([len(species_df) for species_df in
+ this_species_msa_dfs if
+ isinstance(species_df, pd.DataFrame)]) > 600):
+ continue
+
+ paired_msa_rows = _match_rows_by_sequence_similarity(this_species_msa_dfs)
+ all_paired_msa_rows.extend(paired_msa_rows)
+ all_paired_msa_rows_dict[species_dfs_present].extend(paired_msa_rows)
+ all_paired_msa_rows_dict = {
+ num_examples: np.array(paired_msa_rows) for
+ num_examples, paired_msa_rows in all_paired_msa_rows_dict.items()
+ }
+ return all_paired_msa_rows_dict
+
+
+def reorder_paired_rows(all_paired_msa_rows_dict: Dict[int, np.ndarray]
+ ) -> np.ndarray:
+ """Creates a list of indices of paired MSA rows across chains.
+
+ Args:
+ all_paired_msa_rows_dict: a mapping from the number of paired chains to the
+ paired indices.
+
+ Returns:
+ a list of lists, each containing indices of paired MSA rows across chains.
+ The paired-index lists are ordered by:
+ 1) the number of chains in the paired alignment, i.e, all-chain pairings
+ will come first.
+ 2) e-values
+ """
+ all_paired_msa_rows = []
+
+ for num_pairings in sorted(all_paired_msa_rows_dict, reverse=True):
+ paired_rows = all_paired_msa_rows_dict[num_pairings]
+ paired_rows_product = abs(np.array([np.prod(rows) for rows in paired_rows]))
+ paired_rows_sort_index = np.argsort(paired_rows_product)
+ all_paired_msa_rows.extend(paired_rows[paired_rows_sort_index])
+
+ return np.array(all_paired_msa_rows)
+
+
+def block_diag(*arrs: np.ndarray, pad_value: float = 0.0) -> np.ndarray:
+ """Like scipy.linalg.block_diag but with an optional padding value."""
+ ones_arrs = [np.ones_like(x) for x in arrs]
+ off_diag_mask = 1.0 - scipy.linalg.block_diag(*ones_arrs)
+ diag = scipy.linalg.block_diag(*arrs)
+ diag += (off_diag_mask * pad_value).astype(diag.dtype)
+ return diag
+
+
+def _correct_post_merged_feats(
+ np_example: Mapping[str, np.ndarray],
+ np_chains_list: Sequence[Mapping[str, np.ndarray]],
+ pair_msa_sequences: bool
+) -> Mapping[str, np.ndarray]:
+ """Adds features that need to be computed/recomputed post merging."""
+
+ np_example['seq_length'] = np.asarray(
+ np_example['aatype'].shape[0],
+ dtype=np.int32
+ )
+ np_example['num_alignments'] = np.asarray(
+ np_example['msa'].shape[0],
+ dtype=np.int32
+ )
+
+ if not pair_msa_sequences:
+ # Generate a bias that is 1 for the first row of every block in the
+ # block diagonal MSA - i.e. make sure the cluster stack always includes
+ # the query sequences for each chain (since the first row is the query
+ # sequence).
+ cluster_bias_masks = []
+ for chain in np_chains_list:
+ mask = np.zeros(chain['msa'].shape[0])
+ mask[0] = 1
+ cluster_bias_masks.append(mask)
+
+ np_example['cluster_bias_mask'] = np.concatenate(cluster_bias_masks)
+
+ # Initialize Bert mask with masked out off diagonals.
+ msa_masks = [
+ np.ones(x['msa'].shape, dtype=np.float32)
+ for x in np_chains_list
+ ]
+
+ np_example['bert_mask'] = block_diag(
+ *msa_masks, pad_value=0
+ )
+ else:
+ np_example['cluster_bias_mask'] = np.zeros(np_example['msa'].shape[0])
+ np_example['cluster_bias_mask'][0] = 1
+
+ # Initialize Bert mask with masked out off diagonals.
+ msa_masks = [
+ np.ones(x['msa'].shape, dtype=np.float32) for
+ x in np_chains_list
+ ]
+ msa_masks_all_seq = [
+ np.ones(x['msa_all_seq'].shape, dtype=np.float32) for
+ x in np_chains_list
+ ]
+
+ msa_mask_block_diag = block_diag(
+ *msa_masks, pad_value=0
+ )
+ msa_mask_all_seq = np.concatenate(msa_masks_all_seq, axis=1)
+ np_example['bert_mask'] = np.concatenate(
+ [msa_mask_all_seq, msa_mask_block_diag],
+ axis=0
+ )
+
+ return np_example
+
+
+def _pad_templates(chains: Sequence[Mapping[str, np.ndarray]],
+ max_templates: int) -> Sequence[Mapping[str, np.ndarray]]:
+ """For each chain pad the number of templates to a fixed size.
+
+ Args:
+ chains: A list of protein chains.
+ max_templates: Each chain will be padded to have this many templates.
+
+ Returns:
+ The list of chains, updated to have template features padded to
+ max_templates.
+ """
+ for chain in chains:
+ for k, v in chain.items():
+ if k in TEMPLATE_FEATURES:
+ padding = np.zeros_like(v.shape)
+ padding[0] = max_templates - v.shape[0]
+ padding = [(0, p) for p in padding]
+ chain[k] = np.pad(v, padding, mode='constant')
+ return chains
+
+
+def _merge_features_from_multiple_chains(
+ chains: Sequence[Mapping[str, np.ndarray]],
+ pair_msa_sequences: bool) -> Mapping[str, np.ndarray]:
+ """Merge features from multiple chains.
+
+ Args:
+ chains: A list of feature dictionaries that we want to merge.
+ pair_msa_sequences: Whether to concatenate MSA features along the
+ num_res dimension (if True), or to block diagonalize them (if False).
+
+ Returns:
+ A feature dictionary for the merged example.
+ """
+ merged_example = {}
+ for feature_name in chains[0]:
+ feats = [x[feature_name] for x in chains]
+ feature_name_split = feature_name.split('_all_seq')[0]
+ if feature_name_split in MSA_FEATURES:
+ if pair_msa_sequences or '_all_seq' in feature_name:
+ merged_example[feature_name] = np.concatenate(feats, axis=1)
+ else:
+ merged_example[feature_name] = block_diag(
+ *feats, pad_value=MSA_PAD_VALUES[feature_name])
+ elif feature_name_split in SEQ_FEATURES:
+ merged_example[feature_name] = np.concatenate(feats, axis=0)
+ elif feature_name_split in TEMPLATE_FEATURES:
+ merged_example[feature_name] = np.concatenate(feats, axis=1)
+ elif feature_name_split in CHAIN_FEATURES:
+ merged_example[feature_name] = np.sum(x for x in feats).astype(np.int32)
+ else:
+ merged_example[feature_name] = feats[0]
+ return merged_example
+
+
+def _merge_homomers_dense_msa(
+ chains: Iterable[Mapping[str, np.ndarray]]) -> Sequence[Mapping[str, np.ndarray]]:
+ """Merge all identical chains, making the resulting MSA dense.
+
+ Args:
+ chains: An iterable of features for each chain.
+
+ Returns:
+ A list of feature dictionaries. All features with the same entity_id
+ will be merged - MSA features will be concatenated along the num_res
+ dimension - making them dense.
+ """
+ entity_chains = collections.defaultdict(list)
+ for chain in chains:
+ entity_id = chain['entity_id'][0]
+ entity_chains[entity_id].append(chain)
+
+ grouped_chains = []
+ for entity_id in sorted(entity_chains):
+ chains = entity_chains[entity_id]
+ grouped_chains.append(chains)
+ chains = [
+ _merge_features_from_multiple_chains(chains, pair_msa_sequences=True)
+ for chains in grouped_chains]
+ return chains
+
+
+def _concatenate_paired_and_unpaired_features(
+ example: Mapping[str, np.ndarray]) -> Mapping[str, np.ndarray]:
+ """Merges paired and block-diagonalised features."""
+ features = MSA_FEATURES
+ for feature_name in features:
+ if feature_name in example:
+ feat = example[feature_name]
+ feat_all_seq = example[feature_name + '_all_seq']
+ merged_feat = np.concatenate([feat_all_seq, feat], axis=0)
+ example[feature_name] = merged_feat
+ example['num_alignments'] = np.array(example['msa'].shape[0],
+ dtype=np.int32)
+ return example
+
+
+def merge_chain_features(np_chains_list: List[Mapping[str, np.ndarray]],
+ pair_msa_sequences: bool,
+ max_templates: int) -> Mapping[str, np.ndarray]:
+ """Merges features for multiple chains to single FeatureDict.
+
+ Args:
+ np_chains_list: List of FeatureDicts for each chain.
+ pair_msa_sequences: Whether to merge paired MSAs.
+ max_templates: The maximum number of templates to include.
+
+ Returns:
+ Single FeatureDict for entire complex.
+ """
+ np_chains_list = _pad_templates(
+ np_chains_list, max_templates=max_templates)
+ np_chains_list = _merge_homomers_dense_msa(np_chains_list)
+ # Unpaired MSA features will be always block-diagonalised; paired MSA
+ # features will be concatenated.
+ np_example = _merge_features_from_multiple_chains(
+ np_chains_list, pair_msa_sequences=False)
+ if pair_msa_sequences:
+ np_example = _concatenate_paired_and_unpaired_features(np_example)
+ np_example = _correct_post_merged_feats(
+ np_example=np_example,
+ np_chains_list=np_chains_list,
+ pair_msa_sequences=pair_msa_sequences)
+
+ return np_example
+
+
+def deduplicate_unpaired_sequences(
+ np_chains: List[Mapping[str, np.ndarray]]) -> List[Mapping[str, np.ndarray]]:
+ """Removes unpaired sequences which duplicate a paired sequence."""
+
+ feature_names = np_chains[0].keys()
+ msa_features = MSA_FEATURES
+
+ for chain in np_chains:
+ # Convert the msa_all_seq numpy array to a tuple for hashing.
+ sequence_set = set(tuple(s) for s in chain['msa_all_seq'])
+ keep_rows = []
+ # Go through unpaired MSA seqs and remove any rows that correspond to the
+ # sequences that are already present in the paired MSA.
+ for row_num, seq in enumerate(chain['msa']):
+ if tuple(seq) not in sequence_set:
+ keep_rows.append(row_num)
+ for feature_name in feature_names:
+ if feature_name in msa_features:
+ chain[feature_name] = chain[feature_name][keep_rows]
+ chain['num_alignments'] = np.array(chain['msa'].shape[0], dtype=np.int32)
+ return np_chains
diff --git a/openfold/data/parsers.py b/openfold/data/parsers.py
index b3ccc53f..8def128a 100644
--- a/openfold/data/parsers.py
+++ b/openfold/data/parsers.py
@@ -16,14 +16,43 @@
"""Functions for parsing various file formats."""
import collections
import dataclasses
+import itertools
import re
import string
-from typing import Dict, Iterable, List, Optional, Sequence, Tuple
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple, Set
DeletionMatrix = Sequence[Sequence[int]]
+@dataclasses.dataclass(frozen=True)
+class Msa:
+ """Class representing a parsed MSA file"""
+ sequences: Sequence[str]
+ deletion_matrix: DeletionMatrix
+ descriptions: Optional[Sequence[str]]
+
+ def __post_init__(self):
+ if(not (
+ len(self.sequences) ==
+ len(self.deletion_matrix) ==
+ len(self.descriptions)
+ )):
+ raise ValueError(
+ "All fields for an MSA must have the same length"
+ )
+
+ def __len__(self):
+ return len(self.sequences)
+
+ def truncate(self, max_seqs: int):
+ return Msa(
+ sequences=self.sequences[:max_seqs],
+ deletion_matrix=self.deletion_matrix[:max_seqs],
+ descriptions=self.descriptions[:max_seqs],
+ )
+
+
@dataclasses.dataclass(frozen=True)
class TemplateHit:
"""Class representing a template hit."""
@@ -31,7 +60,7 @@ class TemplateHit:
index: int
name: str
aligned_cols: int
- sum_probs: float
+ sum_probs: Optional[float]
query: str
hit_sequence: str
indices_query: List[int]
@@ -69,9 +98,7 @@ def parse_fasta(fasta_string: str) -> Tuple[Sequence[str], Sequence[str]]:
return sequences, descriptions
-def parse_stockholm(
- stockholm_string: str,
-) -> Tuple[Sequence[str], DeletionMatrix, Sequence[str]]:
+def parse_stockholm(stockholm_string: str) -> Msa:
"""Parses sequences and deletion matrix from stockholm format alignment.
Args:
@@ -126,10 +153,14 @@ def parse_stockholm(
deletion_count = 0
deletion_matrix.append(deletion_vec)
- return msa, deletion_matrix, list(name_to_sequence.keys())
+ return Msa(
+ sequences=msa,
+ deletion_matrix=deletion_matrix,
+ descriptions=list(name_to_sequence.keys())
+ )
-def parse_a3m(a3m_string: str) -> Tuple[Sequence[str], DeletionMatrix]:
+def parse_a3m(a3m_string: str) -> Msa:
"""Parses sequences and deletion matrix from a3m format alignment.
Args:
@@ -144,7 +175,7 @@ def parse_a3m(a3m_string: str) -> Tuple[Sequence[str], DeletionMatrix]:
at `deletion_matrix[i][j]` is the number of residues deleted from
the aligned sequence i at residue position j.
"""
- sequences, _ = parse_fasta(a3m_string)
+ sequences, descriptions = parse_fasta(a3m_string)
deletion_matrix = []
for msa_sequence in sequences:
deletion_vec = []
@@ -160,7 +191,11 @@ def parse_a3m(a3m_string: str) -> Tuple[Sequence[str], DeletionMatrix]:
# Make the MSA matrix out of aligned (deletion-free) sequences.
deletion_table = str.maketrans("", "", string.ascii_lowercase)
aligned_sequences = [s.translate(deletion_table) for s in sequences]
- return aligned_sequences, deletion_matrix
+ return Msa(
+ sequences=aligned_sequences,
+ deletion_matrix=deletion_matrix,
+ descriptions=descriptions
+ )
def _convert_sto_seq_to_a3m(
@@ -174,7 +209,9 @@ def _convert_sto_seq_to_a3m(
def convert_stockholm_to_a3m(
- stockholm_format: str, max_sequences: Optional[int] = None
+ stockholm_format: str,
+ max_sequences: Optional[int] = None,
+ remove_first_row_gaps: bool = True,
) -> str:
"""Converts MSA in Stockholm format to the A3M format."""
descriptions = {}
@@ -212,13 +249,19 @@ def convert_stockholm_to_a3m(
# Convert sto format to a3m line by line
a3m_sequences = {}
- # query_sequence is assumed to be the first sequence
- query_sequence = next(iter(sequences.values()))
- query_non_gaps = [res != "-" for res in query_sequence]
+ if(remove_first_row_gaps):
+ # query_sequence is assumed to be the first sequence
+ query_sequence = next(iter(sequences.values()))
+ query_non_gaps = [res != "-" for res in query_sequence]
for seqname, sto_sequence in sequences.items():
- a3m_sequences[seqname] = "".join(
- _convert_sto_seq_to_a3m(query_non_gaps, sto_sequence)
- )
+ # Dots are optional in a3m format and are commonly removed.
+ out_sequence = sto_sequence.replace('.', '')
+ if(remove_first_row_gaps):
+ out_sequence = ''.join(
+ _convert_sto_seq_to_a3m(query_non_gaps, out_sequence)
+ )
+ a3m_sequences[seqname] = out_sequence
+
fasta_chunks = (
f">{k} {descriptions.get(k, '')}\n{a3m_sequences[k]}"
@@ -227,6 +270,124 @@ def convert_stockholm_to_a3m(
return "\n".join(fasta_chunks) + "\n" # Include terminating newline.
+def _keep_line(line: str, seqnames: Set[str]) -> bool:
+ """Function to decide which lines to keep."""
+ if not line.strip():
+ return True
+ if line.strip() == '//': # End tag
+ return True
+ if line.startswith('# STOCKHOLM'): # Start tag
+ return True
+ if line.startswith('#=GC RF'): # Reference Annotation Line
+ return True
+ if line[:4] == '#=GS': # Description lines - keep if sequence in list.
+ _, seqname, _ = line.split(maxsplit=2)
+ return seqname in seqnames
+ elif line.startswith('#'): # Other markup - filter out
+ return False
+ else: # Alignment data - keep if sequence in list.
+ seqname = line.partition(' ')[0]
+ return seqname in seqnames
+
+
+def truncate_stockholm_msa(stockholm_msa_path: str, max_sequences: int) -> str:
+ """Reads + truncates a Stockholm file while preventing excessive RAM usage."""
+ seqnames = set()
+ filtered_lines = []
+
+ with open(stockholm_msa_path) as f:
+ for line in f:
+ if line.strip() and not line.startswith(('#', '//')):
+ # Ignore blank lines, markup and end symbols - remainder are alignment
+ # sequence parts.
+ seqname = line.partition(' ')[0]
+ seqnames.add(seqname)
+ if len(seqnames) >= max_sequences:
+ break
+
+ f.seek(0)
+ for line in f:
+ if _keep_line(line, seqnames):
+ filtered_lines.append(line)
+
+ return ''.join(filtered_lines)
+
+
+def remove_empty_columns_from_stockholm_msa(stockholm_msa: str) -> str:
+ """Removes empty columns (dashes-only) from a Stockholm MSA."""
+ processed_lines = {}
+ unprocessed_lines = {}
+ for i, line in enumerate(stockholm_msa.splitlines()):
+ if line.startswith('#=GC RF'):
+ reference_annotation_i = i
+ reference_annotation_line = line
+ # Reached the end of this chunk of the alignment. Process chunk.
+ _, _, first_alignment = line.rpartition(' ')
+ mask = []
+ for j in range(len(first_alignment)):
+ for _, unprocessed_line in unprocessed_lines.items():
+ prefix, _, alignment = unprocessed_line.rpartition(' ')
+ if alignment[j] != '-':
+ mask.append(True)
+ break
+ else: # Every row contained a hyphen - empty column.
+ mask.append(False)
+ # Add reference annotation for processing with mask.
+ unprocessed_lines[reference_annotation_i] = reference_annotation_line
+
+ if not any(mask): # All columns were empty. Output empty lines for chunk.
+ for line_index in unprocessed_lines:
+ processed_lines[line_index] = ''
+ else:
+ for line_index, unprocessed_line in unprocessed_lines.items():
+ prefix, _, alignment = unprocessed_line.rpartition(' ')
+ masked_alignment = ''.join(itertools.compress(alignment, mask))
+ processed_lines[line_index] = f'{prefix} {masked_alignment}'
+
+ # Clear raw_alignments.
+ unprocessed_lines = {}
+ elif line.strip() and not line.startswith(('#', '//')):
+ unprocessed_lines[i] = line
+ else:
+ processed_lines[i] = line
+ return '\n'.join((processed_lines[i] for i in range(len(processed_lines))))
+
+
+def deduplicate_stockholm_msa(stockholm_msa: str) -> str:
+ """Remove duplicate sequences (ignoring insertions wrt query)."""
+ sequence_dict = collections.defaultdict(str)
+
+ # First we must extract all sequences from the MSA.
+ for line in stockholm_msa.splitlines():
+ # Only consider the alignments - ignore reference annotation, empty lines,
+ # descriptions or markup.
+ if line.strip() and not line.startswith(('#', '//')):
+ line = line.strip()
+ seqname, alignment = line.split()
+ sequence_dict[seqname] += alignment
+
+ seen_sequences = set()
+ seqnames = set()
+ # First alignment is the query.
+ query_align = next(iter(sequence_dict.values()))
+ mask = [c != '-' for c in query_align] # Mask is False for insertions.
+ for seqname, alignment in sequence_dict.items():
+ # Apply mask to remove all insertions from the string.
+ masked_alignment = ''.join(itertools.compress(alignment, mask))
+ if masked_alignment in seen_sequences:
+ continue
+ else:
+ seen_sequences.add(masked_alignment)
+ seqnames.add(seqname)
+
+ filtered_lines = []
+ for line in stockholm_msa.splitlines():
+ if _keep_line(line, seqnames):
+ filtered_lines.append(line)
+
+ return '\n'.join(filtered_lines) + '\n'
+
+
def _get_hhr_line_regex_groups(
regex_pattern: str, line: str
) -> Sequence[Optional[str]]:
@@ -280,7 +441,7 @@ def _parse_hhr_hit(detailed_lines: Sequence[str]) -> TemplateHit:
"Could not parse section: %s. Expected this: \n%s to contain summary."
% (detailed_lines, detailed_lines[2])
)
- (prob_true, e_value, _, aligned_cols, _, _, sum_probs, neff) = [
+ (_, _, _, aligned_cols, _, _, sum_probs, _) = [
float(x) for x in match.groups()
]
@@ -388,3 +549,115 @@ def parse_e_values_from_tblout(tblout: str) -> Dict[str, float]:
target_name = fields[0]
e_values[target_name] = float(e_value)
return e_values
+
+
+def _get_indices(sequence: str, start: int) -> List[int]:
+ """Returns indices for non-gap/insert residues starting at the given index."""
+ indices = []
+ counter = start
+ for symbol in sequence:
+ # Skip gaps but add a placeholder so that the alignment is preserved.
+ if symbol == '-':
+ indices.append(-1)
+ # Skip deleted residues, but increase the counter.
+ elif symbol.islower():
+ counter += 1
+ # Normal aligned residue. Increase the counter and append to indices.
+ else:
+ indices.append(counter)
+ counter += 1
+ return indices
+
+
+@dataclasses.dataclass(frozen=True)
+class HitMetadata:
+ pdb_id: str
+ chain: str
+ start: int
+ end: int
+ length: int
+ text: str
+
+
+def _parse_hmmsearch_description(description: str) -> HitMetadata:
+ """Parses the hmmsearch A3M sequence description line."""
+ # Example 1: >4pqx_A/2-217 [subseq from] mol:protein length:217 Free text
+ # Example 2: >5g3r_A/1-55 [subseq from] mol:protein length:352
+ match = re.match(
+ r'^>?([a-z0-9]+)_(\w+)/([0-9]+)-([0-9]+).*protein length:([0-9]+) *(.*)$',
+ description.strip())
+
+ if not match:
+ raise ValueError(f'Could not parse description: "{description}".')
+
+ return HitMetadata(
+ pdb_id=match[1],
+ chain=match[2],
+ start=int(match[3]),
+ end=int(match[4]),
+ length=int(match[5]),
+ text=match[6]
+ )
+
+
+def parse_hmmsearch_a3m(
+ query_sequence: str,
+ a3m_string: str,
+ skip_first: bool = True
+) -> Sequence[TemplateHit]:
+ """Parses an a3m string produced by hmmsearch.
+
+ Args:
+ query_sequence: The query sequence.
+ a3m_string: The a3m string produced by hmmsearch.
+ skip_first: Whether to skip the first sequence in the a3m string.
+
+ Returns:
+ A sequence of `TemplateHit` results.
+ """
+ # Zip the descriptions and MSAs together, skip the first query sequence.
+ parsed_a3m = list(zip(*parse_fasta(a3m_string)))
+ if skip_first:
+ parsed_a3m = parsed_a3m[1:]
+
+ indices_query = _get_indices(query_sequence, start=0)
+
+ hits = []
+ for i, (hit_sequence, hit_description) in enumerate(parsed_a3m, start=1):
+ if 'mol:protein' not in hit_description:
+ continue # Skip non-protein chains.
+ metadata = _parse_hmmsearch_description(hit_description)
+ # Aligned columns are only the match states.
+ aligned_cols = sum([r.isupper() and r != '-' for r in hit_sequence])
+ indices_hit = _get_indices(hit_sequence, start=metadata.start - 1)
+
+ hit = TemplateHit(
+ index=i,
+ name=f'{metadata.pdb_id}_{metadata.chain}',
+ aligned_cols=aligned_cols,
+ sum_probs=None,
+ query=query_sequence,
+ hit_sequence=hit_sequence.upper(),
+ indices_query=indices_query,
+ indices_hit=indices_hit,
+ )
+ hits.append(hit)
+
+ return hits
+
+
+def parse_hmmsearch_sto(
+ output_string: str,
+ input_sequence: str
+) -> Sequence[TemplateHit]:
+ """Gets parsed template hits from the raw string output by the tool."""
+ a3m_string = convert_stockholm_to_a3m(
+ output_string,
+ remove_first_row_gaps=False
+ )
+ template_hits = parse_hmmsearch_a3m(
+ query_sequence=input_sequence,
+ a3m_string=a3m_string,
+ skip_first=False
+ )
+ return template_hits
diff --git a/openfold/data/templates.py b/openfold/data/templates.py
index a5ae3860..8c55b5f3 100644
--- a/openfold/data/templates.py
+++ b/openfold/data/templates.py
@@ -14,8 +14,10 @@
# limitations under the License.
"""Functions for getting templates and calculating template features."""
+import abc
import dataclasses
import datetime
+import functools
import glob
import json
import logging
@@ -65,10 +67,6 @@ class DateError(PrefilterError):
"""An error indicating that the hit date was after the max allowed date."""
-class PdbIdError(PrefilterError):
- """An error indicating that the hit PDB ID was identical to the query."""
-
-
class AlignRatioError(PrefilterError):
"""An error indicating that the hit align ratio to the query was too small."""
@@ -91,6 +89,24 @@ class LengthError(PrefilterError):
}
+def empty_template_feats(n_res):
+ return {
+ "template_aatype": np.zeros(
+ (0, n_res, len(residue_constants.restypes_with_x_and_gap)),
+ np.float32
+ ),
+ "template_all_atom_mask": np.zeros(
+ (0, n_res, residue_constants.atom_type_num), np.float32
+ ),
+ "template_all_atom_positions": np.zeros(
+ (0, n_res, residue_constants.atom_type_num, 3), np.float32
+ ),
+ "template_domain_names": np.array([''.encode()], dtype=object),
+ "template_sequence": np.array([''.encode()], dtype=object),
+ "template_sum_probs": np.zeros((0, 1), dtype=np.float32),
+ }
+
+
def _get_pdb_id_and_chain(hit: parsers.TemplateHit) -> Tuple[str, str]:
"""Returns PDB id and chain id for an HHSearch Hit."""
# PDB ID: 4 letters. Chain ID: 1+ alphanumeric letters or "." if unknown.
@@ -204,7 +220,6 @@ def _assess_hhsearch_hit(
hit: parsers.TemplateHit,
hit_pdb_code: str,
query_sequence: str,
- query_pdb_code: Optional[str],
release_dates: Mapping[str, datetime.datetime],
release_date_cutoff: datetime.datetime,
max_subsequence_ratio: float = 0.95,
@@ -218,7 +233,6 @@ def _assess_hhsearch_hit(
different from the value in the actual hit since the original pdb might
have become obsolete.
query_sequence: Amino acid sequence of the query.
- query_pdb_code: 4 letter pdb code of the query.
release_dates: Dictionary mapping pdb codes to their structure release
dates.
release_date_cutoff: Max release date that is valid for this query.
@@ -230,7 +244,6 @@ def _assess_hhsearch_hit(
Raises:
DateError: If the hit date was after the max allowed date.
- PdbIdError: If the hit PDB ID was identical to the query.
AlignRatioError: If the hit align ratio to the query was too small.
DuplicateError: If the hit was an exact subsequence of the query.
LengthError: If the hit was too short.
@@ -241,13 +254,6 @@ def _assess_hhsearch_hit(
template_sequence = hit.hit_sequence.replace("-", "")
length_ratio = float(len(template_sequence)) / len(query_sequence)
- # Check whether the template is a large subsequence or duplicate of original
- # query. This can happen due to duplicate entries in the PDB database.
- duplicate = (
- template_sequence in query_sequence
- and length_ratio > max_subsequence_ratio
- )
-
if _is_after_cutoff(hit_pdb_code, release_dates, release_date_cutoff):
date = release_dates[hit_pdb_code.upper()]
raise DateError(
@@ -255,16 +261,19 @@ def _assess_hhsearch_hit(
f"({release_date_cutoff})."
)
- if query_pdb_code is not None:
- if query_pdb_code.lower() == hit_pdb_code.lower():
- raise PdbIdError("PDB code identical to Query PDB code.")
-
if align_ratio <= min_align_ratio:
raise AlignRatioError(
"Proportion of residues aligned to query too small. "
f"Align ratio: {align_ratio}."
)
+ # Check whether the template is a large subsequence or duplicate of original
+ # query. This can happen due to duplicate entries in the PDB database.
+ duplicate = (
+ template_sequence in query_sequence
+ and length_ratio > max_subsequence_ratio
+ )
+
if duplicate:
raise DuplicateError(
"Template is an exact subsequence of query with large "
@@ -424,9 +433,10 @@ def _realign_pdb_template_to_query(
)
try:
- (old_aligned_template, new_aligned_template), _ = parsers.parse_a3m(
+ parsed_a3m = parsers.parse_a3m(
aligner.align([old_template_sequence, new_template_sequence])
)
+ old_aligned_template, new_aligned_template = parsed_a3m.sequences
except Exception as e:
raise QueryToTemplateAlignError(
"Could not align old template %s to template %s (%s_%s). Error: %s"
@@ -768,7 +778,6 @@ class SingleHitResult:
def _prefilter_hit(
query_sequence: str,
- query_pdb_code: Optional[str],
hit: parsers.TemplateHit,
max_template_date: datetime.datetime,
release_dates: Mapping[str, datetime.datetime],
@@ -789,17 +798,14 @@ def _prefilter_hit(
hit=hit,
hit_pdb_code=hit_pdb_code,
query_sequence=query_sequence,
- query_pdb_code=query_pdb_code,
release_dates=release_dates,
release_date_cutoff=max_template_date,
)
except PrefilterError as e:
hit_name = f"{hit_pdb_code}_{hit_chain_id}"
msg = f"hit {hit_name} did not pass prefilter: {str(e)}"
- logging.info("%s: %s", query_pdb_code, msg)
- if strict_error_check and isinstance(
- e, (DateError, PdbIdError, DuplicateError)
- ):
+ logging.info(msg)
+ if strict_error_check and isinstance(e, (DateError, DuplicateError)):
# In strict mode we treat some prefilter cases as errors.
return PrefilterResult(valid=False, error=msg, warning=None)
@@ -808,9 +814,16 @@ def _prefilter_hit(
return PrefilterResult(valid=True, error=None, warning=None)
+@functools.lru_cache(16, typed=False)
+def _read_file(path):
+ with open(path, 'r') as f:
+ file_data = f.read()
+
+ return file_data
+
+
def _process_single_hit(
query_sequence: str,
- query_pdb_code: Optional[str],
hit: parsers.TemplateHit,
mmcif_dir: str,
max_template_date: datetime.datetime,
@@ -847,9 +860,9 @@ def _process_single_hit(
query_sequence,
template_sequence,
)
+
# Fail if we can't find the mmCIF file.
- with open(cif_path, "r") as cif_file:
- cif_string = cif_file.read()
+ cif_string = _read_file(cif_path)
parsing_result = mmcif_parsing.parse(
file_id=hit_pdb_code, mmcif_string=cif_string
@@ -882,7 +895,11 @@ def _process_single_hit(
kalign_binary_path=kalign_binary_path,
_zero_center_positions=_zero_center_positions,
)
- features["template_sum_probs"] = [hit.sum_probs]
+
+ if hit.sum_probs is None:
+ features["template_sum_probs"] = [0]
+ else:
+ features["template_sum_probs"] = [hit.sum_probs]
# It is possible there were some errors when parsing the other chains in the
# mmCIF file, but the template features for the chain we want were still
@@ -903,7 +920,7 @@ def _process_single_hit(
% (
hit_pdb_code,
hit_chain_id,
- hit.sum_probs,
+ hit.sum_probs if hit.sum_probs else 0.,
hit.index,
str(e),
parsing_result.errors,
@@ -920,7 +937,7 @@ def _process_single_hit(
% (
hit_pdb_code,
hit_chain_id,
- hit.sum_probs,
+ hit.sum_probs if hit.sum_probs else 0.,
hit.index,
str(e),
parsing_result.errors,
@@ -986,8 +1003,8 @@ class TemplateSearchResult:
warnings: Sequence[str]
-class TemplateHitFeaturizer:
- """A class for turning hhr hits to template features."""
+class TemplateHitFeaturizer(abc.ABC):
+ """An abstract base class for turning template hits to features."""
def __init__(
self,
mmcif_dir: str,
@@ -1036,7 +1053,7 @@ def __init__(
raise ValueError(
"max_template_date must be set and have format YYYY-MM-DD."
)
- self.max_hits = max_hits
+ self._max_hits = max_hits
self._kalign_binary_path = kalign_binary_path
self._strict_error_check = strict_error_check
@@ -1059,31 +1076,29 @@ def __init__(
self._shuffle_top_k_prefiltered = _shuffle_top_k_prefiltered
self._zero_center_positions = _zero_center_positions
+ @abc.abstractmethod
+ def get_templates(
+ self,
+ query_sequence: str,
+ hits: Sequence[parsers.TemplateHit]
+ ) -> TemplateSearchResult:
+ """ Computes the templates for a given query sequence """
+
+
+class HhsearchHitFeaturizer(TemplateHitFeaturizer):
def get_templates(
self,
query_sequence: str,
- query_pdb_code: Optional[str],
- query_release_date: Optional[datetime.datetime],
hits: Sequence[parsers.TemplateHit],
) -> TemplateSearchResult:
"""Computes the templates for given query sequence (more details above)."""
- logging.info("Searching for template for: %s", query_pdb_code)
+ logging.info("Searching for template for: %s", query_sequence)
template_features = {}
for template_feature_name in TEMPLATE_FEATURES:
template_features[template_feature_name] = []
- # Always use a max_template_date. Set to query_release_date minus 60 days
- # if that's earlier.
- template_cutoff_date = self._max_template_date
- if query_release_date:
- delta = datetime.timedelta(days=60)
- if query_release_date - delta < template_cutoff_date:
- template_cutoff_date = query_release_date - delta
- assert template_cutoff_date < query_release_date
- assert template_cutoff_date <= self._max_template_date
-
- num_hits = 0
+ already_seen = set()
errors = []
warnings = []
@@ -1091,9 +1106,8 @@ def get_templates(
for hit in hits:
prefilter_result = _prefilter_hit(
query_sequence=query_sequence,
- query_pdb_code=query_pdb_code,
hit=hit,
- max_template_date=template_cutoff_date,
+ max_template_date=self._max_template_date,
release_dates=self._release_dates,
obsolete_pdbs=self._obsolete_pdbs,
strict_error_check=self._strict_error_check,
@@ -1119,17 +1133,16 @@ def get_templates(
for i in idx:
# We got all the templates we wanted, stop processing hits.
- if num_hits >= self.max_hits:
+ if len(already_seen) >= self._max_hits:
break
hit = filtered[i]
result = _process_single_hit(
query_sequence=query_sequence,
- query_pdb_code=query_pdb_code,
hit=hit,
mmcif_dir=self._mmcif_dir,
- max_template_date=template_cutoff_date,
+ max_template_date=self._max_template_date,
release_dates=self._release_dates,
obsolete_pdbs=self._obsolete_pdbs,
strict_error_check=self._strict_error_check,
@@ -1153,22 +1166,124 @@ def get_templates(
result.warning,
)
else:
- # Increment the hit counter, since we got features out of this hit.
- num_hits += 1
+ already_seen_key = result.features["template_sequence"]
+ if(already_seen_key in already_seen):
+ continue
+ already_seen.add(already_seen_key)
for k in template_features:
template_features[k].append(result.features[k])
- for name in template_features:
- if num_hits > 0:
+ if already_seen:
+ for name in template_features:
template_features[name] = np.stack(
template_features[name], axis=0
).astype(TEMPLATE_FEATURES[name])
- else:
- # Make sure the feature has correct dtype even if empty.
- template_features[name] = np.array(
- [], dtype=TEMPLATE_FEATURES[name]
- )
+ else:
+ num_res = len(query_sequence)
+ # Construct a default template with all zeros.
+ template_features = empty_template_feats(num_res)
return TemplateSearchResult(
features=template_features, errors=errors, warnings=warnings
)
+
+
+class HmmsearchHitFeaturizer(TemplateHitFeaturizer):
+ def get_templates(
+ self,
+ query_sequence: str,
+ hits: Sequence[parsers.TemplateHit]
+ ) -> TemplateSearchResult:
+ logging.info("Searching for template for: %s", query_sequence)
+
+ template_features = {}
+ for template_feature_name in TEMPLATE_FEATURES:
+ template_features[template_feature_name] = []
+
+ already_seen = set()
+ errors = []
+ warnings = []
+
+ # DISCREPANCY: This filtering scheme that saves time
+ filtered = []
+ for hit in hits:
+ prefilter_result = _prefilter_hit(
+ query_sequence=query_sequence,
+ hit=hit,
+ max_template_date=self._max_template_date,
+ release_dates=self._release_dates,
+ obsolete_pdbs=self._obsolete_pdbs,
+ strict_error_check=self._strict_error_check,
+ )
+
+ if prefilter_result.error:
+ errors.append(prefilter_result.error)
+
+ if prefilter_result.warning:
+ warnings.append(prefilter_result.warning)
+
+ if prefilter_result.valid:
+ filtered.append(hit)
+
+ filtered = list(
+ sorted(
+ filtered, key=lambda x: x.sum_probs if x.sum_probs else 0., reverse=True
+ )
+ )
+ idx = list(range(len(filtered)))
+ if(self._shuffle_top_k_prefiltered):
+ stk = self._shuffle_top_k_prefiltered
+ idx[:stk] = np.random.permutation(idx[:stk])
+
+ for i in idx:
+ if(len(already_seen) >= self._max_hits):
+ break
+
+ hit = filtered[i]
+
+ result = _process_single_hit(
+ query_sequence=query_sequence,
+ hit=hit,
+ mmcif_dir=self._mmcif_dir,
+ max_template_date = self._max_template_date,
+ release_dates = self._release_dates,
+ obsolete_pdbs = self._obsolete_pdbs,
+ strict_error_check = self._strict_error_check,
+ kalign_binary_path = self._kalign_binary_path
+ )
+
+ if result.error:
+ errors.append(result.error)
+
+ if result.warning:
+ warnings.append(result.warning)
+
+ if result.features is None:
+ logging.debug(
+ "Skipped invalid hit %s, error: %s, warning: %s",
+ hit.name, result.error, result.warning,
+ )
+ else:
+ already_seen_key = result.features["template_sequence"]
+ if(already_seen_key in already_seen):
+ continue
+ # Increment the hit counter, since we got features out of this hit.
+ already_seen.add(already_seen_key)
+ for k in template_features:
+ template_features[k].append(result.features[k])
+
+ if already_seen:
+ for name in template_features:
+ template_features[name] = np.stack(
+ template_features[name], axis=0
+ ).astype(TEMPLATE_FEATURES[name])
+ else:
+ num_res = len(query_sequence)
+ # Construct a default template with all zeros.
+ template_features = empty_template_feats(num_res)
+
+ return TemplateSearchResult(
+ features=template_features,
+ errors=errors,
+ warnings=warnings,
+ )
diff --git a/openfold/data/tools/hhblits.py b/openfold/data/tools/hhblits.py
index 3e1c1e7b..bcd3013b 100644
--- a/openfold/data/tools/hhblits.py
+++ b/openfold/data/tools/hhblits.py
@@ -18,7 +18,7 @@
import logging
import os
import subprocess
-from typing import Any, Mapping, Optional, Sequence
+from typing import Any, List, Mapping, Optional, Sequence
from openfold.data.tools import utils
@@ -99,9 +99,9 @@ def __init__(
self.p = p
self.z = z
- def query(self, input_fasta_path: str) -> Mapping[str, Any]:
+ def query(self, input_fasta_path: str) -> List[Mapping[str, Any]]:
"""Queries the database using HHblits."""
- with utils.tmpdir_manager(base_dir="/tmp") as query_tmp_dir:
+ with utils.tmpdir_manager() as query_tmp_dir:
a3m_path = os.path.join(query_tmp_dir, "output.a3m")
db_cmd = []
@@ -172,4 +172,4 @@ def query(self, input_fasta_path: str) -> Mapping[str, Any]:
n_iter=self.n_iter,
e_value=self.e_value,
)
- return raw_output
+ return [raw_output]
diff --git a/openfold/data/tools/hhsearch.py b/openfold/data/tools/hhsearch.py
index d62e3b41..42a6edf6 100644
--- a/openfold/data/tools/hhsearch.py
+++ b/openfold/data/tools/hhsearch.py
@@ -18,8 +18,9 @@
import logging
import os
import subprocess
-from typing import Sequence
+from typing import Sequence, Optional
+from openfold.data import parsers
from openfold.data.tools import utils
@@ -62,11 +63,20 @@ def __init__(
f"Could not find HHsearch database {database_path}"
)
- def query(self, a3m: str) -> str:
+ @property
+ def output_format(self) -> str:
+ return 'hhr'
+
+ @property
+ def input_format(self) -> str:
+ return 'a3m'
+
+ def query(self, a3m: str, output_dir: Optional[str] = None) -> str:
"""Queries the database using HHsearch using a given a3m."""
- with utils.tmpdir_manager(base_dir="/tmp") as query_tmp_dir:
+ with utils.tmpdir_manager() as query_tmp_dir:
input_path = os.path.join(query_tmp_dir, "query.a3m")
- hhr_path = os.path.join(query_tmp_dir, "output.hhr")
+ output_dir = query_tmp_dir if output_dir is None else output_dir
+ hhr_path = os.path.join(output_dir, "hhsearch_output.hhr")
with open(input_path, "w") as f:
f.write(a3m)
@@ -104,3 +114,12 @@ def query(self, a3m: str) -> str:
with open(hhr_path) as f:
hhr = f.read()
return hhr
+
+ @staticmethod
+ def get_template_hits(
+ output_string: str,
+ input_sequence: str
+ ) -> Sequence[parsers.TemplateHit]:
+ """Gets parsed template hits from the raw string output by the tool"""
+ del input_sequence # Used by hmmsearch but not needed for hhsearch
+ return parsers.parse_hhr(output_string)
diff --git a/openfold/data/tools/hmmbuild.py b/openfold/data/tools/hmmbuild.py
new file mode 100644
index 00000000..e554ed23
--- /dev/null
+++ b/openfold/data/tools/hmmbuild.py
@@ -0,0 +1,137 @@
+# 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.
+
+"""A Python wrapper for hmmbuild - construct HMM profiles from MSA."""
+
+import os
+import re
+import subprocess
+
+from absl import logging
+from openfold.data.tools import utils
+
+
+class Hmmbuild(object):
+ """Python wrapper of the hmmbuild binary."""
+
+ def __init__(self,
+ *,
+ binary_path: str,
+ singlemx: bool = False):
+ """Initializes the Python hmmbuild wrapper.
+
+ Args:
+ binary_path: The path to the hmmbuild executable.
+ singlemx: Whether to use --singlemx flag. If True, it forces HMMBuild to
+ just use a common substitution score matrix.
+
+ Raises:
+ RuntimeError: If hmmbuild binary not found within the path.
+ """
+ self.binary_path = binary_path
+ self.singlemx = singlemx
+
+ def build_profile_from_sto(self, sto: str, model_construction='fast') -> str:
+ """Builds a HHM for the aligned sequences given as an A3M string.
+
+ Args:
+ sto: A string with the aligned sequences in the Stockholm format.
+ model_construction: Whether to use reference annotation in the msa to
+ determine consensus columns ('hand') or default ('fast').
+
+ Returns:
+ A string with the profile in the HMM format.
+
+ Raises:
+ RuntimeError: If hmmbuild fails.
+ """
+ return self._build_profile(sto, model_construction=model_construction)
+
+ def build_profile_from_a3m(self, a3m: str) -> str:
+ """Builds a HHM for the aligned sequences given as an A3M string.
+
+ Args:
+ a3m: A string with the aligned sequences in the A3M format.
+
+ Returns:
+ A string with the profile in the HMM format.
+
+ Raises:
+ RuntimeError: If hmmbuild fails.
+ """
+ lines = []
+ for line in a3m.splitlines():
+ if not line.startswith('>'):
+ line = re.sub('[a-z]+', '', line) # Remove inserted residues.
+ lines.append(line + '\n')
+ msa = ''.join(lines)
+ return self._build_profile(msa, model_construction='fast')
+
+ def _build_profile(self, msa: str, model_construction: str = 'fast') -> str:
+ """Builds a HMM for the aligned sequences given as an MSA string.
+
+ Args:
+ msa: A string with the aligned sequences, in A3M or STO format.
+ model_construction: Whether to use reference annotation in the msa to
+ determine consensus columns ('hand') or default ('fast').
+
+ Returns:
+ A string with the profile in the HMM format.
+
+ Raises:
+ RuntimeError: If hmmbuild fails.
+ ValueError: If unspecified arguments are provided.
+ """
+ if model_construction not in {'hand', 'fast'}:
+ raise ValueError(f'Invalid model_construction {model_construction} - only'
+ 'hand and fast supported.')
+
+ with utils.tmpdir_manager() as query_tmp_dir:
+ input_query = os.path.join(query_tmp_dir, 'query.msa')
+ output_hmm_path = os.path.join(query_tmp_dir, 'output.hmm')
+
+ with open(input_query, 'w') as f:
+ f.write(msa)
+
+ cmd = [self.binary_path]
+ # If adding flags, we have to do so before the output and input:
+
+ if model_construction == 'hand':
+ cmd.append(f'--{model_construction}')
+ if self.singlemx:
+ cmd.append('--singlemx')
+ cmd.extend([
+ '--amino',
+ output_hmm_path,
+ input_query,
+ ])
+
+ logging.info('Launching subprocess %s', cmd)
+ process = subprocess.Popen(cmd, stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE)
+
+ with utils.timing('hmmbuild query'):
+ stdout, stderr = process.communicate()
+ retcode = process.wait()
+ logging.info('hmmbuild stdout:\n%s\n\nstderr:\n%s\n',
+ stdout.decode('utf-8'), stderr.decode('utf-8'))
+
+ if retcode:
+ raise RuntimeError('hmmbuild failed\nstdout:\n%s\n\nstderr:\n%s\n'
+ % (stdout.decode('utf-8'), stderr.decode('utf-8')))
+
+ with open(output_hmm_path, encoding='utf-8') as f:
+ hmm = f.read()
+
+ return hmm
diff --git a/openfold/data/tools/hmmsearch.py b/openfold/data/tools/hmmsearch.py
new file mode 100644
index 00000000..adba37f3
--- /dev/null
+++ b/openfold/data/tools/hmmsearch.py
@@ -0,0 +1,137 @@
+# 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.
+
+"""A Python wrapper for hmmsearch - search profile against a sequence db."""
+
+import os
+import subprocess
+from typing import Optional, Sequence
+
+from absl import logging
+from openfold.data import parsers
+from openfold.data.tools import hmmbuild
+from openfold.data.tools import utils
+
+
+class Hmmsearch(object):
+ """Python wrapper of the hmmsearch binary."""
+
+ def __init__(self,
+ *,
+ binary_path: str,
+ hmmbuild_binary_path: str,
+ database_path: str,
+ flags: Optional[Sequence[str]] = None
+ ):
+ """Initializes the Python hmmsearch wrapper.
+
+ Args:
+ binary_path: The path to the hmmsearch executable.
+ hmmbuild_binary_path: The path to the hmmbuild executable. Used to build
+ an hmm from an input a3m.
+ database_path: The path to the hmmsearch database (FASTA format).
+ flags: List of flags to be used by hmmsearch.
+
+ Raises:
+ RuntimeError: If hmmsearch binary not found within the path.
+ """
+ self.binary_path = binary_path
+ self.hmmbuild_runner = hmmbuild.Hmmbuild(binary_path=hmmbuild_binary_path)
+ self.database_path = database_path
+ if flags is None:
+ # Default hmmsearch run settings.
+ flags = ['--F1', '0.1',
+ '--F2', '0.1',
+ '--F3', '0.1',
+ '--incE', '100',
+ '-E', '100',
+ '--domE', '100',
+ '--incdomE', '100']
+ self.flags = flags
+
+ if not os.path.exists(self.database_path):
+ logging.error('Could not find hmmsearch database %s', database_path)
+ raise ValueError(f'Could not find hmmsearch database {database_path}')
+
+ @property
+ def output_format(self) -> str:
+ return 'sto'
+
+ @property
+ def input_format(self) -> str:
+ return 'sto'
+
+ def query(self, msa_sto: str, output_dir: Optional[str] = None) -> str:
+ """Queries the database using hmmsearch using a given stockholm msa."""
+ hmm = self.hmmbuild_runner.build_profile_from_sto(
+ msa_sto,
+ model_construction='hand'
+ )
+ return self.query_with_hmm(hmm, output_dir)
+
+ def query_with_hmm(self,
+ hmm: str,
+ output_dir: Optional[str] = None
+ ) -> str:
+ """Queries the database using hmmsearch using a given hmm."""
+ with utils.tmpdir_manager() as query_tmp_dir:
+ hmm_input_path = os.path.join(query_tmp_dir, 'query.hmm')
+ output_dir = query_tmp_dir if output_dir is None else output_dir
+ out_path = os.path.join(output_dir, 'hmm_output.sto')
+ with open(hmm_input_path, 'w') as f:
+ f.write(hmm)
+
+ cmd = [
+ self.binary_path,
+ '--noali', # Don't include the alignment in stdout.
+ '--cpu', '8'
+ ]
+ # If adding flags, we have to do so before the output and input:
+ if self.flags:
+ cmd.extend(self.flags)
+ cmd.extend([
+ '-A', out_path,
+ hmm_input_path,
+ self.database_path,
+ ])
+
+ logging.info('Launching sub-process %s', cmd)
+ process = subprocess.Popen(
+ cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+ with utils.timing(
+ f'hmmsearch ({os.path.basename(self.database_path)}) query'):
+ stdout, stderr = process.communicate()
+ retcode = process.wait()
+
+ if retcode:
+ raise RuntimeError(
+ 'hmmsearch failed:\nstdout:\n%s\n\nstderr:\n%s\n' % (
+ stdout.decode('utf-8'), stderr.decode('utf-8')))
+
+ with open(out_path) as f:
+ out_msa = f.read()
+
+ return out_msa
+
+ @staticmethod
+ def get_template_hits(
+ output_string: str,
+ input_sequence: str
+ ) -> Sequence[parsers.TemplateHit]:
+ """Gets parsed template hits from the raw string output by the tool."""
+ template_hits = parsers.parse_hmmsearch_sto(
+ output_string,
+ input_sequence,
+ )
+ return template_hits
diff --git a/openfold/data/tools/jackhmmer.py b/openfold/data/tools/jackhmmer.py
index fca93edf..67544d18 100644
--- a/openfold/data/tools/jackhmmer.py
+++ b/openfold/data/tools/jackhmmer.py
@@ -23,6 +23,7 @@
from typing import Any, Callable, Mapping, Optional, Sequence
from urllib import request
+from openfold.data import parsers
from openfold.data.tools import utils
@@ -93,10 +94,13 @@ def __init__(
self.streaming_callback = streaming_callback
def _query_chunk(
- self, input_fasta_path: str, database_path: str
+ self,
+ input_fasta_path: str,
+ database_path: str,
+ max_sequences: Optional[int] = None
) -> Mapping[str, Any]:
"""Queries the database chunk using Jackhmmer."""
- with utils.tmpdir_manager(base_dir="/tmp") as query_tmp_dir:
+ with utils.tmpdir_manager() as query_tmp_dir:
sto_path = os.path.join(query_tmp_dir, "output.sto")
# The F1/F2/F3 are the expected proportion to pass each of the filtering
@@ -167,8 +171,11 @@ def _query_chunk(
with open(tblout_path) as f:
tbl = f.read()
- with open(sto_path) as f:
- sto = f.read()
+ if(max_sequences is None):
+ with open(sto_path) as f:
+ sto = f.read()
+ else:
+ sto = parsers.truncate_stockholm_msa(sto_path, max_sequences)
raw_output = dict(
sto=sto,
@@ -179,11 +186,26 @@ def _query_chunk(
)
return raw_output
+
+ def query(self,
+ input_fasta_path: str,
+ max_sequences: Optional[int] = None
+ ) -> Sequence[Sequence[Mapping[str, Any]]]:
+ return self.query_multiple([input_fasta_path], max_sequences)
- def query(self, input_fasta_path: str) -> Sequence[Mapping[str, Any]]:
+ def query_multiple(self,
+ input_fasta_paths: Sequence[str],
+ max_sequences: Optional[int] = None
+ ) -> Sequence[Sequence[Mapping[str, Any]]]:
"""Queries the database using Jackhmmer."""
if self.num_streamed_chunks is None:
- return [self._query_chunk(input_fasta_path, self.database_path)]
+ single_chunk_results = []
+ for input_fasta_path in input_fasta_paths:
+ single_chunk_result = self._query_chunk(
+ input_fasta_path, self.database_path, max_sequences,
+ )
+ single_chunk_results.append(single_chunk_result)
+ return single_chunk_results
db_basename = os.path.basename(self.database_path)
db_remote_chunk = lambda db_idx: f"{self.database_path}.{db_idx}"
@@ -198,7 +220,7 @@ def query(self, input_fasta_path: str) -> Sequence[Mapping[str, Any]]:
# Download the (i+1)-th chunk while Jackhmmer is running on the i-th chunk
with futures.ThreadPoolExecutor(max_workers=2) as executor:
- chunked_output = []
+ chunked_outputs = [[] for _ in range(len(input_fasta_paths))]
for i in range(1, self.num_streamed_chunks + 1):
# Copy the chunk locally
if i == 1:
@@ -216,13 +238,21 @@ def query(self, input_fasta_path: str) -> Sequence[Mapping[str, Any]]:
# Run Jackhmmer with the chunk
future.result()
- chunked_output.append(
- self._query_chunk(input_fasta_path, db_local_chunk(i))
- )
+ for fasta_idx, input_fasta_path in enumerate(input_fasta_paths):
+ chunked_outputs[fasta_idx].append(
+ self._query_chunk(
+ input_fasta_path,
+ db_local_chunk(i),
+ max_sequences
+ )
+ )
# Remove the local copy of the chunk
os.remove(db_local_chunk(i))
- future = next_future
+ # Do not set next_future for the last chunk so that this works
+ # even for databases with only 1 chunk
+ if(i < self.num_streamed_chunks):
+ future = next_future
if self.streaming_callback:
self.streaming_callback(i)
- return chunked_output
+ return chunked_outputs
diff --git a/openfold/data/tools/kalign.py b/openfold/data/tools/kalign.py
index 5306f864..a9f8aec8 100644
--- a/openfold/data/tools/kalign.py
+++ b/openfold/data/tools/kalign.py
@@ -72,7 +72,7 @@ def align(self, sequences: Sequence[str]) -> str:
"residues long. Got %s (%d residues)." % (s, len(s))
)
- with utils.tmpdir_manager(base_dir="/tmp") as query_tmp_dir:
+ with utils.tmpdir_manager() as query_tmp_dir:
input_fasta_path = os.path.join(query_tmp_dir, "input.fasta")
output_a3m_path = os.path.join(query_tmp_dir, "output.a3m")
diff --git a/openfold/data/tools/parse_msa_files.py b/openfold/data/tools/parse_msa_files.py
new file mode 100644
index 00000000..04450c22
--- /dev/null
+++ b/openfold/data/tools/parse_msa_files.py
@@ -0,0 +1,57 @@
+import os, argparse, pickle, tempfile, concurrent
+from openfold.data import parsers
+from concurrent.futures import ProcessPoolExecutor
+
+
+def parse_stockholm_file(alignment_dir: str, stockholm_file: str):
+ path = os.path.join(alignment_dir, stockholm_file)
+ file_name,_ = os.path.splitext(stockholm_file)
+ with open(path, "r") as infile:
+ msa = parsers.parse_stockholm(infile.read())
+ infile.close()
+ return {file_name: msa}
+
+
+def parse_a3m_file(alignment_dir: str, a3m_file: str):
+ path = os.path.join(alignment_dir, a3m_file)
+ file_name,_ = os.path.splitext(a3m_file)
+ with open(path, "r") as infile:
+ msa = parsers.parse_a3m(infile.read())
+ infile.close()
+ return {file_name: msa}
+
+
+def run_parse_all_msa_files_multiprocessing(stockholm_files: list, a3m_files: list, alignment_dir:str):
+ # Number of workers based on the tasks
+ msa_results={}
+ a3m_tasks = [(alignment_dir, f) for f in a3m_files]
+ sto_tasks = [(alignment_dir, f) for f in stockholm_files]
+ with ProcessPoolExecutor(max_workers = len(a3m_tasks) + len(sto_tasks)) as executor:
+ a3m_futures = {executor.submit(parse_a3m_file, *task): task for task in a3m_tasks}
+ sto_futures = {executor.submit(parse_stockholm_file, *task): task for task in sto_tasks}
+
+ for future in concurrent.futures.as_completed(a3m_futures | sto_futures):
+ try:
+ result = future.result()
+ msa_results.update(result)
+ except Exception as exc:
+ print(f'Task generated an exception: {exc}')
+ return msa_results
+
+
+def main():
+ parser = argparse.ArgumentParser(description='Process msa files in parallel')
+ parser.add_argument('--alignment_dir', type=str, help='path to alignment dir')
+ args = parser.parse_args()
+ alignment_dir = args.alignment_dir
+ stockholm_files = [i for i in os.listdir(alignment_dir)
+ if all([i.endswith('.sto'), "hmm_output" not in i, "uniprot_hits" not in i])]
+ a3m_files = [i for i in os.listdir(alignment_dir) if i.endswith('.a3m')]
+ msa_data = run_parse_all_msa_files_multiprocessing(stockholm_files, a3m_files, alignment_dir)
+ with tempfile.NamedTemporaryFile('wb', suffix='.pkl', delete=False) as outfile:
+ pickle.dump(msa_data, outfile)
+ print(outfile.name)
+
+
+if __name__ == "__main__":
+ main()
\ No newline at end of file
diff --git a/openfold/model/embedders.py b/openfold/model/embedders.py
index f7cd94c0..69cf4951 100644
--- a/openfold/model/embedders.py
+++ b/openfold/model/embedders.py
@@ -13,12 +13,26 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from functools import partial
+
import torch
import torch.nn as nn
from typing import Tuple, Optional
+from openfold.utils import all_atom_multimer
+from openfold.utils.feats import (
+ pseudo_beta_fn,
+ dgram_from_positions,
+ build_template_angle_feat,
+ build_template_pair_feat,
+)
from openfold.model.primitives import Linear, LayerNorm
-from openfold.utils.tensor_utils import add, one_hot
+from openfold.model.template import (
+ TemplatePairStack,
+ TemplatePointwiseAttention,
+)
+from openfold.utils import geometry
+from openfold.utils.tensor_utils import add, one_hot, tensor_tree_map, dict_multimap
class InputEmbedder(nn.Module):
@@ -99,12 +113,13 @@ def forward(
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Args:
- tf:
- "target_feat" features of shape [*, N_res, tf_dim]
- ri:
- "residue_index" features of shape [*, N_res]
- msa:
- "msa_feat" features of shape [*, N_clust, N_res, msa_dim]
+ batch: Dict containing
+ "target_feat":
+ Features of shape [*, N_res, tf_dim]
+ "residue_index":
+ Features of shape [*, N_res]
+ "msa_feat":
+ Features of shape [*, N_clust, N_res, msa_dim]
Returns:
msa_emb:
[*, N_clust, N_res, C_m] MSA embedding
@@ -139,6 +154,161 @@ def forward(
return msa_emb, pair_emb
+class InputEmbedderMultimer(nn.Module):
+ """
+ Embeds a subset of the input features.
+
+ Implements Algorithms 3 (InputEmbedder) and 4 (relpos).
+ """
+
+ def __init__(
+ self,
+ tf_dim: int,
+ msa_dim: int,
+ c_z: int,
+ c_m: int,
+ max_relative_idx: int,
+ use_chain_relative: bool,
+ max_relative_chain: int,
+ **kwargs,
+ ):
+ """
+ Args:
+ tf_dim:
+ Final dimension of the target features
+ msa_dim:
+ Final dimension of the MSA features
+ c_z:
+ Pair embedding dimension
+ c_m:
+ MSA embedding dimension
+ relpos_k:
+ Window size used in relative positional encoding
+ """
+ super(InputEmbedderMultimer, self).__init__()
+
+ self.tf_dim = tf_dim
+ self.msa_dim = msa_dim
+
+ self.c_z = c_z
+ self.c_m = c_m
+
+ self.linear_tf_z_i = Linear(tf_dim, c_z)
+ self.linear_tf_z_j = Linear(tf_dim, c_z)
+ self.linear_tf_m = Linear(tf_dim, c_m)
+ self.linear_msa_m = Linear(msa_dim, c_m)
+
+ # RPE stuff
+ self.max_relative_idx = max_relative_idx
+ self.use_chain_relative = use_chain_relative
+ self.max_relative_chain = max_relative_chain
+ if(self.use_chain_relative):
+ self.no_bins = (
+ 2 * max_relative_idx + 2 +
+ 1 +
+ 2 * max_relative_chain + 2
+ )
+ else:
+ self.no_bins = 2 * max_relative_idx + 1
+ self.linear_relpos = Linear(self.no_bins, c_z)
+
+ def relpos(self, batch):
+ pos = batch["residue_index"]
+ asym_id = batch["asym_id"]
+ asym_id_same = (asym_id[..., None] == asym_id[..., None, :])
+ offset = pos[..., None] - pos[..., None, :]
+
+ clipped_offset = torch.clamp(
+ offset + self.max_relative_idx, 0, 2 * self.max_relative_idx
+ )
+
+ rel_feats = []
+ if(self.use_chain_relative):
+ final_offset = torch.where(
+ asym_id_same,
+ clipped_offset,
+ (2 * self.max_relative_idx + 1) *
+ torch.ones_like(clipped_offset)
+ )
+ boundaries = torch.arange(
+ start=0, end=2 * self.max_relative_idx + 2, device=final_offset.device
+ )
+ rel_pos = one_hot(
+ final_offset,
+ boundaries,
+ )
+
+ rel_feats.append(rel_pos)
+
+ entity_id = batch["entity_id"]
+ entity_id_same = (entity_id[..., None] == entity_id[..., None, :])
+ rel_feats.append(entity_id_same[..., None].to(dtype=rel_pos.dtype))
+
+ sym_id = batch["sym_id"]
+ rel_sym_id = sym_id[..., None] - sym_id[..., None, :]
+
+ max_rel_chain = self.max_relative_chain
+ clipped_rel_chain = torch.clamp(
+ rel_sym_id + max_rel_chain,
+ 0,
+ 2 * max_rel_chain,
+ )
+
+ final_rel_chain = torch.where(
+ entity_id_same,
+ clipped_rel_chain,
+ (2 * max_rel_chain + 1) *
+ torch.ones_like(clipped_rel_chain)
+ )
+
+ boundaries = torch.arange(
+ start=0, end=2 * max_rel_chain + 2, device=final_rel_chain.device
+ )
+ rel_chain = one_hot(
+ final_rel_chain,
+ boundaries,
+ )
+
+ rel_feats.append(rel_chain)
+ else:
+ boundaries = torch.arange(
+ start=0, end=2 * self.max_relative_idx + 1, device=clipped_offset.device
+ )
+ rel_pos = one_hot(
+ clipped_offset, boundaries,
+ )
+ rel_feats.append(rel_pos)
+
+ rel_feat = torch.cat(rel_feats, dim=-1).to(
+ self.linear_relpos.weight.dtype
+ )
+
+ return self.linear_relpos(rel_feat)
+
+ def forward(self, batch) -> Tuple[torch.Tensor, torch.Tensor]:
+ tf = batch["target_feat"]
+ msa = batch["msa_feat"]
+
+ # [*, N_res, c_z]
+ tf_emb_i = self.linear_tf_z_i(tf)
+ tf_emb_j = self.linear_tf_z_j(tf)
+
+ # [*, N_res, N_res, c_z]
+ pair_emb = tf_emb_i[..., None, :] + tf_emb_j[..., None, :, :]
+ pair_emb = pair_emb + self.relpos(batch)
+
+ # [*, N_clust, N_res, c_m]
+ n_clust = msa.shape[-3]
+ tf_m = (
+ self.linear_tf_m(tf)
+ .unsqueeze(-3)
+ .expand(((-1,) * len(tf.shape[:-2]) + (n_clust, -1, -1)))
+ )
+ msa_emb = self.linear_msa_m(msa) + tf_m
+
+ return msa_emb, pair_emb
+
+
class PreembeddingEmbedder(nn.Module):
"""
Embeds the sequence pre-embedding passed to the model and the target_feat features.
@@ -335,7 +505,7 @@ def forward(
return m_update, z_update
-class TemplateAngleEmbedder(nn.Module):
+class TemplateSingleEmbedder(nn.Module):
"""
Embeds the "template_angle_feat" feature.
@@ -355,7 +525,7 @@ def __init__(
c_out:
Output channel dimension
"""
- super(TemplateAngleEmbedder, self).__init__()
+ super(TemplateSingleEmbedder, self).__init__()
self.c_out = c_out
self.c_in = c_in
@@ -459,3 +629,356 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.linear(x)
return x
+
+
+class TemplateEmbedder(nn.Module):
+ def __init__(self, config):
+ super(TemplateEmbedder, self).__init__()
+
+ self.config = config
+ self.template_single_embedder = TemplateSingleEmbedder(
+ **config["template_single_embedder"],
+ )
+ self.template_pair_embedder = TemplatePairEmbedder(
+ **config["template_pair_embedder"],
+ )
+ self.template_pair_stack = TemplatePairStack(
+ **config["template_pair_stack"],
+ )
+ self.template_pointwise_att = TemplatePointwiseAttention(
+ **config["template_pointwise_attention"],
+ )
+
+ def forward(
+ self,
+ batch,
+ z,
+ pair_mask,
+ templ_dim,
+ chunk_size,
+ _mask_trans=True,
+ use_deepspeed_evo_attention=False,
+ use_lma=False,
+ inplace_safe=False
+ ):
+ # Embed the templates one at a time (with a poor man's vmap)
+ pair_embeds = []
+ n = z.shape[-2]
+ n_templ = batch["template_aatype"].shape[templ_dim]
+
+ if (inplace_safe):
+ # We'll preallocate the full pair tensor now to avoid manifesting
+ # a second copy during the stack later on
+ t_pair = z.new_zeros(
+ z.shape[:-3] +
+ (n_templ, n, n, self.config.template_pair_embedder.c_out)
+ )
+
+ for i in range(n_templ):
+ idx = batch["template_aatype"].new_tensor(i)
+ single_template_feats = tensor_tree_map(
+ lambda t: torch.index_select(t, templ_dim, idx).squeeze(templ_dim),
+ batch,
+ )
+
+ # [*, N, N, C_t]
+ t = build_template_pair_feat(
+ single_template_feats,
+ use_unit_vector=self.config.use_unit_vector,
+ inf=self.config.inf,
+ eps=self.config.eps,
+ **self.config.distogram,
+ ).to(z.dtype)
+ t = self.template_pair_embedder(t)
+
+ if (inplace_safe):
+ t_pair[..., i, :, :, :] = t
+ else:
+ pair_embeds.append(t)
+
+ del t
+
+ if (not inplace_safe):
+ t_pair = torch.stack(pair_embeds, dim=templ_dim)
+
+ del pair_embeds
+
+ # [*, S_t, N, N, C_z]
+ t = self.template_pair_stack(
+ t_pair,
+ pair_mask.unsqueeze(-3).to(dtype=z.dtype),
+ chunk_size=chunk_size,
+ use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+ use_lma=use_lma,
+ inplace_safe=inplace_safe,
+ _mask_trans=_mask_trans,
+ )
+ del t_pair
+
+ # [*, N, N, C_z]
+ t = self.template_pointwise_att(
+ t,
+ z,
+ template_mask=batch["template_mask"].to(dtype=z.dtype),
+ use_lma=use_lma,
+ )
+
+ t_mask = torch.sum(batch["template_mask"], dim=-1) > 0
+ # Append singletons
+ t_mask = t_mask.reshape(
+ *t_mask.shape, *([1] * (len(t.shape) - len(t_mask.shape)))
+ )
+
+ if (inplace_safe):
+ t *= t_mask
+ else:
+ t = t * t_mask
+
+ ret = {}
+
+ ret.update({"template_pair_embedding": t})
+
+ del t
+
+ if self.config.embed_angles:
+ template_angle_feat = build_template_angle_feat(
+ batch
+ )
+
+ # [*, S_t, N, C_m]
+ a = self.template_single_embedder(template_angle_feat)
+
+ ret["template_single_embedding"] = a
+
+ return ret
+
+
+class TemplatePairEmbedderMultimer(nn.Module):
+ def __init__(self,
+ c_in: int,
+ c_out: int,
+ c_dgram: int,
+ c_aatype: int,
+ ):
+ super(TemplatePairEmbedderMultimer, self).__init__()
+
+ self.dgram_linear = Linear(c_dgram, c_out, init='relu')
+ self.aatype_linear_1 = Linear(c_aatype, c_out, init='relu')
+ self.aatype_linear_2 = Linear(c_aatype, c_out, init='relu')
+ self.query_embedding_layer_norm = LayerNorm(c_in)
+ self.query_embedding_linear = Linear(c_in, c_out, init='relu')
+
+ self.pseudo_beta_mask_linear = Linear(1, c_out, init='relu')
+ self.x_linear = Linear(1, c_out, init='relu')
+ self.y_linear = Linear(1, c_out, init='relu')
+ self.z_linear = Linear(1, c_out, init='relu')
+ self.backbone_mask_linear = Linear(1, c_out, init='relu')
+
+ def forward(self,
+ template_dgram: torch.Tensor,
+ aatype_one_hot: torch.Tensor,
+ query_embedding: torch.Tensor,
+ pseudo_beta_mask: torch.Tensor,
+ backbone_mask: torch.Tensor,
+ multichain_mask_2d: torch.Tensor,
+ unit_vector: geometry.Vec3Array,
+ ) -> torch.Tensor:
+ act = 0.
+
+ pseudo_beta_mask_2d = (
+ pseudo_beta_mask[..., None] * pseudo_beta_mask[..., None, :]
+ )
+ pseudo_beta_mask_2d *= multichain_mask_2d
+ template_dgram *= pseudo_beta_mask_2d[..., None]
+ act += self.dgram_linear(template_dgram)
+ act += self.pseudo_beta_mask_linear(pseudo_beta_mask_2d[..., None])
+
+ aatype_one_hot = aatype_one_hot.to(template_dgram.dtype)
+ act += self.aatype_linear_1(aatype_one_hot[..., None, :, :])
+ act += self.aatype_linear_2(aatype_one_hot[..., None, :])
+
+ backbone_mask_2d = (
+ backbone_mask[..., None] * backbone_mask[..., None, :]
+ )
+ backbone_mask_2d *= multichain_mask_2d
+ x, y, z = [(coord * backbone_mask_2d).to(dtype=query_embedding.dtype) for coord in unit_vector]
+ act += self.x_linear(x[..., None])
+ act += self.y_linear(y[..., None])
+ act += self.z_linear(z[..., None])
+
+ act += self.backbone_mask_linear(backbone_mask_2d[..., None].to(dtype=query_embedding.dtype))
+
+ query_embedding = self.query_embedding_layer_norm(query_embedding)
+ act += self.query_embedding_linear(query_embedding)
+
+ return act
+
+
+class TemplateSingleEmbedderMultimer(nn.Module):
+ def __init__(self,
+ c_in: int,
+ c_out: int,
+ ):
+ super(TemplateSingleEmbedderMultimer, self).__init__()
+ self.template_single_embedder = Linear(c_in, c_out)
+ self.template_projector = Linear(c_out, c_out)
+
+ def forward(self,
+ batch,
+ atom_pos,
+ aatype_one_hot,
+ ):
+ out = {}
+
+ dtype = batch["template_all_atom_positions"].dtype
+
+ template_chi_angles, template_chi_mask = (
+ all_atom_multimer.compute_chi_angles(
+ atom_pos,
+ batch["template_all_atom_mask"],
+ batch["template_aatype"],
+ )
+ )
+
+ template_features = torch.cat(
+ [
+ aatype_one_hot,
+ torch.sin(template_chi_angles) * template_chi_mask,
+ torch.cos(template_chi_angles) * template_chi_mask,
+ template_chi_mask,
+ ],
+ dim=-1,
+ ).to(dtype=dtype)
+
+ template_mask = template_chi_mask[..., 0].to(dtype=dtype)
+
+ template_activations = self.template_single_embedder(
+ template_features
+ )
+ template_activations = torch.nn.functional.relu(
+ template_activations
+ )
+ template_activations = self.template_projector(
+ template_activations,
+ )
+
+ out["template_single_embedding"] = (
+ template_activations
+ )
+ out["template_mask"] = template_mask
+
+ return out
+
+
+class TemplateEmbedderMultimer(nn.Module):
+ def __init__(self, config):
+ super(TemplateEmbedderMultimer, self).__init__()
+
+ self.config = config
+ self.template_pair_embedder = TemplatePairEmbedderMultimer(
+ **config["template_pair_embedder"],
+ )
+ self.template_single_embedder = TemplateSingleEmbedderMultimer(
+ **config["template_single_embedder"],
+ )
+ self.template_pair_stack = TemplatePairStack(
+ **config["template_pair_stack"],
+ )
+
+ self.linear_t = Linear(config.c_t, config.c_z)
+
+ def forward(self,
+ batch,
+ z,
+ padding_mask_2d,
+ templ_dim,
+ chunk_size,
+ multichain_mask_2d,
+ _mask_trans=True,
+ use_deepspeed_evo_attention=False,
+ use_lma=False,
+ inplace_safe=False
+ ):
+ template_embeds = []
+ n_templ = batch["template_aatype"].shape[templ_dim]
+ for i in range(n_templ):
+ idx = batch["template_aatype"].new_tensor(i)
+ single_template_feats = tensor_tree_map(
+ lambda t: torch.index_select(t, templ_dim, idx),
+ batch,
+ )
+
+ single_template_embeds = {}
+ act = 0.
+
+ template_positions, pseudo_beta_mask = pseudo_beta_fn(
+ single_template_feats["template_aatype"],
+ single_template_feats["template_all_atom_positions"],
+ single_template_feats["template_all_atom_mask"])
+
+ template_dgram = dgram_from_positions(
+ template_positions,
+ inf=self.config.inf,
+ **self.config.distogram,
+ )
+
+ aatype_one_hot = torch.nn.functional.one_hot(
+ single_template_feats["template_aatype"], 22,
+ )
+
+ raw_atom_pos = single_template_feats["template_all_atom_positions"]
+
+ # Vec3Arrays are required to be float32
+ atom_pos = geometry.Vec3Array.from_array(raw_atom_pos.to(dtype=torch.float32))
+
+ rigid, backbone_mask = all_atom_multimer.make_backbone_affine(
+ atom_pos,
+ single_template_feats["template_all_atom_mask"],
+ single_template_feats["template_aatype"],
+ )
+ points = rigid.translation
+ rigid_vec = rigid[..., None].inverse().apply_to_point(points)
+ unit_vector = rigid_vec.normalized()
+
+ pair_act = self.template_pair_embedder(
+ template_dgram,
+ aatype_one_hot,
+ z,
+ pseudo_beta_mask,
+ backbone_mask,
+ multichain_mask_2d,
+ unit_vector,
+ )
+
+ single_template_embeds["template_pair_embedding"] = pair_act
+ single_template_embeds.update(
+ self.template_single_embedder(
+ single_template_feats,
+ atom_pos,
+ aatype_one_hot,
+ )
+ )
+ template_embeds.append(single_template_embeds)
+
+ template_embeds = dict_multimap(
+ partial(torch.cat, dim=templ_dim),
+ template_embeds,
+ )
+
+ # [*, S_t, N, N, C_z]
+ t = self.template_pair_stack(
+ template_embeds["template_pair_embedding"],
+ padding_mask_2d.unsqueeze(-3).to(dtype=z.dtype),
+ chunk_size=chunk_size,
+ use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+ use_lma=use_lma,
+ inplace_safe=inplace_safe,
+ _mask_trans=_mask_trans,
+ )
+ # [*, N, N, C_z]
+ t = torch.sum(t, dim=-4) / n_templ
+ t = torch.nn.functional.relu(t)
+ t = self.linear_t(t)
+ template_embeds["template_pair_embedding"] = t
+
+ return template_embeds
diff --git a/openfold/model/evoformer.py b/openfold/model/evoformer.py
index 8ab3752d..eefbf0bf 100644
--- a/openfold/model/evoformer.py
+++ b/openfold/model/evoformer.py
@@ -18,6 +18,7 @@
import torch.nn as nn
from typing import Tuple, Sequence, Optional
from functools import partial
+from abc import ABC, abstractmethod
from openfold.model.primitives import Linear, LayerNorm
from openfold.model.dropout import DropoutRowwise, DropoutColumnwise
@@ -36,6 +37,8 @@
from openfold.model.triangular_multiplicative_update import (
TriangleMultiplicationOutgoing,
TriangleMultiplicationIncoming,
+ FusedTriangleMultiplicationIncoming,
+ FusedTriangleMultiplicationOutgoing
)
from openfold.utils.checkpointing import checkpoint_blocks, get_checkpoint_fn
from openfold.utils.chunk_utils import chunk_layer, ChunkSizeTuner
@@ -117,43 +120,39 @@ def forward(
return m
-class EvoformerBlockCore(nn.Module):
+class PairStack(nn.Module):
def __init__(
self,
- c_m: int,
c_z: int,
- c_hidden_opm: int,
c_hidden_mul: int,
c_hidden_pair_att: int,
- no_heads_msa: int,
no_heads_pair: int,
transition_n: int,
pair_dropout: float,
+ fuse_projection_weights: bool,
inf: float,
- eps: float,
- _is_extra_msa_stack: bool = False,
+ eps: float
):
- super(EvoformerBlockCore, self).__init__()
-
- self.msa_transition = MSATransition(
- c_m=c_m,
- n=transition_n,
- )
+ super(PairStack, self).__init__()
- self.outer_product_mean = OuterProductMean(
- c_m,
- c_z,
- c_hidden_opm,
- )
-
- self.tri_mul_out = TriangleMultiplicationOutgoing(
- c_z,
- c_hidden_mul,
- )
- self.tri_mul_in = TriangleMultiplicationIncoming(
- c_z,
- c_hidden_mul,
- )
+ if fuse_projection_weights:
+ self.tri_mul_out = FusedTriangleMultiplicationOutgoing(
+ c_z,
+ c_hidden_mul,
+ )
+ self.tri_mul_in = FusedTriangleMultiplicationIncoming(
+ c_z,
+ c_hidden_mul,
+ )
+ else:
+ self.tri_mul_out = TriangleMultiplicationOutgoing(
+ c_z,
+ c_hidden_mul,
+ )
+ self.tri_mul_in = TriangleMultiplicationIncoming(
+ c_z,
+ c_hidden_mul,
+ )
self.tri_att_start = TriangleAttention(
c_z,
@@ -176,56 +175,22 @@ def __init__(
self.ps_dropout_row_layer = DropoutRowwise(pair_dropout)
def forward(self,
- input_tensors: Sequence[torch.Tensor],
- msa_mask: torch.Tensor,
+ z: torch.Tensor,
pair_mask: torch.Tensor,
chunk_size: Optional[int] = None,
use_deepspeed_evo_attention: bool = False,
use_lma: bool = False,
inplace_safe: bool = False,
_mask_trans: bool = True,
- _attn_chunk_size: Optional[int] = None,
- _offload_inference: bool = False,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ _attn_chunk_size: Optional[int] = None
+ ) -> torch.Tensor:
# DeepMind doesn't mask these transitions in the source, so _mask_trans
# should be disabled to better approximate the exact activations of
# the original.
- msa_trans_mask = msa_mask if _mask_trans else None
pair_trans_mask = pair_mask if _mask_trans else None
-
- if(_attn_chunk_size is None):
- _attn_chunk_size = chunk_size
- m, z = input_tensors
-
- m = add(
- m,
- self.msa_transition(
- m, mask=msa_trans_mask, chunk_size=chunk_size,
- ),
- inplace=inplace_safe,
- )
-
- if(_offload_inference and inplace_safe):
- del m, z
- assert(sys.getrefcount(input_tensors[1]) == 2)
- input_tensors[1] = input_tensors[1].cpu()
- torch.cuda.empty_cache()
- m, z = input_tensors
-
- opm = self.outer_product_mean(
- m, mask=msa_mask, chunk_size=chunk_size, inplace_safe=inplace_safe
- )
-
- if(_offload_inference and inplace_safe):
- del m, z
- assert(sys.getrefcount(input_tensors[0]) == 2)
- input_tensors[0] = input_tensors[0].cpu()
- input_tensors[1] = input_tensors[1].to(opm.device)
- m, z = input_tensors
-
- z = add(z, opm, inplace=inplace_safe)
- del opm
+ if (_attn_chunk_size is None):
+ _attn_chunk_size = chunk_size
tmu_update = self.tri_mul_out(
z,
@@ -233,11 +198,11 @@ def forward(self,
inplace_safe=inplace_safe,
_add_with_inplace=True,
)
- if(not inplace_safe):
+ if (not inplace_safe):
z = z + self.ps_dropout_row_layer(tmu_update)
else:
z = tmu_update
-
+
del tmu_update
tmu_update = self.tri_mul_in(
@@ -246,73 +211,63 @@ def forward(self,
inplace_safe=inplace_safe,
_add_with_inplace=True,
)
- if(not inplace_safe):
+ if (not inplace_safe):
z = z + self.ps_dropout_row_layer(tmu_update)
else:
z = tmu_update
-
+
del tmu_update
- z = add(z,
- self.ps_dropout_row_layer(
- self.tri_att_start(
- z,
- mask=pair_mask,
- chunk_size=_attn_chunk_size,
- use_memory_efficient_kernel=False,
- use_deepspeed_evo_attention=use_deepspeed_evo_attention,
- use_lma=use_lma,
- inplace_safe=inplace_safe,
+ z = add(z,
+ self.ps_dropout_row_layer(
+ self.tri_att_start(
+ z,
+ mask=pair_mask,
+ chunk_size=_attn_chunk_size,
+ use_memory_efficient_kernel=False,
+ use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+ use_lma=use_lma,
+ inplace_safe=inplace_safe,
+ )
+ ),
+ inplace=inplace_safe,
)
- ),
- inplace=inplace_safe,
- )
z = z.transpose(-2, -3)
- if(inplace_safe):
- input_tensors[1] = z.contiguous()
- z = input_tensors[1]
+ if (inplace_safe):
+ z = z.contiguous()
z = add(z,
- self.ps_dropout_row_layer(
- self.tri_att_end(
- z,
- mask=pair_mask.transpose(-1, -2),
- chunk_size=_attn_chunk_size,
- use_memory_efficient_kernel=False,
- use_deepspeed_evo_attention=use_deepspeed_evo_attention,
- use_lma=use_lma,
- inplace_safe=inplace_safe,
+ self.ps_dropout_row_layer(
+ self.tri_att_end(
+ z,
+ mask=pair_mask.transpose(-1, -2),
+ chunk_size=_attn_chunk_size,
+ use_memory_efficient_kernel=False,
+ use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+ use_lma=use_lma,
+ inplace_safe=inplace_safe,
+ )
+ ),
+ inplace=inplace_safe,
)
- ),
- inplace=inplace_safe,
- )
z = z.transpose(-2, -3)
- if(inplace_safe):
- input_tensors[1] = z.contiguous()
- z = input_tensors[1]
+ if (inplace_safe):
+ z = z.contiguous()
z = add(z,
- self.pair_transition(
- z, mask=pair_trans_mask, chunk_size=chunk_size,
- ),
- inplace=inplace_safe,
+ self.pair_transition(
+ z, mask=pair_trans_mask, chunk_size=chunk_size,
+ ),
+ inplace=inplace_safe,
)
- if(_offload_inference and inplace_safe):
- device = z.device
- del m, z
- assert(sys.getrefcount(input_tensors[0]) == 2)
- assert(sys.getrefcount(input_tensors[1]) == 2)
- input_tensors[0] = input_tensors[0].to(device)
- input_tensors[1] = input_tensors[1].to(device)
- m, z = input_tensors
-
- return m, z
+ return z
-class EvoformerBlock(nn.Module):
+class MSABlock(nn.Module, ABC):
+ @abstractmethod
def __init__(self,
c_m: int,
c_z: int,
@@ -325,11 +280,14 @@ def __init__(self,
transition_n: int,
msa_dropout: float,
pair_dropout: float,
- no_column_attention: bool,
+ opm_first: bool,
+ fuse_projection_weights: bool,
inf: float,
eps: float,
):
- super(EvoformerBlock, self).__init__()
+ super(MSABlock, self).__init__()
+
+ self.opm_first = opm_first
self.msa_att_row = MSARowAttentionWithPairBias(
c_m=c_m,
@@ -339,32 +297,65 @@ def __init__(self,
inf=inf,
)
- # Specifically, seqemb mode does not use column attention
- self.no_column_attention = no_column_attention
- if not self.no_column_attention:
- self.msa_att_col = MSAColumnAttention(
- c_m,
- c_hidden_msa_att,
- no_heads_msa,
- inf=inf,
- )
-
self.msa_dropout_layer = DropoutRowwise(msa_dropout)
- self.core = EvoformerBlockCore(
+ self.msa_transition = MSATransition(
c_m=c_m,
+ n=transition_n,
+ )
+
+ self.outer_product_mean = OuterProductMean(
+ c_m,
+ c_z,
+ c_hidden_opm,
+ )
+
+ self.pair_stack = PairStack(
c_z=c_z,
- c_hidden_opm=c_hidden_opm,
c_hidden_mul=c_hidden_mul,
c_hidden_pair_att=c_hidden_pair_att,
- no_heads_msa=no_heads_msa,
no_heads_pair=no_heads_pair,
transition_n=transition_n,
pair_dropout=pair_dropout,
+ fuse_projection_weights=fuse_projection_weights,
inf=inf,
- eps=eps,
+ eps=eps
+ )
+
+ def _compute_opm(self,
+ input_tensors: Sequence[torch.Tensor],
+ msa_mask: torch.Tensor,
+ chunk_size: Optional[int] = None,
+ inplace_safe: bool = False,
+ _offload_inference: bool = False
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+ m, z = input_tensors
+
+ if (_offload_inference and inplace_safe):
+ # m: GPU, z: CPU
+ del m, z
+ assert (sys.getrefcount(input_tensors[1]) == 2)
+ input_tensors[1] = input_tensors[1].cpu()
+ m, z = input_tensors
+
+ opm = self.outer_product_mean(
+ m, mask=msa_mask, chunk_size=chunk_size, inplace_safe=inplace_safe
)
+ if (_offload_inference and inplace_safe):
+ # m: GPU, z: GPU
+ del m, z
+ assert (sys.getrefcount(input_tensors[0]) == 2)
+ input_tensors[1] = input_tensors[1].to(opm.device)
+ m, z = input_tensors
+
+ z = add(z, opm, inplace=inplace_safe)
+ del opm
+
+ return m, z
+
+ @abstractmethod
def forward(self,
m: Optional[torch.Tensor],
z: Optional[torch.Tensor],
@@ -380,6 +371,73 @@ def forward(self,
_offload_inference: bool = False,
_offloadable_inputs: Optional[Sequence[torch.Tensor]] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
+ pass
+
+
+class EvoformerBlock(MSABlock):
+ def __init__(self,
+ c_m: int,
+ c_z: int,
+ c_hidden_msa_att: int,
+ c_hidden_opm: int,
+ c_hidden_mul: int,
+ c_hidden_pair_att: int,
+ no_heads_msa: int,
+ no_heads_pair: int,
+ transition_n: int,
+ msa_dropout: float,
+ pair_dropout: float,
+ no_column_attention: bool,
+ opm_first: bool,
+ fuse_projection_weights: bool,
+ inf: float,
+ eps: float,
+ ):
+ super(EvoformerBlock, self).__init__(c_m=c_m,
+ c_z=c_z,
+ c_hidden_msa_att=c_hidden_msa_att,
+ c_hidden_opm=c_hidden_opm,
+ c_hidden_mul=c_hidden_mul,
+ c_hidden_pair_att=c_hidden_pair_att,
+ no_heads_msa=no_heads_msa,
+ no_heads_pair=no_heads_pair,
+ transition_n=transition_n,
+ msa_dropout=msa_dropout,
+ pair_dropout=pair_dropout,
+ opm_first=opm_first,
+ fuse_projection_weights=fuse_projection_weights,
+ inf=inf,
+ eps=eps)
+
+ # Specifically, seqemb mode does not use column attention
+ self.no_column_attention = no_column_attention
+
+ if not self.no_column_attention:
+ self.msa_att_col = MSAColumnAttention(
+ c_m,
+ c_hidden_msa_att,
+ no_heads_msa,
+ inf=inf,
+ )
+
+ def forward(self,
+ m: Optional[torch.Tensor],
+ z: Optional[torch.Tensor],
+ msa_mask: torch.Tensor,
+ pair_mask: torch.Tensor,
+ chunk_size: Optional[int] = None,
+ use_deepspeed_evo_attention: bool = False,
+ use_lma: bool = False,
+ use_flash: bool = False,
+ inplace_safe: bool = False,
+ _mask_trans: bool = True,
+ _attn_chunk_size: Optional[int] = None,
+ _offload_inference: bool = False,
+ _offloadable_inputs: Optional[Sequence[torch.Tensor]] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+ msa_trans_mask = msa_mask if _mask_trans else None
+
if(_attn_chunk_size is None):
_attn_chunk_size = chunk_size
@@ -391,57 +449,110 @@ def forward(self,
m, z = input_tensors
- m = add(m,
- self.msa_dropout_layer(
- self.msa_att_row(
- m,
- z=z,
- mask=msa_mask,
- chunk_size=_attn_chunk_size,
- use_memory_efficient_kernel=False,
- use_deepspeed_evo_attention=use_deepspeed_evo_attention,
- use_lma=use_lma,
+ if self.opm_first:
+ del m, z
+
+ m, z = self._compute_opm(input_tensors=input_tensors,
+ msa_mask=msa_mask,
+ chunk_size=chunk_size,
+ inplace_safe=inplace_safe,
+ _offload_inference=_offload_inference)
+
+ m = add(m,
+ self.msa_dropout_layer(
+ self.msa_att_row(
+ m,
+ z=z,
+ mask=msa_mask,
+ chunk_size=_attn_chunk_size,
+ use_memory_efficient_kernel=False,
+ use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+ use_lma=use_lma,
+ )
+ ),
+ inplace=inplace_safe,
)
- ),
- inplace=inplace_safe,
- )
+
+ if (_offload_inference and inplace_safe):
+ # m: GPU, z: CPU
+ del m, z
+ assert (sys.getrefcount(input_tensors[1]) == 2)
+ input_tensors[1] = input_tensors[1].cpu()
+ torch.cuda.empty_cache()
+ m, z = input_tensors
# Specifically, column attention is not used in seqemb mode.
if not self.no_column_attention:
m = add(m,
- self.msa_att_col(
- m,
- mask=msa_mask,
- chunk_size=chunk_size,
- use_deepspeed_evo_attention=use_deepspeed_evo_attention,
- use_lma=use_lma,
- use_flash=use_flash,
- ),
- inplace=inplace_safe,
- )
+ self.msa_att_col(
+ m,
+ mask=msa_mask,
+ chunk_size=chunk_size,
+ use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+ use_lma=use_lma,
+ use_flash=use_flash,
+ ),
+ inplace=inplace_safe,
+ )
+
+ m = add(
+ m,
+ self.msa_transition(
+ m, mask=msa_trans_mask, chunk_size=chunk_size,
+ ),
+ inplace=inplace_safe,
+ )
+
+ if not self.opm_first:
+ if (not inplace_safe):
+ input_tensors = [m, z]
+
+ del m, z
+
+ m, z = self._compute_opm(input_tensors=input_tensors,
+ msa_mask=msa_mask,
+ chunk_size=chunk_size,
+ inplace_safe=inplace_safe,
+ _offload_inference=_offload_inference)
+
+ if (_offload_inference and inplace_safe):
+ # m: CPU, z: GPU
+ del m, z
+ assert (sys.getrefcount(input_tensors[0]) == 2)
+ device = input_tensors[0].device
+ input_tensors[0] = input_tensors[0].cpu()
+ input_tensors[1] = input_tensors[1].to(device)
+ m, z = input_tensors
+
+ if (not inplace_safe):
+ input_tensors = [m, z]
- if(not inplace_safe):
- input_tensors = [m, input_tensors[1]]
-
del m, z
- m, z = self.core(
- input_tensors,
- msa_mask=msa_mask,
- pair_mask=pair_mask,
+ z = self.pair_stack(
+ z=input_tensors[1],
+ pair_mask=pair_mask,
chunk_size=chunk_size,
use_deepspeed_evo_attention=use_deepspeed_evo_attention,
use_lma=use_lma,
inplace_safe=inplace_safe,
_mask_trans=_mask_trans,
- _attn_chunk_size=_attn_chunk_size,
- _offload_inference=_offload_inference,
+ _attn_chunk_size=_attn_chunk_size
)
+ if (_offload_inference and inplace_safe):
+ # m: GPU, z: GPU
+ device = z.device
+ assert (sys.getrefcount(input_tensors[0]) == 2)
+ input_tensors[0] = input_tensors[0].to(device)
+ m, _ = input_tensors
+ else:
+ m = input_tensors[0]
+
return m, z
-class ExtraMSABlock(nn.Module):
+class ExtraMSABlock(MSABlock):
"""
Almost identical to the standard EvoformerBlock, except in that the
ExtraMSABlock uses GlobalAttention for MSA column attention and
@@ -460,21 +571,29 @@ def __init__(self,
transition_n: int,
msa_dropout: float,
pair_dropout: float,
+ opm_first: bool,
+ fuse_projection_weights: bool,
inf: float,
eps: float,
ckpt: bool,
):
- super(ExtraMSABlock, self).__init__()
-
- self.ckpt = ckpt
+ super(ExtraMSABlock, self).__init__(c_m=c_m,
+ c_z=c_z,
+ c_hidden_msa_att=c_hidden_msa_att,
+ c_hidden_opm=c_hidden_opm,
+ c_hidden_mul=c_hidden_mul,
+ c_hidden_pair_att=c_hidden_pair_att,
+ no_heads_msa=no_heads_msa,
+ no_heads_pair=no_heads_pair,
+ transition_n=transition_n,
+ msa_dropout=msa_dropout,
+ pair_dropout=pair_dropout,
+ opm_first=opm_first,
+ fuse_projection_weights=fuse_projection_weights,
+ inf=inf,
+ eps=eps)
- self.msa_att_row = MSARowAttentionWithPairBias(
- c_m=c_m,
- c_z=c_z,
- c_hidden=c_hidden_msa_att,
- no_heads=no_heads_msa,
- inf=inf,
- )
+ self.ckpt = ckpt
self.msa_att_col = MSAColumnGlobalAttention(
c_in=c_m,
@@ -484,22 +603,6 @@ def __init__(self,
eps=eps,
)
- self.msa_dropout_layer = DropoutRowwise(msa_dropout)
-
- self.core = EvoformerBlockCore(
- c_m=c_m,
- c_z=c_z,
- c_hidden_opm=c_hidden_opm,
- c_hidden_mul=c_hidden_mul,
- c_hidden_pair_att=c_hidden_pair_att,
- no_heads_msa=no_heads_msa,
- no_heads_pair=no_heads_pair,
- transition_n=transition_n,
- pair_dropout=pair_dropout,
- inf=inf,
- eps=eps,
- )
-
def forward(self,
m: Optional[torch.Tensor],
z: Optional[torch.Tensor],
@@ -513,7 +616,7 @@ def forward(self,
_attn_chunk_size: Optional[int] = None,
_offload_inference: bool = False,
_offloadable_inputs: Optional[Sequence[torch.Tensor]] = None,
- ) -> Tuple[torch.Tensor, torch.Tensor]:
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
if(_attn_chunk_size is None):
_attn_chunk_size = chunk_size
@@ -525,7 +628,16 @@ def forward(self,
m, z = input_tensors
- m = add(m,
+ if self.opm_first:
+ del m, z
+
+ m, z = self._compute_opm(input_tensors=input_tensors,
+ msa_mask=msa_mask,
+ chunk_size=chunk_size,
+ inplace_safe=inplace_safe,
+ _offload_inference=_offload_inference)
+
+ m = add(m,
self.msa_dropout_layer(
self.msa_att_row(
m.clone() if torch.is_grad_enabled() else m,
@@ -542,43 +654,89 @@ def forward(self,
inplace=inplace_safe,
)
- if(not inplace_safe):
+ if (not inplace_safe):
input_tensors = [m, z]
del m, z
- def fn(input_tensors):
- m = add(input_tensors[0],
- self.msa_att_col(
- input_tensors[0],
- mask=msa_mask,
- chunk_size=chunk_size,
- use_lma=use_lma,
+ def fn(input_tensors):
+ m, z = input_tensors
+
+ if (_offload_inference and inplace_safe):
+ # m: GPU, z: CPU
+ del m, z
+ assert (sys.getrefcount(input_tensors[1]) == 2)
+ input_tensors[1] = input_tensors[1].cpu()
+ torch.cuda.empty_cache()
+ m, z = input_tensors
+
+ m = add(m,
+ self.msa_att_col(
+ m,
+ mask=msa_mask,
+ chunk_size=chunk_size,
+ use_lma=use_lma,
+ ),
+ inplace=inplace_safe,
+ )
+
+ m = add(
+ m,
+ self.msa_transition(
+ m, mask=msa_mask, chunk_size=chunk_size,
),
inplace=inplace_safe,
)
- if(not inplace_safe):
- input_tensors = [m, input_tensors[1]]
+ if not self.opm_first:
+ if (not inplace_safe):
+ input_tensors = [m, z]
- del m
+ del m, z
- m, z = self.core(
- input_tensors,
- msa_mask=msa_mask,
- pair_mask=pair_mask,
+ m, z = self._compute_opm(input_tensors=input_tensors,
+ msa_mask=msa_mask,
+ chunk_size=chunk_size,
+ inplace_safe=inplace_safe,
+ _offload_inference=_offload_inference)
+
+ if (_offload_inference and inplace_safe):
+ # m: CPU, z: GPU
+ del m, z
+ assert (sys.getrefcount(input_tensors[0]) == 2)
+ device = input_tensors[0].device
+ input_tensors[0] = input_tensors[0].cpu()
+ input_tensors[1] = input_tensors[1].to(device)
+ m, z = input_tensors
+
+ if (not inplace_safe):
+ input_tensors = [m, z]
+
+ del m, z
+
+ z = self.pair_stack(
+ input_tensors[1],
+ pair_mask=pair_mask,
chunk_size=chunk_size,
use_deepspeed_evo_attention=use_deepspeed_evo_attention,
use_lma=use_lma,
inplace_safe=inplace_safe,
_mask_trans=_mask_trans,
- _attn_chunk_size=_attn_chunk_size,
- _offload_inference=_offload_inference,
+ _attn_chunk_size=_attn_chunk_size
)
-
+
+ m = input_tensors[0]
+ if (_offload_inference and inplace_safe):
+ # m: GPU, z: GPU
+ device = z.device
+ del m
+ assert (sys.getrefcount(input_tensors[0]) == 2)
+ input_tensors[0] = input_tensors[0].to(device)
+ m, _ = input_tensors
+
return m, z
- if(torch.is_grad_enabled() and self.ckpt):
+ if (torch.is_grad_enabled() and self.ckpt):
checkpoint_fn = get_checkpoint_fn()
m, z = checkpoint_fn(fn, input_tensors)
else:
@@ -609,8 +767,10 @@ def __init__(
transition_n: int,
msa_dropout: float,
pair_dropout: float,
- blocks_per_ckpt: int,
no_column_attention: bool,
+ opm_first: bool,
+ fuse_projection_weights: bool,
+ blocks_per_ckpt: int,
inf: float,
eps: float,
clear_cache_between_blocks: bool = False,
@@ -646,11 +806,18 @@ def __init__(
Dropout rate for MSA activations
pair_dropout:
Dropout used for pair activations
- blocks_per_ckpt:
- Number of Evoformer blocks in each activation checkpoint
no_column_attention:
When True, doesn't use column attention. Required for running
sequence embedding mode
+ opm_first:
+ When True, Outer Product Mean is performed at the beginning of
+ the Evoformer block instead of after the MSA Stack.
+ Used in Multimer pipeline.
+ fuse_projection_weights:
+ When True, uses FusedTriangleMultiplicativeUpdate variant in
+ the Pair Stack. Used in Multimer pipeline.
+ blocks_per_ckpt:
+ Number of Evoformer blocks in each activation checkpoint
clear_cache_between_blocks:
Whether to clear CUDA's GPU memory cache between blocks of the
stack. Slows down each block but can reduce fragmentation
@@ -678,6 +845,8 @@ def __init__(
msa_dropout=msa_dropout,
pair_dropout=pair_dropout,
no_column_attention=no_column_attention,
+ opm_first=opm_first,
+ fuse_projection_weights=fuse_projection_weights,
inf=inf,
eps=eps,
)
@@ -873,6 +1042,8 @@ def __init__(self,
transition_n: int,
msa_dropout: float,
pair_dropout: float,
+ opm_first: bool,
+ fuse_projection_weights: bool,
inf: float,
eps: float,
ckpt: bool,
@@ -898,6 +1069,8 @@ def __init__(self,
transition_n=transition_n,
msa_dropout=msa_dropout,
pair_dropout=pair_dropout,
+ opm_first=opm_first,
+ fuse_projection_weights=fuse_projection_weights,
inf=inf,
eps=eps,
ckpt=False,
diff --git a/openfold/model/heads.py b/openfold/model/heads.py
index e08addd0..459dce93 100644
--- a/openfold/model/heads.py
+++ b/openfold/model/heads.py
@@ -76,9 +76,17 @@ def forward(self, outputs):
if self.config.tm.enabled:
tm_logits = self.tm(outputs["pair"])
aux_out["tm_logits"] = tm_logits
- aux_out["predicted_tm_score"] = compute_tm(
+ aux_out["ptm_score"] = compute_tm(
tm_logits, **self.config.tm
)
+ asym_id = outputs.get("asym_id")
+ if asym_id is not None:
+ aux_out["iptm_score"] = compute_tm(
+ tm_logits, asym_id=asym_id, interface=True, **self.config.tm
+ )
+ aux_out["weighted_ptm_score"] = (self.config.tm["iptm_weight"] * aux_out["iptm_score"]
+ + self.config.tm["ptm_weight"] * aux_out["ptm_score"])
+
aux_out.update(
compute_predicted_aligned_error(
tm_logits,
diff --git a/openfold/model/model.py b/openfold/model/model.py
index bfe47f5c..52cfda45 100644
--- a/openfold/model/model.py
+++ b/openfold/model/model.py
@@ -18,11 +18,20 @@
import torch
import torch.nn as nn
+from openfold.data import data_transforms_multimer
+from openfold.utils.feats import (
+ pseudo_beta_fn,
+ build_extra_msa_feat,
+ dgram_from_positions,
+ atom14_to_atom37,
+)
+from openfold.utils.tensor_utils import masked_mean
from openfold.model.embedders import (
InputEmbedder,
+ InputEmbedderMultimer,
RecyclingEmbedder,
- TemplateAngleEmbedder,
- TemplatePairEmbedder,
+ TemplateEmbedder,
+ TemplateEmbedderMultimer,
ExtraMSAEmbedder,
PreembeddingEmbedder,
)
@@ -75,9 +84,13 @@ def __init__(self, config):
self.seqemb_mode = config.globals.seqemb_mode_enabled
# Main trunk + structure module
- # If using seqemb mode, embed the sequence embeddings passed
- # to the model ("preembeddings") instead of embedding the sequence
- if self.seqemb_mode:
+ if self.globals.is_multimer:
+ self.input_embedder = InputEmbedderMultimer(
+ **self.config["input_embedder"]
+ )
+ elif self.seqemb_mode:
+ # If using seqemb mode, embed the sequence embeddings passed
+ # to the model ("preembeddings") instead of embedding the sequence
self.input_embedder = PreembeddingEmbedder(
**self.config["preembedding_embedder"],
)
@@ -85,142 +98,113 @@ def __init__(self, config):
self.input_embedder = InputEmbedder(
**self.config["input_embedder"],
)
+
self.recycling_embedder = RecyclingEmbedder(
**self.config["recycling_embedder"],
)
-
- if(self.template_config.enabled):
- self.template_angle_embedder = TemplateAngleEmbedder(
- **self.template_config["template_angle_embedder"],
- )
- self.template_pair_embedder = TemplatePairEmbedder(
- **self.template_config["template_pair_embedder"],
- )
- self.template_pair_stack = TemplatePairStack(
- **self.template_config["template_pair_stack"],
- )
- self.template_pointwise_att = TemplatePointwiseAttention(
- **self.template_config["template_pointwise_attention"],
- )
-
- if(self.extra_msa_config.enabled):
+
+ if self.template_config.enabled:
+ if self.globals.is_multimer:
+ self.template_embedder = TemplateEmbedderMultimer(
+ self.template_config,
+ )
+ else:
+ self.template_embedder = TemplateEmbedder(
+ self.template_config,
+ )
+
+ if self.extra_msa_config.enabled:
self.extra_msa_embedder = ExtraMSAEmbedder(
**self.extra_msa_config["extra_msa_embedder"],
)
self.extra_msa_stack = ExtraMSAStack(
**self.extra_msa_config["extra_msa_stack"],
)
-
+
self.evoformer = EvoformerStack(
**self.config["evoformer_stack"],
)
+
self.structure_module = StructureModule(
+ is_multimer=self.globals.is_multimer,
**self.config["structure_module"],
)
self.aux_heads = AuxiliaryHeads(
self.config["heads"],
)
- def embed_templates(self, batch, z, pair_mask, templ_dim, inplace_safe):
- if(self.template_config.offload_templates):
- return embed_templates_offload(self,
- batch, z, pair_mask, templ_dim, inplace_safe=inplace_safe,
+ def embed_templates(self, batch, feats, z, pair_mask, templ_dim, inplace_safe):
+ if self.globals.is_multimer:
+ asym_id = feats["asym_id"]
+ multichain_mask_2d = (
+ asym_id[..., None] == asym_id[..., None, :]
)
- elif(self.template_config.average_templates):
- return embed_templates_average(self,
- batch, z, pair_mask, templ_dim, inplace_safe=inplace_safe,
+ template_embeds = self.template_embedder(
+ batch,
+ z,
+ pair_mask.to(dtype=z.dtype),
+ templ_dim,
+ chunk_size=self.globals.chunk_size,
+ multichain_mask_2d=multichain_mask_2d,
+ use_deepspeed_evo_attention=self.globals.use_deepspeed_evo_attention,
+ use_lma=self.globals.use_lma,
+ inplace_safe=inplace_safe,
+ _mask_trans=self.config._mask_trans
)
-
- # Embed the templates one at a time (with a poor man's vmap)
- pair_embeds = []
- n = z.shape[-2]
- n_templ = batch["template_aatype"].shape[templ_dim]
-
- if(inplace_safe):
- # We'll preallocate the full pair tensor now to avoid manifesting
- # a second copy during the stack later on
- t_pair = z.new_zeros(
- z.shape[:-3] +
- (n_templ, n, n, self.globals.c_t)
+ feats["template_torsion_angles_mask"] = (
+ template_embeds["template_mask"]
)
-
- for i in range(n_templ):
- idx = batch["template_aatype"].new_tensor(i)
- single_template_feats = tensor_tree_map(
- lambda t: torch.index_select(t, templ_dim, idx).squeeze(templ_dim),
+ else:
+ if self.template_config.offload_templates:
+ return embed_templates_offload(self,
+ batch, z, pair_mask, templ_dim, inplace_safe=inplace_safe,
+ )
+ elif self.template_config.average_templates:
+ return embed_templates_average(self,
+ batch, z, pair_mask, templ_dim, inplace_safe=inplace_safe,
+ )
+
+ template_embeds = self.template_embedder(
batch,
+ z,
+ pair_mask.to(dtype=z.dtype),
+ templ_dim,
+ chunk_size=self.globals.chunk_size,
+ use_deepspeed_evo_attention=self.globals.use_deepspeed_evo_attention,
+ use_lma=self.globals.use_lma,
+ inplace_safe=inplace_safe,
+ _mask_trans=self.config._mask_trans
)
- # [*, N, N, C_t]
- t = build_template_pair_feat(
- single_template_feats,
- use_unit_vector=self.config.template.use_unit_vector,
- inf=self.config.template.inf,
- eps=self.config.template.eps,
- **self.config.template.distogram,
- ).to(z.dtype)
- t = self.template_pair_embedder(t)
-
- if(inplace_safe):
- t_pair[..., i, :, :, :] = t
- else:
- pair_embeds.append(t)
-
- del t
-
- if(not inplace_safe):
- t_pair = torch.stack(pair_embeds, dim=templ_dim)
-
- del pair_embeds
-
- # [*, S_t, N, N, C_z]
- t = self.template_pair_stack(
- t_pair,
- pair_mask.unsqueeze(-3).to(dtype=z.dtype),
- chunk_size=self.globals.chunk_size,
- use_deepspeed_evo_attention=self.globals.use_deepspeed_evo_attention,
- use_lma=self.globals.use_lma,
- inplace_safe=inplace_safe,
- _mask_trans=self.config._mask_trans,
- )
- del t_pair
+ return template_embeds
- # [*, N, N, C_z]
- t = self.template_pointwise_att(
- t,
- z,
- template_mask=batch["template_mask"].to(dtype=z.dtype),
- use_lma=self.globals.use_lma,
- )
-
- t_mask = torch.sum(batch["template_mask"], dim=-1) > 0
- # Append singletons
- t_mask = t_mask.reshape(
- *t_mask.shape, *([1] * (len(t.shape) - len(t_mask.shape)))
- )
-
- if(inplace_safe):
- t *= t_mask
- else:
- t = t * t_mask
-
- ret = {}
-
- ret.update({"template_pair_embedding": t})
-
- del t
-
- if self.config.template.embed_angles:
- template_angle_feat = build_template_angle_feat(
- batch
- )
+ def tolerance_reached(self, prev_pos, next_pos, mask, eps=1e-8) -> bool:
+ """
+ Early stopping criteria based on criteria used in
+ AF2Complex: https://www.nature.com/articles/s41467-022-29394-2
+ Args:
+ prev_pos: Previous atom positions in atom37/14 representation
+ next_pos: Current atom positions in atom37/14 representation
+ mask: 1-D sequence mask
+ eps: Epsilon used in square root calculation
+ Returns:
+ Whether to stop recycling early based on the desired tolerance.
+ """
- # [*, S_t, N, C_m]
- a = self.template_angle_embedder(template_angle_feat)
+ def distances(points):
+ """Compute all pairwise distances for a set of points."""
+ d = points[..., None, :] - points[..., None, :, :]
+ return torch.sqrt(torch.sum(d ** 2, dim=-1))
- ret["template_angle_embedding"] = a
+ if self.config.recycle_early_stop_tolerance < 0:
+ return False
- return ret
+ ca_idx = residue_constants.atom_order['CA']
+ sq_diff = (distances(prev_pos[..., ca_idx, :]) - distances(next_pos[..., ca_idx, :])) ** 2
+ mask = mask[..., None] * mask[..., None, :]
+ sq_diff = masked_mean(mask=mask, value=sq_diff, dim=list(range(len(mask.shape))))
+ diff = torch.sqrt(sq_diff + eps).item()
+ return diff <= self.config.recycle_early_stop_tolerance
def iteration(self, feats, prevs, _recycle=True):
# Primary output dictionary
@@ -229,7 +213,7 @@ def iteration(self, feats, prevs, _recycle=True):
# This needs to be done manually for DeepSpeed's sake
dtype = next(self.parameters()).dtype
for k in feats:
- if(feats[k].dtype == torch.float32):
+ if feats[k].dtype == torch.float32:
feats[k] = feats[k].to(dtype=dtype)
# Grab some data about the input
@@ -238,7 +222,7 @@ def iteration(self, feats, prevs, _recycle=True):
n = feats["target_feat"].shape[-2]
n_seq = feats["msa_feat"].shape[-3]
device = feats["target_feat"].device
-
+
# Controls whether the model uses in-place operations throughout
# The dual condition accounts for activation checkpoints
inplace_safe = not (self.training or torch.is_grad_enabled())
@@ -248,18 +232,22 @@ def iteration(self, feats, prevs, _recycle=True):
pair_mask = seq_mask[..., None] * seq_mask[..., None, :]
msa_mask = feats["msa_mask"]
- ## Initialize the SingleSeq and pair representations
- # m: [*, 1, N, C_m]
- # z: [*, N, N, C_z]
- if self.seqemb_mode:
+ if self.globals.is_multimer:
+ # Initialize the MSA and pair representations
+ # m: [*, S_c, N, C_m]
+ # z: [*, N, N, C_z]
+ m, z = self.input_embedder(feats)
+ elif self.seqemb_mode:
+ # Initialize the SingleSeq and pair representations
+ # m: [*, 1, N, C_m]
+ # z: [*, N, N, C_z]
m, z = self.input_embedder(
feats["target_feat"],
feats["residue_index"],
feats["seq_embedding"]
)
-
else:
- ## Initialize the MSA and pair representations
+ # Initialize the MSA and pair representations
# m: [*, S_c, N, C_m]
# z: [*, N, N, C_z]
m, z = self.input_embedder(
@@ -293,12 +281,12 @@ def iteration(self, feats, prevs, _recycle=True):
requires_grad=False,
)
- x_prev = pseudo_beta_fn(
+ pseudo_beta_x_prev = pseudo_beta_fn(
feats["aatype"], x_prev, None
).to(dtype=z.dtype)
# The recycling embedder is memory-intensive, so we offload first
- if(self.globals.offload_inference and inplace_safe):
+ if self.globals.offload_inference and inplace_safe:
m = m.cpu()
z = z.cpu()
@@ -307,11 +295,13 @@ def iteration(self, feats, prevs, _recycle=True):
m_1_prev_emb, z_prev_emb = self.recycling_embedder(
m_1_prev,
z_prev,
- x_prev,
+ pseudo_beta_x_prev,
inplace_safe=inplace_safe,
)
- if(self.globals.offload_inference and inplace_safe):
+ del pseudo_beta_x_prev
+
+ if self.globals.offload_inference and inplace_safe:
m = m.to(m_1_prev_emb.device)
z = z.to(z_prev.device)
@@ -324,15 +314,17 @@ def iteration(self, feats, prevs, _recycle=True):
# Deletions like these become significant for inference with large N,
# where they free unused tensors and remove references to others such
# that they can be offloaded later
- del m_1_prev, z_prev, x_prev, m_1_prev_emb, z_prev_emb
+ del m_1_prev, z_prev, m_1_prev_emb, z_prev_emb
# Embed the templates + merge with MSA/pair embeddings
- if self.config.template.enabled:
+ if self.config.template.enabled:
template_feats = {
k: v for k, v in feats.items() if k.startswith("template_")
}
+
template_embeds = self.embed_templates(
template_feats,
+ feats,
z,
pair_mask.to(dtype=z.dtype),
no_batch_dims,
@@ -341,35 +333,49 @@ def iteration(self, feats, prevs, _recycle=True):
# [*, N, N, C_z]
z = add(z,
- template_embeds.pop("template_pair_embedding"),
- inplace_safe,
- )
+ template_embeds.pop("template_pair_embedding"),
+ inplace_safe,
+ )
- if "template_angle_embedding" in template_embeds:
+ if (
+ "template_single_embedding" in template_embeds
+ ):
# [*, S = S_c + S_t, N, C_m]
m = torch.cat(
- [m, template_embeds["template_angle_embedding"]],
+ [m, template_embeds["template_single_embedding"]],
dim=-3
)
# [*, S, N]
- torsion_angles_mask = feats["template_torsion_angles_mask"]
- msa_mask = torch.cat(
- [feats["msa_mask"], torsion_angles_mask[..., 2]],
- dim=-2
- )
+ if not self.globals.is_multimer:
+ torsion_angles_mask = feats["template_torsion_angles_mask"]
+ msa_mask = torch.cat(
+ [feats["msa_mask"], torsion_angles_mask[..., 2]],
+ dim=-2
+ )
+ else:
+ msa_mask = torch.cat(
+ [feats["msa_mask"], template_embeds["template_mask"]],
+ dim=-2,
+ )
# Embed extra MSA features + merge with pairwise embeddings
if self.config.extra_msa.enabled:
+ if self.globals.is_multimer:
+ extra_msa_fn = data_transforms_multimer.build_extra_msa_feat
+ else:
+ extra_msa_fn = build_extra_msa_feat
+
# [*, S_e, N, C_e]
- a = self.extra_msa_embedder(build_extra_msa_feat(feats))
+ extra_msa_feat = extra_msa_fn(feats).to(dtype=z.dtype)
+ a = self.extra_msa_embedder(extra_msa_feat)
- if(self.globals.offload_inference):
+ if self.globals.offload_inference:
# To allow the extra MSA stack (and later the evoformer) to
# offload its inputs, we remove all references to them here
input_tensors = [a, z]
del a, z
-
+
# [*, N, N, C_z]
z = self.extra_msa_stack._forward_offload(
input_tensors,
@@ -380,7 +386,7 @@ def iteration(self, feats, prevs, _recycle=True):
pair_mask=pair_mask.to(dtype=m.dtype),
_mask_trans=self.config._mask_trans,
)
-
+
del input_tensors
else:
# [*, N, N, C_z]
@@ -399,7 +405,7 @@ def iteration(self, feats, prevs, _recycle=True):
# m: [*, S, N, C_m]
# z: [*, N, N, C_z]
# s: [*, N, C_s]
- if(self.globals.offload_inference):
+ if self.globals.offload_inference:
input_tensors = [m, z]
del m, z
m, z, s = self.evoformer._forward_offload(
@@ -411,7 +417,7 @@ def iteration(self, feats, prevs, _recycle=True):
use_lma=self.globals.use_lma,
_mask_trans=self.config._mask_trans,
)
-
+
del input_tensors
else:
m, z, s = self.evoformer(
@@ -455,10 +461,34 @@ def iteration(self, feats, prevs, _recycle=True):
# [*, N, N, C_z]
z_prev = outputs["pair"]
+ early_stop = False
+ if self.globals.is_multimer:
+ early_stop = self.tolerance_reached(x_prev, outputs["final_atom_positions"], seq_mask)
+
+ del x_prev
+
# [*, N, 3]
x_prev = outputs["final_atom_positions"]
- return outputs, m_1_prev, z_prev, x_prev
+ return outputs, m_1_prev, z_prev, x_prev, early_stop
+
+ def _disable_activation_checkpointing(self):
+ self.template_embedder.template_pair_stack.blocks_per_ckpt = None
+ self.evoformer.blocks_per_ckpt = None
+
+ for b in self.extra_msa_stack.blocks:
+ b.ckpt = False
+
+ def _enable_activation_checkpointing(self):
+ self.template_embedder.template_pair_stack.blocks_per_ckpt = (
+ self.config.template.template_pair_stack.blocks_per_ckpt
+ )
+ self.evoformer.blocks_per_ckpt = (
+ self.config.evoformer_stack.blocks_per_ckpt
+ )
+
+ for b in self.extra_msa_stack.blocks:
+ b.ckpt = self.config.extra_msa.extra_msa_stack.ckpt
def forward(self, batch):
"""
@@ -519,13 +549,15 @@ def forward(self, batch):
# Main recycling loop
num_iters = batch["aatype"].shape[-1]
- for cycle_no in range(num_iters):
+ early_stop = False
+ num_recycles = 0
+ for cycle_no in range(num_iters):
# Select the features for the current recycling cycle
fetch_cur_batch = lambda t: t[..., cycle_no]
feats = tensor_tree_map(fetch_cur_batch, batch)
# Enable grad iff we're training and it's the final recycling layer
- is_final_iter = cycle_no == (num_iters - 1)
+ is_final_iter = cycle_no == (num_iters - 1) or early_stop
with torch.set_grad_enabled(is_grad_enabled and is_final_iter):
if is_final_iter:
# Sidestep AMP bug (PyTorch issue #65766)
@@ -533,16 +565,25 @@ def forward(self, batch):
torch.clear_autocast_cache()
# Run the next iteration of the model
- outputs, m_1_prev, z_prev, x_prev = self.iteration(
+ outputs, m_1_prev, z_prev, x_prev, early_stop = self.iteration(
feats,
prevs,
_recycle=(num_iters > 1)
)
- if(not is_final_iter):
+ num_recycles += 1
+
+ if not is_final_iter:
del outputs
prevs = [m_1_prev, z_prev, x_prev]
del m_1_prev, z_prev, x_prev
+ else:
+ break
+
+ outputs["num_recycles"] = torch.tensor(num_recycles, device=feats["aatype"].device)
+
+ if "asym_id" in batch:
+ outputs["asym_id"] = feats["asym_id"]
# Run auxiliary heads
outputs.update(self.aux_heads(outputs))
diff --git a/openfold/model/primitives.py b/openfold/model/primitives.py
index d18e3af8..ea38cb34 100644
--- a/openfold/model/primitives.py
+++ b/openfold/model/primitives.py
@@ -131,6 +131,7 @@ def __init__(
bias: bool = True,
init: str = "default",
init_fn: Optional[Callable[[torch.Tensor, torch.Tensor], None]] = None,
+ precision=None
):
"""
Args:
@@ -182,6 +183,28 @@ def __init__(
else:
raise ValueError("Invalid init string.")
+ self.precision = precision
+
+ def forward(self, input: torch.Tensor) -> torch.Tensor:
+ d = input.dtype
+ deepspeed_is_initialized = (
+ deepspeed_is_installed and
+ deepspeed.comm.comm.is_initialized()
+ )
+ if self.precision is not None:
+ with torch.cuda.amp.autocast(enabled=False):
+ bias = self.bias.to(dtype=self.precision) if self.bias is not None else None
+ return nn.functional.linear(input.to(dtype=self.precision),
+ self.weight.to(dtype=self.precision),
+ bias).to(dtype=d)
+
+ if d is torch.bfloat16 and not deepspeed_is_initialized:
+ with torch.cuda.amp.autocast(enabled=False):
+ bias = self.bias.to(dtype=d) if self.bias is not None else None
+ return nn.functional.linear(input, self.weight.to(dtype=d), bias)
+
+ return nn.functional.linear(input, self.weight, self.bias)
+
class LayerNorm(nn.Module):
def __init__(self, c_in, eps=1e-5):
diff --git a/openfold/model/structure_module.py b/openfold/model/structure_module.py
index 2136554e..2bd4c162 100644
--- a/openfold/model/structure_module.py
+++ b/openfold/model/structure_module.py
@@ -20,7 +20,7 @@
import torch
import torch.nn as nn
-from typing import Optional, Tuple, Sequence
+from typing import Optional, Tuple, Sequence, Union
from openfold.model.primitives import Linear, LayerNorm, ipa_point_weights_init_
from openfold.np.residue_constants import (
@@ -29,6 +29,9 @@
restype_atom14_mask,
restype_atom14_rigid_group_positions,
)
+from openfold.utils.geometry.quat_rigid import QuatRigid
+from openfold.utils.geometry.rigid_matrix_vector import Rigid3Array
+from openfold.utils.geometry.vector import Vec3Array, square_euclidean_distance
from openfold.utils.feats import (
frames_and_literature_positions_to_atom14_pos,
torsion_angles_to_frames,
@@ -158,6 +161,51 @@ def forward(
return unnormalized_s, s
+class PointProjection(nn.Module):
+ def __init__(self,
+ c_hidden: int,
+ num_points: int,
+ no_heads: int,
+ is_multimer: bool,
+ return_local_points: bool = False,
+ ):
+ super().__init__()
+ self.return_local_points = return_local_points
+ self.no_heads = no_heads
+ self.num_points = num_points
+ self.is_multimer = is_multimer
+
+ # Multimer requires this to be run with fp32 precision during training
+ precision = torch.float32 if self.is_multimer else None
+ self.linear = Linear(c_hidden, no_heads * 3 * num_points, precision=precision)
+
+ def forward(self,
+ activations: torch.Tensor,
+ rigids: Union[Rigid, Rigid3Array],
+ ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+ # TODO: Needs to run in high precision during training
+ points_local = self.linear(activations)
+ out_shape = points_local.shape[:-1] + (self.no_heads, self.num_points, 3)
+
+ if self.is_multimer:
+ points_local = points_local.view(
+ points_local.shape[:-1] + (self.no_heads, -1)
+ )
+
+ points_local = torch.split(
+ points_local, points_local.shape[-1] // 3, dim=-1
+ )
+
+ points_local = torch.stack(points_local, dim=-1).view(out_shape)
+
+ points_global = rigids[..., None, None].apply(points_local)
+
+ if(self.return_local_points):
+ return points_global, points_local
+
+ return points_global
+
+
class InvariantPointAttention(nn.Module):
"""
Implements Algorithm 22.
@@ -172,6 +220,7 @@ def __init__(
no_v_points: int,
inf: float = 1e5,
eps: float = 1e-8,
+ is_multimer: bool = False,
):
"""
Args:
@@ -198,22 +247,46 @@ def __init__(
self.no_v_points = no_v_points
self.inf = inf
self.eps = eps
+ self.is_multimer = is_multimer
# These linear layers differ from their specifications in the
# supplement. There, they lack bias and use Glorot initialization.
# Here as in the official source, they have bias and use the default
# Lecun initialization.
hc = self.c_hidden * self.no_heads
- self.linear_q = Linear(self.c_s, hc)
- self.linear_kv = Linear(self.c_s, 2 * hc)
+ self.linear_q = Linear(self.c_s, hc, bias=(not is_multimer))
- hpq = self.no_heads * self.no_qk_points * 3
- self.linear_q_points = Linear(self.c_s, hpq)
+ self.linear_q_points = PointProjection(
+ self.c_s,
+ self.no_qk_points,
+ self.no_heads,
+ self.is_multimer
+ )
- hpkv = self.no_heads * (self.no_qk_points + self.no_v_points) * 3
- self.linear_kv_points = Linear(self.c_s, hpkv)
+ if(is_multimer):
+ self.linear_k = Linear(self.c_s, hc, bias=False)
+ self.linear_v = Linear(self.c_s, hc, bias=False)
+ self.linear_k_points = PointProjection(
+ self.c_s,
+ self.no_qk_points,
+ self.no_heads,
+ self.is_multimer
+ )
- hpv = self.no_heads * self.no_v_points * 3
+ self.linear_v_points = PointProjection(
+ self.c_s,
+ self.no_v_points,
+ self.no_heads,
+ self.is_multimer
+ )
+ else:
+ self.linear_kv = Linear(self.c_s, 2 * hc)
+ self.linear_kv_points = PointProjection(
+ self.c_s,
+ self.no_qk_points + self.no_v_points,
+ self.no_heads,
+ self.is_multimer
+ )
self.linear_b = Linear(self.c_z, self.no_heads)
@@ -231,8 +304,8 @@ def __init__(
def forward(
self,
s: torch.Tensor,
- z: Optional[torch.Tensor],
- r: Rigid,
+ z: torch.Tensor,
+ r: Union[Rigid, Rigid3Array],
mask: torch.Tensor,
inplace_safe: bool = False,
_offload_inference: bool = False,
@@ -251,69 +324,68 @@ def forward(
Returns:
[*, N_res, C_s] single representation update
"""
- if(_offload_inference and inplace_safe):
+ if (_offload_inference and inplace_safe):
z = _z_reference_list
else:
z = [z]
-
+
#######################################
# Generate scalar and point activations
#######################################
# [*, N_res, H * C_hidden]
q = self.linear_q(s)
- kv = self.linear_kv(s)
# [*, N_res, H, C_hidden]
q = q.view(q.shape[:-1] + (self.no_heads, -1))
- # [*, N_res, H, 2 * C_hidden]
- kv = kv.view(kv.shape[:-1] + (self.no_heads, -1))
+ # [*, N_res, H, P_qk]
+ q_pts = self.linear_q_points(s, r)
- # [*, N_res, H, C_hidden]
- k, v = torch.split(kv, self.c_hidden, dim=-1)
+ # The following two blocks are equivalent
+ # They're separated only to preserve compatibility with old AF weights
+ if(self.is_multimer):
+ # [*, N_res, H * C_hidden]
+ k = self.linear_k(s)
+ v = self.linear_v(s)
- # [*, N_res, H * P_q * 3]
- q_pts = self.linear_q_points(s)
+ # [*, N_res, H, C_hidden]
+ k = k.view(k.shape[:-1] + (self.no_heads, -1))
+ v = v.view(v.shape[:-1] + (self.no_heads, -1))
- # This is kind of clunky, but it's how the original does it
- # [*, N_res, H * P_q, 3]
- q_pts = torch.split(q_pts, q_pts.shape[-1] // 3, dim=-1)
- q_pts = torch.stack(q_pts, dim=-1)
- q_pts = r[..., None].apply(q_pts)
+ # [*, N_res, H, P_qk, 3]
+ k_pts = self.linear_k_points(s, r)
- # [*, N_res, H, P_q, 3]
- q_pts = q_pts.view(
- q_pts.shape[:-2] + (self.no_heads, self.no_qk_points, 3)
- )
+ # [*, N_res, H, P_v, 3]
+ v_pts = self.linear_v_points(s, r)
+ else:
+ # [*, N_res, H * 2 * C_hidden]
+ kv = self.linear_kv(s)
- # [*, N_res, H * (P_q + P_v) * 3]
- kv_pts = self.linear_kv_points(s)
+ # [*, N_res, H, 2 * C_hidden]
+ kv = kv.view(kv.shape[:-1] + (self.no_heads, -1))
- # [*, N_res, H * (P_q + P_v), 3]
- kv_pts = torch.split(kv_pts, kv_pts.shape[-1] // 3, dim=-1)
- kv_pts = torch.stack(kv_pts, dim=-1)
- kv_pts = r[..., None].apply(kv_pts)
+ # [*, N_res, H, C_hidden]
+ k, v = torch.split(kv, self.c_hidden, dim=-1)
- # [*, N_res, H, (P_q + P_v), 3]
- kv_pts = kv_pts.view(kv_pts.shape[:-2] + (self.no_heads, -1, 3))
+ kv_pts = self.linear_kv_points(s, r)
- # [*, N_res, H, P_q/P_v, 3]
- k_pts, v_pts = torch.split(
- kv_pts, [self.no_qk_points, self.no_v_points], dim=-2
- )
+ # [*, N_res, H, P_q/P_v, 3]
+ k_pts, v_pts = torch.split(
+ kv_pts, [self.no_qk_points, self.no_v_points], dim=-2
+ )
##########################
# Compute attention scores
##########################
# [*, N_res, N_res, H]
b = self.linear_b(z[0])
-
- if(_offload_inference):
- assert(sys.getrefcount(z[0]) == 2)
+
+ if (_offload_inference):
+ assert (sys.getrefcount(z[0]) == 2)
z[0] = z[0].cpu()
# [*, H, N_res, N_res]
- if(is_fp16_enabled()):
+ if (is_fp16_enabled()):
with torch.cuda.amp.autocast(enabled=False):
a = torch.matmul(
permute_final_dims(q.float(), (1, 0, 2)), # [*, H, N_res, C_hidden]
@@ -324,40 +396,43 @@ def forward(
permute_final_dims(q, (1, 0, 2)), # [*, H, N_res, C_hidden]
permute_final_dims(k, (1, 2, 0)), # [*, H, C_hidden, N_res]
)
-
+
a *= math.sqrt(1.0 / (3 * self.c_hidden))
a += (math.sqrt(1.0 / 3) * permute_final_dims(b, (2, 0, 1)))
# [*, N_res, N_res, H, P_q, 3]
pt_att = q_pts.unsqueeze(-4) - k_pts.unsqueeze(-5)
- if(inplace_safe):
+
+ if (inplace_safe):
pt_att *= pt_att
else:
pt_att = pt_att ** 2
- # [*, N_res, N_res, H, P_q]
pt_att = sum(torch.unbind(pt_att, dim=-1))
+
head_weights = self.softplus(self.head_weights).view(
*((1,) * len(pt_att.shape[:-2]) + (-1, 1))
)
head_weights = head_weights * math.sqrt(
1.0 / (3 * (self.no_qk_points * 9.0 / 2))
)
- if(inplace_safe):
+
+ if (inplace_safe):
pt_att *= head_weights
else:
pt_att = pt_att * head_weights
# [*, N_res, N_res, H]
pt_att = torch.sum(pt_att, dim=-1) * (-0.5)
+
# [*, N_res, N_res]
square_mask = mask.unsqueeze(-1) * mask.unsqueeze(-2)
square_mask = self.inf * (square_mask - 1)
# [*, H, N_res, N_res]
pt_att = permute_final_dims(pt_att, (2, 0, 1))
-
- if(inplace_safe):
+
+ if (inplace_safe):
a += pt_att
del pt_att
a += square_mask.unsqueeze(-3)
@@ -368,7 +443,7 @@ def forward(
a.shape[-1],
)
else:
- a = a + pt_att
+ a = a + pt_att
a = a + square_mask.unsqueeze(-3)
a = self.softmax(a)
@@ -383,19 +458,19 @@ def forward(
# [*, N_res, H * C_hidden]
o = flatten_final_dims(o, 2)
- # [*, H, 3, N_res, P_v]
- if(inplace_safe):
+ # [*, H, 3, N_res, P_v]
+ if (inplace_safe):
v_pts = permute_final_dims(v_pts, (1, 3, 0, 2))
o_pt = [
- torch.matmul(a, v.to(a.dtype))
+ torch.matmul(a, v.to(a.dtype))
for v in torch.unbind(v_pts, dim=-3)
]
o_pt = torch.stack(o_pt, dim=-3)
else:
o_pt = torch.sum(
(
- a[..., None, :, :, None]
- * permute_final_dims(v_pts, (1, 3, 0, 2))[..., None, :, :]
+ a[..., None, :, :, None]
+ * permute_final_dims(v_pts, (1, 3, 0, 2))[..., None, :, :]
),
dim=-2,
)
@@ -411,8 +486,9 @@ def forward(
# [*, N_res, H * P_v, 3]
o_pt = o_pt.reshape(*o_pt.shape[:-3], -1, 3)
+ o_pt = torch.unbind(o_pt, dim=-1)
- if(_offload_inference):
+ if (_offload_inference):
z[0] = z[0].to(o_pt.device)
# [*, N_res, H, C_z]
@@ -424,10 +500,236 @@ def forward(
# [*, N_res, C_s]
s = self.linear_out(
torch.cat(
- (o, *torch.unbind(o_pt, dim=-1), o_pt_norm, o_pair), dim=-1
+ (o, *o_pt, o_pt_norm, o_pair), dim=-1
+ ).to(dtype=z[0].dtype)
+ )
+
+ return s
+
+
+#TODO: This module follows the refactoring done in IPA for multimer. Running the regular IPA above
+# in multimer mode should be equivalent, but tests do not pass unless using this version. Determine
+# whether or not the increase in test error matters in practice.
+class InvariantPointAttentionMultimer(nn.Module):
+ """
+ Implements Algorithm 22.
+ """
+ def __init__(
+ self,
+ c_s: int,
+ c_z: int,
+ c_hidden: int,
+ no_heads: int,
+ no_qk_points: int,
+ no_v_points: int,
+ inf: float = 1e5,
+ eps: float = 1e-8,
+ is_multimer: bool = True,
+ ):
+ """
+ Args:
+ c_s:
+ Single representation channel dimension
+ c_z:
+ Pair representation channel dimension
+ c_hidden:
+ Hidden channel dimension
+ no_heads:
+ Number of attention heads
+ no_qk_points:
+ Number of query/key points to generate
+ no_v_points:
+ Number of value points to generate
+ """
+ super(InvariantPointAttentionMultimer, self).__init__()
+
+ self.c_s = c_s
+ self.c_z = c_z
+ self.c_hidden = c_hidden
+ self.no_heads = no_heads
+ self.no_qk_points = no_qk_points
+ self.no_v_points = no_v_points
+ self.inf = inf
+ self.eps = eps
+
+ # These linear layers differ from their specifications in the
+ # supplement. There, they lack bias and use Glorot initialization.
+ # Here as in the official source, they have bias and use the default
+ # Lecun initialization.
+ hc = self.c_hidden * self.no_heads
+ self.linear_q = Linear(self.c_s, hc, bias=False)
+
+ self.linear_q_points = PointProjection(
+ self.c_s,
+ self.no_qk_points,
+ self.no_heads,
+ is_multimer=True
+ )
+
+ self.linear_k = Linear(self.c_s, hc, bias=False)
+ self.linear_v = Linear(self.c_s, hc, bias=False)
+ self.linear_k_points = PointProjection(
+ self.c_s,
+ self.no_qk_points,
+ self.no_heads,
+ is_multimer=True
+ )
+
+ self.linear_v_points = PointProjection(
+ self.c_s,
+ self.no_v_points,
+ self.no_heads,
+ is_multimer=True
+ )
+
+ self.linear_b = Linear(self.c_z, self.no_heads)
+
+ self.head_weights = nn.Parameter(torch.zeros((no_heads)))
+ ipa_point_weights_init_(self.head_weights)
+
+ concat_out_dim = self.no_heads * (
+ self.c_z + self.c_hidden + self.no_v_points * 4
+ )
+ self.linear_out = Linear(concat_out_dim, self.c_s, init="final")
+
+ self.softmax = nn.Softmax(dim=-2)
+
+ def forward(
+ self,
+ s: torch.Tensor,
+ z: Optional[torch.Tensor],
+ r: Union[Rigid, Rigid3Array],
+ mask: torch.Tensor,
+ inplace_safe: bool = False,
+ _offload_inference: bool = False,
+ _z_reference_list: Optional[Sequence[torch.Tensor]] = None,
+ ) -> torch.Tensor:
+ """
+ Args:
+ s:
+ [*, N_res, C_s] single representation
+ z:
+ [*, N_res, N_res, C_z] pair representation
+ r:
+ [*, N_res] transformation object
+ mask:
+ [*, N_res] mask
+ Returns:
+ [*, N_res, C_s] single representation update
+ """
+ if(_offload_inference and inplace_safe):
+ z = _z_reference_list
+ else:
+ z = [z]
+
+ a = 0.
+
+ point_variance = (max(self.no_qk_points, 1) * 9.0 / 2)
+ point_weights = math.sqrt(1.0 / point_variance)
+
+ softplus = lambda x: torch.logaddexp(x, torch.zeros_like(x))
+
+ head_weights = softplus(self.head_weights)
+ point_weights = point_weights * head_weights
+
+ #######################################
+ # Generate scalar and point activations
+ #######################################
+
+ # [*, N_res, H, P_qk]
+ q_pts = Vec3Array.from_array(self.linear_q_points(s, r))
+
+ # [*, N_res, H, P_qk, 3]
+ k_pts = Vec3Array.from_array(self.linear_k_points(s, r))
+
+ pt_att = square_euclidean_distance(q_pts.unsqueeze(-3), k_pts.unsqueeze(-4), epsilon=0.)
+ pt_att = torch.sum(pt_att * point_weights[..., None], dim=-1) * (-0.5)
+ pt_att = pt_att.to(dtype=s.dtype)
+ a = a + pt_att
+
+ scalar_variance = max(self.c_hidden, 1) * 1.
+ scalar_weights = math.sqrt(1.0 / scalar_variance)
+
+ # [*, N_res, H * C_hidden]
+ q = self.linear_q(s)
+ k = self.linear_k(s)
+
+ # [*, N_res, H, C_hidden]
+ q = q.view(q.shape[:-1] + (self.no_heads, -1))
+ k = k.view(k.shape[:-1] + (self.no_heads, -1))
+
+ q = q * scalar_weights
+ a = a + torch.einsum('...qhc,...khc->...qkh', q, k)
+
+ ##########################
+ # Compute attention scores
+ ##########################
+ # [*, N_res, N_res, H]
+ b = self.linear_b(z[0])
+
+ if (_offload_inference):
+ assert (sys.getrefcount(z[0]) == 2)
+ z[0] = z[0].cpu()
+
+ a = a + b
+
+ # [*, N_res, N_res]
+ square_mask = mask.unsqueeze(-1) * mask.unsqueeze(-2)
+ square_mask = self.inf * (square_mask - 1)
+
+ a = a + square_mask.unsqueeze(-1)
+ a = a * math.sqrt(1. / 3) # Normalize by number of logit terms (3)
+ a = self.softmax(a)
+
+ # [*, N_res, H * C_hidden]
+ v = self.linear_v(s)
+
+ # [*, N_res, H, C_hidden]
+ v = v.view(v.shape[:-1] + (self.no_heads, -1))
+
+ o = torch.einsum('...qkh, ...khc->...qhc', a, v)
+
+ # [*, N_res, H * C_hidden]
+ o = flatten_final_dims(o, 2)
+
+ # [*, N_res, H, P_v, 3]
+ v_pts = Vec3Array.from_array(self.linear_v_points(s, r))
+
+ # [*, N_res, H, P_v]
+ o_pt = v_pts[..., None, :, :, :] * a.unsqueeze(-1)
+ o_pt = o_pt.sum(dim=-3)
+ # o_pt = Vec3Array(
+ # torch.sum(a.unsqueeze(-1) * v_pts[..., None, :, :, :].x, dim=-3),
+ # torch.sum(a.unsqueeze(-1) * v_pts[..., None, :, :, :].y, dim=-3),
+ # torch.sum(a.unsqueeze(-1) * v_pts[..., None, :, :, :].z, dim=-3),
+ # )
+
+ # [*, N_res, H * P_v, 3]
+ o_pt = o_pt.reshape(o_pt.shape[:-2] + (-1,))
+
+ # [*, N_res, H, P_v]
+ o_pt = r[..., None].apply_inverse_to_point(o_pt)
+ o_pt_flat = [o_pt.x, o_pt.y, o_pt.z]
+ o_pt_flat = [x.to(dtype=a.dtype) for x in o_pt_flat]
+
+ # [*, N_res, H * P_v]
+ o_pt_norm = o_pt.norm(epsilon=1e-8)
+
+ if (_offload_inference):
+ z[0] = z[0].to(o_pt.x.device)
+
+ o_pair = torch.einsum('...ijh, ...ijc->...ihc', a, z[0].to(dtype=a.dtype))
+
+ # [*, N_res, H * C_z]
+ o_pair = flatten_final_dims(o_pair, 2)
+
+ # [*, N_res, C_s]
+ s = self.linear_out(
+ torch.cat(
+ (o, *o_pt_flat, o_pt_norm, o_pair), dim=-1
).to(dtype=z[0].dtype)
)
-
+
return s
@@ -530,6 +832,7 @@ def __init__(
trans_scale_factor,
epsilon,
inf,
+ is_multimer=False,
**kwargs,
):
"""
@@ -583,6 +886,7 @@ def __init__(
self.trans_scale_factor = trans_scale_factor
self.epsilon = epsilon
self.inf = inf
+ self.is_multimer = is_multimer
# Buffers to be lazily initialized later
# self.default_frames
@@ -595,7 +899,8 @@ def __init__(
self.linear_in = Linear(self.c_s, self.c_s)
- self.ipa = InvariantPointAttention(
+ ipa = InvariantPointAttention if not self.is_multimer else InvariantPointAttentionMultimer
+ self.ipa = ipa(
self.c_s,
self.c_z,
self.c_ipa,
@@ -604,6 +909,7 @@ def __init__(
self.no_v_points,
inf=self.inf,
eps=self.epsilon,
+ is_multimer=self.is_multimer,
)
self.ipa_dropout = nn.Dropout(self.dropout_rate)
@@ -615,7 +921,10 @@ def __init__(
self.dropout_rate,
)
- self.bb_update = BackboneUpdate(self.c_s)
+ if self.is_multimer:
+ self.bb_update = QuatRigid(self.c_s, full_quat=False)
+ else:
+ self.bb_update = BackboneUpdate(self.c_s)
self.angle_resnet = AngleResnet(
self.c_s,
@@ -625,7 +934,7 @@ def __init__(
self.epsilon,
)
- def forward(
+ def _forward_monomer(
self,
evoformer_output_dict,
aatype,
@@ -649,7 +958,7 @@ def forward(
A dictionary of outputs
"""
s = evoformer_output_dict["single"]
-
+
if mask is None:
# [*, N]
mask = s.new_ones(s.shape[:-1])
@@ -661,8 +970,8 @@ def forward(
z = self.layer_norm_z(evoformer_output_dict["pair"])
z_reference_list = None
- if(_offload_inference):
- assert(sys.getrefcount(evoformer_output_dict["pair"]) == 2)
+ if (_offload_inference):
+ assert (sys.getrefcount(evoformer_output_dict["pair"]) == 2)
evoformer_output_dict["pair"] = evoformer_output_dict["pair"].cpu()
z_reference_list = [z]
z = None
@@ -743,8 +1052,8 @@ def forward(
rigids = rigids.stop_rot_gradient()
del z, z_reference_list
-
- if(_offload_inference):
+
+ if (_offload_inference):
evoformer_output_dict["pair"] = (
evoformer_output_dict["pair"].to(s.device)
)
@@ -754,6 +1063,129 @@ def forward(
return outputs
+ def _forward_multimer(
+ self,
+ evoformer_output_dict,
+ aatype,
+ mask=None,
+ inplace_safe=False,
+ _offload_inference=False,
+ ):
+ s = evoformer_output_dict["single"]
+
+ if mask is None:
+ # [*, N]
+ mask = s.new_ones(s.shape[:-1])
+
+ # [*, N, C_s]
+ s = self.layer_norm_s(s)
+
+ # [*, N, N, C_z]
+ z = self.layer_norm_z(evoformer_output_dict["pair"])
+
+ z_reference_list = None
+ if (_offload_inference):
+ assert (sys.getrefcount(evoformer_output_dict["pair"]) == 2)
+ evoformer_output_dict["pair"] = evoformer_output_dict["pair"].cpu()
+ z_reference_list = [z]
+ z = None
+
+ # [*, N, C_s]
+ s_initial = s
+ s = self.linear_in(s)
+
+ # [*, N]
+ rigids = Rigid3Array.identity(
+ s.shape[:-1],
+ s.device,
+ )
+ outputs = []
+ for i in range(self.no_blocks):
+ # [*, N, C_s]
+ s = s + self.ipa(
+ s,
+ z,
+ rigids,
+ mask,
+ inplace_safe=inplace_safe,
+ _offload_inference=_offload_inference,
+ _z_reference_list=z_reference_list
+ )
+ s = self.ipa_dropout(s)
+ s = self.layer_norm_ipa(s)
+ s = self.transition(s)
+
+ # [*, N]
+ rigids = rigids @ self.bb_update(s)
+
+ # [*, N, 7, 2]
+ unnormalized_angles, angles = self.angle_resnet(s, s_initial)
+
+ all_frames_to_global = self.torsion_angles_to_frames(
+ rigids.scale_translation(self.trans_scale_factor),
+ angles,
+ aatype,
+ )
+
+ pred_xyz = self.frames_and_literature_positions_to_atom14_pos(
+ all_frames_to_global,
+ aatype,
+ )
+
+ preds = {
+ "frames": rigids.scale_translation(self.trans_scale_factor).to_tensor(),
+ "sidechain_frames": all_frames_to_global.to_tensor_4x4(),
+ "unnormalized_angles": unnormalized_angles,
+ "angles": angles,
+ "positions": pred_xyz,
+ }
+
+ preds = {k: v.to(dtype=s.dtype) for k, v in preds.items()}
+
+ outputs.append(preds)
+
+ rigids = rigids.stop_rot_gradient()
+
+ del z, z_reference_list
+
+ if (_offload_inference):
+ evoformer_output_dict["pair"] = (
+ evoformer_output_dict["pair"].to(s.device)
+ )
+
+ outputs = dict_multimap(torch.stack, outputs)
+ outputs["single"] = s
+
+ return outputs
+
+ def forward(
+ self,
+ evoformer_output_dict,
+ aatype,
+ mask=None,
+ inplace_safe=False,
+ _offload_inference=False,
+ ):
+ """
+ Args:
+ s:
+ [*, N_res, C_s] single representation
+ z:
+ [*, N_res, N_res, C_z] pair representation
+ aatype:
+ [*, N_res] amino acid indices
+ mask:
+ Optional [*, N_res] sequence mask
+ Returns:
+ A dictionary of outputs
+ """
+ if(self.is_multimer):
+ outputs = self._forward_multimer(evoformer_output_dict, aatype, mask, inplace_safe, _offload_inference)
+ else:
+ outputs = self._forward_monomer(evoformer_output_dict, aatype, mask, inplace_safe, _offload_inference)
+
+ return outputs
+
def _init_residue_constants(self, float_dtype, device):
if not hasattr(self, "default_frames"):
self.register_buffer(
@@ -806,10 +1238,10 @@ def torsion_angles_to_frames(self, r, alpha, f):
return torsion_angles_to_frames(r, alpha, f, self.default_frames)
def frames_and_literature_positions_to_atom14_pos(
- self, r, f # [*, N, 8] # [*, N]
+ self, r, f # [*, N, 8] # [*, N]
):
# Lazily initialize the residue constants on the correct device
- self._init_residue_constants(r.get_rots().dtype, r.get_rots().device)
+ self._init_residue_constants(r.dtype, r.device)
return frames_and_literature_positions_to_atom14_pos(
r,
f,
diff --git a/openfold/model/template.py b/openfold/model/template.py
index 9b158231..411ad027 100644
--- a/openfold/model/template.py
+++ b/openfold/model/template.py
@@ -33,6 +33,8 @@
from openfold.model.triangular_multiplicative_update import (
TriangleMultiplicationOutgoing,
TriangleMultiplicationIncoming,
+ FusedTriangleMultiplicationOutgoing,
+ FusedTriangleMultiplicationIncoming
)
from openfold.utils.checkpointing import checkpoint_blocks
from openfold.utils.chunk_utils import (
@@ -54,6 +56,7 @@ class TemplatePointwiseAttention(nn.Module):
"""
Implements Algorithm 17.
"""
+
def __init__(self, c_t, c_z, c_hidden, no_heads, inf, **kwargs):
"""
Args:
@@ -82,12 +85,12 @@ def __init__(self, c_t, c_z, c_hidden, no_heads, inf, **kwargs):
)
def _chunk(self,
- z: torch.Tensor,
- t: torch.Tensor,
- biases: List[torch.Tensor],
- chunk_size: int,
- use_lma: bool = False,
- ) -> torch.Tensor:
+ z: torch.Tensor,
+ t: torch.Tensor,
+ biases: List[torch.Tensor],
+ chunk_size: int,
+ use_lma: bool = False,
+ ) -> torch.Tensor:
mha_inputs = {
"q_x": z,
"kv_x": t,
@@ -100,15 +103,14 @@ def _chunk(self,
no_batch_dims=len(z.shape[:-2]),
)
-
- def forward(self,
- t: torch.Tensor,
- z: torch.Tensor,
- template_mask: Optional[torch.Tensor] = None,
- # This module suffers greatly from a small chunk size
- chunk_size: Optional[int] = 256,
- use_lma: bool = False,
- ) -> torch.Tensor:
+ def forward(self,
+ t: torch.Tensor,
+ z: torch.Tensor,
+ template_mask: Optional[torch.Tensor] = None,
+ # This module suffers greatly from a small chunk size
+ chunk_size: Optional[int] = 256,
+ use_lma: bool = False,
+ ) -> torch.Tensor:
"""
Args:
t:
@@ -153,6 +155,8 @@ def __init__(
no_heads: int,
pair_transition_n: int,
dropout_rate: float,
+ tri_mul_first: bool,
+ fuse_projection_weights: bool,
inf: float,
**kwargs,
):
@@ -165,6 +169,7 @@ def __init__(
self.pair_transition_n = pair_transition_n
self.dropout_rate = dropout_rate
self.inf = inf
+ self.tri_mul_first = tri_mul_first
self.dropout_row = DropoutRowwise(self.dropout_rate)
self.dropout_col = DropoutColumnwise(self.dropout_rate)
@@ -182,31 +187,110 @@ def __init__(
inf=inf,
)
- self.tri_mul_out = TriangleMultiplicationOutgoing(
- self.c_t,
- self.c_hidden_tri_mul,
- )
- self.tri_mul_in = TriangleMultiplicationIncoming(
- self.c_t,
- self.c_hidden_tri_mul,
- )
+ if fuse_projection_weights:
+ self.tri_mul_out = FusedTriangleMultiplicationOutgoing(
+ self.c_t,
+ self.c_hidden_tri_mul,
+ )
+ self.tri_mul_in = FusedTriangleMultiplicationIncoming(
+ self.c_t,
+ self.c_hidden_tri_mul,
+ )
+ else:
+ self.tri_mul_out = TriangleMultiplicationOutgoing(
+ self.c_t,
+ self.c_hidden_tri_mul,
+ )
+ self.tri_mul_in = TriangleMultiplicationIncoming(
+ self.c_t,
+ self.c_hidden_tri_mul,
+ )
self.pair_transition = PairTransition(
self.c_t,
self.pair_transition_n,
)
- def forward(self,
- z: torch.Tensor,
- mask: torch.Tensor,
- chunk_size: Optional[int] = None,
- use_deepspeed_evo_attention: bool = False,
- use_lma: bool = False,
- inplace_safe: bool = False,
- _mask_trans: bool = True,
- _attn_chunk_size: Optional[int] = None,
- ):
- if(_attn_chunk_size is None):
+ def tri_att_start_end(self,
+ single: torch.Tensor,
+ _attn_chunk_size: Optional[int],
+ single_mask: torch.Tensor,
+ use_deepspeed_evo_attention: bool,
+ use_lma: bool,
+ inplace_safe: bool):
+ single = add(single,
+ self.dropout_row(
+ self.tri_att_start(
+ single,
+ chunk_size=_attn_chunk_size,
+ mask=single_mask,
+ use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+ use_lma=use_lma,
+ inplace_safe=inplace_safe,
+ )
+ ),
+ inplace_safe,
+ )
+
+ single = add(single,
+ self.dropout_col(
+ self.tri_att_end(
+ single,
+ chunk_size=_attn_chunk_size,
+ mask=single_mask,
+ use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+ use_lma=use_lma,
+ inplace_safe=inplace_safe,
+ )
+ ),
+ inplace_safe,
+ )
+
+ return single
+
+ def tri_mul_out_in(self,
+ single: torch.Tensor,
+ single_mask: torch.Tensor,
+ inplace_safe: bool):
+ tmu_update = self.tri_mul_out(
+ single,
+ mask=single_mask,
+ inplace_safe=inplace_safe,
+ _add_with_inplace=True,
+ )
+ if not inplace_safe:
+ single = single + self.dropout_row(tmu_update)
+ else:
+ single = tmu_update
+
+ del tmu_update
+
+ tmu_update = self.tri_mul_in(
+ single,
+ mask=single_mask,
+ inplace_safe=inplace_safe,
+ _add_with_inplace=True,
+ )
+ if not inplace_safe:
+ single = single + self.dropout_row(tmu_update)
+ else:
+ single = tmu_update
+
+ del tmu_update
+
+ return single
+
+ def forward(self,
+ z: torch.Tensor,
+ mask: torch.Tensor,
+ chunk_size: Optional[int] = None,
+ use_deepspeed_evo_attention: bool = False,
+ use_lma: bool = False,
+ inplace_safe: bool = False,
+ _mask_trans: bool = True,
+ _attn_chunk_size: Optional[int] = None,
+ ):
+ if _attn_chunk_size is None:
_attn_chunk_size = chunk_size
single_templates = [
@@ -219,74 +303,40 @@ def forward(self,
for i in range(len(single_templates)):
single = single_templates[i]
single_mask = single_templates_masks[i]
-
- single = add(single,
- self.dropout_row(
- self.tri_att_start(
- single,
- chunk_size=_attn_chunk_size,
- mask=single_mask,
- use_deepspeed_evo_attention=use_deepspeed_evo_attention,
- use_lma=use_lma,
- inplace_safe=inplace_safe,
- )
- ),
- inplace_safe,
- )
-
- single = add(single,
- self.dropout_col(
- self.tri_att_end(
- single,
- chunk_size=_attn_chunk_size,
- mask=single_mask,
- use_deepspeed_evo_attention=use_deepspeed_evo_attention,
- use_lma=use_lma,
- inplace_safe=inplace_safe,
- )
- ),
- inplace_safe,
- )
- tmu_update = self.tri_mul_out(
- single,
- mask=single_mask,
- inplace_safe=inplace_safe,
- _add_with_inplace=True,
- )
- if(not inplace_safe):
- single = single + self.dropout_row(tmu_update)
+ if self.tri_mul_first:
+ single = self.tri_att_start_end(single=self.tri_mul_out_in(single=single,
+ single_mask=single_mask,
+ inplace_safe=inplace_safe),
+ _attn_chunk_size=_attn_chunk_size,
+ single_mask=single_mask,
+ use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+ use_lma=use_lma,
+ inplace_safe=inplace_safe)
else:
- single = tmu_update
-
- del tmu_update
+ single = self.tri_mul_out_in(
+ single=self.tri_att_start_end(single=single,
+ _attn_chunk_size=_attn_chunk_size,
+ single_mask=single_mask,
+ use_deepspeed_evo_attention=use_deepspeed_evo_attention,
+ use_lma=use_lma,
+ inplace_safe=inplace_safe),
+ single_mask=single_mask,
+ inplace_safe=inplace_safe)
- tmu_update = self.tri_mul_in(
- single,
- mask=single_mask,
- inplace_safe=inplace_safe,
- _add_with_inplace=True,
- )
- if(not inplace_safe):
- single = single + self.dropout_row(tmu_update)
- else:
- single = tmu_update
-
- del tmu_update
-
single = add(single,
- self.pair_transition(
- single,
- mask=single_mask if _mask_trans else None,
- chunk_size=chunk_size,
- ),
- inplace_safe,
- )
-
- if(not inplace_safe):
+ self.pair_transition(
+ single,
+ mask=single_mask if _mask_trans else None,
+ chunk_size=chunk_size,
+ ),
+ inplace_safe,
+ )
+
+ if not inplace_safe:
single_templates[i] = single
- if(not inplace_safe):
+ if not inplace_safe:
z = torch.cat(single_templates, dim=-4)
return z
@@ -296,6 +346,7 @@ class TemplatePairStack(nn.Module):
"""
Implements Algorithm 16.
"""
+
def __init__(
self,
c_t,
@@ -305,6 +356,8 @@ def __init__(
no_heads,
pair_transition_n,
dropout_rate,
+ tri_mul_first,
+ fuse_projection_weights,
blocks_per_ckpt,
tune_chunk_size: bool = False,
inf=1e9,
@@ -341,6 +394,8 @@ def __init__(
no_heads=no_heads,
pair_transition_n=pair_transition_n,
dropout_rate=dropout_rate,
+ tri_mul_first=tri_mul_first,
+ fuse_projection_weights=fuse_projection_weights,
inf=inf,
)
self.blocks.append(block)
@@ -349,7 +404,7 @@ def __init__(
self.tune_chunk_size = tune_chunk_size
self.chunk_size_tuner = None
- if(tune_chunk_size):
+ if tune_chunk_size:
self.chunk_size_tuner = ChunkSizeTuner()
def forward(
@@ -371,7 +426,7 @@ def forward(
Returns:
[*, N_templ, N_res, N_res, C_t] template embedding update
"""
- if(mask.shape[-3] == 1):
+ if mask.shape[-3] == 1:
expand_idx = list(mask.shape)
expand_idx[-3] = t.shape[-4]
mask = mask.expand(*expand_idx)
@@ -389,18 +444,18 @@ def forward(
for b in self.blocks
]
- if(chunk_size is not None and self.chunk_size_tuner is not None):
- assert(not self.training)
+ if chunk_size is not None and self.chunk_size_tuner is not None:
+ assert (not self.training)
tuned_chunk_size = self.chunk_size_tuner.tune_chunk_size(
representative_fn=blocks[0],
args=(t.clone(),),
min_chunk_size=chunk_size,
)
blocks = [
- partial(b,
- chunk_size=tuned_chunk_size,
- _attn_chunk_size=max(chunk_size, tuned_chunk_size // 4),
- ) for b in blocks
+ partial(b,
+ chunk_size=tuned_chunk_size,
+ _attn_chunk_size=max(chunk_size, tuned_chunk_size // 4),
+ ) for b in blocks
]
t, = checkpoint_blocks(
@@ -415,11 +470,11 @@ def forward(
def embed_templates_offload(
- model,
- batch,
- z,
- pair_mask,
- templ_dim,
+ model,
+ batch,
+ z,
+ pair_mask,
+ templ_dim,
template_chunk_size=256,
inplace_safe=False,
):
@@ -470,7 +525,7 @@ def embed_templates_offload(
# [*, 1, N, N, C_z]
t = model.template_pair_stack(
t.unsqueeze(templ_dim),
- pair_mask.unsqueeze(-3).to(dtype=z.dtype),
+ pair_mask.unsqueeze(-3).to(dtype=z.dtype),
chunk_size=model.globals.chunk_size,
use_deepspeed_evo_attention=model.globals.use_deepspeed_evo_attention,
use_lma=model.globals.use_lma,
@@ -478,7 +533,7 @@ def embed_templates_offload(
_mask_trans=model.config._mask_trans,
)
- assert(sys.getrefcount(t) == 2)
+ assert (sys.getrefcount(t) == 2)
pair_embeds_cpu.append(t.cpu())
@@ -501,10 +556,10 @@ def embed_templates_offload(
)
t[..., i: i + template_chunk_size, :, :] = att_chunk
-
+
del pair_chunks
- if(inplace_safe):
+ if inplace_safe:
t = t * (torch.sum(batch["template_mask"], dim=-1) > 0)
else:
t *= (torch.sum(batch["template_mask"], dim=-1) > 0)
@@ -516,9 +571,9 @@ def embed_templates_offload(
)
# [*, N, C_m]
- a = model.template_angle_embedder(template_angle_feat)
-
- ret["template_angle_embedding"] = a
+ a = model.template_single_embedder(template_angle_feat)
+
+ ret["template_single_embedding"] = a
ret.update({"template_pair_embedding": t})
@@ -526,10 +581,10 @@ def embed_templates_offload(
def embed_templates_average(
- model,
- batch,
- z,
- pair_mask,
+ model,
+ batch,
+ z,
+ pair_mask,
templ_dim,
templ_group_size=2,
inplace_safe=False,
@@ -565,12 +620,12 @@ def embed_templates_average(
n = z.shape[-2]
n_templ = batch["template_aatype"].shape[templ_dim]
out_tensor = z.new_zeros(z.shape)
- for i in range(0, n_templ, templ_group_size):
+ for i in range(0, n_templ, templ_group_size):
def slice_template_tensor(t):
s = [slice(None) for _ in t.shape]
s[templ_dim] = slice(i, i + templ_group_size)
return t[s]
-
+
template_feats = tensor_tree_map(
slice_template_tensor,
batch,
@@ -588,8 +643,8 @@ def slice_template_tensor(t):
# [*, S_t, N, N, C_z]
t = model.template_pair_embedder(t)
t = model.template_pair_stack(
- t,
- pair_mask.unsqueeze(-3).to(dtype=z.dtype),
+ t,
+ pair_mask.unsqueeze(-3).to(dtype=z.dtype),
chunk_size=model.globals.chunk_size,
use_deepspeed_evo_attention=model.globals.use_deepspeed_evo_attention,
use_lma=model.globals.use_lma,
@@ -605,19 +660,19 @@ def slice_template_tensor(t):
)
denom = math.ceil(n_templ / templ_group_size)
- if(inplace_safe):
+ if inplace_safe:
t /= denom
else:
t = t / denom
- if(inplace_safe):
+ if inplace_safe:
out_tensor += t
else:
out_tensor = out_tensor + t
del t
- if(inplace_safe):
+ if inplace_safe:
out_tensor *= (torch.sum(batch["template_mask"], dim=-1) > 0)
else:
out_tensor = out_tensor * (torch.sum(batch["template_mask"], dim=-1) > 0)
@@ -629,9 +684,9 @@ def slice_template_tensor(t):
)
# [*, N, C_m]
- a = model.template_angle_embedder(template_angle_feat)
-
- ret["template_angle_embedding"] = a
+ a = model.template_single_embedder(template_angle_feat)
+
+ ret["template_single_embedding"] = a
ret.update({"template_pair_embedding": out_tensor})
diff --git a/openfold/model/triangular_multiplicative_update.py b/openfold/model/triangular_multiplicative_update.py
index 56a0e070..4365d5a0 100644
--- a/openfold/model/triangular_multiplicative_update.py
+++ b/openfold/model/triangular_multiplicative_update.py
@@ -15,6 +15,7 @@
from functools import partialmethod
from typing import Optional
+from abc import ABC, abstractmethod
import torch
import torch.nn as nn
@@ -25,11 +26,12 @@
from openfold.utils.tensor_utils import add, permute_final_dims
-class TriangleMultiplicativeUpdate(nn.Module):
+class BaseTriangleMultiplicativeUpdate(nn.Module, ABC):
"""
Implements Algorithms 11 and 12.
"""
- def __init__(self, c_z, c_hidden, _outgoing=True):
+ @abstractmethod
+ def __init__(self, c_z, c_hidden, _outgoing):
"""
Args:
c_z:
@@ -37,15 +39,11 @@ def __init__(self, c_z, c_hidden, _outgoing=True):
c:
Hidden channel dimension
"""
- super(TriangleMultiplicativeUpdate, self).__init__()
+ super(BaseTriangleMultiplicativeUpdate, self).__init__()
self.c_z = c_z
self.c_hidden = c_hidden
self._outgoing = _outgoing
- self.linear_a_p = Linear(self.c_z, self.c_hidden)
- self.linear_a_g = Linear(self.c_z, self.c_hidden, init="gating")
- self.linear_b_p = Linear(self.c_z, self.c_hidden)
- self.linear_b_g = Linear(self.c_z, self.c_hidden, init="gating")
self.linear_g = Linear(self.c_z, self.c_z, init="gating")
self.linear_z = Linear(self.c_hidden, self.c_z, init="final")
@@ -84,6 +82,46 @@ def _combine_projections(self,
return permute_final_dims(p, (1, 2, 0))
+ @abstractmethod
+ def forward(self,
+ z: torch.Tensor,
+ mask: Optional[torch.Tensor] = None,
+ inplace_safe: bool = False,
+ _add_with_inplace: bool = False
+ ) -> torch.Tensor:
+ """
+ Args:
+ x:
+ [*, N_res, N_res, C_z] input tensor
+ mask:
+ [*, N_res, N_res] input mask
+ Returns:
+ [*, N_res, N_res, C_z] output tensor
+ """
+ pass
+
+
+class TriangleMultiplicativeUpdate(BaseTriangleMultiplicativeUpdate):
+ """
+ Implements Algorithms 11 and 12.
+ """
+ def __init__(self, c_z, c_hidden, _outgoing=True):
+ """
+ Args:
+ c_z:
+ Input channel dimension
+ c:
+ Hidden channel dimension
+ """
+ super(TriangleMultiplicativeUpdate, self).__init__(c_z=c_z,
+ c_hidden=c_hidden,
+ _outgoing=_outgoing)
+
+ self.linear_a_p = Linear(self.c_z, self.c_hidden)
+ self.linear_a_g = Linear(self.c_z, self.c_hidden, init="gating")
+ self.linear_b_p = Linear(self.c_z, self.c_hidden)
+ self.linear_b_g = Linear(self.c_z, self.c_hidden, init="gating")
+
def _inference_forward(self,
z: torch.Tensor,
mask: Optional[torch.Tensor] = None,
@@ -397,7 +435,7 @@ def forward(self,
# reduced-precision modes
a_std = a.std()
b_std = b.std()
- if(a_std != 0. and b_std != 0.):
+ if(is_fp16_enabled() and a_std != 0. and b_std != 0.):
a = a / a.std()
b = b / b.std()
@@ -428,3 +466,152 @@ class TriangleMultiplicationIncoming(TriangleMultiplicativeUpdate):
Implements Algorithm 12.
"""
__init__ = partialmethod(TriangleMultiplicativeUpdate.__init__, _outgoing=False)
+
+
+class FusedTriangleMultiplicativeUpdate(BaseTriangleMultiplicativeUpdate):
+ """
+ Implements Algorithms 11 and 12.
+ """
+
+ def __init__(self, c_z, c_hidden, _outgoing=True):
+ """
+ Args:
+ c_z:
+ Input channel dimension
+ c:
+ Hidden channel dimension
+ """
+ super(FusedTriangleMultiplicativeUpdate, self).__init__(c_z=c_z,
+ c_hidden=c_hidden,
+ _outgoing=_outgoing)
+
+ self.linear_ab_p = Linear(self.c_z, self.c_hidden * 2)
+ self.linear_ab_g = Linear(self.c_z, self.c_hidden * 2, init="gating")
+
+ def _inference_forward(self,
+ z: torch.Tensor,
+ mask: Optional[torch.Tensor] = None,
+ _inplace_chunk_size: Optional[int] = None,
+ with_add: bool = True,
+ ):
+ """
+ Args:
+ z:
+ A [*, N, N, C_z] pair representation
+ mask:
+ A [*, N, N] pair mask
+ with_add:
+ If True, z is overwritten with (z + update). Otherwise, it is
+ overwritten with (update).
+ Returns:
+ A reference to the overwritten z
+ """
+ if mask is None:
+ mask = z.new_ones(z.shape[:-1])
+
+ mask = mask.unsqueeze(-1)
+
+ def compute_projection_helper(pair, mask):
+ p = self.linear_ab_g(pair)
+ p.sigmoid_()
+ p *= self.linear_ab_p(pair)
+ p *= mask
+
+ return p
+
+ def compute_projection(pair, mask):
+ p = compute_projection_helper(pair, mask)
+ left = p[..., :self.c_hidden]
+ right = p[..., self.c_hidden:]
+
+ return left, right
+
+ z_norm_in = self.layer_norm_in(z)
+ a, b = compute_projection(z_norm_in, mask)
+ x = self._combine_projections(a, b, _inplace_chunk_size=_inplace_chunk_size)
+ x = self.layer_norm_out(x)
+ x = self.linear_z(x)
+ g = self.linear_g(z_norm_in)
+ g.sigmoid_()
+ x *= g
+ if (with_add):
+ z += x
+ else:
+ z = x
+
+ return z
+
+ def forward(self,
+ z: torch.Tensor,
+ mask: Optional[torch.Tensor] = None,
+ inplace_safe: bool = False,
+ _add_with_inplace: bool = False,
+ _inplace_chunk_size: Optional[int] = 256
+ ) -> torch.Tensor:
+ """
+ Args:
+ x:
+ [*, N_res, N_res, C_z] input tensor
+ mask:
+ [*, N_res, N_res] input mask
+ Returns:
+ [*, N_res, N_res, C_z] output tensor
+ """
+ if (inplace_safe):
+ x = self._inference_forward(
+ z,
+ mask,
+ _inplace_chunk_size=_inplace_chunk_size,
+ with_add=_add_with_inplace,
+ )
+ return x
+
+ if mask is None:
+ mask = z.new_ones(z.shape[:-1])
+
+ mask = mask.unsqueeze(-1)
+
+ z = self.layer_norm_in(z)
+ ab = mask
+ ab = ab * self.sigmoid(self.linear_ab_g(z))
+ ab = ab * self.linear_ab_p(z)
+
+ a = ab[..., :self.c_hidden]
+ b = ab[..., self.c_hidden:]
+
+ # Prevents overflow of torch.matmul in combine projections in
+ # reduced-precision modes
+ a_std = a.std()
+ b_std = b.std()
+ if (is_fp16_enabled() and a_std != 0. and b_std != 0.):
+ a = a / a.std()
+ b = b / b.std()
+
+ if (is_fp16_enabled()):
+ with torch.cuda.amp.autocast(enabled=False):
+ x = self._combine_projections(a.float(), b.float())
+ else:
+ x = self._combine_projections(a, b)
+
+ del a, b
+ x = self.layer_norm_out(x)
+ x = self.linear_z(x)
+ g = self.sigmoid(self.linear_g(z))
+ x = x * g
+
+ return x
+
+
+class FusedTriangleMultiplicationOutgoing(FusedTriangleMultiplicativeUpdate):
+ """
+ Implements Algorithm 11.
+ """
+ __init__ = partialmethod(FusedTriangleMultiplicativeUpdate.__init__, _outgoing=True)
+
+
+class FusedTriangleMultiplicationIncoming(FusedTriangleMultiplicativeUpdate):
+ """
+ Implements Algorithm 12.
+ """
+ __init__ = partialmethod(FusedTriangleMultiplicativeUpdate.__init__, _outgoing=False)
+
diff --git a/openfold/np/protein.py b/openfold/np/protein.py
index 72bb87a1..352eb752 100644
--- a/openfold/np/protein.py
+++ b/openfold/np/protein.py
@@ -36,6 +36,11 @@
ModelOutput = Mapping[str, Any] # Is a nested dict.
PICO_TO_ANGSTROM = 0.01
+PDB_CHAIN_IDS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
+PDB_MAX_CHAINS = len(PDB_CHAIN_IDS)
+assert(PDB_MAX_CHAINS == 62)
+
+
@dataclasses.dataclass(frozen=True)
class Protein:
"""Protein structure representation."""
@@ -73,6 +78,13 @@ class Protein:
# Chain corresponding to each parent
parents_chain_index: Optional[Sequence[int]] = None
+ def __post_init__(self):
+ if(len(np.unique(self.chain_index)) > PDB_MAX_CHAINS):
+ raise ValueError(
+ f"Cannot build an instance with more than {PDB_MAX_CHAINS} "
+ "chains because these cannot be written to PDB format"
+ )
+
def from_pdb_string(pdb_str: str, chain_id: Optional[str] = None) -> Protein:
"""Takes a PDB string and constructs a Protein object.
@@ -108,6 +120,7 @@ def from_pdb_string(pdb_str: str, chain_id: Optional[str] = None) -> Protein:
for chain in model:
if(chain_id is not None and chain.id != chain_id):
continue
+
for res in chain:
if res.id[2] != " ":
raise ValueError(
@@ -132,6 +145,7 @@ def from_pdb_string(pdb_str: str, chain_id: Optional[str] = None) -> Protein:
if np.sum(mask) < 0.5:
# If no known atom positions are reported for the residue then skip it.
continue
+
aatype.append(restype_idx)
atom_positions.append(pos)
atom_mask.append(mask)
@@ -224,6 +238,14 @@ def from_proteinnet_string(proteinnet_str: str) -> Protein:
)
+def _chain_end(atom_index, end_resname, chain_name, residue_index) -> str:
+ chain_end = 'TER'
+ return(
+ f'{chain_end:<6}{atom_index:>5} {end_resname:>3} '
+ f'{chain_name:>1}{residue_index:>4}'
+ )
+
+
def get_pdb_headers(prot: Protein, chain_id: int = 0) -> Sequence[str]:
pdb_headers = []
@@ -316,21 +338,46 @@ def to_pdb(prot: Protein) -> str:
atom_positions = prot.atom_positions
residue_index = prot.residue_index.astype(np.int32)
b_factors = prot.b_factors
- chain_index = prot.chain_index
+ chain_index = prot.chain_index.astype(np.int32)
if np.any(aatype > residue_constants.restype_num):
raise ValueError("Invalid aatypes.")
+ # Construct a mapping from chain integer indices to chain ID strings.
+ chain_ids = {}
+ for i in np.unique(chain_index): # np.unique gives sorted output.
+ if i >= PDB_MAX_CHAINS:
+ raise ValueError(
+ f"The PDB format supports at most {PDB_MAX_CHAINS} chains."
+ )
+ chain_ids[i] = PDB_CHAIN_IDS[i]
+
headers = get_pdb_headers(prot)
- if(len(headers) > 0):
+ if (len(headers) > 0):
pdb_lines.extend(headers)
+ pdb_lines.append("MODEL 1")
n = aatype.shape[0]
atom_index = 1
+ last_chain_index = chain_index[0]
prev_chain_index = 0
chain_tags = string.ascii_uppercase
+
# Add all atom sites.
- for i in range(n):
+ for i in range(aatype.shape[0]):
+ # Close the previous chain if in a multichain PDB.
+ if last_chain_index != chain_index[i]:
+ pdb_lines.append(
+ _chain_end(
+ atom_index,
+ res_1to3(aatype[i - 1]),
+ chain_ids[chain_index[i - 1]],
+ residue_index[i - 1]
+ )
+ )
+ last_chain_index = chain_index[i]
+ atom_index += 1 # Atom index increases at the TER symbol.
+
res_name_3 = res_1to3(aatype[i])
for atom_name, pos, mask, b_factor in zip(
atom_types, atom_positions[i], atom_mask[i], b_factors[i]
@@ -355,6 +402,8 @@ def to_pdb(prot: Protein) -> str:
# PDB is a columnar format, every space matters here!
atom_line = (
f"{record_type:<6}{atom_index:>5} {name:<4}{alt_loc:>1}"
+ #TODO: check this refactor, chose main branch version
+ #f"{res_name_3:>3} {chain_ids[chain_index[i]]:>1}"
f"{res_name_3:>3} {chain_tag:>1}"
f"{residue_index[i]:>4}{insertion_code:>1} "
f"{pos[0]:>8.3f}{pos[1]:>8.3f}{pos[2]:>8.3f}"
@@ -386,9 +435,12 @@ def to_pdb(prot: Protein) -> str:
# each new chain.
pdb_lines.extend(get_pdb_headers(prot, prev_chain_index))
+ pdb_lines.append("ENDMDL")
pdb_lines.append("END")
- pdb_lines.append("")
- return "\n".join(pdb_lines)
+
+ # Pad all lines to 80 characters
+ pdb_lines = [line.ljust(80) for line in pdb_lines]
+ return '\n'.join(pdb_lines) + '\n' # Add terminating newline.
def to_modelcif(prot: Protein) -> str:
@@ -539,7 +591,7 @@ def from_prediction(
features: FeatureDict,
result: ModelOutput,
b_factors: Optional[np.ndarray] = None,
- chain_index: Optional[np.ndarray] = None,
+ remove_leading_feature_dimension: bool = True,
remark: Optional[str] = None,
parents: Optional[Sequence[str]] = None,
parents_chain_index: Optional[Sequence[int]] = None
@@ -550,20 +602,32 @@ def from_prediction(
features: Dictionary holding model inputs.
result: Dictionary holding model outputs.
b_factors: (Optional) B-factors to use for the protein.
+ remove_leading_feature_dimension: Whether to remove the leading dimension
+ of the `features` values
chain_index: (Optional) Chain indices for multi-chain predictions
remark: (Optional) Remark about the prediction
parents: (Optional) List of template names
Returns:
A protein instance.
"""
+ def _maybe_remove_leading_dim(arr: np.ndarray) -> np.ndarray:
+ return arr[0] if remove_leading_feature_dimension else arr
+
+ if 'asym_id' in features:
+ chain_index = _maybe_remove_leading_dim(features["asym_id"]) - 1
+ else:
+ chain_index = np.zeros_like(
+ _maybe_remove_leading_dim(features["aatype"])
+ )
+
if b_factors is None:
b_factors = np.zeros_like(result["final_atom_mask"])
return Protein(
- aatype=features["aatype"],
+ aatype=_maybe_remove_leading_dim(features["aatype"]),
atom_positions=result["final_atom_positions"],
atom_mask=result["final_atom_mask"],
- residue_index=features["residue_index"] + 1,
+ residue_index=_maybe_remove_leading_dim(features["residue_index"]) + 1,
b_factors=b_factors,
chain_index=chain_index,
remark=remark,
diff --git a/openfold/np/relax/amber_minimize.py b/openfold/np/relax/amber_minimize.py
index 3cfbf221..02816bb8 100644
--- a/openfold/np/relax/amber_minimize.py
+++ b/openfold/np/relax/amber_minimize.py
@@ -563,60 +563,3 @@ def run_pipeline(
)
iteration += 1
return ret
-
-
-def get_initial_energies(
- pdb_strs: Sequence[str],
- stiffness: float = 0.0,
- restraint_set: str = "non_hydrogen",
- exclude_residues: Optional[Sequence[int]] = None,
-):
- """Returns initial potential energies for a sequence of PDBs.
-
- Assumes the input PDBs are ready for minimization, and all have the same
- topology.
- Allows time to be saved by not pdbfixing / rebuilding the system.
-
- Args:
- pdb_strs: List of PDB strings.
- stiffness: kcal/mol A**2, spring constant of heavy atom restraining
- potential.
- restraint_set: Which atom types to restrain.
- exclude_residues: An optional list of zero-indexed residues to exclude from
- restraints.
-
- Returns:
- A list of initial energies in the same order as pdb_strs.
- """
- exclude_residues = exclude_residues or []
-
- openmm_pdbs = [
- openmm_app.PDBFile(PdbStructure(io.StringIO(p))) for p in pdb_strs
- ]
- force_field = openmm_app.ForceField("amber99sb.xml")
- system = force_field.createSystem(
- openmm_pdbs[0].topology, constraints=openmm_app.HBonds
- )
- stiffness = stiffness * ENERGY / (LENGTH ** 2)
- if stiffness > 0 * ENERGY / (LENGTH ** 2):
- _add_restraints(
- system, openmm_pdbs[0], stiffness, restraint_set, exclude_residues
- )
- simulation = openmm_app.Simulation(
- openmm_pdbs[0].topology,
- system,
- openmm.LangevinIntegrator(0, 0.01, 0.0),
- openmm.Platform.getPlatformByName("CPU"),
- )
- energies = []
- for pdb in openmm_pdbs:
- try:
- simulation.context.setPositions(pdb.positions)
- state = simulation.context.getState(getEnergy=True)
- energies.append(state.getPotentialEnergy().value_in_unit(ENERGY))
- except Exception as e: # pylint: disable=broad-except
- logging.error(
- "Error getting initial energy, returning large value %s", e
- )
- energies.append(unit.Quantity(1e20, ENERGY))
- return energies
diff --git a/openfold/np/relax/utils.py b/openfold/np/relax/utils.py
index fc19a91b..422d21bf 100644
--- a/openfold/np/relax/utils.py
+++ b/openfold/np/relax/utils.py
@@ -79,7 +79,7 @@ def assert_equal_nonterminal_atom_types(
"""Checks that pre- and post-minimized proteins have same atom set."""
# Ignore any terminal OXT atoms which may have been added by minimization.
oxt = residue_constants.atom_order["OXT"]
- no_oxt_mask = np.ones(shape=atom_mask.shape, dtype=np.bool)
+ no_oxt_mask = np.ones(shape=atom_mask.shape, dtype=bool)
no_oxt_mask[..., oxt] = False
np.testing.assert_almost_equal(
ref_atom_mask[no_oxt_mask], atom_mask[no_oxt_mask]
diff --git a/openfold/np/residue_constants.py b/openfold/np/residue_constants.py
index dcb0c6ed..9e9a1d18 100644
--- a/openfold/np/residue_constants.py
+++ b/openfold/np/residue_constants.py
@@ -17,14 +17,13 @@
import collections
import functools
+import os
from typing import Mapping, List, Tuple
from importlib import resources
import numpy as np
import tree
-# Internal import (35fd).
-
# Distance from one CA to next CA [trans configuration: omega = 180].
ca_ca = 3.80209737096
@@ -450,9 +449,9 @@ def load_stereo_chemical_props() -> Tuple[
("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
+ 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
"""
# TODO: this file should be downloaded in a setup script
stereo_chemical_props = resources.read_text("openfold.resources", "stereo_chemical_props.txt")
@@ -1310,3 +1309,179 @@ def aatype_to_str_sequence(aatype):
restypes_with_x[aatype[i]]
for i in range(len(aatype))
])
+
+
+### ALPHAFOLD MULTIMER STUFF ###
+def _make_chi_atom_indices():
+ """Returns atom indices needed to compute chi angles for all residue types.
+
+ Returns:
+ A tensor of shape [residue_types=21, chis=4, atoms=4]. The residue types are
+ in the order specified in residue_constants.restypes + unknown residue type
+ at the end. For chi angles which are not defined on the residue, the
+ positions indices are by default set to 0.
+ """
+ chi_atom_indices = []
+ for residue_name in restypes:
+ residue_name = restype_1to3[residue_name]
+ residue_chi_angles = chi_angles_atoms[residue_name]
+ atom_indices = []
+ for chi_angle in residue_chi_angles:
+ atom_indices.append(
+ [atom_order[atom] for atom in chi_angle])
+ for _ in range(4 - len(atom_indices)):
+ atom_indices.append([0, 0, 0, 0]) # For chi angles not defined on the AA.
+ chi_atom_indices.append(atom_indices)
+
+ chi_atom_indices.append([[0, 0, 0, 0]] * 4) # For UNKNOWN residue.
+
+ return np.array(chi_atom_indices)
+
+
+def _make_renaming_matrices():
+ """Matrices to map atoms to symmetry partners in ambiguous case."""
+ # As the atom naming is ambiguous for 7 of the 20 amino acids, provide
+ # alternative groundtruth coordinates where the naming is swapped
+ restype_3 = [
+ restype_1to3[res] for res in restypes
+ ]
+ restype_3 += ['UNK']
+ # Matrices for renaming ambiguous atoms.
+ all_matrices = {res: np.eye(14, dtype=np.float32) for res in restype_3}
+ for resname, swap in residue_atom_renaming_swaps.items():
+ correspondences = np.arange(14)
+ for source_atom_swap, target_atom_swap in swap.items():
+ source_index = restype_name_to_atom14_names[
+ resname].index(source_atom_swap)
+ target_index = restype_name_to_atom14_names[
+ resname].index(target_atom_swap)
+ correspondences[source_index] = target_index
+ correspondences[target_index] = source_index
+ renaming_matrix = np.zeros((14, 14), dtype=np.float32)
+ for index, correspondence in enumerate(correspondences):
+ renaming_matrix[index, correspondence] = 1.
+ all_matrices[resname] = renaming_matrix.astype(np.float32)
+ renaming_matrices = np.stack([all_matrices[restype] for restype in restype_3])
+ return renaming_matrices
+
+
+def _make_restype_atom37_mask():
+ """Mask of which atoms are present for which residue type in atom37."""
+ # create the corresponding mask
+ restype_atom37_mask = np.zeros([21, 37], dtype=np.float32)
+ 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]
+ restype_atom37_mask[restype, atom_type] = 1
+ return restype_atom37_mask
+
+
+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. if name else 0.) for name in atom_names])
+
+ restype_atom14_mask.append([0.] * 14)
+ restype_atom14_mask = np.array(restype_atom14_mask, dtype=np.float32)
+ return restype_atom14_mask
+
+
+def _make_restype_atom37_to_atom14():
+ """Map from atom37 to atom14 per residue type."""
+ restype_atom37_to_atom14 = [] # mapping (restype, atom37) --> atom14
+ for rt in restypes:
+ atom_names = restype_name_to_atom14_names[
+ restype_1to3[rt]]
+ atom_name_to_idx14 = {name: i for i, name in enumerate(atom_names)}
+ restype_atom37_to_atom14.append([
+ (atom_name_to_idx14[name] if name in atom_name_to_idx14 else 0)
+ for name in atom_types
+ ])
+
+ restype_atom37_to_atom14.append([0] * 37)
+ restype_atom37_to_atom14 = np.array(restype_atom37_to_atom14, dtype=np.int32)
+ return restype_atom37_to_atom14
+
+
+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
+
+
+def _make_restype_atom14_is_ambiguous():
+ """Mask which atoms are ambiguous in atom14."""
+ # create an ambiguous atoms mask. shape: (21, 14)
+ restype_atom14_is_ambiguous = np.zeros((21, 14), dtype=np.float32)
+ for resname, swap in residue_atom_renaming_swaps.items():
+ for atom_name1, atom_name2 in swap.items():
+ restype = restype_order[
+ restype_3to1[resname]]
+ atom_idx1 = restype_name_to_atom14_names[resname].index(
+ atom_name1)
+ atom_idx2 = restype_name_to_atom14_names[resname].index(
+ atom_name2)
+ restype_atom14_is_ambiguous[restype, atom_idx1] = 1
+ restype_atom14_is_ambiguous[restype, atom_idx2] = 1
+
+ return restype_atom14_is_ambiguous
+
+
+def _make_restype_rigidgroup_base_atom37_idx():
+ """Create Map from rigidgroups to atom37 indices."""
+ # Create an array with the atom names.
+ # shape (num_restypes, num_rigidgroups, 3_atoms): (21, 8, 3)
+ base_atom_names = np.full([21, 8, 3], '', dtype=object)
+
+ # 0: backbone frame
+ base_atom_names[:, 0, :] = ['C', 'CA', 'N']
+
+ # 3: 'psi-group'
+ base_atom_names[:, 3, :] = ['CA', 'C', 'O']
+
+ # 4,5,6,7: 'chi1,2,3,4-group'
+ for restype, restype_letter in enumerate(restypes):
+ resname = restype_1to3[restype_letter]
+ for chi_idx in range(4):
+ if chi_angles_mask[restype][chi_idx]:
+ atom_names = chi_angles_atoms[resname][chi_idx]
+ base_atom_names[restype, chi_idx + 4, :] = atom_names[1:]
+
+ # Translate atom names into atom37 indices.
+ lookuptable = atom_order.copy()
+ lookuptable[''] = 0
+ restype_rigidgroup_base_atom37_idx = np.vectorize(lambda x: lookuptable[x])(
+ base_atom_names)
+ return restype_rigidgroup_base_atom37_idx
+
+
+CHI_ATOM_INDICES = _make_chi_atom_indices()
+RENAMING_MATRICES = _make_renaming_matrices()
+RESTYPE_ATOM14_TO_ATOM37 = _make_restype_atom14_to_atom37()
+RESTYPE_ATOM37_TO_ATOM14 = _make_restype_atom37_to_atom14()
+RESTYPE_ATOM37_MASK = _make_restype_atom37_mask()
+RESTYPE_ATOM14_MASK = _make_restype_atom14_mask()
+RESTYPE_ATOM14_IS_AMBIGUOUS = _make_restype_atom14_is_ambiguous()
+RESTYPE_RIGIDGROUP_BASE_ATOM37_IDX = _make_restype_rigidgroup_base_atom37_idx()
+
+# Create mask for existing rigid groups.
+RESTYPE_RIGIDGROUP_MASK = np.zeros([21, 8], dtype=np.float32)
+RESTYPE_RIGIDGROUP_MASK[:, 0] = 1
+RESTYPE_RIGIDGROUP_MASK[:, 3] = 1
+RESTYPE_RIGIDGROUP_MASK[:20, 4:] = chi_angles_mask
diff --git a/openfold/utils/all_atom_multimer.py b/openfold/utils/all_atom_multimer.py
new file mode 100644
index 00000000..6ffd7ac1
--- /dev/null
+++ b/openfold/utils/all_atom_multimer.py
@@ -0,0 +1,494 @@
+# 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.
+"""Ops for all atom representations."""
+
+from functools import partial
+from typing import Dict, Text, Tuple
+
+import torch
+
+from openfold.np import residue_constants as rc
+from openfold.utils import geometry, tensor_utils
+from openfold.utils.geometry.rigid_matrix_vector import Rigid3Array
+import numpy as np
+
+
+def squared_difference(x, y):
+ return np.square(x - y)
+
+
+def get_rc_tensor(rc_np, aatype):
+ return torch.tensor(rc_np, device=aatype.device)[aatype]
+
+
+def atom14_to_atom37(
+ atom14_data: torch.Tensor, # (*, N, 14, ...)
+ aatype: torch.Tensor # (*, N)
+) -> Tuple: # (*, N, 37, ...)
+ """Convert atom14 to atom37 representation."""
+ idx_atom37_to_atom14 = get_rc_tensor(rc.RESTYPE_ATOM37_TO_ATOM14, aatype).long()
+ no_batch_dims = len(aatype.shape) - 1
+ atom37_data = tensor_utils.batched_gather(
+ atom14_data,
+ idx_atom37_to_atom14,
+ dim=no_batch_dims + 1,
+ no_batch_dims=no_batch_dims + 1
+ )
+ atom37_mask = get_rc_tensor(rc.RESTYPE_ATOM37_MASK, aatype)
+ if len(atom14_data.shape) == no_batch_dims + 2:
+ atom37_data *= atom37_mask
+ elif len(atom14_data.shape) == no_batch_dims + 3:
+ atom37_data *= atom37_mask[..., None].to(dtype=atom37_data.dtype)
+ else:
+ raise ValueError("Incorrectly shaped data")
+ return atom37_data, atom37_mask
+
+
+def atom37_to_atom14(aatype, all_atom_pos, all_atom_mask):
+ """Convert Atom37 positions to Atom14 positions."""
+ residx_atom14_to_atom37 = get_rc_tensor(
+ rc.RESTYPE_ATOM14_TO_ATOM37, aatype
+ )
+ no_batch_dims = len(aatype.shape)
+ atom14_mask = tensor_utils.batched_gather(
+ all_atom_mask,
+ residx_atom14_to_atom37,
+ dim=no_batch_dims + 1,
+ no_batch_dims=no_batch_dims + 1,
+ ).to(all_atom_pos.dtype)
+ # create a mask for known groundtruth positions
+ atom14_mask *= get_rc_tensor(rc.RESTYPE_ATOM14_MASK, aatype)
+ # gather the groundtruth positions
+ atom14_positions = tensor_utils.batched_gather(
+ all_atom_pos,
+ residx_atom14_to_atom37,
+ dim=no_batch_dims + 1,
+ no_batch_dims=no_batch_dims + 1,
+ ),
+ atom14_positions = atom14_mask * atom14_positions
+ return atom14_positions, atom14_mask
+
+
+def get_alt_atom14(aatype, positions: torch.Tensor, mask):
+ """Get alternative atom14 positions."""
+ # pick the transformation matrices for the given residue sequence
+ # shape (num_res, 14, 14)
+ renaming_transform = get_rc_tensor(rc.RENAMING_MATRICES, aatype)
+ alternative_positions = torch.sum(
+ positions[..., None, :] * renaming_transform[..., None],
+ dim=-2
+ )
+
+ # Create the mask for the alternative ground truth (differs from the
+ # ground truth mask, if only one of the atoms in an ambiguous pair has a
+ # ground truth position)
+ alternative_mask = torch.sum(mask[..., None] * renaming_transform, dim=-2)
+
+ return alternative_positions, alternative_mask
+
+
+def atom37_to_frames(
+ aatype: torch.Tensor, # (...)
+ all_atom_positions: torch.Tensor, # (..., 37)
+ all_atom_mask: torch.Tensor, # (..., 37)
+) -> Dict[Text, torch.Tensor]:
+ """Computes the frames for the up to 8 rigid groups for each residue."""
+ # 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'
+
+ no_batch_dims = len(aatype.shape) - 1
+
+ # Compute the gather indices for all residues in the chain.
+ # shape (N, 8, 3)
+ residx_rigidgroup_base_atom37_idx = get_rc_tensor(
+ rc.RESTYPE_RIGIDGROUP_BASE_ATOM37_IDX, aatype
+ )
+
+ # Gather the base atom positions for each rigid group.
+ base_atom_pos = tensor_utils.batched_gather(
+ all_atom_positions,
+ residx_rigidgroup_base_atom37_idx,
+ dim = no_batch_dims + 1,
+ batch_dims = no_batch_dims + 1,
+ )
+
+ # Compute the Rigids.
+ point_on_neg_x_axis = base_atom_pos[..., :, :, 0]
+ origin = base_atom_pos[..., :, :, 1]
+ point_on_xy_plane = base_atom_pos[..., :, :, 2]
+ gt_rotation = geometry.Rot3Array.from_two_vectors(
+ origin - point_on_neg_x_axis, point_on_xy_plane - origin
+ )
+
+ gt_frames = geometry.Rigid3Array(gt_rotation, origin)
+
+ # Compute a mask whether the group exists.
+ # (N, 8)
+ group_exists = get_rc_tensor(rc.RESTYPE_RIGIDGROUP_MASK, aatype)
+
+ # Compute a mask whether ground truth exists for the group
+ gt_atoms_exist = tensor_utils.batched_gather( # shape (N, 8, 3)
+ all_atom_mask.to(dtype=all_atom_positions.dtype),
+ residx_rigidgroup_base_atom37_idx,
+ batch_dims=no_batch_dims + 1,
+ )
+ gt_exists = torch.min(gt_atoms_exist, dim=-1) * group_exists # (N, 8)
+
+ # Adapt backbone frame to old convention (mirror x-axis and z-axis).
+ rots = np.tile(np.eye(3, dtype=all_atom_positions.dtype), [8, 1, 1])
+ rots[0, 0, 0] = -1
+ rots[0, 2, 2] = -1
+ gt_frames = gt_frames.compose_rotation(
+ geometry.Rot3Array.from_array(
+ torch.tensor(rots, device=aatype.device)
+ )
+ )
+
+ # The frames for ambiguous rigid groups are just rotated by 180 degree around
+ # the x-axis. The ambiguous group is always the last chi-group.
+ restype_rigidgroup_is_ambiguous = np.zeros([21, 8], dtype=all_atom_positions.dtype)
+ restype_rigidgroup_rots = np.tile(
+ np.eye(3, dtype=all_atom_positions.dtype), [21, 8, 1, 1]
+ )
+
+ for resname, _ in rc.residue_atom_renaming_swaps.items():
+ restype = rc.restype_order[
+ rc.restype_3to1[resname]
+ ]
+ chi_idx = int(sum(rc.chi_angles_mask[restype]) - 1)
+ restype_rigidgroup_is_ambiguous[restype, chi_idx + 4] = 1
+ restype_rigidgroup_rots[restype, chi_idx + 4, 1, 1] = -1
+ restype_rigidgroup_rots[restype, chi_idx + 4, 2, 2] = -1
+
+ # Gather the ambiguity information for each residue.
+ residx_rigidgroup_is_ambiguous = torch.tensor(
+ restype_rigidgroup_is_ambiguous,
+ device=aatype.device,
+ )[aatype]
+ ambiguity_rot = torch.tensor(
+ restype_rigidgroup_rots,
+ device=aatype.device,
+ )[aatype]
+ ambiguity_rot = geometry.Rot3Array.from_array(
+ torch.Tensor(ambiguity_rot, device=aatype.device)
+ )
+
+ # Create the alternative ground truth frames.
+ alt_gt_frames = gt_frames.compose_rotation(ambiguity_rot)
+
+ fix_shape = lambda x: x.reshape(x.shape[:-2] + (8,))
+
+ # reshape back to original residue layout
+ gt_frames = fix_shape(gt_frames)
+ gt_exists = fix_shape(gt_exists)
+ group_exists = fix_shape(group_exists)
+ residx_rigidgroup_is_ambiguous = fix_shape(residx_rigidgroup_is_ambiguous)
+ alt_gt_frames = fix_shape(alt_gt_frames)
+
+ return {
+ 'rigidgroups_gt_frames': gt_frames, # Rigid (..., 8)
+ 'rigidgroups_gt_exists': gt_exists, # (..., 8)
+ 'rigidgroups_group_exists': group_exists, # (..., 8)
+ 'rigidgroups_group_is_ambiguous':
+ residx_rigidgroup_is_ambiguous, # (..., 8)
+ 'rigidgroups_alt_gt_frames': alt_gt_frames, # Rigid (..., 8)
+ }
+
+
+def torsion_angles_to_frames(
+ aatype: torch.Tensor, # (N)
+ backb_to_global: geometry.Rigid3Array, # (N)
+ torsion_angles_sin_cos: torch.Tensor # (N, 7, 2)
+) -> geometry.Rigid3Array: # (N, 8)
+ """Compute rigid group frames from torsion angles."""
+ # Gather the default frames for all rigid groups.
+ # geometry.Rigid3Array with shape (N, 8)
+ m = get_rc_tensor(rc.restype_rigid_group_default_frame, aatype)
+ default_frames = geometry.Rigid3Array.from_array4x4(m)
+
+ # Create the rotation matrices according to the given angles (each frame is
+ # defined such that its rotation is around the x-axis).
+ sin_angles = torsion_angles_sin_cos[..., 0]
+ cos_angles = torsion_angles_sin_cos[..., 1]
+
+ # insert zero rotation for backbone group.
+ num_residues = aatype.shape[-1]
+ sin_angles = torch.cat(
+ [
+ torch.zeros_like(aatype).unsqueeze(dim=-1),
+ sin_angles,
+ ],
+ dim=-1)
+ cos_angles = torch.cat(
+ [
+ torch.ones_like(aatype).unsqueeze(dim=-1),
+ cos_angles
+ ],
+ dim=-1
+ )
+ zeros = torch.zeros_like(sin_angles)
+ ones = torch.ones_like(sin_angles)
+
+ # all_rots are geometry.Rot3Array with shape (..., N, 8)
+ all_rots = geometry.Rot3Array(
+ ones, zeros, zeros,
+ zeros, cos_angles, -sin_angles,
+ zeros, sin_angles, cos_angles
+ )
+
+ # Apply rotations to the frames.
+ all_frames = default_frames.compose_rotation(all_rots)
+
+ # chi2, chi3, and chi4 frames do not transform to the backbone frame but to
+ # the previous frame. So chain them up accordingly.
+
+ chi1_frame_to_backb = all_frames[..., 4]
+ chi2_frame_to_backb = chi1_frame_to_backb @ all_frames[..., 5]
+ chi3_frame_to_backb = chi2_frame_to_backb @ all_frames[..., 6]
+ chi4_frame_to_backb = chi3_frame_to_backb @ all_frames[..., 7]
+
+ all_frames_to_backb = Rigid3Array.cat(
+ [
+ all_frames[..., 0:5],
+ chi2_frame_to_backb[..., None],
+ chi3_frame_to_backb[..., None],
+ chi4_frame_to_backb[..., None]
+ ],
+ dim=-1
+ )
+
+ # Create the global frames.
+ # shape (N, 8)
+ all_frames_to_global = backb_to_global[..., None] @ all_frames_to_backb
+
+ return all_frames_to_global
+
+
+def frames_and_literature_positions_to_atom14_pos(
+ aatype: torch.Tensor, # (*, N)
+ all_frames_to_global: geometry.Rigid3Array # (N, 8)
+) -> geometry.Vec3Array: # (*, N, 14)
+ """Put atom literature positions (atom14 encoding) in each rigid group."""
+ # Pick the appropriate transform for every atom.
+ residx_to_group_idx = get_rc_tensor(
+ rc.restype_atom14_to_rigid_group,
+ aatype
+ )
+ group_mask = torch.nn.functional.one_hot(
+ residx_to_group_idx,
+ num_classes=8
+ ) # shape (*, N, 14, 8)
+
+ # geometry.Rigid3Array with shape (N, 14)
+ map_atoms_to_global = all_frames_to_global[..., None, :] * group_mask
+ map_atoms_to_global = map_atoms_to_global.map_tensor_fn(
+ partial(torch.sum, dim=-1)
+ )
+
+ # Gather the literature atom positions for each residue.
+ # geometry.Vec3Array with shape (N, 14)
+ lit_positions = geometry.Vec3Array.from_array(
+ get_rc_tensor(
+ rc.restype_atom14_rigid_group_positions,
+ aatype
+ )
+ )
+
+ # Transform each atom from its local frame to the global frame.
+ # geometry.Vec3Array with shape (N, 14)
+ pred_positions = map_atoms_to_global.apply_to_point(lit_positions)
+
+ # Mask out non-existing atoms.
+ mask = get_rc_tensor(rc.restype_atom14_mask, aatype)
+ pred_positions = pred_positions * mask
+
+ return pred_positions
+
+
+def extreme_ca_ca_distance_violations(
+ positions: geometry.Vec3Array, # (N, 37(14))
+ mask: torch.Tensor, # (N, 37(14))
+ residue_index: torch.Tensor, # (N)
+ max_angstrom_tolerance=1.5,
+ eps: float = 1e-6
+) -> torch.Tensor:
+ """Counts residues whose Ca is a large distance from its neighbor."""
+ this_ca_pos = positions[..., :-1, 1] # (N - 1,)
+ this_ca_mask = mask[..., :-1, 1] # (N - 1)
+ next_ca_pos = positions[..., 1:, 1] # (N - 1,)
+ next_ca_mask = mask[..., 1:, 1] # (N - 1)
+ has_no_gap_mask = (
+ (residue_index[..., 1:] - residue_index[..., :-1]) == 1.0
+ ).astype(positions.x.dtype)
+ ca_ca_distance = geometry.euclidean_distance(this_ca_pos, next_ca_pos, eps)
+ violations = (ca_ca_distance - rc.ca_ca) > max_angstrom_tolerance
+ mask = this_ca_mask * next_ca_mask * has_no_gap_mask
+ return tensor_utils.masked_mean(mask=mask, value=violations, dim=-1)
+
+
+def get_chi_atom_indices(device: torch.device):
+ """Returns atom indices needed to compute chi angles for all residue types.
+
+ Returns:
+ A tensor of shape [residue_types=21, chis=4, atoms=4]. The residue types are
+ in the order specified in rc.restypes + unknown residue type
+ at the end. For chi angles which are not defined on the residue, the
+ positions indices are by default set to 0.
+ """
+ chi_atom_indices = []
+ for residue_name in rc.restypes:
+ residue_name = rc.restype_1to3[residue_name]
+ residue_chi_angles = rc.chi_angles_atoms[residue_name]
+ atom_indices = []
+ for chi_angle in residue_chi_angles:
+ atom_indices.append(
+ [rc.atom_order[atom] for atom in chi_angle]
+ )
+ for _ in range(4 - len(atom_indices)):
+ atom_indices.append([0, 0, 0, 0]) # For chi angles not defined on the AA.
+ chi_atom_indices.append(atom_indices)
+
+ chi_atom_indices.append([[0, 0, 0, 0]] * 4) # For UNKNOWN residue.
+ return torch.tensor(chi_atom_indices, device=device)
+
+
+def compute_chi_angles(
+ positions: geometry.Vec3Array,
+ mask: torch.Tensor,
+ aatype: torch.Tensor
+):
+ """Computes the chi angles given all atom positions and the amino acid type.
+
+ Args:
+ positions: A Vec3Array of shape
+ [num_res, rc.atom_type_num], with positions of
+ atoms needed to calculate chi angles. Supports up to 1 batch dimension.
+ mask: An optional tensor of shape
+ [num_res, rc.atom_type_num] that masks which atom
+ positions are set for each residue. If given, then the chi mask will be
+ set to 1 for a chi angle only if the amino acid has that chi angle and all
+ the chi atoms needed to calculate that chi angle are set. If not given
+ (set to None), the chi mask will be set to 1 for a chi angle if the amino
+ acid has that chi angle and whether the actual atoms needed to calculate
+ it were set will be ignored.
+ aatype: A tensor of shape [num_res] with amino acid type integer
+ code (0 to 21). Supports up to 1 batch dimension.
+
+ Returns:
+ A tuple of tensors (chi_angles, mask), where both have shape
+ [num_res, 4]. The mask masks out unused chi angles for amino acid
+ types that have less than 4 chi angles. If atom_positions_mask is set, the
+ chi mask will also mask out uncomputable chi angles.
+ """
+
+ # Don't assert on the num_res and batch dimensions as they might be unknown.
+ assert positions.shape[-1] == rc.atom_type_num
+ assert mask.shape[-1] == rc.atom_type_num
+ no_batch_dims = len(aatype.shape) - 1
+
+ # Compute the table of chi angle indices. Shape: [restypes, chis=4, atoms=4].
+ chi_atom_indices = get_chi_atom_indices(aatype.device)
+
+ # DISCREPANCY: DeepMind doesn't remove the gaps here. I don't know why
+ # theirs works.
+ aatype_gapless = torch.clamp(aatype, max=20)
+
+ # Select atoms to compute chis. Shape: [*, num_res, chis=4, atoms=4].
+ atom_indices = chi_atom_indices[aatype_gapless]
+ # Gather atom positions. Shape: [num_res, chis=4, atoms=4, xyz=3].
+ chi_angle_atoms = positions.map_tensor_fn(
+ partial(
+ tensor_utils.batched_gather,
+ inds=atom_indices,
+ dim=-1,
+ no_batch_dims=no_batch_dims + 1
+ )
+ )
+
+ a, b, c, d = [chi_angle_atoms[..., i] for i in range(4)]
+
+ chi_angles = geometry.dihedral_angle(a, b, c, d)
+
+ # Copy the chi angle mask, add the UNKNOWN residue. Shape: [restypes, 4].
+ chi_angles_mask = list(rc.chi_angles_mask)
+ chi_angles_mask.append([0.0, 0.0, 0.0, 0.0])
+ chi_angles_mask = torch.tensor(chi_angles_mask, device=aatype.device)
+ # Compute the chi angle mask. Shape [num_res, chis=4].
+ chi_mask = chi_angles_mask[aatype_gapless]
+
+ # The chi_mask is set to 1 only when all necessary chi angle atoms were set.
+ # Gather the chi angle atoms mask. Shape: [num_res, chis=4, atoms=4].
+ chi_angle_atoms_mask = tensor_utils.batched_gather(
+ mask,
+ atom_indices,
+ dim=-1,
+ no_batch_dims=no_batch_dims + 1
+ )
+ # Check if all 4 chi angle atoms were set. Shape: [num_res, chis=4].
+ chi_angle_atoms_mask = torch.prod(chi_angle_atoms_mask, dim=-1)
+ chi_mask = chi_mask * chi_angle_atoms_mask.to(chi_angles.dtype)
+
+ return chi_angles, chi_mask
+
+
+def make_transform_from_reference(
+ a_xyz: geometry.Vec3Array,
+ b_xyz: geometry.Vec3Array,
+ c_xyz: geometry.Vec3Array
+) -> geometry.Rigid3Array:
+ """Returns rotation and translation matrices to convert from reference.
+
+ Note that this method does not take care of symmetries. If you provide the
+ coordinates in the non-standard way, the A atom will end up in the negative
+ y-axis rather than in the positive y-axis. You need to take care of such
+ cases in your code.
+
+ Args:
+ a_xyz: A Vec3Array.
+ b_xyz: A Vec3Array.
+ c_xyz: A Vec3Array.
+
+ Returns:
+ A Rigid3Array which, when applied to coordinates in a canonicalized
+ reference frame, will give coordinates approximately equal
+ the original coordinates (in the global frame).
+ """
+ rotation = geometry.Rot3Array.from_two_vectors(
+ c_xyz - b_xyz,
+ a_xyz - b_xyz
+ )
+ return geometry.Rigid3Array(rotation, b_xyz)
+
+
+def make_backbone_affine(
+ positions: geometry.Vec3Array,
+ mask: torch.Tensor,
+ aatype: torch.Tensor,
+) -> Tuple[geometry.Rigid3Array, torch.Tensor]:
+ a = rc.atom_order['N']
+ b = rc.atom_order['CA']
+ c = rc.atom_order['C']
+
+ rigid_mask = (mask[..., a] * mask[..., b] * mask[..., c])
+
+ rigid = make_transform_from_reference(
+ a_xyz=positions[..., a],
+ b_xyz=positions[..., b],
+ c_xyz=positions[..., c],
+ )
+
+ return rigid, rigid_mask
diff --git a/openfold/utils/argparse.py b/openfold/utils/argparse_utils.py
similarity index 100%
rename from openfold/utils/argparse.py
rename to openfold/utils/argparse_utils.py
diff --git a/openfold/utils/feats.py b/openfold/utils/feats.py
index 527141ea..75845d0f 100644
--- a/openfold/utils/feats.py
+++ b/openfold/utils/feats.py
@@ -18,10 +18,11 @@
import numpy as np
import torch
import torch.nn as nn
-from typing import Dict
+from typing import Dict, Union
from openfold.np import protein
import openfold.np.residue_constants as rc
+from openfold.utils.geometry import rigid_matrix_vector, rotation_matrix, vector
from openfold.utils.rigid_utils import Rotation, Rigid
from openfold.utils.tensor_utils import (
batched_gather,
@@ -89,6 +90,23 @@ def build_template_angle_feat(template_feats):
return template_angle_feat
+def dgram_from_positions(
+ pos: torch.Tensor,
+ min_bin: float = 3.25,
+ max_bin: float = 50.75,
+ no_bins: float = 39,
+ inf: float = 1e8,
+):
+ dgram = torch.sum(
+ (pos[..., None, :] - pos[..., None, :, :]) ** 2, dim=-1, keepdim=True
+ )
+ lower = torch.linspace(min_bin, max_bin, no_bins, device=pos.device) ** 2
+ upper = torch.cat([lower[1:], lower.new_tensor([inf])], dim=-1)
+ dgram = ((dgram > lower) * (dgram < upper)).type(dgram.dtype)
+
+ return dgram
+
+
def build_template_pair_feat(
batch,
min_bin, max_bin, no_bins,
@@ -100,12 +118,7 @@ def build_template_pair_feat(
# Compute distogram (this seems to differ slightly from Alg. 5)
tpb = batch["template_pseudo_beta"]
- dgram = torch.sum(
- (tpb[..., None, :] - tpb[..., None, :, :]) ** 2, dim=-1, keepdim=True
- )
- lower = torch.linspace(min_bin, max_bin, no_bins, device=tpb.device) ** 2
- upper = torch.cat([lower[1:], lower.new_tensor([inf])], dim=-1)
- dgram = ((dgram > lower) * (dgram < upper)).type(dgram.dtype)
+ dgram = dgram_from_positions(tpb, min_bin, max_bin, no_bins, inf)
to_concat = [dgram, template_mask_2d[..., None]]
@@ -170,18 +183,21 @@ def build_extra_msa_feat(batch):
def torsion_angles_to_frames(
- r: Rigid,
+ r: Union[Rigid, rigid_matrix_vector.Rigid3Array],
alpha: torch.Tensor,
aatype: torch.Tensor,
rrgdf: torch.Tensor,
):
+
+ rigid_type = type(r)
+
# [*, N, 8, 4, 4]
default_4x4 = rrgdf[aatype, ...]
# [*, N, 8] transformations, i.e.
# One [*, N, 8, 3, 3] rotation matrix and
# One [*, N, 8, 3] translation matrix
- default_r = r.from_tensor_4x4(default_4x4)
+ default_r = rigid_type.from_tensor_4x4(default_4x4)
bb_rot = alpha.new_zeros((*((1,) * len(alpha.shape[:-1])), 2))
bb_rot[..., 1] = 1
@@ -201,14 +217,13 @@ def torsion_angles_to_frames(
# This follows the original code rather than the supplement, which uses
# different indices.
- all_rots = alpha.new_zeros(default_r.get_rots().get_rot_mats().shape)
+ all_rots = alpha.new_zeros(default_r.shape + (4, 4))
all_rots[..., 0, 0] = 1
all_rots[..., 1, 1] = alpha[..., 1]
all_rots[..., 1, 2] = -alpha[..., 0]
- all_rots[..., 2, 1:] = alpha
-
- all_rots = Rigid(Rotation(rot_mats=all_rots), None)
+ all_rots[..., 2, 1:3] = alpha
+ all_rots = rigid_type.from_tensor_4x4(all_rots)
all_frames = default_r.compose(all_rots)
chi2_frame_to_frame = all_frames[..., 5]
@@ -220,7 +235,7 @@ def torsion_angles_to_frames(
chi3_frame_to_bb = chi2_frame_to_bb.compose(chi3_frame_to_frame)
chi4_frame_to_bb = chi3_frame_to_bb.compose(chi4_frame_to_frame)
- all_frames_to_bb = Rigid.cat(
+ all_frames_to_bb = rigid_type.cat(
[
all_frames[..., :5],
chi2_frame_to_bb.unsqueeze(-1),
@@ -229,14 +244,14 @@ def torsion_angles_to_frames(
],
dim=-1,
)
-
+
all_frames_to_global = r[..., None].compose(all_frames_to_bb)
return all_frames_to_global
def frames_and_literature_positions_to_atom14_pos(
- r: Rigid,
+ r: Union[Rigid, rigid_matrix_vector.Rigid3Array],
aatype: torch.Tensor,
default_frames,
group_idx,
@@ -263,7 +278,7 @@ def frames_and_literature_positions_to_atom14_pos(
lambda x: torch.sum(x, dim=-1)
)
- # [*, N, 14, 1]
+ # [*, N, 14]
atom_mask = atom_mask[aatype, ...].unsqueeze(-1)
# [*, N, 14, 3]
diff --git a/openfold/utils/geometry/__init__.py b/openfold/utils/geometry/__init__.py
new file mode 100644
index 00000000..85c32bf0
--- /dev/null
+++ b/openfold/utils/geometry/__init__.py
@@ -0,0 +1,28 @@
+# 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.
+"""Geometry Module."""
+
+from openfold.utils.geometry import rigid_matrix_vector
+from openfold.utils.geometry import rotation_matrix
+from openfold.utils.geometry import vector
+
+Rot3Array = rotation_matrix.Rot3Array
+Rigid3Array = rigid_matrix_vector.Rigid3Array
+
+Vec3Array = vector.Vec3Array
+square_euclidean_distance = vector.square_euclidean_distance
+euclidean_distance = vector.euclidean_distance
+dihedral_angle = vector.dihedral_angle
+dot = vector.dot
+cross = vector.cross
diff --git a/openfold/utils/geometry/quat_rigid.py b/openfold/utils/geometry/quat_rigid.py
new file mode 100644
index 00000000..5d53798d
--- /dev/null
+++ b/openfold/utils/geometry/quat_rigid.py
@@ -0,0 +1,38 @@
+import torch
+import torch.nn as nn
+
+from openfold.model.primitives import Linear
+from openfold.utils.geometry.rigid_matrix_vector import Rigid3Array
+from openfold.utils.geometry.rotation_matrix import Rot3Array
+from openfold.utils.geometry.vector import Vec3Array
+
+
+class QuatRigid(nn.Module):
+ def __init__(self, c_hidden, full_quat):
+ super().__init__()
+ self.full_quat = full_quat
+ if self.full_quat:
+ rigid_dim = 7
+ else:
+ rigid_dim = 6
+
+ self.linear = Linear(c_hidden, rigid_dim, init="final", precision=torch.float32)
+
+ def forward(self, activations: torch.Tensor) -> Rigid3Array:
+ # NOTE: During training, this needs to be run in higher precision
+ rigid_flat = self.linear(activations)
+
+ rigid_flat = torch.unbind(rigid_flat, dim=-1)
+ if(self.full_quat):
+ qw, qx, qy, qz = rigid_flat[:4]
+ translation = rigid_flat[4:]
+ else:
+ qx, qy, qz = rigid_flat[:3]
+ qw = torch.ones_like(qx)
+ translation = rigid_flat[3:]
+
+ rotation = Rot3Array.from_quaternion(
+ qw, qx, qy, qz, normalize=True,
+ )
+ translation = Vec3Array(*translation)
+ return Rigid3Array(rotation, translation)
diff --git a/openfold/utils/geometry/rigid_matrix_vector.py b/openfold/utils/geometry/rigid_matrix_vector.py
new file mode 100644
index 00000000..62ceff95
--- /dev/null
+++ b/openfold/utils/geometry/rigid_matrix_vector.py
@@ -0,0 +1,181 @@
+# 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.
+"""Rigid3Array Transformations represented by a Matrix and a Vector."""
+
+from __future__ import annotations
+import dataclasses
+from typing import Union, List
+
+import torch
+
+from openfold.utils.geometry import rotation_matrix
+from openfold.utils.geometry import vector
+
+
+Float = Union[float, torch.Tensor]
+
+
+@dataclasses.dataclass(frozen=True)
+class Rigid3Array:
+ """Rigid Transformation, i.e. element of special euclidean group."""
+
+ rotation: rotation_matrix.Rot3Array
+ translation: vector.Vec3Array
+
+ def __matmul__(self, other: Rigid3Array) -> Rigid3Array:
+ new_rotation = self.rotation @ other.rotation # __matmul__
+ new_translation = self.apply_to_point(other.translation)
+ return Rigid3Array(new_rotation, new_translation)
+
+ def __getitem__(self, index) -> Rigid3Array:
+ return Rigid3Array(
+ self.rotation[index],
+ self.translation[index],
+ )
+
+ def __mul__(self, other: torch.Tensor) -> Rigid3Array:
+ return Rigid3Array(
+ self.rotation * other,
+ self.translation * other,
+ )
+
+ def map_tensor_fn(self, fn) -> Rigid3Array:
+ return Rigid3Array(
+ self.rotation.map_tensor_fn(fn),
+ self.translation.map_tensor_fn(fn),
+ )
+
+ def inverse(self) -> Rigid3Array:
+ """Return Rigid3Array corresponding to inverse transform."""
+ inv_rotation = self.rotation.inverse()
+ inv_translation = inv_rotation.apply_to_point(-self.translation)
+ return Rigid3Array(inv_rotation, inv_translation)
+
+ def apply_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
+ """Apply Rigid3Array transform to point."""
+ return self.rotation.apply_to_point(point) + self.translation
+
+ def apply(self, point: torch.Tensor) -> torch.Tensor:
+ return self.apply_to_point(vector.Vec3Array.from_array(point)).to_tensor()
+
+ def apply_inverse_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
+ """Apply inverse Rigid3Array transform to point."""
+ new_point = point - self.translation
+ return self.rotation.apply_inverse_to_point(new_point)
+
+ def invert_apply(self, point: torch.Tensor) -> torch.Tensor:
+ return self.apply_inverse_to_point(vector.Vec3Array.from_array(point)).to_tensor()
+
+ def compose_rotation(self, other_rotation):
+ rot = self.rotation @ other_rotation
+ return Rigid3Array(rot, self.translation.clone())
+
+ def compose(self, other_rigid):
+ return self @ other_rigid
+
+ def unsqueeze(self, dim: int):
+ return Rigid3Array(
+ self.rotation.unsqueeze(dim),
+ self.translation.unsqueeze(dim),
+ )
+
+ @property
+ def shape(self) -> torch.Size:
+ return self.rotation.xx.shape
+
+ @property
+ def dtype(self) -> torch.dtype:
+ return self.rotation.xx.dtype
+
+ @property
+ def device(self) -> torch.device:
+ return self.rotation.xx.device
+
+ @classmethod
+ def identity(cls, shape, device) -> Rigid3Array:
+ """Return identity Rigid3Array of given shape."""
+ return cls(
+ rotation_matrix.Rot3Array.identity(shape, device),
+ vector.Vec3Array.zeros(shape, device)
+ )
+
+ @classmethod
+ def cat(cls, rigids: List[Rigid3Array], dim: int) -> Rigid3Array:
+ return cls(
+ rotation_matrix.Rot3Array.cat(
+ [r.rotation for r in rigids], dim=dim
+ ),
+ vector.Vec3Array.cat(
+ [r.translation for r in rigids], dim=dim
+ ),
+ )
+
+ def scale_translation(self, factor: Float) -> Rigid3Array:
+ """Scale translation in Rigid3Array by 'factor'."""
+ return Rigid3Array(self.rotation, self.translation * factor)
+
+ def to_tensor(self) -> torch.Tensor:
+ rot_array = self.rotation.to_tensor()
+ vec_array = self.translation.to_tensor()
+ array = torch.zeros(
+ rot_array.shape[:-2] + (4, 4),
+ device=rot_array.device,
+ dtype=rot_array.dtype
+ )
+ array[..., :3, :3] = rot_array
+ array[..., :3, 3] = vec_array
+ array[..., 3, 3] = 1.
+ return array
+
+ def to_tensor_4x4(self) -> torch.Tensor:
+ return self.to_tensor()
+
+ def reshape(self, new_shape) -> Rigid3Array:
+ rots = self.rotation.reshape(new_shape)
+ trans = self.translation.reshape(new_shape)
+ return Rigid3Array(rots, trans)
+
+ def stop_rot_gradient(self) -> Rigid3Array:
+ return Rigid3Array(
+ self.rotation.stop_gradient(),
+ self.translation,
+ )
+
+ @classmethod
+ def from_array(cls, array):
+ rot = rotation_matrix.Rot3Array.from_array(
+ array[..., :3, :3],
+ )
+ vec = vector.Vec3Array.from_array(array[..., :3, 3])
+ return cls(rot, vec)
+
+ @classmethod
+ def from_tensor_4x4(cls, array):
+ return cls.from_array(array)
+
+ @classmethod
+ def from_array4x4(cls, array: torch.tensor) -> Rigid3Array:
+ """Construct Rigid3Array from homogeneous 4x4 array."""
+ rotation = rotation_matrix.Rot3Array(
+ array[..., 0, 0], array[..., 0, 1], array[..., 0, 2],
+ array[..., 1, 0], array[..., 1, 1], array[..., 1, 2],
+ array[..., 2, 0], array[..., 2, 1], array[..., 2, 2]
+ )
+ translation = vector.Vec3Array(
+ array[..., 0, 3], array[..., 1, 3], array[..., 2, 3]
+ )
+ return cls(rotation, translation)
+
+ def cuda(self) -> Rigid3Array:
+ return Rigid3Array.from_tensor_4x4(self.to_tensor_4x4().cuda())
diff --git a/openfold/utils/geometry/rotation_matrix.py b/openfold/utils/geometry/rotation_matrix.py
new file mode 100644
index 00000000..cf012483
--- /dev/null
+++ b/openfold/utils/geometry/rotation_matrix.py
@@ -0,0 +1,208 @@
+# 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.
+"""Rot3Array Matrix Class."""
+
+from __future__ import annotations
+import dataclasses
+from typing import List
+
+import torch
+
+from openfold.utils.geometry import utils
+from openfold.utils.geometry import vector
+from openfold.utils.tensor_utils import tensor_tree_map
+
+
+COMPONENTS = ['xx', 'xy', 'xz', 'yx', 'yy', 'yz', 'zx', 'zy', 'zz']
+
+@dataclasses.dataclass(frozen=True)
+class Rot3Array:
+ """Rot3Array Matrix in 3 dimensional Space implemented as struct of arrays."""
+ xx: torch.Tensor = dataclasses.field(metadata={'dtype': torch.float32})
+ xy: torch.Tensor
+ xz: torch.Tensor
+ yx: torch.Tensor
+ yy: torch.Tensor
+ yz: torch.Tensor
+ zx: torch.Tensor
+ zy: torch.Tensor
+ zz: torch.Tensor
+
+ __array_ufunc__ = None
+
+ def __getitem__(self, index):
+ field_names = utils.get_field_names(Rot3Array)
+ return Rot3Array(
+ **{
+ name: getattr(self, name)[index]
+ for name in field_names
+ }
+ )
+
+ def __mul__(self, other: torch.Tensor):
+ field_names = utils.get_field_names(Rot3Array)
+ return Rot3Array(
+ **{
+ name: getattr(self, name) * other
+ for name in field_names
+ }
+ )
+
+ def __matmul__(self, other: Rot3Array) -> Rot3Array:
+ """Composes two Rot3Arrays."""
+ c0 = self.apply_to_point(vector.Vec3Array(other.xx, other.yx, other.zx))
+ c1 = self.apply_to_point(vector.Vec3Array(other.xy, other.yy, other.zy))
+ c2 = self.apply_to_point(vector.Vec3Array(other.xz, other.yz, other.zz))
+ return Rot3Array(c0.x, c1.x, c2.x, c0.y, c1.y, c2.y, c0.z, c1.z, c2.z)
+
+ def map_tensor_fn(self, fn) -> Rot3Array:
+ field_names = utils.get_field_names(Rot3Array)
+ return Rot3Array(
+ **{
+ name: fn(getattr(self, name))
+ for name in field_names
+ }
+ )
+
+ def inverse(self) -> Rot3Array:
+ """Returns inverse of Rot3Array."""
+ return Rot3Array(
+ self.xx, self.yx, self.zx,
+ self.xy, self.yy, self.zy,
+ self.xz, self.yz, self.zz
+ )
+
+ def apply_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
+ """Applies Rot3Array to point."""
+ return vector.Vec3Array(
+ self.xx * point.x + self.xy * point.y + self.xz * point.z,
+ self.yx * point.x + self.yy * point.y + self.yz * point.z,
+ self.zx * point.x + self.zy * point.y + self.zz * point.z
+ )
+
+ def apply_inverse_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
+ """Applies inverse Rot3Array to point."""
+ return self.inverse().apply_to_point(point)
+
+
+ def unsqueeze(self, dim: int):
+ return Rot3Array(
+ *tensor_tree_map(
+ lambda t: t.unsqueeze(dim),
+ [getattr(self, c) for c in COMPONENTS]
+ )
+ )
+
+ def stop_gradient(self) -> Rot3Array:
+ return Rot3Array(
+ *[getattr(self, c).detach() for c in COMPONENTS]
+ )
+
+ @classmethod
+ def identity(cls, shape, device) -> Rot3Array:
+ """Returns identity of given shape."""
+ ones = torch.ones(shape, dtype=torch.float32, device=device)
+ zeros = torch.zeros(shape, dtype=torch.float32, device=device)
+ return cls(ones, zeros, zeros, zeros, ones, zeros, zeros, zeros, ones)
+
+ @classmethod
+ def from_two_vectors(
+ cls, e0: vector.Vec3Array,
+ e1: vector.Vec3Array
+ ) -> Rot3Array:
+ """Construct Rot3Array from two Vectors.
+
+ Rot3Array is constructed such that in the corresponding frame 'e0' lies on
+ the positive x-Axis and 'e1' lies in the xy plane with positive sign of y.
+
+ Args:
+ e0: Vector
+ e1: Vector
+ Returns:
+ Rot3Array
+ """
+ # Normalize the unit vector for the x-axis, e0.
+ e0 = e0.normalized()
+ # make e1 perpendicular to e0.
+ c = e1.dot(e0)
+ e1 = (e1 - c * e0).normalized()
+ # Compute e2 as cross product of e0 and e1.
+ e2 = e0.cross(e1)
+ return cls(e0.x, e1.x, e2.x, e0.y, e1.y, e2.y, e0.z, e1.z, e2.z)
+
+ @classmethod
+ def from_array(cls, array: torch.Tensor) -> Rot3Array:
+ """Construct Rot3Array Matrix from array of shape. [..., 3, 3]."""
+ rows = torch.unbind(array, dim=-2)
+ rc = [torch.unbind(e, dim=-1) for e in rows]
+ return cls(*[e for row in rc for e in row])
+
+ def to_tensor(self) -> torch.Tensor:
+ """Convert Rot3Array to array of shape [..., 3, 3]."""
+ return torch.stack(
+ [
+ torch.stack([self.xx, self.xy, self.xz], dim=-1),
+ torch.stack([self.yx, self.yy, self.yz], dim=-1),
+ torch.stack([self.zx, self.zy, self.zz], dim=-1)
+ ],
+ dim=-2
+ )
+
+ @classmethod
+ def from_quaternion(cls,
+ w: torch.Tensor,
+ x: torch.Tensor,
+ y: torch.Tensor,
+ z: torch.Tensor,
+ normalize: bool = True,
+ eps: float = 1e-6
+ ) -> Rot3Array:
+ """Construct Rot3Array from components of quaternion."""
+ if normalize:
+ inv_norm = torch.rsqrt(torch.clamp(w**2 + x**2 + y**2 + z**2, min=eps))
+ w = w * inv_norm
+ x = x * inv_norm
+ y = y * inv_norm
+ z = z * inv_norm
+ xx = 1.0 - 2.0 * (y ** 2 + z ** 2)
+ xy = 2.0 * (x * y - w * z)
+ xz = 2.0 * (x * z + w * y)
+ yx = 2.0 * (x * y + w * z)
+ yy = 1.0 - 2.0 * (x ** 2 + z ** 2)
+ yz = 2.0 * (y * z - w * x)
+ zx = 2.0 * (x * z - w * y)
+ zy = 2.0 * (y * z + w * x)
+ zz = 1.0 - 2.0 * (x ** 2 + y ** 2)
+ return cls(xx, xy, xz, yx, yy, yz, zx, zy, zz)
+
+ def reshape(self, new_shape):
+ field_names = utils.get_field_names(Rot3Array)
+ reshape_fn = lambda t: t.reshape(new_shape)
+ return Rot3Array(
+ **{
+ name: reshape_fn(getattr(self, name))
+ for name in field_names
+ }
+ )
+
+ @classmethod
+ def cat(cls, rots: List[Rot3Array], dim: int) -> Rot3Array:
+ field_names = utils.get_field_names(Rot3Array)
+ cat_fn = lambda l: torch.cat(l, dim=dim)
+ return cls(
+ **{
+ name: cat_fn([getattr(r, name) for r in rots])
+ for name in field_names
+ }
+ )
diff --git a/openfold/utils/geometry/test_utils.py b/openfold/utils/geometry/test_utils.py
new file mode 100644
index 00000000..d2d11e2c
--- /dev/null
+++ b/openfold/utils/geometry/test_utils.py
@@ -0,0 +1,97 @@
+# 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.
+"""Shared utils for tests."""
+
+import dataclasses
+import torch
+
+from openfold.utils.geometry import rigid_matrix_vector
+from openfold.utils.geometry import rotation_matrix
+from openfold.utils.geometry import vector
+
+
+def assert_rotation_matrix_equal(matrix1: rotation_matrix.Rot3Array,
+ matrix2: rotation_matrix.Rot3Array):
+ for field in dataclasses.fields(rotation_matrix.Rot3Array):
+ field = field.name
+ assert torch.equal(
+ getattr(matrix1, field), getattr(matrix2, field))
+
+
+def assert_rotation_matrix_close(mat1: rotation_matrix.Rot3Array,
+ mat2: rotation_matrix.Rot3Array):
+ assert torch.allclose(mat1.to_tensor(), mat2.to_tensor(), atol=1e-6)
+
+
+def assert_array_equal_to_rotation_matrix(array: torch.Tensor,
+ matrix: rotation_matrix.Rot3Array):
+ """Check that array and Matrix match."""
+ assert torch.equal(matrix.xx, array[..., 0, 0])
+ assert torch.equal(matrix.xy, array[..., 0, 1])
+ assert torch.equal(matrix.xz, array[..., 0, 2])
+ assert torch.equal(matrix.yx, array[..., 1, 0])
+ assert torch.equal(matrix.yy, array[..., 1, 1])
+ assert torch.equal(matrix.yz, array[..., 1, 2])
+ assert torch.equal(matrix.zx, array[..., 2, 0])
+ assert torch.equal(matrix.zy, array[..., 2, 1])
+ assert torch.equal(matrix.zz, array[..., 2, 2])
+
+
+def assert_array_close_to_rotation_matrix(array: torch.Tensor,
+ matrix: rotation_matrix.Rot3Array):
+ assert torch.allclose(matrix.to_tensor(), array, atol=1e-6)
+
+
+def assert_vectors_equal(vec1: vector.Vec3Array, vec2: vector.Vec3Array):
+ assert torch.equal(vec1.x, vec2.x)
+ assert torch.equal(vec1.y, vec2.y)
+ assert torch.equal(vec1.z, vec2.z)
+
+
+def assert_vectors_close(vec1: vector.Vec3Array, vec2: vector.Vec3Array):
+ assert torch.allclose(vec1.x, vec2.x, atol=1e-6, rtol=0.)
+ assert torch.allclose(vec1.y, vec2.y, atol=1e-6, rtol=0.)
+ assert torch.allclose(vec1.z, vec2.z, atol=1e-6, rtol=0.)
+
+
+def assert_array_close_to_vector(array: torch.Tensor, vec: vector.Vec3Array):
+ assert torch.allclose(vec.to_tensor(), array, atol=1e-6, rtol=0.)
+
+
+def assert_array_equal_to_vector(array: torch.Tensor, vec: vector.Vec3Array):
+ assert torch.equal(vec.to_tensor(), array)
+
+
+def assert_rigid_equal_to_rigid(rigid1: rigid_matrix_vector.Rigid3Array,
+ rigid2: rigid_matrix_vector.Rigid3Array):
+ assert_rot_trans_equal_to_rigid(rigid1.rotation, rigid1.translation, rigid2)
+
+
+def assert_rigid_close_to_rigid(rigid1: rigid_matrix_vector.Rigid3Array,
+ rigid2: rigid_matrix_vector.Rigid3Array):
+ assert_rot_trans_close_to_rigid(rigid1.rotation, rigid1.translation, rigid2)
+
+
+def assert_rot_trans_equal_to_rigid(rot: rotation_matrix.Rot3Array,
+ trans: vector.Vec3Array,
+ rigid: rigid_matrix_vector.Rigid3Array):
+ assert_rotation_matrix_equal(rot, rigid.rotation)
+ assert_vectors_equal(trans, rigid.translation)
+
+
+def assert_rot_trans_close_to_rigid(rot: rotation_matrix.Rot3Array,
+ trans: vector.Vec3Array,
+ rigid: rigid_matrix_vector.Rigid3Array):
+ assert_rotation_matrix_close(rot, rigid.rotation)
+ assert_vectors_close(trans, rigid.translation)
diff --git a/openfold/utils/geometry/utils.py b/openfold/utils/geometry/utils.py
new file mode 100644
index 00000000..6c4d52ba
--- /dev/null
+++ b/openfold/utils/geometry/utils.py
@@ -0,0 +1,22 @@
+# 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.
+"""Utils for geometry library."""
+
+import dataclasses
+
+
+def get_field_names(cls):
+ fields = dataclasses.fields(cls)
+ field_names = [f.name for f in fields]
+ return field_names
diff --git a/openfold/utils/geometry/vector.py b/openfold/utils/geometry/vector.py
new file mode 100644
index 00000000..c25fad8f
--- /dev/null
+++ b/openfold/utils/geometry/vector.py
@@ -0,0 +1,261 @@
+# 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.
+"""Vec3Array Class."""
+
+from __future__ import annotations
+import dataclasses
+from typing import Union, List
+
+import torch
+
+Float = Union[float, torch.Tensor]
+
+@dataclasses.dataclass(frozen=True)
+class Vec3Array:
+ x: torch.Tensor = dataclasses.field(metadata={'dtype': torch.float32})
+ y: torch.Tensor
+ z: torch.Tensor
+
+ def __post_init__(self):
+ if hasattr(self.x, 'dtype'):
+ assert self.x.dtype == self.y.dtype
+ assert self.x.dtype == self.z.dtype
+ assert all([x == y for x, y in zip(self.x.shape, self.y.shape)])
+ assert all([x == z for x, z in zip(self.x.shape, self.z.shape)])
+
+ def __add__(self, other: Vec3Array) -> Vec3Array:
+ return Vec3Array(
+ self.x + other.x,
+ self.y + other.y,
+ self.z + other.z,
+ )
+
+ def __sub__(self, other: Vec3Array) -> Vec3Array:
+ return Vec3Array(
+ self.x - other.x,
+ self.y - other.y,
+ self.z - other.z,
+ )
+
+ def __mul__(self, other: Float) -> Vec3Array:
+ return Vec3Array(
+ self.x * other,
+ self.y * other,
+ self.z * other,
+ )
+
+ def __rmul__(self, other: Float) -> Vec3Array:
+ return self * other
+
+ def __truediv__(self, other: Float) -> Vec3Array:
+ return Vec3Array(
+ self.x / other,
+ self.y / other,
+ self.z / other,
+ )
+
+ def __neg__(self) -> Vec3Array:
+ return self * -1
+
+ def __pos__(self) -> Vec3Array:
+ return self * 1
+
+ def __getitem__(self, index) -> Vec3Array:
+ return Vec3Array(
+ self.x[index],
+ self.y[index],
+ self.z[index],
+ )
+
+ def __iter__(self):
+ return iter((self.x, self.y, self.z))
+
+ @property
+ def shape(self):
+ return self.x.shape
+
+ def map_tensor_fn(self, fn) -> Vec3Array:
+ return Vec3Array(
+ fn(self.x),
+ fn(self.y),
+ fn(self.z),
+ )
+
+ def cross(self, other: Vec3Array) -> Vec3Array:
+ """Compute cross product between 'self' and 'other'."""
+ new_x = self.y * other.z - self.z * other.y
+ new_y = self.z * other.x - self.x * other.z
+ new_z = self.x * other.y - self.y * other.x
+ return Vec3Array(new_x, new_y, new_z)
+
+ def dot(self, other: Vec3Array) -> Float:
+ """Compute dot product between 'self' and 'other'."""
+ return self.x * other.x + self.y * other.y + self.z * other.z
+
+ def norm(self, epsilon: float = 1e-6) -> Float:
+ """Compute Norm of Vec3Array, clipped to epsilon."""
+ # To avoid NaN on the backward pass, we must use maximum before the sqrt
+ norm2 = self.dot(self)
+ if epsilon:
+ norm2 = torch.clamp(norm2, min=epsilon**2)
+ return torch.sqrt(norm2)
+
+ def norm2(self):
+ return self.dot(self)
+
+ def normalized(self, epsilon: float = 1e-6) -> Vec3Array:
+ """Return unit vector with optional clipping."""
+ return self / self.norm(epsilon)
+
+ def clone(self) -> Vec3Array:
+ return Vec3Array(
+ self.x.clone(),
+ self.y.clone(),
+ self.z.clone(),
+ )
+
+ def reshape(self, new_shape) -> Vec3Array:
+ x = self.x.reshape(new_shape)
+ y = self.y.reshape(new_shape)
+ z = self.z.reshape(new_shape)
+
+ return Vec3Array(x, y, z)
+
+ def sum(self, dim: int) -> Vec3Array:
+ return Vec3Array(
+ torch.sum(self.x, dim=dim),
+ torch.sum(self.y, dim=dim),
+ torch.sum(self.z, dim=dim),
+ )
+
+ def unsqueeze(self, dim: int):
+ return Vec3Array(
+ self.x.unsqueeze(dim),
+ self.y.unsqueeze(dim),
+ self.z.unsqueeze(dim),
+ )
+
+ @classmethod
+ def zeros(cls, shape, device="cpu"):
+ """Return Vec3Array corresponding to zeros of given shape."""
+ return cls(
+ torch.zeros(shape, dtype=torch.float32, device=device),
+ torch.zeros(shape, dtype=torch.float32, device=device),
+ torch.zeros(shape, dtype=torch.float32, device=device)
+ )
+
+ def to_tensor(self) -> torch.Tensor:
+ return torch.stack([self.x, self.y, self.z], dim=-1)
+
+ @classmethod
+ def from_array(cls, tensor):
+ return cls(*torch.unbind(tensor, dim=-1))
+
+ @classmethod
+ def cat(cls, vecs: List[Vec3Array], dim: int) -> Vec3Array:
+ return cls(
+ torch.cat([v.x for v in vecs], dim=dim),
+ torch.cat([v.y for v in vecs], dim=dim),
+ torch.cat([v.z for v in vecs], dim=dim),
+ )
+
+
+def square_euclidean_distance(
+ vec1: Vec3Array,
+ vec2: Vec3Array,
+ epsilon: float = 1e-6
+) -> Float:
+ """Computes square of euclidean distance between 'vec1' and 'vec2'.
+
+ Args:
+ vec1: Vec3Array to compute distance to
+ vec2: Vec3Array to compute distance from, should be
+ broadcast compatible with 'vec1'
+ epsilon: distance is clipped from below to be at least epsilon
+
+ Returns:
+ Array of square euclidean distances;
+ shape will be result of broadcasting 'vec1' and 'vec2'
+ """
+ difference = vec1 - vec2
+ distance = difference.dot(difference)
+ if epsilon:
+ distance = torch.clamp(distance, min=epsilon)
+ return distance
+
+
+def dot(vector1: Vec3Array, vector2: Vec3Array) -> Float:
+ return vector1.dot(vector2)
+
+
+def cross(vector1: Vec3Array, vector2: Vec3Array) -> Float:
+ return vector1.cross(vector2)
+
+
+def norm(vector: Vec3Array, epsilon: float = 1e-6) -> Float:
+ return vector.norm(epsilon)
+
+
+def normalized(vector: Vec3Array, epsilon: float = 1e-6) -> Vec3Array:
+ return vector.normalized(epsilon)
+
+
+def euclidean_distance(
+ vec1: Vec3Array,
+ vec2: Vec3Array,
+ epsilon: float = 1e-6
+) -> Float:
+ """Computes euclidean distance between 'vec1' and 'vec2'.
+
+ Args:
+ vec1: Vec3Array to compute euclidean distance to
+ vec2: Vec3Array to compute euclidean distance from, should be
+ broadcast compatible with 'vec1'
+ epsilon: distance is clipped from below to be at least epsilon
+
+ Returns:
+ Array of euclidean distances;
+ shape will be result of broadcasting 'vec1' and 'vec2'
+ """
+ distance_sq = square_euclidean_distance(vec1, vec2, epsilon**2)
+ distance = torch.sqrt(distance_sq)
+ return distance
+
+
+def dihedral_angle(a: Vec3Array, b: Vec3Array, c: Vec3Array,
+ d: Vec3Array) -> Float:
+ """Computes torsion angle for a quadruple of points.
+
+ For points (a, b, c, d), this is the angle between the planes defined by
+ points (a, b, c) and (b, c, d). It is also known as the dihedral angle.
+
+ Arguments:
+ a: A Vec3Array of coordinates.
+ b: A Vec3Array of coordinates.
+ c: A Vec3Array of coordinates.
+ d: A Vec3Array of coordinates.
+
+ Returns:
+ A tensor of angles in radians: [-pi, pi].
+ """
+ v1 = a - b
+ v2 = b - c
+ v3 = d - c
+
+ c1 = v1.cross(v2)
+ c2 = v3.cross(v2)
+ c3 = c2.cross(c1)
+
+ v2_mag = v2.norm()
+ return torch.atan2(c3.dot(v2), v2_mag * c1.dot(c2))
diff --git a/openfold/utils/import_weights.py b/openfold/utils/import_weights.py
index b1e0638a..c698d9d1 100644
--- a/openfold/utils/import_weights.py
+++ b/openfold/utils/import_weights.py
@@ -13,6 +13,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+import re
from enum import Enum
from dataclasses import dataclass
from functools import partial
@@ -27,7 +28,7 @@
# With Param, a poor man's enum with attributes (Rust-style)
class ParamType(Enum):
LinearWeight = partial( # hack: partial prevents fns from becoming methods
- lambda w: w.transpose(-1, -2)
+ lambda w: w.unsqueeze(-1) if len(w.shape) == 1 else w.transpose(-1, -2)
)
LinearWeightMHA = partial(
lambda w: w.reshape(*w.shape[:-2], -1).transpose(-1, -2)
@@ -39,6 +40,13 @@ class ParamType(Enum):
LinearWeightOPM = partial(
lambda w: w.reshape(*w.shape[:-3], -1, w.shape[-1]).transpose(-1, -2)
)
+ LinearWeightMultimer = partial(
+ lambda w: w.unsqueeze(-1) if len(w.shape) == 1 else
+ w.reshape(w.shape[0], -1).transpose(-1, -2)
+ )
+ LinearBiasMultimer = partial(
+ lambda w: w.reshape(-1)
+ )
Other = partial(lambda w: w)
def __init__(self, fn):
@@ -50,6 +58,7 @@ class Param:
param: Union[torch.Tensor, List[torch.Tensor]]
param_type: ParamType = ParamType.Other
stacked: bool = False
+ swap: bool = False
def process_translation_dict(d, top_layer=True):
@@ -93,6 +102,7 @@ def stacked(param_dict_list, out=None):
param=[param.param for param in v],
param_type=v[0].param_type,
stacked=True,
+ swap=v[0].swap
)
out[k] = stacked_param
@@ -114,7 +124,12 @@ def assign(translation_dict, orig_weights):
try:
weights = list(map(param_type.transformation, weights))
for p, w in zip(ref, weights):
- p.copy_(w)
+ if param.swap:
+ index = p.shape[0] // 2
+ p[:index].copy_(w[index:])
+ p[index:].copy_(w[:index])
+ else:
+ p.copy_(w)
except:
print(k)
print(ref[0].shape)
@@ -122,26 +137,44 @@ def assign(translation_dict, orig_weights):
raise
-def generate_translation_dict(model, version):
+def generate_translation_dict(model, version, is_multimer=False):
#######################
# Some templates
#######################
-
LinearWeight = lambda l: (Param(l, param_type=ParamType.LinearWeight))
-
LinearBias = lambda l: (Param(l))
-
LinearWeightMHA = lambda l: (Param(l, param_type=ParamType.LinearWeightMHA))
-
LinearBiasMHA = lambda b: (Param(b, param_type=ParamType.LinearBiasMHA))
-
LinearWeightOPM = lambda l: (Param(l, param_type=ParamType.LinearWeightOPM))
+ LinearWeightMultimer = lambda l: (
+ Param(l, param_type=ParamType.LinearWeightMultimer)
+ )
+ LinearBiasMultimer = lambda l: (
+ Param(l, param_type=ParamType.LinearBiasMultimer)
+ )
+ LinearWeightSwap = lambda l: (Param(l, param_type=ParamType.LinearWeight, swap=True))
+ LinearBiasSwap = lambda l: (Param(l, swap=True))
LinearParams = lambda l: {
"weights": LinearWeight(l.weight),
"bias": LinearBias(l.bias),
}
+ LinearParamsMHA = lambda l: {
+ "weights": LinearWeightMHA(l.weight),
+ "bias": LinearBiasMHA(l.bias),
+ }
+
+ LinearParamsSwap = lambda l: {
+ "weights": LinearWeightSwap(l.weight),
+ "bias": LinearBiasSwap(l.bias),
+ }
+
+ LinearParamsMultimer = lambda l: {
+ "weights": LinearWeightMultimer(l.weight),
+ "bias": LinearBiasMultimer(l.bias),
+ }
+
LayerNormParams = lambda l: {
"scale": Param(l.weight),
"offset": Param(l.bias),
@@ -178,31 +211,48 @@ def generate_translation_dict(model, version):
"attention": AttentionGatedParams(tri_att.mha),
}
- TriMulOutParams = lambda tri_mul: {
- "layer_norm_input": LayerNormParams(tri_mul.layer_norm_in),
- "left_projection": LinearParams(tri_mul.linear_a_p),
- "right_projection": LinearParams(tri_mul.linear_b_p),
- "left_gate": LinearParams(tri_mul.linear_a_g),
- "right_gate": LinearParams(tri_mul.linear_b_g),
- "center_layer_norm": LayerNormParams(tri_mul.layer_norm_out),
- "output_projection": LinearParams(tri_mul.linear_z),
- "gating_linear": LinearParams(tri_mul.linear_g),
- }
+ def TriMulOutParams(tri_mul, outgoing=True):
+ if re.fullmatch("^model_[1-5]_multimer_v3$", version):
+ lin_param_type = LinearParams if outgoing else LinearParamsSwap
+ d = {
+ "left_norm_input": LayerNormParams(tri_mul.layer_norm_in),
+ "projection": lin_param_type(tri_mul.linear_ab_p),
+ "gate": lin_param_type(tri_mul.linear_ab_g),
+ "center_norm": LayerNormParams(tri_mul.layer_norm_out),
+ }
+ else:
+ # see commit b88f8da on the Alphafold repo
+ # Alphafold swaps the pseudocode's a and b between the incoming/outcoming
+ # iterations of triangle multiplication, which is confusing and not
+ # reproduced in our implementation.
+ if outgoing:
+ left_projection = LinearParams(tri_mul.linear_a_p)
+ right_projection = LinearParams(tri_mul.linear_b_p)
+ left_gate = LinearParams(tri_mul.linear_a_g)
+ right_gate = LinearParams(tri_mul.linear_b_g)
+ else:
+ left_projection = LinearParams(tri_mul.linear_b_p)
+ right_projection = LinearParams(tri_mul.linear_a_p)
+ left_gate = LinearParams(tri_mul.linear_b_g)
+ right_gate = LinearParams(tri_mul.linear_a_g)
+
+ d = {
+ "layer_norm_input": LayerNormParams(tri_mul.layer_norm_in),
+ "left_projection": left_projection,
+ "right_projection": right_projection,
+ "left_gate": left_gate,
+ "right_gate": right_gate,
+ "center_layer_norm": LayerNormParams(tri_mul.layer_norm_out),
+ }
- # see commit b88f8da on the Alphafold repo
- # Alphafold swaps the pseudocode's a and b between the incoming/outcoming
- # iterations of triangle multiplication, which is confusing and not
- # reproduced in our implementation.
- TriMulInParams = lambda tri_mul: {
- "layer_norm_input": LayerNormParams(tri_mul.layer_norm_in),
- "left_projection": LinearParams(tri_mul.linear_b_p),
- "right_projection": LinearParams(tri_mul.linear_a_p),
- "left_gate": LinearParams(tri_mul.linear_b_g),
- "right_gate": LinearParams(tri_mul.linear_a_g),
- "center_layer_norm": LayerNormParams(tri_mul.layer_norm_out),
- "output_projection": LinearParams(tri_mul.linear_z),
- "gating_linear": LinearParams(tri_mul.linear_g),
- }
+ d.update({
+ "output_projection": LinearParams(tri_mul.linear_z),
+ "gating_linear": LinearParams(tri_mul.linear_g),
+ })
+
+ return d
+
+ TriMulInParams = partial(TriMulOutParams, outgoing=False)
PairTransitionParams = lambda pt: {
"input_layer_norm": LayerNormParams(pt.layer_norm),
@@ -236,8 +286,46 @@ def generate_translation_dict(model, version):
IPAParams = lambda ipa: {
"q_scalar": LinearParams(ipa.linear_q),
"kv_scalar": LinearParams(ipa.linear_kv),
- "q_point_local": LinearParams(ipa.linear_q_points),
- "kv_point_local": LinearParams(ipa.linear_kv_points),
+ "q_point_local": LinearParams(ipa.linear_q_points.linear),
+ "kv_point_local": LinearParams(ipa.linear_kv_points.linear),
+ "trainable_point_weights": Param(
+ param=ipa.head_weights, param_type=ParamType.Other
+ ),
+ "attention_2d": LinearParams(ipa.linear_b),
+ "output_projection": LinearParams(ipa.linear_out),
+ }
+
+ PointProjectionParams = lambda pp: {
+ "point_projection": LinearParamsMHA(
+ pp.linear,
+ ),
+ }
+
+ IPAParamsMultimer = lambda ipa: {
+ "q_scalar_projection": {
+ "weights": LinearWeightMHA(
+ ipa.linear_q.weight,
+ ),
+ },
+ "k_scalar_projection": {
+ "weights": LinearWeightMHA(
+ ipa.linear_k.weight,
+ ),
+ },
+ "v_scalar_projection": {
+ "weights": LinearWeightMHA(
+ ipa.linear_v.weight,
+ ),
+ },
+ "q_point_projection": PointProjectionParams(
+ ipa.linear_q_points
+ ),
+ "k_point_projection": PointProjectionParams(
+ ipa.linear_k_points
+ ),
+ "v_point_projection": PointProjectionParams(
+ ipa.linear_v_points
+ ),
"trainable_point_weights": Param(
param=ipa.head_weights, param_type=ParamType.Other
),
@@ -280,109 +368,183 @@ def EvoformerBlockParams(b, is_extra_msa=False):
b.msa_att_row
),
col_att_name: msa_col_att_params,
- "msa_transition": MSATransitionParams(b.core.msa_transition),
+ "msa_transition": MSATransitionParams(b.msa_transition),
"outer_product_mean":
- OuterProductMeanParams(b.core.outer_product_mean),
+ OuterProductMeanParams(b.outer_product_mean),
"triangle_multiplication_outgoing":
- TriMulOutParams(b.core.tri_mul_out),
+ TriMulOutParams(b.pair_stack.tri_mul_out),
"triangle_multiplication_incoming":
- TriMulInParams(b.core.tri_mul_in),
+ TriMulInParams(b.pair_stack.tri_mul_in),
"triangle_attention_starting_node":
- TriAttParams(b.core.tri_att_start),
+ TriAttParams(b.pair_stack.tri_att_start),
"triangle_attention_ending_node":
- TriAttParams(b.core.tri_att_end),
+ TriAttParams(b.pair_stack.tri_att_end),
"pair_transition":
- PairTransitionParams(b.core.pair_transition),
+ PairTransitionParams(b.pair_stack.pair_transition),
}
return d
ExtraMSABlockParams = partial(EvoformerBlockParams, is_extra_msa=True)
- FoldIterationParams = lambda sm: {
- "invariant_point_attention": IPAParams(sm.ipa),
- "attention_layer_norm": LayerNormParams(sm.layer_norm_ipa),
- "transition": LinearParams(sm.transition.layers[0].linear_1),
- "transition_1": LinearParams(sm.transition.layers[0].linear_2),
- "transition_2": LinearParams(sm.transition.layers[0].linear_3),
- "transition_layer_norm": LayerNormParams(sm.transition.layer_norm),
- "affine_update": LinearParams(sm.bb_update.linear),
- "rigid_sidechain": {
- "input_projection": LinearParams(sm.angle_resnet.linear_in),
- "input_projection_1": LinearParams(sm.angle_resnet.linear_initial),
- "resblock1": LinearParams(sm.angle_resnet.layers[0].linear_1),
- "resblock2": LinearParams(sm.angle_resnet.layers[0].linear_2),
- "resblock1_1": LinearParams(sm.angle_resnet.layers[1].linear_1),
- "resblock2_1": LinearParams(sm.angle_resnet.layers[1].linear_2),
- "unnormalized_angles": LinearParams(sm.angle_resnet.linear_out),
- },
- }
+ def FoldIterationParams(sm):
+ d = {
+ "invariant_point_attention":
+ IPAParamsMultimer(sm.ipa) if is_multimer else IPAParams(sm.ipa),
+ "attention_layer_norm": LayerNormParams(sm.layer_norm_ipa),
+ "transition": LinearParams(sm.transition.layers[0].linear_1),
+ "transition_1": LinearParams(sm.transition.layers[0].linear_2),
+ "transition_2": LinearParams(sm.transition.layers[0].linear_3),
+ "transition_layer_norm": LayerNormParams(sm.transition.layer_norm),
+ "affine_update": LinearParams(sm.bb_update.linear),
+ "rigid_sidechain": {
+ "input_projection": LinearParams(sm.angle_resnet.linear_in),
+ "input_projection_1":
+ LinearParams(sm.angle_resnet.linear_initial),
+ "resblock1": LinearParams(sm.angle_resnet.layers[0].linear_1),
+ "resblock2": LinearParams(sm.angle_resnet.layers[0].linear_2),
+ "resblock1_1":
+ LinearParams(sm.angle_resnet.layers[1].linear_1),
+ "resblock2_1":
+ LinearParams(sm.angle_resnet.layers[1].linear_2),
+ "unnormalized_angles":
+ LinearParams(sm.angle_resnet.linear_out),
+ },
+ }
+
+ if(is_multimer):
+ d.pop("affine_update")
+ d["quat_rigid"] = {
+ "rigid": LinearParams(
+ sm.bb_update.linear
+ )
+ }
+
+ return d
############################
# translations dict overflow
############################
-
ems_blocks = model.extra_msa_stack.blocks
ems_blocks_params = stacked([ExtraMSABlockParams(b) for b in ems_blocks])
evo_blocks = model.evoformer.blocks
evo_blocks_params = stacked([EvoformerBlockParams(b) for b in evo_blocks])
- translations = {
- "evoformer": {
- "preprocess_1d": LinearParams(model.input_embedder.linear_tf_m),
- "preprocess_msa": LinearParams(model.input_embedder.linear_msa_m),
- "left_single": LinearParams(model.input_embedder.linear_tf_z_i),
- "right_single": LinearParams(model.input_embedder.linear_tf_z_j),
- "prev_pos_linear": LinearParams(model.recycling_embedder.linear),
- "prev_msa_first_row_norm": LayerNormParams(
- model.recycling_embedder.layer_norm_m
- ),
- "prev_pair_norm": LayerNormParams(
- model.recycling_embedder.layer_norm_z
- ),
- "pair_activiations": LinearParams(
- model.input_embedder.linear_relpos
- ),
- "extra_msa_activations": LinearParams(
- model.extra_msa_embedder.linear
- ),
- "extra_msa_stack": ems_blocks_params,
- "evoformer_iteration": evo_blocks_params,
- "single_activations": LinearParams(model.evoformer.linear),
- },
- "structure_module": {
- "single_layer_norm": LayerNormParams(
- model.structure_module.layer_norm_s
- ),
- "initial_projection": LinearParams(
- model.structure_module.linear_in
- ),
- "pair_layer_norm": LayerNormParams(
- model.structure_module.layer_norm_z
- ),
- "fold_iteration": FoldIterationParams(model.structure_module),
- },
- "predicted_lddt_head": {
- "input_layer_norm": LayerNormParams(
- model.aux_heads.plddt.layer_norm
- ),
- "act_0": LinearParams(model.aux_heads.plddt.linear_1),
- "act_1": LinearParams(model.aux_heads.plddt.linear_2),
- "logits": LinearParams(model.aux_heads.plddt.linear_3),
- },
- "distogram_head": {
- "half_logits": LinearParams(model.aux_heads.distogram.linear),
- },
- "experimentally_resolved_head": {
- "logits": LinearParams(
- model.aux_heads.experimentally_resolved.linear
- ),
- },
- "masked_msa_head": {
- "logits": LinearParams(model.aux_heads.masked_msa.linear),
- },
- }
+ if(not is_multimer):
+ translations = {
+ "evoformer": {
+ "preprocess_1d": LinearParams(model.input_embedder.linear_tf_m),
+ "preprocess_msa": LinearParams(model.input_embedder.linear_msa_m),
+ "left_single": LinearParams(model.input_embedder.linear_tf_z_i),
+ "right_single": LinearParams(model.input_embedder.linear_tf_z_j),
+ "prev_pos_linear": LinearParams(model.recycling_embedder.linear),
+ "prev_msa_first_row_norm": LayerNormParams(
+ model.recycling_embedder.layer_norm_m
+ ),
+ "prev_pair_norm": LayerNormParams(
+ model.recycling_embedder.layer_norm_z
+ ),
+ "pair_activiations": LinearParams(
+ model.input_embedder.linear_relpos
+ ),
+ "extra_msa_activations": LinearParams(
+ model.extra_msa_embedder.linear
+ ),
+ "extra_msa_stack": ems_blocks_params,
+ "evoformer_iteration": evo_blocks_params,
+ "single_activations": LinearParams(model.evoformer.linear),
+ },
+ "structure_module": {
+ "single_layer_norm": LayerNormParams(
+ model.structure_module.layer_norm_s
+ ),
+ "initial_projection": LinearParams(
+ model.structure_module.linear_in
+ ),
+ "pair_layer_norm": LayerNormParams(
+ model.structure_module.layer_norm_z
+ ),
+ "fold_iteration": FoldIterationParams(model.structure_module),
+ },
+ "predicted_lddt_head": {
+ "input_layer_norm": LayerNormParams(
+ model.aux_heads.plddt.layer_norm
+ ),
+ "act_0": LinearParams(model.aux_heads.plddt.linear_1),
+ "act_1": LinearParams(model.aux_heads.plddt.linear_2),
+ "logits": LinearParams(model.aux_heads.plddt.linear_3),
+ },
+ "distogram_head": {
+ "half_logits": LinearParams(model.aux_heads.distogram.linear),
+ },
+ "experimentally_resolved_head": {
+ "logits": LinearParams(
+ model.aux_heads.experimentally_resolved.linear
+ ),
+ },
+ "masked_msa_head": {
+ "logits": LinearParams(model.aux_heads.masked_msa.linear),
+ },
+ }
+ else:
+ translations = {
+ "evoformer": {
+ "preprocess_1d": LinearParams(model.input_embedder.linear_tf_m),
+ "preprocess_msa": LinearParams(model.input_embedder.linear_msa_m),
+ "left_single": LinearParams(model.input_embedder.linear_tf_z_i),
+ "right_single": LinearParams(model.input_embedder.linear_tf_z_j),
+ "prev_pos_linear": LinearParams(model.recycling_embedder.linear),
+ "prev_msa_first_row_norm": LayerNormParams(
+ model.recycling_embedder.layer_norm_m
+ ),
+ "prev_pair_norm": LayerNormParams(
+ model.recycling_embedder.layer_norm_z
+ ),
+ "~_relative_encoding": {
+ "position_activations": LinearParams(
+ model.input_embedder.linear_relpos
+ ),
+ },
+ "extra_msa_activations": LinearParams(
+ model.extra_msa_embedder.linear
+ ),
+ "extra_msa_stack": ems_blocks_params,
+ "evoformer_iteration": evo_blocks_params,
+ "single_activations": LinearParams(model.evoformer.linear),
+ },
+ "structure_module": {
+ "single_layer_norm": LayerNormParams(
+ model.structure_module.layer_norm_s
+ ),
+ "initial_projection": LinearParams(
+ model.structure_module.linear_in
+ ),
+ "pair_layer_norm": LayerNormParams(
+ model.structure_module.layer_norm_z
+ ),
+ "fold_iteration": FoldIterationParams(model.structure_module),
+ },
+ "predicted_lddt_head": {
+ "input_layer_norm": LayerNormParams(
+ model.aux_heads.plddt.layer_norm
+ ),
+ "act_0": LinearParams(model.aux_heads.plddt.linear_1),
+ "act_1": LinearParams(model.aux_heads.plddt.linear_2),
+ "logits": LinearParams(model.aux_heads.plddt.linear_3),
+ },
+ "distogram_head": {
+ "half_logits": LinearParams(model.aux_heads.distogram.linear),
+ },
+ "experimentally_resolved_head": {
+ "logits": LinearParams(
+ model.aux_heads.experimentally_resolved.linear
+ ),
+ },
+ "masked_msa_head": {
+ "logits": LinearParams(model.aux_heads.masked_msa.linear),
+ },
+ }
no_templ = [
"model_3",
@@ -394,48 +556,98 @@ def EvoformerBlockParams(b, is_extra_msa=False):
]
if version not in no_templ:
- tps_blocks = model.template_pair_stack.blocks
+ tps_blocks = model.template_embedder.template_pair_stack.blocks
tps_blocks_params = stacked(
[TemplatePairBlockParams(b) for b in tps_blocks]
- )
- template_param_dict = {
- "template_embedding": {
- "single_template_embedding": {
- "embedding2d": LinearParams(
- model.template_pair_embedder.linear
- ),
- "template_pair_stack": {
- "__layer_stack_no_state": tps_blocks_params,
+ )
+ if (not is_multimer):
+ template_param_dict = {
+ "template_embedding": {
+ "single_template_embedding": {
+ "embedding2d": LinearParams(
+ model.template_embedder.template_pair_embedder.linear
+ ),
+ "template_pair_stack": {
+ "__layer_stack_no_state": tps_blocks_params,
+ },
+ "output_layer_norm": LayerNormParams(
+ model.template_embedder.template_pair_stack.layer_norm
+ ),
},
- "output_layer_norm": LayerNormParams(
- model.template_pair_stack.layer_norm
+ "attention": AttentionParams(model.template_embedder.template_pointwise_att.mha),
+ },
+ "template_single_embedding": LinearParams(
+ model.template_embedder.template_single_embedder.linear_1
+ ),
+ "template_projection": LinearParams(
+ model.template_embedder.template_single_embedder.linear_2
+ ),
+ }
+ else:
+ temp_embedder = model.template_embedder
+ template_param_dict = {
+ "template_embedding": {
+ "single_template_embedding": {
+ "query_embedding_norm": LayerNormParams(
+ temp_embedder.template_pair_embedder.query_embedding_layer_norm
+ ),
+ "template_pair_embedding_0": LinearParams(
+ temp_embedder.template_pair_embedder.dgram_linear
+ ),
+ "template_pair_embedding_1": LinearParams(
+ temp_embedder.template_pair_embedder.pseudo_beta_mask_linear
+ ),
+ "template_pair_embedding_2": LinearParams(
+ temp_embedder.template_pair_embedder.aatype_linear_1
+ ),
+ "template_pair_embedding_3": LinearParams(
+ temp_embedder.template_pair_embedder.aatype_linear_2
+ ),
+ "template_pair_embedding_4": LinearParams(
+ temp_embedder.template_pair_embedder.x_linear
+ ),
+ "template_pair_embedding_5": LinearParams(
+ temp_embedder.template_pair_embedder.y_linear
+ ),
+ "template_pair_embedding_6": LinearParams(
+ temp_embedder.template_pair_embedder.z_linear
+ ),
+ "template_pair_embedding_7": LinearParams(
+ temp_embedder.template_pair_embedder.backbone_mask_linear
+ ),
+ "template_pair_embedding_8": LinearParams(
+ temp_embedder.template_pair_embedder.query_embedding_linear
+ ),
+ "template_embedding_iteration": tps_blocks_params,
+ "output_layer_norm": LayerNormParams(
+ temp_embedder.template_pair_stack.layer_norm
+ ),
+ },
+ "output_linear": LinearParams(
+ temp_embedder.linear_t
),
},
- "attention": AttentionParams(model.template_pointwise_att.mha),
- },
- "template_single_embedding": LinearParams(
- model.template_angle_embedder.linear_1
- ),
- "template_projection": LinearParams(
- model.template_angle_embedder.linear_2
- ),
- }
-
- translations["evoformer"].update(template_param_dict)
+ "template_projection": LinearParams(
+ temp_embedder.template_single_embedder.template_projector,
+ ),
+ "template_single_embedding": LinearParams(
+ temp_embedder.template_single_embedder.template_single_embedder,
+ ),
+ }
- if "_ptm" in version:
+ translations["evoformer"].update(template_param_dict)
+
+ if is_multimer or "_ptm" in version:
translations["predicted_aligned_error_head"] = {
"logits": LinearParams(model.aux_heads.tm.linear)
}
-
return translations
def import_jax_weights_(model, npz_path, version="model_1"):
data = np.load(npz_path)
-
- translations = generate_translation_dict(model, version)
+ translations = generate_translation_dict(model, version, is_multimer=("multimer" in version))
# Flatten keys and insert missing key prefixes
flat = process_translation_dict(translations)
@@ -453,3 +665,45 @@ def import_jax_weights_(model, npz_path, version="model_1"):
# Set weights
assign(flat, data)
+
+
+def convert_deprecated_v1_keys(state_dict):
+ """Update older OpenFold model weight names to match the current model code."""
+
+ replacements = {
+ 'template_angle_embedder': 'template_single_embedder',
+ 'core.msa_transition': 'msa_transition',
+ 'core.outer_product_mean': 'outer_product_mean',
+ 'core.tri_': 'pair_stack.tri_',
+ 'core.pair_transition': 'pair_stack.pair_transition',
+ 'ipa.linear_q_points': 'ipa.linear_q_points.linear',
+ 'ipa.linear_kv_points': 'ipa.linear_kv_points.linear'
+ }
+
+ convert_key_re = re.compile("(%s)" % "|".join(map(re.escape, replacements.keys())))
+
+ converted_state_dict = {}
+ for key, value in state_dict.items():
+ # For each match, look-up replacement value in the dictionary
+ new_key = convert_key_re.sub(lambda m: replacements[m.group()], key)
+
+ # Add prefix for template modules
+ if new_key.startswith('template'):
+ new_key = f'template_embedder.{new_key}'
+
+ converted_state_dict[new_key] = value
+
+ return converted_state_dict
+
+
+def import_openfold_weights_(model, state_dict):
+ """
+ Import model weights. Several parts of the model were refactored in the process
+ of adding support for Multimer. The state dicts of older models are translated
+ to match the refactored model code.
+ """
+ try:
+ model.load_state_dict(state_dict)
+ except RuntimeError:
+ converted_state_dict = convert_deprecated_v1_keys(state_dict)
+ model.load_state_dict(converted_state_dict)
diff --git a/openfold/utils/loss.py b/openfold/utils/loss.py
index ab422aa8..395e3475 100644
--- a/openfold/utils/loss.py
+++ b/openfold/utils/loss.py
@@ -13,25 +13,25 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from functools import partial
-import logging
import ml_collections
import numpy as np
import torch
import torch.nn as nn
-from torch.distributions.bernoulli import Bernoulli
from typing import Dict, Optional, Tuple
from openfold.np import residue_constants
-from openfold.utils import feats
from openfold.utils.rigid_utils import Rotation, Rigid
+from openfold.utils.geometry.vector import Vec3Array, euclidean_distance
+from openfold.utils.all_atom_multimer import get_rc_tensor
from openfold.utils.tensor_utils import (
tree_map,
- tensor_tree_map,
masked_mean,
permute_final_dims,
- batched_gather,
)
+import logging
+from openfold.utils.tensor_utils import tensor_tree_map
+
+logger = logging.getLogger(__name__)
def softmax_cross_entropy(logits, labels):
@@ -87,6 +87,7 @@ def compute_fape(
target_positions: torch.Tensor,
positions_mask: torch.Tensor,
length_scale: float,
+ pair_mask: Optional[torch.Tensor] = None,
l1_clamp_distance: Optional[float] = None,
eps=1e-8,
) -> torch.Tensor:
@@ -108,6 +109,9 @@ def compute_fape(
[*, N_pts] positions mask
length_scale:
Length scale by which the loss is divided
+ pair_mask:
+ [*, N_frames, N_pts] mask to use for
+ separating intra- from inter-chain losses.
l1_clamp_distance:
Cutoff above which distance errors are disregarded
eps:
@@ -134,21 +138,30 @@ def compute_fape(
normed_error = normed_error * frames_mask[..., None]
normed_error = normed_error * positions_mask[..., None, :]
- # FP16-friendly averaging. Roughly equivalent to:
- #
- # norm_factor = (
- # torch.sum(frames_mask, dim=-1) *
- # torch.sum(positions_mask, dim=-1)
- # )
- # normed_error = torch.sum(normed_error, dim=(-1, -2)) / (eps + norm_factor)
- #
- # ("roughly" because eps is necessarily duplicated in the latter)
- normed_error = torch.sum(normed_error, dim=-1)
- normed_error = (
- normed_error / (eps + torch.sum(frames_mask, dim=-1))[..., None]
- )
- normed_error = torch.sum(normed_error, dim=-1)
- normed_error = normed_error / (eps + torch.sum(positions_mask, dim=-1))
+ if pair_mask is not None:
+ normed_error = normed_error * pair_mask
+ normed_error = torch.sum(normed_error, dim=(-1, -2))
+
+ mask = frames_mask[..., None] * positions_mask[..., None, :] * pair_mask
+ norm_factor = torch.sum(mask, dim=(-2, -1))
+
+ normed_error = normed_error / (eps + norm_factor)
+ else:
+ # FP16-friendly averaging. Roughly equivalent to:
+ #
+ # norm_factor = (
+ # torch.sum(frames_mask, dim=-1) *
+ # torch.sum(positions_mask, dim=-1)
+ # )
+ # normed_error = torch.sum(normed_error, dim=(-1, -2)) / (eps + norm_factor)
+ #
+ # ("roughly" because eps is necessarily duplicated in the latter)
+ normed_error = torch.sum(normed_error, dim=-1)
+ normed_error = (
+ normed_error / (eps + torch.sum(frames_mask, dim=-1))[..., None]
+ )
+ normed_error = torch.sum(normed_error, dim=-1)
+ normed_error = normed_error / (eps + torch.sum(positions_mask, dim=-1))
return normed_error
@@ -157,13 +170,19 @@ def backbone_loss(
backbone_rigid_tensor: torch.Tensor,
backbone_rigid_mask: torch.Tensor,
traj: torch.Tensor,
+ pair_mask: Optional[torch.Tensor] = None,
use_clamped_fape: Optional[torch.Tensor] = None,
clamp_distance: float = 10.0,
loss_unit_distance: float = 10.0,
eps: float = 1e-4,
**kwargs,
) -> torch.Tensor:
- pred_aff = Rigid.from_tensor_7(traj)
+ ### need to check if the traj belongs to 4*4 matrix or a tensor_7
+ if traj.shape[-1] == 7:
+ pred_aff = Rigid.from_tensor_7(traj)
+ elif traj.shape[-1] == 4:
+ pred_aff = Rigid.from_tensor_4x4(traj)
+
pred_aff = Rigid(
Rotation(rot_mats=pred_aff.get_rots().get_rot_mats(), quats=None),
pred_aff.get_trans(),
@@ -184,6 +203,7 @@ def backbone_loss(
pred_aff.get_trans(),
gt_aff[None].get_trans(),
backbone_rigid_mask[None],
+ pair_mask=pair_mask,
l1_clamp_distance=clamp_distance,
length_scale=loss_unit_distance,
eps=eps,
@@ -196,6 +216,7 @@ def backbone_loss(
pred_aff.get_trans(),
gt_aff[None].get_trans(),
backbone_rigid_mask[None],
+ pair_mask=pair_mask,
l1_clamp_distance=None,
length_scale=loss_unit_distance,
eps=eps,
@@ -226,10 +247,10 @@ def sidechain_loss(
**kwargs,
) -> torch.Tensor:
renamed_gt_frames = (
- 1.0 - alt_naming_is_better[..., None, None, None]
- ) * rigidgroups_gt_frames + alt_naming_is_better[
- ..., None, None, None
- ] * rigidgroups_alt_gt_frames
+ 1.0 - alt_naming_is_better[..., None, None, None]
+ ) * rigidgroups_gt_frames + alt_naming_is_better[
+ ..., None, None, None
+ ] * rigidgroups_alt_gt_frames
# Steamroll the inputs
sidechain_frames = sidechain_frames[-1]
@@ -253,6 +274,7 @@ def sidechain_loss(
sidechain_atom_pos,
renamed_atom14_gt_positions,
renamed_atom14_gt_exists,
+ pair_mask=None,
l1_clamp_distance=clamp_distance,
length_scale=length_scale,
eps=eps,
@@ -266,10 +288,28 @@ def fape_loss(
batch: Dict[str, torch.Tensor],
config: ml_collections.ConfigDict,
) -> torch.Tensor:
- bb_loss = backbone_loss(
- traj=out["sm"]["frames"],
- **{**batch, **config.backbone},
- )
+ traj = out["sm"]["frames"]
+ asym_id = batch.get("asym_id")
+ if asym_id is not None:
+ intra_chain_mask = (asym_id[..., None] == asym_id[..., None, :]).to(dtype=traj.dtype)
+ intra_chain_bb_loss = backbone_loss(
+ traj=traj,
+ pair_mask=intra_chain_mask,
+ **{**batch, **config.intra_chain_backbone},
+ )
+ interface_bb_loss = backbone_loss(
+ traj=traj,
+ pair_mask=1. - intra_chain_mask,
+ **{**batch, **config.interface_backbone},
+ )
+ weighted_bb_loss = (intra_chain_bb_loss * config.intra_chain_backbone.weight
+ + interface_bb_loss * config.interface_backbone.weight)
+ else:
+ bb_loss = backbone_loss(
+ traj=traj,
+ **{**batch, **config.backbone},
+ )
+ weighted_bb_loss = bb_loss * config.backbone.weight
sc_loss = sidechain_loss(
out["sm"]["sidechain_frames"],
@@ -277,8 +317,8 @@ def fape_loss(
**{**batch, **config.sidechain},
)
- loss = config.backbone.weight * bb_loss + config.sidechain.weight * sc_loss
-
+ loss = weighted_bb_loss + config.sidechain.weight * sc_loss
+
# Average over the batch dimension
loss = torch.mean(loss)
@@ -340,7 +380,7 @@ def supervised_chi_loss(
(true_chi_shifted - pred_angles) ** 2, dim=-1
)
sq_chi_error = torch.minimum(sq_chi_error, sq_chi_error_shifted)
-
+
# The ol' switcheroo
sq_chi_error = sq_chi_error.permute(
*range(len(sq_chi_error.shape))[1:-2], 0, -2, -1
@@ -452,7 +492,7 @@ def lddt_ca(
ca_pos = residue_constants.atom_order["CA"]
all_atom_pred_pos = all_atom_pred_pos[..., ca_pos, :]
all_atom_positions = all_atom_positions[..., ca_pos, :]
- all_atom_mask = all_atom_mask[..., ca_pos : (ca_pos + 1)] # keep dim
+ all_atom_mask = all_atom_mask[..., ca_pos: (ca_pos + 1)] # keep dim
return lddt(
all_atom_pred_pos,
@@ -482,18 +522,21 @@ def lddt_loss(
ca_pos = residue_constants.atom_order["CA"]
all_atom_pred_pos = all_atom_pred_pos[..., ca_pos, :]
all_atom_positions = all_atom_positions[..., ca_pos, :]
- all_atom_mask = all_atom_mask[..., ca_pos : (ca_pos + 1)] # keep dim
+ all_atom_mask = all_atom_mask[..., ca_pos: (ca_pos + 1)] # keep dim
score = lddt(
- all_atom_pred_pos,
- all_atom_positions,
- all_atom_mask,
- cutoff=cutoff,
+ all_atom_pred_pos,
+ all_atom_positions,
+ all_atom_mask,
+ cutoff=cutoff,
eps=eps
)
- score = score.detach()
+ # TODO: Remove after initial pipeline testing
+ score = torch.nan_to_num(score, nan=torch.nanmean(score))
+ score[score < 0] = 0
+ score = score.detach()
bin_index = torch.floor(score * no_bins).long()
bin_index = torch.clamp(bin_index, max=(no_bins - 1))
lddt_ca_one_hot = torch.nn.functional.one_hot(
@@ -533,7 +576,7 @@ def distogram_loss(
device=logits.device,
)
boundaries = boundaries ** 2
-
+
dists = torch.sum(
(pseudo_beta[..., None, :] - pseudo_beta[..., None, :, :]) ** 2,
dim=-1,
@@ -627,6 +670,8 @@ def compute_predicted_aligned_error(
def compute_tm(
logits: torch.Tensor,
residue_weights: Optional[torch.Tensor] = None,
+ asym_id: Optional[torch.Tensor] = None,
+ interface: bool = False,
max_bin: int = 31,
no_bins: int = 64,
eps: float = 1e-8,
@@ -649,11 +694,26 @@ def compute_tm(
tm_per_bin = 1.0 / (1 + (bin_centers ** 2) / (d0 ** 2))
predicted_tm_term = torch.sum(probs * tm_per_bin, dim=-1)
- normed_residue_mask = residue_weights / (eps + residue_weights.sum())
+ n = residue_weights.shape[-1]
+ pair_mask = residue_weights.new_ones((n, n), dtype=torch.int32)
+ if interface and (asym_id is not None):
+ if len(asym_id.shape) > 1:
+ assert len(asym_id.shape) <= 2
+ batch_size = asym_id.shape[0]
+ pair_mask = residue_weights.new_ones((batch_size, n, n), dtype=torch.int32)
+ pair_mask *= (asym_id[..., None] != asym_id[..., None, :]).to(dtype=pair_mask.dtype)
+
+ predicted_tm_term *= pair_mask
+
+ pair_residue_weights = pair_mask * (
+ residue_weights[..., None, :] * residue_weights[..., :, None]
+ )
+ denom = eps + torch.sum(pair_residue_weights, dim=-1, keepdims=True)
+ normed_residue_mask = pair_residue_weights / denom
per_alignment = torch.sum(predicted_tm_term * normed_residue_mask, dim=-1)
weighted = per_alignment * residue_weights
-
+
argmax = (weighted == torch.max(weighted)).nonzero()[0]
return per_alignment[tuple(argmax)]
@@ -671,7 +731,11 @@ def tm_loss(
eps=1e-8,
**kwargs,
):
- pred_affine = Rigid.from_tensor_7(final_affine_tensor)
+ # first check whether this is a tensor_7 or tensor_4*4
+ if final_affine_tensor.shape[-1] == 7:
+ pred_affine = Rigid.from_tensor_7(final_affine_tensor)
+ elif final_affine_tensor.shape[-1] == 4:
+ pred_affine = Rigid.from_tensor_4x4(final_affine_tensor)
backbone_rigid = Rigid.from_tensor_4x4(backbone_rigid_tensor)
def _points(affine):
@@ -709,7 +773,7 @@ def _points(affine):
(resolution >= min_resolution) & (resolution <= max_resolution)
)
- # Average over the loss dimension
+ # Average over the batch dimension
loss = torch.mean(loss)
return loss
@@ -771,19 +835,20 @@ def between_residue_bond_loss(
# The C-N bond to proline has slightly different length because of the ring.
next_is_proline = aatype[..., 1:] == residue_constants.resname_to_idx["PRO"]
gt_length = (
- ~next_is_proline
- ) * residue_constants.between_res_bond_length_c_n[
- 0
- ] + next_is_proline * residue_constants.between_res_bond_length_c_n[
- 1
- ]
+ ~next_is_proline
+ ) * residue_constants.between_res_bond_length_c_n[
+ 0
+ ] + next_is_proline * residue_constants.between_res_bond_length_c_n[
+ 1
+ ]
gt_stddev = (
- ~next_is_proline
- ) * residue_constants.between_res_bond_length_stddev_c_n[
- 0
- ] + next_is_proline * residue_constants.between_res_bond_length_stddev_c_n[
- 1
- ]
+ ~next_is_proline
+ ) * residue_constants.between_res_bond_length_stddev_c_n[
+ 0
+ ] + next_is_proline * residue_constants.between_res_bond_length_stddev_c_n[
+ 1
+ ]
+
c_n_bond_length_error = torch.sqrt(eps + (c_n_bond_length - gt_length) ** 2)
c_n_loss_per_residue = torch.nn.functional.relu(
c_n_bond_length_error - tolerance_factor_soft * gt_stddev
@@ -879,6 +944,7 @@ def between_residue_clash_loss(
atom14_atom_exists: torch.Tensor,
atom14_atom_radius: torch.Tensor,
residue_index: torch.Tensor,
+ asym_id: Optional[torch.Tensor] = None,
overlap_tolerance_soft=1.5,
overlap_tolerance_hard=1.5,
eps=1e-10,
@@ -908,7 +974,6 @@ def between_residue_clash_loss(
shape (N, 14)
"""
fp_type = atom14_pred_positions.dtype
-
# Create the distance matrix.
# (N, N, 14, 14)
dists = torch.sqrt(
@@ -954,9 +1019,13 @@ def between_residue_clash_loss(
)
n_one_hot = n_one_hot.type(fp_type)
- neighbour_mask = (
- residue_index[..., :, None, None, None] + 1
- ) == residue_index[..., None, :, None, None]
+ neighbour_mask = (residue_index[..., :, None] + 1) == residue_index[..., None, :]
+
+ if asym_id is not None:
+ neighbour_mask = neighbour_mask & (asym_id[..., :, None] == asym_id[..., None, :])
+
+ neighbour_mask = neighbour_mask[..., None, None]
+
c_n_bonds = (
neighbour_mask
* c_one_hot[..., None, None, :, None]
@@ -998,7 +1067,7 @@ def between_residue_clash_loss(
# Compute the per atom loss sum.
# shape (N, 14)
per_atom_loss_sum = torch.sum(dists_to_low_error, dim=(-4, -2)) + torch.sum(
- dists_to_low_error, axis=(-3, -1)
+ dists_to_low_error, dim=(-3, -1)
)
# Compute the hard clash mask.
@@ -1007,17 +1076,20 @@ def between_residue_clash_loss(
dists < (dists_lower_bound - overlap_tolerance_hard)
)
+ per_atom_num_clash = torch.sum(clash_mask, dim=(-4, -2)) + torch.sum(clash_mask, dim=(-3, -1))
+
# Compute the per atom clash.
# shape (N, 14)
per_atom_clash_mask = torch.maximum(
- torch.amax(clash_mask, axis=(-4, -2)),
- torch.amax(clash_mask, axis=(-3, -1)),
+ torch.amax(clash_mask, dim=(-4, -2)),
+ torch.amax(clash_mask, dim=(-3, -1)),
)
return {
"mean_loss": mean_loss, # shape ()
"per_atom_loss_sum": per_atom_loss_sum, # shape (N, 14)
"per_atom_clash_mask": per_atom_clash_mask, # shape (N, 14)
+ "per_atom_num_clash": per_atom_num_clash # shape (N, 14)
}
@@ -1097,6 +1169,8 @@ def within_residue_violations(
(dists < atom14_dists_lower_bound) | (dists > atom14_dists_upper_bound)
)
+ per_atom_num_clash = torch.sum(violations, dim=-2) + torch.sum(violations, dim=-1)
+
# Compute the per atom violations.
per_atom_violations = torch.maximum(
torch.max(violations, dim=-2)[0], torch.max(violations, axis=-1)[0]
@@ -1105,6 +1179,7 @@ def within_residue_violations(
return {
"per_atom_loss_sum": per_atom_loss_sum,
"per_atom_violations": per_atom_violations,
+ "per_atom_num_clash": per_atom_num_clash
}
@@ -1134,11 +1209,24 @@ def find_structural_violations(
residue_constants.van_der_waals_radius[name[0]]
for name in residue_constants.atom_types
]
+
atomtype_radius = atom14_pred_positions.new_tensor(atomtype_radius)
- atom14_atom_radius = (
- batch["atom14_atom_exists"]
- * atomtype_radius[batch["residx_atom14_to_atom37"]]
- )
+
+ # TODO: Consolidate monomer/multimer modes
+ asym_id = batch.get("asym_id")
+ if asym_id is not None:
+ residx_atom14_to_atom37 = get_rc_tensor(
+ residue_constants.RESTYPE_ATOM14_TO_ATOM37, batch["aatype"]
+ )
+ atom14_atom_radius = (
+ batch["atom14_atom_exists"]
+ * atomtype_radius[residx_atom14_to_atom37.long()]
+ )
+ else:
+ atom14_atom_radius = (
+ batch["atom14_atom_exists"]
+ * atomtype_radius[batch["residx_atom14_to_atom37"]]
+ )
# Compute the between residue clash loss.
between_residue_clashes = between_residue_clash_loss(
@@ -1146,6 +1234,7 @@ def find_structural_violations(
atom14_atom_exists=batch["atom14_atom_exists"],
atom14_atom_radius=atom14_atom_radius,
residue_index=batch["residue_index"],
+ asym_id=asym_id,
overlap_tolerance_soft=clash_overlap_tolerance,
overlap_tolerance_hard=clash_overlap_tolerance,
)
@@ -1208,6 +1297,9 @@ def find_structural_violations(
"clashes_per_atom_clash_mask": between_residue_clashes[
"per_atom_clash_mask"
], # (N, 14)
+ "clashes_per_atom_num_clash": between_residue_clashes[
+ "per_atom_num_clash"
+ ], # (N, 14)
},
"within_residues": {
"per_atom_loss_sum": residue_violations[
@@ -1216,6 +1308,9 @@ def find_structural_violations(
"per_atom_violations": residue_violations[
"per_atom_violations"
], # (N, 14),
+ "per_atom_num_clash": residue_violations[
+ "per_atom_num_clash"
+ ], # (N, 14)
},
"total_per_residue_violations_mask": per_residue_violations_mask, # (N)
}
@@ -1268,8 +1363,8 @@ def extreme_ca_ca_distance_violations(
eps + torch.sum((this_ca_pos - next_ca_pos) ** 2, dim=-1)
)
violations = (
- ca_ca_distance - residue_constants.ca_ca
- ) > max_angstrom_tolerance
+ ca_ca_distance - residue_constants.ca_ca
+ ) > max_angstrom_tolerance
mask = this_ca_mask * next_ca_mask * has_no_gap_mask
mean = masked_mean(mask, violations, -1)
return mean
@@ -1337,15 +1432,21 @@ def compute_violation_metrics_np(
def violation_loss(
violations: Dict[str, torch.Tensor],
atom14_atom_exists: torch.Tensor,
+ average_clashes: bool = False,
eps=1e-6,
**kwargs,
) -> torch.Tensor:
num_atoms = torch.sum(atom14_atom_exists)
- l_clash = torch.sum(
- violations["between_residues"]["clashes_per_atom_loss_sum"]
- + violations["within_residues"]["per_atom_loss_sum"]
- )
- l_clash = l_clash / (eps + num_atoms)
+
+ per_atom_clash = (violations["between_residues"]["clashes_per_atom_loss_sum"] +
+ violations["within_residues"]["per_atom_loss_sum"])
+
+ if average_clashes:
+ num_clash = (violations["between_residues"]["clashes_per_atom_num_clash"] +
+ violations["within_residues"]["per_atom_num_clash"])
+ per_atom_clash = per_atom_clash / (num_clash + eps)
+
+ l_clash = torch.sum(per_atom_clash) / (eps + num_atoms)
loss = (
violations["between_residues"]["bonds_c_n_loss_mean"]
+ violations["between_residues"]["angles_ca_c_n_loss_mean"]
@@ -1449,16 +1550,16 @@ def compute_renamed_ground_truth(
alt_naming_is_better = (alt_per_res_lddt < per_res_lddt).type(fp_type)
renamed_atom14_gt_positions = (
- 1.0 - alt_naming_is_better[..., None, None]
- ) * atom14_gt_positions + alt_naming_is_better[
- ..., None, None
- ] * atom14_alt_gt_positions
+ 1.0 - alt_naming_is_better[..., None, None]
+ ) * atom14_gt_positions + alt_naming_is_better[
+ ..., None, None
+ ] * atom14_alt_gt_positions
renamed_atom14_gt_mask = (
- 1.0 - alt_naming_is_better[..., None]
- ) * atom14_gt_exists + alt_naming_is_better[..., None] * batch[
- "atom14_alt_gt_exists"
- ]
+ 1.0 - alt_naming_is_better[..., None]
+ ) * atom14_gt_exists + alt_naming_is_better[..., None] * batch[
+ "atom14_alt_gt_exists"
+ ]
return {
"alt_naming_is_better": alt_naming_is_better,
@@ -1481,17 +1582,17 @@ def experimentally_resolved_loss(
loss = torch.sum(errors * atom37_atom_exists, dim=-1)
loss = loss / (eps + torch.sum(atom37_atom_exists, dim=(-1, -2)).unsqueeze(-1))
loss = torch.sum(loss, dim=-1)
-
+
loss = loss * (
(resolution >= min_resolution) & (resolution <= max_resolution)
)
loss = torch.mean(loss)
-
+
return loss
-def masked_msa_loss(logits, true_msa, bert_mask, eps=1e-8, **kwargs):
+def masked_msa_loss(logits, true_msa, bert_mask, num_classes, eps=1e-8, **kwargs):
"""
Computes BERT-style masked MSA loss. Implements subsection 1.9.9.
@@ -1503,7 +1604,7 @@ def masked_msa_loss(logits, true_msa, bert_mask, eps=1e-8, **kwargs):
Masked MSA loss
"""
errors = softmax_cross_entropy(
- logits, torch.nn.functional.one_hot(true_msa, num_classes=23)
+ logits, torch.nn.functional.one_hot(true_msa, num_classes=num_classes)
)
# FP16-friendly averaging. Equivalent to:
@@ -1524,13 +1625,75 @@ def masked_msa_loss(logits, true_msa, bert_mask, eps=1e-8, **kwargs):
return loss
+def chain_center_of_mass_loss(
+ all_atom_pred_pos: torch.Tensor,
+ all_atom_positions: torch.Tensor,
+ all_atom_mask: torch.Tensor,
+ asym_id: torch.Tensor,
+ clamp_distance: float = -4.0,
+ weight: float = 0.05,
+ eps: float = 1e-10, **kwargs
+) -> torch.Tensor:
+ """
+ Computes chain centre-of-mass loss. Implements section 2.5, eqn 1 in the Multimer paper.
+
+ Args:
+ all_atom_pred_pos:
+ [*, N_pts, 37, 3] All-atom predicted atom positions
+ all_atom_positions:
+ [*, N_pts, 37, 3] Ground truth all-atom positions
+ all_atom_mask:
+ [*, N_pts, 37] All-atom positions mask
+ asym_id:
+ [*, N_pts] Chain asym IDs
+ clamp_distance:
+ Cutoff above which distance errors are disregarded
+ weight:
+ Weight for loss
+ eps:
+ Small value used to regularize denominators
+ Returns:
+ [*] loss tensor
+ """
+ ca_pos = residue_constants.atom_order["CA"]
+ all_atom_pred_pos = all_atom_pred_pos[..., ca_pos, :]
+ all_atom_positions = all_atom_positions[..., ca_pos, :]
+ all_atom_mask = all_atom_mask[..., ca_pos: (ca_pos + 1)] # keep dim
+
+ one_hot = torch.nn.functional.one_hot(asym_id.long()).to(dtype=all_atom_mask.dtype)
+ one_hot = one_hot * all_atom_mask
+ chain_pos_mask = one_hot.transpose(-2, -1)
+ chain_exists = torch.any(chain_pos_mask, dim=-1).to(dtype=all_atom_positions.dtype)
+
+ def get_chain_center_of_mass(pos):
+ center_sum = (chain_pos_mask[..., None] * pos[..., None, :, :]).sum(dim=-2)
+ centers = center_sum / (torch.sum(chain_pos_mask, dim=-1, keepdim=True) + eps)
+ return Vec3Array.from_array(centers)
+
+ pred_centers = get_chain_center_of_mass(all_atom_pred_pos) # [B, NC, 3]
+ true_centers = get_chain_center_of_mass(all_atom_positions) # [B, NC, 3]
+
+ pred_dists = euclidean_distance(pred_centers[..., None, :], pred_centers[..., :, None], epsilon=eps)
+ true_dists = euclidean_distance(true_centers[..., None, :], true_centers[..., :, None], epsilon=eps)
+ losses = torch.clamp((weight * (pred_dists - true_dists - clamp_distance)), max=0) ** 2
+ loss_mask = chain_exists[..., :, None] * chain_exists[..., None, :]
+
+ loss = masked_mean(loss_mask, losses, dim=(-1, -2))
+ return loss
+
+
class AlphaFoldLoss(nn.Module):
"""Aggregation of the various losses described in the supplement"""
+
def __init__(self, config):
super(AlphaFoldLoss, self).__init__()
self.config = config
- def forward(self, out, batch, _return_breakdown=False):
+ def loss(self, out, batch, _return_breakdown=False):
+ """
+ Rename previous forward() as loss()
+ so that can be reused in the subclass
+ """
if "violation" not in out.keys():
out["violation"] = find_structural_violations(
batch,
@@ -1576,31 +1739,36 @@ def forward(self, out, batch, _return_breakdown=False):
),
"violation": lambda: violation_loss(
out["violation"],
- **batch,
+ **{**batch, **self.config.violation},
),
}
- if(self.config.tm.enabled):
+ if self.config.tm.enabled:
loss_fns["tm"] = lambda: tm_loss(
logits=out["tm_logits"],
**{**batch, **out, **self.config.tm},
)
+ if self.config.chain_center_of_mass.enabled:
+ loss_fns["chain_center_of_mass"] = lambda: chain_center_of_mass_loss(
+ all_atom_pred_pos=out["final_atom_positions"],
+ **{**batch, **self.config.chain_center_of_mass},
+ )
+
cum_loss = 0.
losses = {}
for loss_name, loss_fn in loss_fns.items():
weight = self.config[loss_name].weight
loss = loss_fn()
- if(torch.isnan(loss) or torch.isinf(loss)):
- #for k,v in batch.items():
- # if(torch.any(torch.isnan(v)) or torch.any(torch.isinf(v))):
+ if torch.isnan(loss) or torch.isinf(loss):
+ # for k,v in batch.items():
+ # if torch.any(torch.isnan(v)) or torch.any(torch.isinf(v)):
# logging.warning(f"{k}: is nan")
- #logging.warning(f"{loss_name}: {loss}")
+ # logging.warning(f"{loss_name}: {loss}")
logging.warning(f"{loss_name} loss is NaN. Skipping...")
loss = loss.new_tensor(0., requires_grad=True)
cum_loss = cum_loss + weight * loss
losses[loss_name] = loss.detach().clone()
-
losses["unscaled_loss"] = cum_loss.detach().clone()
# Scale the loss by the square root of the minimum of the crop size and
@@ -1611,7 +1779,15 @@ def forward(self, out, batch, _return_breakdown=False):
losses["loss"] = cum_loss.detach().clone()
- if(not _return_breakdown):
+ if not _return_breakdown:
return cum_loss
-
+
return cum_loss, losses
+
+ def forward(self, out, batch, _return_breakdown=False):
+ if not _return_breakdown:
+ cum_loss = self.loss(out, batch, _return_breakdown)
+ return cum_loss
+ else:
+ cum_loss, losses = self.loss(out, batch, _return_breakdown)
+ return cum_loss, losses
diff --git a/openfold/utils/multi_chain_permutation.py b/openfold/utils/multi_chain_permutation.py
new file mode 100644
index 00000000..c950862a
--- /dev/null
+++ b/openfold/utils/multi_chain_permutation.py
@@ -0,0 +1,421 @@
+import logging
+import random
+import torch
+
+from openfold.np import residue_constants as rc
+
+logger = logging.getLogger(__name__)
+
+
+def compute_rmsd(
+ true_atom_pos: torch.Tensor,
+ pred_atom_pos: torch.Tensor,
+ atom_mask: torch.Tensor = None,
+ eps: float = 1e-6,
+) -> torch.Tensor:
+ sq_diff = torch.square(true_atom_pos - pred_atom_pos).sum(dim=-1, keepdim=False)
+ if atom_mask is not None:
+ sq_diff = torch.masked_select(sq_diff, atom_mask.to(sq_diff.device))
+ msd = torch.mean(sq_diff)
+ msd = torch.nan_to_num(msd, nan=1e8)
+ return torch.sqrt(msd + eps) # prevent sqrt 0
+
+
+def kabsch_rotation(P, Q):
+ """
+ Calculate the best rotation that minimises the RMSD between P and Q.
+
+ The optimal rotation matrix was calculated using Kabsch algorithm:
+ https://en.wikipedia.org/wiki/Kabsch_algorithm
+
+ Args:
+ P: [N * 3] Nres is the number of atoms and each row corresponds to the atom's x,y,z coordinates
+ Q: [N * 3] the same dimension as P
+
+ return:
+ A 3*3 rotation matrix
+ """
+ assert P.shape == torch.Size([Q.shape[0], Q.shape[1]])
+
+ # Firstly, compute SVD of P.T * Q
+ u, _, vt = torch.linalg.svd(torch.matmul(P.to(torch.float32).T,
+ Q.to(torch.float32)))
+ # Then construct s matrix
+ s = torch.eye(P.shape[1], device=P.device)
+ # correct the rotation matrix to ensure a right-handed coordinate
+ s[-1, -1] = torch.sign(torch.linalg.det(torch.matmul(u, vt)))
+ # finally compute the rotation matrix
+ r_opt = torch.matmul(torch.matmul(u, s), vt)
+ assert r_opt.shape == torch.Size([3,3])
+ return r_opt.to(device=P.device, dtype=P.dtype)
+
+
+def get_optimal_transform(
+ src_atoms: torch.Tensor,
+ tgt_atoms: torch.Tensor,
+ mask: torch.Tensor = None,
+):
+ """
+ src_atoms: predicted CA positions, shape:[num_res,3]
+ tgt_atoms: ground-truth CA positions, shape:[num_res,3]
+ mask: a vector of boolean values, shape:[num_res]
+ """
+ assert src_atoms.shape == tgt_atoms.shape, (src_atoms.shape, tgt_atoms.shape)
+ assert src_atoms.shape[-1] == 3
+ if mask is not None:
+ assert len(mask.shape) == 1, "mask should have the shape of [num_res]"
+ if torch.isnan(src_atoms).any() or torch.isinf(src_atoms).any():
+ #
+ # sometimes using fake test inputs generates NaN in the predicted atom positions
+ # #
+ logging.warning(f"src_atom has nan or inf")
+ src_atoms = torch.nan_to_num(src_atoms, nan=0.0, posinf=1.0, neginf=1.0)
+
+ if mask is not None:
+ assert mask.dtype == torch.bool
+ assert mask.shape[-1] == src_atoms.shape[-2]
+ if mask.sum() == 0:
+ src_atoms = torch.zeros((1, 3), device=src_atoms.device, dtype=src_atoms.dtype)
+ tgt_atoms = src_atoms
+ else:
+ src_atoms = src_atoms[mask, :]
+ tgt_atoms = tgt_atoms[mask, :]
+ src_center = src_atoms.mean(-2, keepdim=True, dtype=src_atoms.dtype)
+ tgt_center = tgt_atoms.mean(-2, keepdim=True, dtype=src_atoms.dtype)
+ r = kabsch_rotation(src_atoms, tgt_atoms)
+
+ x = tgt_center - src_center @ r
+ return r, x
+
+
+def get_least_asym_entity_or_longest_length(batch, input_asym_id):
+ """
+ First check how many subunit(s) one sequence has. Select the subunit that is less
+ common, e.g. if the protein was AABBB then select one of the A as anchor
+
+ If there is a tie, e.g. AABB, first check which sequence is the longer/longest,
+ then choose one of the corresponding subunits as anchor
+
+ Args:
+ batch: in this function batch is the full ground truth features
+ input_asym_id: A list of asym_ids that are in the cropped input features
+
+ Return:
+ anchor_gt_asym_id: Tensor(int) selected ground truth asym_id
+ anchor_pred_asym_ids: list(Tensor(int)) a list of all possible pred anchor candidates
+ """
+ entity_2_asym_list = get_entity_2_asym_list(batch)
+ unique_entity_ids = torch.unique(batch["entity_id"])
+ entity_asym_count = {}
+ entity_length = {}
+
+ for entity_id in unique_entity_ids:
+ asym_ids = torch.unique(batch["asym_id"][batch["entity_id"] == entity_id])
+
+ # Make sure some asym IDs associated with ground truth entity ID exist in cropped prediction
+ asym_ids_in_pred = [a for a in asym_ids if a in input_asym_id]
+ if not asym_ids_in_pred:
+ continue
+
+ entity_asym_count[int(entity_id)] = len(asym_ids)
+
+ # Calculate entity length
+ entity_mask = (batch["entity_id"] == entity_id)
+ entity_length[int(entity_id)] = entity_mask.sum().item()
+
+ min_asym_count = min(entity_asym_count.values())
+ least_asym_entities = [entity for entity, count in entity_asym_count.items() if count == min_asym_count]
+
+ # If multiple entities have the least asym_id count, return those with the longest length
+ if len(least_asym_entities) > 1:
+ max_length = max([entity_length[entity] for entity in least_asym_entities])
+ least_asym_entities = [entity for entity in least_asym_entities if entity_length[entity] == max_length]
+
+ # If still multiple entities, return a random one
+ if len(least_asym_entities) > 1:
+ least_asym_entities = [random.choice(least_asym_entities)]
+
+ assert len(least_asym_entities) == 1
+ least_asym_entities = least_asym_entities[0]
+
+ anchor_gt_asym_id = random.choice(entity_2_asym_list[least_asym_entities])
+ anchor_pred_asym_ids = [asym_id for asym_id in entity_2_asym_list[least_asym_entities] if asym_id in input_asym_id]
+
+ return anchor_gt_asym_id, anchor_pred_asym_ids
+
+
+def greedy_align(
+ batch,
+ per_asym_residue_index,
+ entity_2_asym_list,
+ pred_ca_pos,
+ pred_ca_mask,
+ true_ca_poses,
+ true_ca_masks,
+):
+ """
+ Implement Algorithm 4 in the Supplementary Information of AlphaFold-Multimer paper:
+ Evans,R et al., 2022 Protein complex prediction with AlphaFold-Multimer, bioRxiv 2021.10.04.463034; doi: https://doi.org/10.1101/2021.10.04.463034
+ """
+ used = [False for _ in range(len(true_ca_poses))]
+ align = []
+ unique_asym_ids = [i for i in torch.unique(batch["asym_id"]) if i != 0]
+ for cur_asym_id in unique_asym_ids:
+ i = int(cur_asym_id - 1)
+ asym_mask = batch["asym_id"] == cur_asym_id
+ cur_entity_ids = batch["entity_id"][asym_mask][0]
+ best_rmsd = torch.inf
+ best_idx = None
+ cur_asym_list = entity_2_asym_list[int(cur_entity_ids)]
+ cur_residue_index = per_asym_residue_index[int(cur_asym_id)]
+ cur_pred_pos = pred_ca_pos[asym_mask]
+ cur_pred_mask = pred_ca_mask[asym_mask]
+ for next_asym_id in cur_asym_list:
+ j = int(next_asym_id - 1)
+ if not used[j]: # possible candidate
+ cropped_pos = torch.index_select(true_ca_poses[j], 1, cur_residue_index)
+ mask = torch.index_select(true_ca_masks[j], 1, cur_residue_index)
+ rmsd = compute_rmsd(
+ torch.squeeze(cropped_pos, 0), torch.squeeze(cur_pred_pos, 0),
+ (cur_pred_mask * mask).bool()
+ )
+ if rmsd is not None and rmsd < best_rmsd:
+ best_rmsd = rmsd
+ best_idx = j
+
+ assert best_idx is not None
+ used[best_idx] = True
+ align.append((i, best_idx))
+ return align
+
+
+def pad_features(feature_tensor, nres_pad, pad_dim):
+ """Pad input feature tensor"""
+ pad_shape = list(feature_tensor.shape)
+ pad_shape[pad_dim] = nres_pad
+ padding_tensor = feature_tensor.new_zeros(pad_shape, device=feature_tensor.device)
+ return torch.concat((feature_tensor, padding_tensor), dim=pad_dim)
+
+
+def merge_labels(per_asym_residue_index, labels, align, original_nres):
+ """
+ Merge ground truth labels according to the permutation results
+
+ labels: list of original ground truth feats
+ align: list of tuples, each entry specify the corresponding label of the asym.
+
+ modified based on UniFold:
+ https://github.com/dptech-corp/Uni-Fold/blob/b1c89a2cebd4e4ee4c47b4e443f92beeb9138fbb/unifold/losses/chain_align.py#L176C1-L176C1
+ """
+ outs = {}
+ for k, v in labels[0].items():
+ cur_out = {}
+ for i, j in align:
+ label = labels[j][k]
+ # to 1-based
+ cur_residue_index = per_asym_residue_index[i + 1]
+ if len(v.shape) <= 1 or "template" in k or "row_mask" in k:
+ continue
+ else:
+ dimension_to_merge = 1
+ cur_out[i] = label.index_select(dimension_to_merge, cur_residue_index)
+ cur_out = [x[1] for x in sorted(cur_out.items())]
+ if len(cur_out) > 0:
+ new_v = torch.concat(cur_out, dim=dimension_to_merge)
+ # below check whether padding is needed
+ if new_v.shape[dimension_to_merge] != original_nres:
+ nres_pad = original_nres - new_v.shape[dimension_to_merge]
+ new_v = pad_features(new_v, nres_pad, pad_dim=dimension_to_merge)
+ outs[k] = new_v
+ return outs
+
+
+def split_ground_truth_labels(gt_features):
+ """
+ Splits ground truth features according to chains
+
+ Returns:
+ a list of feature dictionaries with only necessary ground truth features
+ required to finish multi-chain permutation
+ """
+ unique_asym_ids, asym_id_counts = torch.unique(gt_features["asym_id"], sorted=True, return_counts=True)
+ n_res = gt_features["asym_id"].shape[-1]
+
+ def split_dim(shape):
+ return next(iter(i for i, size in enumerate(shape) if size == n_res), None)
+
+ labels = list(map(dict, zip(*[[(k, v) for v in torch.split(v_all, asym_id_counts.tolist(),
+ dim=split_dim(v_all.shape))]
+ for k, v_all in gt_features.items()
+ if n_res in v_all.shape])))
+ return labels
+
+
+def get_per_asym_residue_index(features):
+ unique_asym_ids = [i for i in torch.unique(features["asym_id"]) if i != 0]
+ per_asym_residue_index = {}
+ for cur_asym_id in unique_asym_ids:
+ asym_mask = (features["asym_id"] == cur_asym_id).bool()
+ per_asym_residue_index[int(cur_asym_id)] = torch.masked_select(features["residue_index"], asym_mask)
+
+ return per_asym_residue_index
+
+
+def get_entity_2_asym_list(batch):
+ """
+ Generates a dictionary mapping unique entity IDs to lists of unique asymmetry IDs (asym_id) for each entity.
+
+ Args:
+ batch (dict): A dictionary containing data batches, including "entity_id" and "asym_id" tensors.
+
+ Returns:
+ entity_2_asym_list (dict): A dictionary where keys are unique entity IDs, and values are lists of unique asymmetry IDs
+ associated with each entity.
+ """
+ entity_2_asym_list = {}
+ unique_entity_ids = torch.unique(batch["entity_id"])
+ for cur_ent_id in unique_entity_ids:
+ ent_mask = batch["entity_id"] == cur_ent_id
+ cur_asym_id = torch.unique(batch["asym_id"][ent_mask])
+ entity_2_asym_list[int(cur_ent_id)] = cur_asym_id
+ return entity_2_asym_list
+
+
+def calculate_input_mask(true_ca_masks, anchor_gt_idx, anchor_gt_residue,
+ asym_mask, pred_ca_mask):
+ """
+ Calculate an input mask for downstream optimal transformation computation
+
+ Args:
+ true_ca_masks (Tensor): ca mask from ground truth.
+ anchor_gt_idx (Tensor): The index of selected ground truth anchor.
+ asym_mask (Tensor): Boolean tensor indicating which regions are selected predicted anchor.
+ pred_ca_mask (Tensor): ca mask from predicted structure.
+
+ Returns:
+ input_mask (Tensor): A boolean mask
+ """
+ pred_ca_mask = torch.squeeze(pred_ca_mask, 0)
+ asym_mask = torch.squeeze(asym_mask, 0)
+ anchor_pred_mask = pred_ca_mask[asym_mask]
+ anchor_true_mask = torch.index_select(true_ca_masks[anchor_gt_idx], 1, anchor_gt_residue)
+ input_mask = (anchor_true_mask * anchor_pred_mask).bool()
+ return input_mask
+
+
+def calculate_optimal_transform(true_ca_poses,
+ anchor_gt_idx, anchor_gt_residue,
+ true_ca_masks, pred_ca_mask,
+ asym_mask,
+ pred_ca_pos):
+ input_mask = calculate_input_mask(true_ca_masks,
+ anchor_gt_idx,
+ anchor_gt_residue,
+ asym_mask,
+ pred_ca_mask)
+ input_mask = torch.squeeze(input_mask, 0)
+ pred_ca_pos = torch.squeeze(pred_ca_pos, 0)
+ asym_mask = torch.squeeze(asym_mask, 0)
+ anchor_true_pos = torch.index_select(true_ca_poses[anchor_gt_idx], 1, anchor_gt_residue)
+ anchor_pred_pos = pred_ca_pos[asym_mask]
+ r, x = get_optimal_transform(
+ anchor_pred_pos, torch.squeeze(anchor_true_pos, 0),
+ mask=input_mask
+ )
+
+ return r, x
+
+
+def compute_permutation_alignment(out, features, ground_truth):
+ """
+ A class method that first permutate chains in ground truth first
+ before calculating the loss.
+
+ Details are described in Section 7.3 in the Supplementary of AlphaFold-Multimer paper:
+ https://www.biorxiv.org/content/10.1101/2021.10.04.463034v2
+ """
+ unique_asym_ids = set(torch.unique(features['asym_id']).tolist())
+ unique_asym_ids.discard(0) # Remove padding asym_id
+ is_monomer = len(unique_asym_ids) == 1
+
+ per_asym_residue_index = get_per_asym_residue_index(features)
+
+ if is_monomer:
+ best_align = list(enumerate(range(len(per_asym_residue_index))))
+ return best_align, per_asym_residue_index
+
+ best_rmsd = float('inf')
+ best_align = None
+ # First select anchors from predicted structures and ground truths
+ anchor_gt_asym, anchor_pred_asym_ids = get_least_asym_entity_or_longest_length(ground_truth,
+ features['asym_id'])
+ entity_2_asym_list = get_entity_2_asym_list(ground_truth)
+ labels = split_ground_truth_labels(ground_truth)
+ assert isinstance(labels, list)
+ anchor_gt_idx = int(anchor_gt_asym) - 1
+ # Then calculate optimal transform by aligning anchors
+ ca_idx = rc.atom_order["CA"]
+ pred_ca_pos = out["final_atom_positions"][..., ca_idx, :] # [bsz, nres, 3]
+ pred_ca_mask = out["final_atom_mask"][..., ca_idx].to(dtype=pred_ca_pos.dtype) # [bsz, nres]
+
+ true_ca_poses = [
+ l["all_atom_positions"][..., ca_idx, :] for l in labels
+ ] # list([nres, 3])
+ true_ca_masks = [
+ l["all_atom_mask"][..., ca_idx].long() for l in labels
+ ] # list([nres,])
+ for candidate_pred_anchor in anchor_pred_asym_ids:
+ asym_mask = (features["asym_id"] == candidate_pred_anchor).bool()
+ anchor_gt_residue = per_asym_residue_index[candidate_pred_anchor.item()]
+ r, x = calculate_optimal_transform(true_ca_poses,
+ anchor_gt_idx,
+ anchor_gt_residue,
+ true_ca_masks,
+ pred_ca_mask,
+ asym_mask,
+ pred_ca_pos)
+ aligned_true_ca_poses = [ca.to(r.dtype) @ r + x for ca in true_ca_poses] # apply transforms
+ align = greedy_align(
+ features,
+ per_asym_residue_index,
+ entity_2_asym_list,
+ pred_ca_pos,
+ pred_ca_mask,
+ aligned_true_ca_poses,
+ true_ca_masks,
+ )
+ merged_labels = merge_labels(per_asym_residue_index, labels, align,
+ original_nres=features['aatype'].shape[-1])
+ rmsd = compute_rmsd(
+ true_atom_pos=merged_labels['all_atom_positions'][..., ca_idx, :].to(r.dtype) @ r + x,
+ pred_atom_pos=pred_ca_pos,
+ atom_mask=(pred_ca_mask * merged_labels['all_atom_mask'][..., ca_idx].long()).bool())
+ if rmsd < best_rmsd:
+ best_rmsd = rmsd
+ best_align = align
+
+ return best_align, per_asym_residue_index
+
+
+def multi_chain_permutation_align(out, features, ground_truth):
+ """Compute multi-chain permutation alignment.
+
+ Args:
+ out: The output of model.forward()
+ features: Input features
+ ground_truth: Ground truth features
+ """
+
+ labels = split_ground_truth_labels(ground_truth)
+
+ # Then permute ground truth chains before calculating the loss
+ align, per_asym_residue_index = compute_permutation_alignment(out=out,
+ features=features,
+ ground_truth=ground_truth)
+
+ # reorder ground truth labels according to permutation results
+ labels = merge_labels(per_asym_residue_index, labels, align,
+ original_nres=features['aatype'].shape[-1])
+
+ features.update(labels)
+ return features
diff --git a/openfold/utils/rigid_utils.py b/openfold/utils/rigid_utils.py
index 24cfbd75..6220c608 100644
--- a/openfold/utils/rigid_utils.py
+++ b/openfold/utils/rigid_utils.py
@@ -978,6 +978,16 @@ def device(self) -> torch.device:
"""
return self._trans.device
+ @property
+ def dtype(self) -> torch.dtype:
+ """
+ Returns the dtype of the Rigid tensors.
+
+ Returns:
+ The dtype of the Rigid tensors
+ """
+ return self._rots.dtype
+
def get_rots(self) -> Rotation:
"""
Getter for the rotation.
diff --git a/openfold/utils/script_utils.py b/openfold/utils/script_utils.py
index 0d098a67..facc103d 100644
--- a/openfold/utils/script_utils.py
+++ b/openfold/utils/script_utils.py
@@ -12,6 +12,7 @@
from openfold.np.relax import relax
from openfold.utils.import_weights import (
import_jax_weights_,
+ import_openfold_weights_
)
from pytorch_lightning.utilities.deepspeed import (
@@ -90,7 +91,7 @@ def load_models_from_command_line(config, model_device, openfold_checkpoint_path
ckpt_path,
)
d = torch.load(ckpt_path)
- model.load_state_dict(d["ema"]["params"])
+ import_openfold_weights_(model=model, state_dict=d["ema"]["params"])
else:
ckpt_path = path
d = torch.load(ckpt_path)
@@ -98,7 +99,7 @@ def load_models_from_command_line(config, model_device, openfold_checkpoint_path
if "ema" in d:
# The public weights have had this done to them already
d = d["ema"]["params"]
- model.load_state_dict(d)
+ import_openfold_weights_(model=model, state_dict=d)
model = model.to(model_device)
logger.info(
@@ -122,7 +123,7 @@ def parse_fasta(data):
][1:]
tags, seqs = lines[::2], lines[1::2]
- tags = [t.split()[0] for t in tags]
+ tags = [re.split('\W| \|', t)[0] for t in tags]
return tags, seqs
@@ -219,7 +220,7 @@ def prep_output(out, batch, feature_dict, feature_processor, config_preset, mult
features=batch,
result=out,
b_factors=plddt_b_factors,
- chain_index=chain_index,
+ remove_leading_feature_dimension=False,
remark=remark,
parents=template_domain_names,
parents_chain_index=template_chain_index,
diff --git a/run_pretrained_openfold.py b/run_pretrained_openfold.py
index bdd985eb..10f991d7 100644
--- a/run_pretrained_openfold.py
+++ b/run_pretrained_openfold.py
@@ -17,24 +17,19 @@
import math
import numpy as np
import os
-
-from openfold.utils.script_utils import load_models_from_command_line, parse_fasta, run_model, prep_output, \
- update_timings, relax_protein
+import pickle
+import random
+import time
logging.basicConfig()
logger = logging.getLogger(__file__)
logger.setLevel(level=logging.INFO)
-import pickle
-
-import random
-import time
import torch
-
torch_versions = torch.__version__.split(".")
torch_major_version = int(torch_versions[0])
torch_minor_version = int(torch_versions[1])
-if(
+if (
torch_major_version > 1 or
(torch_major_version == 1 and torch_minor_version >= 12)
):
@@ -45,16 +40,16 @@
from openfold.config import model_config
from openfold.data import templates, feature_pipeline, data_pipeline
-from openfold.np import residue_constants, protein
-import openfold.np.relax.relax as relax
-
-from openfold.utils.tensor_utils import (
- tensor_tree_map,
-)
+from openfold.data.tools import hhsearch, hmmsearch
+from openfold.np import protein
+from openfold.utils.script_utils import (load_models_from_command_line, parse_fasta, run_model,
+ prep_output, relax_protein)
+from openfold.utils.tensor_utils import tensor_tree_map
from openfold.utils.trace_utils import (
pad_feature_dict_seq,
trace_model_,
)
+
from scripts.precompute_embeddings import EmbeddingGenerator
from scripts.utils import add_data_args
@@ -69,18 +64,30 @@ def precompute_alignments(tags, seqs, alignment_dir, args):
fp.write(f">{tag}\n{seq}")
local_alignment_dir = os.path.join(alignment_dir, tag)
- if(args.use_precomputed_alignments is None and not os.path.isdir(local_alignment_dir)):
+
+ if args.use_precomputed_alignments is None:
logger.info(f"Generating alignments for {tag}...")
- os.makedirs(local_alignment_dir)
+ os.makedirs(local_alignment_dir, exist_ok=True)
+
+ if "multimer" in args.config_preset:
+ template_searcher = hmmsearch.Hmmsearch(
+ binary_path=args.hmmsearch_binary_path,
+ hmmbuild_binary_path=args.hmmbuild_binary_path,
+ database_path=args.pdb_seqres_database_path,
+ )
+ else:
+ template_searcher = hhsearch.HHSearch(
+ binary_path=args.hhsearch_binary_path,
+ databases=[args.pdb70_database_path],
+ )
# In seqemb mode, use AlignmentRunner only to generate templates
if args.use_single_seq_mode:
alignment_runner = data_pipeline.AlignmentRunner(
jackhmmer_binary_path=args.jackhmmer_binary_path,
- hhsearch_binary_path=args.hhsearch_binary_path,
uniref90_database_path=args.uniref90_database_path,
- pdb70_database_path=args.pdb70_database_path,
+ template_searcher=template_searcher,
no_cpus=args.cpus,
)
embedding_generator = EmbeddingGenerator()
@@ -89,14 +96,17 @@ def precompute_alignments(tags, seqs, alignment_dir, args):
alignment_runner = data_pipeline.AlignmentRunner(
jackhmmer_binary_path=args.jackhmmer_binary_path,
hhblits_binary_path=args.hhblits_binary_path,
- hhsearch_binary_path=args.hhsearch_binary_path,
uniref90_database_path=args.uniref90_database_path,
mgnify_database_path=args.mgnify_database_path,
bfd_database_path=args.bfd_database_path,
+ uniref30_database_path=args.uniref30_database_path,
uniclust30_database_path=args.uniclust30_database_path,
- pdb70_database_path=args.pdb70_database_path,
- no_cpus=args.cpus,
+ uniprot_database_path=args.uniprot_database_path,
+ template_searcher=template_searcher,
+ use_small_bfd=args.bfd_database_path is None,
+ no_cpus=args.cpus
)
+
alignment_runner.run(
tmp_fasta_path, local_alignment_dir
)
@@ -133,6 +143,14 @@ def generate_feature_dict(
alignment_dir=local_alignment_dir,
seqemb_mode=args.use_single_seq_mode,
)
+ elif "multimer" in args.config_preset:
+ with open(tmp_fasta_path, "w") as fp:
+ fp.write(
+ '\n'.join([f">{tag}\n{seq}" for tag, seq in zip(tags, seqs)])
+ )
+ feature_dict = data_processor.process_fasta(
+ fasta_path=tmp_fasta_path, alignment_dir=alignment_dir,
+ )
else:
with open(tmp_fasta_path, "w") as fp:
fp.write(
@@ -147,6 +165,7 @@ def generate_feature_dict(
return feature_dict
+
def list_files_with_extensions(dir, extensions):
return [f for f in os.listdir(dir) if f.endswith(extensions)]
@@ -157,31 +176,49 @@ def main(args):
if args.config_preset.startswith("seq"):
args.use_single_seq_mode = True
+
config = model_config(args.config_preset, long_sequence_inference=args.long_sequence_inference)
- if(args.trace_model):
- if(not config.data.predict.fixed_size):
+ if args.trace_model:
+ if not config.data.predict.fixed_size:
raise ValueError(
"Tracing requires that fixed_size mode be enabled in the config"
)
- template_featurizer = templates.TemplateHitFeaturizer(
- mmcif_dir=args.template_mmcif_dir,
- max_template_date=args.max_template_date,
- max_hits=config.data.predict.max_templates,
- kalign_binary_path=args.kalign_binary_path,
- release_dates_path=args.release_dates_path,
- obsolete_pdbs_path=args.obsolete_pdbs_path
- )
+ is_multimer = "multimer" in args.config_preset
+
+ if is_multimer:
+ template_featurizer = templates.HmmsearchHitFeaturizer(
+ mmcif_dir=args.template_mmcif_dir,
+ max_template_date=args.max_template_date,
+ max_hits=config.data.predict.max_templates,
+ kalign_binary_path=args.kalign_binary_path,
+ release_dates_path=args.release_dates_path,
+ obsolete_pdbs_path=args.obsolete_pdbs_path
+ )
+ else:
+ template_featurizer = templates.HhsearchHitFeaturizer(
+ mmcif_dir=args.template_mmcif_dir,
+ max_template_date=args.max_template_date,
+ max_hits=config.data.predict.max_templates,
+ kalign_binary_path=args.kalign_binary_path,
+ release_dates_path=args.release_dates_path,
+ obsolete_pdbs_path=args.obsolete_pdbs_path
+ )
data_processor = data_pipeline.DataPipeline(
template_featurizer=template_featurizer,
)
+ if is_multimer:
+ data_processor = data_pipeline.DataPipelineMultimer(
+ monomer_data_pipeline=data_processor,
+ )
+
output_dir_base = args.output_dir
random_seed = args.data_random_seed
if random_seed is None:
- random_seed = random.randrange(2**32)
+ random_seed = random.randrange(2 ** 32)
np.random.seed(random_seed)
torch.manual_seed(random_seed + 1)
@@ -198,10 +235,19 @@ def main(args):
seq_list = []
for fasta_file in list_files_with_extensions(args.fasta_dir, (".fasta", ".fa")):
# Gather input sequences
- with open(os.path.join(args.fasta_dir, fasta_file), "r") as fp:
+ fasta_path = os.path.join(args.fasta_dir, fasta_file)
+ with open(fasta_path, "r") as fp:
data = fp.read()
tags, seqs = parse_fasta(data)
+
+ if not is_multimer and len(tags) != 1:
+ print(
+ f"{fasta_path} contains more than one sequence but "
+ f"multimer mode is not enabled. Skipping..."
+ )
+ continue
+
# assert len(tags) == len(set(tags)), "All FASTA tags must be unique"
tag = '-'.join(tags)
@@ -217,6 +263,7 @@ def main(args):
args.openfold_checkpoint_path,
args.jax_param_path,
args.output_dir)
+
for model, output_directory in model_generator:
cur_tracing_interval = 0
for (tag, tags), seqs in sorted_targets:
@@ -228,7 +275,7 @@ def main(args):
precompute_alignments(tags, seqs, alignment_dir, args)
feature_dict = feature_dicts.get(tag, None)
- if(feature_dict is None):
+ if feature_dict is None:
feature_dict = generate_feature_dict(
tags,
seqs,
@@ -237,7 +284,7 @@ def main(args):
args,
)
- if(args.trace_model):
+ if args.trace_model:
n = feature_dict["aatype"].shape[-2]
rounded_seqlen = round_up_seqlen(n)
feature_dict = pad_feature_dict_seq(
@@ -247,16 +294,16 @@ def main(args):
feature_dicts[tag] = feature_dict
processed_feature_dict = feature_processor.process_features(
- feature_dict, mode='predict',
+ feature_dict, mode='predict', is_multimer=is_multimer
)
processed_feature_dict = {
- k:torch.as_tensor(v, device=args.model_device)
- for k,v in processed_feature_dict.items()
+ k: torch.as_tensor(v, device=args.model_device)
+ for k, v in processed_feature_dict.items()
}
- if(args.trace_model):
- if(rounded_seqlen > cur_tracing_interval):
+ if args.trace_model:
+ if rounded_seqlen > cur_tracing_interval:
logger.info(
f"Tracing model at {rounded_seqlen} residues..."
)
@@ -305,7 +352,8 @@ def main(args):
if not args.skip_relaxation:
# Relax the prediction.
logger.info(f"Running relaxation on {unrelaxed_output_path}...")
- relax_protein(config, args.model_device, unrelaxed_protein, output_directory, output_name, args.cif_output)
+ relax_protein(config, args.model_device, unrelaxed_protein, output_directory, output_name,
+ args.cif_output)
if args.save_outputs:
output_dict_path = os.path.join(
@@ -407,13 +455,13 @@ def main(args):
add_data_args(parser)
args = parser.parse_args()
- if(args.jax_param_path is None and args.openfold_checkpoint_path is None):
+ if args.jax_param_path is None and args.openfold_checkpoint_path is None:
args.jax_param_path = os.path.join(
"openfold", "resources", "params",
"params_" + args.config_preset + ".npz"
)
- if(args.model_device == "cpu" and torch.cuda.is_available()):
+ if args.model_device == "cpu" and torch.cuda.is_available():
logging.warning(
"""The model is being run on CPU. Consider specifying
--model_device for better performance"""
diff --git a/scripts/__init__.py b/scripts/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/scripts/convert_of_weights_to_jax.py b/scripts/convert_of_weights_to_jax.py
index 8665fdd3..9a16d0d3 100644
--- a/scripts/convert_of_weights_to_jax.py
+++ b/scripts/convert_of_weights_to_jax.py
@@ -26,6 +26,7 @@
ParamType,
generate_translation_dict,
process_translation_dict,
+ import_openfold_weights_
)
from openfold.utils.tensor_utils import tree_map
@@ -63,7 +64,7 @@ def main(args):
config = model_config(args.config_preset)
model = AlphaFold(config)
- model.load_state_dict(d)
+ import_openfold_weights_(model=model, state_dict=d)
translation = generate_translation_dict(model, args.config_preset)
translation = process_translation_dict(translation)
diff --git a/scripts/data_dir_to_fasta.py b/scripts/data_dir_to_fasta.py
index f7271973..5769ddfb 100644
--- a/scripts/data_dir_to_fasta.py
+++ b/scripts/data_dir_to_fasta.py
@@ -1,7 +1,8 @@
import argparse
import logging
import os
-
+import string
+from collections import defaultdict
from openfold.data import mmcif_parsing
from openfold.np import protein, residue_constants
@@ -22,7 +23,7 @@ def main(args):
if(mmcif.mmcif_object is None):
logging.warning(f'Failed to parse {fname}...')
if(args.raise_errors):
- raise list(mmcif.errors.values())[0]
+ raise Exception(list(mmcif.errors.values())[0])
else:
continue
@@ -31,6 +32,25 @@ def main(args):
chain_id = '_'.join([basename, chain])
fasta.append(f">{chain_id}")
fasta.append(seq)
+ elif(ext == ".pdb"):
+ with open(fpath, 'r') as fp:
+ pdb_str = fp.read()
+
+ protein_object = protein.from_pdb_string(pdb_str)
+ aatype = protein_object.aatype
+ chain_index = protein_object.chain_index
+
+ last_chain_index = chain_index[0]
+ chain_dict = defaultdict(list)
+ for i in range(aatype.shape[0]):
+ chain_dict[chain_index[i]].append(residue_constants.restypes_with_x[aatype[i]])
+
+ chain_tags = string.ascii_uppercase
+ for chain, seq in chain_dict.items():
+ chain_id = '_'.join([basename, chain_tags[chain]])
+ fasta.append(f">{chain_id}")
+ fasta.append(''.join(seq))
+
elif(ext == ".core"):
with open(fpath, 'r') as fp:
core_str = fp.read()
diff --git a/scripts/deepspeed_inference_test.py b/scripts/deepspeed_inference_test.py
new file mode 100644
index 00000000..cbe14f57
--- /dev/null
+++ b/scripts/deepspeed_inference_test.py
@@ -0,0 +1,54 @@
+import copy
+import os
+
+import torch
+import deepspeed
+
+
+local_rank = int(os.getenv('LOCAL_RANK', '0'))
+world_size = int(os.getenv('WORLD_SIZE', '1'))
+
+class Model(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+
+ self.ml = torch.nn.ModuleList()
+ for _ in range(4000):
+ self.ml.append(torch.nn.Linear(500, 500))
+
+ def forward(self, batch):
+ for i, l in enumerate(self.ml):
+ # print(f"{i}: {l.weight.device}")
+ batch = l(batch)
+
+ return batch
+
+
+class DummyDataset(torch.utils.data.Dataset):
+ def __init__(self):
+ self.batch = torch.rand(500, 500)
+
+ def __getitem__(self, idx):
+ return copy.deepcopy(self.batch)
+
+ def __len__(self):
+ return 1000
+
+dd = DummyDataset()
+dl = torch.utils.data.DataLoader(dd)
+example = next(iter(dl)).to(f"cuda:{local_rank}")
+
+model = Model()
+model = model.to(f"cuda:{local_rank}")
+
+model = deepspeed.init_inference(
+ model,
+ mp_size=world_size,
+ checkpoint=None,
+ replace_method=None,
+ #replace_method="auto"
+)
+
+out = model(example)
+#if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
+# print(out)
diff --git a/scripts/download_alphafold_dbs.sh b/scripts/download_alphafold_dbs.sh
index 885430d9..87570b08 100755
--- a/scripts/download_alphafold_dbs.sh
+++ b/scripts/download_alphafold_dbs.sh
@@ -56,10 +56,16 @@ bash "${SCRIPT_DIR}/download_pdb70.sh" "${DOWNLOAD_DIR}"
echo "Downloading PDB mmCIF files..."
bash "${SCRIPT_DIR}/download_pdb_mmcif.sh" "${DOWNLOAD_DIR}"
-echo "Downloading Uniclust30..."
-bash "${SCRIPT_DIR}/download_uniclust30.sh" "${DOWNLOAD_DIR}"
+echo "Downloading Uniref30..."
+bash "${SCRIPT_DIR}/download_uniref30.sh" "${DOWNLOAD_DIR}"
echo "Downloading Uniref90..."
bash "${SCRIPT_DIR}/download_uniref90.sh" "${DOWNLOAD_DIR}"
+echo "Downloading UniProt..."
+bash "${SCRIPT_DIR}/download_uniprot.sh" "${DOWNLOAD_DIR}"
+
+echo "Downloading PDB SeqRes..."
+bash "${SCRIPT_DIR}/download_pdb_seqres.sh" "${DOWNLOAD_DIR}"
+
echo "All data downloaded."
diff --git a/scripts/download_alphafold_params.sh b/scripts/download_alphafold_params.sh
index 51966e6f..ff96ff4f 100755
--- a/scripts/download_alphafold_params.sh
+++ b/scripts/download_alphafold_params.sh
@@ -31,7 +31,7 @@ fi
DOWNLOAD_DIR="$1"
ROOT_DIR="${DOWNLOAD_DIR}/params"
-SOURCE_URL="https://storage.googleapis.com/alphafold/alphafold_params_2022-01-19.tar"
+SOURCE_URL="https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar"
BASENAME=$(basename "${SOURCE_URL}")
mkdir --parents "${ROOT_DIR}"
diff --git a/scripts/download_mgnify.sh b/scripts/download_mgnify.sh
index 739956f3..8424015f 100755
--- a/scripts/download_mgnify.sh
+++ b/scripts/download_mgnify.sh
@@ -32,8 +32,8 @@ fi
DOWNLOAD_DIR="$1"
ROOT_DIR="${DOWNLOAD_DIR}/mgnify"
# Mirror of:
-# ftp://ftp.ebi.ac.uk/pub/databases/metagenomics/peptide_database/2018_12/mgy_clusters.fa.gz
-SOURCE_URL="https://storage.googleapis.com/alphafold-databases/casp14_versions/mgy_clusters_2018_12.fa.gz"
+# ftp://ftp.ebi.ac.uk/pub/databases/metagenomics/peptide_database/2022_05/mgy_clusters.fa.gz
+SOURCE_URL="https://storage.googleapis.com/alphafold-databases/v2.3/mgy_clusters_2022_05.fa.gz"
BASENAME=$(basename "${SOURCE_URL}")
mkdir --parents "${ROOT_DIR}"
diff --git a/scripts/download_pdb_seqres.sh b/scripts/download_pdb_seqres.sh
new file mode 100755
index 00000000..5e42b5dd
--- /dev/null
+++ b/scripts/download_pdb_seqres.sh
@@ -0,0 +1,42 @@
+#!/bin/bash
+#
+# 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.
+#
+# Downloads and unzips the PDB SeqRes database for AlphaFold.
+#
+# Usage: bash download_pdb_seqres.sh /path/to/download/directory
+set -e
+
+if [[ $# -eq 0 ]]; then
+ echo "Error: download directory must be provided as an input argument."
+ exit 1
+fi
+
+if ! command -v aria2c &> /dev/null ; then
+ echo "Error: aria2c could not be found. Please install aria2c (sudo apt install aria2)."
+ exit 1
+fi
+
+DOWNLOAD_DIR="$1"
+ROOT_DIR="${DOWNLOAD_DIR}/pdb_seqres"
+SOURCE_URL="ftp://ftp.wwpdb.org/pub/pdb/derived_data/pdb_seqres.txt"
+BASENAME=$(basename "${SOURCE_URL}")
+
+mkdir --parents "${ROOT_DIR}"
+aria2c "${SOURCE_URL}" --dir="${ROOT_DIR}"
+
+# Keep only protein sequences.
+grep --after-context=1 --no-group-separator '>.* mol:protein' "${ROOT_DIR}/pdb_seqres.txt" > "${ROOT_DIR}/pdb_seqres_filtered.txt"
+mv "${ROOT_DIR}/pdb_seqres_filtered.txt" "${ROOT_DIR}/pdb_seqres.txt"
diff --git a/scripts/download_uniprot.sh b/scripts/download_uniprot.sh
new file mode 100755
index 00000000..e815e3c3
--- /dev/null
+++ b/scripts/download_uniprot.sh
@@ -0,0 +1,55 @@
+#!/bin/bash
+#
+# 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.
+#
+# Downloads, unzips and merges the SwissProt and TrEMBL databases for
+# AlphaFold-Multimer.
+#
+# Usage: bash download_uniprot.sh /path/to/download/directory
+set -e
+
+if [[ $# -eq 0 ]]; then
+ echo "Error: download directory must be provided as an input argument."
+ exit 1
+fi
+
+if ! command -v aria2c &> /dev/null ; then
+ echo "Error: aria2c could not be found. Please install aria2c (sudo apt install aria2)."
+ exit 1
+fi
+
+DOWNLOAD_DIR="$1"
+ROOT_DIR="${DOWNLOAD_DIR}/uniprot"
+
+TREMBL_SOURCE_URL="ftp://ftp.ebi.ac.uk/pub/databases/uniprot/current_release/knowledgebase/complete/uniprot_trembl.fasta.gz"
+TREMBL_BASENAME=$(basename "${TREMBL_SOURCE_URL}")
+TREMBL_UNZIPPED_BASENAME="${TREMBL_BASENAME%.gz}"
+
+SPROT_SOURCE_URL="ftp://ftp.ebi.ac.uk/pub/databases/uniprot/current_release/knowledgebase/complete/uniprot_sprot.fasta.gz"
+SPROT_BASENAME=$(basename "${SPROT_SOURCE_URL}")
+SPROT_UNZIPPED_BASENAME="${SPROT_BASENAME%.gz}"
+
+mkdir --parents "${ROOT_DIR}"
+aria2c "${TREMBL_SOURCE_URL}" --dir="${ROOT_DIR}"
+aria2c "${SPROT_SOURCE_URL}" --dir="${ROOT_DIR}"
+pushd "${ROOT_DIR}"
+gunzip "${ROOT_DIR}/${TREMBL_BASENAME}"
+gunzip "${ROOT_DIR}/${SPROT_BASENAME}"
+
+# Concatenate TrEMBL and SwissProt, rename to uniprot and clean up.
+cat "${ROOT_DIR}/${SPROT_UNZIPPED_BASENAME}" >> "${ROOT_DIR}/${TREMBL_UNZIPPED_BASENAME}"
+mv "${ROOT_DIR}/${TREMBL_UNZIPPED_BASENAME}" "${ROOT_DIR}/uniprot.fasta"
+rm "${ROOT_DIR}/${SPROT_UNZIPPED_BASENAME}"
+popd
diff --git a/scripts/download_uniref30.sh b/scripts/download_uniref30.sh
index 12e52072..8177695d 100755
--- a/scripts/download_uniref30.sh
+++ b/scripts/download_uniref30.sh
@@ -30,10 +30,15 @@ if ! command -v aria2c &> /dev/null ; then
fi
DOWNLOAD_DIR="$1"
-ROOT_DIR="${DOWNLOAD_DIR}"
-SOURCE_URL="http://wwwuser.gwdg.de/~compbiol/colabfold/uniref30_2103.tar.gz"
+ROOT_DIR="${DOWNLOAD_DIR}/uniref30"
+# Mirror of:
+# https://wwwuser.gwdg.de/~compbiol/uniclust/2021_03/UniRef30_2021_03.tar.gz
+SOURCE_URL="https://storage.googleapis.com/alphafold-databases/v2.3/UniRef30_2021_03.tar.gz"
BASENAME=$(basename "${SOURCE_URL}")
mkdir --parents "${ROOT_DIR}"
aria2c "${SOURCE_URL}" --dir="${ROOT_DIR}" -x 4 --check-certificate=false
-gunzip "${ROOT_DIR}/${BASENAME}"
+tar --extract --verbose --file="${ROOT_DIR}/${BASENAME}" \
+ --directory="${ROOT_DIR}"
+rm "${ROOT_DIR}/${BASENAME}"
+
diff --git a/scripts/flatten_roda.sh b/scripts/flatten_roda.sh
index d5c4fa7d..074736a6 100755
--- a/scripts/flatten_roda.sh
+++ b/scripts/flatten_roda.sh
@@ -26,7 +26,12 @@ mkdir -p "${ALIGNMENT_DIR}"
for chain_dir in $(ls "${RODA_DIR}"); do
CHAIN_DIR_PATH="${RODA_DIR}/${chain_dir}"
for subdir in $(ls "${CHAIN_DIR_PATH}"); do
- if [[ $subdir = "pdb" ]] || [[ $subdir = "cif" ]]; then
+ if [[ ! -d "$subdir" ]]; then
+ echo "$subdir is not directory"
+ continue
+ elif [[ -z $(ls "${subdir}")]]; then
+ continue
+ elif [[ $subdir = "pdb" ]] || [[ $subdir = "cif" ]]; then
mv "${CHAIN_DIR_PATH}/${subdir}"/* "${DATA_DIR}"
else
CHAIN_ALIGNMENT_DIR="${ALIGNMENT_DIR}/${chain_dir}"
diff --git a/scripts/generate_alphafold_feature_dict.py b/scripts/generate_alphafold_feature_dict.py
index 7c7be611..bde5c76e 100644
--- a/scripts/generate_alphafold_feature_dict.py
+++ b/scripts/generate_alphafold_feature_dict.py
@@ -2,35 +2,62 @@
import os
import pickle
-from alphafold.data import pipeline, templates
+from alphafold.data import pipeline, pipeline_multimer, templates
+from alphafold.data.tools import hmmsearch, hhsearch
from scripts.utils import add_data_args
def main(args):
- template_featurizer = templates.TemplateHitFeaturizer(
- mmcif_dir=args.mmcif_dir,
- max_template_date=args.max_template_date,
- max_hits=20,
- kalign_binary_path=args.kalign_binary_path,
- release_dates_path=None,
- obsolete_pdbs_path=args.obsolete_pdbs_path,
- )
+ if (args.multimer):
+ template_searcher = hmmsearch.Hmmsearch(
+ binary_path=args.hmmsearch_binary_path,
+ hmmbuild_binary_path=args.hmmbuild_binary_path,
+ database_path=args.pdb_seqres_database_path,
+ )
+
+ template_featurizer = templates.HmmsearchHitFeaturizer(
+ mmcif_dir=args.template_mmcif_dir,
+ max_template_date=args.max_template_date,
+ max_hits=20,
+ kalign_binary_path=args.kalign_binary_path,
+ release_dates_path=args.release_dates_path,
+ obsolete_pdbs_path=args.obsolete_pdbs_path
+ )
+ else:
+ template_searcher = hhsearch.HHSearch(
+ binary_path=args.hhsearch_binary_path,
+ databases=[args.pdb70_database_path],
+ )
+
+ template_featurizer = templates.HhsearchHitFeaturizer(
+ mmcif_dir=args.template_mmcif_dir,
+ max_template_date=args.max_template_date,
+ max_hits=20,
+ kalign_binary_path=args.kalign_binary_path,
+ release_dates_path=None,
+ obsolete_pdbs_path=args.obsolete_pdbs_path
+ )
data_pipeline = pipeline.DataPipeline(
jackhmmer_binary_path=args.jackhmmer_binary_path,
hhblits_binary_path=args.hhblits_binary_path,
- hhsearch_binary_path=args.hhsearch_binary_path,
uniref90_database_path=args.uniref90_database_path,
mgnify_database_path=args.mgnify_database_path,
bfd_database_path=args.bfd_database_path,
- uniclust30_database_path=args.uniclust30_database_path,
- pdb70_database_path=args.pdb70_database_path,
+ uniref30_database_path=args.uniref30_database_path,
small_bfd_database_path=None,
template_featurizer=template_featurizer,
+ template_searcher=template_searcher,
use_small_bfd=False,
)
+ if (args.multimer):
+ data_pipeline = pipeline_multimer.DataPipeline(
+ monomer_data_pipeline=data_pipeline,
+ jackhmmer_binary_path=args.jackhmmer_binary_path,
+ uniprot_database_path=args.uniprot_database_path)
+
feature_dict = data_pipeline.process(
input_fasta_path=args.fasta_path,
msa_output_dir=args.output_dir,
@@ -44,6 +71,7 @@ def main(args):
parser.add_argument("fasta_path", type=str)
parser.add_argument("mmcif_dir", type=str)
parser.add_argument("output_dir", type=str)
+ parser.add_argument("--multimer", action='store_true')
add_data_args(parser)
args = parser.parse_args()
diff --git a/scripts/generate_chain_data_cache.py b/scripts/generate_chain_data_cache.py
index d26968e0..72d9c2da 100644
--- a/scripts/generate_chain_data_cache.py
+++ b/scripts/generate_chain_data_cache.py
@@ -4,10 +4,11 @@
import logging
from multiprocessing import Pool
import os
-
+import string
import sys
sys.path.append(".") # an innocent hack to get this to run from the top level
+from collections import defaultdict
from tqdm import tqdm
from openfold.data.mmcif_parsing import parse
@@ -49,20 +50,27 @@ def parse_file(
pdb_string = fp.read()
protein_object = protein.from_pdb_string(pdb_string, None)
+ aatype = protein_object.aatype
+ chain_index = protein_object.chain_index
- chain_dict = {}
- chain_dict["seq"] = residue_constants.aatype_to_str_sequence(
- protein_object.aatype,
- )
- chain_dict["resolution"] = 0.
-
- if(chain_cluster_size_dict is not None):
- cluster_size = chain_cluster_size_dict.get(
- full_name.upper(), -1
- )
- chain_dict["cluster_size"] = cluster_size
+ chain_dict = defaultdict(list)
+ for i in range(aatype.shape[0]):
+ chain_dict[chain_index[i]].append(residue_constants.restypes_with_x[aatype[i]])
- out = {file_id: chain_dict}
+ out = {}
+ chain_tags = string.ascii_uppercase
+ for chain, seq in chain_dict.items():
+ full_name = "_".join([file_id, chain_tags[chain]])
+ out[full_name] = {}
+ local_data = out[full_name]
+ local_data["resolution"] = 0.
+ local_data["seq"] = ''.join(seq)
+
+ if(chain_cluster_size_dict is not None):
+ cluster_size = chain_cluster_size_dict.get(
+ full_name.upper(), -1
+ )
+ local_data["cluster_size"] = cluster_size
return out
diff --git a/scripts/generate_mmcif_cache.py b/scripts/generate_mmcif_cache.py
index d7a90d84..4a79ba5f 100644
--- a/scripts/generate_mmcif_cache.py
+++ b/scripts/generate_mmcif_cache.py
@@ -13,7 +13,7 @@
from openfold.data.mmcif_parsing import parse
-def parse_file(f, args):
+def parse_file(f, args, chain_cluster_size_dict=None):
with open(os.path.join(args.mmcif_dir, f), "r") as fp:
mmcif_string = fp.read()
file_id = os.path.splitext(f)[0]
@@ -28,6 +28,18 @@ def parse_file(f, args):
local_data["release_date"] = mmcif.header["release_date"]
chain_ids, seqs = list(zip(*mmcif.chain_to_seqres.items()))
+
+ if chain_cluster_size_dict is not None:
+ cluster_sizes = []
+ for chain_id in chain_ids:
+ full_name = "_".join([file_id, chain_id])
+ cluster_size = chain_cluster_size_dict.get(
+ full_name.upper(), -1
+ )
+ cluster_sizes.append(cluster_size)
+
+ local_data["cluster_sizes"] = cluster_sizes
+
local_data["chain_ids"] = chain_ids
local_data["seqs"] = seqs
local_data["no_chains"] = len(chain_ids)
@@ -38,8 +50,21 @@ def parse_file(f, args):
def main(args):
+ chain_cluster_size_dict = None
+ if args.cluster_file is not None:
+ chain_cluster_size_dict = {}
+ with open(args.cluster_file, "r") as fp:
+ clusters = [l.strip() for l in fp.readlines()]
+
+ for cluster in clusters:
+ chain_ids = cluster.split()
+ cluster_len = len(chain_ids)
+ for chain_id in chain_ids:
+ chain_id = chain_id.upper()
+ chain_cluster_size_dict[chain_id] = cluster_len
+
files = [f for f in os.listdir(args.mmcif_dir) if ".cif" in f]
- fn = partial(parse_file, args=args)
+ fn = partial(parse_file, args=args, chain_cluster_size_dict=chain_cluster_size_dict)
data = {}
with Pool(processes=args.no_workers) as p:
with tqdm(total=len(files)) as pbar:
@@ -63,6 +88,15 @@ def main(args):
"--no_workers", type=int, default=4,
help="Number of workers to use for parsing"
)
+ parser.add_argument(
+ "--cluster_file", type=str, default=None,
+ help=(
+ "Path to a cluster file (e.g. PDB40), one cluster "
+ "({PROT1_ID}_{CHAIN_ID} {PROT2_ID}_{CHAIN_ID} ...) per line. "
+ "Chains not in this cluster file will NOT be filtered by cluster "
+ "size."
+ )
+ )
parser.add_argument(
"--chunksize", type=int, default=10,
help="How many files should be distributed to each worker at a time"
diff --git a/scripts/precompute_alignments.py b/scripts/precompute_alignments.py
index e1509a61..1bef41a9 100644
--- a/scripts/precompute_alignments.py
+++ b/scripts/precompute_alignments.py
@@ -11,6 +11,7 @@
import openfold.data.mmcif_parsing as mmcif_parsing
from openfold.data.data_pipeline import AlignmentRunner
from openfold.data.parsers import parse_fasta
+from openfold.data.tools import hhsearch, hmmsearch
from openfold.np import protein, residue_constants
from utils import add_data_args
@@ -39,7 +40,8 @@ def run_seq_group_alignments(seq_groups, alignment_runner, args):
alignment_runner.run(
fasta_path, alignment_dir
)
- except:
+ except Exception as e:
+ logging.warning(e)
logging.warning(f"Failed to run alignments for {first_name}. Skipping...")
os.remove(fasta_path)
os.rmdir(alignment_dir)
@@ -114,15 +116,30 @@ def parse_and_align(files, alignment_runner, args):
def main(args):
# Build the alignment tool runner
+ if args.hmmsearch_binary_path is not None and args.pdb_seqres_database_path is not None:
+ template_searcher = hmmsearch.Hmmsearch(
+ binary_path=args.hmmsearch_binary_path,
+ hmmbuild_binary_path=args.hmmbuild_binary_path,
+ database_path=args.pdb_seqres_database_path,
+ )
+ elif args.hhsearch_binary_path is not None and args.pdb70_database_path is not None:
+ template_searcher = hhsearch.HHSearch(
+ binary_path=args.hhsearch_binary_path,
+ databases=[args.pdb70_database_path],
+ )
+ else:
+ template_searcher = None
+
alignment_runner = AlignmentRunner(
jackhmmer_binary_path=args.jackhmmer_binary_path,
hhblits_binary_path=args.hhblits_binary_path,
- hhsearch_binary_path=args.hhsearch_binary_path,
uniref90_database_path=args.uniref90_database_path,
mgnify_database_path=args.mgnify_database_path,
bfd_database_path=args.bfd_database_path,
+ uniref30_database_path=args.uniref30_database_path,
uniclust30_database_path=args.uniclust30_database_path,
- pdb70_database_path=args.pdb70_database_path,
+ uniprot_database_path=args.uniprot_database_path,
+ template_searcher=template_searcher,
use_small_bfd=args.bfd_database_path is None,
no_cpus=args.cpus_per_task,
)
diff --git a/scripts/utils.py b/scripts/utils.py
index 01eae33b..1d81333d 100644
--- a/scripts/utils.py
+++ b/scripts/utils.py
@@ -14,9 +14,18 @@ def add_data_args(parser: argparse.ArgumentParser):
parser.add_argument(
'--pdb70_database_path', type=str, default=None,
)
+ parser.add_argument(
+ '--pdb_seqres_database_path', type=str, default=None,
+ )
+ parser.add_argument(
+ '--uniref30_database_path', type=str, default=None,
+ )
parser.add_argument(
'--uniclust30_database_path', type=str, default=None,
)
+ parser.add_argument(
+ '--uniprot_database_path', type=str, default=None,
+ )
parser.add_argument(
'--bfd_database_path', type=str, default=None,
)
@@ -29,6 +38,12 @@ def add_data_args(parser: argparse.ArgumentParser):
parser.add_argument(
'--hhsearch_binary_path', type=str, default='/usr/bin/hhsearch'
)
+ parser.add_argument(
+ '--hmmsearch_binary_path', type=str, default='/usr/bin/hmmsearch'
+ )
+ parser.add_argument(
+ '--hmmbuild_binary_path', type=str, default='/usr/bin/hmmbuild'
+ )
parser.add_argument(
'--kalign_binary_path', type=str, default='/usr/bin/kalign'
)
diff --git a/setup.py b/setup.py
index 0ef74ef4..bec98625 100644
--- a/setup.py
+++ b/setup.py
@@ -75,6 +75,8 @@ def get_cuda_bare_metal_version(cuda_dir):
extra_cuda_flags += cc_flag
+cc_flag = ['-gencode', 'arch=compute_70,code=sm_70']
+
if bare_metal_major != -1:
modules = [CUDAExtension(
name="attn_core_inplace_cuda",
@@ -111,10 +113,10 @@ def get_cuda_bare_metal_version(cuda_dir):
setup(
name='openfold',
- version='1.0.1',
+ version='2.0.0',
description='A PyTorch reimplementation of DeepMind\'s AlphaFold 2',
- author='Gustaf Ahdritz & DeepMind',
- author_email='gahdritz@gmail.com',
+ author='OpenFold Team',
+ author_email='jennifer.wei@omsf.io',
license='Apache License, Version 2.0',
url='https://github.com/aqlaboratory/openfold',
packages=find_packages(exclude=["tests", "scripts"]),
diff --git a/tests/compare_utils.py b/tests/compare_utils.py
index af04b10d..326f5e25 100644
--- a/tests/compare_utils.py
+++ b/tests/compare_utils.py
@@ -6,11 +6,14 @@
import unittest
import numpy as np
+import torch
from openfold.config import model_config
from openfold.model.model import AlphaFold
from openfold.utils.import_weights import import_jax_weights_
+from tests.config import consts
+
# Give JAX some GPU memory discipline
# (by default it hogs 90% of GPU memory. This disables that behavior and also
# forces it to proactively free memory that it allocates)
@@ -57,26 +60,27 @@ def import_alphafold():
def get_alphafold_config():
- config = alphafold.model.config.model_config("model_1_ptm") # noqa
+ config = alphafold.model.config.model_config(consts.model) # noqa
config.model.global_config.deterministic = True
return config
-_param_path = "openfold/resources/params/params_model_1_ptm.npz"
+dir_path = os.path.dirname(os.path.realpath(__file__))
+_param_path = os.path.join(dir_path, "..", f"openfold/resources/params/params_{consts.model}.npz")
_model = None
def get_global_pretrained_openfold():
global _model
if _model is None:
- _model = AlphaFold(model_config("model_1_ptm"))
+ _model = AlphaFold(model_config(consts.model))
_model = _model.eval()
if not os.path.exists(_param_path):
raise FileNotFoundError(
"""Cannot load pretrained parameters. Make sure to run the
installation script before running tests."""
)
- import_jax_weights_(_model, _param_path, version="model_1_ptm")
+ import_jax_weights_(_model, _param_path, version=consts.model)
_model = _model.cuda()
return _model
@@ -97,7 +101,7 @@ def _remove_key_prefix(d, prefix):
for k, v in list(d.items()):
if k.startswith(prefix):
d.pop(k)
- d[k[len(prefix) :]] = v
+ d[k[len(prefix):]] = v
def fetch_alphafold_module_weights(weight_path):
@@ -106,7 +110,6 @@ def fetch_alphafold_module_weights(weight_path):
if "/" in weight_path:
spl = weight_path.split("/")
spl = spl if len(spl[-1]) != 0 else spl[:-1]
- module_name = spl[-1]
prefix = "/".join(spl[:-1]) + "/"
_remove_key_prefix(params, prefix)
@@ -117,3 +120,20 @@ def fetch_alphafold_module_weights(weight_path):
"Make sure to call import_alphafold before running this function"
)
return params
+
+
+def _assert_abs_diff_small_base(compare_func, expected, actual, eps):
+ # Helper function for comparing absolute differences of two torch tensors.
+ abs_diff = torch.abs(expected - actual)
+ err = compare_func(abs_diff)
+ zero_tensor = torch.tensor(0, dtype=err.dtype)
+ rtol = 1.6e-2 if err.dtype == torch.bfloat16 else 1.3e-6
+ torch.testing.assert_close(err, zero_tensor, atol=eps, rtol=rtol)
+
+
+def assert_max_abs_diff_small(expected, actual, eps):
+ _assert_abs_diff_small_base(torch.max, expected, actual, eps)
+
+
+def assert_mean_abs_diff_small(expected, actual, eps):
+ _assert_abs_diff_small_base(torch.mean, expected, actual, eps)
diff --git a/tests/config.py b/tests/config.py
index d6bef73e..60c875e7 100644
--- a/tests/config.py
+++ b/tests/config.py
@@ -1,9 +1,13 @@
import ml_collections as mlc
-consts = mlc.ConfigDict(
+
+monomer_consts = mlc.ConfigDict(
{
+ "model": "model_1_ptm", # monomer:model_1_ptm, multimer: model_1_multimer_v3
+ "is_multimer": False, # monomer: False, multimer: True
+ "chunk_size": 4,
"batch_size": 2,
- "n_res": 11,
+ "n_res": 22,
"n_seq": 13,
"n_templ": 3,
"n_extra": 17,
@@ -16,9 +20,37 @@
"c_s": 384,
"c_t": 64,
"c_e": 64,
+ "msa_logits": 23, # monomer: 23, multimer: 22
+ "template_mmcif_dir": None # Set for test_multimer_datamodule
}
)
+multimer_consts = mlc.ConfigDict(
+ {
+ "model": "model_1_multimer_v3", # monomer:model_1_ptm, multimer: model_1_multimer_v3
+ "is_multimer": True, # monomer: False, multimer: True
+ "chunk_size": 4,
+ "batch_size": 2,
+ "n_res": 22,
+ "n_seq": 13,
+ "n_templ": 3,
+ "n_extra": 17,
+ "n_heads_extra_msa": 8,
+ "eps": 5e-4,
+ # For compatibility with DeepMind's pretrained weights, it's easiest for
+ # everyone if these take their real values.
+ "c_m": 256,
+ "c_z": 128,
+ "c_s": 384,
+ "c_t": 64,
+ "c_e": 64,
+ "msa_logits": 22, # monomer: 23, multimer: 22
+ "template_mmcif_dir": None # Set for test_multimer_datamodule
+ }
+)
+
+consts = monomer_consts
+
config = mlc.ConfigDict(
{
"data": {
diff --git a/tests/data_utils.py b/tests/data_utils.py
index 25563a3c..a44ed3d3 100644
--- a/tests/data_utils.py
+++ b/tests/data_utils.py
@@ -12,10 +12,37 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from random import randint
import torch
import numpy as np
from scipy.spatial.transform import Rotation
+from tests.config import consts
+
+
+def random_asym_ids(n_res, split_chains=True, min_chain_len=4):
+ n_chain = randint(1, n_res // min_chain_len) if consts.is_multimer else 1
+
+ if not split_chains:
+ return [0] * n_res
+
+ assert n_res >= n_chain
+
+ pieces = []
+ asym_ids = []
+ final_idx = n_chain - 1
+ for idx in range(n_chain - 1):
+ n_stop = (n_res - sum(pieces) - n_chain + idx - min_chain_len)
+ if n_stop <