diff --git a/harbor/analysis/cross_docking.py b/harbor/analysis/cross_docking.py index df4acd3..6fa93da 100644 --- a/harbor/analysis/cross_docking.py +++ b/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,6 +2200,15 @@ 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]: @@ -2161,6 +2216,7 @@ def create_success_rate_evaluator(self) -> [BinaryEvaluation]: 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, ) ] diff --git a/harbor/pli/calculate_plip_interactions.py b/harbor/pli/calculate_plip_interactions.py index ff9fbde..a4c166a 100644 --- a/harbor/pli/calculate_plip_interactions.py +++ b/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() diff --git a/harbor/pli/plip_analysis_schema.py b/harbor/pli/plip_analysis_schema.py index 3b242ca..93a9d0f 100644 --- a/harbor/pli/plip_analysis_schema.py +++ b/harbor/pli/plip_analysis_schema.py @@ -194,7 +194,7 @@ def plip_constructor(plip: PLInteraction) -> list[ProteinLigandInteraction]: output_interaction_type = InteractionType.HydrophobicInteraction elif interaction_type == "pication": - output_interaction_type = InteractionType.PiStacking + output_interaction_type = InteractionType.PiCation protein_charge = ( FormalCharge.Positive if interaction.protcharged