| """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): |
| 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()) |
|
|
