Merge pull request #386 from KosinskiLab/update-analysis-pipeline

Fix KeyError 'iptm+ptm' not found in analysis pipeline
KosinskiLab · Jul 22, 2024 · a80beae · a80beae
2 parents 80e71b1 + 8ddb328
commit a80beae
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 5 deletions.
diff --git a/alphapulldown/analysis_pipeline/alpha_analysis_jax0.3.def b/alphapulldown/analysis_pipeline/alpha_analysis_jax0.3.def
@@ -11,4 +11,4 @@ From: kosinskilab/alpha_analysis_basis_jax0.3:latest
     chmod +x /app/run_get_good_pae.sh
     chmod +x /app/run_execute_notebook.sh
     chmod +x /app/run_pi_score.sh
-    chmod +x /app/alpha-analysis/get_good_inter_pae.py
+    chmod +x /app/alpha-analysis/get_good_inter_pae.py
diff --git a/alphapulldown/analysis_pipeline/alpha_analysis_jax0.4.def b/alphapulldown/analysis_pipeline/alpha_analysis_jax0.4.def
@@ -11,4 +11,4 @@ From: kosinskilab/alpha_analysis_basis_jax0.4:latest
     chmod +x /app/run_get_good_pae.sh
     chmod +x /app/run_execute_notebook.sh
     chmod +x /app/run_pi_score.sh
-    chmod +x /app/alpha-analysis/get_good_inter_pae.py
+    chmod +x /app/alpha-analysis/get_good_inter_pae.py
diff --git a/alphapulldown/analysis_pipeline/create_notebook.py b/alphapulldown/analysis_pipeline/create_notebook.py
@@ -111,7 +111,7 @@ def main(argv):
                 open(os.path.join(result_subdir, "ranking_debug.json"), "rb")
             )
 
-            if "iptm" in data.keys() or "iptm+ptm" in data.keys():
+            if "iptm+ptm" in data.keys():
                 iptm_ptm_score = data["iptm+ptm"][best_model]   
                 pae_mtx,iptm_score = obtain_pae_and_iptm(result_subdir=os.path.join(FLAGS.output_dir, job),
                                               best_model=best_model)

diff --git a/alphapulldown/analysis_pipeline/get_good_inter_pae.py b/alphapulldown/analysis_pipeline/get_good_inter_pae.py
@@ -131,7 +131,7 @@ def main(argv):
                         best_model = ranking_data['order'][0]
                         data = ranking_data
 
-                        if "iptm" in data or "iptm+ptm" in data:
+                        if "iptm+ptm" in data:
                             iptm_ptm_score = data['iptm+ptm'][best_model]
 
                             try:
@@ -154,6 +154,8 @@ def main(argv):
                                     output_df = pd.concat([score_df, output_df])
                             except Exception as e:
                                 logging.error(f"Error processing PAE and iPTM for job {job}: {e}")
+                        else:
+                            logging.warning(f"{job} does not seem to be a multimeric model. iptm+ptm scores are not in the ranking_debug.json Skipped")
                     except Exception as e:
                         logging.error(f"Error getting the best model name from ranking_debug.json for job :{job} : {e}")
                 except Exception as e:

diff --git a/alphapulldown/folding_backend/alphafold_backend.py b/alphapulldown/folding_backend/alphafold_backend.py
@@ -624,7 +624,15 @@ def postprocess(
             if f'relax_{model_name}' in timings:
                 continue
             t_0 = time.time()
-            unrelaxed_protein = prediction_results[model_name]['unrelaxed_protein']
+            if 'unrelaxed_protein' in prediction_results[model_name].keys():
+                unrelaxed_protein = prediction_results[model_name]['unrelaxed_protein']
+            else:
+                unrelaxed_pdb_path = os.path.join(output_dir, f"unrelaxed_{model_name}.pdb")
+                if not os.path.exists(unrelaxed_pdb_path):
+                    logging.error(f"Cannot find {unrelaxed_pdb_path} for relaxation! Skipping...")
+                    continue
+                unrelaxed_pdb_string = open(unrelaxed_pdb_path, 'r').read()
+                unrelaxed_protein = protein.from_pdb_string(unrelaxed_pdb_string)
             relaxed_pdb_str, _, violations = amber_relaxer.process(
                 prot=unrelaxed_protein)
             relax_metrics[model_name] = {