|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| import argparse
|
| import json
|
| import logging
|
| import os
|
| import re
|
|
|
| from colabdesign import clear_mem, mk_afdesign_model
|
| from colabdesign.shared.utils import copy_dict
|
|
|
| from pxdbench.globals import AF2_PARAMS_PATH
|
| from pxdbench.metrics.Kalign import Binder_align_and_calculate_rmsd
|
| from pxdbench.tools.af2.af2_utils import add_cyclic_offset
|
| from pxdbench.utils import concat_dict_values, seed_everything
|
|
|
| logger = logging.getLogger(__name__)
|
|
|
|
|
| def predict_binder_monomer(
|
| prediction_model,
|
| sequence: str,
|
| design_name: str,
|
| model_indices: list[int],
|
| save_dir: str,
|
| ):
|
| """
|
| Predict monomeric binder structure using AlphaFold2 and compute monomer-specific metrics.
|
|
|
| Args:
|
| prediction_model: Initialized ColabDesign AFDesign model instance (hallucination protocol).
|
| sequence (str): Amino acid sequence of the binder monomer.
|
| design_name (str): Unique identifier for the design (e.g., "pdbname_seq0_MONOMER_ONLY").
|
| model_indices (list[int]): List of AlphaFold2 model indices to use (0-4).
|
| save_dir (str): Directory to save predicted PDB files and metrics.
|
|
|
| Returns:
|
| dict: Prediction statistics (pLDDT, pTM, pAE) for each model index.
|
| """
|
| sequence = re.sub(r"[^A-Z]", "", sequence.upper())
|
| prediction_model.set_seq(sequence)
|
| prediction_stats = {}
|
|
|
| for model_num in model_indices:
|
| output_name = f"{design_name}_model{model_num + 1}"
|
| output_pdb = os.path.join(save_dir, f"{output_name}.pdb")
|
| output_stats_json = os.path.join(save_dir, f"{output_name}.json")
|
|
|
| if os.path.exists(output_pdb) and os.path.exists(output_stats_json):
|
| print(
|
| f"Found existing {output_pdb} and {output_stats_json}. Will load from them."
|
| )
|
|
|
| with open(output_stats_json, "r") as f:
|
| stats = json.load(f)
|
| print(f"Loaded {output_stats_json}.")
|
| else:
|
| prediction_model.predict(models=[model_num], num_recycles=3, verbose=True)
|
| metrics = copy_dict(prediction_model.aux["log"])
|
| stats = {
|
| "pLDDT_MONOMER": round(metrics["plddt"], 3),
|
| "pTM_MONOMER": round(metrics["ptm"], 3),
|
| "pAE_MONOMER": round(metrics["pae"], 3),
|
| }
|
|
|
| prediction_model.save_pdb(output_pdb)
|
| with open(output_stats_json, "w") as f:
|
| json.dump(stats, f)
|
| prediction_stats[model_num] = stats
|
|
|
| return prediction_stats
|
|
|
|
|
| def binder_only_prediction(
|
| save_dir,
|
| design_pdb_dir,
|
| data_list,
|
| af2_cfg,
|
| binder_chain="B",
|
| verbose=True,
|
| is_cyclic=False,
|
| ):
|
| """
|
| Run batch prediction for monomeric binders and compute comparative structural metrics.
|
|
|
| Workflow:
|
| 1. Initializes AlphaFold2 model for monomer hallucination.
|
| 2. Processes each design in data_list, optimizing model compilation by reusing length-specific inputs.
|
| 3. Predicts monomer structures and computes pLDDT, pTM, pAE.
|
| 4. Calculates RMSD between:
|
| - Predicted monomer and predicted complex (if complex PDB exists).
|
| - Predicted monomer and original design template (if template PDB exists).
|
| 5. Aggregates metrics across AlphaFold2 models.
|
|
|
| Args:
|
| save_dir (str): Directory to save monomer prediction outputs (PDBs, metrics).
|
| design_pdb_dir (str): Directory containing original design template PDBs.
|
| data_list (list[dict]): List of design data with keys "name", "sequence", "seq_idx".
|
| af2_cfg (dict): AlphaFold2 configuration (model indices, multimer usage, etc.).
|
| binder_chain (str, optional): Chain ID of the binder in original design templates. Defaults to "B".
|
| verbose (bool, optional): Whether to print progress. Defaults to True.
|
| is_cyclic (bool, optional): Whether the binder is cyclic (adds cyclic offset). Defaults to False.
|
|
|
| Returns:
|
| list[dict]: Aggregated monomer prediction metrics for each design in data_list.
|
| """
|
| clear_mem()
|
| prediction_model = mk_afdesign_model(
|
| protocol="hallucination",
|
| use_templates=False,
|
| initial_guess=False,
|
| use_initial_atom_pos=False,
|
| num_recycles=3,
|
| data_dir=AF2_PARAMS_PATH,
|
| use_multimer=af2_cfg["use_multimer"],
|
| )
|
|
|
| os.makedirs(save_dir, exist_ok=True)
|
|
|
| results = []
|
| length_prev = -1
|
| for item in data_list:
|
| name = item["name"]
|
| seq = item["sequence"]
|
| seq_idx = item["seq_idx"]
|
| binder_len = len(seq)
|
|
|
|
|
| if length_prev != binder_len:
|
| prediction_model.prep_inputs(
|
| length=binder_len,
|
| )
|
| length_prev = binder_len
|
| if is_cyclic:
|
| add_cyclic_offset(prediction_model)
|
| design_name = f"{name}_seq{seq_idx}_MONOMER_ONLY"
|
| design_complex_name = f"{name}_seq{seq_idx}"
|
| stats = predict_binder_monomer(
|
| prediction_model, seq, design_name, af2_cfg["model_ids"], save_dir
|
| )
|
|
|
| stat_list = []
|
| for model_id in af2_cfg["model_ids"]:
|
| s = stats[model_id]
|
| pred_binder_pdb = os.path.join(
|
| save_dir, f"{design_name}_model{model_id + 1}.pdb"
|
| )
|
| pred_complex_pdb = os.path.join(
|
| save_dir, f"{design_complex_name}_model{model_id + 1}.pdb"
|
| )
|
| if os.path.isfile(pred_complex_pdb):
|
| bound_unbound_RMSD = round(
|
| Binder_align_and_calculate_rmsd(
|
| pred_binder_pdb, pred_complex_pdb, binder_chain
|
| ),
|
| 2,
|
| )
|
| else:
|
| bound_unbound_RMSD = None
|
|
|
|
|
| ori_design_pdb = os.path.join(design_pdb_dir, name + ".pdb")
|
| if os.path.isfile(ori_design_pdb):
|
| binder_rmsd = round(
|
| Binder_align_and_calculate_rmsd(
|
| pred_binder_pdb, ori_design_pdb, binder_chain
|
| ),
|
| 2,
|
| )
|
| else:
|
| binder_rmsd = None
|
|
|
| s["bound_unbound_RMSD"] = bound_unbound_RMSD
|
| s["af2_binder_pred_design_rmsd"] = binder_rmsd
|
| stat_list.append(s)
|
| stat = concat_dict_values(stat_list)
|
| if verbose:
|
| print(f"{name}-seq{seq_idx}, {stat}")
|
| results.append(stat)
|
| return results
|
|
|
|
|
| def main():
|
| parser = argparse.ArgumentParser(description="AF2 Binder Monomer Prediction")
|
| parser.add_argument("--input", type=str, required=True, help="Input JSON file")
|
| parser.add_argument("--output", type=str, required=True, help="Output JSON file")
|
| parser.add_argument("--seed", type=int, default=None)
|
|
|
| args = parser.parse_args()
|
|
|
| with open(args.input, "r") as f:
|
| input_data = json.load(f)
|
|
|
|
|
|
|
|
|
| if args.seed is not None:
|
| seed_everything(args.seed, deterministic=False)
|
|
|
| try:
|
| results = binder_only_prediction(
|
| save_dir=input_data["save_dir"],
|
| design_pdb_dir=input_data["design_pdb_dir"],
|
| data_list=input_data["data_list"],
|
| binder_chain=input_data["binder_chain"],
|
| af2_cfg=input_data["af2_cfg"],
|
| verbose=True,
|
| is_cyclic=input_data["is_cyclic"],
|
| )
|
|
|
| with open(args.output, "w") as f:
|
| json.dump(results, f)
|
|
|
| print(f"Successfully completed AF2 binder only prediction!")
|
|
|
| except Exception as e:
|
| print(f"Error: {str(e)}")
|
| import traceback
|
|
|
| traceback.print_exc()
|
| exit(1)
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
