|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| import json
|
| import os
|
|
|
| import pandas as pd
|
|
|
| from pxdbench.tasks.base import BaseTask
|
| from pxdbench.tools.protmpnn.main_mpnn import get_gt_sequence
|
| from pxdbench.tools.protmpnn.mpnn_predictor import MPNNPredictor
|
| from pxdbench.utils import save_eval_results
|
|
|
| from .registry import register_task
|
|
|
|
|
| @register_task("binder")
|
| class BinderTask(BaseTask):
|
| def __init__(self, input_data, cfg, device_id: int, seed: int):
|
| """
|
| Initialize a BinderTask instance.
|
|
|
| Args:
|
| input_data (dict): Task input parameters including PDB paths and chain specifications.
|
| cfg (dict): Configuration dictionary with task settings.
|
| device_id (int): GPU device ID (-1 for CPU).
|
| seed (int): Random seed for reproducibility.
|
|
|
| Validates:
|
| - Exactly one binder chain is specified (multiple binder chains not supported).
|
| """
|
| self.task_type = "binder"
|
| self.task_name = input_data.get("name", "binder")
|
| assert "cond_chains" in input_data
|
| assert "binder_chains" in input_data
|
| self.cond_chains = input_data["cond_chains"]
|
| self.binder_chains = input_data["binder_chains"]
|
| self.pdb_name_to_binder_seq_list = input_data.get(
|
| "pdb_name_to_binder_seq_list", None
|
| )
|
| if input_data.get("orig_seqs_json", None) is not None:
|
| with open(input_data["orig_seqs_json"], "r") as f:
|
| self.orig_seqs = json.load(f)
|
| elif input_data.get("orig_seqs", None) is not None:
|
| self.orig_seqs = input_data["orig_seqs"]
|
| else:
|
| self.orig_seqs = None
|
|
|
|
|
| self.use_binder_seq_list = cfg.get("use_binder_seq_list", False)
|
| self.eval_diversity = cfg.get("eval_diversity", False)
|
| self.eval_binder_monomer = cfg.get("eval_binder_monomer", True)
|
| self.eval_complex = cfg.get("eval_complex", True)
|
| self.eval_protenix_mini = cfg.get("eval_protenix_mini", True)
|
| self.eval_protenix = cfg.get("eval_protenix", False)
|
|
|
|
|
| assert (
|
| len(self.binder_chains) == 1
|
| ), f"Get {len(self.binder_chains)} binder chains, but only 1 is allowed."
|
|
|
| super().__init__(input_data, cfg, device_id, seed)
|
|
|
| def prepare_data_from_seq_list(self):
|
| datas = []
|
| for name in self.pdb_names:
|
| binder_seq_list = self.pdb_name_to_binder_seq_list[name]
|
| for i, seq in enumerate(binder_seq_list):
|
| data = {"name": name, "seq_idx": i, "sequence": seq}
|
| datas.append(data)
|
| return datas
|
|
|
| def design_sequence(self, verbose=True):
|
| """
|
| Generates binder sequences based on task configuration.
|
|
|
| Supports three modes:
|
| 1. Use pre-provided sequence lists (self.use_binder_seq_list)
|
| 2. Use ground truth sequences from PDB files (self.use_gt_seq)
|
| 3. De novo design using MPNN (default)
|
|
|
| Args:
|
| verbose (bool, optional): Whether to print detailed progress. Defaults to True.
|
|
|
| Returns:
|
| list[dict]: List of design results with keys "name", "seq_idx", and "sequence".
|
| """
|
| if self.use_binder_seq_list:
|
| results = self.prepare_data_from_seq_list()
|
| elif self.use_gt_seq:
|
| results = get_gt_sequence(
|
| self.pdb_dir, self.pdb_names, self.binder_chains[0]
|
| )
|
| else:
|
| mpnn_predictor = MPNNPredictor(
|
| self.cfg.tools.mpnn,
|
| device_id=self.device_id,
|
| verbose=verbose,
|
| seed=self.seed,
|
| )
|
| results = mpnn_predictor.design_binder(
|
| self.pdb_dir,
|
| self.pdb_names,
|
| self.num_seqs,
|
| binder_chains=self.binder_chains,
|
| cond_chains=self.cond_chains,
|
| )
|
| return results
|
|
|
| def run(self):
|
| """
|
| Executes the complete binder design evaluation workflow.
|
|
|
| Workflow steps:
|
| 1. Designs sequences via design_sequence()
|
| 2. Runs structure predictions (AF2 complex/monomer, Protenix) based on config
|
| 3. Calculates secondary structure and diversity metrics
|
| 4. Saves sample-level results to CSV and summary metrics to JSON
|
|
|
| Returns:
|
| dict: Dictionary with task metadata and output file paths.
|
| """
|
| results = self.design_sequence()
|
| self.check_results(results)
|
| binder_chain = self.binder_chains[0]
|
|
|
| af2_pred_path = os.path.join(self.out_dir, "af2_pred")
|
| if self.eval_complex:
|
| self.af2_complex_predict(results, af2_pred_path)
|
|
|
| if self.eval_binder_monomer:
|
| self.af2_monomer_predict(results, af2_pred_path)
|
|
|
| if self.eval_protenix_mini:
|
| self.protenix_predict(results, orig_seqs=self.orig_seqs)
|
|
|
| if self.eval_protenix:
|
| self.protenix_predict(results, orig_seqs=self.orig_seqs, is_large=True)
|
|
|
| self.cal_secondary(results, binder_chain)
|
| div = self.cal_diversity()
|
| sample_df = pd.DataFrame(results)
|
| sample_df = sample_df.sort_values(by=["name", "seq_idx"])
|
| self.compute_success_rate(self.cfg.filters, sample_df)
|
| summary_dict = {"task": self.task_type, "name": self.task_name}
|
| summary_dict.update(
|
| self.summary_from_df(sample_df, other_metrics={"diversity": div})
|
| )
|
| sample_save_path, summary_save_path = save_eval_results(
|
| sample_df, summary_dict, self.out_dir, self.sample_fn, self.summary_fn
|
| )
|
| print(
|
| f"Eval done! Results are saved in {sample_save_path} and {summary_save_path}"
|
| )
|
| return {
|
| "task": self.task_type,
|
| "name": self.task_name,
|
| "sample_save_path": sample_save_path,
|
| "summary_save_path": summary_save_path,
|
| }
|
|
|
| def check_results(self, results):
|
| """
|
| Validates design results for consistency and correctness.
|
|
|
| Checks:
|
| 1. No duplicate entries (by structure name + sequence index)
|
| 2. Correct number of sequences per structure (when not using pre-provided lists)
|
|
|
| Args:
|
| results (list[dict]): List of design results from design_sequence()
|
|
|
| Raises:
|
| ValueError: If duplicates are found or sequence count is incorrect.
|
| """
|
| result_names = [
|
| result["name"] + f"_seq{result['seq_idx']}" for result in results
|
| ]
|
| if len(result_names) != len(set(result_names)):
|
| raise ValueError(f"Found duplicate names in results: {result_names}.")
|
| if self.use_binder_seq_list or self.use_gt_seq:
|
| pass
|
| elif len(result_names) != len(self.pdb_names) * self.num_seqs:
|
| raise ValueError(
|
| f"Found {len(result_names)} results, but {len(self.pdb_names)} pdb_names, each with {self.num_seqs} seqs are provided."
|
| )
|
| return
|
|
|
