| """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) |
|
|
|
|
| |
| 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() |
|
|
| |
| 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 |
|
|
| |
| 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"] |
| 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")))) |
|
|
| |
| 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"]} |
| 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} |
|
|
|
|
| |
| def mlp(din, dout=512, 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 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 |
| 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() |
|
|
|
|
| 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) |
| |
| 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) |
|
|
| |
| 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()) |
| |
| 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()}) |
|
|
