| """RNA encoder comparison + fusion under the final protocol (real pairs only; virtual data only at test time). |
| |
| RNA encoders (frozen inputs, cache/rnafm/<enc>_<set>.npy from 43_*.py; "hvg" = log1p CP10k of the 2000 HEST HVGs, |
| identical to 42_train_regimes.py): |
| hvg, scgpt, geneformer, nicheformer, tahoex1, aidocell, se600m, c2s, c2s-last, genept-bge, genept-bmc (whatever exists) |
| cat:<a>+<b>+... early fusion: concatenation of the z-scored inputs, each block scaled by 1/sqrt(its dim) |
| ens:<a>+<b>+... late fusion: mean of the L2-normalised outputs of the separately trained single-encoder heads |
| (possible because every head maps into the same frozen UNI2 space) |
| Head: linear or MLP (42_'s mlp) from the encoder input to the frozen UNI2 space, InfoNCE, AdamW 3e-4 / wd 0.05 / |
| OneCycle, batch 256, 2700 steps, early stopping on val R@1 every 50 steps (= 42_ S1/S4 RNA recipe). |
| FM inputs are z-scored with the HEST TRAINING-spot mean / std of fold k (applied unchanged to virtual / external sets). |
| Per fold k (splits/final/hest_fold<k>.csv, ORION splits/patient_5fold, TCGA splits/final/tcga_fold<k>.csv): |
| trained H&E<->RNA HEST test spots (real) |
| emergent mIF<->RNA(virtual) ORION test tiles, mIF = S4 mIF MLP head of fold k (trained here once, cached) |
| emergent text<->RNA(virtual) TCGA test images, text = S4 linear text head of fold k (clean real text) |
| calib H&E<->RNA(virtual) HEST test spots with STPath RNA instead of the real one (circular, reference only) |
| --dlpfc: additionally train on HEST WITHOUT the spatialLIBD study (val = 10% of groups) and test H&E<->RNA on the |
| external DLPFC set (482 spots; FM inputs z-scored per dataset as in 38_), mixed galleries G=100. |
| Usage: python 44_rnafm_eval.py [--encs ...] [--heads linear mlp] [--folds 0 1 2 3 4] [--dlpfc] |
| Output results/rnafm_long.csv (resumable), results/rnafm_summary.csv |
| """ |
| import argparse |
| 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, SPLITS |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--encs", nargs="+", default=None) |
| ap.add_argument("--heads", nargs="+", default=["linear", "mlp"]) |
| ap.add_argument("--folds", nargs="+", type=int, default=[0, 1, 2, 3, 4]) |
| ap.add_argument("--dlpfc", action="store_true") |
| args = ap.parse_args() |
| dev = "cuda" |
| H, RF = os.path.join(CACHE, "hest"), os.path.join(CACHE, "rnafm") |
| G100, DRAWS, BS, LR, EVAL_EVERY, STEPS = 100, 50, 256, 3e-4, 50, 2700 |
| f32 = lambda a: torch.from_numpy(np.asarray(a, dtype=np.float32)) |
| SETS = ["hest", "hest_v", "orion_v", "tcga_v", "dlpfc"] |
|
|
| |
| 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) |
| 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) |
| HE_D = HE_D[:, 0] if HE_D.dim() == 3 else HE_D |
| TXT = f32(np.load(os.path.join(CACHE, "alltxt_text_biomedclip.npy"))[:, 0]).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")), "hest_v": virt("hest_val"), "orion_v": virt("orion"), |
| "tcga_v": virt("tcga_text"), "dlpfc": d} |
|
|
|
|
| def load_enc(e): |
| e = e.split("@")[0] |
| if e == "hvg": |
| return {s: f32(x).to(dev) for s, x in hvg_inputs().items()} |
| return {s: f32(np.load(os.path.join(RF, f"{e}_{s}.npy"))).to(dev) for s in SETS} |
|
|
|
|
| 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): |
| """modality head X -> frozen anchor space A (rows paired). A_* may be (n, V, 1536) -> random view per sample.""" |
| 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 = F.normalize(A_va[:, 0] if A_va.dim() == 3 else A_va, dim=-1) |
| 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) * F.normalize(a, dim=-1) @ F.normalize(h(X_tr[b]), dim=-1).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, F.normalize(h(X_va), dim=-1)) |
| 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, best |
|
|
|
|
| class Fold: |
| def __init__(self, k): |
| s = pd.read_csv(os.path.join(SPLITS, "patient_5fold", f"fold{k}.csv")) |
| self.o = {x: s.tile_id[s.split == x].values for x in ["train", "val", "test"]} |
| mu, sd = IF_O[self.o["train"]].mean((0, 1)), IF_O[self.o["train"]].std((0, 1)) + 1e-3 |
| self.IF = (IF_O - mu) / sd |
| h = pd.read_csv(os.path.join(SPLITS, "final", f"hest_fold{k}.csv")) |
| self.h = {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 |
| self.t = {x: T.index[tsp == x].values for x in ["train", "val", "test"]} |
| self.t_va = one_text_per_image(self.t["val"], np.random.default_rng(0)) |
| self.t_te = one_text_per_image(self.t["test"], np.random.default_rng(k)) |
| |
| 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 |
| self.d = {"train": pool[~isva], "val": pool[isva]} |
|
|
| def other_heads(self, k): |
| """S4 heads of fold k for mIF (MLP) and text (linear); cached in runs/rnafm/other_fold<k>.pt""" |
| f = os.path.join(ROOT, "runs", "rnafm", f"other_fold{k}.pt") |
| hm, ht = mlp(512).to(dev), nn.Linear(512, 1536).to(dev) |
| if os.path.exists(f): |
| sd = torch.load(f, map_location=dev); hm.load_state_dict(sd["mif"]); ht.load_state_dict(sd["text"]) |
| return hm.eval(), ht.eval() |
| o, t = self.o, self.t |
| hm, _ = train_head(self.IF[o["train"]].flatten(0, 1), HE_O[o["train"]].flatten(0, 1), |
| self.IF[o["val"], 0], HE_O[o["val"], 0], "mlp", 400, k) |
| ii, tt = self.t_va |
| ht, _ = train_head(TXT[t["train"]], HE_T[T_IMG[torch.as_tensor(t["train"], device=dev)]], |
| TXT[torch.as_tensor(tt, device=dev)], HE_T[torch.as_tensor(ii, device=dev), 0], |
| "linear", 1000, k, pos_tr=T_IMG[torch.as_tensor(t["train"], device=dev)]) |
| os.makedirs(os.path.dirname(f), exist_ok=True) |
| torch.save({"mif": hm.state_dict(), "text": ht.state_dict()}, f) |
| return hm, ht |
|
|
|
|
| def prep(X, tr, per_dataset_ext=False): |
| """z-score with HEST-train stats (hvg: unchanged, = 42_). Returns dict set -> tensor.""" |
| if X is None: |
| return None |
| mu, sd = X["hest"][tr].mean(0), X["hest"][tr].std(0) + 1e-6 |
| out = {s: (x - mu) / sd for s, x in X.items()} |
| if per_dataset_ext: |
| out["dlpfc"] = (X["dlpfc"] - X["dlpfc"].mean(0)) / (X["dlpfc"].std(0) + 1e-6) |
| return out |
|
|
|
|
| def build(enc, fold_tr, dl_tr=None): |
| """input tensors per set for encoder spec enc; hvg kept raw (42_), FMs / fusions z-scored.""" |
| parts = enc.split(":", 1)[1].split("+") if enc.startswith("cat:") else [enc] |
| main, ext = [], [] |
| for p in parts: |
| X = CACHE_ENC[p] |
| if p.split("@")[0] == "hvg" and len(parts) == 1: |
| main.append(X) |
| if dl_tr is not None: |
| ext.append(prep(X, dl_tr, True)) |
| continue |
| w = 1 / np.sqrt(X["hest"].shape[1]) if len(parts) > 1 else 1.0 |
| main.append({s: v * w for s, v in prep(X, fold_tr).items()}) |
| if dl_tr is not None: |
| ext.append({s: v * w for s, v in prep(X, dl_tr, True).items()}) |
| cat = lambda L: {s: torch.cat([d[s] for d in L], 1) for s in SETS} |
| return cat(main), (cat(ext) if dl_tr is not None else None) |
|
|
|
|
| |
| avail = ["hvg"] + sorted({f.rsplit("_", 1)[0] for f in os.listdir(RF) if f.endswith("_hest.npy") |
| and all(os.path.exists(os.path.join(RF, f.replace("_hest.npy", f"_{s}.npy"))) for s in SETS)}) |
| encs = sorted(args.encs or avail, key=lambda e: e.startswith("ens:")) |
| singles = sorted({p for e in encs for p in (e.split(":", 1)[1].split("+") if ":" in e else [e])}) |
| CACHE_ENC = {e: load_enc(e) for e in singles} |
| print("encoders:", {e: CACHE_ENC[e]["hest"].shape[1] for e in singles}, flush=True) |
| f_out = os.path.join(ROOT, "results", "rnafm_long.csv") |
| rows = pd.read_csv(f_out).to_dict("records") if os.path.exists(f_out) else [] |
| done = {(r["enc"], r["head"], r["fold"]) for r in rows} |
| nz = lambda x: F.normalize(x, dim=-1) |
| for k in args.folds: |
| F_ = Fold(k) |
| hm, ht = F_.other_heads(k) |
| te_o = F_.o["test"]; ii, tt = (torch.as_tensor(x, device=dev) for x in F_.t_te) |
| with torch.no_grad(): |
| z_mif = nz(hm(F_.IF[te_o, 0])); z_txt = nz(ht(TXT[tt])) |
| heads_out = {} |
| for enc in encs: |
| for hk in args.heads: |
| parts = enc.split(":", 1)[1].split("+") if enc.startswith("ens:") else None |
| need = any(enc in e.split(":", 1)[1].split("+") for e in encs if e.startswith("ens:")) |
| if (enc, hk, k) in done and parts is None and not need: |
| continue |
| tr, va, te = F_.h["train"], F_.h["val"], F_.h["test"] |
| if parts is not None: |
| if not all((p, hk) in heads_out for p in parts) or (enc, hk, k) in done: |
| continue |
| Z = {s: nz(sum(heads_out[(p, hk)][s] for p in parts)) for s in heads_out[(parts[0], hk)]} |
| best = float("nan") |
| else: |
| X, Xd = build(enc, tr, F_.d["train"] if args.dlpfc else None) |
| sd_ = k + 1000 * int(enc.split("@")[1]) if "@" in enc else k |
| h, best = train_head(X["hest"][tr], HE_H[tr], X["hest"][va], HE_H[va], hk, STEPS, sd_) |
| with torch.no_grad(): |
| Z = {"hest": nz(h(X["hest"][te])), "hest_v": nz(h(X["hest_v"][te])), |
| "orion_v": nz(h(X["orion_v"][te_o])), "tcga_v": nz(h(X["tcga_v"][ii]))} |
| if args.dlpfc: |
| hd, _ = train_head(Xd["hest"][F_.d["train"]], HE_H[F_.d["train"]], Xd["hest"][F_.d["val"]], |
| HE_H[F_.d["val"]], hk, STEPS, sd_) |
| with torch.no_grad(): |
| Z["dlpfc"] = nz(hd(Xd["dlpfc"])) |
| heads_out[(enc, hk)] = Z |
| res = {("H&E<->RNA (HEST)", "trained"): g100(nz(HE_H[te]), Z["hest"]), |
| ("mIF<->RNA(virtual) (ORION)", "emergent"): g100(z_mif, Z["orion_v"]), |
| ("text<->RNA(virtual) (TCGA)", "emergent"): g100(z_txt, Z["tcga_v"]), |
| ("H&E<->RNA(virtual) (HEST)", "calib"): g100(nz(HE_H[te]), Z["hest_v"])} |
| if "dlpfc" in Z: |
| res[("H&E<->RNA (DLPFC external)", "external")] = g100(nz(HE_D), Z["dlpfc"]) |
| if (enc, hk, k) in done: |
| continue |
| for (pair, kind), r in res.items(): |
| rows.append({"enc": enc, "head": hk, "fold": k, "pair": pair, "kind": kind, "val_R@1": best, **r}) |
| pd.DataFrame(rows).to_csv(f_out, index=False) |
| print(f"fold {k} {enc:<34} {hk:<6} val {best:.3f} " + |
| " ".join(f"{p.split(' (')[0]} {r['R@1']:.3f}" for (p, _), r in res.items()), flush=True) |
|
|
| R = pd.DataFrame(rows) |
| S = R.groupby(["enc", "head", "kind", "pair"], sort=False)[["R@1", "R@10", "pair_acc"]].agg(["mean", "std"]).round(4) |
| S.to_csv(os.path.join(ROOT, "results", "rnafm_summary.csv")) |
| pd.set_option("display.width", 250) |
| print(R.pivot_table(index=["enc", "head"], columns="pair", values="R@1", sort=False).round(3).to_string()) |
|
|
