Merge remote-tracking branch 'upstream/main'

alphafold/model/tf/protein_features_test.py

@@ -27,6 +27,10 @@ def _random_bytes():
 
 class FeaturesTest(parameterized.TestCase, tf.test.TestCase):
 
+  def setUp(self):
+    super().setUp()
+    tf.disable_v2_behavior()
+
   def testFeatureNames(self):
     self.assertEqual(len(protein_features.FEATURE_SIZES),
                      len(protein_features.FEATURE_TYPES))
@@ -47,5 +51,4 @@ def testReplacement(self):
 
 
 if __name__ == '__main__':
-  tf.disable_v2_behavior()
   absltest.main()

alphafold/model/tf/shape_helpers_test.py

@@ -21,6 +21,10 @@
 
 class ShapeTest(tf.test.TestCase):
 
+  def setUp(self):
+    super().setUp()
+    tf.disable_v2_behavior()
+
   def test_shape_list(self):
     """Test that shape_list can allow for reshaping to dynamic shapes."""
     a = tf.zeros([10, 4, 4, 2])
@@ -35,5 +39,4 @@ def test_shape_list(self):
 
 
 if __name__ == '__main__':
-  tf.disable_v2_behavior()
   tf.test.main()

alphafold/model/utils.py

@@ -54,7 +54,7 @@ def mask_mean(mask, value, axis=None, drop_mask_channel=False, eps=1e-10):
     axis = [axis]
   elif axis is None:
     axis = list(range(len(mask_shape)))
-  assert isinstance(axis, collections.Iterable), (
+  assert isinstance(axis, collections.abc.Iterable), (
       'axis needs to be either an iterable, integer or "None"')
 
   broadcast_factor = 1.

alphafold/notebooks/notebook_utils.py

@@ -169,14 +169,25 @@ def empty_placeholder_template_features(
 
 
 def get_pae_json(pae: np.ndarray, max_pae: float) -> str:
-  """Returns the PAE in the same format as is used in the AFDB."""
+  """Returns the PAE in the same format as is used in the AFDB.
+
+  Note that the values are presented as floats to 1 decimal place,
+  whereas AFDB returns integer values.
+
+  Args:
+    pae: The n_res x n_res PAE array.
+    max_pae: The maximum possible PAE value.
+  Returns:
+    PAE output format as a JSON string.
+  """
+  # Check the PAE array is the correct shape.
+  if (pae.ndim != 2 or pae.shape[0] != pae.shape[1]):
+    raise ValueError(f'PAE must be a square matrix, got {pae.shape}')
+
+  # Round the predicted aligned errors to 1 decimal place.
   rounded_errors = np.round(pae.astype(np.float64), decimals=1)
-  indices = np.indices((len(rounded_errors), len(rounded_errors))) + 1
-  indices_1 = indices[0].flatten().tolist()
-  indices_2 = indices[1].flatten().tolist()
-  return json.dumps(
-      [{'residue1': indices_1,
-        'residue2': indices_2,
-        'distance': rounded_errors.flatten().tolist(),
-        'max_predicted_aligned_error': max_pae}],
-      indent=None, separators=(',', ':'))
+  formatted_output = [{
+      'predicted_aligned_error': rounded_errors.tolist(),
+      'max_predicted_aligned_error': max_pae
+  }]
+  return json.dumps(formatted_output, indent=None, separators=(',', ':'))

alphafold/notebooks/notebook_utils_test.py

@@ -194,9 +194,8 @@ def test_get_pae_json(self):
     pae = np.array([[0.01, 13.12345], [20.0987, 0.0]])
     pae_json = notebook_utils.get_pae_json(pae=pae, max_pae=31.75)
     self.assertEqual(
-        pae_json,
-        '[{"residue1":[1,1,2,2],"residue2":[1,2,1,2],"distance":'
-        '[0.0,13.1,20.1,0.0],"max_predicted_aligned_error":31.75}]')
+        pae_json, '[{"predicted_aligned_error":[[0.0,13.1],[20.1,0.0]],'
+        '"max_predicted_aligned_error":31.75}]')
 
 
 if __name__ == '__main__':

docker/Dockerfile

@@ -57,7 +57,7 @@ RUN wget -q -P /tmp \
 
 # Install conda packages.
 ENV PATH="/opt/conda/bin:$PATH"
-RUN conda update -qy conda \
+RUN conda install -qy conda==4.13.0 \
     && conda install -y -c conda-forge \
       openmm=7.5.1 \
       python=3.8 \
@@ -75,8 +75,8 @@ RUN pip3 install --upgrade --no-cache-dir pip \
     && pip3 install --no-cache-dir -r /app/alphafold/requirements.txt \
     && pip3 install --upgrade --no-cache-dir \
       jax==0.2.14 \
-      jaxlib==0.1.69+cuda${CUDA_JAXLIB/./} \
-      -f https://storage.googleapis.com/jax-releases/jax_releases.html
+      jaxlib==0.1.69+cuda$(cut -f1,2 -d. <<< ${CUDA_JAXLIB} | sed 's/\.//g') \
+      -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
 # Apply OpenMM patch.
 WORKDIR /opt/conda/lib/python3.8/site-packages

notebooks/AlphaFold.ipynb

@@ -18,6 +18,8 @@
         "\n",
         "Please note that this Colab notebook is provided as an early-access prototype and is not a finished product. It is provided for theoretical modelling only and caution should be exercised in its use. \n",
         "\n",
+        "The **PAE file format** has been updated to match AFDB. Please see the [AFDB FAQ](https://alphafold.ebi.ac.uk/faq/#faq-7) for a description of the new format.\n",
+        "\n",
         "**Citing this work**\n",
         "\n",
         "Any publication that discloses findings arising from using this notebook should [cite](https://github.com/deepmind/alphafold/#citing-this-work) the [AlphaFold paper](https://doi.org/10.1038/s41586-021-03819-2).\n",
@@ -102,7 +104,7 @@
         "\n",
         "      PATH=%env PATH\n",
         "      %env PATH=/opt/conda/bin:{PATH}\n",
-        "      %shell conda update -qy conda \\\n",
+        "      %shell conda install -qy conda==4.13.0 \\\n",
         "          \u0026\u0026 conda install -qy -c conda-forge \\\n",
         "            python=3.7 \\\n",
         "            openmm=7.5.1 \\\n",

run_alphafold.py

@@ -21,7 +21,7 @@
 import shutil
 import sys
 import time
-from typing import Dict, Union, Optional
+from typing import Dict, Union
 
 from absl import app
 from absl import flags

setup.py

@@ -18,7 +18,7 @@
 
 setup(
     name='alphafold',
-    version='2.2.2',
+    version='2.2.3',
     description='An implementation of the inference pipeline of AlphaFold v2.0.'
     'This is a completely new model that was entered as AlphaFold2 in CASP14 '
     'and published in Nature.',
@@ -42,7 +42,10 @@
         'scipy',
         'tensorflow-cpu',
     ],
-    tests_require=['mock'],
+    tests_require=[
+        'matplotlib',  # For notebook_utils_test.
+        'mock',
+    ],
     classifiers=[
         'Development Status :: 5 - Production/Stable',
         'Intended Audience :: Science/Research',