| """Final four-modality evaluation with multi-encoder late fusion (frozen UNI2 anchor; real pairs only; virtual / |
| synthetic data only at test time). Combines the winners of 44_rnafm_eval.py (RNA) and 46_text_enc_eval.py (text). |
| |
| Every head maps one frozen encoder into the frozen UNI2 space (S4 recipe of 42_: MLP for mIF, linear for RNA / text), so |
| the heads of several encoders of one modality can be averaged at test time: z_mod = normalise(sum_e normalise(head_e(x_e))). |
| Configurations (mIF = VirTues 16-ch + MLP head in all of them): |
| S4 RNA hvg text BiomedCLIP (current method) |
| single RNA AIDO.Cell text KEEP (best single encoders) |
| ens-small RNA hvg + Tahoe-x1 + SE-600M + AIDO.Cell text KEEP + BiomedCLIP + Qwen3-Emb-0.6B |
| ens-large RNA all 10 encoders text KEEP + BiomedCLIP + Qwen3-0.6B + PubMedBERT-emb + MedCPT |
| Pairs per fold k (same splits / test items as 42_ / 44_ / 46_): trained H&E<->mIF (ORION), H&E<->RNA (HEST), |
| H&E<->text (TCGA); emergent mIF<->RNA(v), mIF<->text(s) (ORION), text<->RNA(v), text<->mIF(MIPHEI) (TCGA); |
| external DLPFC (RNA heads retrained without spatialLIBD, inputs z-scored per dataset): H&E<->RNA, RNA<->text(s), |
| RNA<->mIF(MIPHEI). |
| Output results/final_ens_long.csv, results/final_ens_summary.csv, results/final_ens_tests.csv (paired vs S4) |
| """ |
| import os |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from scipy import stats |
| from scipy.optimize import linear_sum_assignment |
|
|
| from common import CACHE, ROOT, SPLITS |
|
|
| 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) |
| RNA_ALL = ["hvg", "tahoex1", "se600m", "scgpt", "nicheformer", "aidocell", "geneformer", "c2s", "genept-bge", "genept-bmc"] |
| CONF = {"S4": (["hvg"], ["biomedclip"]), |
| "single": (["aidocell"], ["keep"]), |
| "ens-small": (["hvg", "tahoex1", "se600m", "aidocell"], ["keep", "biomedclip", "qwen3-0.6b"]), |
| "ens-large": (RNA_ALL, ["keep", "biomedclip", "qwen3-0.6b", "pubmedbert", "medcpt"])} |
|
|
| |
| HE_O = f32(np.load(os.path.join(CACHE, "he_uni2.npy"))).to(dev) |
| 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) |
| IF_V_D = f32(np.load(os.path.join(CACHE, "if_virtues_m16_virtual_miphei_dlpfc.npy"))[:, 0]).to(dev) |
| HE_H = f32(np.load(os.path.join(H, "he_uni2.npy"))[:, 0]).to(dev) |
| HE_T = f32(np.load(os.path.join(CACHE, "allimg_he_uni2.npy"))).to(dev) |
| HE_D = f32(np.load(os.path.join(CACHE, "external", "dlpfc_he_uni2.npy"))).to(dev) |
| 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] |
| HM = pd.read_csv(os.path.join(H, "meta.csv")) |
| hvg = pd.read_csv(os.path.join(H, "hvg_genes.csv")).gene.values |
|
|
|
|
| def hvg_inputs(): |
| def virt(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] |
| eg = pd.read_csv(os.path.join(CACHE, "external", "dlpfc_genes_full.csv")).gene.values |
| ec = np.load(os.path.join(CACHE, "external", "dlpfc_counts_full.npy")).astype(np.float32) |
| ei = pd.Series(range(len(eg)), index=eg); ei = ei[~ei.index.duplicated()] |
| d = np.log1p(ec / ec.sum(1, keepdims=True).clip(1) * 1e4) |
| d = np.stack([d[:, ei[g]] if g in ei.index else np.zeros(len(d), np.float32) for g in hvg], 1) |
| return {"hest": np.load(os.path.join(H, "hvg.npy")), "orion_v": virt("orion"), "tcga_v": virt("tcga_text"), "dlpfc": d} |
|
|
|
|
| RNA_IN = {e: ({s: f32(x).to(dev) for s, x in hvg_inputs().items()} 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", "dlpfc"]}) |
| for e in RNA_ALL} |
| TXT_ENCS = sorted({e for _, x in CONF.values() for e in x}) |
| 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", "dlpfc_gentext"]} for e in TXT_ENCS} |
|
|
|
|
| def mlp(din, dout=1536, hid=512): |
| return nn.Sequential(nn.Linear(din, hid), nn.GELU(), nn.Dropout(0.1), nn.Linear(hid, hid), nn.GELU(), |
| nn.LayerNorm(hid), nn.Linear(hid, dout)) |
|
|
|
|
| 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): |
| 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 train_head(X_tr, A_tr, X_va, A_va, kind, steps, seed, pos_tr=None): |
| torch.manual_seed(seed); rng = np.random.default_rng(seed) |
| h = (mlp(X_tr.shape[1]) if kind == "mlp" else nn.Linear(X_tr.shape[1], 1536)).to(dev) |
| logit = nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32, device=dev)) |
| 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 = nz(A_va[:, 0] if A_va.dim() == 3 else A_va) |
| best, state = -1, None |
| for step in range(1, steps + 1): |
| h.train() |
| b = torch.as_tensor(rng.choice(len(X_tr), BS, replace=False), device=dev) |
| a = A_tr[b, torch.randint(0, A_tr.shape[1], (BS,), device=dev)] if A_tr.dim() == 3 else A_tr[b] |
| lg = logit.exp().clamp(max=100) * nz(a) @ nz(h(X_tr[b])).T |
| if pos_tr is None: |
| t = torch.arange(BS, device=dev) |
| loss = (F.cross_entropy(lg, t) + F.cross_entropy(lg.T, t)) / 2 |
| else: |
| ii = pos_tr[b]; pos = (ii[:, None] == ii[None, :]).float(); tgt = pos / pos.sum(1, keepdim=True) |
| loss = (torch.sum(-tgt * F.log_softmax(lg, 1), 1).mean() + torch.sum(-tgt * F.log_softmax(lg.T, 1), 1).mean()) / 2 |
| 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()} |
| h.load_state_dict(state); h.eval() |
| return h |
|
|
|
|
| def zs(X, ref_set, idx, raw, ext=None): |
| """z-score every set with the training rows of ref_set (raw = keep the baseline input unchanged, as in 42_); |
| ext: sets z-scored on themselves (external data, = 38_ / 44_)""" |
| if raw: |
| out = dict(X) |
| else: |
| mu, sd = X[ref_set][idx].mean(0), X[ref_set][idx].std(0) + 1e-6 |
| out = {s: (v - mu) / sd for s, v in X.items()} |
| for s in ext or []: |
| out[s] = (X[s] - X[s].mean(0)) / (X[s].std(0) + 1e-6) |
| return out |
|
|
|
|
| rows = [] |
| for k in range(5): |
| s = pd.read_csv(os.path.join(SPLITS, "patient_5fold", f"fold{k}.csv")) |
| o = {x: s.tile_id[s.split == x].values for x in ["train", "val", "test"]} |
| nif = lambda x: (x - IF_O[o["train"]].mean((0, 1))) / (IF_O[o["train"]].std((0, 1)) + 1e-3) |
| h = pd.read_csv(os.path.join(SPLITS, "final", f"hest_fold{k}.csv")) |
| hs = {x: h.spot_idx[h.split == x].values for x in ["train", "val", "test"]} |
| 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 |
| ts = {x: T.index[tsp == x].values for x in ["train", "val", "test"]} |
| vi, vt = (torch.as_tensor(x, device=dev) for x in one_text_per_image(ts["val"], np.random.default_rng(0))) |
| ii, tt = (torch.as_tensor(x, device=dev) for x in one_text_per_image(ts["test"], np.random.default_rng(k))) |
| trt = torch.as_tensor(ts["train"], device=dev) |
| pool = np.flatnonzero(~HM.group.str.startswith("spatialLIBD").values) |
| grp = np.random.default_rng(k).permutation(sorted(HM.group.iloc[pool].unique())) |
| isva = HM.group.iloc[pool].isin(set(grp[: max(1, len(grp) // 10)])).values |
| d_tr, d_va = pool[~isva], pool[isva] |
|
|
| hm = mlp(512).to(dev) |
| hm.load_state_dict(torch.load(os.path.join(ROOT, "runs", "rnafm", f"other_fold{k}.pt"), map_location=dev)["mif"]) |
| hm.eval() |
| with torch.no_grad(): |
| Zm = {"he_o": nz(HE_O[o["test"], 0]), "orion": nz(hm(nif(IF_O[o["test"], 0]))), |
| "tcga_v": nz(hm(nif(IF_V_T[ii]))), "dlpfc_v": nz(hm(nif(IF_V_D)))} |
| Zr, Zt = {}, {} |
| for e in RNA_ALL: |
| X = zs(RNA_IN[e], "hest", hs["train"], e == "hvg") |
| hr = train_head(X["hest"][hs["train"]], HE_H[hs["train"]], X["hest"][hs["val"]], HE_H[hs["val"]], "linear", 2700, k) |
| Xd = zs(RNA_IN[e], "hest", d_tr, False, ext=["dlpfc"]) |
| hd = train_head(Xd["hest"][d_tr], HE_H[d_tr], Xd["hest"][d_va], HE_H[d_va], "linear", 2700, k) |
| with torch.no_grad(): |
| Zr[e] = {"hest": nz(hr(X["hest"][hs["test"]])), "orion_v": nz(hr(X["orion_v"][o["test"]])), |
| "tcga_v": nz(hr(X["tcga_v"][ii])), "dlpfc": nz(hd(Xd["dlpfc"]))} |
| print(f"fold {k} RNA head {e} done", flush=True) |
| for e in TXT_ENCS: |
| X = zs(TXT_IN[e], "alltxt", trt, e == "biomedclip") |
| ht = train_head(X["alltxt"][trt], HE_T[T_IMG[trt]], X["alltxt"][vt], HE_T[vi, 0], "linear", 1000, k, pos_tr=T_IMG[trt]) |
| |
| with torch.no_grad(): |
| Zt[e] = {"tcga": nz(ht(X["alltxt"][tt])), "orion_s": nz(ht(X["orion_gentext"][o["test"]])), |
| "dlpfc_s": nz(ht(X["dlpfc_gentext"]))} |
| he_h, he_t, he_d = nz(HE_H[hs["test"]]), nz(HE_T[ii, 0]), nz(HE_D) |
| for cname, (re_, te_) in CONF.items(): |
| R = {s_: nz(sum(Zr[e][s_] for e in re_)) for s_ in Zr[re_[0]]} |
| Tx = {s_: nz(sum(Zt[e][s_] for e in te_)) for s_ in Zt[te_[0]]} |
| res = {("H&E<->mIF (ORION)", "trained"): g100(Zm["he_o"], Zm["orion"]), |
| ("H&E<->RNA (HEST)", "trained"): g100(he_h, R["hest"]), |
| ("H&E<->text (TCGA)", "trained"): g100(he_t, Tx["tcga"]), |
| ("mIF<->RNA(virtual) (ORION)", "emergent"): g100(Zm["orion"], R["orion_v"]), |
| ("mIF<->text(synthetic) (ORION)", "emergent"): g100(Zm["orion"], Tx["orion_s"]), |
| ("text<->RNA(virtual) (TCGA)", "emergent"): g100(Tx["tcga"], R["tcga_v"]), |
| ("text<->mIF(MIPHEI) (TCGA)", "emergent"): g100(Tx["tcga"], Zm["tcga_v"]), |
| ("H&E<->RNA (DLPFC)", "external"): g100(he_d, R["dlpfc"]), |
| ("RNA<->text(synthetic) (DLPFC)", "external"): g100(R["dlpfc"], Tx["dlpfc_s"]), |
| ("RNA<->mIF(MIPHEI) (DLPFC)", "external"): g100(R["dlpfc"], Zm["dlpfc_v"])} |
| for (pair, kind), r in res.items(): |
| rows.append({"config": cname, "fold": k, "pair": pair, "kind": kind, **r}) |
| print(f"fold {k} {cname:<10} " + " ".join(f"{p.split(' (')[0]} {r['R@1']:.3f}" for (p, _), r in res.items()), flush=True) |
| pd.DataFrame(rows).to_csv(os.path.join(ROOT, "results", "final_ens_long.csv"), index=False) |
|
|
| R = pd.DataFrame(rows) |
| S = R.groupby(["config", "kind", "pair"], sort=False)[["R@1", "R@10", "pair_acc"]].agg(["mean", "std"]).round(4) |
| S.to_csv(os.path.join(ROOT, "results", "final_ens_summary.csv")) |
| tests = [] |
| for c in CONF: |
| if c == "S4": |
| continue |
| for p in R.pair.unique(): |
| a = R[(R.config == c) & (R.pair == p)].sort_values("fold")["R@1"].values |
| b = R[(R.config == "S4") & (R.pair == p)].sort_values("fold")["R@1"].values |
| tests.append({"config": c, "pair": p, "S4": b.mean(), "new": a.mean(), "delta": (a - b).mean(), |
| "folds_won": int((a > b).sum()), "p_paired_t": stats.ttest_rel(a, b).pvalue}) |
| pd.DataFrame(tests).round(4).to_csv(os.path.join(ROOT, "results", "final_ens_tests.csv"), index=False) |
| pd.set_option("display.width", 250) |
| print(R.pivot_table(index="config", columns="pair", values="R@1", sort=False).round(3).T.to_string()) |
| print(pd.DataFrame(tests).round(4).to_string()) |
|
|
