scripts/reconstruct.py

"""Claims 1 & 2 — reconstruction RMSD on CATH (test split) and CAMEO.

Full-token reconstruction protocol (matches the authors' run_example.py, with an
explicit true_length so proteins longer than the 128-token budget still decode to
their real residue count):

    encode -> keep first min(L,128) tokens -> decode(true_length=L) -> Kabsch RMSD (CA, Å)

Datasets:
  cath  : cctien/protein_backbone_cath_4.3  (test split, CA coords)
  cameo : genbio-ai/casp14-casp15-cameo-test-proteins  (cameo_structure_gts/*_target.pdb)

Proteins are filtered to <=256 residues (paper's training/eval regime).
"""
from __future__ import annotations

import argparse
import json
import os
from pathlib import Path

import numpy as np
import torch

from apt.models import APTTokenizer
from apt.utils import kabsch_rmsd

CKPT = Path(os.environ.get("CKPT_DIR", "/tmp/apt_ckpts"))
TOK = CKPT / "tokenizer128.pt"
OUT = Path(os.environ.get("OUT_DIR", "/tmp/out"))
MAXLEN = 256


def tmscore(recon_L3: np.ndarray, native_L3: np.ndarray) -> float:
    """TM-score after Kabsch superposition (recon already aligned upstream)."""
    L = native_L3.shape[0]
    d0 = 1.24 * (max(L, 19) - 15) ** (1.0 / 3.0) - 1.8
    d0 = max(d0, 0.5)
    # Kabsch-align recon onto native
    P = recon_L3 - recon_L3.mean(0)
    Q = native_L3 - native_L3.mean(0)
    H = P.T @ Q
    U, S, Vt = np.linalg.svd(H)
    d = np.sign(np.linalg.det(Vt.T @ U.T))
    D = np.diag([1, 1, d])
    R = Vt.T @ D @ U.T
    P_al = P @ R.T
    di = np.linalg.norm(P_al - Q, axis=1)
    return float((1.0 / (1.0 + (di / d0) ** 2)).mean())


def load_cath(n_max: int | None):
    from huggingface_hub import hf_hub_download

    fp = hf_hub_download(
        repo_id="cctien/protein_backbone_cath_4.3",
        filename="data/test-00000-of-00001.parquet",
        repo_type="dataset",
    )
    import pyarrow.parquet as pq

    tbl = pq.read_table(fp)
    names = tbl.column("name").to_pylist()
    coords = tbl.column("coords").to_pylist()
    out = []
    for nm, c in zip(names, coords):
        ca = np.asarray(c["CA"], dtype=np.float32)  # (L,3)
        if ca.ndim != 2 or ca.shape[1] != 3:
            continue
        if np.isnan(ca).any():
            continue
        if not (1 <= ca.shape[0] <= MAXLEN):
            continue
        out.append((nm, ca))
    if n_max:
        out = out[:n_max]
    return out


def load_cameo(n_max: int | None):
    from huggingface_hub import snapshot_download

    root = snapshot_download(
        repo_id="genbio-ai/casp14-casp15-cameo-test-proteins",
        repo_type="dataset",
        allow_patterns=["cameo_structure_gts/*"],
    )
    pdb_dir = Path(root) / "cameo_structure_gts"
    out = []
    for pdb in sorted(pdb_dir.glob("*_target.pdb")):
        ca = []
        seen = set()
        with pdb.open() as fh:
            for line in fh:
                if not line.startswith("ATOM"):
                    continue
                if line[12:16].strip() != "CA":
                    continue
                if line[16:17] not in (" ", "A"):
                    continue
                resid = line[21:27]
                if resid in seen:
                    continue
                seen.add(resid)
                ca.append([float(line[30:38]), float(line[38:46]), float(line[46:54])])
        ca = np.asarray(ca, dtype=np.float32)
        if not (1 <= ca.shape[0] <= MAXLEN):
            continue
        out.append((pdb.stem, ca))
    if n_max:
        out = out[:n_max]
    return out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--dataset", choices=["cath", "cameo"], required=True)
    ap.add_argument("--n_max", type=int, default=0)
    ap.add_argument("--cap_tokens", type=int, default=0, help="0 = use all L tokens; else cap")
    ap.add_argument("--noise", type=float, default=0.45)
    ap.add_argument("--n_steps", type=int, default=100)
    args = ap.parse_args()

    dev = "cuda" if torch.cuda.is_available() else "cpu"
    torch.manual_seed(0)
    tok = APTTokenizer.from_pretrained(TOK).to(dev).eval()
    max_toks = tok.cfg.n_tokens

    data = (load_cath if args.dataset == "cath" else load_cameo)(args.n_max or None)
    print(f"{args.dataset}: {len(data)} structures (<= {MAXLEN} res), device={dev}")

    rows = []
    for i, (name, ca) in enumerate(data):
        L = ca.shape[0]
        x = torch.from_numpy(ca).float()
        x = x - x.mean(0, keepdim=True)
        x = (x / 10.0).unsqueeze(0).to(dev)
        try:
            keep = L if args.cap_tokens == 0 else min(L, args.cap_tokens)
            with torch.no_grad():
                _, _, idx_BL = tok.encode(x)
                idx_BL = idx_BL[:, :keep]
                recon = tok.decode(idx_BL, true_length=L, n_steps=args.n_steps, noise_weight=args.noise)
                rmsd = float(kabsch_rmsd(recon.cpu(), x.cpu()).item()) * 10.0
            r_np = recon.squeeze(0).cpu().numpy() * 10.0
            n_np = ca - ca.mean(0)
            tm = tmscore(r_np, n_np)
        except Exception as e:  # noqa: BLE001
            print(f"  skip {name} (L={L}): {e}")
            continue
        rows.append({"name": name, "L": L, "n_tokens": keep, "rmsd": rmsd, "tm": tm})
        if i % 50 == 0:
            print(f"  [{i}/{len(data)}] {name} L={L} rmsd={rmsd:.3f} tm={tm:.3f}")

    rmsds = np.array([r["rmsd"] for r in rows])
    tms = np.array([r["tm"] for r in rows])

    def subset(pred):
        rs = np.array([r["rmsd"] for r in rows if pred(r)])
        ts = np.array([r["tm"] for r in rows if pred(r)])
        if not len(rs):
            return None
        return {"n": len(rs), "rmsd_mean": round(float(rs.mean()), 4),
                "rmsd_median": round(float(np.median(rs)), 4), "tm_mean": round(float(ts.mean()), 4)}

    summary = {
        "dataset": args.dataset,
        "cap_tokens": args.cap_tokens, "noise": args.noise, "n_steps": args.n_steps,
        "n": len(rows),
        "rmsd_mean": round(float(rmsds.mean()), 4),
        "rmsd_median": round(float(np.median(rmsds)), 4),
        "tm_mean": round(float(tms.mean()), 4),
        "frac_tm_gt_0.5": round(float((tms > 0.5).mean()), 4),
        "frac_full_token": round(float(np.mean([r["L"] <= max_toks for r in rows])), 4),
        "subset_L_le_128": subset(lambda r: r["L"] <= 128),
        "subset_L_gt_128": subset(lambda r: r["L"] > 128),
    }
    print("RESULT_JSON " + json.dumps(summary))

    OUT.mkdir(parents=True, exist_ok=True)
    (OUT / f"recon_{args.dataset}_summary.json").write_text(json.dumps(summary, indent=2))
    import csv

    with (OUT / f"recon_{args.dataset}_per_structure.csv").open("w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=["name", "L", "n_tokens", "rmsd", "tm"])
        w.writeheader()
        w.writerows(rows)


if __name__ == "__main__":
    main()