File size: 8,959 Bytes
e8fa066 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | """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())
|