Changed PAE json output to new format that matches the format used in AFDB.

PiperOrigin-RevId: 467198544
Change-Id: I504dc91c993b10dbb1cef448ef8eaa1161186acb

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 @@ class NotebookUtilsTest(parameterized.TestCase):
     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__':

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",