"""Virtual modality completion: does training text <-> RNA / text <-> protein DIRECTLY on virtual pairs beat the emergent (H&E-anchor) route? Training pairs (TCGA all-pairs pool, 6,144 patches, every real text of a patch is a positive): rna : real text <-> STPath virtual RNA of the same patch (cache/stpath/tcga_text_pred.npy -> log1p CP10k, 2000 HEST HVGs) protein : real text <-> GigaTIME virtual mIF of the same patch (cache/gigatime/tcga_text_means.npy, 11 channels shared with ORION) Model: frozen BiomedCLIP text tower + MLP head and an MLP head on the modality vector -> shared 512-d space (CLIP-style, multi-positive InfoNCE), patient GroupKFold(5) over TCGA, 10% of training patients for early stopping. Each feature is rank-gaussianised within its own dataset (virtual TCGA / real HEST / real ORION), because virtual values (STPath expected counts, GigaTIME probabilities) and real measurements live on different scales. Evaluation only on REAL measurements (never virtual): rna : HEST real RNA, concept prompts vs RNA of the concept gene (protocol of 24_eval_gene_emergent.py (B): AUROC top quartile, per HEST study+patient fold, all spots / bowel spots) protein : ORION real mIF tile means, concept prompts vs marker+ cell fraction (probe of 18_eval_fm.py, per patient_5fold test fold); also ORION GigaTIME virtual means (same tiles) to separate the virtual->real gap Baselines (emergent, read from results/): text head trained on PathGen H&E-caption pairs, scored against the RNA (23_) / VirTues mIF (17_) embeddings bound to UNI2. Outputs results/virtual_completion.csv """ import glob import json import os import numpy as np import pandas as pd import scipy.sparse as sp import torch import torch.nn as nn import torch.nn.functional as F from scipy.special import ndtri from scipy.stats import rankdata, spearmanr from sklearn.metrics import roc_auc_score from sklearn.model_selection import GroupKFold from common import CACHE, ROOT, RUNS, SPLITS, load_all_tiles dev = "cuda" OUT = os.path.join(ROOT, "results") TXT_ENC = "biomedclip" prompts = json.load(open(os.path.join(ROOT, "prompts.json"))) PKEYS = list(prompts) def rankgauss(x): r = np.apply_along_axis(rankdata, 0, x) return ndtri((r - 0.5) / len(x)).astype(np.float32) # ---------------- TCGA training pool (texts, images, virtual modalities) I = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv")) T = pd.read_csv(os.path.join(CACHE, "alltxt_meta.csv")) TX = torch.from_numpy(np.load(os.path.join(CACHE, f"alltxt_text_{TXT_ENC}.npy"))[:, 0].astype(np.float32)).to(dev) PR = torch.from_numpy(np.load(os.path.join(CACHE, f"prompts_text_{TXT_ENC}.npy"))[:, 0].astype(np.float32)).to(dev) t_img = torch.as_tensor(T.img_idx.values, device=dev) assert (pd.read_csv(os.path.join(CACHE, "stpath", "tcga_text_meta.csv")).key.values == I.key.values).all() assert (pd.read_csv(os.path.join(CACHE, "gigatime", "tcga_text_meta.csv")).key.values == I.key.values).all() # RNA: STPath log1p(counts) -> CP10k -> log1p, HEST HVGs H = os.path.join(CACHE, "hest") hvg = pd.read_csv(os.path.join(H, "hvg_genes.csv")).gene.values sg = pd.read_csv(os.path.join(CACHE, "stpath", "tcga_text_genes.csv")).gene.values c = np.expm1(np.load(os.path.join(CACHE, "stpath", "tcga_text_pred.npy")).astype(np.float32)).clip(0) c = np.log1p(c / c.sum(1, keepdims=True) * 1e4) gix = pd.Series(range(len(sg)), index=sg).loc[hvg].values V_RNA = rankgauss(c[:, gix]) del c # protein: GigaTIME channels shared with ORION (Actin-D + Transgelin -> SMA) ch = json.load(open(os.path.join(CACHE, "gigatime", "channels.json"))) G_CH, SHARED = ch["channels"], ch["orion_shared"] 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_) P_ORION = sorted(set(SHARED.values()), key=ORION_CH.index) def giga_vec(means): return np.stack([np.mean([means[:, G_CH.index(g)] for g, o in SHARED.items() if o == oc], 0) for oc in P_ORION], 1) V_PROT = rankgauss(giga_vec(np.load(os.path.join(CACHE, "gigatime", "tcga_text_means.npy")))) # ---------------- real evaluation sets M = pd.read_csv(os.path.join(H, "meta.csv")) hest_x = rankgauss(np.load(os.path.join(H, "hvg.npy")).astype(np.float32)) genes = pd.read_csv(os.path.join(H, "genes.csv")).gene.values X = sp.load_npz(os.path.join(H, "counts.npz")).tocsr() L = sp.diags(1e4 / np.asarray(X.sum(1)).ravel()) @ X L.data = np.log1p(L.data) gpos = {g: i for i, g in enumerate(genes)} rna = lambda gs: np.asarray(L[:, [gpos[g] for g in gs]].mean(1)).ravel() CONCEPT_RNA = {"immune": ["PTPRC"], "tumor_epith": ["KRT8", "EPCAM"], "stroma": ["ACTA2"], "proliferation": ["MKI67"], "vessels": ["PECAM1"], "macrophage": ["CD68"]} # as 24_ hest_fold = np.full(len(M), -1) for f in glob.glob(os.path.join(RUNS, "gene", "uni2_mlp_fold*", "test.pt")): hest_fold[torch.load(f, weights_only=False)["idx"]] = int(f.split("_fold")[1].split(os.sep)[0]) assert (hest_fold >= 0).all() df = load_all_tiles() nn_ = pd.read_csv(os.path.join(OUT, "tile_stats.csv")).set_index("tile_id").loc[df.tile_id, "n_nuclei"].values df["n_nuclei"] = nn_ orion_real = rankgauss(np.load(os.path.join(CACHE, "gigatime", "orion_real_if_means.npy"))[:, [ORION_CH.index(o) for o in P_ORION]]) orion_virt = rankgauss(giga_vec(np.load(os.path.join(CACHE, "gigatime", "orion_means.npy")))) assert (pd.read_csv(os.path.join(CACHE, "gigatime", "orion_meta.csv")).tile_id.values == df.tile_id.values).all() orion_fold = np.full(len(df), -1) for k in range(5): s = pd.read_csv(os.path.join(SPLITS, "patient_5fold", f"fold{k}.csv")) orion_fold[s.tile_id[s.split == "test"].values] = k def orion_probe(scores, sel, key, label_col): """18_eval_fm.probe restricted to tile mask `sel`""" sub = sel if key == "cellularity" else sel & (nn_ >= 20) y = df[label_col].values[sub].astype(float) if label_col != "n_nuclei": y = y / np.maximum(nn_[sub], 1) lab = y >= np.quantile(y, 0.75) return {"spearman": spearmanr(scores[sub], y).correlation, "auroc": roc_auc_score(lab, scores[sub]) if 0 < lab.sum() < lab.size else np.nan} # ---------------- model def mlp(din, dout=512, hid=512): # same head as 20_pairing.py 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 r1(a, b): s = a @ b.T return float(((s > s.diag()[:, None]).sum(1) == 0).float().mean()) def train(V, tr_t, va_t, epochs=60, bs=256, lr=3e-4, seed=0): torch.manual_seed(seed); rng = np.random.default_rng(seed) th, mh = mlp(TX.shape[1]).to(dev), mlp(V.shape[1]).to(dev) logit = nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32, device=dev)) opt = torch.optim.AdamW(list(th.parameters()) + list(mh.parameters()) + [logit], lr=lr, weight_decay=0.05) sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=lr, total_steps=epochs * max(1, len(tr_t) // bs), pct_start=0.05) va = T.loc[va_t].groupby("img_idx").head(1).index.values # one text per val image va_i, va_t1 = torch.as_tensor(T.img_idx.values[va], device=dev), torch.as_tensor(va, device=dev) best, state = -1, None for ep in range(epochs): th.train(); mh.train() perm = rng.permutation(tr_t) for s in range(0, len(perm) - bs + 1, bs): b = torch.as_tensor(perm[s:s + bs], device=dev) ii = t_img[b] lg = logit.exp().clamp(max=100) * F.normalize(mh(V[ii]), dim=-1) @ F.normalize(th(TX[b]), dim=-1).T 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(); sched.step() th.eval(); mh.eval() with torch.no_grad(): zv, zt = F.normalize(mh(V[va_i]), dim=-1), F.normalize(th(TX[va_t1]), dim=-1) v = (r1(zv, zt) + r1(zt, zv)) / 2 if v > best: best, state = v, ({k: x.clone() for k, x in th.state_dict().items()}, {k: x.clone() for k, x in mh.state_dict().items()}) th.load_state_dict(state[0]); mh.load_state_dict(state[1]); th.eval(); mh.eval() return th, mh, best @torch.no_grad() def score(th, mh, x): zp = F.normalize(th(PR), dim=-1) z = F.normalize(mh(torch.from_numpy(x).to(dev)), dim=-1) return (z @ zp.T).cpu().numpy() # (n, n_prompts) rows = [] folds = list(GroupKFold(5).split(I, groups=I.patient)) for mod, Vnp in [("rna", V_RNA), ("protein", V_PROT)]: V = torch.from_numpy(Vnp).to(dev) for k, (trv, te) in enumerate(folds): pats = np.random.default_rng(k).permutation(sorted(I.patient.iloc[trv].unique())) va_p = set(pats[: max(1, len(pats) // 10)]) isplit = pd.Series("test", index=I.index) isplit.iloc[trv] = np.where(I.patient.iloc[trv].isin(va_p), "val", "train") ts = isplit.loc[T.img_idx].values tr_t, va_t, te_t = (T.index[ts == s].values for s in ["train", "val", "test"]) th, mh, vbest = train(V, tr_t, va_t) # in-domain sanity: held-out TCGA patients, text <-> virtual modality, 1 text per image, R@1 on 100-galleries te1 = T.loc[te_t].groupby("img_idx").head(1).index.values with torch.no_grad(): zv = F.normalize(mh(V[T.img_idx.values[te1]]), dim=-1).cpu() zt = F.normalize(th(TX[te1]), dim=-1).cpu() g = np.random.default_rng(0) r = np.mean([(r1(zv[s], zt[s]) + r1(zt[s], zv[s])) / 2 for s in (g.choice(len(te1), 100, replace=False) for _ in range(50))]) rows.append({"modality": mod, "model_fold": k, "eval": "tcga_virtual_R@1_g100", "concept": "-", "auroc": np.nan, "spearman": r, "val_r1": vbest}) if mod == "rna": S = score(th, mh, hest_x) for c, gs in CONCEPT_RNA.items(): y_all = rna(gs) for gal in ["all", "bowel"]: gm = np.ones(len(M), bool) if gal == "all" else (M.organ == "Bowel").values for hk in range(5): sel = gm & (hest_fold == hk) y = y_all[sel]; lab = y >= np.quantile(y, 0.75) s = S[sel, PKEYS.index(c)] rows.append({"modality": mod, "model_fold": k, "eval": f"hest_real_{gal}", "concept": c, "eval_fold": hk, "auroc": roc_auc_score(lab, s) if 0 < lab.sum() < lab.size else np.nan, "spearman": spearmanr(s, y).correlation}) else: for src, x in [("orion_real", orion_real), ("orion_virtual", orion_virt)]: S = score(th, mh, x) for i, (c, pr) in enumerate(prompts.items()): for ok in range(5): rows.append({"modality": mod, "model_fold": k, "eval": src, "concept": c, "eval_fold": ok, **orion_probe(S[:, i], orion_fold == ok, c, pr["label"])}) print(mod, "fold", k, "val", round(vbest, 3), "tcga R@1_g100", round(r, 3), flush=True) R = pd.DataFrame(rows) R.to_csv(os.path.join(OUT, "virtual_completion.csv"), index=False) # ---------------- comparison with the emergent route (same prompts, same text encoder, same real labels) pd.set_option("display.width", 250) print(R[R["eval"] == "tcga_virtual_R@1_g100"].groupby("modality").spearman.agg(["mean", "std"]).round(3)) D = R[R["eval"] != "tcga_virtual_R@1_g100"] print(D.groupby(["modality", "eval", "concept"]).auroc.mean().unstack("concept").round(3).to_string()) print(D.groupby(["modality", "eval"]).auroc.mean().round(3).to_string()) E = pd.read_csv(os.path.join(OUT, "gene_emergent_text.csv")) E = E[E.text_enc == TXT_ENC] print("emergent text<->RNA (HEST real):\n", E.groupby(["gallery", "modality", "concept"]).auroc.mean().unstack("concept").round(3).to_string()) A = pd.read_csv(os.path.join(OUT, "anchor_text_probes.csv")) A = A[(A.text_enc == TXT_ENC) & (A.split == "patient_5fold") & (A.anchor == "uni2")] print("emergent text<->protein (ORION real, VirTues):\n", A.groupby(["modality", "concept"]).auroc.mean().unstack("concept").round(3).to_string()) # single-marker reference on ORION real: the matched channel itself as the score REF = {"immune": "CD3e", "tumor_epith": "Pan-CK", "stroma": "SMA", "proliferation": "Ki67", "macrophage": "CD68", "cellularity": "Hoechst"} ref = {c: np.mean([orion_probe(orion_real[:, P_ORION.index(ch_)], orion_fold == ok, c, prompts[c]["label"])["auroc"] for ok in range(5)]) for c, ch_ in REF.items()} print("single real channel as score (reference):", {c: round(v, 3) for c, v in ref.items()})