model/PXDesignBench/pxdbench/tools/af2/main_af2_complex.py

# Copyright 2025 ByteDance and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

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 align_and_calculate_rmsd
from pxdbench.permutation import permute_generated_min_complex_rmsd
from pxdbench.tools.af2.af2_utils import add_cyclic_offset, renumber_by_rebuilding
from pxdbench.utils import concat_dict_values, seed_everything

logger = logging.getLogger(__name__)


def predict_binder_structure(
    prediction_model,
    sequence: str,
    design_name: str,
    ori_design_pdb: str,
    model_indices: list[int],
    save_dir: str,
    design_chain_layout: str,
):
    """
    Predict binder structure using AlphaFold2 and compute structural metrics.

    Args:
        prediction_model: Initialized ColabDesign AFDesign model instance.
        sequence (str): Amino acid sequence of the binder to predict.
        design_name (str): Unique identifier for the design (e.g., "pdbname_seq0").
        ori_design_pdb: Path to designed pdb.
        model_indices (list[int]): List of AlphaFold2 model indices to use (0-4).
        save_dir (str): Directory to save predicted PDB files and metrics.
        design_chain_layout (str): "cond_first" or "cond_last".

    Returns:
        dict: Prediction statistics (pLDDT, pTM, i_pTM, etc.) for each model index.
    """
    sequence = re.sub(r"[^A-Z]", "", sequence.upper())
    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."
            )
            # load stats
            with open(output_stats_json, "r") as f:
                stats = json.load(f)
            print(f"Loaded {output_stats_json}.")

        else:
            prediction_model.predict(
                seq=sequence, models=[model_num], num_recycles=3, verbose=True
            )
            metrics = copy_dict(prediction_model.aux["log"])
            stats = {
                "pLDDT": round(metrics["plddt"], 2),
                "pTM": round(metrics["ptm"], 2),
                "i_pTM": round(metrics["i_ptm"], 2),
                "pAE": round(metrics["pae"], 2),
                "i_pAE": round(metrics["i_pae"], 2),  # i_pae divdied by 31
                "unscaled_i_pAE": round(metrics["i_pae"] * 31, 2),  # raw i_pae
            }
            # save pdb and stats
            prediction_model.save_pdb(output_pdb)
            # renumber
            renumber_by_rebuilding(
                ori_design_pdb, output_pdb, output_pdb, ref_layout=design_chain_layout
            )
            permute_generated_min_complex_rmsd(output_pdb, ori_design_pdb, output_pdb)
            with open(output_stats_json, "w") as f:
                json.dump(stats, f)

        prediction_stats[model_num] = stats

    return prediction_stats


def complex_prediction(
    input_dir: str,
    save_dir: str,
    design_pdb_dir: str,
    data_list: list[dict],
    cond_chain: str,
    binder_chain: str,
    af2_cfg,
    verbose=True,
    is_cyclic=False,
):
    """
    Run batch prediction for binder complexes using AlphaFold2.

    Args:
        input_dir (str): Directory containing input PDB files for target structures.
        save_dir (str): Directory to save prediction outputs (PDBs, metrics).
        design_pdb_dir (str): Directory to save designed pdbs.
        data_list (list[dict]): List of design data with keys "name", "sequence", "seq_idx".
        cond_chain (str): Chain ID(s) of the target (conditioning) structure(s).
        binder_chain (str): Chain ID of the binder to design/predict.
        af2_cfg (dict): AlphaFold2 configuration (model indices, multimer usage, etc.).
        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 prediction statistics for each design in data_list.
    """
    use_binder_template = af2_cfg["use_binder_template"]
    logger.info(f"Input use_binder_template: {use_binder_template}")

    clear_mem()
    prediction_model = mk_afdesign_model(
        protocol="binder",
        num_recycles=3,
        data_dir=AF2_PARAMS_PATH,
        use_multimer=af2_cfg["use_multimer"],
        use_initial_guess=af2_cfg["use_initial_guess"],
        use_initial_atom_pos=af2_cfg["use_initial_atom_pos"],
    )

    os.makedirs(save_dir, exist_ok=True)

    results = []
    for item in data_list:
        name = item["name"]
        seq = item["sequence"]
        seq_idx = item["seq_idx"]
        pdb_file = os.path.join(input_dir, f"{name}.pdb")
        if not os.path.exists(pdb_file):
            print(f"ERROR: {pdb_file} not found")
            continue

        prediction_model.prep_inputs(
            pdb_filename=pdb_file,
            chain=cond_chain,
            binder_chain=binder_chain,
            use_binder_template=use_binder_template,
            rm_target_seq=True,
            rm_target_sc=False,
            rm_template_ic=True,
        )
        if is_cyclic:
            add_cyclic_offset(prediction_model)

        design_name = f"{name}_seq{seq_idx}"
        ori_design_pdb = os.path.join(design_pdb_dir, name + ".pdb")
        stats = predict_binder_structure(
            prediction_model,
            seq,
            design_name,
            ori_design_pdb,
            af2_cfg["model_ids"],
            save_dir,
            design_chain_layout="cond_last" if "A" in binder_chain else "cond_first",
        )
        stat_list = []
        for model_id in af2_cfg["model_ids"]:
            s = stats[model_id]

            # compute predict-design RMSD
            pred_complex_pdb = os.path.join(
                save_dir, f"{design_name}_model{model_id + 1}.pdb"
            )
            if os.path.isfile(ori_design_pdb):
                complex_rmsd = align_and_calculate_rmsd(
                    pred_complex_pdb, ori_design_pdb
                )
                if complex_rmsd is not None:
                    complex_rmsd = round(complex_rmsd, 2)
            else:
                complex_rmsd = None
            s["af2_complex_pred_design_rmsd"] = complex_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 Complex 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)

    # args = parser.parse_args()
    # model_ids = [int(x) for x in args.model_ids.split(",")]

    if args.seed is not None:
        seed_everything(args.seed, deterministic=False)

    try:
        results = complex_prediction(
            input_dir=input_data["input_dir"],
            save_dir=input_data["save_dir"],
            design_pdb_dir=input_data["design_pdb_dir"],
            data_list=input_data["data_list"],
            cond_chain=input_data["cond_chain"],
            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 complex prediction!")

    except Exception as e:
        print(f"Error: {str(e)}")
        import traceback

        traceback.print_exc()
        exit(1)


if __name__ == "__main__":
    main()