fix rescore.py bugs: missing holo->apo Kabsch (Q_apo=0 for all designs), residue-ID alignment on renumbered chains (PXDesign/Protenix), HETATM chain-A dropped, single-seed -> 3-seed ensemble, add .cif support

Replaces old v2 rescore script with the paper's canonical v5-S2 3-seed ensemble pipeline (seeds 1024/5555/789). Fixes:
1. Adds second Kabsch (holo->apo) on binder coords so the binder sits next to the apo receptor slot; previously Q_apo=0 for every design because 1CFD sits ~34 A from 3CLN in world frame.
2. Sequence-anchored ordinal Kabsch instead of residue-ID intersection — the ID intersection paired the wrong CAs whenever chain A was renumbered from 1 (PXDesign, Protenix, Proteina).
3. HETATM-safe receptor extraction — PXDesign writes chain A as HETATM and get_residues(only_standard=True) dropped every one.
4. 3-seed ensemble (seeds 1024/5555/789) with per-seed and mean scores.
5. Accepts .cif in addition to .pdb.
6. --design_dir / --target CLI args instead of hardcoded internal paths.

code/scripts/rescore.py

@@ -1,177 +1,265 @@
 #!/usr/bin/env python3
-"""
-Re-score binder PDB designs with a Q_theta checkpoint.
-
-Walks a directory of designs (binder PDB + sibling holo / apo receptor PDBs),
-runs each through DifferentiableQTheta, and writes per-design
-S = Q_theta(holo) - Q_theta(apo) plus the raw holo/apo scores to JSON.
+"""Re-score binder designs with the paper's canonical Q_theta scoring pipeline
+(v5-S2 TS-S2 3-seed ensemble: seeds 1024 / 5555 / 789).
+
+Fixes vs the older `rescore_v2.py`:
+
+  (1) Sequence-anchored ordinal Kabsch — residue-ID intersection fails when a
+      generator (PXDesign, Protenix, Proteina) renumbers chain A from 1 while
+      the canonical holo PDB is numbered from its crystallographic start.
+  (2) Second Kabsch to move the binder from holo frame into apo frame before
+      scoring against the apo receptor — otherwise the binder sits 30+ Å from
+      the apo receptor slot and Q_apo collapses to 0 for every design.
+  (3) HETATM-safe receptor extraction — PXDesign writes chain A as HETATM;
+      `get_residues(..., only_standard=True)` filters HETATM and drops every
+      chain-A residue, so the whole design gets skipped.
+  (4) 3-seed ensemble — a single checkpoint under-reports; paper averages
+      seeds 1024 / 5555 / 789.
+  (5) `.cif` inputs supported alongside `.pdb`.
+  (6) `--design_dir` argument instead of hardcoded internal directory tree.
 
 Usage:
-    python code/scripts/rescore.py \\
-        --checkpoint checkpoints/Q_theta_phase2.pt \\
-        --gpu 0
+    python code/scripts/rescore.py \
+        --design_dir path/to/designs \
+        --target cam \
+        --gpu 0 \
+        --out results/rescore_out.json
 """
-import os, sys, json, argparse, glob, logging
+import os
+import sys
+import json
+import argparse
+import logging
+from pathlib import Path
+
 import numpy as np
 import torch
-from pathlib import Path
 
-logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
 logger = logging.getLogger(__name__)
 
-BASE = str(Path(__file__).resolve().parent.parent.parent)
-sys.path.insert(0, os.path.join(BASE, 'code'))
-sys.path.insert(0, BASE)
+BASE = Path(__file__).resolve().parent.parent.parent
+sys.path.insert(0, str(BASE / "code"))
+sys.path.insert(0, str(BASE))
 
 from models.differentiable_features import DifferentiableQTheta
-from utils.pdb_utils import load_structure, get_residues, get_backbone_coords, get_aa_indices, align_structures
-
-HOLO_PDB = os.path.join(BASE, 'data/pdbs/cam_holo/3CLN.pdb')
-APO_PDB = os.path.join(BASE, 'data/pdbs/cam_apo/1CFD.pdb')
-
-
-def score_pdb_list(dq, pdb_list, ref_resnums, ref_coords, device):
-    """Score a list of design PDB files."""
-    results = []
-    for pdb_path in pdb_list:
-        name = os.path.basename(pdb_path).replace(".pdb", "")
-        try:
-            design_model = load_structure(pdb_path)
-            chains = [c.id for c in design_model.get_chains()]
-            rec_chain = 'A' if 'A' in chains else chains[0]
-            binder_chain = 'B' if 'B' in chains else [c for c in chains if c != rec_chain][0]
-
-            rec_res = get_residues(design_model[rec_chain])
-            binder_res = get_residues(design_model[binder_chain])
-            rec_coords_d, _ = get_backbone_coords(rec_res)
-            binder_coords, binder_mask = get_backbone_coords(binder_res)
-            binder_aa_idx = get_aa_indices(binder_res)
-
-            design_resnums = {r.get_id()[1]: i for i, r in enumerate(rec_res)}
-            common = sorted(set(design_resnums.keys()) & set(ref_resnums.keys()))
-            if len(common) < 10:
-                logger.warning(f"  Skip {name}: <10 common residues")
-                continue
-
-            d_ca = rec_coords_d[[design_resnums[r] for r in common], 1]
-            r_ca = ref_coords[[ref_resnums[r] for r in common], 1]
-            mobile_center = d_ca.mean(0)
-            ref_center = r_ca.mean(0)
-            _, R = align_structures(d_ca, r_ca)
-
-            flat = binder_coords.reshape(-1, 3) - mobile_center
-            aligned_binder = (flat @ R.T + ref_center).reshape(-1, 4, 3)
-
-            coords_t = torch.from_numpy(aligned_binder).float().to(device)
-            mask_t = torch.from_numpy(binder_mask).bool().to(device)
-            aa_t = torch.from_numpy(binder_aa_idx).long().to(device)
-
-            with torch.no_grad():
-                q_holo = dq.score(coords_t, mask_t, binder_aa_idx=aa_t,
-                                  receptor_label='holo').item()
-                q_apo = dq.score(coords_t, mask_t, binder_aa_idx=aa_t,
-                                 receptor_label='apo').item()
-            S = q_holo - q_apo
-            results.append({"design": name, "Q_holo": q_holo, "Q_apo": q_apo, "S": S})
-        except Exception as e:
-            logger.warning(f"  Skip {name}: {e}")
-    return results
-
-
-def summarize(results, label):
-    if not results:
-        return {}
-    S = [r["S"] for r in results]
-    return {
-        "method": label, "n": len(S),
-        "S_mean": float(np.mean(S)), "S_std": float(np.std(S)),
-        "S_pos_pct": float(np.mean([s > 0 for s in S]) * 100),
-        "Q_holo_mean": float(np.mean([r["Q_holo"] for r in results])),
-        "Q_apo_mean": float(np.mean([r["Q_apo"] for r in results])),
+from utils.pdb_utils import (
+    load_structure, get_backbone_coords, get_aa_indices, align_structures,
+)
+
+# --------------------------------------------------------------------------
+# 3-seed ensemble (paper: v5-S2 TS-S2, seeds 1024 / 5555 / 789)
+# --------------------------------------------------------------------------
+SEEDS = [1024, 5555, 789]
+CKPT_TEMPLATE = str(BASE / "results/v5_training/10seed/seed_{seed}/best_phase2.pt")
+ESM_DIR = str(BASE / "data/esm2_embeddings")
+
+# Per-target holo/apo canonical PDB pair (extend as needed).
+TARGETS = {
+    "cam":      {"holo": "data/pdbs/cam_holo/3CLN.pdb",   "apo": "data/pdbs/cam_apo/1CFD.pdb",   "chain": "A", "esm_target": "cam"},
+    "bcl2":     {"holo": "data/pdbs/bcl2_holo/2XA0.pdb",  "apo": "data/pdbs/bcl2_apo/1G5M.pdb",  "chain": "A", "esm_target": "bcl2"},
+    "era":      {"holo": "data/pdbs/era_complex/1GWR.pdb","apo": "data/pdbs/era_apo/3ERT.pdb",   "chain": "A", "esm_target": "era"},
+    "mdm2":     {"holo": "data/pdbs/mdm2_holo/1T4E.pdb",  "apo": "data/pdbs/mdm2_apo/1Z1M.pdb",  "chain": "A", "esm_target": "mdm2"},
+    "a2a":      {"holo": "data/pdbs/a2a_holo/6GDG.pdb",   "apo": "data/pdbs/a2a_apo/4EIY.pdb",   "chain": "A", "esm_target": "a2a"},
+    "pai1":     {"holo": "data/pdbs/pai1_holo/1LJ5.pdb",  "apo": "data/pdbs/pai1_apo/1B3K.pdb",  "chain": "A", "esm_target": "pai1"},
+    "ran":      {"holo": "data/pdbs/ran_holo/1RRP.pdb",   "apo": "data/pdbs/ran_apo/1BYT.pdb",   "chain": "A", "esm_target": "ran"},
+    "integrin": {"holo": "data/pdbs/integrin_holo/2VDO.pdb","apo": "data/pdbs/integrin_apo/2VDK.pdb","chain": "A", "esm_target": "integrin"},
+}
+
+_THREE2ONE = {
+    'ALA':'A','ARG':'R','ASN':'N','ASP':'D','CYS':'C','GLU':'E','GLN':'Q',
+    'GLY':'G','HIS':'H','ILE':'I','LEU':'L','LYS':'K','MET':'M','PHE':'F',
+    'PRO':'P','SER':'S','THR':'T','TRP':'W','TYR':'Y','VAL':'V',
+}
+
+
+def _seq_one(residues):
+    return ''.join(_THREE2ONE.get(r.resname.strip(), 'X') for r in residues)
+
+
+def _load_ref(pdb_path, chain):
+    """Canonical holo/apo reference. Standard chain, ATOM-only."""
+    m = load_structure(pdb_path)
+    rs = [r for r in m[chain] if 'CA' in r and r.get_id()[0] == ' ']
+    coords, _ = get_backbone_coords(rs)
+    return rs, coords
+
+
+def _get_design_chains(model):
+    """Return (receptor_residues, binder_residues) for a design PDB/CIF.
+
+    HETATM-safe: PXDesign writes receptor as HETATM in chain A. We include any
+    residue with a CA atom, regardless of het-flag. The receptor is the
+    longest chain (>= 50 residues); the binder is the 5–120 residue chain.
+    """
+    rec_res, binder_res = None, None
+    for chain in model.get_chains():
+        residues = [r for r in chain if 'CA' in r]
+        if not residues:
+            continue
+        if 5 <= len(residues) <= 120:
+            if binder_res is None or len(residues) < len(binder_res):
+                binder_res = residues
+        if len(residues) >= 50:
+            if rec_res is None or len(residues) > len(rec_res):
+                rec_res = residues
+    return rec_res, binder_res
+
+
+def _seq_offset(design_res, ref_res):
+    """Locate where the design's receptor sequence starts within the reference
+    sequence — handles receptor cropping and residue renumbering-from-1.
+    """
+    d_seq = _seq_one(design_res)
+    r_seq = _seq_one(ref_res)
+    probe = d_seq[:30] if len(d_seq) >= 30 else d_seq
+    idx = r_seq.find(probe)
+    return max(0, idx)
+
+
+def _score_one(design_path, holo_res, holo_bb, apo_res, apo_bb):
+    """Extract binder coords in holo AND apo frames. Return None if the design
+    can't be parsed."""
+    try:
+        model = load_structure(str(design_path))
+    except Exception as e:
+        logger.warning(f"  Skip {design_path.name}: load failed ({e})")
+        return None
+
+    rec_res, binder_res = _get_design_chains(model)
+    if binder_res is None:
+        logger.warning(f"  Skip {design_path.name}: no binder chain (5–120 residues) found")
+        return None
+    if rec_res is None:
+        # No receptor in the design — binder assumed to already be in holo frame
+        b_bb, b_mask = get_backbone_coords(binder_res)
+        b_aa = get_aa_indices(binder_res)
+        return b_bb, b_bb, b_mask, b_aa
+
+    d_bb, _ = get_backbone_coords(rec_res)  # design chain A backbone
+
+    # Sequence-anchored ordinal alignment: i-th design residue corresponds to
+    # (offset + i)-th reference residue. Robust to residue renumbering-from-1.
+    offset_h = _seq_offset(rec_res, holo_res)
+    n_align_h = min(len(rec_res), len(holo_res) - offset_h)
+    if n_align_h < 10:
+        logger.warning(f"  Skip {design_path.name}: sequence anchor could not align to holo")
+        return None
+    d_ca_h = d_bb[:n_align_h, 1]
+    h_ca = holo_bb[offset_h:offset_h + n_align_h, 1]
+    mc_h = d_ca_h.mean(0); rc_h = h_ca.mean(0)
+    _, R_h = align_structures(d_ca_h, h_ca)
+
+    b_bb, b_mask = get_backbone_coords(binder_res)
+    b_aa = get_aa_indices(binder_res)
+    flat = b_bb.reshape(-1, 3) - mc_h
+    binder_holo = (flat @ R_h.T + rc_h).reshape(-1, 4, 3)
+
+    # Holo -> apo transfer via residue-ID intersection on the CANONICAL PDBs
+    # (their numbering semantics match: e.g. 3CLN 5–147, 1CFD 5–147 overlap).
+    holo_rn = {r.get_id()[1]: i for i, r in enumerate(holo_res)}
+    apo_rn = {r.get_id()[1]: i for i, r in enumerate(apo_res)}
+    common = sorted(set(holo_rn) & set(apo_rn))
+    if len(common) >= 10:
+        h_ca2 = holo_bb[[holo_rn[k] for k in common], 1]
+        a_ca2 = apo_bb[[apo_rn[k] for k in common], 1]
+        mh2 = h_ca2.mean(0); ma2 = a_ca2.mean(0)
+        _, R_ha = align_structures(h_ca2, a_ca2)
+        flat2 = binder_holo.reshape(-1, 3) - mh2
+        binder_apo = (flat2 @ R_ha.T + ma2).reshape(-1, 4, 3)
+    else:
+        binder_apo = binder_holo  # last resort: same frame
+    return binder_holo, binder_apo, b_mask, b_aa
+
+
+def rescore(target, design_dir, gpu, out_path):
+    cfg = TARGETS[target]
+    device = f"cuda:{gpu}"
+    holo_pdb = str(BASE / cfg["holo"])
+    apo_pdb = str(BASE / cfg["apo"])
+
+    holo_res, holo_bb = _load_ref(holo_pdb, cfg["chain"])
+    apo_res, apo_bb = _load_ref(apo_pdb, cfg["chain"])
+    logger.info(f"[{target}] holo={cfg['holo']} n={len(holo_res)}  apo={cfg['apo']} n={len(apo_res)}")
+
+    designs = sorted([p for p in Path(design_dir).iterdir() if p.suffix.lower() in (".pdb", ".cif", ".mmcif")])
+    logger.info(f"Found {len(designs)} design files in {design_dir}")
+    extracted = []
+    for p in designs:
+        r = _score_one(p, holo_res, holo_bb, apo_res, apo_bb)
+        if r is None:
+            continue
+        b_h, b_a, m, aa = r
+        extracted.append({"id": p.stem, "coords_holo": b_h, "coords_apo": b_a, "mask": m, "aa_idx": aa})
+    logger.info(f"Extracted {len(extracted)} usable designs")
+    if not extracted:
+        raise RuntimeError("No designs successfully extracted")
+
+    results = {d["id"]: {} for d in extracted}
+    for seed in SEEDS:
+        ckpt = CKPT_TEMPLATE.format(seed=seed)
+        if not Path(ckpt).exists():
+            logger.warning(f"Missing checkpoint {ckpt}; skip seed {seed}")
+            continue
+        logger.info(f"Scoring seed {seed}...")
+        dq = DifferentiableQTheta(checkpoint_path=ckpt, device=device, esm_dir=ESM_DIR)
+        dq.load_receptor(holo_pdb, chain=cfg["chain"], label="holo", esm_target=cfg["esm_target"])
+        dq.load_receptor(apo_pdb, chain=cfg["chain"], label="apo", esm_target=cfg["esm_target"])
+        for d in extracted:
+            try:
+                mt = torch.from_numpy(d["mask"]).bool().to(device)
+                at = torch.from_numpy(d["aa_idx"]).long().to(device)
+                with torch.no_grad():
+                    ct_h = torch.from_numpy(d["coords_holo"]).float().to(device)
+                    qh = dq.score(ct_h, mt, binder_aa_idx=at, receptor_label="holo").item()
+                    ct_a = torch.from_numpy(d["coords_apo"]).float().to(device)
+                    qa = dq.score(ct_a, mt, binder_aa_idx=at, receptor_label="apo").item()
+                results[d["id"]][str(seed)] = {"Q_holo": round(qh, 6), "Q_apo": round(qa, 6), "S": round(qh - qa, 6)}
+            except Exception as e:
+                results[d["id"]][str(seed)] = {"error": repr(e)}
+        del dq
+        torch.cuda.empty_cache()
+
+    # Mean over available seeds
+    for d in extracted:
+        by_seed = results[d["id"]]
+        qh = [v["Q_holo"] for v in by_seed.values() if "S" in v]
+        qa = [v["Q_apo"] for v in by_seed.values() if "S" in v]
+        ss = [v["S"] for v in by_seed.values() if "S" in v]
+        if ss:
+            by_seed["mean"] = {
+                "Q_holo": round(float(np.mean(qh)), 6),
+                "Q_apo": round(float(np.mean(qa)), 6),
+                "S": round(float(np.mean(ss)), 6),
+                "S_std": round(float(np.std(ss)), 6),
+                "n_seeds": len(ss),
+            }
+
+    out = {
+        "target": target, "holo_pdb": holo_pdb, "apo_pdb": apo_pdb,
+        "seeds": SEEDS, "checkpoints": {str(s): CKPT_TEMPLATE.format(seed=s) for s in SEEDS},
+        "n_designs": len(extracted), "per_design": results,
     }
+    Path(out_path).parent.mkdir(parents=True, exist_ok=True)
+    Path(out_path).write_text(json.dumps(out, indent=2))
+    logger.info(f"Saved {out_path}")
 
+    s_all = [results[d["id"]]["mean"]["S"] for d in extracted if "mean" in results[d["id"]]]
+    if s_all:
+        s = np.array(s_all)
+        logger.info(f"[{target}] n={len(s)}  mean S={s.mean():+.4f}  S>0: {100*(s>0).mean():.0f}%  ({int((s>0).sum())}/{len(s)})")
 
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--gpu", type=int, default=7)
-    parser.add_argument("--checkpoint", default="checkpoints/Q_theta_phase2.pt")
-    args = parser.parse_args()
-
-    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
-    device = "cuda:0"
-
-    logger.info(f"Loading Q_theta from {args.checkpoint}")
-    dq = DifferentiableQTheta(checkpoint_path=args.checkpoint, device=device,
-                               esm_dir=os.path.join(BASE, "data/esm2_embeddings"))
-    dq.load_receptor(HOLO_PDB, chain='A', label='holo', esm_target='cam')
-    dq.load_receptor(APO_PDB, chain='A', label='apo', esm_target='cam')
-
-    ref_model = load_structure(HOLO_PDB)
-    ref_res = get_residues(ref_model['A'])
-    ref_coords, _ = get_backbone_coords(ref_res)
-    ref_resnums = {r.get_id()[1]: i for i, r in enumerate(ref_res)}
-
-    output_dir = os.path.join(BASE, "results/v2_strict_holdout/scoring")
-    os.makedirs(output_dir, exist_ok=True)
-
-    # Define design directories
-    design_sets = {
-        "vanilla": os.path.join(BASE, "results/independent_validation/vanilla/holo_pdbs"),
-        "langevin": os.path.join(BASE, "results/langevin_refinement/refined_pdbs"),
-        "classifier": os.path.join(BASE, "results/guided_diffusion/guided"),
-        "smc_r3": os.path.join(BASE, "results/smc_guidance/cam/round_3"),
-    }
 
-    # Also check for TDS and PXDesign
-    tds_dirs = glob.glob(os.path.join(BASE, "results/tds_guidance/cam/designs"))
-    if tds_dirs:
-        design_sets["tds"] = tds_dirs[0]
-
-    # PXDesign directories
-    for px_method in ["pxdesign_scoring", "pxdesign_classifier", "pxdesign_tds",
-                       "pxdesign_smc", "pxdesign_langevin"]:
-        px_dir = os.path.join(BASE, f"results_familysplit/design_bd30/{px_method}")
-        if not os.path.exists(px_dir):
-            px_dir = os.path.join(BASE, f"results/{px_method}")
-        if os.path.exists(px_dir):
-            pdbs = glob.glob(os.path.join(px_dir, "*.pdb"))
-            if pdbs:
-                design_sets[px_method] = px_dir
-
-    all_results = {}
-    summaries = []
-
-    for method, pdb_dir in design_sets.items():
-        if not os.path.exists(pdb_dir):
-            logger.warning(f"  {method}: directory not found ({pdb_dir})")
-            continue
-        pdbs = sorted(glob.glob(os.path.join(pdb_dir, "*.pdb")))
-        if not pdbs:
-            logger.warning(f"  {method}: no PDB files")
-            continue
-
-        logger.info(f"\n=== {method} ({len(pdbs)} designs) ===")
-        results = score_pdb_list(dq, pdbs, ref_resnums, ref_coords, device)
-        s = summarize(results, method)
-        if s:
-            summaries.append(s)
-            logger.info(f"  {method}: n={s['n']}, S̄={s['S_mean']:.3f}±{s['S_std']:.3f}, "
-                       f"S>0={s['S_pos_pct']:.0f}%, Q+={s['Q_holo_mean']:.3f}, Q-={s['Q_apo_mean']:.3f}")
-        all_results[method] = {"results": results, "summary": s}
-
-    # Save
-    with open(os.path.join(output_dir, "rescore_v2_all.json"), "w") as f:
-        json.dump(all_results, f, indent=2)
-
-    # Print summary table
-    print("\n" + "=" * 70)
-    print("V2 RESCORING SUMMARY (strict holdout, CaM OOD)")
-    print("=" * 70)
-    print(f"{'Method':20s} {'n':>4s} {'S̄':>8s} {'±σ':>6s} {'S>0%':>6s} {'Q+':>6s} {'Q-':>6s}")
-    print("-" * 70)
-    for s in sorted(summaries, key=lambda x: x['S_mean'], reverse=True):
-        print(f"{s['method']:20s} {s['n']:4d} {s['S_mean']:8.3f} {s['S_std']:6.3f} "
-              f"{s['S_pos_pct']:5.1f}% {s['Q_holo_mean']:6.3f} {s['Q_apo_mean']:6.3f}")
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--design_dir", required=True, help="dir with .pdb / .cif designs (receptor chain A + binder chain B)")
+    ap.add_argument("--target", required=True, choices=list(TARGETS.keys()))
+    ap.add_argument("--gpu", type=int, default=0)
+    ap.add_argument("--out", default=None)
+    args = ap.parse_args()
+    if args.out is None:
+        args.out = str(BASE / f"results/rescore/{args.target}_qtheta_v5s2.json")
+    rescore(args.target, args.design_dir, args.gpu, args.out)
 
 
 if __name__ == "__main__":