"doc/vscode:/vscode.git/clone" did not exist on "2d3aaef8b63a08d3e17cdd30d1cbb851e2898f9c"
Commit 6f6e8c2e authored by DeepMind's avatar DeepMind Committed by Copybara-Service
Browse files

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

PiperOrigin-RevId: 467198544
Change-Id: I504dc91c993b10dbb1cef448ef8eaa1161186acb
parent 22fc0eea
......@@ -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=(',', ':'))
......@@ -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__':
......
......@@ -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",
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment