Update plip analysis for paper

harbor/analysis/cross_docking.py

@@ -1,4 +1,6 @@
 import itertools
+import logging
+import time
 from pydantic import BaseModel, Field, model_validator, field_validator, ConfigDict
 from typing_extensions import Self
 import abc
@@ -14,6 +16,8 @@
 from operator import eq, gt, lt, ge, le, ne
 from pydantic import confloat
 
+logger = logging.getLogger(__name__)
+
 
 class Operator(StrEnum):
     EQ = "eq"
@@ -381,7 +385,7 @@ def to_models(self) -> list["DataFrameModel"]:
             model_dataframe = self.dataframe.groupby(relevant_columns)[
                 relevant_columns
             ].head(1)
-            print(model_name, model_type, relevant_columns)
+            logger.debug(f"{model_name} {model_type} {relevant_columns}")
 
             # rename columns
             param_columns = [
@@ -1216,6 +1220,14 @@ class RMSDScorer(Scorer):
     number_to_return: int = 1
 
 
+class PLIFScorer(Scorer):
+    type_: str = "PLIFScorer"
+    name: str = "PLIF_Recall"
+    variable: str = "PLIFData_plif_tversky_recall"
+    ascending: bool = False
+    number_to_return: int = 1
+
+
 class SuccessRate(ModelBase):
     name: str = "SuccessRate"
     type_: str = "SuccessRate"
@@ -1359,6 +1371,8 @@ def get_class_from_name(name: str):
             return RMSDScorer
         case "POSITScorer":
             return POSITScorer
+        case "PLIFScorer":
+            return PLIFScorer
         case "PoseSelector":
             return PoseSelector
         case "FractionGood":
@@ -1381,10 +1395,26 @@ def _bootstrap_worker(args):
         result = evaluator_copy.process_single_bootstrap(pose_selected_data)
         return bootstrap_idx, result
     except Exception as e:
-        print(f"Error processing bootstrap {bootstrap_idx}: {e}")
+        logger.error(f"Error processing bootstrap {bootstrap_idx}: {e}")
         return bootstrap_idx, None
 
 
+def _bootstrap_chunk_worker(args):
+    """Process a chunk of bootstraps in one worker call to amortize pickle overhead."""
+    indices, evaluator_json, pose_selected_data = args
+    evaluator_data = json.loads(evaluator_json)
+    evaluator_copy = get_class_from_name(evaluator_data["type_"])(**evaluator_data)
+    results = []
+    for idx in indices:
+        try:
+            result = evaluator_copy.process_single_bootstrap(pose_selected_data)
+            results.append((idx, result))
+        except Exception as e:
+            logger.error(f"Error processing bootstrap {idx}: {e}")
+            results.append((idx, None))
+    return results
+
+
 class Evaluator(ModelBase):
     name: str = "Evaluator"
     type_: str = "Evaluator"
@@ -1489,40 +1519,43 @@ def run(self, data: DockingDataModel, n_cpus: int = 1) -> SuccessRate:
 
         if n_cpus == 1:
             # Sequential processing
+            t_bootstrap_total = 0.0
             for bootstrap_idx in range(self.n_bootstraps):
                 try:
+                    t0 = time.perf_counter()
                     result = self.process_single_bootstrap(pose_selected_data)
+                    elapsed = time.perf_counter() - t0
+                    t_bootstrap_total += elapsed
                     if result is not None:
                         all_results.append(result)
                 except Exception as e:
-                    print(f"Error processing bootstrap {bootstrap_idx}: {e}")
+                    logger.error(f"Error processing bootstrap {bootstrap_idx}: {e}")
                     continue
+            logger.info(f"Total time: {t_bootstrap_total:.1f}s  per bootstrap: {t_bootstrap_total/self.n_bootstraps:.3f}s")
         else:
             # Parallel processing
             from concurrent.futures import ProcessPoolExecutor, as_completed
             import multiprocessing as mp
 
             n_cpus = min(n_cpus, mp.cpu_count())
-            print(
-                f"Running {self.n_bootstraps} bootstraps in parallel using {n_cpus} CPUs."
-            )
+            logger.info(f"Running {self.n_bootstraps} bootstraps in parallel using {n_cpus} CPUs.")
 
-            # Create worker arguments
-            worker_args = [
-                (bootstrap_idx, self.to_json_str(), pose_selected_data)
-                for bootstrap_idx in range(self.n_bootstraps)
+            # Chunk bootstraps so pose_selected_data is pickled once per worker, not once per bootstrap
+            indices = list(range(self.n_bootstraps))
+            chunk_size = max(1, len(indices) // n_cpus)
+            chunks = [
+                indices[i : i + chunk_size] for i in range(0, len(indices), chunk_size)
+            ]
+            evaluator_json = self.to_json_str()
+            chunk_args = [
+                (chunk, evaluator_json, pose_selected_data) for chunk in chunks
             ]
 
             with ProcessPoolExecutor(max_workers=n_cpus) as executor:
-                future_to_idx = {
-                    executor.submit(_bootstrap_worker, args): args[0]
-                    for args in worker_args
-                }
-
-                for future in as_completed(future_to_idx):
-                    bootstrap_idx, result = future.result()
-                    if result is not None:
-                        all_results.append(result)
+                for chunk_results in executor.map(_bootstrap_chunk_worker, chunk_args):
+                    for bootstrap_idx, result in chunk_results:
+                        if result is not None:
+                            all_results.append(result)
 
         return SuccessRate.from_replicates(all_results)
 
@@ -1606,11 +1639,13 @@ def calculate_result(
     def calculate_results(
         cls, data: DockingDataModel, evaluators: list[Evaluator], n_cpus: int = 1
     ) -> list["Results"]:
-        data_copies = [data.__deepcopy__() for ev in evaluators]
         results = []
-        for data, ev in tqdm(zip(data_copies, evaluators), total=len(evaluators)):
-            result = ev.run(data, n_cpus=n_cpus)
+        n = len(evaluators)
+        for i, ev in enumerate(evaluators):
+            logger.info(f"Evaluator {i+1}/{n}: {ev.name}")
+            result = ev.run(data.__deepcopy__(), n_cpus=n_cpus)
             results.append(cls(evaluator=ev, success_rate=result))
+            logger.info(f"Evaluator {i+1}/{n} done.")
         return results
 
     @classmethod
@@ -1829,13 +1864,20 @@ class RMSDScorerSettings(EvaluatorSettingsBase):
     rmsd_name: str = Field("RMSD", description="Name of the RMSD score")
 
 
+class PLIFScorerSettings(EvaluatorSettingsBase):
+    use: bool = False
+    plif_column_name: str = "PLIFData_plif_tversky_recall"
+    plif_name: str = Field("PLIF_Recall", description="Name of the PLIF recall score")
+
+
 class ScorerSettings(CompositSettingsBase):
     use: bool = True
     rmsd_scorer_settings: RMSDScorerSettings = RMSDScorerSettings()
     posit_scorer_settings: POSITScorerSettings = POSITScorerSettings()
+    plif_scorer_settings: PLIFScorerSettings = PLIFScorerSettings()
 
     def get_component_settings(self) -> list[EvaluatorSettingsBase]:
-        return [self.rmsd_scorer_settings, self.posit_scorer_settings]
+        return [self.rmsd_scorer_settings, self.posit_scorer_settings, self.plif_scorer_settings]
 
 
 class SuccessRateSettings(EvaluatorSettingsBase):
@@ -1844,6 +1886,10 @@ class SuccessRateSettings(EvaluatorSettingsBase):
     rmsd_cutoff: float = Field(
         2.0, description="RMSD cutoff to label the resulting poses as successful"
     )
+    below_cutoff_is_good: bool = Field(
+        True,
+        description="Whether values below (True) or above (False) the cutoff are successes",
+    )
 
 
 def generate_logarithmic_scale(n_max: int, base: int = 10) -> list[int]:
@@ -2154,13 +2200,23 @@ def create_scorers(self) -> list[Scorer]:
                 )
             )
 
+        if settings.plif_scorer_settings.use:
+            plif_settings = settings.plif_scorer_settings
+            scorers.append(
+                PLIFScorer(
+                    name=plif_settings.plif_name,
+                    variable=plif_settings.plif_column_name,
+                )
+            )
+
         return scorers
 
     def create_success_rate_evaluator(self) -> [BinaryEvaluation]:
         return [
             BinaryEvaluation(
                 variable=self.success_rate_evaluator_settings.success_rate_column,
                 cutoff=self.success_rate_evaluator_settings.rmsd_cutoff,
+                below_cutoff_is_good=self.success_rate_evaluator_settings.below_cutoff_is_good,
             )
         ]
 

harbor/pli/calculate_plip_interactions.py

@@ -1,3 +1,15 @@
+"""
+Expected Usage:
+python calculate_plip_interactions.py --yaml-input input.yaml --output-dir output_directory --ncpus 4
+
+Where `input.yaml` contains a mapping of names to directories containing PDB files, and `output_directory` is where the interaction CSV files will be saved.
+
+i.e. input.yaml:
+----------------
+crystal: 20250313_plip_analysis/crystal
+docked: 20250313_plip_analysis/docked
+"""
+
 from harbor.pli.plip_analysis_schema import PLIntReport
 from pathlib import Path
 import click
@@ -13,7 +25,9 @@ class ProcessingError(Exception):
     pass
 
 
-def analyze_structure(structure: Path, name: str, output_dir: Path) -> Path:
+def analyze_structure(
+    structure: Path, name: str, output_dir: Path, ligand_id: str
+) -> Path:
     """
     Analyze a single structure using PLIP.
 
@@ -40,6 +54,7 @@ def analyze_structure(structure: Path, name: str, output_dir: Path) -> Path:
         outpath = output_dir / f"{name}_{structure.stem}_interactions.csv"
         interactions = PLIntReport.from_complex_path(
             complex_path=structure,
+            ligand_id=ligand_id,
         )
         interactions.to_csv(outpath)
         click.echo(f"Saved interactions to {outpath}")
@@ -51,7 +66,7 @@ def analyze_structure(structure: Path, name: str, output_dir: Path) -> Path:
 
 
 def process_structure_batch(
-    structures: list[Path], name: str, output_dir: Path, ncpus: int
+    structures: list[Path], name: str, output_dir: Path, ncpus: int, ligand_id: str
 ) -> tuple[list[Path], list[str]]:
     """
     Process a batch of structures in parallel.
@@ -68,6 +83,7 @@ def process_structure_batch(
         analyze_structure,
         name=name,
         output_dir=output_dir,
+        ligand_id=ligand_id,
     )
 
     with ProcessPool(max_workers=ncpus) as pool:
@@ -90,71 +106,34 @@ def process_structure_batch(
     return successful_outputs, errors
 
 
-@click.command()
-@click.option(
-    "--pdb-dir",
-    type=click.Path(exists=True, path_type=Path),
-    help="Path to directory containing PDB files",
-    required=False,
-)
-@click.option(
-    "--yaml-input",
-    type=click.Path(exists=True, path_type=Path),
-    help="Path to input yaml file containing name: path pairs",
-    required=False,
-)
-@click.option(
-    "--output-dir",
-    type=click.Path(path_type=Path),
-    default=Path("./"),
-    help="Path to output directory",
-    required=False,
-)
-@click.option(
-    "--ncpus", type=int, default=1, help="Number of cpus to use for parallel processing"
-)
-@click.option(
-    "--error-log",
-    type=click.Path(path_type=Path),
-    help="Path to error log file",
-    default="plip_errors.log",
-)
-def main(
-    pdb_dir: Path, yaml_input: Path, output_dir: Path, ncpus: int, error_log: Path
+def calculate_plip(
+    yaml_input: Path, output_dir: Path, ncpus: int, ligand_id: str, error_log: Path
 ):
-    """
-    Get PLIP interactions
-
-    Basic usage, which create a csv file of the calculated interactions for all the pdb files in this directory:
-    harbor calculate-plip-interactions --pdb-dir directory_with_pdb_files
-
-    For more complex usage, you can provide a YAML file that maps names to directories containing PDB files:
-    harbor calculate-plip-interactions --yaml-input input.yaml --output-dir output_directory --ncpus 4
-
-    Where `input.yaml` contains a mapping of names to directories containing PDB files, and `output_directory` is where the interaction CSV files will be saved.
+    """Main function to calculate PLIP interactions from complexes in folder, indicated in a YAML input file.
 
-    i.e. input.yaml:
-    ----------------
-    crystal: 20250313_plip_analysis/crystal
-    docked: 20250313_plip_analysis/docked
+    Parameters
+    ----------
+    yaml_input : Path
+        Path to the input YAML file containing structure info.
+    output_dir : Path
+        Directory where the interaction CSV files will be saved.
+    ncpus : int
+        Number of CPUs to use for parallel processing.
+    error_log : Path
+        Path to the error log file.
+    ligand_id : str
+        Residue name for the ligand.
     """
     output_dir.mkdir(exist_ok=True)
+    all_errors = []
 
-    if not yaml_input and not pdb_dir:
-        click.echo("Please provide either --pdb-dir or --yaml-input", err=True)
+    try:
+        with open(yaml_input, "r") as f:
+            input_dict = yaml.safe_load(f)
+    except yaml.YAMLError as e:
+        click.echo(f"Error reading YAML file: {e}", err=True)
         raise click.Abort()
 
-    all_errors = []
-    if pdb_dir:
-        input_dict = {"default": pdb_dir}
-    elif yaml_input:
-        try:
-            with open(yaml_input, "r") as f:
-                input_dict = yaml.safe_load(f)
-        except yaml.YAMLError as e:
-            click.echo(f"Error reading YAML file: {e}", err=True)
-            raise click.Abort()
-
     for name, structure_dir in input_dict.items():
         structure_dir = Path(structure_dir)
         if not structure_dir.exists():
@@ -174,7 +153,7 @@ def main(
 
         click.echo(f"Analyzing {len(structures)} structures")
         successful, errors = process_structure_batch(
-            structures, name, output_dir, ncpus
+            structures, name, output_dir, ncpus, ligand_id
         )
 
         if errors:
@@ -194,5 +173,33 @@ def main(
         raise click.Abort()
 
 
+@click.command()
+@click.option(
+    "--yaml-input",
+    type=click.Path(exists=True, path_type=Path),
+    help="Path to input yaml file containing name: path pairs",
+    required=True,
+)
+@click.option(
+    "--output-dir",
+    type=click.Path(path_type=Path),
+    help="Path to output directory",
+    required=True,
+)
+@click.option(
+    "--ncpus", type=int, default=1, help="Number of cpus to use for parallel processing"
+)
+@click.option(
+    "--error-log",
+    type=click.Path(path_type=Path),
+    help="Path to error log file",
+    default="plip_errors.log",
+)
+@click.option("--ligand-id", type=str, help="Residue name of the ligand", default="UNK")
+def main(yaml_input, output_dir, ncpus, error_log, ligand_id):
+    """Get PLIP interactions"""
+    calculate_plip(yaml_input, output_dir, ncpus, ligand_id, error_log)
+
+
 if __name__ == "__main__":
     main()