|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| import json
|
| import logging
|
| import os
|
| from abc import ABC, abstractmethod
|
| from functools import cached_property
|
| from typing import Optional
|
|
|
| import numpy as np
|
| import pandas as pd
|
|
|
| from pxdbench.metrics import diversity, secondary
|
| from pxdbench.tools.af2.af2_predictor import AF2ComplexPredictor, AF2MonomerPredictor
|
| from pxdbench.tools.ptx.interface import ProtenixAPI
|
| from pxdbench.tools.registry import get_backend
|
|
|
| logger = logging.getLogger(__name__)
|
|
|
|
|
| class BaseTask(ABC):
|
| task_type: str
|
| task_name: str
|
| backend_spec: Optional[str] = None
|
|
|
| def __init__(self, input_data, cfg, device_id: int, seed: int):
|
| self.cfg = cfg
|
| self.device_id = device_id
|
| self.seed = seed
|
| assert "pdb_dir" in input_data
|
| assert "pdb_names" in input_data
|
| self.pdb_dir = input_data["pdb_dir"]
|
| self.pdb_names = input_data["pdb_names"]
|
| self.out_dir = input_data.get("out_dir", os.path.dirname(self.pdb_dir))
|
| self.sample_fn = input_data.get("sample_fn", "sample_level_output.csv")
|
| self.summary_fn = input_data.get("summary_fn", "summary_output.json")
|
|
|
|
|
| self.num_seqs = cfg.get("num_seqs", 4)
|
| self.use_gt_seq = cfg.get("use_gt_seq", False)
|
| self._ptx_mini_inst: Optional[ProtenixAPI] = None
|
| self._ptx_inst: Optional[ProtenixAPI] = None
|
|
|
|
|
| self.process_pdb_paths()
|
|
|
| def get_device(self):
|
| if self.device_id >= 0:
|
| return f"cuda:{self.device_id}"
|
| else:
|
| return "cpu"
|
|
|
| @cached_property
|
| def ptx_factory(self):
|
| return get_backend(self.backend_spec)
|
|
|
| def get_ptx(self, is_large: bool = False) -> ProtenixAPI:
|
| if is_large:
|
| if self._ptx_inst is None:
|
| self._ptx_inst = self.ptx_factory(
|
| cfg=self.cfg.tools.ptx,
|
| device=self.get_device(),
|
| )
|
| return self._ptx_inst
|
|
|
| if self._ptx_mini_inst is None:
|
| self._ptx_mini_inst = self.ptx_factory(
|
| cfg=self.cfg.tools.ptx_mini,
|
| device=self.get_device(),
|
| )
|
| return self._ptx_mini_inst
|
|
|
| @abstractmethod
|
| def design_sequence(self):
|
| pass
|
|
|
| @abstractmethod
|
| def run(self):
|
| """Execute the task"""
|
| pass
|
|
|
| @staticmethod
|
| def summary_from_df(
|
| sample_df: pd.DataFrame,
|
| exclude_keys=["name", "seq_idx", "sequence"],
|
| other_metrics={},
|
| ):
|
| """
|
| Compute summary metrics from a DataFrame.
|
|
|
| Args:
|
| sample_df (pd.DataFrame): Input DataFrame containing sample data.
|
| exclude_keys (list, optional): Columns to exclude from summary. Defaults to ["name", "seq_idx", "sequence"].
|
| other_metrics (dict, optional): Additional metrics to include. Defaults to {}.
|
|
|
| Returns:
|
| dict: A dictionary containing computed summary metrics.
|
| """
|
| metrics = {}
|
| for col in sample_df.columns:
|
| if col in exclude_keys:
|
| continue
|
| col_data = sample_df[col].dropna()
|
|
|
|
|
| if pd.api.types.is_numeric_dtype(sample_df[col]):
|
| values = col_data.values
|
| metrics[f"{col}.avg"] = float(np.mean(values))
|
|
|
|
|
|
|
| elif pd.api.types.is_object_dtype(col_data):
|
| if all(isinstance(x, list) for x in col_data):
|
| try:
|
|
|
| per_sample_mean = col_data.apply(lambda x: np.mean(x))
|
| metrics[f"{col}.avg"] = float(np.mean(per_sample_mean))
|
|
|
| except:
|
| pass
|
|
|
| if "_success" in col:
|
| metrics[f"{col}.count"] = int(np.sum(sample_df[col].astype(bool)))
|
|
|
| metrics.update(other_metrics)
|
| return metrics
|
|
|
| @staticmethod
|
| def compute_success_rate(filters_cfg, metrics: pd.DataFrame) -> pd.DataFrame:
|
| """
|
| Compute success rate for each filter based on metrics.
|
|
|
| Args:
|
| filters_cfg (dict): Configuration for filters.
|
| metrics (pd.DataFrame): DataFrame containing metrics.
|
|
|
| Returns:
|
| pd.DataFrame: Updated DataFrame with success rate metrics.
|
| """
|
| for filter_name, filter_details in filters_cfg.items():
|
| missing = [k for k in filter_details.keys() if k not in metrics.columns]
|
|
|
| def row_success(row):
|
| for metric_name, (sym, thres) in filter_details.items():
|
| if metric_name not in row:
|
| continue
|
| value = row[metric_name]
|
| if value is None:
|
| return None
|
| if isinstance(value, list):
|
|
|
| value = min(value) if sym == "<" else max(value)
|
| if sym == "<" and value >= thres:
|
| return 0
|
| if sym == ">" and value <= thres:
|
| return 0
|
| return 1
|
|
|
| if missing:
|
| print(
|
| f"Missing columns {missing} for filter '{filter_name}'. Available columns: {list(metrics.columns)}"
|
| )
|
| metrics[f"{filter_name}_success"] = None
|
| metrics[f"{filter_name}_success_ignore_missing"] = metrics.apply(
|
| row_success, axis=1
|
| )
|
| else:
|
| metrics[f"{filter_name}_success"] = metrics.apply(row_success, axis=1)
|
| metrics[f"{filter_name}_success_ignore_missing"] = metrics[
|
| f"{filter_name}_success"
|
| ]
|
| return metrics
|
|
|
| def process_pdb_paths(self):
|
| """
|
| Validate and filter PDB file paths based on their existence.
|
|
|
| This method checks if each PDB file specified in self.pdb_names exists in the
|
| directory specified by self.pdb_dir. It filters out any PDB names that don't
|
| correspond to existing files and updates self.pdb_names to only contain valid names.
|
|
|
| Logs a warning message for each PDB file that is not found and skipped.
|
| """
|
|
|
| valid_pdb_names = []
|
| for name in self.pdb_names:
|
| pdb_path = os.path.join(self.pdb_dir, name + ".pdb")
|
|
|
| if not os.path.exists(pdb_path):
|
| logger.warning(
|
| f"pdb_path {pdb_path} does not exist. Will skip this file."
|
| )
|
| continue
|
| valid_pdb_names.append(name)
|
| self.pdb_names = list(valid_pdb_names)
|
| return
|
|
|
| def cal_diversity(self, pdb_names=None, binder_chain=None):
|
| """
|
| Calculate diversity of PDB structures. May be slow.
|
|
|
| Args:
|
| pdb_names (list, optional): List of PDB names to consider. Defaults to None.
|
| binder_chain (str, optional): Chain ID of the binder. Defaults to None.
|
|
|
| Returns:
|
| float: Diversity value.
|
| """
|
| if self.eval_diversity:
|
| all_names = self.pdb_names if pdb_names is None else pdb_names
|
| pdb_paths = [
|
| os.path.join(self.pdb_dir, name + ".pdb") for name in all_names
|
| ]
|
| div = diversity.compute_diversity(pdb_paths, binder_chain)
|
| else:
|
| div = -1
|
| return div
|
|
|
| def cal_secondary(self, results, chain_id=None):
|
| """
|
| Calculate secondary structure metrics for PDB structures.
|
|
|
| Args:
|
| results (list): List of dictionaries containing PDB structure information.
|
| chain_id (str, optional): Chain ID of the binder. Defaults to None.
|
| """
|
| for item in results:
|
| pdb_path = os.path.join(self.pdb_dir, item["name"] + ".pdb")
|
| alpha, beta, loop = secondary.cacl_secondary_structure(pdb_path, chain_id)
|
| Rg, ref_ratio = secondary.get_chain_rg(pdb_path, chain_id)
|
| item.update(
|
| {
|
| "alpha": alpha,
|
| "beta": beta,
|
| "loop": loop,
|
| "Rg": Rg,
|
| "ref_ratio": ref_ratio,
|
| }
|
| )
|
|
|
| def af2_complex_predict(self, data_list, save_dir, verbose=True):
|
| """
|
| Run AF2 complex prediction.
|
|
|
| Args:
|
| data_list (list): List of data samples.
|
| save_dir (str): Directory to save predictions.
|
| verbose (bool, optional): Whether to print verbose output. Defaults to True.
|
| """
|
| assert self.task_type in ["binder"]
|
| predictor = AF2ComplexPredictor(
|
| self.cfg.tools.af2,
|
| device_id=self.device_id,
|
| verbose=verbose,
|
| seed=self.seed,
|
| )
|
| predictor.predict(
|
| input_dir=self.pdb_dir,
|
| save_dir=save_dir,
|
| design_pdb_dir=self.pdb_dir,
|
| data_list=data_list,
|
| cond_chain=",".join(self.cond_chains),
|
| binder_chain=",".join(self.binder_chains),
|
| )
|
|
|
| def af2_monomer_predict(self, data_list, save_dir, verbose=True):
|
| """
|
| Run AF2 monomer prediction.
|
|
|
| Args:
|
| data_list (list): List of data samples.
|
| save_dir (str): Directory to save predictions.
|
| verbose (bool, optional): Whether to print verbose output. Defaults to True.
|
| """
|
| assert self.task_type in ["binder", "ligand_binder"]
|
| predictor = AF2MonomerPredictor(
|
| self.cfg.tools.af2,
|
| device_id=self.device_id,
|
| verbose=verbose,
|
| seed=self.seed,
|
| )
|
| predictor.predict(
|
| save_dir=save_dir,
|
| design_pdb_dir=self.pdb_dir,
|
| data_list=data_list,
|
| binder_chain=self.binder_chains[0],
|
| )
|
|
|
| def protenix_predict(self, data_list, orig_seqs=None, is_large=False):
|
| """
|
| Run Protenix prediction.
|
|
|
| Args:
|
| data_list (list): List of data samples.
|
| is_large (bool, optional): Whether to use the large model. Defaults to False.
|
| """
|
| ptx_cfg = self.cfg.tools.ptx if is_large else self.cfg.tools.ptx_mini
|
| ptx_filter = self.get_ptx(is_large)
|
| dump_dir = os.path.join(
|
| self.out_dir, "ptx_pred" if is_large else "ptx_mini_pred"
|
| )
|
|
|
|
|
| binder_chain_idx = 0 if self.binder_chains[0] == "A" else None
|
| json_path = ptx_filter.prepare_json(
|
| self.pdb_dir,
|
| data_list,
|
| dump_dir=dump_dir,
|
| binder_chain_idx=binder_chain_idx,
|
| orig_seqs=orig_seqs,
|
| use_template=ptx_cfg.get("use_template", False),
|
| )
|
| pred_pdb_paths = ptx_filter.predict(
|
| input_json_path=json_path,
|
| design_pdb_dir=self.pdb_dir,
|
| data_list=data_list,
|
| dump_dir=dump_dir,
|
| seed=self.seed,
|
| N_sample=ptx_cfg.N_sample,
|
| N_step=ptx_cfg.N_step,
|
| step_scale_eta=ptx_cfg.step_scale_eta,
|
| gamma0=ptx_cfg.gamma0,
|
| N_cycle=ptx_cfg.N_cycle,
|
| binder_chain_idx=binder_chain_idx,
|
| use_msa=ptx_cfg.get("use_msa", True),
|
| suffix="_mini" if not is_large else "",
|
| )
|
| return pred_pdb_paths
|
|
|
