self_consistency.py

"""AF2 self-consistency evaluation."""

from pathlib import Path

import pandas as pd
from caliby import CalibyModel

from file_utils import _write_zip_from_dir


def _run_self_consistency(
    model: CalibyModel,
    df: pd.DataFrame,
    out_pdb_list: list[str],
    out_dir: Path,
    download_stem: str,
) -> tuple[str, dict[str, str]]:
    from caliby.eval.eval_utils.folding_utils import clear_mem_torch

    clear_mem_torch()

    sc_out_dir = out_dir / "self_consistency"
    id_to_metrics = model.self_consistency_eval(out_pdb_list, out_dir=str(sc_out_dir))

    for metric in ["sc_ca_rmsd", "avg_ca_plddt", "tmalign_score"]:
        df[metric] = [id_to_metrics.get(Path(path).stem, {}).get(metric, float("nan")) for path in out_pdb_list]

    af2_pdb_data = {}
    for path in out_pdb_list:
        af2_path = sc_out_dir / "struct_preds" / f"af2_{Path(path).stem}.pdb"
        if af2_path.exists():
            af2_pdb_data[Path(path).stem] = af2_path.read_text()

    sc_zip_path = _write_zip_from_dir(sc_out_dir, download_stem, "_self_consistency.zip")
    return sc_zip_path, af2_pdb_data