Fix PAE visualization for single chain predictions.

PiperOrigin-RevId: 407773747 Change-Id: I8635dd77ecf366f4f5a362993b964adffc3ff186
google-deepmind · Nov 5, 2021 · 91b4322 · 91b4322
1 parent c128d1a
commit 91b4322
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/notebooks/AlphaFold.ipynb b/notebooks/AlphaFold.ipynb
@@ -648,8 +648,9 @@
         "  total_num_res = best_unrelaxed_prot.residue_index.shape[-1]\n",
         "  chain_ids = best_unrelaxed_prot.chain_index\n",
         "  for chain_boundary in np.nonzero(chain_ids[:-1] - chain_ids[1:]):\n",
-        "    plt.plot([0, total_num_res], [chain_boundary, chain_boundary], color='red')\n",
-        "    plt.plot([chain_boundary, chain_boundary], [0, total_num_res], color='red')\n",
+        "    if chain_boundary.size:\n",
+        "      plt.plot([0, total_num_res], [chain_boundary, chain_boundary], color='red')\n",
+        "      plt.plot([chain_boundary, chain_boundary], [0, total_num_res], color='red')\n",
         "\n",
         "  plt.title('Predicted Aligned Error')\n",
         "  plt.xlabel('Scored residue')\n",