Fix pae plots being generated for non ptm models

sokrypton · Dec 28, 2023 · 2515b3b · 2515b3b
1 parent f80e0d8
commit 2515b3b
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 9 deletions.
diff --git a/AlphaFold2.ipynb b/AlphaFold2.ipynb
@@ -471,7 +471,10 @@
         "    img = f.read()\n",
         "  return prefix + base64.b64encode(img).decode('utf-8')\n",
         "\n",
-        "pae = image_to_data_url(os.path.join(jobname,f\"{jobname}{jobname_prefix}_pae.png\"))\n",
+        "pae = \"\"\n",
+        "pae_file = os.path.join(jobname,f\"{jobname}{jobname_prefix}_pae.png\")\n",
+        "if os.path.isfile(pae_file):\n",
+        "    pae = image_to_data_url(pae_file)\n",
         "cov = image_to_data_url(os.path.join(jobname,f\"{jobname}{jobname_prefix}_coverage.png\"))\n",
         "plddt = image_to_data_url(os.path.join(jobname,f\"{jobname}{jobname_prefix}_plddt.png\"))\n",
         "display(HTML(f\"\"\"\n",
@@ -493,7 +496,7 @@
         "</style>\n",
         "<div style=\"max-width:90%; padding:2em;\">\n",
         "  <h1>Plots for {escape(jobname)}</h1>\n",
-        "  <img src=\"{pae}\" class=\"full\" />\n",
+        "  { '<!--' if pae == '' else '' }<img src=\"{pae}\" class=\"full\" />{ '-->' if pae == '' else '' }\n",
         "  <img src=\"{cov}\" class=\"half\" />\n",
         "  <img src=\"{plddt}\" class=\"half\" />\n",
         "</div>\n",

diff --git a/colabfold/batch.py b/colabfold/batch.py
@@ -1620,13 +1620,13 @@ def run(
                     "max_predicted_aligned_error":scores[0]["max_pae"]}))
                 result_files.append(af_pae_file)
 
-            # make pAE plots
-            paes_plot = plot_paes([np.asarray(x["pae"]) for x in scores],
-                Ls=query_sequence_len_array, dpi=dpi)
-            pae_png = result_dir.joinpath(f"{jobname}_pae.png")
-            paes_plot.savefig(str(pae_png), bbox_inches='tight')
-            paes_plot.close()
-            result_files.append(pae_png)
+                # make pAE plots
+                paes_plot = plot_paes([np.asarray(x["pae"]) for x in scores],
+                    Ls=query_sequence_len_array, dpi=dpi)
+                pae_png = result_dir.joinpath(f"{jobname}_pae.png")
+                paes_plot.savefig(str(pae_png), bbox_inches='tight')
+                paes_plot.close()
+                result_files.append(pae_png)
 
             # make pLDDT plot
             plddt_plot = plot_plddts([np.asarray(x["plddt"]) for x in scores],