Blame for experiments/31_virtual_validation.py

"""Validate the 4-modality ImageBind model with VIRTUAL modalities (validation only; every encoder/head was trained
on real pairs: H&E-mIF ORION (17_), H&E-text PathGen (12_), H&E-RNA HEST (23_); H&E = frozen UNI2).

Virtual data fill the modality pairs that no dataset measures together, giving quadruplets:
  ORION : H&E (real) | mIF (real)        | RNA (STPath, virtual) | text (GPT synthetic caption of the H&E, 27_/32_)
  TCGA  : H&E (real) | text (real)       | mIF (GigaTIME, virtual) | RNA (STPath, virtual)
  HEST  : H&E (real) | RNA (real) + RNA (STPath, virtual)  -> calibration of the virtual RNA
  ORION : mIF real vs mIF virtual (GigaTIME) on the same tiles -> calibration of the virtual mIF
Every embedding lives in the UNI2 space, so any two modalities are compared by cosine similarity.

Heads:  RNA  runs/gene/uni2_mlp_fold<k>/head.pt (MLP on log1p-CP10k of the 2000 HEST HVGs; HEST fold k)
        mIF  runs/bench_fm/patient_5fold_fold<k>_uni2_fm-virtues[_shared]_fm-geneformer_pathgenclip/heads.pt
             (17-channel VirTues for real ORION; 10-channel shared panel (30_) for real-vs-virtual / TCGA)
        text runs/text_pathgen/uni2_biomedclip/head.pt (PathGen captions); TCGA evaluation images exclude every
             patient of its train/val split
Metrics on random 100-item galleries (50 draws): R@1 (tie-aware, both directions averaged) and one-to-one Hungarian
pairing accuracy (chance 0.01). "circular" = both sides are derived from the same H&E image (virtual-from-H&E), so the
score is an upper bound of what real data would give, not evidence of biological agreement.
Outputs results/virtual_validation.csv
"""
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, RUNS, SPLITS, load_all_tiles

OUT = os.path.join(ROOT, "results")
H = os.path.join(CACHE, "hest")
G100, DRAWS = 100, 50


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 load_mlp(sd):
    h = mlp(sd["0.weight"].shape[1]); h.load_state_dict(sd); return h.eval()


def nz(x):
    return F.normalize(torch.as_tensor(np.asarray(x), dtype=torch.float32), dim=-1)


def g100(a, b, seed=0):
    """mean over random 100-galleries of tie-aware R@1 (both directions) and Hungarian pairing accuracy"""
    rng = np.random.default_rng(seed)
    r, p = [], []
    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)[:, None]
        r.append(((sim > d).sum(1) == 0).mean() / 2 + ((sim.T > d).sum(1) == 0).mean() / 2)
        ri, ci = linear_sum_assignment(-sim)
        p.append((ci[np.argsort(ri)] == np.arange(len(s))).mean())
    return float(np.mean(r)), float(np.mean(p))


rows = []


def add(dataset, fold, pair, za, zb, kind, circular):
    r, p = g100(za, zb)
    rows.append({"dataset": dataset, "fold": fold, "pair": pair, "kind": kind, "circular": circular, "n": len(za),
                 "R@1_g100": r, "pair_acc_g100": p})


# ---------------- RNA: STPath log1p(counts) -> log1p CP10k on the HEST HVGs (same input as the HEST-trained head)
hvg = pd.read_csv(os.path.join(H, "hvg_genes.csv")).gene.values


def virtual_rna(st):
    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 torch.from_numpy(c[:, pd.Series(range(len(g)), index=g).loc[hvg].values])


rna_heads = [load_mlp(torch.load(os.path.join(RUNS, "gene", f"uni2_mlp_fold{k}", "head.pt"), map_location="cpu")) for k in range(5)]
text_head = load_mlp(torch.load(os.path.join(RUNS, "text_pathgen", "uni2_biomedclip", "head.pt"), map_location="cpu"))


def mif_heads(enc, k, split="patient_5fold"):
    f = os.path.join(RUNS, "bench_fm", f"{split}_fold{k}_uni2_fm-{enc}_fm-geneformer_pathgenclip", "heads.pt")
    if not os.path.exists(f):
        return None
    d = torch.load(f, map_location="cpu")
    h = load_mlp(d["state"]["if"])
    return lambda x: h((torch.as_tensor(x, dtype=torch.float32) - d["if_mean"]) / d["if_std"])


# ---------------- (A) HEST: H&E <-> real RNA vs H&E <-> virtual RNA, same spots, HEST test folds
M = pd.read_csv(os.path.join(H, "meta.csv"))
hv = pd.read_csv(os.path.join(CACHE, "stpath", "hest_val_meta.csv"))
assert (hv["sample"].values == M["sample"].values).all() and (hv.spot.values == M.spot.values).all()
he_h = np.load(os.path.join(H, "he_uni2.npy"))[:, 0].astype(np.float32)
real_h = torch.from_numpy(np.load(os.path.join(H, "hvg.npy")).astype(np.float32))
virt_h = virtual_rna("hest_val")
with torch.no_grad():
    for k in range(5):
        te = torch.load(os.path.join(RUNS, "gene", f"uni2_mlp_fold{k}", "test.pt"), weights_only=False)["idx"]
        zr, zv, zh = nz(rna_heads[k](real_h[te])), nz(rna_heads[k](virt_h[te])), nz(he_h[te])
        for sub, m in [("all", np.ones(len(te), bool)), ("STPath-heldout", hv.hest_bench_heldout.values[te])]:
            if m.sum() < G100:
                continue
            add(f"HEST ({sub})", k, "H&E<->RNA real", zh[m], zr[m], "trained (real)", False)
            add(f"HEST ({sub})", k, "H&E<->RNA virtual", zh[m], zv[m], "virtual calibration", True)
            add(f"HEST ({sub})", k, "RNA real<->RNA virtual", zr[m], zv[m], "virtual calibration", False)
print("HEST done", flush=True)

# ---------------- (B) ORION: patient 5-fold test tiles
df = load_all_tiles()
assert (pd.read_csv(os.path.join(CACHE, "stpath", "orion_meta.csv")).key.values == df.image_path.values).all()
he_o = np.load(os.path.join(CACHE, "he_uni2.npy"))[:, 0].astype(np.float32)
rna_o = virtual_rna("orion")
if17 = np.load(os.path.join(CACHE, "if_virtues.npy"))[:, 0].astype(np.float32)
f_sh = os.path.join(CACHE, "if_virtues_shared.npy")
f_vo = os.path.join(CACHE, "if_virtues_shared_virtual_orion.npy")
ifs_r = np.load(f_sh)[:, 0].astype(np.float32) if os.path.exists(f_sh) else None
ifs_v = np.load(f_vo)[:, 0].astype(np.float32) if os.path.exists(f_vo) else None
f_gt = os.path.join(CACHE, "orion_gentext_text_biomedclip.npy")  # synthetic GPT captions of the ORION H&E (32_ + 08_)
gt_o = torch.from_numpy(np.load(f_gt)[:, 0].astype(np.float32)) if os.path.exists(f_gt) else None
with torch.no_grad():
    for k in range(5):
        s = pd.read_csv(os.path.join(SPLITS, "patient_5fold", f"fold{k}.csv"))
        te = s.tile_id[s.split == "test"].values
        zh, zr = nz(he_o[te]), nz(rna_heads[k](rna_o[te]))
        add("ORION", k, "H&E<->RNA virtual", zh, zr, "virtual (cross-dataset head)", True)
        h17 = mif_heads("virtues", k)
        if h17 is not None:
            zp = nz(h17(if17[te]))
            add("ORION", k, "H&E<->mIF real (17 ch)", zh, zp, "trained (real)", False)
            add("ORION", k, "mIF real (17 ch)<->RNA virtual", zp, zr, "emergent, real x virtual", False)
        hs = mif_heads("virtues_shared", k)
        if hs is not None and ifs_r is not None:
            zs = nz(hs(ifs_r[te]))
            add("ORION", k, "H&E<->mIF real (10 ch)", zh, zs, "trained (real)", False)
            add("ORION", k, "mIF real (10 ch)<->RNA virtual", zs, zr, "emergent, real x virtual", False)
            if ifs_v is not None:
                zv = nz(hs(ifs_v[te]))
                add("ORION", k, "H&E<->mIF virtual (10 ch)", zh, zv, "virtual calibration", True)
                add("ORION", k, "mIF real<->mIF virtual (10 ch)", zs, zv, "virtual calibration", False)
                add("ORION", k, "mIF virtual<->RNA virtual", zv, zr, "emergent, virtual x virtual", True)
        if gt_o is not None:  # synthetic text is written from the same H&E -> circular with H&E and H&E-derived data
            zt = nz(text_head(gt_o[te]))
            add("ORION", k, "H&E<->text synthetic", zh, zt, "virtual (text head from PathGen)", True)
            add("ORION", k, "text synthetic<->RNA virtual", zt, zr, "emergent, virtual x virtual", True)
            if h17 is not None:
                add("ORION", k, "text synthetic<->mIF real (17 ch)", zt, zp, "emergent, real x virtual", False)
            if hs is not None and ifs_r is not None:
                add("ORION", k, "text synthetic<->mIF real (10 ch)", zt, zs, "emergent, real x virtual", False)
print("ORION done", flush=True)

# ---------------- (C) TCGA quadruplets: H&E, real text, virtual mIF, virtual RNA
I = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv"))
T = pd.read_csv(os.path.join(CACHE, "alltxt_meta.csv"))
P = pd.read_csv(os.path.join(CACHE, "pathgen_meta.csv"))
seen = set(P.patient[P.split.isin(["train", "val"])])
ev = I.index[~I.patient.isin(seen)].values
rng = np.random.default_rng(0)
t1 = T[T.img_idx.isin(ev)].groupby("img_idx").apply(lambda g: g.index[rng.integers(len(g))], include_groups=False)
ev = t1.index.values  # images with >= 1 text, patients unseen by the text head
print(f"TCGA eval: {len(ev)} images, {I.patient.iloc[ev].nunique()} patients (excluded {len(seen)} text-head patients)", flush=True)
he_t = np.load(os.path.join(CACHE, "allimg_he_uni2.npy"))[ev, 0].astype(np.float32)
tx = torch.from_numpy(np.load(os.path.join(CACHE, "alltxt_text_biomedclip.npy"))[t1.values, 0].astype(np.float32))
rna_t = virtual_rna("tcga_text")[ev]
f_vt = os.path.join(CACHE, "if_virtues_shared_virtual_tcga_text.npy")
ifv_t = np.load(f_vt)[ev, 0].astype(np.float32) if os.path.exists(f_vt) else None
with torch.no_grad():
    zh, zt = nz(he_t), nz(text_head(tx))
    for k in range(5):
        zr = nz(rna_heads[k](rna_t))
        Z = {"H&E": zh, "text": zt, "RNA virtual": zr}
        hs = mif_heads("virtues_shared", k)
        if hs is not None and ifv_t is not None:
            Z["mIF virtual"] = nz(hs(ifv_t))
        names = list(Z)
        for i in range(len(names)):
            for j in range(i + 1, len(names)):
                a, b = names[i], names[j]
                if (a, b) == ("H&E", "text") and k > 0:
                    continue  # does not depend on the fold heads
                virt = [x for x in (a, b) if "virtual" in x]
                circ = len(virt) == 2 or (len(virt) == 1 and "H&E" in (a, b))
                kind = "trained (real)" if (a, b) == ("H&E", "text") else \
                    "emergent, real x virtual" if len(virt) == 1 and "H&E" not in (a, b) else \
                    "emergent, virtual x virtual" if len(virt) == 2 else "virtual (cross-dataset head)"
                add("TCGA", k, f"{a}<->{b}", Z[a], Z[b], kind, circ)
print("TCGA done", flush=True)

# ---------------- (D) mIF as 10 channel MEANS (GigaTIME per-tile means are reliable, its pixel maps are not):
# MLP head on rank-gaussianised shared-panel means, trained on REAL ORION (patient_5fold fold k train/val) vs UNI2 H&E;
# applied unchanged to virtual means (rank-gaussianised within their own set)
import json
from scipy.special import ndtri
from scipy.stats import rankdata

rg = lambda x: torch.from_numpy(ndtri((np.apply_along_axis(rankdata, 0, x) - 0.5) / len(x)).astype(np.float32))
ORION_CH = ["Hoechst", "CD31", "CD45", "CD68", "CD4", "FOXP3", "CD8a", "CD45RO", "CD20", "PD-L1", "CD3e",
            "CD163", "E-cadherin", "PD-1", "Ki67", "Pan-CK", "SMA"]  # order of orion_real_if_means.npy (25_)
PANEL = [("Hoechst", ["DAPI"]), ("CD68", ["CD68"]), ("CD4", ["CD4"]), ("CD8a", ["CD8"]), ("CD20", ["CD20"]),
         ("PD-L1", ["PD-L1"]), ("CD3e", ["CD3"]), ("Ki67", ["Ki67"]), ("Pan-CK", ["CK"]), ("SMA", ["Actin-D", "Transgelin"])]  # = 30_
gch = json.load(open(os.path.join(CACHE, "gigatime", "channels.json")))["channels"]
gmean = lambda m: np.stack([m[:, [gch.index(g) for g in gs]].mean(1) for _, gs in PANEL], 1)
mr_o = rg(np.load(os.path.join(CACHE, "gigatime", "orion_real_if_means.npy"))[:, [ORION_CH.index(o) for o, _ in PANEL]])
mv_o = rg(gmean(np.load(os.path.join(CACHE, "gigatime", "orion_means.npy"))))
mv_t = rg(gmean(np.load(os.path.join(CACHE, "gigatime", "tcga_text_means.npy"))))[ev]
HE8 = torch.from_numpy(np.load(os.path.join(CACHE, "he_uni2.npy")).astype(np.float32))


def train_means_head(tr, va, epochs=150, bs=256, lr=1e-3, seed=0):
    torch.manual_seed(seed); rng_ = np.random.default_rng(seed)
    h = mlp(mr_o.shape[1])
    logit = nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32))
    opt = torch.optim.AdamW(list(h.parameters()) + [logit], lr=lr, weight_decay=0.05)
    best, state = -1, None
    for ep in range(epochs):
        h.train()
        perm = rng_.permutation(tr)
        for s0 in range(0, len(perm) - bs + 1, bs):
            b = torch.as_tensor(perm[s0:s0 + bs])
            za = F.normalize(HE8[b, torch.randint(0, 8, (len(b),))], dim=-1)
            lg = logit.exp().clamp(max=100) * za @ F.normalize(h(mr_o[b]), dim=-1).T
            t = torch.arange(len(b))
            loss = (F.cross_entropy(lg, t) + F.cross_entropy(lg.T, t)) / 2
            opt.zero_grad(); loss.backward(); opt.step()
        h.eval()
        with torch.no_grad():
            a, m = nz(HE8[va, 0]), nz(h(mr_o[va]))
            sim = a @ m.T
            v = float(((sim > sim.diag()[:, None]).sum(1) == 0).float().mean())
        if v > best:
            best, state = v, {k_: x.clone() for k_, x in h.state_dict().items()}
    h.load_state_dict(state); h.eval()
    return h


for k in range(5):
    s = pd.read_csv(os.path.join(SPLITS, "patient_5fold", f"fold{k}.csv"))
    tr, va, te = (s.tile_id[s.split == x].values for x in ["train", "val", "test"])
    h = train_means_head(tr, va)
    with torch.no_grad():
        zh, zr = nz(he_o[te]), nz(rna_heads[k](rna_o[te]))
        zmr, zmv = nz(h(mr_o[te])), nz(h(mv_o[te]))
        add("ORION", k, "H&E<->mIF real (10 means)", zh, zmr, "trained (real)", False)
        add("ORION", k, "H&E<->mIF virtual (10 means)", zh, zmv, "virtual calibration", True)
        add("ORION", k, "mIF real<->mIF virtual (10 means)", zmr, zmv, "virtual calibration", False)
        add("ORION", k, "mIF real (10 means)<->RNA virtual", zmr, zr, "emergent, real x virtual", False)
        add("ORION", k, "mIF virtual (10 means)<->RNA virtual", zmv, zr, "emergent, virtual x virtual", True)
        if gt_o is not None:
            zt = nz(text_head(gt_o[te]))
            add("ORION", k, "text synthetic<->mIF real (10 means)", zt, zmr, "emergent, real x virtual", False)
            add("ORION", k, "text synthetic<->mIF virtual (10 means)", zt, zmv, "emergent, virtual x virtual", True)
        zmt = nz(h(mv_t))
        add("TCGA", k, "H&E<->mIF virtual (10 means)", nz(he_t), zmt, "virtual (cross-dataset head)", True)
        add("TCGA", k, "text<->mIF virtual (10 means)", nz(text_head(tx)), zmt, "emergent, real x virtual", False)
        add("TCGA", k, "RNA virtual<->mIF virtual (10 means)", nz(rna_heads[k](rna_t)), zmt, "emergent, virtual x virtual", True)
    print("means head fold", k, flush=True)

# ---------------- (E) calibration of the synthetic text: TCGA patches with REAL text AND a synthetic caption (27_ tcga_calib)
f_cm, f_cf = os.path.join(CACHE, "tcga_calib_meta.csv"), os.path.join(CACHE, "tcga_calib_text_biomedclip.npy")
if os.path.exists(f_cf):
    CM = pd.read_csv(f_cm)
    pos = pd.Series(range(len(ev)), index=ev)
    assert CM.img_idx.isin(ev).all()
    j = pos.loc[CM.img_idx].values  # rows of the TCGA eval arrays
    with torch.no_grad():
        zs = nz(text_head(torch.from_numpy(np.load(f_cf)[:, 0].astype(np.float32))))
        zr_txt, zh_c = nz(text_head(tx))[j], nz(he_t)[j]
        add("TCGA calib", 0, "H&E<->text real", zh_c, zr_txt, "trained (real)", False)
        add("TCGA calib", 0, "H&E<->text synthetic", zh_c, zs, "virtual calibration", True)
        add("TCGA calib", 0, "text real<->text synthetic", zr_txt, zs, "virtual calibration", False)
        for k in range(5):
            zrna = nz(rna_heads[k](rna_t))[j]
            add("TCGA calib", k, "text real<->RNA virtual", zr_txt, zrna, "emergent, real x virtual", False)
            add("TCGA calib", k, "text synthetic<->RNA virtual", zs, zrna, "emergent, virtual x virtual", True)
    print("calibration done", len(CM), flush=True)

R = pd.DataFrame(rows)
R.to_csv(os.path.join(OUT, "virtual_validation.csv"), index=False)
pd.set_option("display.width", 250); pd.set_option("display.max_colwidth", 60)
S = R.groupby(["dataset", "pair", "kind", "circular"])[["R@1_g100", "pair_acc_g100", "n"]].agg(["mean", "std"]).round(3)
print(S.drop(columns=[("n", "std")]).to_string())