Blame for 52_final_model.py

"""Train and save the final four-modality model (EJ-ad+MV, the recommended configuration of 51_) for embed.py.

    .venv/bin/python -I 52_final_model.py cv                 5-fold check of the stage-2 selection rule (3 seeds)
    .venv/bin/python -I 52_final_model.py fold0   [--rule R] -> models/ej_ad_mv_fold0/ (fold-0 train / val; test held out)
    .venv/bin/python -I 52_final_model.py release [--rule R] -> models/ej_ad_mv/       (train = fold-0 train + test, val = fold-0 val)
    --out NAME writes models/NAME/ (cv: results/final_cv_*_NAME.csv) instead, e.g. to retrain without overwriting
               the shipped checkpoints / results

Same data, heads, losses, steps and early stopping as 51_medvtok_eval.py (UNI2 anchor; EJ-ad+MV):
  stage 1  mIF head (VirTues -> MLP, 400 steps), RNA heads (hvg / Tahoe-x1 / SE-600M / AIDO.Cell -> linear, 2700),
           text heads (KEEP / BiomedCLIP / Qwen3-Emb-0.6B -> linear, 1000), each InfoNCE against frozen UNI2 on real
           pairs, early stopping on its val R@1 (deterministic: identical to 51_)
  stage 2  heads frozen; residual adapter H&E = UNI2 + A([UNI2, c * MedVTok]) + temperatures, sum of the three pair
           losses (RNA / text: mean over the encoder heads) + similarity distillation to UNI2 (lambda 1), AdamW lr 1e-4,
           1000 steps (OneCycle)
Stage-2 selection rules (the val curve is flat after ~100 steps, so 51_'s argmax of the noisy val geo-mean picks very
different steps between otherwise identical runs; 'cv' quantifies it):
  es        51_: early stopping on the geometric mean of the three val R@1, one seed
  last      the adapter after the full 1000-step schedule, one seed (val unused)
  ens-es    mean of the residuals of 3 early-stopped adapters (seeds): H&E = UNI2 + mean_s A_s(.)
  ens-last  same with the 1000-step adapters
Real pairs only (ORION H&E-mIF, HEST H&E-RNA, TCGA/PathGen H&E-text); virtual / synthetic data never trained on, used
only in the emergent test pairs. 'fold0' evaluates on its held-out test (51_ protocol, g100, same test items) and
writes test.csv + test_items.npz (for 53_); 'release' has no test split: expected performance = the 5-fold CV.
Output: results/final_cv_{long,summary,curves}.csv (cv); models/<name>/{model.pt, meta.json, val_curve.csv[, test.csv]}
"""
import json
import os
import sys
import time

import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from common import CACHE, ROOT, SPLITS  # noqa: E402
from embed import RNA_E, TXT_E, Bind4, mlp  # noqa: E402

MODE = sys.argv[1]
assert MODE in ("cv", "fold0", "release"), MODE
RULE = sys.argv[sys.argv.index("--rule") + 1] if "--rule" in sys.argv else "ens-es"  # chosen by `cv` (RESULTS.md)
assert RULE in ("es", "last", "ens-es", "ens-last"), RULE
NSEED = 3
dev = "cuda"
H, RF = os.path.join(CACHE, "hest"), os.path.join(CACHE, "rnafm")
G100, DRAWS, BS, LR, EVAL_EVERY = 100, 50, 256, 3e-4, 50
f32 = lambda a: torch.from_numpy(np.asarray(a, dtype=np.float32))
nz = lambda x: F.normalize(x, dim=-1)
t0 = time.time()

# ---------------------------------------------------------------- data (= 51_)
IF_O = f32(np.load(os.path.join(CACHE, "if_virtues_m16.npy"))).to(dev)
IF_V_T = f32(np.load(os.path.join(CACHE, "if_virtues_m16_virtual_miphei_tcga_text.npy"))[:, 0]).to(dev)
U = {"o": np.load(os.path.join(CACHE, "he_uni2.npy")), "h": np.load(os.path.join(H, "he_uni2.npy"))[:, 0],
     "t": np.load(os.path.join(CACHE, "allimg_he_uni2.npy"))}
MV = {"o": np.load(os.path.join(CACHE, "he_medvtok.npy")), "h": np.load(os.path.join(H, "he_medvtok.npy"))[:, 0],
      "t": np.load(os.path.join(CACHE, "allimg_he_medvtok.npy"))}
U = {s: f32(x).to(dev) for s, x in U.items()}
MV = {s: f32(x).to(dev) for s, x in MV.items()}
MISS = pd.read_csv(os.path.join(H, "he_img_missing.csv")).idx.values
T = pd.read_csv(os.path.join(CACHE, "alltxt_meta.csv"))
T_IMG = torch.as_tensor(T.img_idx.values, device=dev)
T = T[~pd.read_csv(os.path.join(CACHE, "alltxt_junk.csv")).junk.values]
hvg = pd.read_csv(os.path.join(H, "hvg_genes.csv")).gene.values


def virt_hvg(st):  # = 47_ / 51_ hvg_inputs (virtual RNA, test only)
    g = pd.read_csv(os.path.join(CACHE, "stpath", f"{st}_genes.csv")).gene.values
    c = np.expm1(np.load(os.path.join(CACHE, "stpath", f"{st}_pred.npy")).astype(np.float32)).clip(0)
    c = np.log1p(c / c.sum(1, keepdims=True) * 1e4)
    return c[:, pd.Series(range(len(g)), index=g).loc[hvg].values]


RNA_IN = {e: ({"hest": f32(np.load(os.path.join(H, "hvg.npy"))).to(dev), "orion_v": f32(virt_hvg("orion")).to(dev),
               "tcga_v": f32(virt_hvg("tcga_text")).to(dev)} if e == "hvg" else
              {s: f32(np.load(os.path.join(RF, f"{e}_{s}.npy"))).to(dev) for s in ["hest", "orion_v", "tcga_v"]})
          for e in RNA_E}
TXT_IN = {e: {s: f32(np.load(os.path.join(CACHE, f"{s}_text_{e}.npy"))[:, 0]).to(dev)
              for s in ["alltxt", "orion_gentext"]} for e in TXT_E}
RAW = {"rna:hvg", "text:biomedclip"}  # encoders used without z-scoring (= 47_ / 51_)
DIMS = {"mif": IF_O.shape[-1], **{f"rna:{e}": RNA_IN[e]["hest"].shape[1] for e in RNA_E},
        **{f"text:{e}": TXT_IN[e]["alltxt"].shape[1] for e in TXT_E}}


def one_text_per_image(t_ids, rng):
    pick = T.loc[t_ids].groupby("img_idx").apply(lambda g: g.index[rng.integers(len(g))], include_groups=False)
    return pick.index.values, pick.values


def g100(a, b, seed=0):  # = 51_
    from scipy.optimize import linear_sum_assignment
    a, b = a.cpu(), b.cpu()
    rng = np.random.default_rng(seed)
    out = {m: [] for m in ["R@1", "R@5", "R@10", "MRR", "pair_acc"]}
    for _ in range(DRAWS):
        s_ = torch.from_numpy(rng.choice(len(a), min(G100, len(a)), replace=False))
        sim = (a[s_] @ b[s_].T).numpy(); d = np.diag(sim)
        for S in (sim, sim.T):
            rank = (S > d[:, None]).sum(1)
            for k_ in (1, 5, 10):
                out[f"R@{k_}"].append((rank < k_).mean() / 2)
            out["MRR"].append((1 / (rank + 1)).mean() / 2)
        ri, ci = linear_sum_assignment(-sim)
        out["pair_acc"].append((ci[np.argsort(ri)] == np.arange(len(s_))).mean())
    return {m: float(np.sum(v) / DRAWS) for m, v in out.items()}


def r1(a, b):
    s_ = a @ b.T
    return float((((s_ > s_.diag()[:, None]).sum(1) == 0).float().mean() + ((s_.T > s_.diag()[:, None]).sum(1) == 0).float().mean()) / 2)


def clip_loss(lg, pos=None):
    if pos is None:
        t = torch.arange(len(lg), device=dev)
        return (F.cross_entropy(lg, t) + F.cross_entropy(lg.T, t)) / 2
    tgt = pos / pos.sum(1, keepdim=True)
    return (torch.sum(-tgt * F.log_softmax(lg, 1), 1).mean() + torch.sum(-tgt * F.log_softmax(lg.T, 1), 1).mean()) / 2


def run(k, release=False, test=True):
    """stage 1 + stage 2 (NSEED seeds) on fold k -> (model, use(rule), test results per rule, val curves, info)"""
    s = pd.read_csv(os.path.join(SPLITS, "patient_5fold", f"fold{k}.csv"))
    hh = pd.read_csv(os.path.join(SPLITS, "final", f"hest_fold{k}.csv"))
    sp = pd.read_csv(os.path.join(SPLITS, "final", f"tcga_fold{k}.csv")).set_index("img_idx").split
    tsp = sp.loc[T.img_idx].values
    o = {x: s.tile_id[s.split == x].values for x in ["train", "val", "test"]}
    hs = {x: np.setdiff1d(hh.spot_idx[hh.split == x].values, MISS) for x in ["train", "val", "test"]}
    ts = {x: T.index[tsp == x].values for x in ["train", "val", "test"]}
    if release:  # everything real except the early-stopping val split
        for d in (o, hs, ts):
            d["train"] = np.sort(np.r_[d["train"], d.pop("test")])
    sizes = {n: {x: len(v) for x, v in d.items()} for n, d in [("orion", o), ("hest", hs), ("tcga_text", ts)]}
    print(f"fold {k}{' (release)' if release else ''}: {sizes}", flush=True)

    # ------------------------------------------------ model + normalisation statistics (train only)
    model = Bind4(DIMS, "cat").to(dev)
    model.set_norm("mif", IF_O[o["train"]].mean((0, 1)), IF_O[o["train"]].std((0, 1)) + 1e-3)
    trt = torch.as_tensor(ts["train"], device=dev)
    for e in RNA_E:
        if f"rna:{e}" not in RAW:
            x = RNA_IN[e]["hest"][hs["train"]]
            model.set_norm(f"rna:{e}", x.mean(0), x.std(0) + 1e-6)
    for e in TXT_E:
        if f"text:{e}" not in RAW:
            x = TXT_IN[e]["alltxt"][trt]
            model.set_norm(f"text:{e}", x.mean(0), x.std(0) + 1e-6)
    model.c_mv.fill_(float(U["o"][o["train"]].norm(dim=-1).mean() / MV["o"][o["train"]].norm(dim=-1).mean()))
    Xr = {e: {st: model.norm(f"rna:{e}", v) for st, v in RNA_IN[e].items()} for e in RNA_E}
    Xt = {e: {st: model.norm(f"text:{e}", v) for st, v in TXT_IN[e].items()} for e in TXT_E}
    Xm = model.norm("mif", IF_O)
    o_tr, h_tr = torch.as_tensor(o["train"], device=dev), torch.as_tensor(hs["train"], device=dev)
    vi, vt = (torch.as_tensor(x, device=dev) for x in one_text_per_image(ts["val"], np.random.default_rng(0)))

    def anc(key, idx, views=None):
        x, m = U[key], MV[key]
        if x.dim() == 2:
            return x[idx], m[idx]
        v = views if views is not None else 0
        return x[idx, v], m[idx, v]

    def batch(task, rng):
        if task == "mif":
            b = o_tr[torch.as_tensor(rng.choice(len(o_tr), BS, replace=False), device=dev)]
            d = torch.randint(0, 8, (BS,), device=dev)
            return anc("o", b, d), {"mif": Xm[b, d]}, None
        if task == "rna":
            b = h_tr[torch.as_tensor(rng.choice(len(h_tr), BS, replace=False), device=dev)]
            return anc("h", b), {e: Xr[e]["hest"][b] for e in RNA_E}, None
        b = trt[torch.as_tensor(rng.choice(len(trt), BS, replace=False), device=dev)]
        im = T_IMG[b]
        return (anc("t", im, torch.randint(0, 8, (BS,), device=dev)), {e: Xt[e]["alltxt"][b] for e in TXT_E},
                (im[:, None] == im[None, :]).float())

    VAL = {"mif": (anc("o", o["val"]), {"mif": Xm[o["val"], 0]}),
           "rna": (anc("h", hs["val"]), {e: Xr[e]["hest"][hs["val"]] for e in RNA_E}),
           "text": (anc("t", vi), {e: Xt[e]["alltxt"][vt] for e in TXT_E})}

    # ------------------------------------------------ stage 1 (= 51_ train_head, UNI2 anchor)
    def train_head(key, task, enc, steps):
        torch.manual_seed(k); rng = np.random.default_rng(k)
        h = model.h(key)  # initialised exactly like 51_: a fresh module created right after the same seed
        h.load_state_dict((mlp(h[0].in_features) if key == "mif" else torch.nn.Linear(h.in_features, h.out_features)).state_dict())
        logit = model.logit[model._k(key)]
        opt = torch.optim.AdamW(list(h.parameters()) + [logit], lr=LR, weight_decay=0.05)
        sch = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=LR, total_steps=steps, pct_start=0.05)
        za_va, x_va = nz(VAL[task][0][0]), VAL[task][1][enc]
        best, state = -1, None
        for step in range(1, steps + 1):
            h.train()
            (u, _), X, pos = batch(task, rng)
            loss = clip_loss(logit.exp().clamp(max=100) * nz(u) @ nz(h(X[enc])).T, pos)
            opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sch.step()
            if step % EVAL_EVERY == 0 or step == steps:
                h.eval()
                with torch.no_grad():
                    v = r1(za_va, nz(h(x_va)))
                if v > best:
                    best, state = v, ({k_: x.detach().clone() for k_, x in h.state_dict().items()}, logit.detach().clone())
        h.load_state_dict(state[0]); h.eval()
        with torch.no_grad():
            logit.copy_(state[1])
        return best

    stage1 = {"mif": train_head("mif", "mif", "mif", 400)}
    for e in RNA_E:
        stage1[f"rna:{e}"] = train_head(f"rna:{e}", "rna", e, 2700)
    for e in TXT_E:
        stage1[f"text:{e}"] = train_head(f"text:{e}", "text", e, 1000)
    for p in model.heads.parameters():
        p.requires_grad_(False)
    model.heads.eval()
    logit0 = {k_: v.detach().clone() for k_, v in model.logit.items()}
    print(f"[{time.time() - t0:.0f}s] fold {k} stage 1 val R@1: {json.dumps({a: round(b, 3) for a, b in stage1.items()})}",
          flush=True)

    # ------------------------------------------------ stage 2 (= 51_ joint, mode 'ad', adapter 'cat'), NSEED seeds
    def fuse_n(mod, encs, X):  # inputs already normalised
        return nz(sum(nz(model.h(f"{mod}:{e}")(X[e])) for e in encs))

    def val_score():
        model.eval()
        with torch.no_grad():
            v = {"mif": r1(nz(model.he(*VAL["mif"][0])), nz(model.h("mif")(VAL["mif"][1]["mif"]))),
                 "rna": r1(nz(model.he(*VAL["rna"][0])), fuse_n("rna", RNA_E, VAL["rna"][1])),
                 "text": r1(nz(model.he(*VAL["text"][0])), fuse_n("text", TXT_E, VAL["text"][1]))}
        return float(np.exp(np.mean(np.log([max(x, 1e-4) for x in v.values()])))), v

    def snap():
        return ({k_: x.detach().clone() for k_, x in model.ah[0].state_dict().items()},
                {k_: x.detach().clone() for k_, x in model.logit.items()})

    def stage2(seed):
        torch.manual_seed(seed); rng = np.random.default_rng(seed)
        model.ah = torch.nn.ModuleList([model.new_adapter()]).to(dev)
        with torch.no_grad():
            for k_, v in model.logit.items():
                v.copy_(logit0[k_])
        opt = torch.optim.AdamW(list(model.ah.parameters()) + list(model.logit.parameters()), lr=1e-4, weight_decay=0.05)
        sch = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=1e-4, total_steps=1000, pct_start=0.05)
        best, es, curve = -1, None, []
        for step in range(1, 1001):
            model.ah.train()
            loss = 0
            for task in ("mif", "rna", "text"):
                (u, mv), X, pos = batch(task, rng)
                za = nz(model.he(u, mv))
                if task == "mif":
                    loss = loss + clip_loss(model.lg("mif") * za @ nz(model.h("mif")(X["mif"])).T, pos)
                else:
                    encs = RNA_E if task == "rna" else TXT_E
                    loss = loss + sum(clip_loss(model.lg(f"{task}:{e}") * za @ nz(model.h(f"{task}:{e}")(X[e])).T, pos)
                                      for e in encs) / len(encs)
                un = nz(u)
                loss = loss + F.mse_loss(za @ za.T, un @ un.T)
            opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sch.step()
            if step % EVAL_EVERY == 0 or step == 1000:
                sc, v = val_score()
                curve.append({"seed": seed, "step": step, "geo": sc, **v})
                if sc > best:
                    best, es = sc, (step, *snap())
        return es, (1000, *snap()), curve

    seeds = [k] + [1000 * i + k for i in range(1, NSEED)]  # the first = 51_'s seed
    runs = [stage2(sd) for sd in seeds]
    curve = pd.DataFrame([c for r in runs for c in r[2]])
    sel = {"es": [runs[0][0]], "last": [runs[0][1]], "ens-es": [r[0] for r in runs], "ens-last": [r[1] for r in runs],
           **{f"es@seed{i}": [r[0]] for i, r in enumerate(runs)}, **{f"last@seed{i}": [r[1]] for i, r in enumerate(runs)}}
    print(f"[{time.time() - t0:.0f}s] fold {k} stage 2: early-stopping step per seed {[r[0][0] for r in runs]}", flush=True)

    def use(rule):
        """set the model's adapter(s) to a selection rule (temperatures: those of the first member; unused at inference)"""
        st = sel[rule]
        model.ah = torch.nn.ModuleList([model.new_adapter() for _ in st]).to(dev)
        for a, (_, sd_, _) in zip(model.ah, st):
            a.load_state_dict(sd_)
        with torch.no_grad():
            for k_, v in model.logit.items():
                v.copy_(st[0][2][k_])
        model.eval()

    res = {}
    if test:
        ii, tt = (torch.as_tensor(x, device=dev) for x in one_text_per_image(ts["test"], np.random.default_rng(k)))
        for rule in ["stage1"] + list(sel):
            if rule == "stage1":
                model.ah = torch.nn.ModuleList([model.new_adapter()]).to(dev)  # zero output: H&E = UNI2 (= E of 51_)
            else:
                use(rule)
            with torch.no_grad():
                R = lambda st, idx: fuse_n("rna", RNA_E, {e: Xr[e][st][idx] for e in RNA_E})
                Tx = lambda st, idx: fuse_n("text", TXT_E, {e: Xt[e][st][idx] for e in TXT_E})
                zm = nz(model.h("mif")(Xm[o["test"], 0]))
                zt = Tx("alltxt", tt)
                res[rule] = {"H&E<->mIF (ORION)": g100(nz(model.he(*anc("o", o["test"]))), zm),
                             "H&E<->RNA (HEST)": g100(nz(model.he(*anc("h", hs["test"]))), R("hest", hs["test"])),
                             "H&E<->text (TCGA)": g100(nz(model.he(*anc("t", ii))), zt),
                             "mIF<->RNA(virtual) (ORION)": g100(zm, R("orion_v", o["test"])),
                             "mIF<->text(synthetic) (ORION)": g100(zm, Tx("orion_gentext", o["test"])),
                             "text<->RNA(virtual) (TCGA)": g100(zt, R("tcga_v", ii)),
                             "text<->mIF(MIPHEI) (TCGA)": g100(zt, nz(model.h("mif")(model.norm("mif", IF_V_T[ii]))))}
        res["_items"] = dict(orion=o["test"], hest=hs["test"], tcga_img=ii.cpu().numpy(), tcga_txt=tt.cpu().numpy())
    info = {"sizes": sizes, "stage1_val_R@1": stage1, "es_steps": [r[0][0] for r in runs], "seeds": seeds}
    return model, use, res, curve, info


if MODE == "cv":
    TAG = f"_{sys.argv[sys.argv.index('--out') + 1]}" if "--out" in sys.argv else ""  # results/final_cv_<x><TAG>.csv
    rows = []
    for k in range(5):
        _, _, res, curve, info = run(k)
        res.pop("_items")
        for rule, rr in res.items():
            for pair, r in rr.items():
                rows.append({"rule": rule, "fold": k, "pair": pair, **r})
        D = pd.DataFrame(rows)
        D.to_csv(os.path.join(ROOT, "results", f"final_cv_long{TAG}.csv"), index=False)
        curve.assign(fold=k).to_csv(os.path.join(ROOT, "results", f"final_cv_curves{TAG}.csv"), mode="w" if k == 0 else "a",
                                    header=k == 0, index=False)
        print(D[D.fold == k].pivot_table(index="rule", columns="pair", values="R@1", sort=False).round(3).to_string(), flush=True)
    D = pd.DataFrame(rows)
    S = D.groupby(["rule", "pair"], sort=False)[["R@1", "pair_acc"]].agg(["mean", "std"]).round(4)
    S.to_csv(os.path.join(ROOT, "results", f"final_cv_summary{TAG}.csv"))
    pd.set_option("display.width", 250)
    print(D.pivot_table(index="rule", columns="pair", values="R@1", sort=False).round(3).to_string())
    for r in ["es", "last"]:  # seed-to-seed spread of the single-seed rules (std over seeds, mean over folds)
        sub = D[D.rule.str.startswith(f"{r}@seed")]
        print(r, "seed std of R@1 (mean over folds):",
              sub.groupby(["pair", "fold"])["R@1"].std().groupby("pair").mean().round(4).to_dict())
    sys.exit()

# ---------------------------------------------------------------- fold0 / release: build, save
release = MODE == "release"
OUT = os.path.join(ROOT, "models", sys.argv[sys.argv.index("--out") + 1] if "--out" in sys.argv else
                   "ej_ad_mv" if release else "ej_ad_mv_fold0")
os.makedirs(OUT, exist_ok=True)
model, use, res, curve, info = run(0, release=release, test=not release)
use(RULE)
curve.round(4).to_csv(os.path.join(OUT, "val_curve.csv"), index=False)
torch.save({k_: v.cpu() for k_, v in model.state_dict().items()}, os.path.join(OUT, "model.pt"))
meta = {"config": "EJ-ad+MV", "stage2_rule": RULE, "n_adapters": len(model.ah), "mode": MODE,
        "created": time.strftime("%Y-%m-%d %H:%M"), "dims": DIMS, "adapter": "cat", "raw_encoders": sorted(RAW),
        "hvg_genes": hvg.tolist(), "c_mv": float(model.c_mv),
        "split": "fold0 of splits/patient_5fold (ORION), splits/final/hest_fold0, tcga_fold0"
                 + ("; train = train + test, val = val (early stopping)" if release else "; test held out"),
        **info, "anchor": "UNI2-h CLS (1536-d), frozen", "image_extra": "MedVTok feats (1152-d), frozen",
        "mif_markers": "embed.MIF_MARKERS", "rna_encoders": RNA_E, "text_encoders": TXT_E}
json.dump(meta, open(os.path.join(OUT, "meta.json"), "w"), indent=1)
print(f"saved {OUT} (rule {RULE}, {len(model.ah)} adapter(s))", flush=True)
if not release:
    items = res.pop("_items")
    np.savez(os.path.join(OUT, "test_items.npz"), **items)
    ref = pd.read_csv(os.path.join(ROOT, "results", "medvtok_long.csv"))
    ref = ref[(ref.config == "EJ-ad+MV") & (ref.fold == 0)].set_index("pair")
    D = pd.DataFrame([{"pair": p, "R@1": r["R@1"], "pair_acc": r["pair_acc"], "R@1_51_fold0": ref.loc[p, "R@1"]}
                      for p, r in res[RULE].items()])
    D.to_csv(os.path.join(OUT, "test.csv"), index=False)
    print(D.round(3).to_string(), flush=True)