"""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//{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_.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)