Blame for experiments/17_bench_fm.py

"""Modality-specific foundation-model encoders for the H&E / mIF / markers / text binding (ImageBind-style).

Same trainer as 09_bench.py, but every modality can use its own *pretrained, modality-specific* foundation model
(frozen; features precomputed) followed by a trainable MLP head:
  --if_enc  fm:<name> -> cache/if_<name>.npy       (N,8,D; 8 dihedral views, aligned with the H&E views)  e.g. fm:virtues
  --mk_enc  fm:<name> -> cache/markers_<name>.npy  (N,1,D)                                                 e.g. fm:geneformer, fm:scgpt
  --txt_enc <name>    -> cache/text_<name>.npy     (N,4,D)                                                 e.g. pathgenclip
  --save_heads        -> runs/bench_fm/<run>/heads.pt (best-epoch encoders + input normalisation, for 31_virtual_validation.py)
The from-scratch encoders of 09_bench.py (resnet18 | convnext_tiny | chanset, linear | mlp | ftt) remain available
as baselines. Runs are written to runs/bench_fm/.

Original 09 description:

With a frozen anchor and anchor-only losses, each trainable encoder receives gradient only from its own
InfoNCE(he, m) term, so one run trains one IF encoder, one marker encoder and one text head *independently*;
each modality keeps its own best epoch (val R@1 of he<->m) and its test embeddings are saved at that epoch.
Cross pairs (if<->markers, if<->text, markers<->text) are then evaluated post hoc for any encoder combination
(10_bench_summary.py).

IF encoders     : resnet18 | convnext_tiny | chanset (shared per-channel ResNet18 + channel tokens + transformer)
marker encoders : linear | mlp | ftt (FT-Transformer, one token per marker)
text encoders   : plip | biomedclip | imagebind   (frozen text tower, trainable MLP head)
"""
import argparse
import json
import os
import time

import numpy as np
import pandas as pd
import timm
import torch
import torch.nn as nn
import torch.nn.functional as F

from common import CACHE, RUNS, SPLITS, dihedral

p = argparse.ArgumentParser()
p.add_argument("--split", default="orion_official")
p.add_argument("--fold", type=int, default=0)
p.add_argument("--anchor", default="uni2")
p.add_argument("--if_enc", default="resnet18", help="resnet18|convnext_tiny|chanset|fm:<name>")
p.add_argument("--mk_enc", default="mlp", help="linear|mlp|ftt|fm:<name>")
p.add_argument("--txt_enc", default="plip", help="any <name> with cache/text_<name>.npy")
p.add_argument("--epochs", type=int, default=40)
p.add_argument("--bs", type=int, default=128)
p.add_argument("--lr", type=float, default=3e-4)
p.add_argument("--seed", type=int, default=0)
p.add_argument("--save_heads", action="store_true", help="also save each modality's encoder at its best epoch (heads.pt)")
args = p.parse_args()
torch.manual_seed(args.seed); np.random.seed(args.seed)
dev = "cuda"
TRAIN = ["if", "markers", "text"]
ENC = {"if": args.if_enc, "markers": args.mk_enc, "text": args.txt_enc}
tag = lambda e: e.replace("fm:", "fm-")
run_dir = os.path.join(RUNS, "bench_fm", f"{args.split}_fold{args.fold}_{args.anchor}_{tag(args.if_enc)}_{tag(args.mk_enc)}_{args.txt_enc}")
os.makedirs(run_dir, exist_ok=True)

# ---------------- data ----------------
sp = pd.read_csv(os.path.join(SPLITS, args.split, f"fold{args.fold}.csv"))
HE = torch.from_numpy(np.load(os.path.join(CACHE, f"he_{args.anchor}.npy")).astype(np.float32))  # N,8,D
D = HE.shape[-1]
IF_FM, MK_FM = args.if_enc.startswith("fm:"), args.mk_enc.startswith("fm:")
if IF_FM:  # frozen mIF foundation-model features, one per dihedral view
    IF = torch.from_numpy(np.load(os.path.join(CACHE, f"if_{args.if_enc[3:]}.npy")).astype(np.float32))
else:
    IF = torch.from_numpy(np.load(os.path.join(CACHE, "if.npy")))
if MK_FM:  # frozen gene-expression foundation-model features of the marker vector
    MK = torch.from_numpy(np.load(os.path.join(CACHE, f"markers_{args.mk_enc[3:]}.npy"))[:, 0].astype(np.float32))
else:
    MK = torch.from_numpy(np.load(os.path.join(CACHE, "markers.npy")))
TX = torch.from_numpy(np.load(os.path.join(CACHE, f"text_{args.txt_enc}.npy")).astype(np.float32))  # N,4,Dt
tr_idx, va_idx, te_idx = (sp.tile_id[sp.split == s].values for s in ["train", "val", "test"])
if IF_FM:
    if_mean, if_std = IF[tr_idx].mean((0, 1)).to(dev), IF[tr_idx].std((0, 1)).to(dev) + 1e-3
else:
    if_mean = IF[tr_idx].float().mean((0, 2, 3)).view(1, -1, 1, 1).to(dev) / 255
    if_std = IF[tr_idx].float().std((0, 2, 3)).view(1, -1, 1, 1).to(dev) / 255 + 1e-3
mk_mean, mk_std = MK[tr_idx].mean(0), MK[tr_idx].std(0) + 1e-3


def batch(idx, train=False):
    idx = torch.as_tensor(idx)
    d = torch.randint(0, 8, (len(idx),)) if train else torch.zeros(len(idx), dtype=torch.long)
    v = torch.randint(0, 4, (len(idx),)) if train else torch.zeros(len(idx), dtype=torch.long)
    if IF_FM:  # precomputed features of the same dihedral view as the H&E anchor
        x = IF[idx, d].to(dev)
    else:
        x = IF[idx].to(dev, non_blocking=True).float() / 255
        x = torch.stack([dihedral(a, int(k)) for a, k in zip(x, d)])
        if train:
            x = x * (1 + 0.1 * torch.randn(len(idx), x.shape[1], 1, 1, device=dev))
    mk = (MK[idx] - mk_mean) / mk_std
    if train and not MK_FM:
        mk = mk + 0.05 * torch.randn_like(mk)
    return {"he": HE[idx, d].to(dev), "if": (x - if_mean) / if_std, "markers": mk.to(dev), "text": TX[idx, v].to(dev)}


# ---------------- encoders ----------------
def head(din):
    return nn.Sequential(nn.LayerNorm(din), nn.Linear(din, D))


class TimmEnc(nn.Module):
    def __init__(self, name):
        super().__init__()
        tag = {"resnet18": "resnet18.a1_in1k", "convnext_tiny": "convnext_tiny.fb_in1k"}[name]
        self.trunk = timm.create_model(tag, pretrained=True, in_chans=17, num_classes=0)
        self.head = head(self.trunk.num_features)

    def forward(self, x):
        return self.head(self.trunk(x))


class ChanSet(nn.Module):
    """Channel-agnostic mIF encoder: each channel -> shared 1-ch ResNet18 (112px) -> token (+ channel embedding)
    -> 2-layer transformer over channel tokens -> mean pool."""
    def __init__(self, C=17, d=256):
        super().__init__()
        self.cnn = timm.create_model("resnet18.a1_in1k", pretrained=True, in_chans=1, num_classes=0)
        self.proj = nn.Linear(512, d)
        self.ch = nn.Parameter(torch.zeros(1, C, d)); nn.init.normal_(self.ch, std=0.02)
        self.tf = nn.TransformerEncoder(nn.TransformerEncoderLayer(d, 4, 4 * d, 0.1, batch_first=True, norm_first=True), 2)
        self.head = head(d)

    def forward(self, x):
        B, C, H, W = x.shape
        x = F.interpolate(x, size=(112, 112), mode="bilinear", antialias=True)
        t = self.proj(self.cnn(x.reshape(B * C, 1, 112, 112))).view(B, C, -1) + self.ch
        return self.head(self.tf(t).mean(1))


class FTT(nn.Module):
    """FT-Transformer: feature tokenizer (value * w_j + b_j) + CLS, 3 layers."""
    def __init__(self, F_in, d=128):
        super().__init__()
        self.w = nn.Parameter(torch.randn(F_in, d) * 0.02)
        self.b = nn.Parameter(torch.zeros(F_in, d))
        self.cls = nn.Parameter(torch.zeros(1, 1, d))
        self.tf = nn.TransformerEncoder(nn.TransformerEncoderLayer(d, 8, 2 * d, 0.1, batch_first=True, norm_first=True), 3)
        self.head = head(d)

    def forward(self, x):
        t = x.unsqueeze(-1) * self.w + self.b
        t = torch.cat([self.cls.expand(len(x), -1, -1), t], 1)
        return self.head(self.tf(t)[:, 0])


def mlp(din, 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, D))


enc = nn.ModuleDict({
    "if": mlp(IF.shape[-1]) if IF_FM else ChanSet() if args.if_enc == "chanset" else TimmEnc(args.if_enc),
    "markers": mlp(MK.shape[1]) if MK_FM else {"linear": lambda: nn.Linear(MK.shape[1], D), "mlp": lambda: mlp(MK.shape[1]),
                                             "ftt": lambda: FTT(MK.shape[1])}[args.mk_enc](),
    "text": mlp(TX.shape[-1]),  # frozen text tower features -> trainable projection
}).to(dev)
logit = nn.ParameterDict({m: nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32)) for m in TRAIN}).to(dev)
params = list(enc.parameters()) + list(logit.parameters())
opt = torch.optim.AdamW(params, lr=args.lr, weight_decay=0.05)
steps = args.epochs * (len(tr_idx) // args.bs)
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=args.lr, total_steps=steps, pct_start=0.05)


def embed(b):
    z = {"he": F.normalize(b["he"], dim=-1)}
    for m in TRAIN:
        z[m] = F.normalize(enc[m](b[m]).float(), dim=-1)
    return z


@torch.no_grad()
def embed_all(idx):
    enc.eval()
    out = {m: [] for m in ["he"] + TRAIN}
    for s in range(0, len(idx), 256):
        with torch.autocast("cuda", dtype=torch.bfloat16):
            z = embed(batch(idx[s:s + 256]))
        for m in out:
            out[m].append(z[m].float().cpu())
    return {m: torch.cat(v) for m, v in out.items()}


def r1(a, b):
    s = a @ b.T
    return float(((s > s.diag()[:, None]).sum(1) == 0).float().mean())


# ---------------- train with per-modality model selection ----------------
best = {m: -1.0 for m in TRAIN}
best_ep = {m: -1 for m in TRAIN}
test_z = {}
best_state = {}
hist = []
t0 = time.time()
for ep in range(args.epochs):
    enc.train()
    perm = np.random.permutation(tr_idx)
    tl = {m: [] for m in TRAIN}
    for s in range(0, len(perm) - args.bs + 1, args.bs):
        b = batch(perm[s:s + args.bs], train=True)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            z = embed(b)
        tgt = torch.arange(args.bs, device=dev)
        loss = 0
        for m in TRAIN:
            lg = logit[m].exp().clamp(max=100) * z["he"] @ z[m].T
            l = (F.cross_entropy(lg, tgt) + F.cross_entropy(lg.T, tgt)) / 2
            loss = loss + l
            tl[m].append(l.item())
        opt.zero_grad(set_to_none=True)
        loss.backward()
        opt.step(); sched.step()
    zv = embed_all(va_idx)
    rec = {"epoch": ep}
    improved = []
    for m in TRAIN:
        v = (r1(zv["he"], zv[m]) + r1(zv[m], zv["he"])) / 2
        rec[f"loss_{m}"], rec[f"val_r1_{m}"] = float(np.mean(tl[m])), v
        if v > best[m]:
            best[m], best_ep[m] = v, ep
            improved.append(m)
    if improved:
        zt = embed_all(te_idx)
        test_z["he"] = zt["he"]
        for m in improved:
            test_z[m] = zt[m]
            if args.save_heads:
                best_state[m] = {k: v.detach().cpu().clone() for k, v in enc[m].state_dict().items()}
    hist.append(rec)
    print(f"ep {ep} " + " ".join(f"{m}:{rec[f'val_r1_{m}']:.3f}" for m in TRAIN) + f" ({time.time() - t0:.0f}s)", flush=True)

torch.save({"z": test_z, "tile_id": te_idx}, os.path.join(run_dir, "test_embeddings.pt"))
if args.save_heads:  # + the input normalisation needed to apply the heads to new data
    torch.save({"state": best_state, "if_mean": if_mean.cpu(), "if_std": if_std.cpu(), "mk_mean": mk_mean, "mk_std": mk_std},
               os.path.join(run_dir, "heads.pt"))
json.dump({"args": vars(args), "best_val": best, "best_epoch": best_ep, "n_train": len(tr_idx), "n_test": len(te_idx),
           "time_s": time.time() - t0, "history": hist}, open(os.path.join(run_dir, "train.json"), "w"), indent=1)
print("done", best, best_ep)