"""Validate the 4-modality ImageBind model with VIRTUAL modalities (validation only; every encoder/head was trained on real pairs: H&E-mIF ORION (17_), H&E-text PathGen (12_), H&E-RNA HEST (23_); H&E = frozen UNI2). Virtual data fill the modality pairs that no dataset measures together, giving quadruplets: ORION : H&E (real) | mIF (real) | RNA (STPath, virtual) | text (GPT synthetic caption of the H&E, 27_/32_) TCGA : H&E (real) | text (real) | mIF (GigaTIME, virtual) | RNA (STPath, virtual) HEST : H&E (real) | RNA (real) + RNA (STPath, virtual) -> calibration of the virtual RNA ORION : mIF real vs mIF virtual (GigaTIME) on the same tiles -> calibration of the virtual mIF Every embedding lives in the UNI2 space, so any two modalities are compared by cosine similarity. Heads: RNA runs/gene/uni2_mlp_fold/head.pt (MLP on log1p-CP10k of the 2000 HEST HVGs; HEST fold k) mIF runs/bench_fm/patient_5fold_fold_uni2_fm-virtues[_shared]_fm-geneformer_pathgenclip/heads.pt (17-channel VirTues for real ORION; 10-channel shared panel (30_) for real-vs-virtual / TCGA) text runs/text_pathgen/uni2_biomedclip/head.pt (PathGen captions); TCGA evaluation images exclude every patient of its train/val split Metrics on random 100-item galleries (50 draws): R@1 (tie-aware, both directions averaged) and one-to-one Hungarian pairing accuracy (chance 0.01). "circular" = both sides are derived from the same H&E image (virtual-from-H&E), so the score is an upper bound of what real data would give, not evidence of biological agreement. Outputs results/virtual_validation.csv """ 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.optimize import linear_sum_assignment from common import CACHE, ROOT, RUNS, SPLITS, load_all_tiles OUT = os.path.join(ROOT, "results") H = os.path.join(CACHE, "hest") G100, DRAWS = 100, 50 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 load_mlp(sd): h = mlp(sd["0.weight"].shape[1]); h.load_state_dict(sd); return h.eval() def nz(x): return F.normalize(torch.as_tensor(np.asarray(x), dtype=torch.float32), dim=-1) def g100(a, b, seed=0): """mean over random 100-galleries of tie-aware R@1 (both directions) and Hungarian pairing accuracy""" rng = np.random.default_rng(seed) r, p = [], [] 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)[:, None] r.append(((sim > d).sum(1) == 0).mean() / 2 + ((sim.T > d).sum(1) == 0).mean() / 2) ri, ci = linear_sum_assignment(-sim) p.append((ci[np.argsort(ri)] == np.arange(len(s))).mean()) return float(np.mean(r)), float(np.mean(p)) rows = [] def add(dataset, fold, pair, za, zb, kind, circular): r, p = g100(za, zb) rows.append({"dataset": dataset, "fold": fold, "pair": pair, "kind": kind, "circular": circular, "n": len(za), "R@1_g100": r, "pair_acc_g100": p}) # ---------------- RNA: STPath log1p(counts) -> log1p CP10k on the HEST HVGs (same input as the HEST-trained head) hvg = pd.read_csv(os.path.join(H, "hvg_genes.csv")).gene.values def virtual_rna(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 torch.from_numpy(c[:, pd.Series(range(len(g)), index=g).loc[hvg].values]) rna_heads = [load_mlp(torch.load(os.path.join(RUNS, "gene", f"uni2_mlp_fold{k}", "head.pt"), map_location="cpu")) for k in range(5)] text_head = load_mlp(torch.load(os.path.join(RUNS, "text_pathgen", "uni2_biomedclip", "head.pt"), map_location="cpu")) def mif_heads(enc, k, split="patient_5fold"): f = os.path.join(RUNS, "bench_fm", f"{split}_fold{k}_uni2_fm-{enc}_fm-geneformer_pathgenclip", "heads.pt") if not os.path.exists(f): return None d = torch.load(f, map_location="cpu") h = load_mlp(d["state"]["if"]) return lambda x: h((torch.as_tensor(x, dtype=torch.float32) - d["if_mean"]) / d["if_std"]) # ---------------- (A) HEST: H&E <-> real RNA vs H&E <-> virtual RNA, same spots, HEST test folds M = pd.read_csv(os.path.join(H, "meta.csv")) hv = pd.read_csv(os.path.join(CACHE, "stpath", "hest_val_meta.csv")) assert (hv["sample"].values == M["sample"].values).all() and (hv.spot.values == M.spot.values).all() he_h = np.load(os.path.join(H, "he_uni2.npy"))[:, 0].astype(np.float32) real_h = torch.from_numpy(np.load(os.path.join(H, "hvg.npy")).astype(np.float32)) virt_h = virtual_rna("hest_val") with torch.no_grad(): for k in range(5): te = torch.load(os.path.join(RUNS, "gene", f"uni2_mlp_fold{k}", "test.pt"), weights_only=False)["idx"] zr, zv, zh = nz(rna_heads[k](real_h[te])), nz(rna_heads[k](virt_h[te])), nz(he_h[te]) for sub, m in [("all", np.ones(len(te), bool)), ("STPath-heldout", hv.hest_bench_heldout.values[te])]: if m.sum() < G100: continue add(f"HEST ({sub})", k, "H&E<->RNA real", zh[m], zr[m], "trained (real)", False) add(f"HEST ({sub})", k, "H&E<->RNA virtual", zh[m], zv[m], "virtual calibration", True) add(f"HEST ({sub})", k, "RNA real<->RNA virtual", zr[m], zv[m], "virtual calibration", False) print("HEST done", flush=True) # ---------------- (B) ORION: patient 5-fold test tiles df = load_all_tiles() assert (pd.read_csv(os.path.join(CACHE, "stpath", "orion_meta.csv")).key.values == df.image_path.values).all() he_o = np.load(os.path.join(CACHE, "he_uni2.npy"))[:, 0].astype(np.float32) rna_o = virtual_rna("orion") if17 = np.load(os.path.join(CACHE, "if_virtues.npy"))[:, 0].astype(np.float32) f_sh = os.path.join(CACHE, "if_virtues_shared.npy") f_vo = os.path.join(CACHE, "if_virtues_shared_virtual_orion.npy") ifs_r = np.load(f_sh)[:, 0].astype(np.float32) if os.path.exists(f_sh) else None ifs_v = np.load(f_vo)[:, 0].astype(np.float32) if os.path.exists(f_vo) else None f_gt = os.path.join(CACHE, "orion_gentext_text_biomedclip.npy") # synthetic GPT captions of the ORION H&E (32_ + 08_) gt_o = torch.from_numpy(np.load(f_gt)[:, 0].astype(np.float32)) if os.path.exists(f_gt) else None with torch.no_grad(): for k in range(5): s = pd.read_csv(os.path.join(SPLITS, "patient_5fold", f"fold{k}.csv")) te = s.tile_id[s.split == "test"].values zh, zr = nz(he_o[te]), nz(rna_heads[k](rna_o[te])) add("ORION", k, "H&E<->RNA virtual", zh, zr, "virtual (cross-dataset head)", True) h17 = mif_heads("virtues", k) if h17 is not None: zp = nz(h17(if17[te])) add("ORION", k, "H&E<->mIF real (17 ch)", zh, zp, "trained (real)", False) add("ORION", k, "mIF real (17 ch)<->RNA virtual", zp, zr, "emergent, real x virtual", False) hs = mif_heads("virtues_shared", k) if hs is not None and ifs_r is not None: zs = nz(hs(ifs_r[te])) add("ORION", k, "H&E<->mIF real (10 ch)", zh, zs, "trained (real)", False) add("ORION", k, "mIF real (10 ch)<->RNA virtual", zs, zr, "emergent, real x virtual", False) if ifs_v is not None: zv = nz(hs(ifs_v[te])) add("ORION", k, "H&E<->mIF virtual (10 ch)", zh, zv, "virtual calibration", True) add("ORION", k, "mIF real<->mIF virtual (10 ch)", zs, zv, "virtual calibration", False) add("ORION", k, "mIF virtual<->RNA virtual", zv, zr, "emergent, virtual x virtual", True) if gt_o is not None: # synthetic text is written from the same H&E -> circular with H&E and H&E-derived data zt = nz(text_head(gt_o[te])) add("ORION", k, "H&E<->text synthetic", zh, zt, "virtual (text head from PathGen)", True) add("ORION", k, "text synthetic<->RNA virtual", zt, zr, "emergent, virtual x virtual", True) if h17 is not None: add("ORION", k, "text synthetic<->mIF real (17 ch)", zt, zp, "emergent, real x virtual", False) if hs is not None and ifs_r is not None: add("ORION", k, "text synthetic<->mIF real (10 ch)", zt, zs, "emergent, real x virtual", False) print("ORION done", flush=True) # ---------------- (C) TCGA quadruplets: H&E, real text, virtual mIF, virtual RNA I = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv")) T = pd.read_csv(os.path.join(CACHE, "alltxt_meta.csv")) P = pd.read_csv(os.path.join(CACHE, "pathgen_meta.csv")) seen = set(P.patient[P.split.isin(["train", "val"])]) ev = I.index[~I.patient.isin(seen)].values rng = np.random.default_rng(0) t1 = T[T.img_idx.isin(ev)].groupby("img_idx").apply(lambda g: g.index[rng.integers(len(g))], include_groups=False) ev = t1.index.values # images with >= 1 text, patients unseen by the text head print(f"TCGA eval: {len(ev)} images, {I.patient.iloc[ev].nunique()} patients (excluded {len(seen)} text-head patients)", flush=True) he_t = np.load(os.path.join(CACHE, "allimg_he_uni2.npy"))[ev, 0].astype(np.float32) tx = torch.from_numpy(np.load(os.path.join(CACHE, "alltxt_text_biomedclip.npy"))[t1.values, 0].astype(np.float32)) rna_t = virtual_rna("tcga_text")[ev] f_vt = os.path.join(CACHE, "if_virtues_shared_virtual_tcga_text.npy") ifv_t = np.load(f_vt)[ev, 0].astype(np.float32) if os.path.exists(f_vt) else None with torch.no_grad(): zh, zt = nz(he_t), nz(text_head(tx)) for k in range(5): zr = nz(rna_heads[k](rna_t)) Z = {"H&E": zh, "text": zt, "RNA virtual": zr} hs = mif_heads("virtues_shared", k) if hs is not None and ifv_t is not None: Z["mIF virtual"] = nz(hs(ifv_t)) names = list(Z) for i in range(len(names)): for j in range(i + 1, len(names)): a, b = names[i], names[j] if (a, b) == ("H&E", "text") and k > 0: continue # does not depend on the fold heads virt = [x for x in (a, b) if "virtual" in x] circ = len(virt) == 2 or (len(virt) == 1 and "H&E" in (a, b)) kind = "trained (real)" if (a, b) == ("H&E", "text") else \ "emergent, real x virtual" if len(virt) == 1 and "H&E" not in (a, b) else \ "emergent, virtual x virtual" if len(virt) == 2 else "virtual (cross-dataset head)" add("TCGA", k, f"{a}<->{b}", Z[a], Z[b], kind, circ) print("TCGA done", flush=True) # ---------------- (D) mIF as 10 channel MEANS (GigaTIME per-tile means are reliable, its pixel maps are not): # MLP head on rank-gaussianised shared-panel means, trained on REAL ORION (patient_5fold fold k train/val) vs UNI2 H&E; # applied unchanged to virtual means (rank-gaussianised within their own set) import json from scipy.special import ndtri from scipy.stats import rankdata rg = lambda x: torch.from_numpy(ndtri((np.apply_along_axis(rankdata, 0, x) - 0.5) / len(x)).astype(np.float32)) ORION_CH = ["Hoechst", "CD31", "CD45", "CD68", "CD4", "FOXP3", "CD8a", "CD45RO", "CD20", "PD-L1", "CD3e", "CD163", "E-cadherin", "PD-1", "Ki67", "Pan-CK", "SMA"] # order of orion_real_if_means.npy (25_) PANEL = [("Hoechst", ["DAPI"]), ("CD68", ["CD68"]), ("CD4", ["CD4"]), ("CD8a", ["CD8"]), ("CD20", ["CD20"]), ("PD-L1", ["PD-L1"]), ("CD3e", ["CD3"]), ("Ki67", ["Ki67"]), ("Pan-CK", ["CK"]), ("SMA", ["Actin-D", "Transgelin"])] # = 30_ gch = json.load(open(os.path.join(CACHE, "gigatime", "channels.json")))["channels"] gmean = lambda m: np.stack([m[:, [gch.index(g) for g in gs]].mean(1) for _, gs in PANEL], 1) mr_o = rg(np.load(os.path.join(CACHE, "gigatime", "orion_real_if_means.npy"))[:, [ORION_CH.index(o) for o, _ in PANEL]]) mv_o = rg(gmean(np.load(os.path.join(CACHE, "gigatime", "orion_means.npy")))) mv_t = rg(gmean(np.load(os.path.join(CACHE, "gigatime", "tcga_text_means.npy"))))[ev] HE8 = torch.from_numpy(np.load(os.path.join(CACHE, "he_uni2.npy")).astype(np.float32)) def train_means_head(tr, va, epochs=150, bs=256, lr=1e-3, seed=0): torch.manual_seed(seed); rng_ = np.random.default_rng(seed) h = mlp(mr_o.shape[1]) logit = nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32)) opt = torch.optim.AdamW(list(h.parameters()) + [logit], lr=lr, weight_decay=0.05) best, state = -1, None for ep in range(epochs): h.train() perm = rng_.permutation(tr) for s0 in range(0, len(perm) - bs + 1, bs): b = torch.as_tensor(perm[s0:s0 + bs]) za = F.normalize(HE8[b, torch.randint(0, 8, (len(b),))], dim=-1) lg = logit.exp().clamp(max=100) * za @ F.normalize(h(mr_o[b]), dim=-1).T t = torch.arange(len(b)) loss = (F.cross_entropy(lg, t) + F.cross_entropy(lg.T, t)) / 2 opt.zero_grad(); loss.backward(); opt.step() h.eval() with torch.no_grad(): a, m = nz(HE8[va, 0]), nz(h(mr_o[va])) sim = a @ m.T v = float(((sim > sim.diag()[:, None]).sum(1) == 0).float().mean()) if v > best: best, state = v, {k_: x.clone() for k_, x in h.state_dict().items()} h.load_state_dict(state); h.eval() return h for k in range(5): s = pd.read_csv(os.path.join(SPLITS, "patient_5fold", f"fold{k}.csv")) tr, va, te = (s.tile_id[s.split == x].values for x in ["train", "val", "test"]) h = train_means_head(tr, va) with torch.no_grad(): zh, zr = nz(he_o[te]), nz(rna_heads[k](rna_o[te])) zmr, zmv = nz(h(mr_o[te])), nz(h(mv_o[te])) add("ORION", k, "H&E<->mIF real (10 means)", zh, zmr, "trained (real)", False) add("ORION", k, "H&E<->mIF virtual (10 means)", zh, zmv, "virtual calibration", True) add("ORION", k, "mIF real<->mIF virtual (10 means)", zmr, zmv, "virtual calibration", False) add("ORION", k, "mIF real (10 means)<->RNA virtual", zmr, zr, "emergent, real x virtual", False) add("ORION", k, "mIF virtual (10 means)<->RNA virtual", zmv, zr, "emergent, virtual x virtual", True) if gt_o is not None: zt = nz(text_head(gt_o[te])) add("ORION", k, "text synthetic<->mIF real (10 means)", zt, zmr, "emergent, real x virtual", False) add("ORION", k, "text synthetic<->mIF virtual (10 means)", zt, zmv, "emergent, virtual x virtual", True) zmt = nz(h(mv_t)) add("TCGA", k, "H&E<->mIF virtual (10 means)", nz(he_t), zmt, "virtual (cross-dataset head)", True) add("TCGA", k, "text<->mIF virtual (10 means)", nz(text_head(tx)), zmt, "emergent, real x virtual", False) add("TCGA", k, "RNA virtual<->mIF virtual (10 means)", nz(rna_heads[k](rna_t)), zmt, "emergent, virtual x virtual", True) print("means head fold", k, flush=True) # ---------------- (E) calibration of the synthetic text: TCGA patches with REAL text AND a synthetic caption (27_ tcga_calib) f_cm, f_cf = os.path.join(CACHE, "tcga_calib_meta.csv"), os.path.join(CACHE, "tcga_calib_text_biomedclip.npy") if os.path.exists(f_cf): CM = pd.read_csv(f_cm) pos = pd.Series(range(len(ev)), index=ev) assert CM.img_idx.isin(ev).all() j = pos.loc[CM.img_idx].values # rows of the TCGA eval arrays with torch.no_grad(): zs = nz(text_head(torch.from_numpy(np.load(f_cf)[:, 0].astype(np.float32)))) zr_txt, zh_c = nz(text_head(tx))[j], nz(he_t)[j] add("TCGA calib", 0, "H&E<->text real", zh_c, zr_txt, "trained (real)", False) add("TCGA calib", 0, "H&E<->text synthetic", zh_c, zs, "virtual calibration", True) add("TCGA calib", 0, "text real<->text synthetic", zr_txt, zs, "virtual calibration", False) for k in range(5): zrna = nz(rna_heads[k](rna_t))[j] add("TCGA calib", k, "text real<->RNA virtual", zr_txt, zrna, "emergent, real x virtual", False) add("TCGA calib", k, "text synthetic<->RNA virtual", zs, zrna, "emergent, virtual x virtual", True) print("calibration done", len(CM), flush=True) R = pd.DataFrame(rows) R.to_csv(os.path.join(OUT, "virtual_validation.csv"), index=False) pd.set_option("display.width", 250); pd.set_option("display.max_colwidth", 60) S = R.groupby(["dataset", "pair", "kind", "circular"])[["R@1_g100", "pair_acc_g100", "n"]].agg(["mean", "std"]).round(3) print(S.drop(columns=[("n", "std")]).to_string())