File size: 5,230 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 | """Bind the gene-expression modality (HEST Visium spots) to the frozen H&E anchor (ImageBind recipe).
anchor (frozen, no head): UNI2-h (same model as the ORION / PathGen anchor) or GigaPath spot features
gene encoders: mlp MLP on log1p-CP10k of the 2000 HVGs (from scratch)
fm:scgpt frozen scGPT-human spot embedding (22_) + MLP head
fm:geneformer frozen Geneformer-V2 spot embedding (22_) + MLP head
loss: symmetric InfoNCE(he, gene); the gene head maps into the raw anchor space, so the same space also holds the
ORION mIF head and the PathGen text head (emergent gene<->protein, gene<->text evaluated in 24_).
CV: 5-fold GroupKFold over study+patient groups (21_), 10% of training groups held out for early stopping.
Outputs runs/gene/<anchor>_<enc>_fold<k>/{head.pt, test.pt}, results/gene_retrieval.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 sklearn.model_selection import GroupKFold
from common import CACHE, ROOT, RUNS, metrics
p = argparse.ArgumentParser()
p.add_argument("--anchors", nargs="+", default=["uni2", "gigapath"])
p.add_argument("--encs", nargs="+", default=["mlp", "fm:scgpt", "fm:geneformer"])
p.add_argument("--epochs", type=int, default=100)
p.add_argument("--bs", type=int, default=256)
p.add_argument("--lr", type=float, default=3e-4)
args = p.parse_args()
dev = "cuda"
H = os.path.join(CACHE, "hest")
M = pd.read_csv(os.path.join(H, "meta.csv"))
def feats(enc):
f = os.path.join(H, "hvg.npy" if enc == "mlp" else f"gene_{enc.split(':')[1]}.npy")
x = np.load(f).astype(np.float32)
return torch.from_numpy(x.reshape(len(x), -1)).to(dev)
def mlp(din, dout, hid=512): # same head as 09_bench / 12_ / 20_
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 r1(a, b):
s = a @ b.T
return float(((s > s.diag()[:, None]).sum(1) == 0).float().mean())
def train(A, G, tr, va, seed=0):
torch.manual_seed(seed); rng = np.random.default_rng(seed)
head = mlp(G.shape[1], A.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=args.lr, weight_decay=0.05)
steps = args.epochs * max(1, len(tr) // args.bs)
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=args.lr, total_steps=steps, pct_start=0.05)
za_va = F.normalize(A[va], dim=-1)
best, state = -1, None
for ep in range(args.epochs):
head.train()
perm = rng.permutation(tr)
for s in range(0, len(perm) - args.bs + 1, args.bs):
b = torch.as_tensor(perm[s:s + args.bs], device=dev)
za, zg = F.normalize(A[b], dim=-1), F.normalize(head(G[b]), dim=-1)
lg = logit.exp().clamp(max=100) * za @ zg.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(G[va]), dim=-1)
v = (r1(za_va, zg) + r1(zg, za_va)) / 2
if v > best:
best, state = v, {k: x.clone() for k, x in head.state_dict().items()}
head.load_state_dict(state); head.eval()
return head, best
folds = list(GroupKFold(5).split(M, groups=M.group))
rows = []
for an in args.anchors:
A = torch.from_numpy(np.load(os.path.join(H, f"he_{an}.npy"))[:, 0].astype(np.float32)).to(dev)
for enc in args.encs:
G = feats(enc)
for k, (trv, te) in enumerate(folds):
grp = np.random.default_rng(k).permutation(sorted(M.group.iloc[trv].unique()))
va_g = set(grp[: max(1, len(grp) // 10)])
is_va = M.group.iloc[trv].isin(va_g).values
tr, va = trv[~is_va], trv[is_va]
head, best = train(A, G, tr, va)
with torch.no_grad():
za, zg = F.normalize(A[te], dim=-1).cpu(), F.normalize(head(G[te]), dim=-1).cpu()
d = os.path.join(RUNS, "gene", f"{an}_{enc.replace(':', '-')}_fold{k}")
os.makedirs(d, exist_ok=True)
torch.save(head.state_dict(), os.path.join(d, "head.pt"))
torch.save({"idx": te, "he": za, "gene": zg}, os.path.join(d, "test.pt"))
for q, g, name in [(za, zg, "he->gene"), (zg, za, "gene->he")]:
rows.append({"anchor": an, "gene_enc": enc, "fold": k, "dir": name, "val_r1": best, **metrics(q, g)})
print(f"{an} {enc} fold{k}: val {best:.3f} test R@1_g100 {np.mean([r['R@1_g100'] for r in rows[-2:]]):.3f} "
f"pair_acc {np.mean([r['pair_acc'] for r in rows[-2:]]):.3f} (n_test={len(te)})", flush=True)
pd.DataFrame(rows).to_csv(os.path.join(ROOT, "results", "gene_retrieval.csv"), index=False)
R = pd.DataFrame(rows)
g = R.groupby(["anchor", "gene_enc", "fold"])[["R@1", "R@1_g100", "R@10", "pair_acc"]].mean()
print(g.groupby(["anchor", "gene_enc"]).agg(["mean", "std"]).round(3).to_string())
|