Fix PAE visualization for single chain predictions.

PiperOrigin-RevId: 407773747
Change-Id: I8635dd77ecf366f4f5a362993b964adffc3ff186

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