add fast af and fixes

af_backprop/alphafold/model/folding.py

@@ -441,22 +441,20 @@ def generate_affines(representations, batch, config, global_config,
       name='pair_layer_norm')(
           representations['pair'])
 
-  def fold_iter(x,_):
-    x["key"], key = x["key"].split()
-    x["act"], out = fold_iteration(
-        x["act"],
+  def fold_iter(act, key):
+    act, out = fold_iteration(
+        act,
         initial_act=initial_act,
         static_feat_2d=act_2d,
-        safe_key=key,
+        safe_key=prng.SafeKey(key),
         sequence_mask=sequence_mask,
         update_affine=True,
         is_training=is_training,
         aatype=batch['aatype'],
         scale_rate=batch["scale_rate"])
-    return x, out
-  x = {"act":activations,"key":safe_key}
-  x, output = hk.scan(fold_iter, x, None, c.num_layer)
-  activations = x["act"]
+    return act, out  
+  keys = jax.random.split(safe_key.get(), c.num_layer)
+  activations, output = hk.scan(fold_iter, activations, keys)
   
   # Include the activations in the output dict for use by the LDDT-Head.
   output['act'] = activations['act']

af_backprop/examples/AlphaFold_single.ipynb

@@ -32,7 +32,11 @@
       "source": [
         "#AlphaFold - single sequence input\n",
         "- WARNING - For DEMO and educational purposes only. \n",
-        "- For natural proteins you often need more than a single sequence to accurately predict the structure. See [ColabFold](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/AlphaFold2.ipynb) notebook if you want to predict the protein structure from a multiple-sequence-alignment. That being said, this notebook could potentially be useful for evaluating *de novo* designed proteins.\n"
+        "- For natural proteins you often need more than a single sequence to accurately predict the structure. See [ColabFold](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/AlphaFold2.ipynb) notebook if you want to predict the protein structure from a multiple-sequence-alignment. That being said, this notebook could be useful for evaluating *de novo* designed proteins and learning the idealized principles of proteins.\n",
+        "\n",
+        "### Tips and Instructions\n",
+        "- click the little ▶ play icon to the left of each cell below.\n",
+        "- hold mouseover aminoacid to get name and position number"
       ],
       "metadata": {
         "id": "VpfCw7IzVHXv"
@@ -42,12 +46,18 @@
       "cell_type": "code",
       "source": [
         "#@title Setup\n",
+        "\n",
+        "# import libraries\n",
         "from IPython.utils import io\n",
         "import os,sys,re\n",
         "import tensorflow as tf\n",
         "import jax\n",
         "import jax.numpy as jnp\n",
         "import numpy as np\n",
+        "import matplotlib\n",
+        "from matplotlib import animation\n",
+        "import matplotlib.pyplot as plt\n",
+        "from IPython.display import HTML\n",
         "\n",
         "with io.capture_output() as captured:\n",
         "  if not os.path.isdir(\"af_backprop\"):\n",
@@ -58,6 +68,7 @@
         "    %shell mkdir params\n",
         "    %shell curl -fsSL https://storage.googleapis.com/alphafold/alphafold_params_2021-07-14.tar | tar x -C params\n",
         "\n",
+        "# configure which device to use\n",
         "try:\n",
         "  # check if TPU is available\n",
         "  import jax.tools.colab_tpu\n",
@@ -75,6 +86,7 @@
         "    tf.config.set_visible_devices([], 'GPU')\n",
         "\n",
         "sys.path.append('/content/af_backprop')\n",
+        "\n",
         "# import libraries\n",
         "from utils import update_seq, update_aatype, get_plddt, get_pae\n",
         "import colabfold as cf\n",
@@ -83,12 +95,12 @@
         "from alphafold.model import data, config, model\n",
         "from alphafold.common import residue_constants\n",
         "\n",
+        "# custom functions\n",
         "def clear_mem():\n",
         "  backend = jax.lib.xla_bridge.get_backend()\n",
         "  for buf in backend.live_buffers(): buf.delete()\n",
         "\n",
         "def setup_model(max_len, model_name=\"model_2_ptm\"):\n",
-        "\n",
         "  clear_mem()\n",
         "\n",
         "  # setup model\n",
@@ -132,121 +144,27 @@
         "            \"init_pos\":outputs['structure_module']['final_atom_positions'][None]}\n",
         "    \n",
         "    aux = {\"final_atom_positions\":outputs[\"structure_module\"][\"final_atom_positions\"],\n",
-        "          \"final_atom_mask\":outputs[\"structure_module\"][\"final_atom_mask\"],\n",
-        "          \"plddt\":get_plddt(outputs),\"pae\":get_pae(outputs),\n",
-        "          \"inputs\":inputs, \"prev\":prev}\n",
+        "           \"final_atom_mask\":outputs[\"structure_module\"][\"final_atom_mask\"],\n",
+        "           \"plddt\":get_plddt(outputs),\"pae\":get_pae(outputs),\n",
+        "           \"inputs\":inputs, \"prev\":prev}\n",
         "    return aux\n",
         "\n",
         "  return jax.jit(runner), {\"inputs\":inputs,\"params\":model_params}\n",
         "\n",
-        "MAX_LEN = 50\n",
-        "RUNNER, OPT = setup_model(MAX_LEN)"
-      ],
-      "metadata": {
-        "cellView": "form",
-        "id": "24ybo88aBiSU"
-      },
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "%%time\n",
-        "#@title Enter the amino acid sequence to fold ⬇️\n",
-        "\n",
-        "sequence = 'GGGGGGGGGGGGGGGGGGGG' #@param {type:\"string\"}\n",
-        "recycles = 0 #@param [\"0\", \"1\", \"2\", \"3\", \"6\", \"12\", \"24\"] {type:\"raw\"}\n",
-        "SEQ = re.sub(\"[^A-Z]\", \"\", sequence.upper())\n",
-        "LEN = len(SEQ)\n",
-        "if LEN > MAX_LEN:\n",
-        "  print(\"recompiling...\")\n",
-        "  MAX_LEN = LEN\n",
-        "  RUNNER, OPT = setup_model(MAX_LEN)\n",
-        "\n",
-        "x = np.array([residue_constants.restype_order.get(aa,0) for aa in SEQ])\n",
-        "x = np.pad(x,[0,MAX_LEN-LEN],constant_values=-1)\n",
-        "x = jax.nn.one_hot(x,20)\n",
-        "\n",
-        "OPT[\"prev\"] = {'init_msa_first_row': np.zeros([1, MAX_LEN, 256]),\n",
-        "               'init_pair': np.zeros([1, MAX_LEN, MAX_LEN, 128]),\n",
-        "               'init_pos': np.zeros([1, MAX_LEN, 37, 3])}\n",
-        "\n",
-        "positions = []\n",
-        "plddts = []\n",
-        "for r in range(recycles+1):\n",
-        "  outs = RUNNER(x, OPT)\n",
-        "  outs = jax.tree_map(lambda x:np.asarray(x), outs)\n",
-        "  positions.append(outs[\"prev\"][\"init_pos\"][0,:LEN])\n",
-        "  plddts.append(outs[\"plddt\"][:LEN])\n",
-        "  OPT[\"prev\"] = outs[\"prev\"]\n",
-        "  if recycles > 0:\n",
-        "    print(r, plddts[-1].mean())"
-      ],
-      "metadata": {
-        "cellView": "form",
-        "id": "cAoC4ar8G7ZH"
-      },
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "#@title Display 3D structure {run: \"auto\"}\n",
-        "color = \"lDDT\" #@param [\"chain\", \"lDDT\", \"rainbow\"]\n",
-        "show_sidechains = True #@param {type:\"boolean\"}\n",
-        "show_mainchains = False #@param {type:\"boolean\"}\n",
-        "#@markdown - TIP - hold mouse over aminoacid to get name and position number\n",
-        "\n",
-        "def save_pdb(outs, filename):\n",
+        "def save_pdb(outs, filename, LEN):\n",
         "  '''save pdb coordinates'''\n",
-        "  p = {\"residue_index\":outs[\"inputs\"][\"residue_index\"][0][:LEN] + 1,\n",
-        "        \"aatype\":outs[\"inputs\"][\"aatype\"].argmax(-1)[0][:LEN],\n",
-        "        \"atom_positions\":outs[\"final_atom_positions\"][:LEN],\n",
-        "        \"atom_mask\":outs[\"final_atom_mask\"][:LEN]}\n",
-        "  b_factors = 100.0 * outs[\"plddt\"][:LEN,None] * p[\"atom_mask\"]\n",
+        "  p = {\"residue_index\":outs[\"inputs\"][\"residue_index\"][0] + 1,\n",
+        "       \"aatype\":outs[\"inputs\"][\"aatype\"].argmax(-1)[0],\n",
+        "       \"atom_positions\":outs[\"final_atom_positions\"],\n",
+        "       \"atom_mask\":outs[\"final_atom_mask\"],\n",
+        "       \"plddt\":outs[\"plddt\"]}\n",
+        "  p = jax.tree_map(lambda x:x[:LEN], p)\n",
+        "  b_factors = 100.0 * p.pop(\"plddt\")[:,None] * p[\"atom_mask\"]\n",
         "  p = protein.Protein(**p,b_factors=b_factors)\n",
         "  pdb_lines = protein.to_pdb(p)\n",
-        "  with open(filename, 'w') as f:\n",
-        "    f.write(pdb_lines)\n",
-        "\n",
-        "save_pdb(outs,\"out.pdb\")\n",
-        "num_res = int(outs[\"inputs\"][\"aatype\"][0].sum())\n",
-        "\n",
-        "v = cf.show_pdb(\"out.pdb\", show_sidechains, show_mainchains, color,\n",
-        "                color_HP=True, size=(800,480))       \n",
-        "v.setHoverable({},\n",
-        "               True,\n",
-        "               '''function(atom,viewer,event,container){if(!atom.label){atom.label=viewer.addLabel(\"      \"+atom.resn+\":\"+atom.resi,{position:atom,backgroundColor:'mintcream',fontColor:'black'});}}''',\n",
-        "               '''function(atom,viewer){if(atom.label){viewer.removeLabel(atom.label);delete atom.label;}}''')\n",
-        "v.show()           \n",
-        "\n",
-        "if color == \"lDDT\":\n",
-        "  cf.plot_plddt_legend().show()  \n",
-        "if \"pae\" in outs:\n",
-        "  cf.plot_confidence(outs[\"plddt\"][:LEN]*100, outs[\"pae\"][:LEN,:LEN]).show()\n",
-        "else:\n",
-        "  cf.plot_confidence(outs[\"plddt\"][:LEN]*100).show()"
-      ],
-      "metadata": {
-        "cellView": "form",
-        "id": "-KbUGG4ZOp0J"
-      },
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "#@title Animate\n",
-        "#@markdown - Animate trajectory if more than 0 recycle(s)\n",
-        "import matplotlib\n",
-        "from matplotlib import animation\n",
-        "import matplotlib.pyplot as plt\n",
-        "from IPython.display import HTML\n",
+        "  with open(filename, 'w') as f: f.write(pdb_lines)\n",
         "\n",
-        "def make_animation(positions, plddts=None, line_w=2.0):\n",
+        "def make_animation(positions, plddts=None, line_w=2.0, dpi=100):\n",
         "\n",
         "  def ca_align_to_last(positions):\n",
         "    def align(P, Q):\n",
@@ -269,7 +187,7 @@
         "  fig.subplots_adjust(top = 0.90, bottom = 0.10, right = 1, left = 0, hspace = 0, wspace = 0)\n",
         "  fig.set_figwidth(13)\n",
         "  fig.set_figheight(5)\n",
-        "  fig.set_dpi(100)\n",
+        "  fig.set_dpi(dpi)\n",
         "\n",
         "  xy_min = pos[...,:2].min() - 1\n",
         "  xy_max = pos[...,:2].max() + 1\n",
@@ -295,8 +213,98 @@
         "    \n",
         "  ani = animation.ArtistAnimation(fig, ims, blit=True, interval=120)\n",
         "  plt.close()\n",
-        "  return ani.to_html5_video()\n",
+        "  return ani.to_html5_video()"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "24ybo88aBiSU"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%%time\n",
+        "#@title Enter the amino acid sequence to fold ⬇️\n",
+        "\n",
+        "# initialize model runner if doesn't exist\n",
+        "if \"runner\" not in dir():\n",
+        "  max_length = 50\n",
+        "  current_seq = \"\"\n",
+        "  r = -1\n",
+        "  runner, I = setup_model(max_length)\n",
+        "\n",
+        "# collect user inputs\n",
+        "sequence = 'GGGGGGGGGGGGGGGGGGGGG' #@param {type:\"string\"}\n",
+        "recycles = 0 #@param [\"0\", \"1\", \"2\", \"3\", \"6\", \"12\", \"24\"] {type:\"raw\"}\n",
+        "sequence = re.sub(\"[^A-Z]\", \"\", sequence.upper())\n",
+        "length = len(sequence)\n",
+        "\n",
+        "# if length greater than max_len, recompile for larger length\n",
+        "if length > max_length:\n",
+        "  max_length = length + 10 # a little buffer\n",
+        "  runner, I = setup_model(max_length)\n",
         "\n",
+        "if sequence != current_seq:\n",
+        "  outs = []\n",
+        "  positions = []\n",
+        "  plddts = []\n",
+        "  paes = []\n",
+        "  r = -1\n",
+        "  # convert sequence to one_hot\n",
+        "  x = np.array([residue_constants.restype_order.get(aa,0) for aa in sequence])\n",
+        "  x = np.pad(x,[0,max_length-length],constant_values=-1)\n",
+        "  x = jax.nn.one_hot(x,20)\n",
+        "\n",
+        "  # restart recycle\n",
+        "  I[\"prev\"] = {'init_msa_first_row': np.zeros([1, max_length, 256]),\n",
+        "               'init_pair': np.zeros([1, max_length, max_length, 128]),\n",
+        "               'init_pos': np.zeros([1, max_length, 37, 3])}\n",
+        "  current_seq = sequence\n",
+        "\n",
+        "# run for defined number of recycles\n",
+        "while r < recycles:\n",
+        "  O = runner(x, I)\n",
+        "  O = jax.tree_map(lambda x:np.asarray(x), O)\n",
+        "  positions.append(O[\"final_atom_positions\"][:length])\n",
+        "  plddts.append(O[\"plddt\"][:length])\n",
+        "  paes.append(O[\"pae\"][:length,:length])\n",
+        "  I[\"prev\"] = O[\"prev\"]\n",
+        "  outs.append(O)\n",
+        "  r += 1\n",
+        "\n",
+        "#@markdown #### Display options\n",
+        "color = \"lDDT\" #@param [\"chain\", \"lDDT\", \"rainbow\"]\n",
+        "show_sidechains = True #@param {type:\"boolean\"}\n",
+        "show_mainchains = False #@param {type:\"boolean\"}\n",
+        "\n",
+        "print(f\"plotting prediction at recycle={recycles}\")\n",
+        "save_pdb(outs[recycles], \"out.pdb\", length)\n",
+        "v = cf.show_pdb(\"out.pdb\", show_sidechains, show_mainchains, color,\n",
+        "                 color_HP=True, size=(800,480))       \n",
+        "v.setHoverable({},\n",
+        "               True,\n",
+        "               '''function(atom,viewer,event,container){if(!atom.label){atom.label=viewer.addLabel(\"      \"+atom.resn+\":\"+atom.resi,{position:atom,backgroundColor:'mintcream',fontColor:'black'});}}''',\n",
+        "               '''function(atom,viewer){if(atom.label){viewer.removeLabel(atom.label);delete atom.label;}}''')\n",
+        "v.show()           \n",
+        "if color == \"lDDT\": cf.plot_plddt_legend().show()\n",
+        "\n",
+        "# add confidence plots\n",
+        "cf.plot_confidence(plddts[recycles]*100, paes[recycles]).show()"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "cAoC4ar8G7ZH"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title Animate\n",
+        "#@markdown - Animate trajectory if more than 0 recycle(s)\n",
         "HTML(make_animation(np.asarray(positions),\n",
         "                    np.asarray(plddts) * 100.0))"
       ],
@@ -308,4 +316,4 @@
       "outputs": []
     }
   ]
-}
+}

alphafold/LICENSE

@@ -1,202 +0,0 @@
-
-                                 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

@@ -1,14 +0,0 @@
-# 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

@@ -1,14 +0,0 @@
-# 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

@@ -1,155 +0,0 @@
-# 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

@@ -1,229 +0,0 @@
-# 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

@@ -1,89 +0,0 @@
-# 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

@@ -1,895 +0,0 @@
-# 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

@@ -1,190 +0,0 @@
-# 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

@@ -1,14 +0,0 @@
-# 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

@@ -1,384 +0,0 @@
-# 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

@@ -1,364 +0,0 @@
-# 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

@@ -1,209 +0,0 @@
-# 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

@@ -1,922 +0,0 @@
-# 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

@@ -1,14 +0,0 @@
-# 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

@@ -1,155 +0,0 @@
-# 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