| """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 = {} |
| 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)) |
|
|
| |
| 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}) |
| |
| 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) |
|
|
| |
| 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()) |
|
|
