| """External test sets (never used for training): DLPFC Visium (brain, 12 sections / 3 donors) and HEST IDC |
| (breast, Xenium pseudo-spots, 4 samples; 280-gene panel). Prepared by 36_external_prep.py -> cache/external/. |
| |
| Real data only for training (HEST-1k, train/val by study|patient group, val = early stopping only): |
| dlpfc head : H&E(UNI2) <-> RNA MLP on the 2000 HEST HVGs, trained on HEST WITHOUT the spatialLIBD study |
| (= DLPFC; those spots are in cache/hest) -> DLPFC is fully held out |
| idc head : same recipe on the IDC panel genes present in the HEST common gene set (IDC samples are not in HEST cache) |
| RNA input : log1p(CP10k over the HEST common genes / over the panel genes), then z-scored per gene with the |
| statistics of its OWN dataset (HEST training spots vs the external set; label-free platform alignment) |
| 5 heads per set (different val groups / seeds) -> mean +- std. |
| Virtual data only at test time: MIPHEI-ViT virtual mIF of the external H&E (35_ + 30_ --m16) through the 16-ch mIF |
| heads trained on REAL ORION (patient 5-fold, fold k). |
| Synthetic text only at test time: GPT captions of the external H&E (27_ --sets dlpfc idc, 32_, 08_ biomedclip) |
| through the TCGA text heads (runs/text_clean/uni2_biomedclip_fold<k>, trained on CLEAN real image-text pairs in 33_). |
| Pairs: H&E<->RNA (real, external), RNA<->mIF(MIPHEI) (emergent, real x virtual), H&E<->mIF(MIPHEI) (circular), |
| RNA<->text(synthetic) (emergent, real x virtual), H&E<->text(synthetic) (circular), |
| mIF(MIPHEI)<->text(synthetic) (virtual x virtual). |
| Metrics: 100-item galleries (50 draws, both directions): R@1/R@5/R@10/MRR + Hungarian pairing accuracy. |
| Usage: python 38_external_eval.py feats (UNI2 features of the external patches, GPU) |
| python 38_external_eval.py eval -> results/external_long.csv, external_summary.csv |
| """ |
| import os |
| import sys |
|
|
| 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.optimize import linear_sum_assignment |
|
|
| from common import CACHE, ROOT, RUNS |
|
|
| E = os.path.join(CACHE, "external") |
| H = os.path.join(CACHE, "hest") |
| OUT = os.path.join(ROOT, "results") |
| SETS = {"dlpfc": "dlpfc_he_05um.npy"} |
| MODE = sys.argv[1] |
|
|
| if MODE == "feats": |
| import timm |
| dev = "cuda" |
| model = timm.create_model( |
| "vit_giant_patch14_224", pretrained=False, img_size=224, patch_size=14, depth=24, num_heads=24, |
| init_values=1e-5, embed_dim=1536, mlp_ratio=2.66667 * 2, num_classes=0, no_embed_class=True, |
| mlp_layer=timm.layers.SwiGLUPacked, act_layer=torch.nn.SiLU, reg_tokens=8, dynamic_img_size=True) |
| model.load_state_dict(torch.load("/diskarray/home/tl688/reg2026/model_alldata/uni2-h.bin", map_location="cpu"), strict=True) |
| model = model.eval().to(dev) |
| mean = torch.tensor([0.485, 0.456, 0.406], device=dev).view(1, 3, 1, 1) |
| std = torch.tensor([0.229, 0.224, 0.225], device=dev).view(1, 3, 1, 1) |
| for st, f in SETS.items(): |
| X = np.load(os.path.join(E, f), mmap_mode="r") |
| out = np.zeros((len(X), 1536), np.float32) |
| with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16): |
| for s in range(0, len(X), 32): |
| x = torch.from_numpy(np.asarray(X[s:s + 32])).to(dev).permute(0, 3, 1, 2).float() / 255 |
| if x.shape[-1] != 224: |
| x = F.interpolate(x, size=(224, 224), mode="bilinear", antialias=True, align_corners=False) |
| out[s:s + 32] = model((x - mean) / std).float().cpu().numpy() |
| np.save(os.path.join(E, f"{st}_he_uni2.npy"), out) |
| print(st, out.shape, flush=True) |
| |
| m, hm = pd.read_csv(os.path.join(E, "dlpfc_meta.csv")), pd.read_csv(os.path.join(H, "meta.csv")) |
| hf = np.load(os.path.join(H, "he_uni2.npy"))[:, 0].astype(np.float32) |
| o = np.load(os.path.join(E, "dlpfc_he_uni2.npy")) |
| j = m.reset_index().merge(hm.reset_index(), left_on=["hest_id", "barcode"], right_on=["sample", "spot"], suffixes=("_e", "_h")) |
| if len(j): |
| a, b = F.normalize(torch.from_numpy(o[j.index_e.values]), dim=-1), F.normalize(torch.from_numpy(hf[j.index_h.values]), dim=-1) |
| rnd = F.normalize(torch.from_numpy(hf[np.random.default_rng(0).choice(len(hf), 200, replace=False)]), dim=-1) |
| print(f"protocol check: {len(j)} shared spots, cosine(ours, HEST cache) = {(a * b).sum(1).numpy().round(3)}; " |
| f"mean cosine to random HEST spots {float((a @ rnd.T).mean()):.3f}") |
| sys.exit(0) |
|
|
| |
| dev = "cuda" if torch.cuda.is_available() else "cpu" |
| G100, DRAWS, K = 100, 50, 5 |
| HM = pd.read_csv(os.path.join(H, "meta.csv")) |
| hgenes = pd.read_csv(os.path.join(H, "genes.csv")).gene.values |
| hcounts = sp.load_npz(os.path.join(H, "counts.npz")).tocsr().astype(np.float32) |
| hvg = pd.read_csv(os.path.join(H, "hvg_genes.csv")).gene.values |
| HE_H = torch.from_numpy(np.load(os.path.join(H, "he_uni2.npy"))[:, 0].astype(np.float32)) |
|
|
|
|
| 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 lognorm(C, total): |
| return np.log1p(C / np.maximum(total, 1)[:, None] * 1e4).astype(np.float32) |
|
|
|
|
| def zscore(x, ref): |
| mu, sd = ref.mean(0), ref.std(0) + 1e-6 |
| return (x - mu) / sd |
|
|
|
|
| def r1(a, b): |
| s = a @ b.T |
| return float(((s > s.diag()[:, None]).sum(1) == 0).float().mean()) |
|
|
|
|
| def train_head(G, tr, va, seed, epochs=100, bs=256, lr=3e-4): |
| torch.manual_seed(seed); rng = np.random.default_rng(seed) |
| A, Gd = HE_H.to(dev), torch.from_numpy(G).to(dev) |
| head = mlp(G.shape[1]).to(dev) |
| logit = nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32, device=dev)) |
| opt = torch.optim.AdamW(list(head.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) // bs), pct_start=0.05) |
| za_va = F.normalize(A[va], dim=-1) |
| best, state = -1, None |
| for ep in range(epochs): |
| head.train() |
| perm = rng.permutation(tr) |
| for s in range(0, len(perm) - bs + 1, bs): |
| b = torch.as_tensor(perm[s:s + bs], device=dev) |
| lg = logit.exp().clamp(max=100) * F.normalize(A[b], dim=-1) @ F.normalize(head(Gd[b]), dim=-1).T |
| t = torch.arange(len(b), device=dev) |
| loss = (F.cross_entropy(lg, t) + F.cross_entropy(lg.T, t)) / 2 |
| opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sched.step() |
| head.eval() |
| with torch.no_grad(): |
| zg = F.normalize(head(Gd[va]), dim=-1) |
| v = (r1(za_va, zg) + r1(zg, za_va)) / 2 |
| if v > best: |
| best, state = v, {k_: x.detach().cpu().clone() for k_, x in head.state_dict().items()} |
| head = head.cpu(); head.load_state_dict(state); head.eval() |
| return head, best |
|
|
|
|
| def g100(a, b, seed=0): |
| 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()} |
|
|
|
|
| nz = lambda x: F.normalize(torch.as_tensor(np.asarray(x), dtype=torch.float32), dim=-1) |
|
|
|
|
| def mif_head(k): |
| d = torch.load(os.path.join(RUNS, "bench_fm", f"patient_5fold_fold{k}_uni2_fm-virtues_m16_fm-geneformer_pathgenclip", "heads.pt"), map_location="cpu") |
| h = mlp(d["state"]["if"]["0.weight"].shape[1]); h.load_state_dict(d["state"]["if"]); h.eval() |
| return lambda x: h((torch.as_tensor(x, dtype=torch.float32) - d["if_mean"]) / d["if_std"]) |
|
|
|
|
| def text_head(k): |
| sd = torch.load(os.path.join(RUNS, "text_clean", f"uni2_biomedclip_fold{k}", "head.pt"), map_location="cpu") |
| h = mlp(sd["0.weight"].shape[1], sd["6.weight"].shape[0]); h.load_state_dict(sd); h.eval() |
| return lambda x: h(torch.as_tensor(x, dtype=torch.float32)) |
|
|
|
|
| rows = [] |
| hset = pd.Series(range(len(hgenes)), index=hgenes) |
| htot = np.asarray(hcounts.sum(1)).ravel() |
| for st in SETS: |
| em = pd.read_csv(os.path.join(E, f"{st}_meta.csv")) |
| he = np.load(os.path.join(E, f"{st}_he_uni2.npy")) |
| if st == "dlpfc": |
| eg = pd.read_csv(os.path.join(E, "dlpfc_genes_full.csv")).gene.values |
| ec = np.load(os.path.join(E, "dlpfc_counts_full.npy")).astype(np.float32) |
| eidx = pd.Series(range(len(eg)), index=eg) |
| eidx = eidx[~eidx.index.duplicated()] |
| common = [g for g in hgenes if g in eidx.index] |
| genes = [g for g in hvg if g in eidx.index] |
| etot = ec[:, eidx.loc[common].values].sum(1) |
| excl = HM.group.str.startswith("spatialLIBD").values |
| htot_ = htot |
| else: |
| eg = pd.read_csv(os.path.join(E, "idc_genes.csv")).gene.values |
| ec = np.load(os.path.join(E, "idc_counts.npy")).astype(np.float32) |
| eidx = pd.Series(range(len(eg)), index=eg) |
| eidx = eidx[~eidx.index.duplicated()] |
| genes = [g for g in eg if g in hset.index] |
| etot = ec[:, eidx.loc[genes].values].sum(1) |
| htot_ = np.asarray(hcounts[:, hset.loc[genes].values].sum(1)).ravel() |
| excl = HM["sample"].isin(em["sample"].unique()).values |
| Gh = lognorm(hcounts[:, hset.loc[genes].values].toarray(), htot_) |
| Ge = lognorm(ec[:, eidx.loc[genes].values], etot) |
| pool = np.flatnonzero(~excl & (htot_ > 0)) |
| print(f"{st}: {len(em)} external spots, {len(genes)} genes, HEST training pool {len(pool)} spots " |
| f"({int(excl.sum())} excluded)", flush=True) |
| f_mv = os.path.join(CACHE, f"if_virtues_m16_virtual_miphei_{st}.npy") |
| mv = np.load(f_mv)[:, 0].astype(np.float32) if os.path.exists(f_mv) else None |
| f_t = os.path.join(CACHE, f"{st}_gentext_text_biomedclip.npy") |
| tx = np.load(f_t)[:, 0].astype(np.float32) if os.path.exists(f_t) else None |
| zh = nz(he) |
| for k in range(K): |
| grp = np.random.default_rng(k).permutation(sorted(HM.group.iloc[pool].unique())) |
| va_g = set(grp[: max(1, len(grp) // 10)]) |
| isva = HM.group.iloc[pool].isin(va_g).values |
| tr, va = pool[~isva], pool[isva] |
| for norm in ["z-scored per dataset", "plain"]: |
| Gtr = zscore(Gh, Gh[tr]) if norm != "plain" else Gh |
| Gex = zscore(Ge, Ge) if norm != "plain" else Ge |
| head, best = train_head(Gtr, tr, va, seed=k) |
| with torch.no_grad(): |
| zg = nz(head(torch.from_numpy(Gex))) |
| rows.append({"set": st, "fold": k, "rna_input": norm, "part": "real", "pair": "H&E<->RNA", |
| "kind": "trained pair (external)", "circular": False, "n": len(zh), "val_R@1": best, **g100(zh, zg)}) |
| if mv is not None and norm == "z-scored per dataset": |
| with torch.no_grad(): |
| zm = nz(mif_head(k)(mv)) |
| for pair, a, b, kind, circ in [("RNA<->mIF(MIPHEI)", zg, zm, "emergent: real x virtual", False), |
| ("H&E<->mIF(MIPHEI)", zh, zm, "virtual from this H&E", True)]: |
| rows.append({"set": st, "fold": k, "rna_input": norm, "part": "virtual", "pair": pair, "kind": kind, |
| "circular": circ, "n": len(zh), **g100(a, b)}) |
| if tx is not None and norm == "z-scored per dataset": |
| with torch.no_grad(): |
| zt = nz(text_head(k)(tx)) |
| zm = nz(mif_head(k)(mv)) if mv is not None else None |
| for pair, a, b, kind, circ in [("RNA<->text(synthetic)", zg, zt, "emergent: real x virtual", False), |
| ("H&E<->text(synthetic)", zh, zt, "virtual from this H&E", True), |
| ("mIF(MIPHEI)<->text(synthetic)", zm, zt, "emergent: virtual x virtual", True)]: |
| if a is not None: |
| rows.append({"set": st, "fold": k, "rna_input": norm, "part": "virtual", "pair": pair, "kind": kind, |
| "circular": circ, "n": len(zh), **g100(a, b)}) |
| if norm == "z-scored per dataset": |
| np.savez(os.path.join(E, f"emb_{st}_fold{k}.npz"), he=zh.numpy(), rna=zg.numpy(), |
| **({"mif": zm.numpy()} if mv is not None else {}), **({"text": zt.numpy()} if tx is not None else {})) |
| print(f" {st} head {k}: val R@1 {best:.3f}", flush=True) |
|
|
| R = pd.DataFrame(rows) |
| R.to_csv(os.path.join(OUT, "external_long.csv"), index=False) |
| cols = ["R@1", "R@5", "R@10", "MRR", "pair_acc"] |
| S = R.groupby(["set", "part", "pair", "rna_input", "kind", "circular"], sort=False)[cols + ["n"]].agg(["mean", "std"]) |
| S.round(4).to_csv(os.path.join(OUT, "external_summary.csv")) |
| pd.set_option("display.width", 250) |
| print(R.groupby(["set", "part", "pair", "rna_input", "circular"], sort=False)[cols + ["n"]].mean().round(3).to_string()) |
|
|
