"""Emergent alignment of the gene-expression modality (HEST, 23_) with protein (ORION mIF, 17_/09_) and text (PathGen, 12_). All heads were trained separately against the same frozen UNI2 H&E space; gene never co-occurs with mIF or text in training, and no dataset measures RNA + protein + text on the same tissue, so 1:1 pairing accuracy cannot be computed for these pairs. Instead: (A) protein <-> gene, cross-dataset neighbour transfer protein->gene: each ORION test tile (mIF embedding) retrieves its k=20 nearest HEST spots by gene embedding; predicted RNA level = mean log1p-CP10k of the matching gene over those spots; score = Spearman(predicted RNA, measured protein = marker+ cells / nuclei) over tiles (>=20 nuclei). gene->protein: each HEST spot (gene embedding) retrieves its k=20 nearest ORION tiles by mIF embedding; Spearman(predicted protein, measured RNA of the spot). references : he->gene (ORION UNI2 vs HEST gene head; trained modality pair, cross-dataset), he->he (UNI2 vs UNI2, image similarity only), and the same with labels shuffled (chance ~0). galleries : all HEST spots, and bowel-only spots (same organ as ORION colorectal cancer). (B) text <-> gene concept probes: prompts.json through the PathGen-trained text head, scored against HEST test-fold gene embeddings (and UNI2 H&E of the same spots); label = RNA of the concept gene; Spearman + AUROC (top quartile). Outputs results/gene_emergent_protein.csv, results/gene_emergent_text.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.functional as F from scipy.stats import spearmanr from sklearn.metrics import roc_auc_score from common import CACHE, ROOT, RUNS, load_all_tiles OUT = os.path.join(ROOT, "results") H = os.path.join(CACHE, "hest") K = 20 M = pd.read_csv(os.path.join(H, "meta.csv")) 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)} def rna(gs): return np.asarray(L[:, [gpos[g] for g in gs]].mean(1)).ravel() PROT2RNA = {"CD45": ["PTPRC"], "Pan-CK": ["KRT8", "EPCAM"], "E-cadherin": ["CDH1"], "SMA": ["ACTA2"], "Ki67": ["MKI67"], "CD31": ["PECAM1"], "CD68": ["CD68"], "CD163": ["CD163"], "CD3e": ["CD3E"], "CD4": ["CD4"], "CD8a": ["CD8A"], "FOXP3": ["FOXP3"], "CD20": ["MS4A1"], "PD-L1": ["CD274"], "PD-1": ["PDCD1"]} RNA = {m: rna(g) for m, g in PROT2RNA.items()} # ---------------- HEST: gene embeddings of every spot from the fold whose head did not train on it hest = {} for f in glob.glob(os.path.join(RUNS, "gene", "uni2_*_fold*", "test.pt")): enc = os.path.basename(os.path.dirname(f)).split("_", 1)[1].rsplit("_fold", 1)[0] d = torch.load(f, weights_only=False) z = hest.setdefault(enc, {"gene": torch.zeros(len(M), d["gene"].shape[1]), "he": torch.zeros(len(M), d["he"].shape[1]), "fold": np.full(len(M), -1)}) z["gene"][d["idx"]], z["he"][d["idx"]] = d["gene"], d["he"] z["fold"][d["idx"]] = int(f.split("_fold")[1].split(os.sep)[0]) hest = {e: z for e, z in hest.items() if (z["fold"] >= 0).all()} print("gene encoders:", list(hest)) # ---------------- ORION: patient 5-fold test embeddings (mIF encoders) in the same UNI2 space 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 orion = {} for root in ["bench", "bench_fm"]: for f in glob.glob(os.path.join(RUNS, root, "patient_5fold_*_uni2_*", "test_embeddings.pt")): a = json.load(open(os.path.join(os.path.dirname(f), "train.json")))["args"] E = torch.load(f, weights_only=False) o = orion.setdefault(a["if_enc"], {"if": torch.zeros(len(df), E["z"]["if"].shape[1]), "he": torch.zeros(len(df), E["z"]["he"].shape[1]), "seen": np.zeros(len(df), bool)}) t = np.asarray(E["tile_id"]) o["if"][t], o["he"][t], o["seen"][t] = E["z"]["if"], E["z"]["he"], True orion = {e: o for e, o in orion.items() if o["seen"].all()} print("mIF encoders:", list(orion)) tissue = nn_ >= 20 PROT = {m: df[f"{m}_count"].values / np.maximum(nn_, 1) for m in PROT2RNA} def knn_mean(q, g, vals, k=K, bs=1024): out = np.zeros(len(q)) for i in range(0, len(q), bs): s = q[i:i + bs] @ g.T idx = s.topk(k, dim=1).indices.numpy() out[i:i + bs] = vals[idx].mean(1) return out rows = [] rng = np.random.default_rng(0) for genc, hz in hest.items(): for gal in ["all", "bowel"]: gm = np.ones(len(M), bool) if gal == "all" else (M.organ == "Bowel").values for ienc, oz in orion.items(): queries = [("protein->gene (emergent)", oz["if"], hz["gene"]), ("he->gene (cross-dataset)", oz["he"], hz["gene"]), ("he->he (image only)", oz["he"], hz["he"])] for kind, q, g in queries: for m in PROT2RNA: pred = knn_mean(q[tissue], g[gm], RNA[m][gm]) y = PROT[m][tissue] rows.append({"gene_enc": genc, "if_enc": ienc, "gallery": gal, "kind": kind, "marker": m, "spearman": spearmanr(pred, y).correlation, "spearman_shuffled": spearmanr(pred, rng.permutation(y)).correlation}) # reverse: each HEST spot -> nearest ORION tiles for kind, q, g in [("gene->protein (emergent)", hz["gene"][gm], oz["if"][tissue]), ("gene->he (cross-dataset)", hz["gene"][gm], oz["he"][tissue])]: for m in PROT2RNA: pred = knn_mean(q, g, PROT[m][tissue]) y = RNA[m][gm] rows.append({"gene_enc": genc, "if_enc": ienc, "gallery": gal, "kind": kind, "marker": m, "spearman": spearmanr(pred, y).correlation, "spearman_shuffled": spearmanr(pred, rng.permutation(y)).correlation}) print("done", genc, flush=True) P = pd.DataFrame(rows) P.to_csv(os.path.join(OUT, "gene_emergent_protein.csv"), index=False) # ---------------- (B) text <-> gene concept probes prompts = json.load(open(os.path.join(ROOT, "prompts.json"))) CONCEPT_RNA = {"immune": ["PTPRC"], "tumor_epith": ["KRT8", "EPCAM"], "stroma": ["ACTA2"], "proliferation": ["MKI67"], "vessels": ["PECAM1"], "macrophage": ["CD68"]} TXT = sorted(os.path.basename(d).split("_", 1)[1] for d in glob.glob(os.path.join(RUNS, "text_pathgen", "uni2_*")) if os.path.exists(os.path.join(CACHE, f"prompts_text_{os.path.basename(d).split('_', 1)[1]}.npy"))) def load_head(enc, din, dout): h = torch.nn.Sequential(torch.nn.Linear(din, 512), torch.nn.GELU(), torch.nn.Dropout(0.1), torch.nn.Linear(512, 512), torch.nn.GELU(), torch.nn.LayerNorm(512), torch.nn.Linear(512, dout)) h.load_state_dict(torch.load(os.path.join(RUNS, "text_pathgen", f"uni2_{enc}", "head.pt"), map_location="cpu")) return h.eval() trow = [] for t in TXT: prm = torch.from_numpy(np.load(os.path.join(CACHE, f"prompts_text_{t}.npy"))[:, 0].astype(np.float32)) with torch.no_grad(): zp = F.normalize(load_head(t, prm.shape[1], 1536)(prm), dim=-1) keys = list(prompts) for genc, hz in hest.items(): embs = [("he (UNI2)", hz["he"]), (f"gene ({genc})", hz["gene"])] if genc == list(hest)[0] else [(f"gene ({genc})", hz["gene"])] for mod, z in embs: for c, gs in CONCEPT_RNA.items(): s = (z @ zp[keys.index(c)]).numpy() for gal in ["all", "bowel"]: gm = np.ones(len(M), bool) if gal == "all" else (M.organ == "Bowel").values for k in range(5): sel = gm & (hz["fold"] == k) y = rna(gs)[sel] lab = y >= np.quantile(y, 0.75) trow.append({"text_enc": t, "modality": mod, "concept": c, "gallery": gal, "fold": k, "spearman": spearmanr(s[sel], y).correlation, "auroc": roc_auc_score(lab, s[sel]) if 0 < lab.sum() < len(lab) else np.nan}) T = pd.DataFrame(trow) T.to_csv(os.path.join(OUT, "gene_emergent_text.csv"), index=False) pd.set_option("display.width", 250); pd.set_option("display.max_rows", 500) print(P.groupby(["gallery", "kind", "if_enc", "gene_enc"])[["spearman", "spearman_shuffled"]].mean().round(3).to_string()) print(P[P.gallery == "bowel"].groupby(["kind", "marker"]).spearman.mean().unstack("kind").round(2).to_string()) print(T.groupby(["gallery", "text_enc", "modality"])[["auroc", "spearman"]].mean().round(3).to_string())