code/scripts/rescore.py

#!/usr/bin/env python3
"""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 \
        --design_dir path/to/designs \
        --target cam \
        --gpu 0 \
        --out results/rescore_out.json
"""
import os
import sys
import json
import argparse
import logging
from pathlib import Path

import numpy as np
import torch

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

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_backbone_coords, get_aa_indices, align_structures,
)

# --------------------------------------------------------------------------
# 3-seed ensemble (paper: v5-S2 TS-S2, seeds 1024 / 5555 / 789)
# --------------------------------------------------------------------------
SEEDS = [1024, 5555, 789]
CKPT_LAYOUTS = [
    str(BASE / "checkpoints/v5s2/seed_{seed}/best_phase2.pt"),         # HF release layout
    str(BASE / "results/v5_training/10seed/seed_{seed}/best_phase2.pt"),  # in-repo dev layout
]
ESM_DIR = str(BASE / "data/esm2_embeddings")


def _resolve_ckpt(seed):
    for tmpl in CKPT_LAYOUTS:
        p = tmpl.format(seed=seed)
        if Path(p).exists():
            return p
    return None

# 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 = _resolve_ckpt(seed)
        if ckpt is None:
            tried = [t.format(seed=seed) for t in CKPT_LAYOUTS]
            logger.warning(f"Missing checkpoint for seed {seed} in any of: {tried}; skip")
            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): _resolve_ckpt(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():
    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__":
    main()