Blame for experiments/46_text_enc_eval.py

"""Text encoder comparison + fusion under the final protocol (clean real TCGA texts only; synthetic / virtual = test only).

Text encoders (frozen, cache/alltxt_text_<enc>.npy and cache/orion_gentext_text_<enc>.npy; 45_ / 45b_ / 08_):
  biomedclip (current; raw features as in 42_), conchclip, plip, medcpt, biolord, pubmedbert, medembed, bge, qwen3-0.6b,
  qwen3-4b, keep, quiltnet (whatever exists for both sets). Non-baseline inputs are z-scored with TCGA-train stats.
  cat:<a>+<b>  early fusion (z-scored blocks scaled by 1/sqrt(dim), one head)
  ens:<a>+<b>  late fusion (mean of the L2-normalised outputs of the single-encoder heads; same frozen UNI2 space)
Head: linear or MLP into the frozen UNI2 space, multi-positive InfoNCE (several texts per image), AdamW 3e-4 / wd 0.05 /
OneCycle, batch 256, 1000 steps, early stopping on val R@1 every 50 steps (= 42_ S4 text recipe).
Per fold k (splits/final/tcga_fold<k>.csv; one random text per test image, = 33_/42_):
  trained   H&E<->text                 TCGA test images (real text)
  emergent  mIF<->text(synthetic)      ORION patient_5fold test tiles, mIF = S4 MLP head (runs/rnafm/other_fold<k>.pt, 44_)
  emergent  text<->RNA(virtual)        TCGA test images, RNA = hvg linear head of fold k (runs/rnafm/rna_hvg_fold<k>.pt)
  emergent  text<->mIF(MIPHEI)         TCGA test images, MIPHEI virtual mIF through the same mIF head
Usage: python 46_text_enc_eval.py [--encs ...] [--heads linear mlp] [--folds 0 1 2 3 4]
Output results/textenc_long.csv (resumable), results/textenc_summary.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 scipy.optimize import linear_sum_assignment

from common import CACHE, ROOT, SPLITS

ap = argparse.ArgumentParser()
ap.add_argument("--encs", nargs="+", default=None)
ap.add_argument("--heads", nargs="+", default=["linear", "mlp"])
ap.add_argument("--folds", nargs="+", type=int, default=[0, 1, 2, 3, 4])
args = ap.parse_args()
dev = "cuda"
H = os.path.join(CACHE, "hest")
G100, DRAWS, BS, LR, EVAL_EVERY, STEPS = 100, 50, 256, 3e-4, 50, 1000
f32 = lambda a: torch.from_numpy(np.asarray(a, dtype=np.float32))
nz = lambda x: F.normalize(x, dim=-1)

HE_O = f32(np.load(os.path.join(CACHE, "he_uni2.npy"))).to(dev)
IF_O = f32(np.load(os.path.join(CACHE, "if_virtues_m16.npy"))).to(dev)
IF_V_T = f32(np.load(os.path.join(CACHE, "if_virtues_m16_virtual_miphei_tcga_text.npy"))[:, 0]).to(dev)
HE_H = f32(np.load(os.path.join(H, "he_uni2.npy"))[:, 0]).to(dev)
RNA_H = f32(np.load(os.path.join(H, "hvg.npy"))).to(dev)
HE_T = f32(np.load(os.path.join(CACHE, "allimg_he_uni2.npy"))).to(dev)
T = pd.read_csv(os.path.join(CACHE, "alltxt_meta.csv"))
T_IMG = torch.as_tensor(T.img_idx.values, device=dev)
T = T[~pd.read_csv(os.path.join(CACHE, "alltxt_junk.csv")).junk.values]
hvg = pd.read_csv(os.path.join(H, "hvg_genes.csv")).gene.values


def virtual_rna(st):  # = 42_
    g = pd.read_csv(os.path.join(CACHE, "stpath", f"{st}_genes.csv")).gene.values
    c = np.expm1(np.load(os.path.join(CACHE, "stpath", f"{st}_pred.npy")).astype(np.float32)).clip(0)
    c = np.log1p(c / c.sum(1, keepdims=True) * 1e4)
    return f32(c[:, pd.Series(range(len(g)), index=g).loc[hvg].values]).to(dev)


RNA_V_T = virtual_rna("tcga_text")


def load_enc(e):
    e = e.split("@")[0]  # "<enc>@<n>": same input, head trained with another seed (seed-ensemble control)
    return {s: f32(np.load(os.path.join(CACHE, f"{s}_text_{e}.npy"))[:, 0]).to(dev) for s in ["alltxt", "orion_gentext"]}


def mlp(din, dout=1536, 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 one_text_per_image(t_ids, rng):
    pick = T.loc[t_ids].groupby("img_idx").apply(lambda g: g.index[rng.integers(len(g))], include_groups=False)
    return pick.index.values, pick.values


def g100(a, b, seed=0):  # = 33_ / 42_
    a, b = a.cpu(), b.cpu()
    rng = np.random.default_rng(seed)
    out = {m: [] for m in ["R@1", "R@5", "R@10", "MRR", "pair_acc"]}
    for _ in range(DRAWS):
        s = torch.from_numpy(rng.choice(len(a), min(G100, len(a)), replace=False))
        sim = (a[s] @ b[s].T).numpy(); d = np.diag(sim)
        for S in (sim, sim.T):
            rank = (S > d[:, None]).sum(1)
            for k_ in (1, 5, 10):
                out[f"R@{k_}"].append((rank < k_).mean() / 2)
            out["MRR"].append((1 / (rank + 1)).mean() / 2)
        ri, ci = linear_sum_assignment(-sim)
        out["pair_acc"].append((ci[np.argsort(ri)] == np.arange(len(s))).mean())
    return {m: float(np.sum(v) / DRAWS) 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() + ((s.T > s.diag()[:, None]).sum(1) == 0).float().mean()) / 2)


def train_head(X_tr, A_tr, X_va, A_va, kind, steps, seed, pos_tr=None):  # = 44_
    torch.manual_seed(seed); rng = np.random.default_rng(seed)
    h = (mlp(X_tr.shape[1]) if kind == "mlp" else nn.Linear(X_tr.shape[1], 1536)).to(dev)
    logit = nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32, device=dev))
    opt = torch.optim.AdamW(list(h.parameters()) + [logit], lr=LR, weight_decay=0.05)
    sch = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=LR, total_steps=steps, pct_start=0.05)
    za_va = nz(A_va[:, 0] if A_va.dim() == 3 else A_va)
    best, state = -1, None
    for step in range(1, steps + 1):
        h.train()
        b = torch.as_tensor(rng.choice(len(X_tr), BS, replace=False), device=dev)
        a = A_tr[b, torch.randint(0, A_tr.shape[1], (BS,), device=dev)] if A_tr.dim() == 3 else A_tr[b]
        lg = logit.exp().clamp(max=100) * nz(a) @ nz(h(X_tr[b])).T
        if pos_tr is None:
            t = torch.arange(BS, device=dev)
            loss = (F.cross_entropy(lg, t) + F.cross_entropy(lg.T, t)) / 2
        else:
            ii = pos_tr[b]; pos = (ii[:, None] == ii[None, :]).float(); tgt = pos / pos.sum(1, keepdim=True)
            loss = (torch.sum(-tgt * F.log_softmax(lg, 1), 1).mean() + torch.sum(-tgt * F.log_softmax(lg.T, 1), 1).mean()) / 2
        opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sch.step()
        if step % EVAL_EVERY == 0 or step == steps:
            h.eval()
            with torch.no_grad():
                v = r1(za_va, nz(h(X_va)))
            if v > best:
                best, state = v, {k_: x.detach().clone() for k_, x in h.state_dict().items()}
    h.load_state_dict(state); h.eval()
    return h, best


class Fold:
    def __init__(self, k):
        s = pd.read_csv(os.path.join(SPLITS, "patient_5fold", f"fold{k}.csv"))
        self.o = {x: s.tile_id[s.split == x].values for x in ["train", "val", "test"]}
        self.if_mu, self.if_sd = IF_O[self.o["train"]].mean((0, 1)), IF_O[self.o["train"]].std((0, 1)) + 1e-3
        h = pd.read_csv(os.path.join(SPLITS, "final", f"hest_fold{k}.csv"))
        self.h = {x: h.spot_idx[h.split == x].values for x in ["train", "val", "test"]}
        sp = pd.read_csv(os.path.join(SPLITS, "final", f"tcga_fold{k}.csv")).set_index("img_idx").split
        tsp = sp.loc[T.img_idx].values
        self.t = {x: T.index[tsp == x].values for x in ["train", "val", "test"]}
        self.t_va = one_text_per_image(self.t["val"], np.random.default_rng(0))
        self.t_te = one_text_per_image(self.t["test"], np.random.default_rng(k))

    def nif(self, x):
        return (x - self.if_mu) / self.if_sd

    def mif_head(self, k):
        f = os.path.join(ROOT, "runs", "rnafm", f"other_fold{k}.pt")  # written by 44_rnafm_eval.py
        hm = mlp(512).to(dev)
        hm.load_state_dict(torch.load(f, map_location=dev)["mif"])
        return hm.eval()

    def rna_head(self, k):
        f = os.path.join(ROOT, "runs", "rnafm", f"rna_hvg_fold{k}.pt")
        hr = nn.Linear(2000, 1536).to(dev)
        if os.path.exists(f):
            hr.load_state_dict(torch.load(f, map_location=dev))
            return hr.eval()
        hr, _ = train_head(RNA_H[self.h["train"]], HE_H[self.h["train"]], RNA_H[self.h["val"]], HE_H[self.h["val"]],
                           "linear", 2700, k)
        torch.save(hr.state_dict(), f)
        return hr


def build(enc, tr):
    parts = enc.split(":", 1)[1].split("+") if enc.startswith("cat:") else [enc]
    blocks = []
    for p in parts:
        X = CACHE_ENC[p]
        if p.split("@")[0] == "biomedclip" and len(parts) == 1:
            blocks.append(X); continue  # = 42_ / 33_ (raw BiomedCLIP features)
        mu, sd = X["alltxt"][tr].mean(0), X["alltxt"][tr].std(0) + 1e-6
        w = 1 / np.sqrt(X["alltxt"].shape[1]) if len(parts) > 1 else 1.0
        blocks.append({s: (v - mu) / sd * w for s, v in X.items()})
    return {s: torch.cat([b[s] for b in blocks], 1) for s in blocks[0]}


avail = sorted({f[len("alltxt_text_"):-4] for f in os.listdir(CACHE) if f.startswith("alltxt_text_") and f.endswith(".npy")
                and os.path.exists(os.path.join(CACHE, "orion_gentext_text_" + f[len("alltxt_text_"):]))})
encs = sorted(args.encs or avail, key=lambda e: e.startswith("ens:"))  # late fusions after their single heads
singles = sorted({p for e in encs for p in (e.split(":", 1)[1].split("+") if ":" in e else [e])})
CACHE_ENC = {e: load_enc(e) for e in singles}
print("encoders:", {e: CACHE_ENC[e]["alltxt"].shape[1] for e in singles}, flush=True)
f_out = os.path.join(ROOT, "results", "textenc_long.csv")
rows = pd.read_csv(f_out).to_dict("records") if os.path.exists(f_out) else []
done = {(r["enc"], r["head"], r["fold"]) for r in rows}
for k in args.folds:
    F_ = Fold(k)
    hm, hr = F_.mif_head(k), F_.rna_head(k)
    te_o = F_.o["test"]; ii, tt = (torch.as_tensor(x, device=dev) for x in F_.t_te)
    vi, vt = (torch.as_tensor(x, device=dev) for x in F_.t_va)
    trt = torch.as_tensor(F_.t["train"], device=dev)
    with torch.no_grad():
        z_mif = nz(hm(F_.nif(IF_O[te_o, 0]))); z_rna = nz(hr(RNA_V_T[ii])); z_vmif = nz(hm(F_.nif(IF_V_T[ii])))
        z_he = nz(HE_T[ii, 0])
    heads_out = {}
    for enc in encs:
        for hk in args.heads:
            parts = enc.split(":", 1)[1].split("+") if enc.startswith("ens:") else None
            need = any(enc in e.split(":", 1)[1].split("+") for e in encs if e.startswith("ens:"))
            if (enc, hk, k) in done and not (parts is None and need):
                continue
            if parts is not None:
                if not all((p, hk) in heads_out for p in parts):
                    continue
                Z = {s: nz(sum(heads_out[(p, hk)][s] for p in parts)) for s in heads_out[(parts[0], hk)]}
                best = float("nan")
            else:
                X = build(enc, trt)
                sd_ = k + 1000 * int(enc.split("@")[1]) if "@" in enc else k
                h, best = train_head(X["alltxt"][trt], HE_T[T_IMG[trt]], X["alltxt"][vt], HE_T[vi, 0], hk, STEPS, sd_,
                                     pos_tr=T_IMG[trt])
                with torch.no_grad():
                    Z = {"tcga": nz(h(X["alltxt"][tt])), "orion": nz(h(X["orion_gentext"][te_o]))}
                heads_out[(enc, hk)] = Z
            if (enc, hk, k) in done:
                continue
            res = {("H&E<->text (TCGA)", "trained"): g100(z_he, Z["tcga"]),
                   ("mIF<->text(synthetic) (ORION)", "emergent"): g100(z_mif, Z["orion"]),
                   ("text<->RNA(virtual) (TCGA)", "emergent"): g100(Z["tcga"], z_rna),
                   ("text<->mIF(MIPHEI) (TCGA)", "emergent"): g100(Z["tcga"], z_vmif)}
            for (pair, kind), r in res.items():
                rows.append({"enc": enc, "head": hk, "fold": k, "pair": pair, "kind": kind, "val_R@1": best, **r})
            pd.DataFrame(rows).to_csv(f_out, index=False)
            print(f"fold {k} {enc:<34} {hk:<6} val {best:.3f}  " +
                  "  ".join(f"{p.split(' (')[0]} {r['R@1']:.3f}" for (p, _), r in res.items()), flush=True)

R = pd.DataFrame(rows)
S = R.groupby(["enc", "head", "kind", "pair"], sort=False)[["R@1", "R@10", "pair_acc"]].agg(["mean", "std"]).round(4)
S.to_csv(os.path.join(ROOT, "results", "textenc_summary.csv"))
pd.set_option("display.width", 250)
print(R.pivot_table(index=["enc", "head"], columns="pair", values="R@1", sort=False).round(3).to_string())