Fix up CASP15 single-seq method support

Author: amorehead

posebench/analysis/complex_alignment.py

@@ -168,7 +168,7 @@ def align_complex_to_protein_only(
     save_ligand: bool = True,
     aligned_filename_suffix: str = "_aligned",
     atom_df_name: str = "ATOM",
-):
+) -> int:
     """Align a predicted protein-ligand structure to a reference protein structure.
 
     :param predicted_protein_pdb: Path to the predicted protein structure in PDB format
@@ -178,6 +178,7 @@ def align_complex_to_protein_only(
     :param save_ligand: Whether to save the aligned ligand structure
     :param aligned_filename_suffix: suffix to append to the aligned files
     :param atom_df_name: Name of the atom dataframe in the PDB file
+    :return: 0 if successful, 1 if unsuccessful
     """
     from biopandas.pdb import PandasPdb
     from rdkit import Chem
@@ -197,14 +198,14 @@ def align_complex_to_protein_only(
         logger.warning(
             f"Unable to parse predicted protein structure {predicted_protein_pdb} due to the error: {e}. Skipping..."
         )
-        return
+        return 1
     try:
         reference_rec = parse_pdb_from_path(reference_protein_pdb)
     except Exception as e:
         logger.warning(
             f"Unable to parse reference protein structure {reference_protein_pdb} due to the error: {e}. Skipping..."
         )
-        return
+        return 1
     if predicted_ligand_sdf is not None:
         predicted_ligand = read_molecule(predicted_ligand_sdf, remove_hs=True, sanitize=True)
     try:
@@ -215,7 +216,7 @@ def align_complex_to_protein_only(
         logger.warning(
             f"Unable to extract predicted protein structure coordinates for input {predicted_protein_pdb} due to the error: {e}. Skipping..."
         )
-        return
+        return 1
     try:
         reference_calpha_coords = extract_receptor_structure(
             reference_rec, None, filter_out_hetero_residues=True
@@ -224,23 +225,23 @@ def align_complex_to_protein_only(
         logger.warning(
             f"Unable to extract reference protein structure coordinates for input {predicted_protein_pdb} due to the error: {e}. Skipping..."
         )
-        return
+        return 1
     if predicted_ligand_sdf is not None:
         try:
             predicted_ligand_conf = predicted_ligand.GetConformer()
         except Exception as e:
             logger.warning(
                 f"Unable to extract predicted ligand conformer for {predicted_ligand_sdf} due to the error: {e}. Skipping..."
             )
-            return
+            return 1
 
     if reference_calpha_coords.shape != predicted_calpha_coords.shape:
         logger.warning(
             f"Receptor structures differ for prediction {predicted_protein_pdb}. Skipping due to shape mismatch:",
             reference_calpha_coords.shape,
             predicted_calpha_coords.shape,
         )
-        return
+        return 1
 
     # Perform the alignment
     rotation, reference_calpha_centroid, predicted_calpha_centroid = align_prediction(
@@ -287,6 +288,8 @@ def align_complex_to_protein_only(
         ) as f:
             f.write(predicted_ligand)
 
+    return 0
+
 
 @hydra.main(
     version_base="1.3",

posebench/models/ensemble_generation.py

@@ -839,6 +839,7 @@ def get_method_predictions(
     cfg: DictConfig,
     binding_site_method: Optional[str] = None,
     input_protein_filepath: Optional[str] = None,
+    is_ss_method: bool = False,
 ) -> List[Tuple[str, str]]:
     """Get the predictions generated by the method.
 
@@ -848,11 +849,13 @@ def get_method_predictions(
     :param binding_site_method: Optional name of the method used to predict AutoDock Vina's binding
         sites.
     :param input_protein_filepath: Optional path to the input protein structure PDB file.
+    :param is_ss_method: Whether the method is a single-sequence method.
     :return: List of method predictions, each as a tuple of the output protein filepath and the
         output ligand filepath.
     """
     pocket_only_suffix = "_pocket_only" if cfg.pocket_only_baseline else ""
     no_ilcl_suffix = "_no_ilcl" if cfg.neuralplexer_no_ilcl else ""
+    single_seq_suffix = "_ss" if is_ss_method else ""
 
     if method == "diffdock":
         ensemble_benchmarking_output_dir = (
@@ -981,7 +984,7 @@ def get_method_predictions(
     elif method == "neuralplexer":
         ensemble_benchmarking_output_dir = (
             Path(cfg.input_dir if cfg.input_dir else cfg.neuralplexer_out_path).parent
-            / f"neuralplexer{pocket_only_suffix}{no_ilcl_suffix}_{cfg.ensemble_benchmarking_dataset}_outputs_{cfg.ensemble_benchmarking_repeat_index}"
+            / f"neuralplexer{single_seq_suffix}{pocket_only_suffix}{no_ilcl_suffix}_{cfg.ensemble_benchmarking_dataset}_outputs_{cfg.ensemble_benchmarking_repeat_index}"
             if cfg.ensemble_benchmarking
             else (cfg.input_dir if cfg.input_dir else cfg.neuralplexer_out_path)
         )
@@ -1078,7 +1081,7 @@ def get_method_predictions(
     elif method == "chai-lab":
         ensemble_benchmarking_output_dir = (
             Path(cfg.input_dir if cfg.input_dir else cfg.chai_out_path).parent
-            / f"chai-lab{pocket_only_suffix}_{cfg.ensemble_benchmarking_dataset}_outputs_{cfg.ensemble_benchmarking_repeat_index}"
+            / f"chai-lab{single_seq_suffix}{pocket_only_suffix}_{cfg.ensemble_benchmarking_dataset}_outputs_{cfg.ensemble_benchmarking_repeat_index}"
             if cfg.ensemble_benchmarking
             else (cfg.input_dir if cfg.input_dir else cfg.chai_out_path)
         )
@@ -1112,7 +1115,7 @@ def get_method_predictions(
     elif method == "alphafold3":
         ensemble_benchmarking_output_dir = (
             Path(cfg.input_dir if cfg.input_dir else cfg.alphafold3_out_path).parent
-            / f"alphafold3{pocket_only_suffix}_{cfg.ensemble_benchmarking_dataset}_outputs_{cfg.ensemble_benchmarking_repeat_index}"
+            / f"alphafold3{single_seq_suffix}{pocket_only_suffix}_{cfg.ensemble_benchmarking_dataset}_outputs_{cfg.ensemble_benchmarking_repeat_index}"
             if cfg.ensemble_benchmarking
             else (cfg.input_dir if cfg.input_dir else cfg.alphafold3_out_path)
         )
@@ -1340,7 +1343,7 @@ def generate_ensemble_predictions(
 
     if cfg.resume:
         ensemble_predictions_dict = {}
-        for method in cfg.ensemble_methods:
+        for method, is_ss_method in zip(cfg.ensemble_methods, cfg.is_ss_ensemble_method):
             if method == "vina":
                 for binding_site_method in cfg.vina_binding_site_methods:
                     method_predictions = get_method_predictions(
@@ -1349,11 +1352,16 @@ def generate_ensemble_predictions(
                         cfg,
                         binding_site_method=binding_site_method,
                         input_protein_filepath=protein_filepath,
+                        is_ss_method=is_ss_method,
                     )
                     ensemble_predictions_dict[f"vina_{binding_site_method}"] = method_predictions
             else:
                 method_predictions = get_method_predictions(
-                    method, input_id, cfg, input_protein_filepath=protein_filepath
+                    method,
+                    input_id,
+                    cfg,
+                    input_protein_filepath=protein_filepath,
+                    is_ss_method=is_ss_method,
                 )
                 ensemble_predictions_dict[method] = method_predictions
 
@@ -1570,11 +1578,16 @@ def rank_ensemble_predictions(
                 and apo_reference_protein_filepath is not None
             ):
                 try:
-                    align_complex_to_protein_only(
+                    alignment_return_code = align_complex_to_protein_only(
                         protein_filepath, ligand_filepath, apo_reference_protein_filepath
                     )
-                    protein_filepath = protein_filepath.replace(".pdb", "_aligned.pdb")
-                    ligand_filepath = ligand_filepath.replace(".sdf", "_aligned.sdf")
+                    if alignment_return_code == 0:
+                        protein_filepath = protein_filepath.replace(".pdb", "_aligned.pdb")
+                        ligand_filepath = ligand_filepath.replace(".sdf", "_aligned.sdf")
+                    else:
+                        logger.warning(
+                            f"Failed to align predicted complex structure {protein_filepath} and ligand structure {ligand_filepath} to the apo protein structure {apo_reference_protein_filepath} from method {method}. Skipping alignment..."
+                        )
                 except Exception as e:
                     logger.warning(
                         f"Failed to align protein-ligand complex {protein_filepath} and {ligand_filepath} to apo protein structure {apo_reference_protein_filepath}. Skipping alignment due to: {e}"
@@ -2247,6 +2260,7 @@ def main(cfg: DictConfig):
     with open_dict(cfg):
         # NOTE: besides their output directories, single-sequence baselines are treated like their multi-sequence counterparts
         output_dir = copy.deepcopy(cfg.output_dir)
+        cfg.is_ss_ensemble_method = [s.endswith("_ss") for s in cfg.ensemble_methods]
         cfg.ensemble_methods = [s.removesuffix("_ss") for s in cfg.ensemble_methods]
         cfg.output_dir = output_dir