Remove old eval/universal_v2_evaluate.py (renamed to main_benchmark)

eval/universal_v2_evaluate.py

@@ -1,341 +0,0 @@
-#!/usr/bin/env python3
-"""universal_v2 §9.1 evaluator.
-
-For each prediction PDB under <root>, computes:
-- mean_plddt_overall, mean_plddt_switch, mean_plddt_core
-- For each state (a, b): Cα RMSD common_core, TM-score, hit_primary (§9.1 binding 3-condition)
-
-Regions come from data/universal_v2/annotations/<case_id>/state_region_FINAL.tsv
-Reference PDBs come from data/universal_v2/structures/<case_id>/{state_a,state_b}.pdb
-
-ID-pattern cases (state_a == state_b): single hit column hit_primary_state_a.
-Per IDP_EVAL_OVERRIDE: treat the whole query region as the target region;
-RMSD ≤3Å AND mean_pLDDT ≥70 → hit_primary=True. switch_region pLDDT check is
-skipped because every residue is marked switch_region.
-
-Usage:
-    python scripts/universal_v2_evaluate.py --case <case_id> --root <pred_dir>
-"""
-from __future__ import annotations
-import argparse
-import csv
-import json
-import os
-import re
-import shutil
-import subprocess
-import sys
-from pathlib import Path
-
-import numpy as np
-import yaml
-from Bio.PDB import PDBParser, Superimposer
-
-# Benchmark data root: env override, else the bundled bench/universal_v2 that
-# ships alongside this script (eval/ and bench/ are siblings at the repo root).
-DATA_ROOT = Path(os.environ.get(
-    "SF_UV2_DATA", Path(__file__).resolve().parents[1] / "bench" / "universal_v2"))
-# TMalign is optional: it only fills the TM-score columns and is NOT part of the
-# hit_primary criterion. Resolved from PATH; None -> those columns are NaN.
-TMALIGN = shutil.which("TMalign")
-
-PRED_NAME_RE = re.compile(
-    r"^(?P<subset>.+?)_unrelaxed_rank_(?P<rank>\d+)_alphafold2_ptm_model_(?P<model>\d+)_seed_(?P<seed>\d+)\.pdb$"
-)
-
-_CASES_CACHE: dict[str, dict] | None = None
-
-
-def load_case(case_id: str) -> dict:
-    global _CASES_CACHE
-    if _CASES_CACHE is None:
-        cases = yaml.safe_load((DATA_ROOT / "cases.yaml").read_text())["cases"]
-        _CASES_CACHE = {c["case_id"]: c for c in cases}
-    return _CASES_CACHE[case_id]
-
-
-def load_regions(case_id: str) -> dict:
-    path = DATA_ROOT / "annotations" / case_id / "state_region_FINAL.tsv"
-    common_core: list[int] = []
-    switch_3a: list[int] = []
-    state_a_resnum: dict[int, int] = {}
-    with path.open() as f:
-        for line in f:
-            if line.startswith("#"):
-                continue
-            line = line.rstrip("\n")
-            if not line:
-                continue
-            parts = line.split("\t")
-            if parts[0] == "residue_index_query":
-                continue
-            qi = int(parts[0])
-            try:
-                resnum_a = int(parts[1])
-            except ValueError:
-                resnum_a = qi
-            cc = int(parts[3])
-            sr = int(parts[4])
-            if cc:
-                common_core.append(qi)
-            if sr:
-                switch_3a.append(qi)
-            state_a_resnum[qi] = resnum_a
-    return {"common_core": common_core, "switch_3A": switch_3a,
-            "state_a_resnum": state_a_resnum}
-
-
-def is_id_case(case_id: str) -> bool:
-    return case_id.startswith("SFB_ID_")
-
-
-def read_ca_per_resid(pdb: Path, chain: str | None = None) -> dict[int, np.ndarray]:
-    s = PDBParser(QUIET=True).get_structure("x", str(pdb))
-    m = next(iter(s))
-    for c in m:
-        if chain is not None and c.id != chain:
-            continue
-        out: dict[int, np.ndarray] = {}
-        for r in c:
-            if r.id[0] != " ":
-                continue
-            if "CA" not in r:
-                continue
-            out[r.id[1]] = r["CA"].coord
-        if out:
-            return out
-    raise RuntimeError(f"no Cα in {pdb} chain={chain}")
-
-
-def read_ca_bfactors(pdb: Path, chain: str | None = None) -> dict[int, float]:
-    s = PDBParser(QUIET=True).get_structure("x", str(pdb))
-    m = next(iter(s))
-    for c in m:
-        if chain is not None and c.id != chain:
-            continue
-        out: dict[int, float] = {}
-        for r in c:
-            if r.id[0] != " ":
-                continue
-            if "CA" not in r:
-                continue
-            out[r.id[1]] = float(r["CA"].bfactor)
-        if out:
-            return out
-    raise RuntimeError(f"no Cα B-factors in {pdb} chain={chain}")
-
-
-def rmsd_on_residues(pred_ca: dict[int, np.ndarray],
-                     ref_ca: dict[int, np.ndarray],
-                     pred_to_ref: dict[int, int],
-                     residues: list[int]) -> tuple[float, int]:
-    pairs = [(qi, pred_to_ref[qi]) for qi in residues
-             if qi in pred_ca and qi in pred_to_ref and pred_to_ref[qi] in ref_ca]
-    if len(pairs) < 3:
-        return float("nan"), len(pairs)
-    from Bio.PDB.Atom import Atom
-    pred_atoms = [Atom("CA", pred_ca[qi], 1.0, 1.0, " ", "CA", 1, "C") for qi, _ in pairs]
-    ref_atoms = [Atom("CA", ref_ca[ri], 1.0, 1.0, " ", "CA", 1, "C") for _, ri in pairs]
-    sup = Superimposer()
-    sup.set_atoms(ref_atoms, pred_atoms)
-    return float(sup.rms), len(pairs)
-
-
-def tmalign(pdb_a: Path, pdb_b: Path) -> tuple[float, float, int]:
-    if not TMALIGN:
-        return float("nan"), float("nan"), 0
-    try:
-        p = subprocess.run([TMALIGN, str(pdb_a), str(pdb_b)],
-                           capture_output=True, text=True, timeout=120)
-    except Exception:
-        return float("nan"), float("nan"), 0
-    t1 = t2 = float("nan"); aln = 0
-    for line in p.stdout.splitlines():
-        if line.startswith("TM-score=") and "normalized by length of Chain_1" in line:
-            try: t1 = float(line.split()[1])
-            except: pass
-        elif line.startswith("TM-score=") and "normalized by length of Chain_2" in line:
-            try: t2 = float(line.split()[1])
-            except: pass
-        elif line.startswith("Aligned length="):
-            parts = line.replace(",", " ").split()
-            for i, q in enumerate(parts):
-                if q == "length=":
-                    try: aln = int(parts[i + 1])
-                    except: pass
-    return t1, t2, aln
-
-
-def get_state_chain(case: dict, which: str) -> str:
-    """Get the chain id within state PDB. The structures/ files use the curated chain
-    from `case_definitions.py` build_v2 — they are usually chain A in the file.
-    Try chain A first; fall back to the chain id from cases.yaml."""
-    return case[f"state_{which}_chain"]
-
-
-def state_chain_in_file(pdb: Path) -> str | None:
-    """Return first chain id present in the PDB file (with Cα atoms)."""
-    s = PDBParser(QUIET=True).get_structure("x", str(pdb))
-    m = next(iter(s))
-    for c in m:
-        for r in c:
-            if r.id[0] == " " and "CA" in r:
-                return c.id
-    return None
-
-
-def evaluate_pdb(pdb: Path, case_id: str) -> dict:
-    case = load_case(case_id)
-    regions = load_regions(case_id)
-    id_case = is_id_case(case_id)
-
-    pred_ca = read_ca_per_resid(pdb, chain="A")
-    pred_plddt = read_ca_bfactors(pdb, chain="A")
-    # Pred residues are 1..L → match regions' residue_index_query directly (it's 1..L too)
-    mean_plddt_overall = float(np.mean(list(pred_plddt.values())))
-    # core / switch pLDDT (None if region empty)
-    core_p = [pred_plddt[r] for r in regions["common_core"] if r in pred_plddt]
-    sw_p = [pred_plddt[r] for r in regions["switch_3A"] if r in pred_plddt]
-    mean_plddt_core = float(np.mean(core_p)) if core_p else float("nan")
-    mean_plddt_switch = float(np.mean(sw_p)) if sw_p else float("nan")
-
-    result = {
-        "pdb": str(pdb),
-        "case_id": case_id,
-        "id_case": id_case,
-        "pred_len": len(pred_ca),
-        "mean_plddt_overall": mean_plddt_overall,
-        "mean_plddt_core": mean_plddt_core,
-        "mean_plddt_switch_3A": mean_plddt_switch,
-        "states": {},
-    }
-
-    # Build pred-resi → state_resi mapping (state_a_resnum comes from FINAL.tsv).
-    # state_b residue mapping = same query residue index (residue_index_query),
-    # because the state_b PDBs were curated to use the query residue numbering too
-    # (see build_v2 case definitions). We discover the actual chain id from the
-    # state PDB file directly.
-    pred_to_state_a = dict(regions["state_a_resnum"])
-
-    for which in ["a", "b"]:
-        pdb_state = DATA_ROOT / "structures" / case_id / f"state_{which}.pdb"
-        if not pdb_state.exists():
-            continue
-        ch = state_chain_in_file(pdb_state)
-        ref_ca = read_ca_per_resid(pdb_state, chain=ch)
-        # Determine pred→ref residue map
-        if which == "a":
-            pred_to_ref = pred_to_state_a
-        else:
-            # state_b: curator built FINAL.tsv via residue-number intersection
-            # (default) or positional pairing (MPT53). Try numbering-intersection
-            # first: pred qi → state_a_resnum, expect those resnums to be present
-            # in state_b PDB. If <50% overlap, fall back to POSITIONAL pairing:
-            # qi=1 → state_b's 1st CA, qi=2 → 2nd, etc.
-            ref_keys = set(ref_ca.keys())
-            num_overlap = sum(1 for rn in pred_to_state_a.values() if rn in ref_keys)
-            if pred_to_state_a and num_overlap >= 0.5 * len(pred_to_state_a):
-                pred_to_ref = pred_to_state_a
-            else:
-                ref_sorted = sorted(ref_ca.keys())
-                pred_sorted = sorted(pred_ca.keys())
-                pred_to_ref = {p: r for p, r in zip(pred_sorted, ref_sorted)}
-
-        rms_core, n_core = rmsd_on_residues(pred_ca, ref_ca, pred_to_ref,
-                                            regions["common_core"])
-        rms_sw, n_sw = rmsd_on_residues(pred_ca, ref_ca, pred_to_ref,
-                                        regions["switch_3A"])
-        t1, t2, aln = tmalign(pdb, pdb_state)
-
-        if id_case:
-            # IDP override: use the WHOLE query region as target (switch_3A = all residues)
-            rms_target = rms_sw if not np.isnan(rms_sw) else rms_core
-            hit = (not np.isnan(rms_target) and rms_target <= 3.0
-                   and mean_plddt_overall >= 70.0)
-        else:
-            # Standard §9.1 binding 3-condition
-            hit = (not np.isnan(rms_core) and rms_core <= 3.0
-                   and mean_plddt_overall >= 70.0
-                   and not np.isnan(mean_plddt_switch)
-                   and mean_plddt_switch >= 70.0)
-
-        result["states"][f"state_{which}"] = {
-            "rmsd_common_core_A": rms_core,
-            "rmsd_switch_3A": rms_sw,
-            "n_common_core_aligned": n_core,
-            "n_switch_3A_aligned": n_sw,
-            "tmalign_tm1": t1,
-            "tmalign_tm2": t2,
-            "tmalign_aligned_len": aln,
-            "hit_primary": hit,
-        }
-
-        if id_case:
-            # ID cases: state_a == state_b, only emit state_a
-            break
-
-    return result
-
-
-def main(argv: list[str] | None = None) -> int:
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--case", required=True)
-    ap.add_argument("--root", required=True, type=Path)
-    ap.add_argument("--out", type=Path, default=None)
-    args = ap.parse_args(argv)
-
-    pdbs = sorted(args.root.rglob("*_unrelaxed_rank_*.pdb"))
-    if not pdbs:
-        print(f"ERROR: no PDBs under {args.root}", file=sys.stderr)
-        return 2
-    out_tsv = args.out or args.root / "evals.tsv"
-    sample = evaluate_pdb(pdbs[0], args.case)
-    state_keys = list(sample["states"].keys())
-    cols = ["subset_id", "model", "seed", "rank",
-            "mean_plddt_overall", "mean_plddt_core", "mean_plddt_switch_3A"]
-    for sk in state_keys:
-        for m in ("rmsd_common_core_A", "rmsd_switch_3A",
-                  "tmalign_tm1", "tmalign_tm2", "hit_primary"):
-            cols.append(f"{sk}__{m}")
-    cols.append("pdb")
-
-    with out_tsv.open("w") as out:
-        out.write("\t".join(cols) + "\n")
-        for pdb in pdbs:
-            m = PRED_NAME_RE.match(pdb.name)
-            sid = m.group("subset") if m else pdb.stem
-            rank = int(m.group("rank")) if m else -1
-            model = int(m.group("model")) if m else -1
-            seed = int(m.group("seed")) if m else -1
-            try:
-                r = evaluate_pdb(pdb, args.case)
-            except Exception as e:
-                print(f"WARN: eval failed {pdb}: {e}", file=sys.stderr)
-                continue
-            (pdb.with_name(pdb.stem + "_eval.json")).write_text(
-                json.dumps(r, indent=2, default=float))
-            row = [sid, model, seed, rank,
-                   f"{r['mean_plddt_overall']:.2f}",
-                   f"{r['mean_plddt_core']:.2f}" if r['mean_plddt_core'] == r['mean_plddt_core'] else "NA",
-                   f"{r['mean_plddt_switch_3A']:.2f}" if r['mean_plddt_switch_3A'] == r['mean_plddt_switch_3A'] else "NA"]
-            for sk in state_keys:
-                sv = r["states"].get(sk, {})
-                for k in ("rmsd_common_core_A", "rmsd_switch_3A",
-                         "tmalign_tm1", "tmalign_tm2", "hit_primary"):
-                    v = sv.get(k)
-                    if v is None or (isinstance(v, float) and v != v):
-                        row.append("NA")
-                    elif isinstance(v, bool):
-                        row.append("1" if v else "0")
-                    elif isinstance(v, float):
-                        row.append(f"{v:.4f}")
-                    else:
-                        row.append(str(v))
-            row.append(str(pdb))
-            out.write("\t".join(str(x) for x in row) + "\n")
-    print(f"wrote {len(pdbs)} preds -> {out_tsv}")
-    return 0
-
-
-if __name__ == "__main__":
-    sys.exit(main())