| """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 |
| from embed import RNA_E, TXT_E, Bind4, mlp |
|
|
| 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" |
| 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() |
|
|
| |
| 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): |
| 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"} |
| 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): |
| 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: |
| 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 = 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})} |
|
|
| |
| def train_head(key, task, enc, steps): |
| torch.manual_seed(k); rng = np.random.default_rng(k) |
| h = model.h(key) |
| 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) |
|
|
| |
| def fuse_n(mod, encs, X): |
| 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)] |
| 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) |
| 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 "" |
| 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"]: |
| 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() |
|
|
| |
| 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) |
|
|
