Blame for experiments/29_virtual_completion.py

"""Virtual modality completion: does training text <-> RNA / text <-> protein DIRECTLY on virtual pairs beat the
emergent (H&E-anchor) route?

Training pairs (TCGA all-pairs pool, 6,144 patches, every real text of a patch is a positive):
  rna     : real text <-> STPath virtual RNA of the same patch (cache/stpath/tcga_text_pred.npy -> log1p CP10k, 2000 HEST HVGs)
  protein : real text <-> GigaTIME virtual mIF of the same patch (cache/gigatime/tcga_text_means.npy, 11 channels shared with ORION)
Model: frozen BiomedCLIP text tower + MLP head and an MLP head on the modality vector -> shared 512-d space
(CLIP-style, multi-positive InfoNCE), patient GroupKFold(5) over TCGA, 10% of training patients for early stopping.
Each feature is rank-gaussianised within its own dataset (virtual TCGA / real HEST / real ORION), because virtual
values (STPath expected counts, GigaTIME probabilities) and real measurements live on different scales.

Evaluation only on REAL measurements (never virtual):
  rna     : HEST real RNA, concept prompts vs RNA of the concept gene (protocol of 24_eval_gene_emergent.py (B):
            AUROC top quartile, per HEST study+patient fold, all spots / bowel spots)
  protein : ORION real mIF tile means, concept prompts vs marker+ cell fraction (probe of 18_eval_fm.py, per
            patient_5fold test fold); also ORION GigaTIME virtual means (same tiles) to separate the virtual->real gap
Baselines (emergent, read from results/): text head trained on PathGen H&E-caption pairs, scored against the RNA
(23_) / VirTues mIF (17_) embeddings bound to UNI2.
Outputs results/virtual_completion.csv
"""
import glob
import json
import os

import numpy as np
import pandas as pd
import scipy.sparse as sp
import torch
import torch.nn as nn
import torch.nn.functional as F
from scipy.special import ndtri
from scipy.stats import rankdata, spearmanr
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import GroupKFold

from common import CACHE, ROOT, RUNS, SPLITS, load_all_tiles

dev = "cuda"
OUT = os.path.join(ROOT, "results")
TXT_ENC = "biomedclip"
prompts = json.load(open(os.path.join(ROOT, "prompts.json")))
PKEYS = list(prompts)


def rankgauss(x):
    r = np.apply_along_axis(rankdata, 0, x)
    return ndtri((r - 0.5) / len(x)).astype(np.float32)


# ---------------- TCGA training pool (texts, images, virtual modalities)
I = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv"))
T = pd.read_csv(os.path.join(CACHE, "alltxt_meta.csv"))
TX = torch.from_numpy(np.load(os.path.join(CACHE, f"alltxt_text_{TXT_ENC}.npy"))[:, 0].astype(np.float32)).to(dev)
PR = torch.from_numpy(np.load(os.path.join(CACHE, f"prompts_text_{TXT_ENC}.npy"))[:, 0].astype(np.float32)).to(dev)
t_img = torch.as_tensor(T.img_idx.values, device=dev)
assert (pd.read_csv(os.path.join(CACHE, "stpath", "tcga_text_meta.csv")).key.values == I.key.values).all()
assert (pd.read_csv(os.path.join(CACHE, "gigatime", "tcga_text_meta.csv")).key.values == I.key.values).all()

# RNA: STPath log1p(counts) -> CP10k -> log1p, HEST HVGs
H = os.path.join(CACHE, "hest")
hvg = pd.read_csv(os.path.join(H, "hvg_genes.csv")).gene.values
sg = pd.read_csv(os.path.join(CACHE, "stpath", "tcga_text_genes.csv")).gene.values
c = np.expm1(np.load(os.path.join(CACHE, "stpath", "tcga_text_pred.npy")).astype(np.float32)).clip(0)
c = np.log1p(c / c.sum(1, keepdims=True) * 1e4)
gix = pd.Series(range(len(sg)), index=sg).loc[hvg].values
V_RNA = rankgauss(c[:, gix])
del c

# protein: GigaTIME channels shared with ORION (Actin-D + Transgelin -> SMA)
ch = json.load(open(os.path.join(CACHE, "gigatime", "channels.json")))
G_CH, SHARED = ch["channels"], ch["orion_shared"]
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_)
P_ORION = sorted(set(SHARED.values()), key=ORION_CH.index)


def giga_vec(means):
    return np.stack([np.mean([means[:, G_CH.index(g)] for g, o in SHARED.items() if o == oc], 0) for oc in P_ORION], 1)


V_PROT = rankgauss(giga_vec(np.load(os.path.join(CACHE, "gigatime", "tcga_text_means.npy"))))

# ---------------- real evaluation sets
M = pd.read_csv(os.path.join(H, "meta.csv"))
hest_x = rankgauss(np.load(os.path.join(H, "hvg.npy")).astype(np.float32))
genes = pd.read_csv(os.path.join(H, "genes.csv")).gene.values
X = sp.load_npz(os.path.join(H, "counts.npz")).tocsr()
L = sp.diags(1e4 / np.asarray(X.sum(1)).ravel()) @ X
L.data = np.log1p(L.data)
gpos = {g: i for i, g in enumerate(genes)}
rna = lambda gs: np.asarray(L[:, [gpos[g] for g in gs]].mean(1)).ravel()
CONCEPT_RNA = {"immune": ["PTPRC"], "tumor_epith": ["KRT8", "EPCAM"], "stroma": ["ACTA2"], "proliferation": ["MKI67"],
               "vessels": ["PECAM1"], "macrophage": ["CD68"]}  # as 24_
hest_fold = np.full(len(M), -1)
for f in glob.glob(os.path.join(RUNS, "gene", "uni2_mlp_fold*", "test.pt")):
    hest_fold[torch.load(f, weights_only=False)["idx"]] = int(f.split("_fold")[1].split(os.sep)[0])
assert (hest_fold >= 0).all()

df = load_all_tiles()
nn_ = pd.read_csv(os.path.join(OUT, "tile_stats.csv")).set_index("tile_id").loc[df.tile_id, "n_nuclei"].values
df["n_nuclei"] = nn_
orion_real = rankgauss(np.load(os.path.join(CACHE, "gigatime", "orion_real_if_means.npy"))[:, [ORION_CH.index(o) for o in P_ORION]])
orion_virt = rankgauss(giga_vec(np.load(os.path.join(CACHE, "gigatime", "orion_means.npy"))))
assert (pd.read_csv(os.path.join(CACHE, "gigatime", "orion_meta.csv")).tile_id.values == df.tile_id.values).all()
orion_fold = np.full(len(df), -1)
for k in range(5):
    s = pd.read_csv(os.path.join(SPLITS, "patient_5fold", f"fold{k}.csv"))
    orion_fold[s.tile_id[s.split == "test"].values] = k


def orion_probe(scores, sel, key, label_col):
    """18_eval_fm.probe restricted to tile mask `sel`"""
    sub = sel if key == "cellularity" else sel & (nn_ >= 20)
    y = df[label_col].values[sub].astype(float)
    if label_col != "n_nuclei":
        y = y / np.maximum(nn_[sub], 1)
    lab = y >= np.quantile(y, 0.75)
    return {"spearman": spearmanr(scores[sub], y).correlation,
            "auroc": roc_auc_score(lab, scores[sub]) if 0 < lab.sum() < lab.size else np.nan}


# ---------------- model
def mlp(din, dout=512, hid=512):  # same head as 20_pairing.py
    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(V, tr_t, va_t, epochs=60, bs=256, lr=3e-4, seed=0):
    torch.manual_seed(seed); rng = np.random.default_rng(seed)
    th, mh = mlp(TX.shape[1]).to(dev), mlp(V.shape[1]).to(dev)
    logit = nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32, device=dev))
    opt = torch.optim.AdamW(list(th.parameters()) + list(mh.parameters()) + [logit], lr=lr, weight_decay=0.05)
    sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=lr, total_steps=epochs * max(1, len(tr_t) // bs), pct_start=0.05)
    va = T.loc[va_t].groupby("img_idx").head(1).index.values  # one text per val image
    va_i, va_t1 = torch.as_tensor(T.img_idx.values[va], device=dev), torch.as_tensor(va, device=dev)
    best, state = -1, None
    for ep in range(epochs):
        th.train(); mh.train()
        perm = rng.permutation(tr_t)
        for s in range(0, len(perm) - bs + 1, bs):
            b = torch.as_tensor(perm[s:s + bs], device=dev)
            ii = t_img[b]
            lg = logit.exp().clamp(max=100) * F.normalize(mh(V[ii]), dim=-1) @ F.normalize(th(TX[b]), dim=-1).T
            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(); sched.step()
        th.eval(); mh.eval()
        with torch.no_grad():
            zv, zt = F.normalize(mh(V[va_i]), dim=-1), F.normalize(th(TX[va_t1]), dim=-1)
            v = (r1(zv, zt) + r1(zt, zv)) / 2
        if v > best:
            best, state = v, ({k: x.clone() for k, x in th.state_dict().items()}, {k: x.clone() for k, x in mh.state_dict().items()})
    th.load_state_dict(state[0]); mh.load_state_dict(state[1]); th.eval(); mh.eval()
    return th, mh, best


@torch.no_grad()
def score(th, mh, x):
    zp = F.normalize(th(PR), dim=-1)
    z = F.normalize(mh(torch.from_numpy(x).to(dev)), dim=-1)
    return (z @ zp.T).cpu().numpy()  # (n, n_prompts)


rows = []
folds = list(GroupKFold(5).split(I, groups=I.patient))
for mod, Vnp in [("rna", V_RNA), ("protein", V_PROT)]:
    V = torch.from_numpy(Vnp).to(dev)
    for k, (trv, te) in enumerate(folds):
        pats = np.random.default_rng(k).permutation(sorted(I.patient.iloc[trv].unique()))
        va_p = set(pats[: max(1, len(pats) // 10)])
        isplit = pd.Series("test", index=I.index)
        isplit.iloc[trv] = np.where(I.patient.iloc[trv].isin(va_p), "val", "train")
        ts = isplit.loc[T.img_idx].values
        tr_t, va_t, te_t = (T.index[ts == s].values for s in ["train", "val", "test"])
        th, mh, vbest = train(V, tr_t, va_t)
        # in-domain sanity: held-out TCGA patients, text <-> virtual modality, 1 text per image, R@1 on 100-galleries
        te1 = T.loc[te_t].groupby("img_idx").head(1).index.values
        with torch.no_grad():
            zv = F.normalize(mh(V[T.img_idx.values[te1]]), dim=-1).cpu()
            zt = F.normalize(th(TX[te1]), dim=-1).cpu()
        g = np.random.default_rng(0)
        r = np.mean([(r1(zv[s], zt[s]) + r1(zt[s], zv[s])) / 2 for s in (g.choice(len(te1), 100, replace=False) for _ in range(50))])
        rows.append({"modality": mod, "model_fold": k, "eval": "tcga_virtual_R@1_g100", "concept": "-", "auroc": np.nan, "spearman": r, "val_r1": vbest})
        if mod == "rna":
            S = score(th, mh, hest_x)
            for c, gs in CONCEPT_RNA.items():
                y_all = rna(gs)
                for gal in ["all", "bowel"]:
                    gm = np.ones(len(M), bool) if gal == "all" else (M.organ == "Bowel").values
                    for hk in range(5):
                        sel = gm & (hest_fold == hk)
                        y = y_all[sel]; lab = y >= np.quantile(y, 0.75)
                        s = S[sel, PKEYS.index(c)]
                        rows.append({"modality": mod, "model_fold": k, "eval": f"hest_real_{gal}", "concept": c, "eval_fold": hk,
                                     "auroc": roc_auc_score(lab, s) if 0 < lab.sum() < lab.size else np.nan, "spearman": spearmanr(s, y).correlation})
        else:
            for src, x in [("orion_real", orion_real), ("orion_virtual", orion_virt)]:
                S = score(th, mh, x)
                for i, (c, pr) in enumerate(prompts.items()):
                    for ok in range(5):
                        rows.append({"modality": mod, "model_fold": k, "eval": src, "concept": c, "eval_fold": ok,
                                     **orion_probe(S[:, i], orion_fold == ok, c, pr["label"])})
        print(mod, "fold", k, "val", round(vbest, 3), "tcga R@1_g100", round(r, 3), flush=True)

R = pd.DataFrame(rows)
R.to_csv(os.path.join(OUT, "virtual_completion.csv"), index=False)

# ---------------- comparison with the emergent route (same prompts, same text encoder, same real labels)
pd.set_option("display.width", 250)
print(R[R["eval"] == "tcga_virtual_R@1_g100"].groupby("modality").spearman.agg(["mean", "std"]).round(3))
D = R[R["eval"] != "tcga_virtual_R@1_g100"]
print(D.groupby(["modality", "eval", "concept"]).auroc.mean().unstack("concept").round(3).to_string())
print(D.groupby(["modality", "eval"]).auroc.mean().round(3).to_string())
E = pd.read_csv(os.path.join(OUT, "gene_emergent_text.csv"))
E = E[E.text_enc == TXT_ENC]
print("emergent text<->RNA (HEST real):\n", E.groupby(["gallery", "modality", "concept"]).auroc.mean().unstack("concept").round(3).to_string())
A = pd.read_csv(os.path.join(OUT, "anchor_text_probes.csv"))
A = A[(A.text_enc == TXT_ENC) & (A.split == "patient_5fold") & (A.anchor == "uni2")]
print("emergent text<->protein (ORION real, VirTues):\n", A.groupby(["modality", "concept"]).auroc.mean().unstack("concept").round(3).to_string())
# single-marker reference on ORION real: the matched channel itself as the score
REF = {"immune": "CD3e", "tumor_epith": "Pan-CK", "stroma": "SMA", "proliferation": "Ki67", "macrophage": "CD68", "cellularity": "Hoechst"}
ref = {c: np.mean([orion_probe(orion_real[:, P_ORION.index(ch_)], orion_fold == ok, c, prompts[c]["label"])["auroc"] for ok in range(5)])
       for c, ch_ in REF.items()}
print("single real channel as score (reference):", {c: round(v, 3) for c, v in ref.items()})