new ui, preps for af2

.gitattributes

@@ -25,3 +25,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+npz filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text

alphafold/LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

alphafold/alphafold/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+"""An implementation of the inference pipeline of AlphaFold v2.0."""

alphafold/alphafold/common/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+"""Common data types and constants used within Alphafold."""

alphafold/alphafold/common/confidence.py

@@ -0,0 +1,155 @@
+# 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.
+
+"""Functions for processing confidence metrics."""
+
+from typing import Dict, Optional, Tuple
+import numpy as np
+import scipy.special
+
+
+def compute_plddt(logits: np.ndarray) -> np.ndarray:
+  """Computes per-residue pLDDT from logits.
+
+  Args:
+    logits: [num_res, num_bins] output from the PredictedLDDTHead.
+
+  Returns:
+    plddt: [num_res] per-residue pLDDT.
+  """
+  num_bins = logits.shape[-1]
+  bin_width = 1.0 / num_bins
+  bin_centers = np.arange(start=0.5 * bin_width, stop=1.0, step=bin_width)
+  probs = scipy.special.softmax(logits, axis=-1)
+  predicted_lddt_ca = np.sum(probs * bin_centers[None, :], axis=-1)
+  return predicted_lddt_ca * 100
+
+
+def _calculate_bin_centers(breaks: np.ndarray):
+  """Gets the bin centers from the bin edges.
+
+  Args:
+    breaks: [num_bins - 1] the error bin edges.
+
+  Returns:
+    bin_centers: [num_bins] the error bin centers.
+  """
+  step = (breaks[1] - breaks[0])
+
+  # Add half-step to get the center
+  bin_centers = breaks + step / 2
+  # Add a catch-all bin at the end.
+  bin_centers = np.concatenate([bin_centers, [bin_centers[-1] + step]],
+                               axis=0)
+  return bin_centers
+
+
+def _calculate_expected_aligned_error(
+    alignment_confidence_breaks: np.ndarray,
+    aligned_distance_error_probs: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+  """Calculates expected aligned distance errors for every pair of residues.
+
+  Args:
+    alignment_confidence_breaks: [num_bins - 1] the error bin edges.
+    aligned_distance_error_probs: [num_res, num_res, num_bins] the predicted
+      probs for each error bin, for each pair of residues.
+
+  Returns:
+    predicted_aligned_error: [num_res, num_res] the expected aligned distance
+      error for each pair of residues.
+    max_predicted_aligned_error: The maximum predicted error possible.
+  """
+  bin_centers = _calculate_bin_centers(alignment_confidence_breaks)
+
+  # Tuple of expected aligned distance error and max possible error.
+  return (np.sum(aligned_distance_error_probs * bin_centers, axis=-1),
+          np.asarray(bin_centers[-1]))
+
+
+def compute_predicted_aligned_error(
+    logits: np.ndarray,
+    breaks: np.ndarray) -> Dict[str, np.ndarray]:
+  """Computes aligned confidence metrics from logits.
+
+  Args:
+    logits: [num_res, num_res, num_bins] the logits output from
+      PredictedAlignedErrorHead.
+    breaks: [num_bins - 1] the error bin edges.
+
+  Returns:
+    aligned_confidence_probs: [num_res, num_res, num_bins] the predicted
+      aligned error probabilities over bins for each residue pair.
+    predicted_aligned_error: [num_res, num_res] the expected aligned distance
+      error for each pair of residues.
+    max_predicted_aligned_error: The maximum predicted error possible.
+  """
+  aligned_confidence_probs = scipy.special.softmax(
+      logits,
+      axis=-1)
+  predicted_aligned_error, max_predicted_aligned_error = (
+      _calculate_expected_aligned_error(
+          alignment_confidence_breaks=breaks,
+          aligned_distance_error_probs=aligned_confidence_probs))
+  return {
+      'aligned_confidence_probs': aligned_confidence_probs,
+      'predicted_aligned_error': predicted_aligned_error,
+      'max_predicted_aligned_error': max_predicted_aligned_error,
+  }
+
+
+def predicted_tm_score(
+    logits: np.ndarray,
+    breaks: np.ndarray,
+    residue_weights: Optional[np.ndarray] = None) -> np.ndarray:
+  """Computes predicted TM alignment score.
+
+  Args:
+    logits: [num_res, num_res, num_bins] the logits output from
+      PredictedAlignedErrorHead.
+    breaks: [num_bins] the error bins.
+    residue_weights: [num_res] the per residue weights to use for the
+      expectation.
+
+  Returns:
+    ptm_score: the predicted TM alignment score.
+  """
+
+  # residue_weights has to be in [0, 1], but can be floating-point, i.e. the
+  # exp. resolved head's probability.
+  if residue_weights is None:
+    residue_weights = np.ones(logits.shape[0])
+
+  bin_centers = _calculate_bin_centers(breaks)
+
+  num_res = np.sum(residue_weights)
+  # Clip num_res to avoid negative/undefined d0.
+  clipped_num_res = max(num_res, 19)
+
+  # Compute d_0(num_res) as defined by TM-score, eqn. (5) in
+  # http://zhanglab.ccmb.med.umich.edu/papers/2004_3.pdf
+  # Yang & Skolnick "Scoring function for automated
+  # assessment of protein structure template quality" 2004
+  d0 = 1.24 * (clipped_num_res - 15) ** (1./3) - 1.8
+
+  # Convert logits to probs
+  probs = scipy.special.softmax(logits, axis=-1)
+
+  # TM-Score term for every bin
+  tm_per_bin = 1. / (1 + np.square(bin_centers) / np.square(d0))
+  # E_distances tm(distance)
+  predicted_tm_term = np.sum(probs * tm_per_bin, axis=-1)
+
+  normed_residue_mask = residue_weights / (1e-8 + residue_weights.sum())
+  per_alignment = np.sum(predicted_tm_term * normed_residue_mask, axis=-1)
+  return np.asarray(per_alignment[(per_alignment * residue_weights).argmax()])

alphafold/alphafold/common/protein.py

@@ -0,0 +1,229 @@
+# 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.
+
+"""Protein data type."""
+import dataclasses
+import io
+from typing import Any, Mapping, Optional
+from alphafold.common import residue_constants
+from Bio.PDB import PDBParser
+import numpy as np
+
+FeatureDict = Mapping[str, np.ndarray]
+ModelOutput = Mapping[str, Any]  # Is a nested dict.
+
+
+@dataclasses.dataclass(frozen=True)
+class Protein:
+  """Protein structure representation."""
+
+  # Cartesian coordinates of atoms in angstroms. The atom types correspond to
+  # residue_constants.atom_types, i.e. the first three are N, CA, CB.
+  atom_positions: np.ndarray  # [num_res, num_atom_type, 3]
+
+  # Amino-acid type for each residue represented as an integer between 0 and
+  # 20, where 20 is 'X'.
+  aatype: np.ndarray  # [num_res]
+
+  # Binary float mask to indicate presence of a particular atom. 1.0 if an atom
+  # is present and 0.0 if not. This should be used for loss masking.
+  atom_mask: np.ndarray  # [num_res, num_atom_type]
+
+  # Residue index as used in PDB. It is not necessarily continuous or 0-indexed.
+  residue_index: np.ndarray  # [num_res]
+
+  # B-factors, or temperature factors, of each residue (in sq. angstroms units),
+  # representing the displacement of the residue from its ground truth mean
+  # value.
+  b_factors: np.ndarray  # [num_res, num_atom_type]
+
+
+def from_pdb_string(pdb_str: str, chain_id: Optional[str] = None) -> Protein:
+  """Takes a PDB string and constructs a Protein object.
+
+  WARNING: All non-standard residue types will be converted into UNK. All
+    non-standard atoms will be ignored.
+
+  Args:
+    pdb_str: The contents of the pdb file
+    chain_id: If None, then the pdb file must contain a single chain (which
+      will be parsed). If chain_id is specified (e.g. A), then only that chain
+      is parsed.
+
+  Returns:
+    A new `Protein` parsed from the pdb contents.
+  """
+  pdb_fh = io.StringIO(pdb_str)
+  parser = PDBParser(QUIET=True)
+  structure = parser.get_structure('none', pdb_fh)
+  models = list(structure.get_models())
+  if len(models) != 1:
+    raise ValueError(
+        f'Only single model PDBs are supported. Found {len(models)} models.')
+  model = models[0]
+
+  if chain_id is not None:
+    chain = model[chain_id]
+  else:
+    chains = list(model.get_chains())
+    if len(chains) != 1:
+      raise ValueError(
+          'Only single chain PDBs are supported when chain_id not specified. '
+          f'Found {len(chains)} chains.')
+    else:
+      chain = chains[0]
+
+  atom_positions = []
+  aatype = []
+  atom_mask = []
+  residue_index = []
+  b_factors = []
+
+  for res in chain:
+    if res.id[2] != ' ':
+      raise ValueError(
+          f'PDB contains an insertion code at chain {chain.id} and residue '
+          f'index {res.id[1]}. These are not supported.')
+    res_shortname = residue_constants.restype_3to1.get(res.resname, 'X')
+    restype_idx = residue_constants.restype_order.get(
+        res_shortname, residue_constants.restype_num)
+    pos = np.zeros((residue_constants.atom_type_num, 3))
+    mask = np.zeros((residue_constants.atom_type_num,))
+    res_b_factors = np.zeros((residue_constants.atom_type_num,))
+    for atom in res:
+      if atom.name not in residue_constants.atom_types:
+        continue
+      pos[residue_constants.atom_order[atom.name]] = atom.coord
+      mask[residue_constants.atom_order[atom.name]] = 1.
+      res_b_factors[residue_constants.atom_order[atom.name]] = atom.bfactor
+    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)
+    residue_index.append(res.id[1])
+    b_factors.append(res_b_factors)
+
+  return Protein(
+      atom_positions=np.array(atom_positions),
+      atom_mask=np.array(atom_mask),
+      aatype=np.array(aatype),
+      residue_index=np.array(residue_index),
+      b_factors=np.array(b_factors))
+
+
+def to_pdb(prot: Protein) -> str:
+  """Converts a `Protein` instance to a PDB string.
+
+  Args:
+    prot: The protein to convert to PDB.
+
+  Returns:
+    PDB string.
+  """
+  restypes = residue_constants.restypes + ['X']
+  res_1to3 = lambda r: residue_constants.restype_1to3.get(restypes[r], 'UNK')
+  atom_types = residue_constants.atom_types
+
+  pdb_lines = []
+
+  atom_mask = prot.atom_mask
+  aatype = prot.aatype
+  atom_positions = prot.atom_positions
+  residue_index = prot.residue_index.astype(np.int32)
+  b_factors = prot.b_factors
+
+  if np.any(aatype > residue_constants.restype_num):
+    raise ValueError('Invalid aatypes.')
+
+  pdb_lines.append('MODEL     1')
+  atom_index = 1
+  chain_id = 'A'
+  # Add all atom sites.
+  for i in range(aatype.shape[0]):
+    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]):
+      if mask < 0.5:
+        continue
+
+      record_type = 'ATOM'
+      name = atom_name if len(atom_name) == 4 else f' {atom_name}'
+      alt_loc = ''
+      insertion_code = ''
+      occupancy = 1.00
+      element = atom_name[0]  # Protein supports only C, N, O, S, this works.
+      charge = ''
+      # PDB is a columnar format, every space matters here!
+      atom_line = (f'{record_type:<6}{atom_index:>5} {name:<4}{alt_loc:>1}'
+                   f'{res_name_3:>3} {chain_id:>1}'
+                   f'{residue_index[i]:>4}{insertion_code:>1}   '
+                   f'{pos[0]:>8.3f}{pos[1]:>8.3f}{pos[2]:>8.3f}'
+                   f'{occupancy:>6.2f}{b_factor:>6.2f}          '
+                   f'{element:>2}{charge:>2}')
+      pdb_lines.append(atom_line)
+      atom_index += 1
+
+  # Close the chain.
+  chain_end = 'TER'
+  chain_termination_line = (
+      f'{chain_end:<6}{atom_index:>5}      {res_1to3(aatype[-1]):>3} '
+      f'{chain_id:>1}{residue_index[-1]:>4}')
+  pdb_lines.append(chain_termination_line)
+  pdb_lines.append('ENDMDL')
+
+  pdb_lines.append('END')
+  pdb_lines.append('')
+  return '\n'.join(pdb_lines)
+
+
+def ideal_atom_mask(prot: Protein) -> np.ndarray:
+  """Computes an ideal atom mask.
+
+  `Protein.atom_mask` typically is defined according to the atoms that are
+  reported in the PDB. This function computes a mask according to heavy atoms
+  that should be present in the given sequence of amino acids.
+
+  Args:
+    prot: `Protein` whose fields are `numpy.ndarray` objects.
+
+  Returns:
+    An ideal atom mask.
+  """
+  return residue_constants.STANDARD_ATOM_MASK[prot.aatype]
+
+
+def from_prediction(features: FeatureDict, result: ModelOutput,
+                    b_factors: Optional[np.ndarray] = None) -> Protein:
+  """Assembles a protein from a prediction.
+
+  Args:
+    features: Dictionary holding model inputs.
+    result: Dictionary holding model outputs.
+    b_factors: (Optional) B-factors to use for the protein.
+
+  Returns:
+    A protein instance.
+  """
+  fold_output = result['structure_module']
+  if b_factors is None:
+    b_factors = np.zeros_like(fold_output['final_atom_mask'])
+
+  return Protein(
+      aatype=features['aatype'][0],
+      atom_positions=fold_output['final_atom_positions'],
+      atom_mask=fold_output['final_atom_mask'],
+      residue_index=features['residue_index'][0] + 1,
+      b_factors=b_factors)

alphafold/alphafold/common/protein_test.py

@@ -0,0 +1,89 @@
+# 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.
+
+"""Tests for protein."""
+
+import os
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from alphafold.common import protein
+from alphafold.common import residue_constants
+import numpy as np
+# Internal import (7716).
+
+TEST_DATA_DIR = 'alphafold/common/testdata/'
+
+
+class ProteinTest(parameterized.TestCase):
+
+  def _check_shapes(self, prot, num_res):
+    """Check that the processed shapes are correct."""
+    num_atoms = residue_constants.atom_type_num
+    self.assertEqual((num_res, num_atoms, 3), prot.atom_positions.shape)
+    self.assertEqual((num_res,), prot.aatype.shape)
+    self.assertEqual((num_res, num_atoms), prot.atom_mask.shape)
+    self.assertEqual((num_res,), prot.residue_index.shape)
+    self.assertEqual((num_res, num_atoms), prot.b_factors.shape)
+
+  @parameterized.parameters(('2rbg.pdb', 'A', 282),
+                            ('2rbg.pdb', 'B', 282))
+  def test_from_pdb_str(self, pdb_file, chain_id, num_res):
+    pdb_file = os.path.join(absltest.get_default_test_srcdir(), TEST_DATA_DIR,
+                            pdb_file)
+    with open(pdb_file) as f:
+      pdb_string = f.read()
+    prot = protein.from_pdb_string(pdb_string, chain_id)
+    self._check_shapes(prot, num_res)
+    self.assertGreaterEqual(prot.aatype.min(), 0)
+    # Allow equal since unknown restypes have index equal to restype_num.
+    self.assertLessEqual(prot.aatype.max(), residue_constants.restype_num)
+
+  def test_to_pdb(self):
+    with open(
+        os.path.join(absltest.get_default_test_srcdir(), TEST_DATA_DIR,
+                     '2rbg.pdb')) as f:
+      pdb_string = f.read()
+    prot = protein.from_pdb_string(pdb_string, chain_id='A')
+    pdb_string_reconstr = protein.to_pdb(prot)
+    prot_reconstr = protein.from_pdb_string(pdb_string_reconstr)
+
+    np.testing.assert_array_equal(prot_reconstr.aatype, prot.aatype)
+    np.testing.assert_array_almost_equal(
+        prot_reconstr.atom_positions, prot.atom_positions)
+    np.testing.assert_array_almost_equal(
+        prot_reconstr.atom_mask, prot.atom_mask)
+    np.testing.assert_array_equal(
+        prot_reconstr.residue_index, prot.residue_index)
+    np.testing.assert_array_almost_equal(
+        prot_reconstr.b_factors, prot.b_factors)
+
+  def test_ideal_atom_mask(self):
+    with open(
+        os.path.join(absltest.get_default_test_srcdir(), TEST_DATA_DIR,
+                     '2rbg.pdb')) as f:
+      pdb_string = f.read()
+    prot = protein.from_pdb_string(pdb_string, chain_id='A')
+    ideal_mask = protein.ideal_atom_mask(prot)
+    non_ideal_residues = set([102] + list(range(127, 285)))
+    for i, (res, atom_mask) in enumerate(
+        zip(prot.residue_index, prot.atom_mask)):
+      if res in non_ideal_residues:
+        self.assertFalse(np.all(atom_mask == ideal_mask[i]), msg=f'{res}')
+      else:
+        self.assertTrue(np.all(atom_mask == ideal_mask[i]), msg=f'{res}')
+
+
+if __name__ == '__main__':
+  absltest.main()

alphafold/alphafold/common/residue_constants.py

@@ -0,0 +1,895 @@
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Constants used in AlphaFold."""
+
+import collections
+import functools
+from typing import List, Mapping, Tuple
+
+import numpy as np
+import tree
+
+# Internal import (35fd).
+
+
+# Distance from one CA to next CA [trans configuration: omega = 180].
+ca_ca = 3.80209737096
+
+# Format: The list for each AA type contains chi1, chi2, chi3, chi4 in
+# this order (or a relevant subset from chi1 onwards). ALA and GLY don't have
+# chi angles so their chi angle lists are empty.
+chi_angles_atoms = {
+    'ALA': [],
+    # Chi5 in arginine is always 0 +- 5 degrees, so ignore it.
+    'ARG': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD'],
+            ['CB', 'CG', 'CD', 'NE'], ['CG', 'CD', 'NE', 'CZ']],
+    'ASN': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'OD1']],
+    'ASP': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'OD1']],
+    'CYS': [['N', 'CA', 'CB', 'SG']],
+    'GLN': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD'],
+            ['CB', 'CG', 'CD', 'OE1']],
+    'GLU': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD'],
+            ['CB', 'CG', 'CD', 'OE1']],
+    'GLY': [],
+    'HIS': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'ND1']],
+    'ILE': [['N', 'CA', 'CB', 'CG1'], ['CA', 'CB', 'CG1', 'CD1']],
+    'LEU': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD1']],
+    'LYS': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD'],
+            ['CB', 'CG', 'CD', 'CE'], ['CG', 'CD', 'CE', 'NZ']],
+    'MET': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'SD'],
+            ['CB', 'CG', 'SD', 'CE']],
+    'PHE': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD1']],
+    'PRO': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD']],
+    'SER': [['N', 'CA', 'CB', 'OG']],
+    'THR': [['N', 'CA', 'CB', 'OG1']],
+    'TRP': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD1']],
+    'TYR': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD1']],
+    'VAL': [['N', 'CA', 'CB', 'CG1']],
+}
+
+# If chi angles given in fixed-length array, this matrix determines how to mask
+# them for each AA type. The order is as per restype_order (see below).
+chi_angles_mask = [
+    [0.0, 0.0, 0.0, 0.0],  # ALA
+    [1.0, 1.0, 1.0, 1.0],  # ARG
+    [1.0, 1.0, 0.0, 0.0],  # ASN
+    [1.0, 1.0, 0.0, 0.0],  # ASP
+    [1.0, 0.0, 0.0, 0.0],  # CYS
+    [1.0, 1.0, 1.0, 0.0],  # GLN
+    [1.0, 1.0, 1.0, 0.0],  # GLU
+    [0.0, 0.0, 0.0, 0.0],  # GLY
+    [1.0, 1.0, 0.0, 0.0],  # HIS
+    [1.0, 1.0, 0.0, 0.0],  # ILE
+    [1.0, 1.0, 0.0, 0.0],  # LEU
+    [1.0, 1.0, 1.0, 1.0],  # LYS
+    [1.0, 1.0, 1.0, 0.0],  # MET
+    [1.0, 1.0, 0.0, 0.0],  # PHE
+    [1.0, 1.0, 0.0, 0.0],  # PRO
+    [1.0, 0.0, 0.0, 0.0],  # SER
+    [1.0, 0.0, 0.0, 0.0],  # THR
+    [1.0, 1.0, 0.0, 0.0],  # TRP
+    [1.0, 1.0, 0.0, 0.0],  # TYR
+    [1.0, 0.0, 0.0, 0.0],  # VAL
+]
+
+# The following chi angles are pi periodic: they can be rotated by a multiple
+# of pi without affecting the structure.
+chi_pi_periodic = [
+    [0.0, 0.0, 0.0, 0.0],  # ALA
+    [0.0, 0.0, 0.0, 0.0],  # ARG
+    [0.0, 0.0, 0.0, 0.0],  # ASN
+    [0.0, 1.0, 0.0, 0.0],  # ASP
+    [0.0, 0.0, 0.0, 0.0],  # CYS
+    [0.0, 0.0, 0.0, 0.0],  # GLN
+    [0.0, 0.0, 1.0, 0.0],  # GLU
+    [0.0, 0.0, 0.0, 0.0],  # GLY
+    [0.0, 0.0, 0.0, 0.0],  # HIS
+    [0.0, 0.0, 0.0, 0.0],  # ILE
+    [0.0, 0.0, 0.0, 0.0],  # LEU
+    [0.0, 0.0, 0.0, 0.0],  # LYS
+    [0.0, 0.0, 0.0, 0.0],  # MET
+    [0.0, 1.0, 0.0, 0.0],  # PHE
+    [0.0, 0.0, 0.0, 0.0],  # PRO
+    [0.0, 0.0, 0.0, 0.0],  # SER
+    [0.0, 0.0, 0.0, 0.0],  # THR
+    [0.0, 0.0, 0.0, 0.0],  # TRP
+    [0.0, 1.0, 0.0, 0.0],  # TYR
+    [0.0, 0.0, 0.0, 0.0],  # VAL
+    [0.0, 0.0, 0.0, 0.0],  # UNK
+]
+
+# Atoms positions relative to the 8 rigid groups, defined by the pre-omega, phi,
+# psi and chi angles:
+# 0: 'backbone group',
+# 1: 'pre-omega-group', (empty)
+# 2: 'phi-group', (currently empty, because it defines only hydrogens)
+# 3: 'psi-group',
+# 4,5,6,7: 'chi1,2,3,4-group'
+# The atom positions are relative to the axis-end-atom of the corresponding
+# rotation axis. The x-axis is in direction of the rotation axis, and the y-axis
+# is defined such that the dihedral-angle-definiting atom (the last entry in
+# chi_angles_atoms above) is in the xy-plane (with a positive y-coordinate).
+# format: [atomname, group_idx, rel_position]
+rigid_group_atom_positions = {
+    'ALA': [
+        ['N', 0, (-0.525, 1.363, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.526, -0.000, -0.000)],
+        ['CB', 0, (-0.529, -0.774, -1.205)],
+        ['O', 3, (0.627, 1.062, 0.000)],
+    ],
+    'ARG': [
+        ['N', 0, (-0.524, 1.362, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.525, -0.000, -0.000)],
+        ['CB', 0, (-0.524, -0.778, -1.209)],
+        ['O', 3, (0.626, 1.062, 0.000)],
+        ['CG', 4, (0.616, 1.390, -0.000)],
+        ['CD', 5, (0.564, 1.414, 0.000)],
+        ['NE', 6, (0.539, 1.357, -0.000)],
+        ['NH1', 7, (0.206, 2.301, 0.000)],
+        ['NH2', 7, (2.078, 0.978, -0.000)],
+        ['CZ', 7, (0.758, 1.093, -0.000)],
+    ],
+    'ASN': [
+        ['N', 0, (-0.536, 1.357, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.526, -0.000, -0.000)],
+        ['CB', 0, (-0.531, -0.787, -1.200)],
+        ['O', 3, (0.625, 1.062, 0.000)],
+        ['CG', 4, (0.584, 1.399, 0.000)],
+        ['ND2', 5, (0.593, -1.188, 0.001)],
+        ['OD1', 5, (0.633, 1.059, 0.000)],
+    ],
+    'ASP': [
+        ['N', 0, (-0.525, 1.362, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.527, 0.000, -0.000)],
+        ['CB', 0, (-0.526, -0.778, -1.208)],
+        ['O', 3, (0.626, 1.062, -0.000)],
+        ['CG', 4, (0.593, 1.398, -0.000)],
+        ['OD1', 5, (0.610, 1.091, 0.000)],
+        ['OD2', 5, (0.592, -1.101, -0.003)],
+    ],
+    'CYS': [
+        ['N', 0, (-0.522, 1.362, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.524, 0.000, 0.000)],
+        ['CB', 0, (-0.519, -0.773, -1.212)],
+        ['O', 3, (0.625, 1.062, -0.000)],
+        ['SG', 4, (0.728, 1.653, 0.000)],
+    ],
+    'GLN': [
+        ['N', 0, (-0.526, 1.361, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.526, 0.000, 0.000)],
+        ['CB', 0, (-0.525, -0.779, -1.207)],
+        ['O', 3, (0.626, 1.062, -0.000)],
+        ['CG', 4, (0.615, 1.393, 0.000)],
+        ['CD', 5, (0.587, 1.399, -0.000)],
+        ['NE2', 6, (0.593, -1.189, -0.001)],
+        ['OE1', 6, (0.634, 1.060, 0.000)],
+    ],
+    'GLU': [
+        ['N', 0, (-0.528, 1.361, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.526, -0.000, -0.000)],
+        ['CB', 0, (-0.526, -0.781, -1.207)],
+        ['O', 3, (0.626, 1.062, 0.000)],
+        ['CG', 4, (0.615, 1.392, 0.000)],
+        ['CD', 5, (0.600, 1.397, 0.000)],
+        ['OE1', 6, (0.607, 1.095, -0.000)],
+        ['OE2', 6, (0.589, -1.104, -0.001)],
+    ],
+    'GLY': [
+        ['N', 0, (-0.572, 1.337, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.517, -0.000, -0.000)],
+        ['O', 3, (0.626, 1.062, -0.000)],
+    ],
+    'HIS': [
+        ['N', 0, (-0.527, 1.360, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.525, 0.000, 0.000)],
+        ['CB', 0, (-0.525, -0.778, -1.208)],
+        ['O', 3, (0.625, 1.063, 0.000)],
+        ['CG', 4, (0.600, 1.370, -0.000)],
+        ['CD2', 5, (0.889, -1.021, 0.003)],
+        ['ND1', 5, (0.744, 1.160, -0.000)],
+        ['CE1', 5, (2.030, 0.851, 0.002)],
+        ['NE2', 5, (2.145, -0.466, 0.004)],
+    ],
+    'ILE': [
+        ['N', 0, (-0.493, 1.373, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.527, -0.000, -0.000)],
+        ['CB', 0, (-0.536, -0.793, -1.213)],
+        ['O', 3, (0.627, 1.062, -0.000)],
+        ['CG1', 4, (0.534, 1.437, -0.000)],
+        ['CG2', 4, (0.540, -0.785, -1.199)],
+        ['CD1', 5, (0.619, 1.391, 0.000)],
+    ],
+    'LEU': [
+        ['N', 0, (-0.520, 1.363, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.525, -0.000, -0.000)],
+        ['CB', 0, (-0.522, -0.773, -1.214)],
+        ['O', 3, (0.625, 1.063, -0.000)],
+        ['CG', 4, (0.678, 1.371, 0.000)],
+        ['CD1', 5, (0.530, 1.430, -0.000)],
+        ['CD2', 5, (0.535, -0.774, 1.200)],
+    ],
+    'LYS': [
+        ['N', 0, (-0.526, 1.362, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.526, 0.000, 0.000)],
+        ['CB', 0, (-0.524, -0.778, -1.208)],
+        ['O', 3, (0.626, 1.062, -0.000)],
+        ['CG', 4, (0.619, 1.390, 0.000)],
+        ['CD', 5, (0.559, 1.417, 0.000)],
+        ['CE', 6, (0.560, 1.416, 0.000)],
+        ['NZ', 7, (0.554, 1.387, 0.000)],
+    ],
+    'MET': [
+        ['N', 0, (-0.521, 1.364, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.525, 0.000, 0.000)],
+        ['CB', 0, (-0.523, -0.776, -1.210)],
+        ['O', 3, (0.625, 1.062, -0.000)],
+        ['CG', 4, (0.613, 1.391, -0.000)],
+        ['SD', 5, (0.703, 1.695, 0.000)],
+        ['CE', 6, (0.320, 1.786, -0.000)],
+    ],
+    'PHE': [
+        ['N', 0, (-0.518, 1.363, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.524, 0.000, -0.000)],
+        ['CB', 0, (-0.525, -0.776, -1.212)],
+        ['O', 3, (0.626, 1.062, -0.000)],
+        ['CG', 4, (0.607, 1.377, 0.000)],
+        ['CD1', 5, (0.709, 1.195, -0.000)],
+        ['CD2', 5, (0.706, -1.196, 0.000)],
+        ['CE1', 5, (2.102, 1.198, -0.000)],
+        ['CE2', 5, (2.098, -1.201, -0.000)],
+        ['CZ', 5, (2.794, -0.003, -0.001)],
+    ],
+    'PRO': [
+        ['N', 0, (-0.566, 1.351, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.527, -0.000, 0.000)],
+        ['CB', 0, (-0.546, -0.611, -1.293)],
+        ['O', 3, (0.621, 1.066, 0.000)],
+        ['CG', 4, (0.382, 1.445, 0.0)],
+        # ['CD', 5, (0.427, 1.440, 0.0)],
+        ['CD', 5, (0.477, 1.424, 0.0)],  # manually made angle 2 degrees larger
+    ],
+    'SER': [
+        ['N', 0, (-0.529, 1.360, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.525, -0.000, -0.000)],
+        ['CB', 0, (-0.518, -0.777, -1.211)],
+        ['O', 3, (0.626, 1.062, -0.000)],
+        ['OG', 4, (0.503, 1.325, 0.000)],
+    ],
+    'THR': [
+        ['N', 0, (-0.517, 1.364, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.526, 0.000, -0.000)],
+        ['CB', 0, (-0.516, -0.793, -1.215)],
+        ['O', 3, (0.626, 1.062, 0.000)],
+        ['CG2', 4, (0.550, -0.718, -1.228)],
+        ['OG1', 4, (0.472, 1.353, 0.000)],
+    ],
+    'TRP': [
+        ['N', 0, (-0.521, 1.363, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.525, -0.000, 0.000)],
+        ['CB', 0, (-0.523, -0.776, -1.212)],
+        ['O', 3, (0.627, 1.062, 0.000)],
+        ['CG', 4, (0.609, 1.370, -0.000)],
+        ['CD1', 5, (0.824, 1.091, 0.000)],
+        ['CD2', 5, (0.854, -1.148, -0.005)],
+        ['CE2', 5, (2.186, -0.678, -0.007)],
+        ['CE3', 5, (0.622, -2.530, -0.007)],
+        ['NE1', 5, (2.140, 0.690, -0.004)],
+        ['CH2', 5, (3.028, -2.890, -0.013)],
+        ['CZ2', 5, (3.283, -1.543, -0.011)],
+        ['CZ3', 5, (1.715, -3.389, -0.011)],
+    ],
+    'TYR': [
+        ['N', 0, (-0.522, 1.362, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.524, -0.000, -0.000)],
+        ['CB', 0, (-0.522, -0.776, -1.213)],
+        ['O', 3, (0.627, 1.062, -0.000)],
+        ['CG', 4, (0.607, 1.382, -0.000)],
+        ['CD1', 5, (0.716, 1.195, -0.000)],
+        ['CD2', 5, (0.713, -1.194, -0.001)],
+        ['CE1', 5, (2.107, 1.200, -0.002)],
+        ['CE2', 5, (2.104, -1.201, -0.003)],
+        ['OH', 5, (4.168, -0.002, -0.005)],
+        ['CZ', 5, (2.791, -0.001, -0.003)],
+    ],
+    'VAL': [
+        ['N', 0, (-0.494, 1.373, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.527, -0.000, -0.000)],
+        ['CB', 0, (-0.533, -0.795, -1.213)],
+        ['O', 3, (0.627, 1.062, -0.000)],
+        ['CG1', 4, (0.540, 1.429, -0.000)],
+        ['CG2', 4, (0.533, -0.776, 1.203)],
+    ],
+}
+
+# A list of atoms (excluding hydrogen) for each AA type. PDB naming convention.
+residue_atoms = {
+    'ALA': ['C', 'CA', 'CB', 'N', 'O'],
+    'ARG': ['C', 'CA', 'CB', 'CG', 'CD', 'CZ', 'N', 'NE', 'O', 'NH1', 'NH2'],
+    'ASP': ['C', 'CA', 'CB', 'CG', 'N', 'O', 'OD1', 'OD2'],
+    'ASN': ['C', 'CA', 'CB', 'CG', 'N', 'ND2', 'O', 'OD1'],
+    'CYS': ['C', 'CA', 'CB', 'N', 'O', 'SG'],
+    'GLU': ['C', 'CA', 'CB', 'CG', 'CD', 'N', 'O', 'OE1', 'OE2'],
+    'GLN': ['C', 'CA', 'CB', 'CG', 'CD', 'N', 'NE2', 'O', 'OE1'],
+    'GLY': ['C', 'CA', 'N', 'O'],
+    'HIS': ['C', 'CA', 'CB', 'CG', 'CD2', 'CE1', 'N', 'ND1', 'NE2', 'O'],
+    'ILE': ['C', 'CA', 'CB', 'CG1', 'CG2', 'CD1', 'N', 'O'],
+    'LEU': ['C', 'CA', 'CB', 'CG', 'CD1', 'CD2', 'N', 'O'],
+    'LYS': ['C', 'CA', 'CB', 'CG', 'CD', 'CE', 'N', 'NZ', 'O'],
+    'MET': ['C', 'CA', 'CB', 'CG', 'CE', 'N', 'O', 'SD'],
+    'PHE': ['C', 'CA', 'CB', 'CG', 'CD1', 'CD2', 'CE1', 'CE2', 'CZ', 'N', 'O'],
+    'PRO': ['C', 'CA', 'CB', 'CG', 'CD', 'N', 'O'],
+    'SER': ['C', 'CA', 'CB', 'N', 'O', 'OG'],
+    'THR': ['C', 'CA', 'CB', 'CG2', 'N', 'O', 'OG1'],
+    'TRP': ['C', 'CA', 'CB', 'CG', 'CD1', 'CD2', 'CE2', 'CE3', 'CZ2', 'CZ3',
+            'CH2', 'N', 'NE1', 'O'],
+    'TYR': ['C', 'CA', 'CB', 'CG', 'CD1', 'CD2', 'CE1', 'CE2', 'CZ', 'N', 'O',
+            'OH'],
+    'VAL': ['C', 'CA', 'CB', 'CG1', 'CG2', 'N', 'O']
+}
+
+# Naming swaps for ambiguous atom names.
+# Due to symmetries in the amino acids the naming of atoms is ambiguous in
+# 4 of the 20 amino acids.
+# (The LDDT paper lists 7 amino acids as ambiguous, but the naming ambiguities
+# in LEU, VAL and ARG can be resolved by using the 3d constellations of
+# the 'ambiguous' atoms and their neighbours)
+residue_atom_renaming_swaps = {
+    'ASP': {'OD1': 'OD2'},
+    'GLU': {'OE1': 'OE2'},
+    'PHE': {'CD1': 'CD2', 'CE1': 'CE2'},
+    'TYR': {'CD1': 'CD2', 'CE1': 'CE2'},
+}
+
+# Van der Waals radii [Angstroem] of the atoms (from Wikipedia)
+van_der_waals_radius = {
+    'C': 1.7,
+    'N': 1.55,
+    'O': 1.52,
+    'S': 1.8,
+}
+
+Bond = collections.namedtuple(
+    'Bond', ['atom1_name', 'atom2_name', 'length', 'stddev'])
+BondAngle = collections.namedtuple(
+    'BondAngle',
+    ['atom1_name', 'atom2_name', 'atom3name', 'angle_rad', 'stddev'])
+
+
+@functools.lru_cache(maxsize=None)
+def load_stereo_chemical_props() -> Tuple[Mapping[str, List[Bond]],
+                                          Mapping[str, List[Bond]],
+                                          Mapping[str, List[BondAngle]]]:
+  """Load stereo_chemical_props.txt into a nice structure.
+
+  Load literature values for bond lengths and bond angles and translate
+  bond angles into the length of the opposite edge of the triangle
+  ("residue_virtual_bonds").
+
+  Returns:
+    residue_bonds:  dict that maps resname --> list of Bond tuples
+    residue_virtual_bonds: dict that maps resname --> list of Bond tuples
+    residue_bond_angles: dict that maps resname --> list of BondAngle tuples
+  """
+  stereo_chemical_props_path = (
+      'alphafold/common/stereo_chemical_props.txt')
+  with open(stereo_chemical_props_path, 'rt') as f:
+    stereo_chemical_props = f.read()
+  lines_iter = iter(stereo_chemical_props.splitlines())
+  # Load bond lengths.
+  residue_bonds = {}
+  next(lines_iter)  # Skip header line.
+  for line in lines_iter:
+    if line.strip() == '-':
+      break
+    bond, resname, length, stddev = line.split()
+    atom1, atom2 = bond.split('-')
+    if resname not in residue_bonds:
+      residue_bonds[resname] = []
+    residue_bonds[resname].append(
+        Bond(atom1, atom2, float(length), float(stddev)))
+  residue_bonds['UNK'] = []
+
+  # Load bond angles.
+  residue_bond_angles = {}
+  next(lines_iter)  # Skip empty line.
+  next(lines_iter)  # Skip header line.
+  for line in lines_iter:
+    if line.strip() == '-':
+      break
+    bond, resname, angle_degree, stddev_degree = line.split()
+    atom1, atom2, atom3 = bond.split('-')
+    if resname not in residue_bond_angles:
+      residue_bond_angles[resname] = []
+    residue_bond_angles[resname].append(
+        BondAngle(atom1, atom2, atom3,
+                  float(angle_degree) / 180. * np.pi,
+                  float(stddev_degree) / 180. * np.pi))
+  residue_bond_angles['UNK'] = []
+
+  def make_bond_key(atom1_name, atom2_name):
+    """Unique key to lookup bonds."""
+    return '-'.join(sorted([atom1_name, atom2_name]))
+
+  # Translate bond angles into distances ("virtual bonds").
+  residue_virtual_bonds = {}
+  for resname, bond_angles in residue_bond_angles.items():
+    # Create a fast lookup dict for bond lengths.
+    bond_cache = {}
+    for b in residue_bonds[resname]:
+      bond_cache[make_bond_key(b.atom1_name, b.atom2_name)] = b
+    residue_virtual_bonds[resname] = []
+    for ba in bond_angles:
+      bond1 = bond_cache[make_bond_key(ba.atom1_name, ba.atom2_name)]
+      bond2 = bond_cache[make_bond_key(ba.atom2_name, ba.atom3name)]
+
+      # Compute distance between atom1 and atom3 using the law of cosines
+      # c^2 = a^2 + b^2 - 2ab*cos(gamma).
+      gamma = ba.angle_rad
+      length = np.sqrt(bond1.length**2 + bond2.length**2
+                       - 2 * bond1.length * bond2.length * np.cos(gamma))
+
+      # Propagation of uncertainty assuming uncorrelated errors.
+      dl_outer = 0.5 / length
+      dl_dgamma = (2 * bond1.length * bond2.length * np.sin(gamma)) * dl_outer
+      dl_db1 = (2 * bond1.length - 2 * bond2.length * np.cos(gamma)) * dl_outer
+      dl_db2 = (2 * bond2.length - 2 * bond1.length * np.cos(gamma)) * dl_outer
+      stddev = np.sqrt((dl_dgamma * ba.stddev)**2 +
+                       (dl_db1 * bond1.stddev)**2 +
+                       (dl_db2 * bond2.stddev)**2)
+      residue_virtual_bonds[resname].append(
+          Bond(ba.atom1_name, ba.atom3name, length, stddev))
+
+  return (residue_bonds,
+          residue_virtual_bonds,
+          residue_bond_angles)
+
+
+# Between-residue bond lengths for general bonds (first element) and for Proline
+# (second element).
+between_res_bond_length_c_n = [1.329, 1.341]
+between_res_bond_length_stddev_c_n = [0.014, 0.016]
+
+# Between-residue cos_angles.
+between_res_cos_angles_c_n_ca = [-0.5203, 0.0353]  # degrees: 121.352 +- 2.315
+between_res_cos_angles_ca_c_n = [-0.4473, 0.0311]  # degrees: 116.568 +- 1.995
+
+# This mapping is used when we need to store atom data in a format that requires
+# fixed atom data size for every residue (e.g. a numpy array).
+atom_types = [
+    'N', 'CA', 'C', 'CB', 'O', 'CG', 'CG1', 'CG2', 'OG', 'OG1', 'SG', 'CD',
+    'CD1', 'CD2', 'ND1', 'ND2', 'OD1', 'OD2', 'SD', 'CE', 'CE1', 'CE2', 'CE3',
+    'NE', 'NE1', 'NE2', 'OE1', 'OE2', 'CH2', 'NH1', 'NH2', 'OH', 'CZ', 'CZ2',
+    'CZ3', 'NZ', 'OXT'
+]
+atom_order = {atom_type: i for i, atom_type in enumerate(atom_types)}
+atom_type_num = len(atom_types)  # := 37.
+
+# A compact atom encoding with 14 columns
+# pylint: disable=line-too-long
+# pylint: disable=bad-whitespace
+restype_name_to_atom14_names = {
+    'ALA': ['N', 'CA', 'C', 'O', 'CB', '',    '',    '',    '',    '',    '',    '',    '',    ''],
+    'ARG': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'NE',  'CZ',  'NH1', 'NH2', '',    '',    ''],
+    'ASN': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'OD1', 'ND2', '',    '',    '',    '',    '',    ''],
+    'ASP': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'OD1', 'OD2', '',    '',    '',    '',    '',    ''],
+    'CYS': ['N', 'CA', 'C', 'O', 'CB', 'SG',  '',    '',    '',    '',    '',    '',    '',    ''],
+    'GLN': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'OE1', 'NE2', '',    '',    '',    '',    ''],
+    'GLU': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'OE1', 'OE2', '',    '',    '',    '',    ''],
+    'GLY': ['N', 'CA', 'C', 'O', '',   '',    '',    '',    '',    '',    '',    '',    '',    ''],
+    'HIS': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'ND1', 'CD2', 'CE1', 'NE2', '',    '',    '',    ''],
+    'ILE': ['N', 'CA', 'C', 'O', 'CB', 'CG1', 'CG2', 'CD1', '',    '',    '',    '',    '',    ''],
+    'LEU': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', '',    '',    '',    '',    '',    ''],
+    'LYS': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'CE',  'NZ',  '',    '',    '',    '',    ''],
+    'MET': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'SD',  'CE',  '',    '',    '',    '',    '',    ''],
+    'PHE': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', 'CE1', 'CE2', 'CZ',  '',    '',    ''],
+    'PRO': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  '',    '',    '',    '',    '',    '',    ''],
+    'SER': ['N', 'CA', 'C', 'O', 'CB', 'OG',  '',    '',    '',    '',    '',    '',    '',    ''],
+    'THR': ['N', 'CA', 'C', 'O', 'CB', 'OG1', 'CG2', '',    '',    '',    '',    '',    '',    ''],
+    'TRP': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', 'NE1', 'CE2', 'CE3', 'CZ2', 'CZ3', 'CH2'],
+    'TYR': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', 'CE1', 'CE2', 'CZ',  'OH',  '',    ''],
+    'VAL': ['N', 'CA', 'C', 'O', 'CB', 'CG1', 'CG2', '',    '',    '',    '',    '',    '',    ''],
+    'UNK': ['',  '',   '',  '',  '',   '',    '',    '',    '',    '',    '',    '',    '',    ''],
+
+}
+# pylint: enable=line-too-long
+# pylint: enable=bad-whitespace
+
+
+# This is the standard residue order when coding AA type as a number.
+# Reproduce it by taking 3-letter AA codes and sorting them alphabetically.
+restypes = [
+    'A', 'R', 'N', 'D', 'C', 'Q', 'E', 'G', 'H', 'I', 'L', 'K', 'M', 'F', 'P',
+    'S', 'T', 'W', 'Y', 'V'
+]
+restype_order = {restype: i for i, restype in enumerate(restypes)}
+restype_num = len(restypes)  # := 20.
+unk_restype_index = restype_num  # Catch-all index for unknown restypes.
+
+restypes_with_x = restypes + ['X']
+restype_order_with_x = {restype: i for i, restype in enumerate(restypes_with_x)}
+
+
+def sequence_to_onehot(
+    sequence: str,
+    mapping: Mapping[str, int],
+    map_unknown_to_x: bool = False) -> np.ndarray:
+  """Maps the given sequence into a one-hot encoded matrix.
+
+  Args:
+    sequence: An amino acid sequence.
+    mapping: A dictionary mapping amino acids to integers.
+    map_unknown_to_x: If True, any amino acid that is not in the mapping will be
+      mapped to the unknown amino acid 'X'. If the mapping doesn't contain
+      amino acid 'X', an error will be thrown. If False, any amino acid not in
+      the mapping will throw an error.
+
+  Returns:
+    A numpy array of shape (seq_len, num_unique_aas) with one-hot encoding of
+    the sequence.
+
+  Raises:
+    ValueError: If the mapping doesn't contain values from 0 to
+      num_unique_aas - 1 without any gaps.
+  """
+  num_entries = max(mapping.values()) + 1
+
+  if sorted(set(mapping.values())) != list(range(num_entries)):
+    raise ValueError('The mapping must have values from 0 to num_unique_aas-1 '
+                     'without any gaps. Got: %s' % sorted(mapping.values()))
+
+  one_hot_arr = np.zeros((len(sequence), num_entries), dtype=np.int32)
+
+  for aa_index, aa_type in enumerate(sequence):
+    if map_unknown_to_x:
+      if aa_type.isalpha() and aa_type.isupper():
+        aa_id = mapping.get(aa_type, mapping['X'])
+      else:
+        raise ValueError(f'Invalid character in the sequence: {aa_type}')
+    else:
+      aa_id = mapping[aa_type]
+    one_hot_arr[aa_index, aa_id] = 1
+
+  return one_hot_arr
+
+
+restype_1to3 = {
+    'A': 'ALA',
+    'R': 'ARG',
+    'N': 'ASN',
+    'D': 'ASP',
+    'C': 'CYS',
+    'Q': 'GLN',
+    'E': 'GLU',
+    'G': 'GLY',
+    'H': 'HIS',
+    'I': 'ILE',
+    'L': 'LEU',
+    'K': 'LYS',
+    'M': 'MET',
+    'F': 'PHE',
+    'P': 'PRO',
+    'S': 'SER',
+    'T': 'THR',
+    'W': 'TRP',
+    'Y': 'TYR',
+    'V': 'VAL',
+}
+
+
+# NB: restype_3to1 differs from Bio.PDB.protein_letters_3to1 by being a simple
+# 1-to-1 mapping of 3 letter names to one letter names. The latter contains
+# many more, and less common, three letter names as keys and maps many of these
+# to the same one letter name (including 'X' and 'U' which we don't use here).
+restype_3to1 = {v: k for k, v in restype_1to3.items()}
+
+# Define a restype name for all unknown residues.
+unk_restype = 'UNK'
+
+resnames = [restype_1to3[r] for r in restypes] + [unk_restype]
+resname_to_idx = {resname: i for i, resname in enumerate(resnames)}
+
+
+# The mapping here uses hhblits convention, so that B is mapped to D, J and O
+# are mapped to X, U is mapped to C, and Z is mapped to E. Other than that the
+# remaining 20 amino acids are kept in alphabetical order.
+# There are 2 non-amino acid codes, X (representing any amino acid) and
+# "-" representing a missing amino acid in an alignment.  The id for these
+# codes is put at the end (20 and 21) so that they can easily be ignored if
+# desired.
+HHBLITS_AA_TO_ID = {
+    'A': 0,
+    'B': 2,
+    'C': 1,
+    'D': 2,
+    'E': 3,
+    'F': 4,
+    'G': 5,
+    'H': 6,
+    'I': 7,
+    'J': 20,
+    'K': 8,
+    'L': 9,
+    'M': 10,
+    'N': 11,
+    'O': 20,
+    'P': 12,
+    'Q': 13,
+    'R': 14,
+    'S': 15,
+    'T': 16,
+    'U': 1,
+    'V': 17,
+    'W': 18,
+    'X': 20,
+    'Y': 19,
+    'Z': 3,
+    '-': 21,
+}
+
+# Partial inversion of HHBLITS_AA_TO_ID.
+ID_TO_HHBLITS_AA = {
+    0: 'A',
+    1: 'C',  # Also U.
+    2: 'D',  # Also B.
+    3: 'E',  # Also Z.
+    4: 'F',
+    5: 'G',
+    6: 'H',
+    7: 'I',
+    8: 'K',
+    9: 'L',
+    10: 'M',
+    11: 'N',
+    12: 'P',
+    13: 'Q',
+    14: 'R',
+    15: 'S',
+    16: 'T',
+    17: 'V',
+    18: 'W',
+    19: 'Y',
+    20: 'X',  # Includes J and O.
+    21: '-',
+}
+
+restypes_with_x_and_gap = restypes + ['X', '-']
+MAP_HHBLITS_AATYPE_TO_OUR_AATYPE = tuple(
+    restypes_with_x_and_gap.index(ID_TO_HHBLITS_AA[i])
+    for i in range(len(restypes_with_x_and_gap)))
+
+
+def _make_standard_atom_mask() -> np.ndarray:
+  """Returns [num_res_types, num_atom_types] mask array."""
+  # +1 to account for unknown (all 0s).
+  mask = np.zeros([restype_num + 1, atom_type_num], dtype=np.int32)
+  for restype, restype_letter in enumerate(restypes):
+    restype_name = restype_1to3[restype_letter]
+    atom_names = residue_atoms[restype_name]
+    for atom_name in atom_names:
+      atom_type = atom_order[atom_name]
+      mask[restype, atom_type] = 1
+  return mask
+
+
+STANDARD_ATOM_MASK = _make_standard_atom_mask()
+
+
+# A one hot representation for the first and second atoms defining the axis
+# of rotation for each chi-angle in each residue.
+def chi_angle_atom(atom_index: int) -> np.ndarray:
+  """Define chi-angle rigid groups via one-hot representations."""
+  chi_angles_index = {}
+  one_hots = []
+
+  for k, v in chi_angles_atoms.items():
+    indices = [atom_types.index(s[atom_index]) for s in v]
+    indices.extend([-1]*(4-len(indices)))
+    chi_angles_index[k] = indices
+
+  for r in restypes:
+    res3 = restype_1to3[r]
+    one_hot = np.eye(atom_type_num)[chi_angles_index[res3]]
+    one_hots.append(one_hot)
+
+  one_hots.append(np.zeros([4, atom_type_num]))  # Add zeros for residue `X`.
+  one_hot = np.stack(one_hots, axis=0)
+  one_hot = np.transpose(one_hot, [0, 2, 1])
+
+  return one_hot
+
+chi_atom_1_one_hot = chi_angle_atom(1)
+chi_atom_2_one_hot = chi_angle_atom(2)
+
+# An array like chi_angles_atoms but using indices rather than names.
+chi_angles_atom_indices = [chi_angles_atoms[restype_1to3[r]] for r in restypes]
+chi_angles_atom_indices = tree.map_structure(
+    lambda atom_name: atom_order[atom_name], chi_angles_atom_indices)
+chi_angles_atom_indices = np.array([
+    chi_atoms + ([[0, 0, 0, 0]] * (4 - len(chi_atoms)))
+    for chi_atoms in chi_angles_atom_indices])
+
+# Mapping from (res_name, atom_name) pairs to the atom's chi group index
+# and atom index within that group.
+chi_groups_for_atom = collections.defaultdict(list)
+for res_name, chi_angle_atoms_for_res in chi_angles_atoms.items():
+  for chi_group_i, chi_group in enumerate(chi_angle_atoms_for_res):
+    for atom_i, atom in enumerate(chi_group):
+      chi_groups_for_atom[(res_name, atom)].append((chi_group_i, atom_i))
+chi_groups_for_atom = dict(chi_groups_for_atom)
+
+
+def _make_rigid_transformation_4x4(ex, ey, translation):
+  """Create a rigid 4x4 transformation matrix from two axes and transl."""
+  # Normalize ex.
+  ex_normalized = ex / np.linalg.norm(ex)
+
+  # make ey perpendicular to ex
+  ey_normalized = ey - np.dot(ey, ex_normalized) * ex_normalized
+  ey_normalized /= np.linalg.norm(ey_normalized)
+
+  # compute ez as cross product
+  eznorm = np.cross(ex_normalized, ey_normalized)
+  m = np.stack([ex_normalized, ey_normalized, eznorm, translation]).transpose()
+  m = np.concatenate([m, [[0., 0., 0., 1.]]], axis=0)
+  return m
+
+
+# create an array with (restype, atomtype) --> rigid_group_idx
+# and an array with (restype, atomtype, coord) for the atom positions
+# and compute affine transformation matrices (4,4) from one rigid group to the
+# previous group
+restype_atom37_to_rigid_group = np.zeros([21, 37], dtype=np.int)
+restype_atom37_mask = np.zeros([21, 37], dtype=np.float32)
+restype_atom37_rigid_group_positions = np.zeros([21, 37, 3], dtype=np.float32)
+restype_atom14_to_rigid_group = np.zeros([21, 14], dtype=np.int)
+restype_atom14_mask = np.zeros([21, 14], dtype=np.float32)
+restype_atom14_rigid_group_positions = np.zeros([21, 14, 3], dtype=np.float32)
+restype_rigid_group_default_frame = np.zeros([21, 8, 4, 4], dtype=np.float32)
+
+
+def _make_rigid_group_constants():
+  """Fill the arrays above."""
+  for restype, restype_letter in enumerate(restypes):
+    resname = restype_1to3[restype_letter]
+    for atomname, group_idx, atom_position in rigid_group_atom_positions[
+        resname]:
+      atomtype = atom_order[atomname]
+      restype_atom37_to_rigid_group[restype, atomtype] = group_idx
+      restype_atom37_mask[restype, atomtype] = 1
+      restype_atom37_rigid_group_positions[restype, atomtype, :] = atom_position
+
+      atom14idx = restype_name_to_atom14_names[resname].index(atomname)
+      restype_atom14_to_rigid_group[restype, atom14idx] = group_idx
+      restype_atom14_mask[restype, atom14idx] = 1
+      restype_atom14_rigid_group_positions[restype,
+                                           atom14idx, :] = atom_position
+
+  for restype, restype_letter in enumerate(restypes):
+    resname = restype_1to3[restype_letter]
+    atom_positions = {name: np.array(pos) for name, _, pos
+                      in rigid_group_atom_positions[resname]}
+
+    # backbone to backbone is the identity transform
+    restype_rigid_group_default_frame[restype, 0, :, :] = np.eye(4)
+
+    # pre-omega-frame to backbone (currently dummy identity matrix)
+    restype_rigid_group_default_frame[restype, 1, :, :] = np.eye(4)
+
+    # phi-frame to backbone
+    mat = _make_rigid_transformation_4x4(
+        ex=atom_positions['N'] - atom_positions['CA'],
+        ey=np.array([1., 0., 0.]),
+        translation=atom_positions['N'])
+    restype_rigid_group_default_frame[restype, 2, :, :] = mat
+
+    # psi-frame to backbone
+    mat = _make_rigid_transformation_4x4(
+        ex=atom_positions['C'] - atom_positions['CA'],
+        ey=atom_positions['CA'] - atom_positions['N'],
+        translation=atom_positions['C'])
+    restype_rigid_group_default_frame[restype, 3, :, :] = mat
+
+    # chi1-frame to backbone
+    if chi_angles_mask[restype][0]:
+      base_atom_names = chi_angles_atoms[resname][0]
+      base_atom_positions = [atom_positions[name] for name in base_atom_names]
+      mat = _make_rigid_transformation_4x4(
+          ex=base_atom_positions[2] - base_atom_positions[1],
+          ey=base_atom_positions[0] - base_atom_positions[1],
+          translation=base_atom_positions[2])
+      restype_rigid_group_default_frame[restype, 4, :, :] = mat
+
+    # chi2-frame to chi1-frame
+    # chi3-frame to chi2-frame
+    # chi4-frame to chi3-frame
+    # luckily all rotation axes for the next frame start at (0,0,0) of the
+    # previous frame
+    for chi_idx in range(1, 4):
+      if chi_angles_mask[restype][chi_idx]:
+        axis_end_atom_name = chi_angles_atoms[resname][chi_idx][2]
+        axis_end_atom_position = atom_positions[axis_end_atom_name]
+        mat = _make_rigid_transformation_4x4(
+            ex=axis_end_atom_position,
+            ey=np.array([-1., 0., 0.]),
+            translation=axis_end_atom_position)
+        restype_rigid_group_default_frame[restype, 4 + chi_idx, :, :] = mat
+
+
+_make_rigid_group_constants()
+
+
+def make_atom14_dists_bounds(overlap_tolerance=1.5,
+                             bond_length_tolerance_factor=15):
+  """compute upper and lower bounds for bonds to assess violations."""
+  restype_atom14_bond_lower_bound = np.zeros([21, 14, 14], np.float32)
+  restype_atom14_bond_upper_bound = np.zeros([21, 14, 14], np.float32)
+  restype_atom14_bond_stddev = np.zeros([21, 14, 14], np.float32)
+  residue_bonds, residue_virtual_bonds, _ = load_stereo_chemical_props()
+  for restype, restype_letter in enumerate(restypes):
+    resname = restype_1to3[restype_letter]
+    atom_list = restype_name_to_atom14_names[resname]
+
+    # create lower and upper bounds for clashes
+    for atom1_idx, atom1_name in enumerate(atom_list):
+      if not atom1_name:
+        continue
+      atom1_radius = van_der_waals_radius[atom1_name[0]]
+      for atom2_idx, atom2_name in enumerate(atom_list):
+        if (not atom2_name) or atom1_idx == atom2_idx:
+          continue
+        atom2_radius = van_der_waals_radius[atom2_name[0]]
+        lower = atom1_radius + atom2_radius - overlap_tolerance
+        upper = 1e10
+        restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
+        restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
+        restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
+        restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper
+
+    # overwrite lower and upper bounds for bonds and angles
+    for b in residue_bonds[resname] + residue_virtual_bonds[resname]:
+      atom1_idx = atom_list.index(b.atom1_name)
+      atom2_idx = atom_list.index(b.atom2_name)
+      lower = b.length - bond_length_tolerance_factor * b.stddev
+      upper = b.length + bond_length_tolerance_factor * b.stddev
+      restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
+      restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
+      restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
+      restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper
+      restype_atom14_bond_stddev[restype, atom1_idx, atom2_idx] = b.stddev
+      restype_atom14_bond_stddev[restype, atom2_idx, atom1_idx] = b.stddev
+  return {'lower_bound': restype_atom14_bond_lower_bound,  # shape (21,14,14)
+          'upper_bound': restype_atom14_bond_upper_bound,  # shape (21,14,14)
+          'stddev': restype_atom14_bond_stddev,  # shape (21,14,14)
+         }

alphafold/alphafold/common/residue_constants_test.py

@@ -0,0 +1,190 @@
+# 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.
+
+"""Test that residue_constants generates correct values."""
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from alphafold.common import residue_constants
+import numpy as np
+
+
+class ResidueConstantsTest(parameterized.TestCase):
+
+  @parameterized.parameters(
+      ('ALA', 0),
+      ('CYS', 1),
+      ('HIS', 2),
+      ('MET', 3),
+      ('LYS', 4),
+      ('ARG', 4),
+  )
+  def testChiAnglesAtoms(self, residue_name, chi_num):
+    chi_angles_atoms = residue_constants.chi_angles_atoms[residue_name]
+    self.assertLen(chi_angles_atoms, chi_num)
+    for chi_angle_atoms in chi_angles_atoms:
+      self.assertLen(chi_angle_atoms, 4)
+
+  def testChiGroupsForAtom(self):
+    for k, chi_groups in residue_constants.chi_groups_for_atom.items():
+      res_name, atom_name = k
+      for chi_group_i, atom_i in chi_groups:
+        self.assertEqual(
+            atom_name,
+            residue_constants.chi_angles_atoms[res_name][chi_group_i][atom_i])
+
+  @parameterized.parameters(
+      ('ALA', 5), ('ARG', 11), ('ASN', 8), ('ASP', 8), ('CYS', 6), ('GLN', 9),
+      ('GLU', 9), ('GLY', 4), ('HIS', 10), ('ILE', 8), ('LEU', 8), ('LYS', 9),
+      ('MET', 8), ('PHE', 11), ('PRO', 7), ('SER', 6), ('THR', 7), ('TRP', 14),
+      ('TYR', 12), ('VAL', 7)
+  )
+  def testResidueAtoms(self, atom_name, num_residue_atoms):
+    residue_atoms = residue_constants.residue_atoms[atom_name]
+    self.assertLen(residue_atoms, num_residue_atoms)
+
+  def testStandardAtomMask(self):
+    with self.subTest('Check shape'):
+      self.assertEqual(residue_constants.STANDARD_ATOM_MASK.shape, (21, 37,))
+
+    with self.subTest('Check values'):
+      str_to_row = lambda s: [c == '1' for c in s]  # More clear/concise.
+      np.testing.assert_array_equal(
+          residue_constants.STANDARD_ATOM_MASK,
+          np.array([
+              # NB This was defined by c+p but looks sane.
+              str_to_row('11111                                '),  # ALA
+              str_to_row('111111     1           1     11 1    '),  # ARG
+              str_to_row('111111         11                    '),  # ASP
+              str_to_row('111111          11                   '),  # ASN
+              str_to_row('11111     1                          '),  # CYS
+              str_to_row('111111     1             11          '),  # GLU
+              str_to_row('111111     1              11         '),  # GLN
+              str_to_row('111 1                                '),  # GLY
+              str_to_row('111111       11     1    1           '),  # HIS
+              str_to_row('11111 11    1                        '),  # ILE
+              str_to_row('111111      11                       '),  # LEU
+              str_to_row('111111     1       1               1 '),  # LYS
+              str_to_row('111111            11                 '),  # MET
+              str_to_row('111111      11      11          1    '),  # PHE
+              str_to_row('111111     1                         '),  # PRO
+              str_to_row('11111   1                            '),  # SER
+              str_to_row('11111  1 1                           '),  # THR
+              str_to_row('111111      11       11 1   1    11  '),  # TRP
+              str_to_row('111111      11      11         11    '),  # TYR
+              str_to_row('11111 11                             '),  # VAL
+              str_to_row('                                     '),  # UNK
+          ]))
+
+    with self.subTest('Check row totals'):
+      # Check each row has the right number of atoms.
+      for row, restype in enumerate(residue_constants.restypes):  # A, R, ...
+        long_restype = residue_constants.restype_1to3[restype]  # ALA, ARG, ...
+        atoms_names = residue_constants.residue_atoms[
+            long_restype]  # ['C', 'CA', 'CB', 'N', 'O'], ...
+        self.assertLen(atoms_names,
+                       residue_constants.STANDARD_ATOM_MASK[row, :].sum(),
+                       long_restype)
+
+  def testAtomTypes(self):
+    self.assertEqual(residue_constants.atom_type_num, 37)
+
+    self.assertEqual(residue_constants.atom_types[0], 'N')
+    self.assertEqual(residue_constants.atom_types[1], 'CA')
+    self.assertEqual(residue_constants.atom_types[2], 'C')
+    self.assertEqual(residue_constants.atom_types[3], 'CB')
+    self.assertEqual(residue_constants.atom_types[4], 'O')
+
+    self.assertEqual(residue_constants.atom_order['N'], 0)
+    self.assertEqual(residue_constants.atom_order['CA'], 1)
+    self.assertEqual(residue_constants.atom_order['C'], 2)
+    self.assertEqual(residue_constants.atom_order['CB'], 3)
+    self.assertEqual(residue_constants.atom_order['O'], 4)
+    self.assertEqual(residue_constants.atom_type_num, 37)
+
+  def testRestypes(self):
+    three_letter_restypes = [
+        residue_constants.restype_1to3[r] for r  in residue_constants.restypes]
+    for restype, exp_restype in zip(
+        three_letter_restypes, sorted(residue_constants.restype_1to3.values())):
+      self.assertEqual(restype, exp_restype)
+    self.assertEqual(residue_constants.restype_num, 20)
+
+  def testSequenceToOneHotHHBlits(self):
+    one_hot = residue_constants.sequence_to_onehot(
+        'ABCDEFGHIJKLMNOPQRSTUVWXYZ-', residue_constants.HHBLITS_AA_TO_ID)
+    exp_one_hot = np.array(
+        [[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
+         [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
+         [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]])
+    np.testing.assert_array_equal(one_hot, exp_one_hot)
+
+  def testSequenceToOneHotStandard(self):
+    one_hot = residue_constants.sequence_to_onehot(
+        'ARNDCQEGHILKMFPSTWYV', residue_constants.restype_order)
+    np.testing.assert_array_equal(one_hot, np.eye(20))
+
+  def testSequenceToOneHotUnknownMapping(self):
+    seq = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
+    expected_out = np.zeros([26, 21])
+    for row, position in enumerate(
+        [0, 20, 4, 3, 6, 13, 7, 8, 9, 20, 11, 10, 12, 2, 20, 14, 5, 1, 15, 16,
+         20, 19, 17, 20, 18, 20]):
+      expected_out[row, position] = 1
+    aa_types = residue_constants.sequence_to_onehot(
+        sequence=seq,
+        mapping=residue_constants.restype_order_with_x,
+        map_unknown_to_x=True)
+    self.assertTrue((aa_types == expected_out).all())
+
+  @parameterized.named_parameters(
+      ('lowercase', 'aaa'),  # Insertions in A3M.
+      ('gaps', '---'),  # Gaps in A3M.
+      ('dots', '...'),  # Gaps in A3M.
+      ('metadata', '>TEST'),  # FASTA metadata line.
+  )
+  def testSequenceToOneHotUnknownMappingError(self, seq):
+    with self.assertRaises(ValueError):
+      residue_constants.sequence_to_onehot(
+          sequence=seq,
+          mapping=residue_constants.restype_order_with_x,
+          map_unknown_to_x=True)
+
+
+if __name__ == '__main__':
+  absltest.main()

alphafold/alphafold/data/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+"""Data pipeline for model features."""

alphafold/alphafold/data/mmcif_parsing.py

@@ -0,0 +1,384 @@
+# 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.
+
+"""Parses the mmCIF file format."""
+import collections
+import dataclasses
+import io
+from typing import Any, Mapping, Optional, Sequence, Tuple
+
+from absl import logging
+from Bio import PDB
+from Bio.Data import SCOPData
+
+# Type aliases:
+ChainId = str
+PdbHeader = Mapping[str, Any]
+PdbStructure = PDB.Structure.Structure
+SeqRes = str
+MmCIFDict = Mapping[str, Sequence[str]]
+
+
+@dataclasses.dataclass(frozen=True)
+class Monomer:
+  id: str
+  num: int
+
+
+# Note - mmCIF format provides no guarantees on the type of author-assigned
+# sequence numbers. They need not be integers.
+@dataclasses.dataclass(frozen=True)
+class AtomSite:
+  residue_name: str
+  author_chain_id: str
+  mmcif_chain_id: str
+  author_seq_num: str
+  mmcif_seq_num: int
+  insertion_code: str
+  hetatm_atom: str
+  model_num: int
+
+
+# Used to map SEQRES index to a residue in the structure.
+@dataclasses.dataclass(frozen=True)
+class ResiduePosition:
+  chain_id: str
+  residue_number: int
+  insertion_code: str
+
+
+@dataclasses.dataclass(frozen=True)
+class ResidueAtPosition:
+  position: Optional[ResiduePosition]
+  name: str
+  is_missing: bool
+  hetflag: str
+
+
+@dataclasses.dataclass(frozen=True)
+class MmcifObject:
+  """Representation of a parsed mmCIF file.
+
+  Contains:
+    file_id: A meaningful name, e.g. a pdb_id. Should be unique amongst all
+      files being processed.
+    header: Biopython header.
+    structure: Biopython structure.
+    chain_to_seqres: Dict mapping chain_id to 1 letter amino acid sequence. E.g.
+      {'A': 'ABCDEFG'}
+    seqres_to_structure: Dict; for each chain_id contains a mapping between
+      SEQRES index and a ResidueAtPosition. e.g. {'A': {0: ResidueAtPosition,
+                                                        1: ResidueAtPosition,
+                                                        ...}}
+    raw_string: The raw string used to construct the MmcifObject.
+  """
+  file_id: str
+  header: PdbHeader
+  structure: PdbStructure
+  chain_to_seqres: Mapping[ChainId, SeqRes]
+  seqres_to_structure: Mapping[ChainId, Mapping[int, ResidueAtPosition]]
+  raw_string: Any
+
+
+@dataclasses.dataclass(frozen=True)
+class ParsingResult:
+  """Returned by the parse function.
+
+  Contains:
+    mmcif_object: A MmcifObject, may be None if no chain could be successfully
+      parsed.
+    errors: A dict mapping (file_id, chain_id) to any exception generated.
+  """
+  mmcif_object: Optional[MmcifObject]
+  errors: Mapping[Tuple[str, str], Any]
+
+
+class ParseError(Exception):
+  """An error indicating that an mmCIF file could not be parsed."""
+
+
+def mmcif_loop_to_list(prefix: str,
+                       parsed_info: MmCIFDict) -> Sequence[Mapping[str, str]]:
+  """Extracts loop associated with a prefix from mmCIF data as a list.
+
+  Reference for loop_ in mmCIF:
+    http://mmcif.wwpdb.org/docs/tutorials/mechanics/pdbx-mmcif-syntax.html
+
+  Args:
+    prefix: Prefix shared by each of the data items in the loop.
+      e.g. '_entity_poly_seq.', where the data items are _entity_poly_seq.num,
+      _entity_poly_seq.mon_id. Should include the trailing period.
+    parsed_info: A dict of parsed mmCIF data, e.g. _mmcif_dict from a Biopython
+      parser.
+
+  Returns:
+    Returns a list of dicts; each dict represents 1 entry from an mmCIF loop.
+  """
+  cols = []
+  data = []
+  for key, value in parsed_info.items():
+    if key.startswith(prefix):
+      cols.append(key)
+      data.append(value)
+
+  assert all([len(xs) == len(data[0]) for xs in data]), (
+      'mmCIF error: Not all loops are the same length: %s' % cols)
+
+  return [dict(zip(cols, xs)) for xs in zip(*data)]
+
+
+def mmcif_loop_to_dict(prefix: str,
+                       index: str,
+                       parsed_info: MmCIFDict,
+                       ) -> Mapping[str, Mapping[str, str]]:
+  """Extracts loop associated with a prefix from mmCIF data as a dictionary.
+
+  Args:
+    prefix: Prefix shared by each of the data items in the loop.
+      e.g. '_entity_poly_seq.', where the data items are _entity_poly_seq.num,
+      _entity_poly_seq.mon_id. Should include the trailing period.
+    index: Which item of loop data should serve as the key.
+    parsed_info: A dict of parsed mmCIF data, e.g. _mmcif_dict from a Biopython
+      parser.
+
+  Returns:
+    Returns a dict of dicts; each dict represents 1 entry from an mmCIF loop,
+    indexed by the index column.
+  """
+  entries = mmcif_loop_to_list(prefix, parsed_info)
+  return {entry[index]: entry for entry in entries}
+
+
+def parse(*,
+          file_id: str,
+          mmcif_string: str,
+          catch_all_errors: bool = True) -> ParsingResult:
+  """Entry point, parses an mmcif_string.
+
+  Args:
+    file_id: A string identifier for this file. Should be unique within the
+      collection of files being processed.
+    mmcif_string: Contents of an mmCIF file.
+    catch_all_errors: If True, all exceptions are caught and error messages are
+      returned as part of the ParsingResult. If False exceptions will be allowed
+      to propagate.
+
+  Returns:
+    A ParsingResult.
+  """
+  errors = {}
+  try:
+    parser = PDB.MMCIFParser(QUIET=True)
+    handle = io.StringIO(mmcif_string)
+    full_structure = parser.get_structure('', handle)
+    first_model_structure = _get_first_model(full_structure)
+    # Extract the _mmcif_dict from the parser, which contains useful fields not
+    # reflected in the Biopython structure.
+    parsed_info = parser._mmcif_dict  # pylint:disable=protected-access
+
+    # Ensure all values are lists, even if singletons.
+    for key, value in parsed_info.items():
+      if not isinstance(value, list):
+        parsed_info[key] = [value]
+
+    header = _get_header(parsed_info)
+
+    # Determine the protein chains, and their start numbers according to the
+    # internal mmCIF numbering scheme (likely but not guaranteed to be 1).
+    valid_chains = _get_protein_chains(parsed_info=parsed_info)
+    if not valid_chains:
+      return ParsingResult(
+          None, {(file_id, ''): 'No protein chains found in this file.'})
+    seq_start_num = {chain_id: min([monomer.num for monomer in seq])
+                     for chain_id, seq in valid_chains.items()}
+
+    # Loop over the atoms for which we have coordinates. Populate two mappings:
+    # -mmcif_to_author_chain_id (maps internal mmCIF chain ids to chain ids used
+    # the authors / Biopython).
+    # -seq_to_structure_mappings (maps idx into sequence to ResidueAtPosition).
+    mmcif_to_author_chain_id = {}
+    seq_to_structure_mappings = {}
+    for atom in _get_atom_site_list(parsed_info):
+      if atom.model_num != '1':
+        # We only process the first model at the moment.
+        continue
+
+      mmcif_to_author_chain_id[atom.mmcif_chain_id] = atom.author_chain_id
+
+      if atom.mmcif_chain_id in valid_chains:
+        hetflag = ' '
+        if atom.hetatm_atom == 'HETATM':
+          # Water atoms are assigned a special hetflag of W in Biopython. We
+          # need to do the same, so that this hetflag can be used to fetch
+          # a residue from the Biopython structure by id.
+          if atom.residue_name in ('HOH', 'WAT'):
+            hetflag = 'W'
+          else:
+            hetflag = 'H_' + atom.residue_name
+        insertion_code = atom.insertion_code
+        if not _is_set(atom.insertion_code):
+          insertion_code = ' '
+        position = ResiduePosition(chain_id=atom.author_chain_id,
+                                   residue_number=int(atom.author_seq_num),
+                                   insertion_code=insertion_code)
+        seq_idx = int(atom.mmcif_seq_num) - seq_start_num[atom.mmcif_chain_id]
+        current = seq_to_structure_mappings.get(atom.author_chain_id, {})
+        current[seq_idx] = ResidueAtPosition(position=position,
+                                             name=atom.residue_name,
+                                             is_missing=False,
+                                             hetflag=hetflag)
+        seq_to_structure_mappings[atom.author_chain_id] = current
+
+    # Add missing residue information to seq_to_structure_mappings.
+    for chain_id, seq_info in valid_chains.items():
+      author_chain = mmcif_to_author_chain_id[chain_id]
+      current_mapping = seq_to_structure_mappings[author_chain]
+      for idx, monomer in enumerate(seq_info):
+        if idx not in current_mapping:
+          current_mapping[idx] = ResidueAtPosition(position=None,
+                                                   name=monomer.id,
+                                                   is_missing=True,
+                                                   hetflag=' ')
+
+    author_chain_to_sequence = {}
+    for chain_id, seq_info in valid_chains.items():
+      author_chain = mmcif_to_author_chain_id[chain_id]
+      seq = []
+      for monomer in seq_info:
+        code = SCOPData.protein_letters_3to1.get(monomer.id, 'X')
+        seq.append(code if len(code) == 1 else 'X')
+      seq = ''.join(seq)
+      author_chain_to_sequence[author_chain] = seq
+
+    mmcif_object = MmcifObject(
+        file_id=file_id,
+        header=header,
+        structure=first_model_structure,
+        chain_to_seqres=author_chain_to_sequence,
+        seqres_to_structure=seq_to_structure_mappings,
+        raw_string=parsed_info)
+
+    return ParsingResult(mmcif_object=mmcif_object, errors=errors)
+  except Exception as e:  # pylint:disable=broad-except
+    errors[(file_id, '')] = e
+    if not catch_all_errors:
+      raise
+    return ParsingResult(mmcif_object=None, errors=errors)
+
+
+def _get_first_model(structure: PdbStructure) -> PdbStructure:
+  """Returns the first model in a Biopython structure."""
+  return next(structure.get_models())
+
+_MIN_LENGTH_OF_CHAIN_TO_BE_COUNTED_AS_PEPTIDE = 21
+
+
+def get_release_date(parsed_info: MmCIFDict) -> str:
+  """Returns the oldest revision date."""
+  revision_dates = parsed_info['_pdbx_audit_revision_history.revision_date']
+  return min(revision_dates)
+
+
+def _get_header(parsed_info: MmCIFDict) -> PdbHeader:
+  """Returns a basic header containing method, release date and resolution."""
+  header = {}
+
+  experiments = mmcif_loop_to_list('_exptl.', parsed_info)
+  header['structure_method'] = ','.join([
+      experiment['_exptl.method'].lower() for experiment in experiments])
+
+  # Note: The release_date here corresponds to the oldest revision. We prefer to
+  # use this for dataset filtering over the deposition_date.
+  if '_pdbx_audit_revision_history.revision_date' in parsed_info:
+    header['release_date'] = get_release_date(parsed_info)
+  else:
+    logging.warning('Could not determine release_date: %s',
+                    parsed_info['_entry.id'])
+
+  header['resolution'] = 0.00
+  for res_key in ('_refine.ls_d_res_high', '_em_3d_reconstruction.resolution',
+                  '_reflns.d_resolution_high'):
+    if res_key in parsed_info:
+      try:
+        raw_resolution = parsed_info[res_key][0]
+        header['resolution'] = float(raw_resolution)
+      except ValueError:
+        logging.warning('Invalid resolution format: %s', parsed_info[res_key])
+
+  return header
+
+
+def _get_atom_site_list(parsed_info: MmCIFDict) -> Sequence[AtomSite]:
+  """Returns list of atom sites; contains data not present in the structure."""
+  return [AtomSite(*site) for site in zip(  # pylint:disable=g-complex-comprehension
+      parsed_info['_atom_site.label_comp_id'],
+      parsed_info['_atom_site.auth_asym_id'],
+      parsed_info['_atom_site.label_asym_id'],
+      parsed_info['_atom_site.auth_seq_id'],
+      parsed_info['_atom_site.label_seq_id'],
+      parsed_info['_atom_site.pdbx_PDB_ins_code'],
+      parsed_info['_atom_site.group_PDB'],
+      parsed_info['_atom_site.pdbx_PDB_model_num'],
+      )]
+
+
+def _get_protein_chains(
+    *, parsed_info: Mapping[str, Any]) -> Mapping[ChainId, Sequence[Monomer]]:
+  """Extracts polymer information for protein chains only.
+
+  Args:
+    parsed_info: _mmcif_dict produced by the Biopython parser.
+
+  Returns:
+    A dict mapping mmcif chain id to a list of Monomers.
+  """
+  # Get polymer information for each entity in the structure.
+  entity_poly_seqs = mmcif_loop_to_list('_entity_poly_seq.', parsed_info)
+
+  polymers = collections.defaultdict(list)
+  for entity_poly_seq in entity_poly_seqs:
+    polymers[entity_poly_seq['_entity_poly_seq.entity_id']].append(
+        Monomer(id=entity_poly_seq['_entity_poly_seq.mon_id'],
+                num=int(entity_poly_seq['_entity_poly_seq.num'])))
+
+  # Get chemical compositions. Will allow us to identify which of these polymers
+  # are proteins.
+  chem_comps = mmcif_loop_to_dict('_chem_comp.', '_chem_comp.id', parsed_info)
+
+  # Get chains information for each entity. Necessary so that we can return a
+  # dict keyed on chain id rather than entity.
+  struct_asyms = mmcif_loop_to_list('_struct_asym.', parsed_info)
+
+  entity_to_mmcif_chains = collections.defaultdict(list)
+  for struct_asym in struct_asyms:
+    chain_id = struct_asym['_struct_asym.id']
+    entity_id = struct_asym['_struct_asym.entity_id']
+    entity_to_mmcif_chains[entity_id].append(chain_id)
+
+  # Identify and return the valid protein chains.
+  valid_chains = {}
+  for entity_id, seq_info in polymers.items():
+    chain_ids = entity_to_mmcif_chains[entity_id]
+
+    # Reject polymers without any peptide-like components, such as DNA/RNA.
+    if any(['peptide' in chem_comps[monomer.id]['_chem_comp.type']
+            for monomer in seq_info]):
+      for chain_id in chain_ids:
+        valid_chains[chain_id] = seq_info
+  return valid_chains
+
+
+def _is_set(data: str) -> bool:
+  """Returns False if data is a special mmCIF character indicating 'unset'."""
+  return data not in ('.', '?')

alphafold/alphafold/data/parsers.py

@@ -0,0 +1,364 @@
+# 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.
+
+"""Functions for parsing various file formats."""
+import collections
+import dataclasses
+import re
+import string
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple
+
+DeletionMatrix = Sequence[Sequence[int]]
+
+
+@dataclasses.dataclass(frozen=True)
+class TemplateHit:
+  """Class representing a template hit."""
+  index: int
+  name: str
+  aligned_cols: int
+  sum_probs: float
+  query: str
+  hit_sequence: str
+  indices_query: List[int]
+  indices_hit: List[int]
+
+
+def parse_fasta(fasta_string: str) -> Tuple[Sequence[str], Sequence[str]]:
+  """Parses FASTA string and returns list of strings with amino-acid sequences.
+
+  Arguments:
+    fasta_string: The string contents of a FASTA file.
+
+  Returns:
+    A tuple of two lists:
+    * A list of sequences.
+    * A list of sequence descriptions taken from the comment lines. In the
+      same order as the sequences.
+  """
+  sequences = []
+  descriptions = []
+  index = -1
+  for line in fasta_string.splitlines():
+    line = line.strip()
+    if line.startswith('>'):
+      index += 1
+      descriptions.append(line[1:])  # Remove the '>' at the beginning.
+      sequences.append('')
+      continue
+    elif not line:
+      continue  # Skip blank lines.
+    sequences[index] += line
+
+  return sequences, descriptions
+
+
+def parse_stockholm(
+    stockholm_string: str
+) -> Tuple[Sequence[str], DeletionMatrix, Sequence[str]]:
+  """Parses sequences and deletion matrix from stockholm format alignment.
+
+  Args:
+    stockholm_string: The string contents of a stockholm file. The first
+      sequence in the file should be the query sequence.
+
+  Returns:
+    A tuple of:
+      * A list of sequences that have been aligned to the query. These
+        might contain duplicates.
+      * The deletion matrix for the alignment as a list of lists. The element
+        at `deletion_matrix[i][j]` is the number of residues deleted from
+        the aligned sequence i at residue position j.
+      * The names of the targets matched, including the jackhmmer subsequence
+        suffix.
+  """
+  name_to_sequence = collections.OrderedDict()
+  for line in stockholm_string.splitlines():
+    line = line.strip()
+    if not line or line.startswith(('#', '//')):
+      continue
+    name, sequence = line.split()
+    if name not in name_to_sequence:
+      name_to_sequence[name] = ''
+    name_to_sequence[name] += sequence
+
+  msa = []
+  deletion_matrix = []
+
+  query = ''
+  keep_columns = []
+  for seq_index, sequence in enumerate(name_to_sequence.values()):
+    if seq_index == 0:
+      # Gather the columns with gaps from the query
+      query = sequence
+      keep_columns = [i for i, res in enumerate(query) if res != '-']
+
+    # Remove the columns with gaps in the query from all sequences.
+    aligned_sequence = ''.join([sequence[c] for c in keep_columns])
+
+    msa.append(aligned_sequence)
+
+    # Count the number of deletions w.r.t. query.
+    deletion_vec = []
+    deletion_count = 0
+    for seq_res, query_res in zip(sequence, query):
+      if seq_res != '-' or query_res != '-':
+        if query_res == '-':
+          deletion_count += 1
+        else:
+          deletion_vec.append(deletion_count)
+          deletion_count = 0
+    deletion_matrix.append(deletion_vec)
+
+  return msa, deletion_matrix, list(name_to_sequence.keys())
+
+
+def parse_a3m(a3m_string: str) -> Tuple[Sequence[str], DeletionMatrix]:
+  """Parses sequences and deletion matrix from a3m format alignment.
+
+  Args:
+    a3m_string: The string contents of a a3m file. The first sequence in the
+      file should be the query sequence.
+
+  Returns:
+    A tuple of:
+      * A list of sequences that have been aligned to the query. These
+        might contain duplicates.
+      * The deletion matrix for the alignment as a list of lists. The element
+        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)
+  deletion_matrix = []
+  for msa_sequence in sequences:
+    deletion_vec = []
+    deletion_count = 0
+    for j in msa_sequence:
+      if j.islower():
+        deletion_count += 1
+      else:
+        deletion_vec.append(deletion_count)
+        deletion_count = 0
+    deletion_matrix.append(deletion_vec)
+
+  # 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
+
+
+def _convert_sto_seq_to_a3m(
+    query_non_gaps: Sequence[bool], sto_seq: str) -> Iterable[str]:
+  for is_query_res_non_gap, sequence_res in zip(query_non_gaps, sto_seq):
+    if is_query_res_non_gap:
+      yield sequence_res
+    elif sequence_res != '-':
+      yield sequence_res.lower()
+
+
+def convert_stockholm_to_a3m(stockholm_format: str,
+                             max_sequences: Optional[int] = None) -> str:
+  """Converts MSA in Stockholm format to the A3M format."""
+  descriptions = {}
+  sequences = {}
+  reached_max_sequences = False
+
+  for line in stockholm_format.splitlines():
+    reached_max_sequences = max_sequences and len(sequences) >= max_sequences
+    if line.strip() and not line.startswith(('#', '//')):
+      # Ignore blank lines, markup and end symbols - remainder are alignment
+      # sequence parts.
+      seqname, aligned_seq = line.split(maxsplit=1)
+      if seqname not in sequences:
+        if reached_max_sequences:
+          continue
+        sequences[seqname] = ''
+      sequences[seqname] += aligned_seq
+
+  for line in stockholm_format.splitlines():
+    if line[:4] == '#=GS':
+      # Description row - example format is:
+      # #=GS UniRef90_Q9H5Z4/4-78            DE [subseq from] cDNA: FLJ22755 ...
+      columns = line.split(maxsplit=3)
+      seqname, feature = columns[1:3]
+      value = columns[3] if len(columns) == 4 else ''
+      if feature != 'DE':
+        continue
+      if reached_max_sequences and seqname not in sequences:
+        continue
+      descriptions[seqname] = value
+      if len(descriptions) == len(sequences):
+        break
+
+  # 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]
+  for seqname, sto_sequence in sequences.items():
+    a3m_sequences[seqname] = ''.join(
+        _convert_sto_seq_to_a3m(query_non_gaps, sto_sequence))
+
+  fasta_chunks = (f">{k} {descriptions.get(k, '')}\n{a3m_sequences[k]}"
+                  for k in a3m_sequences)
+  return '\n'.join(fasta_chunks) + '\n'  # Include terminating newline.
+
+
+def _get_hhr_line_regex_groups(
+    regex_pattern: str, line: str) -> Sequence[Optional[str]]:
+  match = re.match(regex_pattern, line)
+  if match is None:
+    raise RuntimeError(f'Could not parse query line {line}')
+  return match.groups()
+
+
+def _update_hhr_residue_indices_list(
+    sequence: str, start_index: int, indices_list: List[int]):
+  """Computes the relative indices for each residue with respect to the original sequence."""
+  counter = start_index
+  for symbol in sequence:
+    if symbol == '-':
+      indices_list.append(-1)
+    else:
+      indices_list.append(counter)
+      counter += 1
+
+
+def _parse_hhr_hit(detailed_lines: Sequence[str]) -> TemplateHit:
+  """Parses the detailed HMM HMM comparison section for a single Hit.
+
+  This works on .hhr files generated from both HHBlits and HHSearch.
+
+  Args:
+    detailed_lines: A list of lines from a single comparison section between 2
+      sequences (which each have their own HMM's)
+
+  Returns:
+    A dictionary with the information from that detailed comparison section
+
+  Raises:
+    RuntimeError: If a certain line cannot be processed
+  """
+  # Parse first 2 lines.
+  number_of_hit = int(detailed_lines[0].split()[-1])
+  name_hit = detailed_lines[1][1:]
+
+  # Parse the summary line.
+  pattern = (
+      'Probab=(.*)[\t ]*E-value=(.*)[\t ]*Score=(.*)[\t ]*Aligned_cols=(.*)[\t'
+      ' ]*Identities=(.*)%[\t ]*Similarity=(.*)[\t ]*Sum_probs=(.*)[\t '
+      ']*Template_Neff=(.*)')
+  match = re.match(pattern, detailed_lines[2])
+  if match is None:
+    raise RuntimeError(
+        '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) = [float(x) for x in match.groups()]
+
+  # The next section reads the detailed comparisons. These are in a 'human
+  # readable' format which has a fixed length. The strategy employed is to
+  # assume that each block starts with the query sequence line, and to parse
+  # that with a regexp in order to deduce the fixed length used for that block.
+  query = ''
+  hit_sequence = ''
+  indices_query = []
+  indices_hit = []
+  length_block = None
+
+  for line in detailed_lines[3:]:
+    # Parse the query sequence line
+    if (line.startswith('Q ') and not line.startswith('Q ss_dssp') and
+        not line.startswith('Q ss_pred') and
+        not line.startswith('Q Consensus')):
+      # Thus the first 17 characters must be 'Q <query_name> ', and we can parse
+      # everything after that.
+      #              start    sequence       end       total_sequence_length
+      patt = r'[\t ]*([0-9]*) ([A-Z-]*)[\t ]*([0-9]*) \([0-9]*\)'
+      groups = _get_hhr_line_regex_groups(patt, line[17:])
+
+      # Get the length of the parsed block using the start and finish indices,
+      # and ensure it is the same as the actual block length.
+      start = int(groups[0]) - 1  # Make index zero based.
+      delta_query = groups[1]
+      end = int(groups[2])
+      num_insertions = len([x for x in delta_query if x == '-'])
+      length_block = end - start + num_insertions
+      assert length_block == len(delta_query)
+
+      # Update the query sequence and indices list.
+      query += delta_query
+      _update_hhr_residue_indices_list(delta_query, start, indices_query)
+
+    elif line.startswith('T '):
+      # Parse the hit sequence.
+      if (not line.startswith('T ss_dssp') and
+          not line.startswith('T ss_pred') and
+          not line.startswith('T Consensus')):
+        # Thus the first 17 characters must be 'T <hit_name> ', and we can
+        # parse everything after that.
+        #              start    sequence       end     total_sequence_length
+        patt = r'[\t ]*([0-9]*) ([A-Z-]*)[\t ]*[0-9]* \([0-9]*\)'
+        groups = _get_hhr_line_regex_groups(patt, line[17:])
+        start = int(groups[0]) - 1  # Make index zero based.
+        delta_hit_sequence = groups[1]
+        assert length_block == len(delta_hit_sequence)
+
+        # Update the hit sequence and indices list.
+        hit_sequence += delta_hit_sequence
+        _update_hhr_residue_indices_list(delta_hit_sequence, start, indices_hit)
+
+  return TemplateHit(
+      index=number_of_hit,
+      name=name_hit,
+      aligned_cols=int(aligned_cols),
+      sum_probs=sum_probs,
+      query=query,
+      hit_sequence=hit_sequence,
+      indices_query=indices_query,
+      indices_hit=indices_hit,
+  )
+
+
+def parse_hhr(hhr_string: str) -> Sequence[TemplateHit]:
+  """Parses the content of an entire HHR file."""
+  lines = hhr_string.splitlines()
+
+  # Each .hhr file starts with a results table, then has a sequence of hit
+  # "paragraphs", each paragraph starting with a line 'No <hit number>'. We
+  # iterate through each paragraph to parse each hit.
+
+  block_starts = [i for i, line in enumerate(lines) if line.startswith('No ')]
+
+  hits = []
+  if block_starts:
+    block_starts.append(len(lines))  # Add the end of the final block.
+    for i in range(len(block_starts) - 1):
+      hits.append(_parse_hhr_hit(lines[block_starts[i]:block_starts[i + 1]]))
+  return hits
+
+
+def parse_e_values_from_tblout(tblout: str) -> Dict[str, float]:
+  """Parse target to e-value mapping parsed from Jackhmmer tblout string."""
+  e_values = {'query': 0}
+  lines = [line for line in tblout.splitlines() if line[0] != '#']
+  # As per http://eddylab.org/software/hmmer/Userguide.pdf fields are
+  # space-delimited. Relevant fields are (1) target name:  and
+  # (5) E-value (full sequence) (numbering from 1).
+  for line in lines:
+    fields = line.split()
+    e_value = fields[4]
+    target_name = fields[0]
+    e_values[target_name] = float(e_value)
+  return e_values

alphafold/alphafold/data/pipeline.py

@@ -0,0 +1,209 @@
+# 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.
+
+"""Functions for building the input features for the AlphaFold model."""
+
+import os
+from typing import Mapping, Optional, Sequence
+from absl import logging
+from alphafold.common import residue_constants
+from alphafold.data import parsers
+from alphafold.data import templates
+from alphafold.data.tools import hhblits
+from alphafold.data.tools import hhsearch
+from alphafold.data.tools import jackhmmer
+import numpy as np
+
+# Internal import (7716).
+
+FeatureDict = Mapping[str, np.ndarray]
+
+
+def make_sequence_features(
+    sequence: str, description: str, num_res: int) -> FeatureDict:
+  """Constructs a feature dict of sequence features."""
+  features = {}
+  features['aatype'] = residue_constants.sequence_to_onehot(
+      sequence=sequence,
+      mapping=residue_constants.restype_order_with_x,
+      map_unknown_to_x=True)
+  features['between_segment_residues'] = np.zeros((num_res,), dtype=np.int32)
+  features['domain_name'] = np.array([description.encode('utf-8')],
+                                     dtype=np.object_)
+  features['residue_index'] = np.array(range(num_res), dtype=np.int32)
+  features['seq_length'] = np.array([num_res] * num_res, dtype=np.int32)
+  features['sequence'] = np.array([sequence.encode('utf-8')], dtype=np.object_)
+  return features
+
+
+def make_msa_features(
+    msas: Sequence[Sequence[str]],
+    deletion_matrices: Sequence[parsers.DeletionMatrix]) -> FeatureDict:
+  """Constructs a feature dict of MSA features."""
+  if not msas:
+    raise ValueError('At least one MSA must be provided.')
+
+  int_msa = []
+  deletion_matrix = []
+  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):
+      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])
+  num_alignments = len(int_msa)
+  features = {}
+  features['deletion_matrix_int'] = np.array(deletion_matrix, dtype=np.int32)
+  features['msa'] = np.array(int_msa, dtype=np.int32)
+  features['num_alignments'] = np.array(
+      [num_alignments] * num_res, dtype=np.int32)
+  return features
+
+
+class DataPipeline:
+  """Runs the alignment tools and assembles the input features."""
+
+  def __init__(self,
+               jackhmmer_binary_path: str,
+               hhblits_binary_path: str,
+               hhsearch_binary_path: str,
+               uniref90_database_path: str,
+               mgnify_database_path: str,
+               bfd_database_path: Optional[str],
+               uniclust30_database_path: Optional[str],
+               small_bfd_database_path: Optional[str],
+               pdb70_database_path: str,
+               template_featurizer: templates.TemplateHitFeaturizer,
+               use_small_bfd: bool,
+               mgnify_max_hits: int = 501,
+               uniref_max_hits: int = 10000):
+    """Constructs a feature dict for a given FASTA file."""
+    self._use_small_bfd = use_small_bfd
+    self.jackhmmer_uniref90_runner = jackhmmer.Jackhmmer(
+        binary_path=jackhmmer_binary_path,
+        database_path=uniref90_database_path)
+    if use_small_bfd:
+      self.jackhmmer_small_bfd_runner = jackhmmer.Jackhmmer(
+          binary_path=jackhmmer_binary_path,
+          database_path=small_bfd_database_path)
+    else:
+      self.hhblits_bfd_uniclust_runner = hhblits.HHBlits(
+          binary_path=hhblits_binary_path,
+          databases=[bfd_database_path, uniclust30_database_path])
+    self.jackhmmer_mgnify_runner = jackhmmer.Jackhmmer(
+        binary_path=jackhmmer_binary_path,
+        database_path=mgnify_database_path)
+    self.hhsearch_pdb70_runner = hhsearch.HHSearch(
+        binary_path=hhsearch_binary_path,
+        databases=[pdb70_database_path])
+    self.template_featurizer = template_featurizer
+    self.mgnify_max_hits = mgnify_max_hits
+    self.uniref_max_hits = uniref_max_hits
+
+  def process(self, input_fasta_path: str, msa_output_dir: str) -> FeatureDict:
+    """Runs alignment tools on the input sequence and creates features."""
+    with open(input_fasta_path) as f:
+      input_fasta_str = f.read()
+    input_seqs, input_descs = parsers.parse_fasta(input_fasta_str)
+    if len(input_seqs) != 1:
+      raise ValueError(
+          f'More than one input sequence found in {input_fasta_path}.')
+    input_sequence = input_seqs[0]
+    input_description = input_descs[0]
+    num_res = len(input_sequence)
+
+    jackhmmer_uniref90_result = self.jackhmmer_uniref90_runner.query(
+        input_fasta_path)[0]
+    jackhmmer_mgnify_result = self.jackhmmer_mgnify_runner.query(
+        input_fasta_path)[0]
+
+    uniref90_msa_as_a3m = parsers.convert_stockholm_to_a3m(
+        jackhmmer_uniref90_result['sto'], max_sequences=self.uniref_max_hits)
+    hhsearch_result = self.hhsearch_pdb70_runner.query(uniref90_msa_as_a3m)
+
+    uniref90_out_path = os.path.join(msa_output_dir, 'uniref90_hits.sto')
+    with open(uniref90_out_path, 'w') as f:
+      f.write(jackhmmer_uniref90_result['sto'])
+
+    mgnify_out_path = os.path.join(msa_output_dir, 'mgnify_hits.sto')
+    with open(mgnify_out_path, 'w') as f:
+      f.write(jackhmmer_mgnify_result['sto'])
+
+    pdb70_out_path = os.path.join(msa_output_dir, 'pdb70_hits.hhr')
+    with open(pdb70_out_path, 'w') as f:
+      f.write(hhsearch_result)
+
+    uniref90_msa, uniref90_deletion_matrix, _ = parsers.parse_stockholm(
+        jackhmmer_uniref90_result['sto'])
+    mgnify_msa, mgnify_deletion_matrix, _ = parsers.parse_stockholm(
+        jackhmmer_mgnify_result['sto'])
+    hhsearch_hits = parsers.parse_hhr(hhsearch_result)
+    mgnify_msa = mgnify_msa[:self.mgnify_max_hits]
+    mgnify_deletion_matrix = mgnify_deletion_matrix[:self.mgnify_max_hits]
+
+    if self._use_small_bfd:
+      jackhmmer_small_bfd_result = self.jackhmmer_small_bfd_runner.query(
+          input_fasta_path)[0]
+
+      bfd_out_path = os.path.join(msa_output_dir, 'small_bfd_hits.a3m')
+      with open(bfd_out_path, 'w') as f:
+        f.write(jackhmmer_small_bfd_result['sto'])
+
+      bfd_msa, bfd_deletion_matrix, _ = parsers.parse_stockholm(
+          jackhmmer_small_bfd_result['sto'])
+    else:
+      hhblits_bfd_uniclust_result = self.hhblits_bfd_uniclust_runner.query(
+          input_fasta_path)
+
+      bfd_out_path = os.path.join(msa_output_dir, 'bfd_uniclust_hits.a3m')
+      with open(bfd_out_path, 'w') as f:
+        f.write(hhblits_bfd_uniclust_result['a3m'])
+
+      bfd_msa, bfd_deletion_matrix = parsers.parse_a3m(
+          hhblits_bfd_uniclust_result['a3m'])
+
+    templates_result = self.template_featurizer.get_templates(
+        query_sequence=input_sequence,
+        query_pdb_code=None,
+        query_release_date=None,
+        hits=hhsearch_hits)
+
+    sequence_features = make_sequence_features(
+        sequence=input_sequence,
+        description=input_description,
+        num_res=num_res)
+
+    msa_features = make_msa_features(
+        msas=(uniref90_msa, bfd_msa, mgnify_msa),
+        deletion_matrices=(uniref90_deletion_matrix,
+                           bfd_deletion_matrix,
+                           mgnify_deletion_matrix))
+
+    logging.info('Uniref90 MSA size: %d sequences.', len(uniref90_msa))
+    logging.info('BFD MSA size: %d sequences.', len(bfd_msa))
+    logging.info('MGnify MSA size: %d sequences.', len(mgnify_msa))
+    logging.info('Final (deduplicated) MSA size: %d sequences.',
+                 msa_features['num_alignments'][0])
+    logging.info('Total number of templates (NB: this can include bad '
+                 'templates and is later filtered to top 4): %d.',
+                 templates_result.features['template_domain_names'].shape[0])
+
+    return {**sequence_features, **msa_features, **templates_result.features}

alphafold/alphafold/data/templates.py

@@ -0,0 +1,922 @@
+# 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.
+
+"""Functions for getting templates and calculating template features."""
+import dataclasses
+import datetime
+import glob
+import os
+import re
+from typing import Any, Dict, Mapping, Optional, Sequence, Tuple
+
+from absl import logging
+from alphafold.common import residue_constants
+from alphafold.data import mmcif_parsing
+from alphafold.data import parsers
+from alphafold.data.tools import kalign
+import numpy as np
+
+# Internal import (7716).
+
+
+class Error(Exception):
+  """Base class for exceptions."""
+
+
+class NoChainsError(Error):
+  """An error indicating that template mmCIF didn't have any chains."""
+
+
+class SequenceNotInTemplateError(Error):
+  """An error indicating that template mmCIF didn't contain the sequence."""
+
+
+class NoAtomDataInTemplateError(Error):
+  """An error indicating that template mmCIF didn't contain atom positions."""
+
+
+class TemplateAtomMaskAllZerosError(Error):
+  """An error indicating that template mmCIF had all atom positions masked."""
+
+
+class QueryToTemplateAlignError(Error):
+  """An error indicating that the query can't be aligned to the template."""
+
+
+class CaDistanceError(Error):
+  """An error indicating that a CA atom distance exceeds a threshold."""
+
+
+class MultipleChainsError(Error):
+  """An error indicating that multiple chains were found for a given ID."""
+
+
+# Prefilter exceptions.
+class PrefilterError(Exception):
+  """A base class for template prefilter exceptions."""
+
+
+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."""
+
+
+class DuplicateError(PrefilterError):
+  """An error indicating that the hit was an exact subsequence of the query."""
+
+
+class LengthError(PrefilterError):
+  """An error indicating that the hit was too short."""
+
+
+TEMPLATE_FEATURES = {
+    'template_aatype': np.float32,
+    'template_all_atom_masks': np.float32,
+    'template_all_atom_positions': np.float32,
+    'template_domain_names': np.object,
+    'template_sequence': np.object,
+    'template_sum_probs': 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.
+  id_match = re.match(r'[a-zA-Z\d]{4}_[a-zA-Z0-9.]+', hit.name)
+  if not id_match:
+    raise ValueError(f'hit.name did not start with PDBID_chain: {hit.name}')
+  pdb_id, chain_id = id_match.group(0).split('_')
+  return pdb_id.lower(), chain_id
+
+
+def _is_after_cutoff(
+    pdb_id: str,
+    release_dates: Mapping[str, datetime.datetime],
+    release_date_cutoff: Optional[datetime.datetime]) -> bool:
+  """Checks if the template date is after the release date cutoff.
+
+  Args:
+    pdb_id: 4 letter pdb code.
+    release_dates: Dictionary mapping PDB ids to their structure release dates.
+    release_date_cutoff: Max release date that is valid for this query.
+
+  Returns:
+    True if the template release date is after the cutoff, False otherwise.
+  """
+  if release_date_cutoff is None:
+    raise ValueError('The release_date_cutoff must not be None.')
+  if pdb_id in release_dates:
+    return release_dates[pdb_id] > release_date_cutoff
+  else:
+    # Since this is just a quick prefilter to reduce the number of mmCIF files
+    # we need to parse, we don't have to worry about returning True here.
+    logging.warning('Template structure not in release dates dict: %s', pdb_id)
+    return False
+
+
+def _parse_obsolete(obsolete_file_path: str) -> Mapping[str, Optional[str]]:
+  """Parses the data file from PDB that lists which pdb_ids are obsolete."""
+  with open(obsolete_file_path) as f:
+    result = {}
+    for line in f:
+      line = line.strip()
+      # Format:    Date      From     To
+      # 'OBSLTE    06-NOV-19 6G9Y'                - Removed, rare
+      # 'OBSLTE    31-JUL-94 116L     216L'       - Replaced, common
+      # 'OBSLTE    26-SEP-06 2H33     2JM5 2OWI'  - Replaced by multiple, rare
+      if line.startswith('OBSLTE'):
+        if len(line) > 30:
+          # Replaced by at least one structure.
+          from_id = line[20:24].lower()
+          to_id = line[29:33].lower()
+          result[from_id] = to_id
+        elif len(line) == 24:
+          # Removed.
+          from_id = line[20:24].lower()
+          result[from_id] = None
+    return result
+
+
+def _parse_release_dates(path: str) -> Mapping[str, datetime.datetime]:
+  """Parses release dates file, returns a mapping from PDBs to release dates."""
+  if path.endswith('txt'):
+    release_dates = {}
+    with open(path, 'r') as f:
+      for line in f:
+        pdb_id, date = line.split(':')
+        date = date.strip()
+        # Python 3.6 doesn't have datetime.date.fromisoformat() which is about
+        # 90x faster than strptime. However, splitting the string manually is
+        # about 10x faster than strptime.
+        release_dates[pdb_id.strip()] = datetime.datetime(
+            year=int(date[:4]), month=int(date[5:7]), day=int(date[8:10]))
+    return release_dates
+  else:
+    raise ValueError('Invalid format of the release date file %s.' % path)
+
+
+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,
+    min_align_ratio: float = 0.1) -> bool:
+  """Determines if template is valid (without parsing the template mmcif file).
+
+  Args:
+    hit: HhrHit for the template.
+    hit_pdb_code: The 4 letter pdb code of the template hit. This might be
+      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.
+    max_subsequence_ratio: Exclude any exact matches with this much overlap.
+    min_align_ratio: Minimum overlap between the template and query.
+
+  Returns:
+    True if the hit passed the prefilter. Raises an exception otherwise.
+
+  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.
+  """
+  aligned_cols = hit.aligned_cols
+  align_ratio = aligned_cols / len(query_sequence)
+
+  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):
+    raise DateError(f'Date ({release_dates[hit_pdb_code]}) > max template date '
+                    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}.')
+
+  if duplicate:
+    raise DuplicateError('Template is an exact subsequence of query with large '
+                         f'coverage. Length ratio: {length_ratio}.')
+
+  if len(template_sequence) < 10:
+    raise LengthError(f'Template too short. Length: {len(template_sequence)}.')
+
+  return True
+
+
+def _find_template_in_pdb(
+    template_chain_id: str,
+    template_sequence: str,
+    mmcif_object: mmcif_parsing.MmcifObject) -> Tuple[str, str, int]:
+  """Tries to find the template chain in the given pdb file.
+
+  This method tries the three following things in order:
+    1. Tries if there is an exact match in both the chain ID and the sequence.
+       If yes, the chain sequence is returned. Otherwise:
+    2. Tries if there is an exact match only in the sequence.
+       If yes, the chain sequence is returned. Otherwise:
+    3. Tries if there is a fuzzy match (X = wildcard) in the sequence.
+       If yes, the chain sequence is returned.
+  If none of these succeed, a SequenceNotInTemplateError is thrown.
+
+  Args:
+    template_chain_id: The template chain ID.
+    template_sequence: The template chain sequence.
+    mmcif_object: The PDB object to search for the template in.
+
+  Returns:
+    A tuple with:
+    * The chain sequence that was found to match the template in the PDB object.
+    * The ID of the chain that is being returned.
+    * The offset where the template sequence starts in the chain sequence.
+
+  Raises:
+    SequenceNotInTemplateError: If no match is found after the steps described
+      above.
+  """
+  # Try if there is an exact match in both the chain ID and the (sub)sequence.
+  pdb_id = mmcif_object.file_id
+  chain_sequence = mmcif_object.chain_to_seqres.get(template_chain_id)
+  if chain_sequence and (template_sequence in chain_sequence):
+    logging.info(
+        'Found an exact template match %s_%s.', pdb_id, template_chain_id)
+    mapping_offset = chain_sequence.find(template_sequence)
+    return chain_sequence, template_chain_id, mapping_offset
+
+  # Try if there is an exact match in the (sub)sequence only.
+  for chain_id, chain_sequence in mmcif_object.chain_to_seqres.items():
+    if chain_sequence and (template_sequence in chain_sequence):
+      logging.info('Found a sequence-only match %s_%s.', pdb_id, chain_id)
+      mapping_offset = chain_sequence.find(template_sequence)
+      return chain_sequence, chain_id, mapping_offset
+
+  # Return a chain sequence that fuzzy matches (X = wildcard) the template.
+  # Make parentheses unnamed groups (?:_) to avoid the 100 named groups limit.
+  regex = ['.' if aa == 'X' else '(?:%s|X)' % aa for aa in template_sequence]
+  regex = re.compile(''.join(regex))
+  for chain_id, chain_sequence in mmcif_object.chain_to_seqres.items():
+    match = re.search(regex, chain_sequence)
+    if match:
+      logging.info('Found a fuzzy sequence-only match %s_%s.', pdb_id, chain_id)
+      mapping_offset = match.start()
+      return chain_sequence, chain_id, mapping_offset
+
+  # No hits, raise an error.
+  raise SequenceNotInTemplateError(
+      'Could not find the template sequence in %s_%s. Template sequence: %s, '
+      'chain_to_seqres: %s' % (pdb_id, template_chain_id, template_sequence,
+                               mmcif_object.chain_to_seqres))
+
+
+def _realign_pdb_template_to_query(
+    old_template_sequence: str,
+    template_chain_id: str,
+    mmcif_object: mmcif_parsing.MmcifObject,
+    old_mapping: Mapping[int, int],
+    kalign_binary_path: str) -> Tuple[str, Mapping[int, int]]:
+  """Aligns template from the mmcif_object to the query.
+
+  In case PDB70 contains a different version of the template sequence, we need
+  to perform a realignment to the actual sequence that is in the mmCIF file.
+  This method performs such realignment, but returns the new sequence and
+  mapping only if the sequence in the mmCIF file is 90% identical to the old
+  sequence.
+
+  Note that the old_template_sequence comes from the hit, and contains only that
+  part of the chain that matches with the query while the new_template_sequence
+  is the full chain.
+
+  Args:
+    old_template_sequence: The template sequence that was returned by the PDB
+      template search (typically done using HHSearch).
+    template_chain_id: The template chain id was returned by the PDB template
+      search (typically done using HHSearch). This is used to find the right
+      chain in the mmcif_object chain_to_seqres mapping.
+    mmcif_object: A mmcif_object which holds the actual template data.
+    old_mapping: A mapping from the query sequence to the template sequence.
+      This mapping will be used to compute the new mapping from the query
+      sequence to the actual mmcif_object template sequence by aligning the
+      old_template_sequence and the actual template sequence.
+    kalign_binary_path: The path to a kalign executable.
+
+  Returns:
+    A tuple (new_template_sequence, new_query_to_template_mapping) where:
+    * new_template_sequence is the actual template sequence that was found in
+      the mmcif_object.
+    * new_query_to_template_mapping is the new mapping from the query to the
+      actual template found in the mmcif_object.
+
+  Raises:
+    QueryToTemplateAlignError:
+    * If there was an error thrown by the alignment tool.
+    * Or if the actual template sequence differs by more than 10% from the
+      old_template_sequence.
+  """
+  aligner = kalign.Kalign(binary_path=kalign_binary_path)
+  new_template_sequence = mmcif_object.chain_to_seqres.get(
+      template_chain_id, '')
+
+  # Sometimes the template chain id is unknown. But if there is only a single
+  # sequence within the mmcif_object, it is safe to assume it is that one.
+  if not new_template_sequence:
+    if len(mmcif_object.chain_to_seqres) == 1:
+      logging.info('Could not find %s in %s, but there is only 1 sequence, so '
+                   'using that one.',
+                   template_chain_id,
+                   mmcif_object.file_id)
+      new_template_sequence = list(mmcif_object.chain_to_seqres.values())[0]
+    else:
+      raise QueryToTemplateAlignError(
+          f'Could not find chain {template_chain_id} in {mmcif_object.file_id}. '
+          'If there are no mmCIF parsing errors, it is possible it was not a '
+          'protein chain.')
+
+  try:
+    (old_aligned_template, new_aligned_template), _ = parsers.parse_a3m(
+        aligner.align([old_template_sequence, new_template_sequence]))
+  except Exception as e:
+    raise QueryToTemplateAlignError(
+        'Could not align old template %s to template %s (%s_%s). Error: %s' %
+        (old_template_sequence, new_template_sequence, mmcif_object.file_id,
+         template_chain_id, str(e)))
+
+  logging.info('Old aligned template: %s\nNew aligned template: %s',
+               old_aligned_template, new_aligned_template)
+
+  old_to_new_template_mapping = {}
+  old_template_index = -1
+  new_template_index = -1
+  num_same = 0
+  for old_template_aa, new_template_aa in zip(
+      old_aligned_template, new_aligned_template):
+    if old_template_aa != '-':
+      old_template_index += 1
+    if new_template_aa != '-':
+      new_template_index += 1
+    if old_template_aa != '-' and new_template_aa != '-':
+      old_to_new_template_mapping[old_template_index] = new_template_index
+      if old_template_aa == new_template_aa:
+        num_same += 1
+
+  # Require at least 90 % sequence identity wrt to the shorter of the sequences.
+  if float(num_same) / min(
+      len(old_template_sequence), len(new_template_sequence)) < 0.9:
+    raise QueryToTemplateAlignError(
+        'Insufficient similarity of the sequence in the database: %s to the '
+        'actual sequence in the mmCIF file %s_%s: %s. We require at least '
+        '90 %% similarity wrt to the shorter of the sequences. This is not a '
+        'problem unless you think this is a template that should be included.' %
+        (old_template_sequence, mmcif_object.file_id, template_chain_id,
+         new_template_sequence))
+
+  new_query_to_template_mapping = {}
+  for query_index, old_template_index in old_mapping.items():
+    new_query_to_template_mapping[query_index] = (
+        old_to_new_template_mapping.get(old_template_index, -1))
+
+  new_template_sequence = new_template_sequence.replace('-', '')
+
+  return new_template_sequence, new_query_to_template_mapping
+
+
+def _check_residue_distances(all_positions: np.ndarray,
+                             all_positions_mask: np.ndarray,
+                             max_ca_ca_distance: float):
+  """Checks if the distance between unmasked neighbor residues is ok."""
+  ca_position = residue_constants.atom_order['CA']
+  prev_is_unmasked = False
+  prev_calpha = None
+  for i, (coords, mask) in enumerate(zip(all_positions, all_positions_mask)):
+    this_is_unmasked = bool(mask[ca_position])
+    if this_is_unmasked:
+      this_calpha = coords[ca_position]
+      if prev_is_unmasked:
+        distance = np.linalg.norm(this_calpha - prev_calpha)
+        if distance > max_ca_ca_distance:
+          raise CaDistanceError(
+              'The distance between residues %d and %d is %f > limit %f.' % (
+                  i, i + 1, distance, max_ca_ca_distance))
+      prev_calpha = this_calpha
+    prev_is_unmasked = this_is_unmasked
+
+
+def _get_atom_positions(
+    mmcif_object: mmcif_parsing.MmcifObject,
+    auth_chain_id: str,
+    max_ca_ca_distance: float) -> Tuple[np.ndarray, np.ndarray]:
+  """Gets atom positions and mask from a list of Biopython Residues."""
+  num_res = len(mmcif_object.chain_to_seqres[auth_chain_id])
+
+  relevant_chains = [c for c in mmcif_object.structure.get_chains()
+                     if c.id == auth_chain_id]
+  if len(relevant_chains) != 1:
+    raise MultipleChainsError(
+        f'Expected exactly one chain in structure with id {auth_chain_id}.')
+  chain = relevant_chains[0]
+
+  all_positions = np.zeros([num_res, residue_constants.atom_type_num, 3])
+  all_positions_mask = np.zeros([num_res, residue_constants.atom_type_num],
+                                dtype=np.int64)
+  for res_index in range(num_res):
+    pos = np.zeros([residue_constants.atom_type_num, 3], dtype=np.float32)
+    mask = np.zeros([residue_constants.atom_type_num], dtype=np.float32)
+    res_at_position = mmcif_object.seqres_to_structure[auth_chain_id][res_index]
+    if not res_at_position.is_missing:
+      res = chain[(res_at_position.hetflag,
+                   res_at_position.position.residue_number,
+                   res_at_position.position.insertion_code)]
+      for atom in res.get_atoms():
+        atom_name = atom.get_name()
+        x, y, z = atom.get_coord()
+        if atom_name in residue_constants.atom_order.keys():
+          pos[residue_constants.atom_order[atom_name]] = [x, y, z]
+          mask[residue_constants.atom_order[atom_name]] = 1.0
+        elif atom_name.upper() == 'SE' and res.get_resname() == 'MSE':
+          # Put the coordinates of the selenium atom in the sulphur column.
+          pos[residue_constants.atom_order['SD']] = [x, y, z]
+          mask[residue_constants.atom_order['SD']] = 1.0
+
+    all_positions[res_index] = pos
+    all_positions_mask[res_index] = mask
+  _check_residue_distances(
+      all_positions, all_positions_mask, max_ca_ca_distance)
+  return all_positions, all_positions_mask
+
+
+def _extract_template_features(
+    mmcif_object: mmcif_parsing.MmcifObject,
+    pdb_id: str,
+    mapping: Mapping[int, int],
+    template_sequence: str,
+    query_sequence: str,
+    template_chain_id: str,
+    kalign_binary_path: str) -> Tuple[Dict[str, Any], Optional[str]]:
+  """Parses atom positions in the target structure and aligns with the query.
+
+  Atoms for each residue in the template structure are indexed to coincide
+  with their corresponding residue in the query sequence, according to the
+  alignment mapping provided.
+
+  Args:
+    mmcif_object: mmcif_parsing.MmcifObject representing the template.
+    pdb_id: PDB code for the template.
+    mapping: Dictionary mapping indices in the query sequence to indices in
+      the template sequence.
+    template_sequence: String describing the amino acid sequence for the
+      template protein.
+    query_sequence: String describing the amino acid sequence for the query
+      protein.
+    template_chain_id: String ID describing which chain in the structure proto
+      should be used.
+    kalign_binary_path: The path to a kalign executable used for template
+        realignment.
+
+  Returns:
+    A tuple with:
+    * A dictionary containing the extra features derived from the template
+      protein structure.
+    * A warning message if the hit was realigned to the actual mmCIF sequence.
+      Otherwise None.
+
+  Raises:
+    NoChainsError: If the mmcif object doesn't contain any chains.
+    SequenceNotInTemplateError: If the given chain id / sequence can't
+      be found in the mmcif object.
+    QueryToTemplateAlignError: If the actual template in the mmCIF file
+      can't be aligned to the query.
+    NoAtomDataInTemplateError: If the mmcif object doesn't contain
+      atom positions.
+    TemplateAtomMaskAllZerosError: If the mmcif object doesn't have any
+      unmasked residues.
+  """
+  if mmcif_object is None or not mmcif_object.chain_to_seqres:
+    raise NoChainsError('No chains in PDB: %s_%s' % (pdb_id, template_chain_id))
+
+  warning = None
+  try:
+    seqres, chain_id, mapping_offset = _find_template_in_pdb(
+        template_chain_id=template_chain_id,
+        template_sequence=template_sequence,
+        mmcif_object=mmcif_object)
+  except SequenceNotInTemplateError:
+    # If PDB70 contains a different version of the template, we use the sequence
+    # from the mmcif_object.
+    chain_id = template_chain_id
+    warning = (
+        f'The exact sequence {template_sequence} was not found in '
+        f'{pdb_id}_{chain_id}. Realigning the template to the actual sequence.')
+    logging.warning(warning)
+    # This throws an exception if it fails to realign the hit.
+    seqres, mapping = _realign_pdb_template_to_query(
+        old_template_sequence=template_sequence,
+        template_chain_id=template_chain_id,
+        mmcif_object=mmcif_object,
+        old_mapping=mapping,
+        kalign_binary_path=kalign_binary_path)
+    logging.info('Sequence in %s_%s: %s successfully realigned to %s',
+                 pdb_id, chain_id, template_sequence, seqres)
+    # The template sequence changed.
+    template_sequence = seqres
+    # No mapping offset, the query is aligned to the actual sequence.
+    mapping_offset = 0
+
+  try:
+    # Essentially set to infinity - we don't want to reject templates unless
+    # they're really really bad.
+    all_atom_positions, all_atom_mask = _get_atom_positions(
+        mmcif_object, chain_id, max_ca_ca_distance=150.0)
+  except (CaDistanceError, KeyError) as ex:
+    raise NoAtomDataInTemplateError(
+        'Could not get atom data (%s_%s): %s' % (pdb_id, chain_id, str(ex))
+        ) from ex
+
+  all_atom_positions = np.split(all_atom_positions, all_atom_positions.shape[0])
+  all_atom_masks = np.split(all_atom_mask, all_atom_mask.shape[0])
+
+  output_templates_sequence = []
+  templates_all_atom_positions = []
+  templates_all_atom_masks = []
+
+  for _ in query_sequence:
+    # Residues in the query_sequence that are not in the template_sequence:
+    templates_all_atom_positions.append(
+        np.zeros((residue_constants.atom_type_num, 3)))
+    templates_all_atom_masks.append(np.zeros(residue_constants.atom_type_num))
+    output_templates_sequence.append('-')
+
+  for k, v in mapping.items():
+    template_index = v + mapping_offset
+    templates_all_atom_positions[k] = all_atom_positions[template_index][0]
+    templates_all_atom_masks[k] = all_atom_masks[template_index][0]
+    output_templates_sequence[k] = template_sequence[v]
+
+  # Alanine (AA with the lowest number of atoms) has 5 atoms (C, CA, CB, N, O).
+  if np.sum(templates_all_atom_masks) < 5:
+    raise TemplateAtomMaskAllZerosError(
+        'Template all atom mask was all zeros: %s_%s. Residue range: %d-%d' %
+        (pdb_id, chain_id, min(mapping.values()) + mapping_offset,
+         max(mapping.values()) + mapping_offset))
+
+  output_templates_sequence = ''.join(output_templates_sequence)
+
+  templates_aatype = residue_constants.sequence_to_onehot(
+      output_templates_sequence, residue_constants.HHBLITS_AA_TO_ID)
+
+  return (
+      {
+          'template_all_atom_positions': np.array(templates_all_atom_positions),
+          'template_all_atom_masks': np.array(templates_all_atom_masks),
+          'template_sequence': output_templates_sequence.encode(),
+          'template_aatype': np.array(templates_aatype),
+          'template_domain_names': f'{pdb_id.lower()}_{chain_id}'.encode(),
+      },
+      warning)
+
+
+def _build_query_to_hit_index_mapping(
+    hit_query_sequence: str,
+    hit_sequence: str,
+    indices_hit: Sequence[int],
+    indices_query: Sequence[int],
+    original_query_sequence: str) -> Mapping[int, int]:
+  """Gets mapping from indices in original query sequence to indices in the hit.
+
+  hit_query_sequence and hit_sequence are two aligned sequences containing gap
+  characters. hit_query_sequence contains only the part of the original query
+  sequence that matched the hit. When interpreting the indices from the .hhr, we
+  need to correct for this to recover a mapping from original query sequence to
+  the hit sequence.
+
+  Args:
+    hit_query_sequence: The portion of the query sequence that is in the .hhr
+      hit
+    hit_sequence: The portion of the hit sequence that is in the .hhr
+    indices_hit: The indices for each aminoacid relative to the hit sequence
+    indices_query: The indices for each aminoacid relative to the original query
+      sequence
+    original_query_sequence: String describing the original query sequence.
+
+  Returns:
+    Dictionary with indices in the original query sequence as keys and indices
+    in the hit sequence as values.
+  """
+  # If the hit is empty (no aligned residues), return empty mapping
+  if not hit_query_sequence:
+    return {}
+
+  # Remove gaps and find the offset of hit.query relative to original query.
+  hhsearch_query_sequence = hit_query_sequence.replace('-', '')
+  hit_sequence = hit_sequence.replace('-', '')
+  hhsearch_query_offset = original_query_sequence.find(hhsearch_query_sequence)
+
+  # Index of -1 used for gap characters. Subtract the min index ignoring gaps.
+  min_idx = min(x for x in indices_hit if x > -1)
+  fixed_indices_hit = [
+      x - min_idx if x > -1 else -1 for x in indices_hit
+  ]
+
+  min_idx = min(x for x in indices_query if x > -1)
+  fixed_indices_query = [x - min_idx if x > -1 else -1 for x in indices_query]
+
+  # Zip the corrected indices, ignore case where both seqs have gap characters.
+  mapping = {}
+  for q_i, q_t in zip(fixed_indices_query, fixed_indices_hit):
+    if q_t != -1 and q_i != -1:
+      if (q_t >= len(hit_sequence) or
+          q_i + hhsearch_query_offset >= len(original_query_sequence)):
+        continue
+      mapping[q_i + hhsearch_query_offset] = q_t
+
+  return mapping
+
+
+@dataclasses.dataclass(frozen=True)
+class SingleHitResult:
+  features: Optional[Mapping[str, Any]]
+  error: Optional[str]
+  warning: Optional[str]
+
+
+def _process_single_hit(
+    query_sequence: str,
+    query_pdb_code: Optional[str],
+    hit: parsers.TemplateHit,
+    mmcif_dir: str,
+    max_template_date: datetime.datetime,
+    release_dates: Mapping[str, datetime.datetime],
+    obsolete_pdbs: Mapping[str, Optional[str]],
+    kalign_binary_path: str,
+    strict_error_check: bool = False) -> SingleHitResult:
+  """Tries to extract template features from a single HHSearch hit."""
+  # Fail hard if we can't get the PDB ID and chain name from the hit.
+  hit_pdb_code, hit_chain_id = _get_pdb_id_and_chain(hit)
+
+  # This hit has been removed (obsoleted) from PDB, skip it.
+  if hit_pdb_code in obsolete_pdbs and obsolete_pdbs[hit_pdb_code] is None:
+    return SingleHitResult(
+        features=None, error=None, warning=f'Hit {hit_pdb_code} is obsolete.')
+
+  if hit_pdb_code not in release_dates:
+    if hit_pdb_code in obsolete_pdbs:
+      hit_pdb_code = obsolete_pdbs[hit_pdb_code]
+
+  # Pass hit_pdb_code since it might have changed due to the pdb being obsolete.
+  try:
+    _assess_hhsearch_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:
+    msg = f'hit {hit_pdb_code}_{hit_chain_id} did not pass prefilter: {str(e)}'
+    logging.info('%s: %s', query_pdb_code, msg)
+    if strict_error_check and isinstance(
+        e, (DateError, PdbIdError, DuplicateError)):
+      # In strict mode we treat some prefilter cases as errors.
+      return SingleHitResult(features=None, error=msg, warning=None)
+
+    return SingleHitResult(features=None, error=None, warning=None)
+
+  mapping = _build_query_to_hit_index_mapping(
+      hit.query, hit.hit_sequence, hit.indices_hit, hit.indices_query,
+      query_sequence)
+
+  # The mapping is from the query to the actual hit sequence, so we need to
+  # remove gaps (which regardless have a missing confidence score).
+  template_sequence = hit.hit_sequence.replace('-', '')
+
+  cif_path = os.path.join(mmcif_dir, hit_pdb_code + '.cif')
+  logging.info('Reading PDB entry from %s. Query: %s, template: %s',
+               cif_path, 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()
+
+  parsing_result = mmcif_parsing.parse(
+      file_id=hit_pdb_code, mmcif_string=cif_string)
+
+  if parsing_result.mmcif_object is not None:
+    hit_release_date = datetime.datetime.strptime(
+        parsing_result.mmcif_object.header['release_date'], '%Y-%m-%d')
+    if hit_release_date > max_template_date:
+      error = ('Template %s date (%s) > max template date (%s).' %
+               (hit_pdb_code, hit_release_date, max_template_date))
+      if strict_error_check:
+        return SingleHitResult(features=None, error=error, warning=None)
+      else:
+        logging.warning(error)
+        return SingleHitResult(features=None, error=None, warning=None)
+
+  try:
+    features, realign_warning = _extract_template_features(
+        mmcif_object=parsing_result.mmcif_object,
+        pdb_id=hit_pdb_code,
+        mapping=mapping,
+        template_sequence=template_sequence,
+        query_sequence=query_sequence,
+        template_chain_id=hit_chain_id,
+        kalign_binary_path=kalign_binary_path)
+    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
+    # computed. In such case the mmCIF parsing errors are not relevant.
+    return SingleHitResult(
+        features=features, error=None, warning=realign_warning)
+  except (NoChainsError, NoAtomDataInTemplateError,
+          TemplateAtomMaskAllZerosError) as e:
+    # These 3 errors indicate missing mmCIF experimental data rather than a
+    # problem with the template search, so turn them into warnings.
+    warning = ('%s_%s (sum_probs: %.2f, rank: %d): feature extracting errors: '
+               '%s, mmCIF parsing errors: %s'
+               % (hit_pdb_code, hit_chain_id, hit.sum_probs, hit.index,
+                  str(e), parsing_result.errors))
+    if strict_error_check:
+      return SingleHitResult(features=None, error=warning, warning=None)
+    else:
+      return SingleHitResult(features=None, error=None, warning=warning)
+  except Error as e:
+    error = ('%s_%s (sum_probs: %.2f, rank: %d): feature extracting errors: '
+             '%s, mmCIF parsing errors: %s'
+             % (hit_pdb_code, hit_chain_id, hit.sum_probs, hit.index,
+                str(e), parsing_result.errors))
+    return SingleHitResult(features=None, error=error, warning=None)
+
+
+@dataclasses.dataclass(frozen=True)
+class TemplateSearchResult:
+  features: Mapping[str, Any]
+  errors: Sequence[str]
+  warnings: Sequence[str]
+
+
+class TemplateHitFeaturizer:
+  """A class for turning hhr hits to template features."""
+
+  def __init__(
+      self,
+      mmcif_dir: str,
+      max_template_date: str,
+      max_hits: int,
+      kalign_binary_path: str,
+      release_dates_path: Optional[str],
+      obsolete_pdbs_path: Optional[str],
+      strict_error_check: bool = False):
+    """Initializes the Template Search.
+
+    Args:
+      mmcif_dir: Path to a directory with mmCIF structures. Once a template ID
+        is found by HHSearch, this directory is used to retrieve the template
+        data.
+      max_template_date: The maximum date permitted for template structures. No
+        template with date higher than this date will be returned. In ISO8601
+        date format, YYYY-MM-DD.
+      max_hits: The maximum number of templates that will be returned.
+      kalign_binary_path: The path to a kalign executable used for template
+        realignment.
+      release_dates_path: An optional path to a file with a mapping from PDB IDs
+        to their release dates. Thanks to this we don't have to redundantly
+        parse mmCIF files to get that information.
+      obsolete_pdbs_path: An optional path to a file containing a mapping from
+        obsolete PDB IDs to the PDB IDs of their replacements.
+      strict_error_check: If True, then the following will be treated as errors:
+        * If any template date is after the max_template_date.
+        * If any template has identical PDB ID to the query.
+        * If any template is a duplicate of the query.
+        * Any feature computation errors.
+    """
+    self._mmcif_dir = mmcif_dir
+    if not glob.glob(os.path.join(self._mmcif_dir, '*.cif')):
+      logging.error('Could not find CIFs in %s', self._mmcif_dir)
+      raise ValueError(f'Could not find CIFs in {self._mmcif_dir}')
+
+    try:
+      self._max_template_date = datetime.datetime.strptime(
+          max_template_date, '%Y-%m-%d')
+    except ValueError:
+      raise ValueError(
+          'max_template_date must be set and have format YYYY-MM-DD.')
+    self._max_hits = max_hits
+    self._kalign_binary_path = kalign_binary_path
+    self._strict_error_check = strict_error_check
+
+    if release_dates_path:
+      logging.info('Using precomputed release dates %s.', release_dates_path)
+      self._release_dates = _parse_release_dates(release_dates_path)
+    else:
+      self._release_dates = {}
+
+    if obsolete_pdbs_path:
+      logging.info('Using precomputed obsolete pdbs %s.', obsolete_pdbs_path)
+      self._obsolete_pdbs = _parse_obsolete(obsolete_pdbs_path)
+    else:
+      self._obsolete_pdbs = {}
+
+  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)
+
+    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
+    errors = []
+    warnings = []
+
+    for hit in sorted(hits, key=lambda x: x.sum_probs, reverse=True):
+      # We got all the templates we wanted, stop processing hits.
+      if num_hits >= self._max_hits:
+        break
+
+      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,
+          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)
+
+      # There could be an error even if there are some results, e.g. thrown by
+      # other unparsable chains in the same mmCIF file.
+      if result.warning:
+        warnings.append(result.warning)
+
+      if result.features is None:
+        logging.info('Skipped invalid hit %s, error: %s, warning: %s',
+                     hit.name, result.error, result.warning)
+      else:
+        # Increment the hit counter, since we got features out of this hit.
+        num_hits += 1
+        for k in template_features:
+          template_features[k].append(result.features[k])
+
+    for name in template_features:
+      if num_hits > 0:
+        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])
+
+    return TemplateSearchResult(
+        features=template_features, errors=errors, warnings=warnings)

alphafold/alphafold/data/tools/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+"""Python wrappers for third party tools."""

alphafold/alphafold/data/tools/hhblits.py

@@ -0,0 +1,155 @@
+# 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.
+
+"""Library to run HHblits from Python."""
+
+import glob
+import os
+import subprocess
+from typing import Any, Mapping, Optional, Sequence
+
+from absl import logging
+from alphafold.data.tools import utils
+# Internal import (7716).
+
+
+_HHBLITS_DEFAULT_P = 20
+_HHBLITS_DEFAULT_Z = 500
+
+
+class HHBlits:
+  """Python wrapper of the HHblits binary."""
+
+  def __init__(self,
+               *,
+               binary_path: str,
+               databases: Sequence[str],
+               n_cpu: int = 4,
+               n_iter: int = 3,
+               e_value: float = 0.001,
+               maxseq: int = 1_000_000,
+               realign_max: int = 100_000,
+               maxfilt: int = 100_000,
+               min_prefilter_hits: int = 1000,
+               all_seqs: bool = False,
+               alt: Optional[int] = None,
+               p: int = _HHBLITS_DEFAULT_P,
+               z: int = _HHBLITS_DEFAULT_Z):
+    """Initializes the Python HHblits wrapper.
+
+    Args:
+      binary_path: The path to the HHblits executable.
+      databases: A sequence of HHblits database paths. This should be the
+        common prefix for the database files (i.e. up to but not including
+        _hhm.ffindex etc.)
+      n_cpu: The number of CPUs to give HHblits.
+      n_iter: The number of HHblits iterations.
+      e_value: The E-value, see HHblits docs for more details.
+      maxseq: The maximum number of rows in an input alignment. Note that this
+        parameter is only supported in HHBlits version 3.1 and higher.
+      realign_max: Max number of HMM-HMM hits to realign. HHblits default: 500.
+      maxfilt: Max number of hits allowed to pass the 2nd prefilter.
+        HHblits default: 20000.
+      min_prefilter_hits: Min number of hits to pass prefilter.
+        HHblits default: 100.
+      all_seqs: Return all sequences in the MSA / Do not filter the result MSA.
+        HHblits default: False.
+      alt: Show up to this many alternative alignments.
+      p: Minimum Prob for a hit to be included in the output hhr file.
+        HHblits default: 20.
+      z: Hard cap on number of hits reported in the hhr file.
+        HHblits default: 500. NB: The relevant HHblits flag is -Z not -z.
+
+    Raises:
+      RuntimeError: If HHblits binary not found within the path.
+    """
+    self.binary_path = binary_path
+    self.databases = databases
+
+    for database_path in self.databases:
+      if not glob.glob(database_path + '_*'):
+        logging.error('Could not find HHBlits database %s', database_path)
+        raise ValueError(f'Could not find HHBlits database {database_path}')
+
+    self.n_cpu = n_cpu
+    self.n_iter = n_iter
+    self.e_value = e_value
+    self.maxseq = maxseq
+    self.realign_max = realign_max
+    self.maxfilt = maxfilt
+    self.min_prefilter_hits = min_prefilter_hits
+    self.all_seqs = all_seqs
+    self.alt = alt
+    self.p = p
+    self.z = z
+
+  def query(self, input_fasta_path: str) -> Mapping[str, Any]:
+    """Queries the database using HHblits."""
+    with utils.tmpdir_manager(base_dir='/tmp') as query_tmp_dir:
+      a3m_path = os.path.join(query_tmp_dir, 'output.a3m')
+
+      db_cmd = []
+      for db_path in self.databases:
+        db_cmd.append('-d')
+        db_cmd.append(db_path)
+      cmd = [
+          self.binary_path,
+          '-i', input_fasta_path,
+          '-cpu', str(self.n_cpu),
+          '-oa3m', a3m_path,
+          '-o', '/dev/null',
+          '-n', str(self.n_iter),
+          '-e', str(self.e_value),
+          '-maxseq', str(self.maxseq),
+          '-realign_max', str(self.realign_max),
+          '-maxfilt', str(self.maxfilt),
+          '-min_prefilter_hits', str(self.min_prefilter_hits)]
+      if self.all_seqs:
+        cmd += ['-all']
+      if self.alt:
+        cmd += ['-alt', str(self.alt)]
+      if self.p != _HHBLITS_DEFAULT_P:
+        cmd += ['-p', str(self.p)]
+      if self.z != _HHBLITS_DEFAULT_Z:
+        cmd += ['-Z', str(self.z)]
+      cmd += db_cmd
+
+      logging.info('Launching subprocess "%s"', ' '.join(cmd))
+      process = subprocess.Popen(
+          cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+      with utils.timing('HHblits query'):
+        stdout, stderr = process.communicate()
+        retcode = process.wait()
+
+      if retcode:
+        # Logs have a 15k character limit, so log HHblits error line by line.
+        logging.error('HHblits failed. HHblits stderr begin:')
+        for error_line in stderr.decode('utf-8').splitlines():
+          if error_line.strip():
+            logging.error(error_line.strip())
+        logging.error('HHblits stderr end')
+        raise RuntimeError('HHblits failed\nstdout:\n%s\n\nstderr:\n%s\n' % (
+            stdout.decode('utf-8'), stderr[:500_000].decode('utf-8')))
+
+      with open(a3m_path) as f:
+        a3m = f.read()
+
+    raw_output = dict(
+        a3m=a3m,
+        output=stdout,
+        stderr=stderr,
+        n_iter=self.n_iter,
+        e_value=self.e_value)
+    return raw_output

alphafold/alphafold/data/tools/hhsearch.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.
+
+"""Library to run HHsearch from Python."""
+
+import glob
+import os
+import subprocess
+from typing import Sequence
+
+from absl import logging
+
+from alphafold.data.tools import utils
+# Internal import (7716).
+
+
+class HHSearch:
+  """Python wrapper of the HHsearch binary."""
+
+  def __init__(self,
+               *,
+               binary_path: str,
+               databases: Sequence[str],
+               maxseq: int = 1_000_000):
+    """Initializes the Python HHsearch wrapper.
+
+    Args:
+      binary_path: The path to the HHsearch executable.
+      databases: A sequence of HHsearch database paths. This should be the
+        common prefix for the database files (i.e. up to but not including
+        _hhm.ffindex etc.)
+      maxseq: The maximum number of rows in an input alignment. Note that this
+        parameter is only supported in HHBlits version 3.1 and higher.
+
+    Raises:
+      RuntimeError: If HHsearch binary not found within the path.
+    """
+    self.binary_path = binary_path
+    self.databases = databases
+    self.maxseq = maxseq
+
+    for database_path in self.databases:
+      if not glob.glob(database_path + '_*'):
+        logging.error('Could not find HHsearch database %s', database_path)
+        raise ValueError(f'Could not find HHsearch database {database_path}')
+
+  def query(self, a3m: str) -> str:
+    """Queries the database using HHsearch using a given a3m."""
+    with utils.tmpdir_manager(base_dir='/tmp') as query_tmp_dir:
+      input_path = os.path.join(query_tmp_dir, 'query.a3m')
+      hhr_path = os.path.join(query_tmp_dir, 'output.hhr')
+      with open(input_path, 'w') as f:
+        f.write(a3m)
+
+      db_cmd = []
+      for db_path in self.databases:
+        db_cmd.append('-d')
+        db_cmd.append(db_path)
+      cmd = [self.binary_path,
+             '-i', input_path,
+             '-o', hhr_path,
+             '-maxseq', str(self.maxseq)
+             ] + db_cmd
+
+      logging.info('Launching subprocess "%s"', ' '.join(cmd))
+      process = subprocess.Popen(
+          cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+      with utils.timing('HHsearch query'):
+        stdout, stderr = process.communicate()
+        retcode = process.wait()
+
+      if retcode:
+        # Stderr is truncated to prevent proto size errors in Beam.
+        raise RuntimeError(
+            'HHSearch failed:\nstdout:\n%s\n\nstderr:\n%s\n' % (
+                stdout.decode('utf-8'), stderr[:100_000].decode('utf-8')))
+
+      with open(hhr_path) as f:
+        hhr = f.read()
+    return hhr