Migrate release files 151-180

.gitattributes

@@ -92,3 +92,8 @@ examples/tcga_he/05162.png filter=lfs diff=lfs merge=lfs -text
 examples/tcga_he/05256.png filter=lfs diff=lfs merge=lfs -text
 examples/tcga_he/05403.png filter=lfs diff=lfs merge=lfs -text
 examples/tcga_he/05631.png filter=lfs diff=lfs merge=lfs -text
+examples/tcga_he/05632.png filter=lfs diff=lfs merge=lfs -text
+examples/tcga_he/05660.png filter=lfs diff=lfs merge=lfs -text
+examples/tcga_he/05757.png filter=lfs diff=lfs merge=lfs -text
+examples/tcga_he/05776.png filter=lfs diff=lfs merge=lfs -text
+examples/tcga_he/06007.png filter=lfs diff=lfs merge=lfs -text

experiments/01_analyze_data.py

@@ -0,0 +1,83 @@
+"""Descriptive analysis of the ORION tri-modal (+markers) data. Writes results/data_summary.json + figures."""
+import json
+import os
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import tifffile
+from PIL import Image
+
+from common import ORION, ROOT, load_all_tiles, marker_columns
+
+out = os.path.join(ROOT, "results")
+os.makedirs(out, exist_ok=True)
+df = load_all_tiles()
+prop, cnt = marker_columns(df)
+
+summary = {
+    "n_tiles": len(df),
+    "n_slides": int(df.in_slide_name.nunique()),
+    "n_patients": int(df.patient.nunique()),
+    "official_split_tiles": df.orion_split.value_counts().to_dict(),
+    "official_split_slides": df.groupby("orion_split").in_slide_name.nunique().to_dict(),
+    "if_batches_tiles": df.if_batch.value_counts().to_dict(),
+    "if_batches_patients": df.groupby("if_batch").patient.nunique().to_dict(),
+    "tiles_per_slide": df.groupby("in_slide_name").size().describe().round(1).to_dict(),
+}
+
+# per-tile stats from images (all tiles; small set)
+he_mean, if_ch_mean, n_nuc, nuc_cov = [], [], [], []
+for _, r in df.iterrows():
+    he = np.asarray(Image.open(os.path.join(ORION, r.image_path)).convert("RGB"))
+    iff = tifffile.imread(os.path.join(ORION, r.target_path))
+    nuc = tifffile.imread(os.path.join(ORION, r.nuclei_path))
+    he_mean.append(he.mean())
+    if_ch_mean.append(iff.reshape(-1, iff.shape[-1]).mean(0))
+    n_nuc.append(len(np.unique(nuc)) - 1)
+    nuc_cov.append((nuc > 0).mean())
+df["he_mean"], df["n_nuclei"], df["nuc_cov"] = he_mean, n_nuc, nuc_cov
+if_ch_mean = np.stack(if_ch_mean)
+summary["shapes"] = {"he": list(he.shape), "if": list(iff.shape), "nuclei": list(nuc.shape),
+                     "markers": len(prop) + len(cnt)}
+summary["if_channel_names_from_prop_cols"] = [c[:-5] for c in prop]
+summary["n_nuclei_per_tile"] = df.n_nuclei.describe().round(1).to_dict()
+summary["frac_tiles_lt20_nuclei"] = float((df.n_nuclei < 20).mean())
+summary["he_mean_vs_n_nuclei_spearman"] = float(df[["he_mean", "n_nuclei"]].corr("spearman").iloc[0, 1])
+summary["marker_prop_mean"] = df[prop].mean().round(4).to_dict()
+summary["marker_count_mean"] = df[cnt].mean().round(2).to_dict()
+summary["if_channel_mean_intensity"] = dict(zip([c[:-5] for c in prop], if_ch_mean.mean(0).round(2).tolist()))
+df[["tile_id", "in_slide_name", "patient", "if_batch", "orion_split", "n_nuclei", "nuc_cov", "he_mean"]].to_csv(
+    os.path.join(out, "tile_stats.csv"), index=False)
+json.dump(summary, open(os.path.join(out, "data_summary.json"), "w"), indent=1)
+print(json.dumps(summary, indent=1))
+
+# figure: example quadruplets
+rng = np.random.default_rng(0)
+ex = df[df.n_nuclei > 100].sample(4, random_state=0)
+fig, ax = plt.subplots(4, 5, figsize=(15, 12))
+for i, (_, r) in enumerate(ex.iterrows()):
+    he = Image.open(os.path.join(ORION, r.image_path)).convert("RGB")
+    iff = tifffile.imread(os.path.join(ORION, r.target_path))
+    nuc = tifffile.imread(os.path.join(ORION, r.nuclei_path))
+    ax[i, 0].imshow(he); ax[i, 0].set_title(f"H&E  {r.orion_slide_id}")
+    # Hoechst / Pan-CK / CD45 as RGB-ish composite
+    names = [c[:-5] for c in prop]
+    comp = np.stack([iff[..., names.index(n)] for n in ["Pan-CK", "CD45", "Hoechst"]], -1)
+    ax[i, 1].imshow(comp); ax[i, 1].set_title("IF: Pan-CK/CD45/Hoechst")
+    ax[i, 2].imshow(iff[..., names.index("SMA")], cmap="magma"); ax[i, 2].set_title("IF: SMA")
+    ax[i, 3].imshow(nuc > 0, cmap="gray"); ax[i, 3].set_title(f"nuclei ({r.n_nuclei})")
+    ax[i, 4].barh(range(len(cnt)), np.log1p(r[cnt].values.astype(float)))
+    ax[i, 4].set_yticks(range(len(cnt))); ax[i, 4].set_yticklabels([c[:-6] for c in cnt], fontsize=6)
+    ax[i, 4].set_title("markers log1p(count)")
+    for a in ax[i, :4]:
+        a.axis("off")
+plt.tight_layout(); plt.savefig(os.path.join(out, "example_quadruplets.png"), dpi=80)
+
+fig, ax = plt.subplots(1, 3, figsize=(15, 4))
+ax[0].hist(df.n_nuclei, bins=50); ax[0].set_title("nuclei per tile"); ax[0].set_yscale("log")
+df.groupby("if_batch").size().plot.bar(ax=ax[1], title="tiles per IF batch")
+ax[2].imshow(np.corrcoef(np.concatenate([df[prop].values, np.log1p(df[cnt].values)], 1).T), cmap="RdBu_r", vmin=-1, vmax=1)
+ax[2].set_title("marker vector correlation (17 prop + 16 log count)")
+plt.tight_layout(); plt.savefig(os.path.join(out, "data_overview.png"), dpi=80)

experiments/02_make_splits.py

@@ -0,0 +1,51 @@
+"""Define evaluation splits. Each split file: splits/<name>/fold<k>.csv with columns tile_id, split.
+
+  orion_official : the user's train/val/test_dataframe.csv (slide-disjoint)
+  random_tile    : tile-level random 80/10/10 (slides shared across splits -> optimistic upper bound)
+  patient_5fold  : 5-fold CV grouped by patient (CRC33_01/_02 kept together); val = 3 patients from train folds
+  batch_lobo     : leave-one-IF-acquisition-batch-out (C59kX/mX/nX/qX); val = 2 patients from training batches
+"""
+import os
+
+import numpy as np
+
+from common import SPLITS, load_all_tiles
+
+df = load_all_tiles()
+
+
+def save(name, k, split):
+    os.makedirs(os.path.join(SPLITS, name), exist_ok=True)
+    out = df[["tile_id", "in_slide_name", "patient", "if_batch"]].copy()
+    out["split"] = split
+    out.to_csv(os.path.join(SPLITS, name, f"fold{k}.csv"), index=False)
+    print(name, k, out.split.value_counts().to_dict(),
+          {s: out[out.split == s].patient.nunique() for s in ["train", "val", "test"]})
+
+
+# 1. official
+save("orion_official", 0, df.orion_split.values)
+
+# 2. random tile-level
+rng = np.random.default_rng(0)
+u = rng.random(len(df))
+save("random_tile", 0, np.where(u < 0.8, "train", np.where(u < 0.9, "val", "test")))
+
+# 3. patient-grouped 5-fold
+pats = np.array(sorted(df.patient.unique()))
+rng = np.random.default_rng(0)
+perm = rng.permutation(pats)
+folds = np.array_split(perm, 5)
+for k in range(5):
+    test_p = set(folds[k])
+    rest = [p for p in perm if p not in test_p]
+    val_p = set(rest[:3])
+    split = np.where(df.patient.isin(test_p), "test", np.where(df.patient.isin(val_p), "val", "train"))
+    save("patient_5fold", k, split)
+
+# 4. leave-one-IF-batch-out
+for k, b in enumerate(sorted(df.if_batch.unique())):
+    rest = sorted(df[df.if_batch != b].patient.unique())
+    val_p = set(np.random.default_rng(k).choice(rest, 2, replace=False))
+    split = np.where(df.if_batch == b, "test", np.where(df.patient.isin(val_p), "val", "train"))
+    save("batch_lobo", k, split)

experiments/03_cache_data.py

@@ -0,0 +1,35 @@
+"""Resize every modality to 224x224 once and store as .npy (uint8) for fast training."""
+import os
+
+import numpy as np
+import tifffile
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from scipy import ndimage
+
+from common import CACHE, IMG, ORION, load_all_tiles, marker_matrix
+
+os.makedirs(CACHE, exist_ok=True)
+df = load_all_tiles()
+N = len(df)
+he = np.lib.format.open_memmap(os.path.join(CACHE, "he.npy"), "w+", np.uint8, (N, 3, IMG, IMG))
+iff = np.lib.format.open_memmap(os.path.join(CACHE, "if.npy"), "w+", np.uint8, (N, 17, IMG, IMG))
+nuc = np.lib.format.open_memmap(os.path.join(CACHE, "nuclei.npy"), "w+", np.uint8, (N, 2, IMG, IMG))
+
+for i, r in df.iterrows():
+    h = Image.open(os.path.join(ORION, r.image_path)).convert("RGB").resize((IMG, IMG), Image.BICUBIC)
+    he[i] = np.asarray(h).transpose(2, 0, 1)
+    x = torch.from_numpy(tifffile.imread(os.path.join(ORION, r.target_path)).astype(np.float32)).permute(2, 0, 1)[None]
+    iff[i] = F.interpolate(x, size=(IMG, IMG), mode="bilinear", antialias=True, align_corners=False)[0].round().clamp(0, 255).byte().numpy()
+    lab = tifffile.imread(os.path.join(ORION, r.nuclei_path))
+    # boundary = pixels whose label differs from a neighbour (separates touching nuclei)
+    bnd = (ndimage.maximum_filter(lab, 3) != ndimage.minimum_filter(lab, 3)) & (lab > 0)
+    m = torch.from_numpy(np.stack([lab > 0, bnd]).astype(np.float32))[None]
+    nuc[i] = (F.interpolate(m, size=(IMG, IMG), mode="area")[0] * 255).round().byte().numpy()
+    if i % 500 == 0:
+        print(i, N, flush=True)
+
+np.save(os.path.join(CACHE, "markers.npy"), marker_matrix(df))
+he.flush(); iff.flush(); nuc.flush()
+print("done")

experiments/04_extract_anchor.py

@@ -0,0 +1,25 @@
+"""Frozen ImageBind-Huge vision embeddings of every H&E tile under each of the 8 dihedral transforms.
+Output cache/he_imagebind.npy : (N, 8, 1024) float16, L2-normalised (ImageBind postprocessor)."""
+import os
+
+import numpy as np
+import torch
+
+from common import CACHE, dihedral, load_imagebind_vision
+
+MEAN = torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1)
+STD = torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1)
+
+he = np.load(os.path.join(CACHE, "he.npy"), mmap_mode="r")
+model, MT = load_imagebind_vision()
+out = np.zeros((len(he), 8, 1024), np.float16)
+bs = 64
+with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16):
+    for s in range(0, len(he), bs):
+        x = (torch.from_numpy(np.ascontiguousarray(he[s:s + bs])).float() / 255 - MEAN) / STD
+        for d in range(8):
+            e = model({MT.VISION: dihedral(x, d).cuda()})[MT.VISION]
+            out[s:s + bs, d] = e.float().cpu().numpy()
+        print(s, len(he), flush=True)
+np.save(os.path.join(CACHE, "he_imagebind.npy"), out)
+print("saved", out.shape)

experiments/04b_extract_anchor_fm.py

@@ -0,0 +1,94 @@
+"""Frozen pathology foundation-model embeddings of every H&E tile under the 8 dihedral transforms.
+Usage: python 04b_extract_anchor_fm.py {uni2|conch|conchclip|gigapath|plip|phikonv2} [--set <name>]
+Output: cache/he_<name>.npy (N, 8, D) float16 (raw features; L2-normalised later in training).
+
+  uni2     : MahmoodLab/UNI2-h, ViT-H/14 (reg tokens), 224px, ImageNet norm, CLS -> 1536-d
+  gigapath : prov-gigapath/prov-gigapath tile encoder, ViT-g/14, 224px, ImageNet norm -> 1536-d
+  conch    : MahmoodLab/CONCH v1 vision tower (ViT-B/16), 448px, OpenAI-CLIP norm,
+             encode_image(proj_contrast=False, normalize=False) -> 512-d (authors' recommended tile feature)
+  plip     : vinid/plip CLIP ViT-B/32 image features (projected) -> 512-d      [open stand-in for CONCH]
+  phikonv2 : owkin/phikon-v2 ViT-L/16 DINOv2, CLS -> 1024-d                  [open stand-in for GigaPath]
+"""
+import os
+import sys
+
+import numpy as np
+import timm
+import torch
+from PIL import Image, PngImagePlugin
+
+PngImagePlugin.MAX_TEXT_CHUNK = 64 * 1024 * 1024  # some PathGen PNGs embed very large ICC profiles
+
+from common import CACHE, ORION, dihedral, load_all_tiles, load_imagebind_vision
+
+name = sys.argv[1]
+SET = sys.argv[sys.argv.index("--set") + 1] if "--set" in sys.argv else "pathgen" if "--pathgen" in sys.argv else None
+PATHGEN = SET is not None  # encode an image-text pair set (cache/<SET>_meta.csv, e.g. PathGen-1.6M) instead of ORION tiles
+UNI2_LOCAL = "/diskarray/home/tl688/reg2026/model_alldata/uni2-h.bin"
+dev = "cuda"
+IMNET = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+CLIPN = ([0.48145466, 0.4578275, 0.40821073], [0.26862954, 0.26130258, 0.27577711])
+
+if name == "uni2":
+    # HF repo is gated for this account -> use the local copy of the official UNI2-h checkpoint
+    model = timm.create_model(
+        "vit_giant_patch14_224", pretrained=False, img_size=224, patch_size=14, depth=24, num_heads=24,
+        init_values=1e-5, embed_dim=1536, mlp_ratio=2.66667 * 2, num_classes=0, no_embed_class=True,
+        mlp_layer=timm.layers.SwiGLUPacked, act_layer=torch.nn.SiLU, reg_tokens=8, dynamic_img_size=True)
+    model.load_state_dict(torch.load(UNI2_LOCAL, map_location="cpu"), strict=True)
+    size, norm, fwd = 224, IMNET, lambda m, x: m(x)
+elif name == "gigapath":
+    model = timm.create_model("hf_hub:prov-gigapath/prov-gigapath", pretrained=True)
+    size, norm, fwd = 224, IMNET, lambda m, x: m(x)
+elif name == "conch":
+    from conch.open_clip_custom import create_model_from_pretrained
+    model, _ = create_model_from_pretrained("conch_ViT-B-16", "hf_hub:MahmoodLab/conch")
+    size, norm = 448, CLIPN
+    fwd = lambda m, x: m.encode_image(x, proj_contrast=False, normalize=False)
+elif name == "conchclip":  # CONCH image tower projected into its joint image-text space (zero-shot with CONCH text)
+    from conch.open_clip_custom import create_model_from_pretrained
+    model, _ = create_model_from_pretrained("conch_ViT-B-16", "hf_hub:MahmoodLab/conch")
+    size, norm = 448, CLIPN
+    fwd = lambda m, x: m.encode_image(x, proj_contrast=True, normalize=False)
+elif name == "plip":  # open pathology CLIP (OpenPath tweets), stand-in for gated CONCH
+    from transformers import CLIPModel
+    model = CLIPModel.from_pretrained("vinid/plip")
+    size, norm, fwd = 224, CLIPN, lambda m, x: m.get_image_features(pixel_values=x)
+elif name == "phikonv2":  # open DINOv2 ViT-L pathology SSL model, stand-in for gated GigaPath
+    from transformers import AutoModel
+    model = AutoModel.from_pretrained("owkin/phikon-v2")
+    size, norm, fwd = 224, IMNET, lambda m, x: m(pixel_values=x).last_hidden_state[:, 0]
+elif name == "biomedclip":  # BiomedCLIP ViT-B/16 image tower (PMC-15M), projected -> 512; zero-shot baselines
+    import open_clip
+    model, _ = open_clip.create_model_from_pretrained("hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224")
+    size, norm, fwd = 224, CLIPN, lambda m, x: m.encode_image(x)
+elif name == "imagebind":  # used here for PathGen zero-shot baselines
+    model, MT = load_imagebind_vision()
+    size, norm, fwd = 224, CLIPN, lambda m, x: m({MT.VISION: x})[MT.VISION]
+else:
+    raise ValueError(name)
+model.eval().to(dev)
+mean = torch.tensor(norm[0]).view(1, 3, 1, 1)
+std = torch.tensor(norm[1]).view(1, 3, 1, 1)
+
+if PATHGEN:
+    import pandas as pd
+    paths = pd.read_csv(os.path.join(CACHE, f"{SET}_meta.csv")).image_path.values
+else:
+    paths = np.array([os.path.join(ORION, p) for p in load_all_tiles().image_path])
+out = None
+bs = 32
+with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16):
+    for s in range(0, len(paths), bs):
+        ims = [np.asarray(Image.open(p).convert("RGB").resize((size, size), Image.BICUBIC))
+               for p in paths[s:s + bs]]
+        x = (torch.from_numpy(np.stack(ims)).permute(0, 3, 1, 2).float() / 255 - mean) / std
+        for d in range(8):
+            e = fwd(model, dihedral(x, d).to(dev)).float().cpu().numpy()
+            if out is None:
+                out = np.zeros((len(paths), 8, e.shape[1]), np.float16)
+            out[s:s + bs, d] = e
+        if s % 640 == 0:
+            print(name, s, len(paths), flush=True)
+np.save(os.path.join(CACHE, (f"{SET}_" if PATHGEN else "") + f"he_{name}.npy"), out)
+print("saved", name, out.shape, "finite:", np.isfinite(out).all())

experiments/05_train_eval.py

@@ -0,0 +1,229 @@
+"""ImageBind-style binding of IF / nuclei / markers to a frozen H&E (ImageBind-Huge vision) anchor.
+
+loss = sum_m InfoNCE_sym(he, m)       (--loss anchor, ImageBind default; IF<->nuclei etc. are *emergent*)
+     + sum_{m<n} InfoNCE_sym(m, n)    (--loss allpairs ablation)
+
+Evaluation on the test split: cross-modal retrieval for all 12 directed modality pairs
+(R@1/5/10, MRR, median rank) and one-to-one pairing accuracy (Hungarian assignment on similarity).
+"""
+import argparse
+import itertools
+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 scipy.optimize import linear_sum_assignment
+
+from common import CACHE, MODALITIES, RUNS, SPLITS, dihedral
+
+p = argparse.ArgumentParser()
+p.add_argument("--split", default="orion_official")
+p.add_argument("--fold", type=int, default=0)
+p.add_argument("--loss", default="anchor", choices=["anchor", "allpairs"])
+p.add_argument("--anchor", default="imagebind", choices=["imagebind", "uni2", "conch", "gigapath", "plip", "phikonv2"],
+               help="frozen H&E encoder whose features (cache/he_<anchor>.npy) define the shared space")
+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("--backbone", default="resnet18.a1_in1k")
+p.add_argument("--seed", type=int, default=0)
+p.add_argument("--min_nuclei_tissue", type=int, default=20)
+args = p.parse_args()
+torch.manual_seed(args.seed); np.random.seed(args.seed)
+dev = "cuda"
+run_dir = os.path.join(RUNS, f"{args.split}_fold{args.fold}_{args.loss}"
+                       + ("" if args.anchor == "imagebind" else f"_{args.anchor}"))
+os.makedirs(run_dir, exist_ok=True)
+
+# ---------------- data ----------------
+sp = pd.read_csv(os.path.join(SPLITS, args.split, f"fold{args.fold}.csv"))
+he_emb = torch.from_numpy(np.load(os.path.join(CACHE, f"he_{args.anchor}.npy")).astype(np.float32))  # N,8,D
+D = he_emb.shape[-1]  # shared embedding dim = anchor dim (1024 ImageBind/Phikon-v2, 1536 UNI2/GigaPath, 512 CONCH/PLIP)
+IF = torch.from_numpy(np.load(os.path.join(CACHE, "if.npy")))          # N,17,224,224 uint8 (in RAM, ~2.7GB)
+NUC = torch.from_numpy(np.load(os.path.join(CACHE, "nuclei.npy")))     # N,2,224,224
+MK = torch.from_numpy(np.load(os.path.join(CACHE, "markers.npy")))     # N,33
+tr_idx = sp.tile_id[sp.split == "train"].values
+va_idx = sp.tile_id[sp.split == "val"].values
+te_idx = sp.tile_id[sp.split == "test"].values
+# normalisation statistics from training tiles only
+if_mean = IF[tr_idx].float().mean((0, 2, 3)).view(1, -1, 1, 1) / 255
+if_std = IF[tr_idx].float().std((0, 2, 3)).view(1, -1, 1, 1) / 255 + 1e-3
+mk_mean, mk_std = MK[tr_idx].mean(0), MK[tr_idx].std(0) + 1e-3
+
+
+def batch(idx, d=None, train=False):
+    """Returns dict of modality tensors on GPU. d: dihedral id per sample (shared across image modalities)."""
+    idx = torch.as_tensor(idx)
+    if d is None:
+        d = torch.zeros(len(idx), dtype=torch.long)
+    x_if = IF[idx].to(dev, non_blocking=True).float() / 255
+    x_nu = NUC[idx].to(dev, non_blocking=True).float() / 255
+    x_if = torch.stack([dihedral(a, int(k)) for a, k in zip(x_if, d)])
+    x_nu = torch.stack([dihedral(a, int(k)) for a, k in zip(x_nu, d)])
+    if train:  # mild per-channel intensity jitter for IF
+        x_if = x_if * (1 + 0.1 * torch.randn(len(idx), x_if.shape[1], 1, 1, device=dev))
+    x_if = (x_if - if_mean.to(dev)) / if_std.to(dev)
+    x_nu = (x_nu - 0.5) / 0.5
+    mk = ((MK[idx] - mk_mean) / mk_std).to(dev)
+    if train:
+        mk = mk + 0.05 * torch.randn_like(mk)
+    he = he_emb[idx, d].to(dev)
+    return {"he": he, "if": x_if, "nuclei": x_nu, "markers": mk}
+
+
+# ---------------- model ----------------
+class ImgEnc(nn.Module):
+    def __init__(self, in_chans):
+        super().__init__()
+        self.trunk = timm.create_model(args.backbone, pretrained=True, in_chans=in_chans, num_classes=0)
+        self.head = nn.Sequential(nn.LayerNorm(self.trunk.num_features), nn.Linear(self.trunk.num_features, D))
+
+    def forward(self, x):
+        return self.head(self.trunk(x))
+
+
+class MLPEnc(nn.Module):
+    def __init__(self, d_in):
+        super().__init__()
+        self.net = nn.Sequential(nn.Linear(d_in, 512), nn.GELU(), nn.Dropout(0.1), nn.Linear(512, 512), nn.GELU(),
+                                 nn.LayerNorm(512), nn.Linear(512, D))
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class Bind(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.enc = nn.ModuleDict({"if": ImgEnc(17), "nuclei": ImgEnc(2), "markers": MLPEnc(MK.shape[1])})
+        pairs = [("he", m) for m in ["if", "nuclei", "markers"]]
+        if args.loss == "allpairs":
+            pairs += [("if", "nuclei"), ("if", "markers"), ("nuclei", "markers")]
+        self.pairs = pairs
+        # ImageBind: learnable logit scale (init 1/0.07) per pair
+        self.logit_scale = nn.ParameterDict({f"{a}-{b}": nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32))
+                                             for a, b in pairs})
+
+    def embed(self, b):
+        z = {"he": F.normalize(b["he"], dim=-1)}  # frozen anchor
+        for m, e in self.enc.items():
+            z[m] = F.normalize(e(b[m]), dim=-1)
+        return z
+
+    def loss(self, z):
+        tot, logs = 0, {}
+        tgt = torch.arange(len(z["he"]), device=dev)
+        for a, bm in self.pairs:
+            s = self.logit_scale[f"{a}-{bm}"].exp().clamp(max=100)
+            lg = s * z[a] @ z[bm].T
+            l = (F.cross_entropy(lg, tgt) + F.cross_entropy(lg.T, tgt)) / 2
+            tot = tot + l
+            logs[f"{a}-{bm}"] = l.item()
+        return tot, logs
+
+
+model = Bind().to(dev)
+opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.05)
+steps = args.epochs * max(1, len(tr_idx) // args.bs)
+sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=args.lr, total_steps=steps, pct_start=0.05)
+scaler = torch.amp.GradScaler()
+
+
+@torch.no_grad()
+def embed_all(idx):
+    model.eval()
+    zs = {m: [] for m in MODALITIES}
+    for s in range(0, len(idx), 256):
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            z = model.embed(batch(idx[s:s + 256]))
+        for m in MODALITIES:
+            zs[m].append(z[m].float().cpu())
+    return {m: torch.cat(v) for m, v in zs.items()}
+
+
+def retrieval_metrics(zq, zg):
+    sim = (zq @ zg.T).numpy()
+    n = len(sim)
+    rank = (sim > sim[np.arange(n), np.arange(n)][:, None]).sum(1)  # 0 = correct is top-1
+    r, c = linear_sum_assignment(-sim)
+    return {"R@1": float((rank < 1).mean()), "R@5": float((rank < 5).mean()), "R@10": float((rank < 10).mean()),
+            "MRR": float((1 / (rank + 1)).mean()), "MedR": float(np.median(rank + 1)),
+            "pair_acc": float((c[np.argsort(r)] == np.arange(n)).mean()), "N": n, "chance_R@1": 1 / n}
+
+
+def fixed_gallery(zq, zg, size=100, reps=50):
+    """R@1 / pairing accuracy averaged over random galleries of a fixed size (comparable across splits)."""
+    rng = np.random.default_rng(0)
+    n = len(zq)
+    if n <= size:
+        m = retrieval_metrics(zq, zg)
+        return {"R@1_g100": m["R@1"], "pair_acc_g100": m["pair_acc"]}
+    r1, pa = [], []
+    for _ in range(reps):
+        s = torch.from_numpy(rng.choice(n, size, replace=False))
+        m = retrieval_metrics(zq[s], zg[s])
+        r1.append(m["R@1"]); pa.append(m["pair_acc"])
+    return {"R@1_g100": float(np.mean(r1)), "pair_acc_g100": float(np.mean(pa))}
+
+
+def evaluate(idx, fixed=False):
+    z = embed_all(idx)
+    res = {}
+    for a, b in itertools.permutations(MODALITIES, 2):
+        res[f"{a}->{b}"] = retrieval_metrics(z[a], z[b])
+        if fixed:
+            res[f"{a}->{b}"].update(fixed_gallery(z[a], z[b]))
+    return res, z
+
+
+def mean_r1(res):
+    return float(np.mean([v["R@1"] for v in res.values()]))
+
+
+# ---------------- train ----------------
+best, hist = -1, []
+t0 = time.time()
+for ep in range(args.epochs):
+    model.train()
+    perm = np.random.permutation(tr_idx)
+    tl = []
+    for s in range(0, len(perm) - args.bs + 1, args.bs):
+        idx = perm[s:s + args.bs]
+        b = batch(idx, d=torch.randint(0, 8, (len(idx),)), train=True)
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            loss, logs = model.loss(model.embed(b))
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        opt.step(); sched.step()
+        tl.append(loss.item())
+    vres, _ = evaluate(va_idx)
+    v = mean_r1(vres)
+    hist.append({"epoch": ep, "train_loss": float(np.mean(tl)), "val_meanR1": v, **{k: x for k, x in logs.items()}})
+    print(f"ep {ep} loss {np.mean(tl):.3f} val meanR@1 {v:.3f} ({time.time() - t0:.0f}s)", flush=True)
+    if v > best:
+        best = v
+        torch.save(model.state_dict(), os.path.join(run_dir, "best.pt"))
+
+# ---------------- test ----------------
+model.load_state_dict(torch.load(os.path.join(run_dir, "best.pt")))
+tres, z = evaluate(te_idx, fixed=True)
+stats = pd.read_csv(os.path.join(os.path.dirname(RUNS), "results", "tile_stats.csv")).set_index("tile_id")
+tissue = stats.loc[te_idx, "n_nuclei"].values >= args.min_nuclei_tissue
+tres_tissue = {}
+for a, b in itertools.permutations(MODALITIES, 2):
+    tres_tissue[f"{a}->{b}"] = retrieval_metrics(z[a][tissue], z[b][tissue])
+    tres_tissue[f"{a}->{b}"].update(fixed_gallery(z[a][tissue], z[b][tissue]))
+out = {"args": vars(args), "n_train": len(tr_idx), "n_val": len(va_idx), "n_test": len(te_idx),
+       "n_test_tissue": int(tissue.sum()), "best_val_meanR1": best, "test": tres, "test_tissue": tres_tissue,
+       "history": hist}
+json.dump(out, open(os.path.join(run_dir, "metrics.json"), "w"), indent=1)
+torch.save({m: v for m, v in z.items()}, os.path.join(run_dir, "test_embeddings.pt"))
+print("TEST mean R@1 (all pairs):", mean_r1(tres), "| tissue-only:", mean_r1(tres_tissue))
+for k in ["he->if", "he->nuclei", "he->markers", "if->nuclei", "if->markers", "nuclei->markers"]:
+    print(f"  {k:18s} R@1 {tres[k]['R@1']:.3f} R@5 {tres[k]['R@5']:.3f} pair_acc {tres[k]['pair_acc']:.3f} R@1_g100 {tres[k]['R@1_g100']:.3f}")

experiments/06_summarize.py

@@ -0,0 +1,49 @@
+"""Aggregate runs/*/metrics.json into results/summary_*.csv (mean ± std over folds)."""
+import glob
+import json
+import os
+
+import pandas as pd
+
+from common import ROOT, RUNS
+
+PAIRS = ["he->if", "if->he", "he->nuclei", "nuclei->he", "he->markers", "markers->he",
+         "if->nuclei", "nuclei->if", "if->markers", "markers->if", "nuclei->markers", "markers->nuclei"]
+rows = []
+for f in sorted(glob.glob(os.path.join(RUNS, "*", "metrics.json"))):
+    m = json.load(open(f))
+    a = m["args"]
+    for subset in ["test", "test_tissue"]:
+        for pr in PAIRS:
+            r = m[subset][pr]
+            rows.append({"anchor": a.get("anchor", "imagebind"), "split": a["split"], "fold": a["fold"], "loss": a["loss"], "subset": subset, "pair": pr,
+                         "trained_directly": pr.startswith("he") or pr.endswith("he") or a["loss"] == "allpairs",
+                         **{k: r[k] for k in ["N", "chance_R@1", "R@1", "R@5", "R@10", "MRR", "MedR", "pair_acc",
+                                               "R@1_g100", "pair_acc_g100"]}})
+df = pd.DataFrame(rows)
+out = os.path.join(ROOT, "results")
+df.to_csv(os.path.join(out, "all_runs_long.csv"), index=False)
+
+met = ["R@1", "R@5", "R@10", "MRR", "pair_acc", "R@1_g100", "pair_acc_g100"]
+g = df.groupby(["anchor", "loss", "subset", "split", "pair"], sort=False)
+agg = g[met + ["N"]].mean()
+sd = g[met].std().fillna(0)
+tab = agg.copy()
+for c in met:
+    tab[c] = [f"{x:.3f}±{s:.3f}" if n > 1 else f"{x:.3f}" for x, s, n in zip(agg[c], sd[c], g.size())]
+tab["n_folds"] = g.size()
+tab.reset_index().to_csv(os.path.join(out, "summary_by_pair.csv"), index=False)
+
+# compact headline: mean over the 12 directed pairs, separated into anchor-trained vs emergent pairs
+h = df.groupby(["anchor", "loss", "subset", "split", "fold", "trained_directly"])[met].mean().reset_index()
+h = h.groupby(["anchor", "loss", "subset", "split", "trained_directly"])[met].agg(["mean", "std"]).round(3)
+h.to_csv(os.path.join(out, "summary_headline.csv"))
+pd.set_option("display.width", 250)
+print(h[[("R@1", "mean"), ("R@1", "std"), ("R@1_g100", "mean"), ("pair_acc", "mean"), ("pair_acc_g100", "mean")]])
+
+# H&E-encoder ablation (anchor loss, full test set): fixed-gallery R@1 per pair, mean over folds
+ab = df[(df.loss == "anchor") & (df.subset == "test")]
+fm = ab.pivot_table(index=["split", "pair"], columns="anchor", values="R@1_g100", aggfunc="mean").round(3)
+fm.to_csv(os.path.join(out, "anchor_ablation_R1_g100.csv"))
+fm2 = ab.pivot_table(index=["split", "pair"], columns="anchor", values="pair_acc", aggfunc="mean").round(3)
+fm2.to_csv(os.path.join(out, "anchor_ablation_pair_acc.csv"))

experiments/07_make_text.py

@@ -0,0 +1,69 @@
+"""Templated per-tile text for ORION (no real text is paired with ORION tiles).
+
+Each tile gets 4 paraphrases describing cellularity and cell-type composition derived from the
+marker-positive nuclei counts (fraction of all nuclei) — i.e. the text carries a coarse, discretised
+version of the marker information. Kept short (< 77 CLIP tokens).
+Output: cache/text.json  {tile_id: [4 strings]}
+"""
+import json
+import os
+
+import numpy as np
+import pandas as pd
+
+from common import CACHE, ROOT, load_all_tiles
+
+df = load_all_tiles()
+stats = pd.read_csv(os.path.join(ROOT, "results", "tile_stats.csv")).set_index("tile_id")
+df["n_nuclei"] = stats.loc[df.tile_id, "n_nuclei"].values
+
+CELLS = [  # (count column, phrase)
+    ("Pan-CK_count", "Pan-CK+ epithelial cells"), ("SMA_count", "SMA+ stromal cells"),
+    ("CD45_count", "CD45+ immune cells"), ("CD3e_count", "CD3e+ T cells"), ("CD8a_count", "CD8a+ cytotoxic T cells"),
+    ("CD4_count", "CD4+ helper T cells"), ("FOXP3_count", "FOXP3+ regulatory T cells"), ("CD20_count", "CD20+ B cells"),
+    ("CD68_count", "CD68+ macrophages"), ("CD163_count", "CD163+ macrophages"), ("CD31_count", "CD31+ endothelial cells"),
+    ("Ki67_count", "Ki67+ proliferating cells"), ("PD-L1_count", "PD-L1+ cells"), ("PD-1_count", "PD-1+ cells"),
+]
+LEVEL_WORDS = [  # per paraphrase: (few, moderate, abundant)
+    ("few", "moderate", "abundant"), ("scattered", "some", "many"),
+    ("rare", "a moderate number of", "numerous"), ("sparse", "intermediate", "dense"),
+]
+DENSITY = [("acellular", "sparsely cellular", "moderately cellular", "highly cellular"),
+           ("background", "low-density", "medium-density", "dense"),
+           ("nearly empty", "sparse", "moderately dense", "crowded"),
+           ("mostly non-tissue", "hypocellular", "cellular", "hypercellular")]
+OPEN = ["H&E tile, {d} region with {n} nuclei.", "A {d} tissue patch ({n} nuclei).",
+        "Colorectal cancer tissue, {d}, about {n} nuclei.", "{D} area containing {n} nuclei."]
+
+
+def level(frac):
+    return None if frac < 0.02 else 0 if frac < 0.08 else 1 if frac < 0.25 else 2
+
+
+def describe(r, v, rng):
+    n = int(r.n_nuclei)
+    di = 0 if n < 20 else 1 if n < 100 else 2 if n < 250 else 3
+    d = DENSITY[v][di]
+    n_round = int(round(n, -1))
+    parts = [OPEN[v].format(d=d, D=d.capitalize(), n=n_round)]
+    if n >= 20:
+        items = []
+        for col, name in CELLS:
+            lv = level(r[col] / max(n, 1))
+            if lv is not None:
+                items.append((lv, f"{LEVEL_WORDS[v][lv]} {name}"))
+        items.sort(key=lambda x: -x[0])
+        items = [s for _, s in items[:6]]  # most prominent 6 keep text < 77 tokens
+        if v % 2 == 1:
+            rng.shuffle(items)
+        parts.append(("Shows " + ", ".join(items) + ".") if items else "No marker-positive population stands out.")
+    return " ".join(parts)
+
+
+rng = np.random.default_rng(0)
+texts = {int(r.tile_id): [describe(r, v, rng) for v in range(4)] for _, r in df.iterrows()}
+json.dump(texts, open(os.path.join(CACHE, "text.json"), "w"), indent=0)
+u = len(set(t[0] for t in texts.values()))
+print(f"{len(texts)} tiles, {u} unique variant-0 texts")
+for k in [0, 500, 1500]:
+    print(texts[k])

experiments/08_extract_text.py

@@ -0,0 +1,68 @@
+"""Frozen text-tower embeddings of the templated tile descriptions.
+Usage: python 08_extract_text.py {plip|biomedclip|imagebind|conchclip} [--set <name> | --pathgen | --prompts]
+Output: cache/text_<enc>.npy (N, 4 paraphrases, D) float16.
+  plip       : vinid/plip CLIP text tower (77 tokens), projected -> 512
+  biomedclip : microsoft/BiomedCLIP PubMedBERT text tower (256 tokens), projected -> 512
+  imagebind  : ImageBind-Huge text tower (OpenCLIP ViT-H text, 77 tokens) -> 1024
+  conchclip  : MahmoodLab/CONCH text tower (128 tokens), projected -> 512
+"""
+import json
+import os
+import sys
+
+import numpy as np
+import torch
+
+from common import CACHE, ROOT, load_imagebind_vision
+
+name = sys.argv[1]
+SET = sys.argv[sys.argv.index("--set") + 1] if "--set" in sys.argv else "pathgen" if "--pathgen" in sys.argv else None
+PATHGEN = SET is not None  # encode an image-text pair set (cache/<SET>_meta.csv, e.g. PathGen-1.6M) instead of ORION tiles
+PROMPTS = "--prompts" in sys.argv  # encode concept prompts (prompts.json) for zero-shot probing
+if PROMPTS:
+    pr = json.load(open(os.path.join(ROOT, "prompts.json")))
+    flat = [v["text"] for v in pr.values()]
+    ids, NV = list(range(len(flat))), 1
+elif PATHGEN:
+    import pandas as pd
+    flat = pd.read_csv(os.path.join(CACHE, f"{SET}_meta.csv")).caption.tolist()
+    ids, NV = list(range(len(flat))), 1
+else:
+    texts = json.load(open(os.path.join(CACHE, "text.json")))
+    ids, NV = sorted(int(k) for k in texts), 4
+    flat = [texts[str(i)][v] for i in ids for v in range(NV)]
+
+if name == "plip":
+    from transformers import CLIPModel, CLIPTokenizer
+    m = CLIPModel.from_pretrained("vinid/plip").cuda().eval()
+    tok = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")  # PLIP ships no tokenizer; it uses the CLIP B/32 one
+    enc = lambda b: m.get_text_features(**{k: v.cuda() for k, v in tok(b, padding=True, truncation=True, max_length=77, return_tensors="pt").items()})
+elif name == "biomedclip":
+    import open_clip
+    m, _ = open_clip.create_model_from_pretrained("hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224")
+    tok = open_clip.get_tokenizer("hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224")
+    m = m.cuda().eval()
+    enc = lambda b: m.encode_text(tok(b, context_length=256).cuda())
+elif name == "imagebind":
+    model, MT = load_imagebind_vision()
+    from imagebind.models.multimodal_preprocessors import SimpleTokenizer
+    tk = SimpleTokenizer(bpe_path=os.path.join(ROOT, "ImageBind_src", "imagebind", "bpe", "bpe_simple_vocab_16e6.txt.gz"))
+    enc = lambda b: model({MT.TEXT: torch.cat([tk(t).unsqueeze(0) for t in b]).cuda()})[MT.TEXT]
+elif name == "conchclip":  # CONCH text tower (CoCa, 128 tokens), projected -> 512 (pairs with the conchclip image tower)
+    from conch.open_clip_custom import create_model_from_pretrained, get_tokenizer, tokenize
+    m, _ = create_model_from_pretrained("conch_ViT-B-16", "hf_hub:MahmoodLab/conch")
+    m = m.cuda().eval()
+    tok = get_tokenizer()
+    enc = lambda b: m.encode_text(tokenize(tok, b).cuda(), normalize=False)
+else:
+    raise ValueError(name)
+
+out = []
+with torch.no_grad():
+    for s in range(0, len(flat), 256):
+        out.append(enc(flat[s:s + 256]).float().cpu().numpy())
+out = np.concatenate(out).reshape(len(ids), NV, -1).astype(np.float16)
+assert ids == list(range(len(ids)))
+pre = "prompts_" if PROMPTS else f"{SET}_" if PATHGEN else ""
+np.save(os.path.join(CACHE, pre + f"text_{name}.npy"), out)
+print("saved", name, out.shape, np.isfinite(out).all())

experiments/09_bench.py

@@ -0,0 +1,209 @@
+"""Encoder benchmark for the H&E / mIF / markers / text binding (ImageBind-style, frozen H&E anchor).
+
+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", choices=["resnet18", "convnext_tiny", "chanset"])
+p.add_argument("--mk_enc", default="mlp", choices=["linear", "mlp", "ftt"])
+p.add_argument("--txt_enc", default="plip", choices=["plip", "biomedclip", "imagebind"])
+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)
+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}
+run_dir = os.path.join(RUNS, "bench", f"{args.split}_fold{args.fold}_{args.anchor}_{args.if_enc}_{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 = torch.from_numpy(np.load(os.path.join(CACHE, "if.npy")))
+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_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)
+    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:
+        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": ChanSet() if args.if_enc == "chanset" else TimmEnc(args.if_enc),
+    "markers": {"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 = {}
+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]
+    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"))
+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)

experiments/10_bench_summary.py

@@ -0,0 +1,63 @@
+"""Summarise the encoder benchmark (runs/bench/*/test_embeddings.pt).
+
+- per-modality encoder: he<->m retrieval (trained pair)
+- cross-modal (never trained together): if<->markers, if<->text, markers<->text for every encoder combination
+- zero-shot baselines (no training): PLIP image<->PLIP text, ImageBind image<->ImageBind text
+Metrics are tie-aware (duplicate template texts): expected rank under random tie-breaking.
+Outputs results/bench_*.csv
+"""
+import glob
+import itertools
+import json
+import os
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn.functional as F
+
+from common import CACHE, ROOT, RUNS, metrics
+
+OUT = os.path.join(ROOT, "results")
+
+
+runs = {}
+for f in glob.glob(os.path.join(RUNS, "bench", "*", "test_embeddings.pt")):
+    a = json.load(open(os.path.join(os.path.dirname(f), "train.json")))["args"]
+    runs[(a["split"], a["fold"], a["if_enc"], a["mk_enc"], a["txt_enc"])] = torch.load(f, weights_only=False)  # our own files (contain numpy tile ids)
+
+rows = []
+folds = sorted({k[:2] for k in runs})
+for split, fold in folds:
+    rs = {k[2:]: v for k, v in runs.items() if k[:2] == (split, fold)}
+    he = next(iter(rs.values()))["z"]["he"]
+    te = next(iter(rs.values()))["tile_id"]
+    by = {"if": {}, "markers": {}, "text": {}}
+    for (ie, me, txe), v in rs.items():
+        by["if"][ie], by["markers"][me], by["text"][txe] = v["z"]["if"], v["z"]["markers"], v["z"]["text"]
+    # trained pairs
+    for m, encs in by.items():
+        for e, z in encs.items():
+            for q, g, name in [(he, z, f"he->{m}"), (z, he, f"{m}->he")]:
+                rows.append({"split": split, "fold": fold, "kind": "trained", "pair": name, "enc": e, **metrics(q, g)})
+    # cross pairs: all combinations
+    for m1, m2 in [("if", "markers"), ("if", "text"), ("markers", "text")]:
+        for e1, e2 in itertools.product(by[m1], by[m2]):
+            for q, g, name in [(by[m1][e1], by[m2][e2], f"{m1}->{m2}"), (by[m2][e2], by[m1][e1], f"{m2}->{m1}")]:
+                rows.append({"split": split, "fold": fold, "kind": "cross", "pair": name, "enc": f"{e1}+{e2}", **metrics(q, g)})
+    # zero-shot baselines (variant-0 text, untransformed H&E view)
+    for fm in ["plip", "imagebind"]:
+        hz = F.normalize(torch.from_numpy(np.load(os.path.join(CACHE, f"he_{fm}.npy"))[te, 0].astype(np.float32)), dim=-1)
+        tz = F.normalize(torch.from_numpy(np.load(os.path.join(CACHE, f"text_{fm}.npy"))[te, 0].astype(np.float32)), dim=-1)
+        for q, g, name in [(hz, tz, "he->text"), (tz, hz, "text->he")]:
+            rows.append({"split": split, "fold": fold, "kind": "zeroshot", "pair": name, "enc": f"{fm} (zero-shot)", **metrics(q, g)})
+
+df = pd.DataFrame(rows)
+df.to_csv(os.path.join(OUT, "bench_long.csv"), index=False)
+met = ["R@1", "R@5", "MRR", "pair_acc", "R@1_g100"]
+grp = df.assign(split=df.split.replace({"patient_5fold": "patient_5fold"}))
+agg = grp.groupby(["kind", "pair", "enc", "split"])[met].agg(["mean", "std"]).round(3)
+agg.to_csv(os.path.join(OUT, "bench_summary.csv"))
+pd.set_option("display.width", 250); pd.set_option("display.max_rows", 500)
+show = grp.groupby(["kind", "pair", "enc", "split"])[["R@1_g100", "pair_acc"]].mean().unstack("split").round(3)
+print(show.to_string())

experiments/11_pathgen_prep.py

@@ -0,0 +1,24 @@
+"""PathGen-1.6M H&E-caption pairs (real text) -> cache/pathgen_meta.csv with a slide-level 80/10/10 split.
+Used to bind the text modality to the H&E anchor (ImageBind: text only needs to be paired with images)."""
+import json
+import os
+
+import numpy as np
+import pandas as pd
+
+from common import CACHE, ROOT
+
+PG = os.path.join(ROOT, "..", "PathGen-1.6M")
+d = json.load(open(os.path.join(PG, "Data", "PathGen-1.6M.json")))
+m = pd.DataFrame([{"wsi_id": x["wsi_id"], "caption": x["caption"],
+                   "image_path": os.path.join(PG, "Data", "images", f"{x['wsi_id']}_{x['position'][0]}_{x['position'][1]}.png")}
+                  for x in d])
+assert m.image_path.map(os.path.exists).all()
+m["patient"] = m.wsi_id.str[:12]                         # TCGA-XX-XXXX: keep a patient's slides together
+pats = np.random.default_rng(0).permutation(sorted(m.patient.unique()))
+n = len(pats)
+role = {p: "train" if i < 0.8 * n else "val" if i < 0.9 * n else "test" for i, p in enumerate(pats)}
+m["split"] = m.patient.map(role)
+m["project_code"] = m.wsi_id.str[5:7]
+m.to_csv(os.path.join(CACHE, "pathgen_meta.csv"), index_label="idx")
+print(m.split.value_counts().to_dict(), "patients:", n)

experiments/12_train_text_pathgen.py

@@ -0,0 +1,104 @@
+"""Bind real pathology text to the frozen H&E anchor using PathGen-1.6M H&E-caption pairs (ImageBind recipe:
+text is paired only with images; it never sees ORION, IF or markers).
+
+frozen text tower features (cache/pathgen_text_<enc>.npy) -> trainable MLP head -> anchor space
+loss: symmetric InfoNCE vs frozen anchor image features (cache/pathgen_he_<anchor>.npy, 8 dihedral views)
+Usage: python 12_train_text_pathgen.py --txt_enc plip|biomedclip|imagebind [--anchor uni2]
+Saves runs/text_pathgen/<anchor>_<enc>/{head.pt, metrics.json}
+"""
+import argparse
+import json
+import os
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from common import CACHE, RUNS
+
+p = argparse.ArgumentParser()
+p.add_argument("--txt_enc", default="plip")
+p.add_argument("--anchor", default="uni2")
+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)
+p.add_argument("--seed", type=int, default=0)
+args = p.parse_args()
+torch.manual_seed(args.seed); np.random.seed(args.seed)
+dev = "cuda"
+out_dir = os.path.join(RUNS, "text_pathgen", f"{args.anchor}_{args.txt_enc}")
+os.makedirs(out_dir, exist_ok=True)
+
+meta = pd.read_csv(os.path.join(CACHE, "pathgen_meta.csv"))
+HE = torch.from_numpy(np.load(os.path.join(CACHE, f"pathgen_he_{args.anchor}.npy")).astype(np.float32))
+TX = torch.from_numpy(np.load(os.path.join(CACHE, f"pathgen_text_{args.txt_enc}.npy"))[:, 0].astype(np.float32))
+D = HE.shape[-1]
+idx = {s: meta.idx[meta.split == s].values for s in ["train", "val", "test"]}
+
+
+def make_head(din, dout, hid=512):  # same architecture as the ORION benchmark text head
+    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))
+
+
+head = make_head(TX.shape[1], D).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(idx["train"]) // args.bs)
+sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=args.lr, total_steps=steps, pct_start=0.05)
+
+
+@torch.no_grad()
+def retrieval(ids, size=None, reps=50):
+    head.eval()
+    zi = F.normalize(HE[ids, 0], dim=-1)
+    zt = F.normalize(head(TX[ids].to(dev)).float().cpu(), dim=-1)
+
+    def r1(a, b):
+        s = a @ b.T
+        return float(((s > s.diag()[:, None]).sum(1) == 0).float().mean())
+
+    if size is None or len(ids) <= size:
+        return {"he->text": r1(zi, zt), "text->he": r1(zt, zi), "N": len(ids)}
+    rng = np.random.default_rng(0)
+    a, b = [], []
+    for _ in range(reps):
+        s = torch.from_numpy(rng.choice(len(ids), size, replace=False))
+        a.append(r1(zi[s], zt[s])); b.append(r1(zt[s], zi[s]))
+    return {"he->text": float(np.mean(a)), "text->he": float(np.mean(b)), "N": size}
+
+
+best, hist = -1, []
+for ep in range(args.epochs):
+    head.train()
+    perm = np.random.permutation(idx["train"])
+    for s in range(0, len(perm) - args.bs + 1, args.bs):
+        b = perm[s:s + args.bs]
+        zi = F.normalize(HE[b, torch.randint(0, 8, (len(b),))].to(dev), dim=-1)
+        zt = F.normalize(head(TX[b].to(dev)), dim=-1)
+        lg = logit.exp().clamp(max=100) * zi @ zt.T
+        tgt = torch.arange(len(b), device=dev)
+        loss = (F.cross_entropy(lg, tgt) + F.cross_entropy(lg.T, tgt)) / 2
+        opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sched.step()
+    v = retrieval(idx["val"])
+    vm = (v["he->text"] + v["text->he"]) / 2
+    hist.append({"epoch": ep, "loss": loss.item(), "val_r1": vm})
+    if vm > best:
+        best = vm
+        torch.save(head.state_dict(), os.path.join(out_dir, "head.pt"))
+
+head.load_state_dict(torch.load(os.path.join(out_dir, "head.pt")))
+res = {"args": vars(args), "best_val_r1": best, "test_full": retrieval(idx["test"]),
+       "test_g100": retrieval(idx["test"], size=100), "n": {k: len(v) for k, v in idx.items()}, "history": hist}
+# zero-shot baseline on the same test pairs: the text tower's own image tower, no training
+zs_img = os.path.join(CACHE, f"pathgen_he_{args.txt_enc}.npy")
+if os.path.exists(zs_img):
+    zi = F.normalize(torch.from_numpy(np.load(zs_img)[idx["test"], 0].astype(np.float32)), dim=-1)
+    zt = F.normalize(TX[idx["test"]], dim=-1)
+    s = zi @ zt.T
+    res["zeroshot_test_full"] = {"he->text": float(((s > s.diag()[:, None]).sum(1) == 0).float().mean()),
+                                 "text->he": float(((s.T > s.diag()[:, None]).sum(1) == 0).float().mean())}
+json.dump(res, open(os.path.join(out_dir, "metrics.json"), "w"), indent=1)
+print(args.txt_enc, "val", round(best, 3), "test full", res["test_full"], "g100", res["test_g100"])

experiments/13_eval_text_emergent.py

@@ -0,0 +1,88 @@
+"""Emergent text alignment on ORION. The text head was trained ONLY on PathGen H&E-caption pairs (12_), the IF /
+marker encoders ONLY on ORION H&E pairs (09_, --if_enc convnext_tiny --mk_enc mlp). Text and IF/markers never
+co-occur in training; any text<->IF / text<->markers agreement is emergent through the shared UNI2 H&E space.
+
+(A) concept probes: prompts.json (PathGen-style sentences) scored against every ORION test tile embedding
+    (he / if / markers); ground truth = marker-positive fraction (or nuclei count). Metrics: Spearman rho and
+    AUROC (top quartile vs rest), on tissue tiles (>=20 nuclei; all tiles for 'cellularity').
+    Baselines: zero-shot PLIP and ImageBind (their own image/text towers, H&E only, no training).
+(B) retrieval with templated ORION descriptions (07_) as text queries through the PathGen-trained head
+    (domain-shifted: template style != PathGen captions), tie-aware R@1 on 100-tile galleries.
+Outputs results/text_emergent_probes.csv, results/text_emergent_retrieval.csv
+"""
+import glob
+import json
+import os
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from scipy.stats import spearmanr
+from sklearn.metrics import roc_auc_score
+
+from common import CACHE, ROOT, RUNS, load_all_tiles, metrics
+
+TXT = ["plip", "biomedclip", "imagebind"]
+OUT = os.path.join(ROOT, "results")
+prompts = json.load(open(os.path.join(ROOT, "prompts.json")))
+df = load_all_tiles()
+stats = pd.read_csv(os.path.join(OUT, "tile_stats.csv")).set_index("tile_id")
+df["n_nuclei"] = stats.loc[df.tile_id, "n_nuclei"].values
+
+
+def load_head(enc, din, dout):
+    h = torch.nn.Sequential(torch.nn.Linear(din, 512), torch.nn.GELU(), torch.nn.Dropout(0.1), torch.nn.Linear(512, 512),
+                            torch.nn.GELU(), torch.nn.LayerNorm(512), torch.nn.Linear(512, dout))
+    h.load_state_dict(torch.load(os.path.join(RUNS, "text_pathgen", f"uni2_{enc}", "head.pt"), map_location="cpu"))
+    return h.eval()
+
+
+def probe(scores, tiles, key, label_col):
+    sub = tiles if key == "cellularity" else tiles[df.loc[tiles, "n_nuclei"].values >= 20]
+    pos = np.isin(tiles, sub)
+    s = scores[pos]
+    y = df.loc[sub, label_col].values.astype(float)
+    if label_col != "n_nuclei":
+        y = y / np.maximum(df.loc[sub, "n_nuclei"].values, 1)
+    lab = y >= np.quantile(y, 0.75)
+    return {"spearman": float(spearmanr(s, y).correlation),
+            "auroc": float(roc_auc_score(lab, s)) if 0 < lab.sum() < len(lab) else np.nan, "n": int(len(sub))}
+
+
+probe_rows, ret_rows = [], []
+for f in sorted(glob.glob(os.path.join(RUNS, "bench", "*_uni2_convnext_tiny_mlp_*", "test_embeddings.pt"))):
+    run = os.path.basename(os.path.dirname(f))
+    split, fold = run.split("_fold")[0], int(run.split("_fold")[1].split("_")[0])
+    E = torch.load(f, weights_only=False)
+    tiles, Z = np.asarray(E["tile_id"]), E["z"]
+    for enc in TXT:
+        prm = torch.from_numpy(np.load(os.path.join(CACHE, f"prompts_text_{enc}.npy"))[:, 0].astype(np.float32))
+        orion_txt = torch.from_numpy(np.load(os.path.join(CACHE, f"text_{enc}.npy"))[tiles, 0].astype(np.float32))
+        head = load_head(enc, prm.shape[1], Z["he"].shape[1])
+        with torch.no_grad():
+            zp = F.normalize(head(prm), dim=-1)
+            zt = F.normalize(head(orion_txt), dim=-1)
+        for i, (key, pr) in enumerate(prompts.items()):
+            for m in ["he", "if", "markers"]:
+                sc = (Z[m] @ zp[i]).numpy()
+                probe_rows.append({"split": split, "fold": fold, "text_enc": enc, "concept": key, "modality": m,
+                                   "kind": "bound (PathGen text head)", **probe(sc, tiles, key, pr["label"])})
+        for m in ["he", "if", "markers"]:
+            for q, g, name in [(zt, Z[m], f"text->{m}"), (Z[m], zt, f"{m}->text")]:
+                ret_rows.append({"split": split, "fold": fold, "text_enc": enc, "pair": name, **metrics(q, g)})
+    # zero-shot baselines on H&E with the model's own towers (computed once per fold)
+    for fm in ["plip", "imagebind"]:
+        hz = F.normalize(torch.from_numpy(np.load(os.path.join(CACHE, f"he_{fm}.npy"))[tiles, 0].astype(np.float32)), dim=-1)
+        pz = F.normalize(torch.from_numpy(np.load(os.path.join(CACHE, f"prompts_text_{fm}.npy"))[:, 0].astype(np.float32)), dim=-1)
+        for i, (key, pr) in enumerate(prompts.items()):
+            probe_rows.append({"split": split, "fold": fold, "text_enc": fm, "concept": key, "modality": "he",
+                               "kind": "zero-shot (own image tower)", **probe((hz @ pz[i]).numpy(), tiles, key, pr["label"])})
+
+P = pd.DataFrame(probe_rows); R = pd.DataFrame(ret_rows)
+P.to_csv(os.path.join(OUT, "text_emergent_probes.csv"), index=False)
+R.to_csv(os.path.join(OUT, "text_emergent_retrieval.csv"), index=False)
+pd.set_option("display.width", 250); pd.set_option("display.max_rows", 300)
+print(P.groupby(["kind", "text_enc", "modality", "split"])[["auroc", "spearman"]].mean().round(3).unstack("split").to_string())
+print(P.groupby(["kind", "text_enc", "modality", "concept"])["auroc"].mean().unstack("concept").round(2).to_string())
+print(R.groupby(["text_enc", "pair", "split"])["R@1_g100"].mean().unstack("split").round(3).to_string())

experiments/15_extract_marker_fm.py

@@ -0,0 +1,338 @@
+"""Frozen single-cell foundation-model embeddings of the per-tile marker vector.
+
+    python 15_extract_marker_fm.py geneformer [--value prop|count] [--pool mean|cls]
+    python 15_extract_marker_fm.py scgpt      [--value prop|count]
+
+Each ORION tile is treated as one pseudo-cell: the "expression" of each mapped gene is that
+marker's positive-cell fraction (default, `*_prop`) or positive-cell count (`--value count`) in
+the tile; every other gene is unexpressed.  Both models only see the *within-cell ordering* of
+expressed genes (Geneformer: rank-value encoding; scGPT: per-cell quantile binning), so global
+scale (how many cells / how much signal the tile has) is not representable — only composition.
+
+Marker -> gene (protein panel -> HGNC symbol -> Ensembl via Geneformer gc104M dictionary):
+  CD31->PECAM1, CD45->PTPRC, CD45RO->PTPRC (isoform of the same gene: merged, value=max(CD45,CD45RO)),
+  CD68->CD68, CD4->CD4, FOXP3->FOXP3, CD8a->CD8A, CD20->MS4A1, PD-L1->CD274, CD3e->CD3E,
+  CD163->CD163, E-cadherin->CDH1, PD-1->PDCD1, Ki67->MKI67, Pan-CK->KRT8+KRT18+KRT19 (same value
+  for all three; the pan-cytokeratin antibody cocktail targets these simple-epithelium keratins),
+  SMA->ACTA2.  Hoechst (DNA stain) has no gene -> dropped.
+  PD-1_prop is degenerate in the ORION dataframes (>=0.998 in every tile, Spearman 0.02 with
+  PD-1_count), so with --value prop the PDCD1 value is PD-1_count / n_nuclei (results/tile_stats.csv),
+  clipped to [0, 1].
+
+Outputs  cache/markers_<name>.npy  (3187, 1, D) float16, row order = tile_id
+         cache/markers_<name>_meta.json  (mapping, choices, sanity numbers)
+<name> = geneformer | scgpt, with "-count" / "-cls" suffixes for the non-default options.
+"""
+import argparse
+import glob
+import json
+import os
+import pickle
+import time
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+
+from common import CACHE, ROOT, load_all_tiles, marker_columns, marker_matrix
+
+MARKER2GENES = {
+    "CD31": ["PECAM1"], "CD45": ["PTPRC"], "CD45RO": ["PTPRC"], "CD68": ["CD68"], "CD4": ["CD4"],
+    "FOXP3": ["FOXP3"], "CD8a": ["CD8A"], "CD20": ["MS4A1"], "PD-L1": ["CD274"], "CD3e": ["CD3E"],
+    "CD163": ["CD163"], "E-cadherin": ["CDH1"], "PD-1": ["PDCD1"], "Ki67": ["MKI67"],
+    "Pan-CK": ["KRT8", "KRT18", "KRT19"], "SMA": ["ACTA2"],
+}
+DROPPED = {"Hoechst": "nuclear DNA stain, no gene product"}
+GF_REPO, GF_REV = "ctheodoris/Geneformer", None
+SCGPT_REPO = "wanglab/scGPT-human"   # official bowang-lab whole-human pretrained checkpoint + vocab
+
+
+def gene_expression(df, value):
+    """(N, G) pseudo-cell expression over the mapped genes (unique, ordered)."""
+    genes = list(dict.fromkeys(g for gs in MARKER2GENES.values() for g in gs))
+    X = pd.DataFrame(0.0, index=df.index, columns=genes)
+    notes = {}
+    for m, gs in MARKER2GENES.items():
+        if value == "prop":
+            if m == "PD-1":
+                ts = pd.read_csv(os.path.join(ROOT, "results", "tile_stats.csv"), index_col="tile_id")
+                nn_ = ts.loc[df.tile_id, "n_nuclei"].values.astype(float)
+                v = np.clip(df["PD-1_count"].values / np.maximum(nn_, 1.0), 0, 1)
+                notes["PD-1"] = "PD-1_count / max(n_nuclei,1), clipped to [0,1] (PD-1_prop degenerate ~1.0)"
+            else:
+                v = df[f"{m}_prop"].values.astype(float)
+        else:
+            v = df[f"{m}_count"].values.astype(float)
+        for g in gs:
+            X[g] = np.maximum(X[g].values, v)   # PTPRC = max(CD45, CD45RO)
+    return X, notes
+
+
+# ----------------------------------------------------------------------------- Geneformer
+def hf_file(repo, f):
+    from huggingface_hub import hf_hub_download
+    return hf_hub_download(repo, f)
+
+
+def run_geneformer(X, pool, device, bs=256):
+    from transformers import BertModel
+    tokd = pickle.load(open(hf_file(GF_REPO, "geneformer/token_dictionary_gc104M.pkl"), "rb"))
+    medd = pickle.load(open(hf_file(GF_REPO, "geneformer/gene_median_dictionary_gc104M.pkl"), "rb"))
+    n2id = pickle.load(open(hf_file(GF_REPO, "geneformer/gene_name_id_dict_gc104M.pkl"), "rb"))
+    mdir = os.path.dirname(hf_file(GF_REPO, "Geneformer-V2-104M/model.safetensors"))
+    hf_file(GF_REPO, "Geneformer-V2-104M/config.json")
+    ens = [n2id[g] for g in X.columns]
+    tok = np.array([tokd[e] for e in ens])
+    med = np.array([medd[e] for e in ens])
+    CLS, EOS, PAD = tokd["<cls>"], tokd["<eos>"], tokd["<pad>"]
+
+    model, info = BertModel.from_pretrained(mdir, output_loading_info=True, add_pooling_layer=False)
+    model.eval().to(device)
+    # Geneformer tokenizer: X / n_counts * 1e4 / gene_median, rank descending, drop zeros.
+    # The per-cell n_counts factor does not change the within-cell order, so it is omitted.
+    V = X.values / med[None]
+    seqs = []
+    for row in V:
+        nz = np.nonzero(row > 0)[0]
+        order = nz[np.argsort(-row[nz], kind="stable")]
+        seqs.append([CLS] + tok[order].tolist() + [EOS])
+    L = max(map(len, seqs))
+    embs = []
+    with torch.no_grad():
+        for i in range(0, len(seqs), bs):
+            b = seqs[i:i + bs]
+            ids = torch.full((len(b), L), PAD, dtype=torch.long)
+            for j, s in enumerate(b):
+                ids[j, :len(s)] = torch.tensor(s)
+            ids = ids.to(device)
+            att = (ids != PAD).long()
+            hs = model(input_ids=ids, attention_mask=att, output_hidden_states=True).hidden_states
+            h = hs[-2].float()          # penultimate layer (Geneformer EmbExtractor default emb_layer=-1)
+            if pool == "cls":
+                e = h[:, 0]
+            else:                       # mean over gene tokens (exclude <cls>, <eos>, pads)
+                gm = att.clone().float()
+                gm[:, 0] = 0
+                lens = att.sum(1)
+                gm[torch.arange(len(b)), lens - 1] = 0
+                empty = gm.sum(1) == 0  # no expressed gene -> mean over <cls>,<eos>
+                gm[empty] = att[empty].float()
+                e = (h * gm[..., None]).sum(1) / gm.sum(1, keepdim=True)
+            embs.append(e.cpu())
+    E = torch.cat(embs).numpy()
+    n_empty = int(sum(len(s) == 2 for s in seqs))
+    meta = {"model": f"{GF_REPO}/Geneformer-V2-104M", "hidden": int(E.shape[1]),
+            "missing_keys": info["missing_keys"], "unexpected_keys": info["unexpected_keys"],
+            "ensembl": dict(zip(X.columns, ens)), "token_ids": dict(zip(X.columns, tok.tolist())),
+            "gene_median_gc104M": dict(zip(X.columns, med.round(4).tolist())),
+            "input": "<cls> + genes ranked by value/gene_median (desc, nonzero only) + <eos>",
+            "layer": "penultimate (hidden_states[-2])",
+            "pool": "CLS token" if pool == "cls" else "mean over gene tokens (excl. <cls>/<eos>/pad); "
+                    "tiles with no expressed gene: mean over <cls>,<eos>",
+            "n_tiles_no_expressed_gene": n_empty,
+            "seq_len_mean": float(np.mean([len(s) - 2 for s in seqs]))}
+    return E, meta
+
+
+# ----------------------------------------------------------------------------- scGPT
+class GeneEncoder(nn.Module):
+    def __init__(self, n, d, pad):
+        super().__init__()
+        self.embedding = nn.Embedding(n, d, padding_idx=pad)
+        self.enc_norm = nn.LayerNorm(d)
+
+    def forward(self, x):
+        return self.enc_norm(self.embedding(x))
+
+
+class ContinuousValueEncoder(nn.Module):
+    def __init__(self, d, max_value=512):
+        super().__init__()
+        self.linear1, self.linear2, self.norm = nn.Linear(1, d), nn.Linear(d, d), nn.LayerNorm(d)
+        self.max_value = max_value
+
+    def forward(self, x):
+        x = torch.clamp(x.unsqueeze(-1), max=self.max_value)
+        return self.norm(self.linear2(torch.relu(self.linear1(x))))
+
+
+class ExprDecoder(nn.Module):
+    def __init__(self, d):
+        super().__init__()
+        self.fc = nn.Sequential(nn.Linear(d, d), nn.LeakyReLU(), nn.Linear(d, d), nn.LeakyReLU(), nn.Linear(d, 1))
+
+    def forward(self, x):
+        return self.fc(x).squeeze(-1)
+
+
+class MVCDecoder(nn.Module):
+    def __init__(self, d):
+        super().__init__()
+        self.gene2query, self.W = nn.Linear(d, d), nn.Linear(d, d, bias=False)
+
+
+class ScGPT(nn.Module):
+    """Re-implementation of scgpt.model.TransformerModel for the whole-human checkpoint
+    (continuous value encoder, post-norm encoder, no DSBN/batch tokens); the flash-attn packed
+    Wqkv weights map 1:1 onto nn.MultiheadAttention.in_proj_*."""
+
+    def __init__(self, ntoken, d=512, nhead=8, d_hid=512, nlayers=12, pad=0):
+        super().__init__()
+        self.encoder = GeneEncoder(ntoken, d, pad)
+        self.flag_encoder = nn.Embedding(2, d)
+        self.value_encoder = ContinuousValueEncoder(d)
+        layer = nn.TransformerEncoderLayer(d, nhead, d_hid, dropout=0.0, batch_first=True, norm_first=False)
+        self.transformer_encoder = nn.TransformerEncoder(layer, nlayers, enable_nested_tensor=False)
+        self.decoder = ExprDecoder(d)
+        self.mvc_decoder = MVCDecoder(d)
+
+    def encode(self, genes, values, pad_mask):
+        x = self.encoder(genes) + self.value_encoder(values)
+        return self.transformer_encoder(x, src_key_padding_mask=pad_mask)
+
+
+def scgpt_binning(row, n_bins, rng):
+    """scgpt.preprocess.binning for a row of strictly positive values (random tie-splitting _digitize)."""
+    bins = np.quantile(row, np.linspace(0, 1, n_bins - 1))
+    left, right = np.digitize(row, bins), np.digitize(row, bins, right=True)
+    return np.ceil(rng.random(len(row)) * (right - left) + left).astype(np.int64)
+
+
+def run_scgpt(X, device, bs=256, seed=0):
+    from huggingface_hub import hf_hub_download
+    args = json.load(open(hf_hub_download(SCGPT_REPO, "args.json")))
+    vocab = json.load(open(hf_hub_download(SCGPT_REPO, "vocab.json")))
+    sd = torch.load(hf_hub_download(SCGPT_REPO, "best_model.pt"), map_location="cpu", weights_only=True)
+    sd = {k.replace("self_attn.Wqkv.", "self_attn.in_proj_"): v for k, v in sd.items()}
+    PAD, CLS = vocab[args["pad_token"]], vocab["<cls>"]
+    model = ScGPT(len(vocab), args["embsize"], args["nheads"], args["d_hid"], args["nlayers"], PAD)
+    model.load_state_dict(sd, strict=True)
+    model.eval().to(device)
+    n_bins, pad_value = args["n_bins"], args["pad_value"]
+    gid = np.array([vocab[g] for g in X.columns])
+    rng = np.random.default_rng(seed)
+    seqs = []
+    for row in X.values:   # scgpt.tasks.cell_emb: nonzero genes only, <cls> first with value=pad_value, then binning
+        nz = np.nonzero(row > 0)[0]
+        vals = scgpt_binning(row[nz], n_bins, rng).astype(np.float32) if len(nz) else np.zeros(0, np.float32)
+        seqs.append((np.r_[CLS, gid[nz]], np.r_[np.float32(pad_value), vals]))
+    L = max(len(g) for g, _ in seqs)
+    embs, rec = [], []
+    with torch.no_grad():
+        for i in range(0, len(seqs), bs):
+            b = seqs[i:i + bs]
+            g = torch.full((len(b), L), PAD, dtype=torch.long)
+            v = torch.full((len(b), L), float(pad_value))
+            for j, (gg, vv) in enumerate(b):
+                g[j, :len(gg)], v[j, :len(vv)] = torch.from_numpy(gg), torch.from_numpy(vv)
+            g, v = g.to(device), v.to(device)
+            mask = g.eq(PAD)
+            h = model.encode(g, v, mask)
+            embs.append(h[:, 0].float().cpu())
+            pred = model.decoder(h)        # expression decoder: should track the visible input bins
+            for j, (gg, vv) in enumerate(b):
+                if len(gg) > 3:
+                    rec.append(np.corrcoef(pred[j, 1:len(gg)].cpu().numpy(), vv[1:])[0, 1])
+    E = torch.cat(embs).numpy()
+    E = E / np.linalg.norm(E, axis=1, keepdims=True)   # as scgpt.tasks.embed_data
+    meta = {"model": SCGPT_REPO + " (best_model.pt, args.json, vocab.json; official whole-human pretrained)",
+            "hidden": int(E.shape[1]), "strict_load": True,
+            "note_tdc": "tdc/scGPT weights differ from the official checkpoint and its original_model.bin "
+                        "contains a cls_decoder (fine-tuned) -> not used",
+            "vocab_ids": dict(zip(X.columns, gid.tolist())),
+            "input": "<cls>(value=pad_value=-2) + expressed genes, values = scGPT per-cell quantile binning "
+                     f"(n_bins={n_bins}, random tie-split, seed={seed}); padding masked",
+            "pool": "<cls> token output of last layer, L2-normalised (scgpt.tasks.embed_data)",
+            "n_tiles_no_expressed_gene": int(sum(len(gg) == 1 for gg, _ in seqs)),
+            "decoder_recon_pearson_mean": float(np.nanmean(rec))}
+    return E, meta
+
+
+# ----------------------------------------------------------------------------- sanity
+def ridge_cv_spearman(Z, Y, seed=0):
+    from scipy.stats import spearmanr
+    from sklearn.linear_model import RidgeCV
+    from sklearn.model_selection import KFold, cross_val_predict
+    from sklearn.pipeline import make_pipeline
+    from sklearn.preprocessing import StandardScaler
+    m = make_pipeline(StandardScaler(), RidgeCV(alphas=np.logspace(-2, 5, 15), alpha_per_target=True))
+    P = cross_val_predict(m, Z, Y, cv=KFold(5, shuffle=True, random_state=seed))
+    P, Y = P.reshape(len(Y), -1), Y.reshape(len(Y), -1)
+    return np.array([spearmanr(P[:, k], Y[:, k])[0] for k in range(Y.shape[1])])
+
+
+def sanity(E, df):
+    Z = E.astype(np.float32)
+    out = {"finite": bool(np.isfinite(Z).all()), "n_unique_rows": int(len(np.unique(E, axis=0)))}
+    Zn = Z / np.linalg.norm(Z, axis=1, keepdims=True)
+    S = Zn @ Zn.T
+    iu = np.triu_indices(len(Z), 1)
+    out["mean_pairwise_cos"] = float(S[iu].mean())
+    Zc = Z - Z.mean(0)
+    Zc /= np.linalg.norm(Zc, axis=1, keepdims=True) + 1e-8
+    out["mean_pairwise_cos_centered"] = float((Zc @ Zc.T)[iu].mean())
+    ev = np.linalg.svd(Z - Z.mean(0), compute_uv=False) ** 2
+    out["pcs_for_90pct_var"] = int(np.searchsorted(np.cumsum(ev) / ev.sum(), 0.9) + 1)
+    prop, _ = marker_columns(df)
+    Y = df[prop].values.astype(np.float64)
+    informative = [i for i, c in enumerate(prop) if c != "PD-1_prop"]
+    raw = np.load(os.path.join(CACHE, "markers.npy")) if os.path.exists(os.path.join(CACHE, "markers.npy")) \
+        else marker_matrix(df)
+    for tag, X in [("emb", Z), ("raw_marker_baseline", raw)]:
+        r = ridge_cv_spearman(X, Y)
+        out[f"{tag}_prop_spearman_mean17"] = float(np.nanmean(r))
+        out[f"{tag}_prop_spearman_mean16_noPD1"] = float(np.nanmean(r[informative]))
+        out[f"{tag}_prop_spearman_per_marker"] = {c: round(float(x), 3) for c, x in zip(prop, r)}
+        ts_p = os.path.join(ROOT, "results", "tile_stats.csv")
+        if os.path.exists(ts_p):
+            nn_ = pd.read_csv(ts_p, index_col="tile_id").loc[df.tile_id, "n_nuclei"].values
+            out[f"{tag}_log1p_nnuclei_spearman"] = float(ridge_cv_spearman(X, np.log1p(nn_)[:, None])[0])
+    return out
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("model", choices=["geneformer", "scgpt"])
+    ap.add_argument("--value", choices=["prop", "count"], default="prop")
+    ap.add_argument("--pool", choices=["mean", "cls"], default="mean", help="geneformer only")
+    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    a = ap.parse_args()
+    t0 = time.time()
+    df = load_all_tiles()
+    X, notes = gene_expression(df, a.value)
+    name = a.model + ("-count" if a.value == "count" else "") + ("-cls" if a.model == "geneformer" and a.pool == "cls" else "")
+    if a.model == "geneformer":
+        E, meta = run_geneformer(X, a.pool, a.device)
+    else:
+        E, meta = run_scgpt(X, a.device)
+    t_emb = time.time() - t0
+    if a.device.startswith("cuda"):
+        meta["peak_gpu_mem_GB"] = round(torch.cuda.max_memory_allocated() / 1e9, 3)
+    E16 = E.astype(np.float16)[:, None, :]
+    assert E16.shape[0] == len(df) == 3187
+    np.save(os.path.join(CACHE, f"markers_{name}.npy"), E16)
+    meta.update({"name": name, "shape": list(E16.shape), "dtype": "float16", "row_order": "tile_id",
+                 "value_source": "*_prop (positive fraction)" if a.value == "prop" else "*_count (positive cells)",
+                 "value_notes": notes, "marker_to_genes": MARKER2GENES, "dropped_markers": DROPPED,
+                 "merges": {"PTPRC": "max(CD45, CD45RO) — CD45RO is an isoform of PTPRC",
+                            "KRT8/KRT18/KRT19": "all = Pan-CK"},
+                 "scale_note": "both models encode only within-cell gene order -> tile-level abundance/cellularity "
+                               "is not representable",
+                 "runtime_embed_s": round(t_emb, 1)})
+    s = sanity(E16[:, 0], df)
+    meta["sanity"] = s
+    meta["runtime_total_s"] = round(time.time() - t0, 1)
+    json.dump(meta, open(os.path.join(CACHE, f"markers_{name}_meta.json"), "w"), indent=1)
+    print(f"[{name}] saved {E16.shape} float16  embed {t_emb:.1f}s  total {meta['runtime_total_s']}s")
+    for k in ["n_tiles_no_expressed_gene", "decoder_recon_pearson_mean", "missing_keys", "unexpected_keys", "peak_gpu_mem_GB"]:
+        if k in meta:
+            print(f"  {k}: {meta[k]}")
+    for k, v in s.items():
+        if not k.endswith("per_marker"):
+            print(f"  {k}: {v:.4f}" if isinstance(v, float) else f"  {k}: {v}")
+    print("  per-marker Spearman (emb):", s["emb_prop_spearman_per_marker"])
+
+
+if __name__ == "__main__":
+    main()

experiments/16_extract_pathgenclip.py

@@ -0,0 +1,107 @@
+"""PathGen-CLIP-L (jamessyx/PathGen-CLIP-L, CLIP trained on PathGen-1.6M) text + image towers.
+Usage: python 16_extract_pathgenclip.py {text|image} [--pathgen|--prompts]
+  text            -> cache/text_pathgenclip.npy          (N, 4, 768)  ORION templated texts (cache/text.json)
+  text --pathgen  -> cache/pathgen_text_pathgenclip.npy  (N, 1, 768)  PathGen captions (cache/pathgen_meta.csv)
+  text --prompts  -> cache/prompts_text_pathgenclip.npy  (P, 1, 768)  concept prompts (prompts.json)
+  image           -> cache/he_pathgenclip.npy            (N, 8, 768)  ORION tiles, 8 dihedral views
+  image --pathgen -> cache/pathgen_he_pathgenclip.npy    (N, 8, 768)  PathGen images, 8 dihedral views
+All float16, projected CLIP embeddings (encode_text / encode_image), NOT L2-normalised -- same as 04b / 08.
+
+Architecture: the README's 'ViT-B-16' is wrong for -L. The checkpoint has visual.conv1 (1024,3,14,14), 24 visual
+blocks, visual.positional_embedding (577,1024) = 24x24+1 -> 336px, text width 768 / 12 blocks, proj -> 768.
+That is OpenAI ViT-L/14@336; the authors' HF conversion (jamessyx/pathgenclip-vit-large-patch14-hf) is
+initialised from openai/clip-vit-large-patch14-336 with hidden_act=quick_gelu -> open_clip 'ViT-L-14-336-quickgelu',
+loaded strictly. Preprocess: 336px bicubic, OpenAI-CLIP mean/std (all ORION / PathGen images are square, so the
+open_clip Resize+CenterCrop equals the plain square resize used in 04b). Tokenizer: OpenAI CLIP BPE, 77 tokens.
+"""
+import json
+import os
+import sys
+
+import numpy as np
+import torch
+from PIL import Image, PngImagePlugin
+
+PngImagePlugin.MAX_TEXT_CHUNK = 64 * 1024 * 1024  # some PathGen PNGs embed very large ICC profiles
+
+import open_clip
+from huggingface_hub import hf_hub_download
+
+from common import CACHE, ORION, ROOT, dihedral, load_all_tiles
+
+ARCH = "ViT-L-14-336-quickgelu"
+NAME = "pathgenclip"
+mode = sys.argv[1]
+SET = sys.argv[sys.argv.index("--set") + 1] if "--set" in sys.argv else "pathgen" if "--pathgen" in sys.argv else None
+PATHGEN = SET is not None  # encode an image-text pair set (cache/<SET>_meta.csv, e.g. PathGen-1.6M) instead of ORION tiles
+PROMPTS = "--prompts" in sys.argv
+dev = "cuda"
+CLIPN = ([0.48145466, 0.4578275, 0.40821073], [0.26862954, 0.26130258, 0.27577711])
+
+
+def load_model():
+    ckpt = hf_hub_download("jamessyx/PathGen-CLIP-L", "pathgen-clip-l.pt")
+    model = open_clip.create_model(ARCH, pretrained=None)
+    sd = torch.load(ckpt, map_location="cpu")
+    sd = sd.get("state_dict", sd)
+    sd = {k[len("module."):] if k.startswith("module.") else k: v for k, v in sd.items()}
+    model.load_state_dict(sd, strict=True)
+    return model.eval().to(dev)
+
+
+model = load_model()
+
+if mode == "text":
+    tok = open_clip.get_tokenizer(ARCH)
+    if PROMPTS:
+        pr = json.load(open(os.path.join(ROOT, "prompts.json")))
+        flat = [v["text"] for v in pr.values()]
+        ids, NV = list(range(len(flat))), 1
+    elif PATHGEN:
+        import pandas as pd
+        flat = pd.read_csv(os.path.join(CACHE, f"{SET}_meta.csv")).caption.tolist()
+        ids, NV = list(range(len(flat))), 1
+    else:
+        texts = json.load(open(os.path.join(CACHE, "text.json")))
+        ids, NV = sorted(int(k) for k in texts), 4
+        flat = [texts[str(i)][v] for i in ids for v in range(NV)]
+    out = []
+    with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16):
+        for s in range(0, len(flat), 256):
+            out.append(model.encode_text(tok(flat[s:s + 256]).to(dev)).float().cpu().numpy())
+    out = np.concatenate(out).reshape(len(ids), NV, -1).astype(np.float16)
+    assert ids == list(range(len(ids)))
+    pre = "prompts_" if PROMPTS else f"{SET}_" if PATHGEN else ""
+    fn = os.path.join(CACHE, pre + f"text_{NAME}.npy")
+
+elif mode == "image":
+    size = model.visual.image_size[0] if isinstance(model.visual.image_size, (tuple, list)) else model.visual.image_size
+    assert size == 336, size
+    mean = torch.tensor(CLIPN[0]).view(1, 3, 1, 1)
+    std = torch.tensor(CLIPN[1]).view(1, 3, 1, 1)
+    if PATHGEN:
+        import pandas as pd
+        paths = pd.read_csv(os.path.join(CACHE, f"{SET}_meta.csv")).image_path.values
+    else:
+        paths = np.array([os.path.join(ORION, p) for p in load_all_tiles().image_path])
+    out = None
+    bs = 32
+    with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16):
+        for s in range(0, len(paths), bs):
+            ims = [np.asarray(Image.open(p).convert("RGB").resize((size, size), Image.BICUBIC))
+                   for p in paths[s:s + bs]]
+            x = (torch.from_numpy(np.stack(ims)).permute(0, 3, 1, 2).float() / 255 - mean) / std
+            for d in range(8):
+                e = model.encode_image(dihedral(x, d).to(dev)).float().cpu().numpy()
+                if out is None:
+                    out = np.zeros((len(paths), 8, e.shape[1]), np.float16)
+                out[s:s + bs, d] = e
+            if s % 640 == 0:
+                print(NAME, s, len(paths), flush=True)
+    fn = os.path.join(CACHE, (f"{SET}_" if PATHGEN else "") + f"he_{NAME}.npy")
+else:
+    raise ValueError(mode)
+
+np.save(fn, out)
+print("saved", fn, out.shape, out.dtype, "finite:", np.isfinite(out).all(),
+      "peak GPU MB:", torch.cuda.max_memory_allocated() // 2**20)

experiments/17_bench_fm.py

@@ -0,0 +1,241 @@
+"""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)

experiments/18_eval_fm.py

@@ -0,0 +1,94 @@
+"""Evaluate modality-specific encoders: from-scratch baselines (runs/bench, 09_bench.py) vs frozen modality-specific
+foundation models + MLP head (runs/bench_fm, 17_bench_fm.py), all bound to the frozen UNI2 H&E anchor.
+
+(A) retrieval, tie-aware R@1 on 100-tile galleries + Hungarian pairing acc on the full test set:
+    trained pairs  he<->if, he<->markers                          (per encoder)
+    emergent pairs if<->markers (never trained together)           (every if_enc x mk_enc combination of a fold)
+(B) emergent text: the text head trained ONLY on PathGen H&E-caption pairs (12_train_text_pathgen.py, one per text
+    encoder) scores the concept prompts (prompts.json) against ORION he / if / markers embeddings; Spearman + AUROC
+    (same `probe` as 13_eval_text_emergent.py). Text never co-occurs with IF or markers in training.
+Outputs results/fm_retrieval.csv, results/fm_text_probes.csv
+"""
+import glob
+import itertools
+import json
+import os
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from scipy.stats import spearmanr
+from sklearn.metrics import roc_auc_score
+
+from common import CACHE, ROOT, RUNS, load_all_tiles, metrics
+
+OUT = os.path.join(ROOT, "results")
+prompts = json.load(open(os.path.join(ROOT, "prompts.json")))
+df = load_all_tiles()
+stats = pd.read_csv(os.path.join(OUT, "tile_stats.csv")).set_index("tile_id")
+df["n_nuclei"] = stats.loc[df.tile_id, "n_nuclei"].values
+TXT = sorted(os.path.basename(d).split("_", 1)[1] for d in glob.glob(os.path.join(RUNS, "text_pathgen", "uni2_*"))
+             if os.path.exists(os.path.join(d, "head.pt")) and os.path.exists(os.path.join(CACHE, f"prompts_text_{os.path.basename(d).split('_', 1)[1]}.npy")))
+
+
+def load_head(enc, din, dout):
+    h = torch.nn.Sequential(torch.nn.Linear(din, 512), torch.nn.GELU(), torch.nn.Dropout(0.1), torch.nn.Linear(512, 512),
+                            torch.nn.GELU(), torch.nn.LayerNorm(512), torch.nn.Linear(512, dout))
+    h.load_state_dict(torch.load(os.path.join(RUNS, "text_pathgen", f"uni2_{enc}", "head.pt"), map_location="cpu"))
+    return h.eval()
+
+
+def probe(scores, tiles, key, label_col):
+    sub = tiles if key == "cellularity" else tiles[df.loc[tiles, "n_nuclei"].values >= 20]
+    s = scores[np.isin(tiles, sub)]
+    y = df.loc[sub, label_col].values.astype(float)
+    if label_col != "n_nuclei":
+        y = y / np.maximum(df.loc[sub, "n_nuclei"].values, 1)
+    lab = y >= np.quantile(y, 0.75)
+    return {"spearman": float(spearmanr(s, y).correlation),
+            "auroc": float(roc_auc_score(lab, s)) if 0 < lab.sum() < len(lab) else np.nan, "n": int(len(sub))}
+
+
+# collect test embeddings: (split, fold) -> {"if": {enc: z}, "markers": {enc: z}, "he": z, "tile_id": ids}
+folds = {}
+for root, src in [("bench", "scratch"), ("bench_fm", None)]:
+    for f in glob.glob(os.path.join(RUNS, root, "*", "test_embeddings.pt")):
+        a = json.load(open(os.path.join(os.path.dirname(f), "train.json")))["args"]
+        if a["anchor"] != "uni2":
+            continue
+        E = torch.load(f, weights_only=False)  # our own files (contain numpy tile ids)
+        d = folds.setdefault((a["split"], a["fold"]), {"if": {}, "markers": {}, "he": E["z"]["he"], "tile_id": np.asarray(E["tile_id"])})
+        assert (d["tile_id"] == np.asarray(E["tile_id"])).all()
+        d["if"][a["if_enc"]] = E["z"]["if"]
+        d["markers"][a["mk_enc"]] = E["z"]["markers"]
+
+ret, prb = [], []
+for (split, fold), d in sorted(folds.items()):
+    he, tiles = d["he"], d["tile_id"]
+    for m in ["if", "markers"]:
+        for e, z in d[m].items():
+            for q, g, name in [(he, z, f"he->{m}"), (z, he, f"{m}->he")]:
+                ret.append({"split": split, "fold": fold, "kind": "trained", "pair": name, "enc": e, **metrics(q, g)})
+    for ei, em in itertools.product(d["if"], d["markers"]):
+        zi, zm = d["if"][ei], d["markers"][em]
+        for q, g, name in [(zi, zm, "if->markers"), (zm, zi, "markers->if")]:
+            ret.append({"split": split, "fold": fold, "kind": "emergent", "pair": name, "enc": f"{ei} + {em}", **metrics(q, g)})
+    for t in TXT:
+        prm = torch.from_numpy(np.load(os.path.join(CACHE, f"prompts_text_{t}.npy"))[:, 0].astype(np.float32))
+        with torch.no_grad():
+            zp = F.normalize(load_head(t, prm.shape[1], he.shape[1])(prm), dim=-1)
+        embs = [("he", "uni2", he)] + [(m, e, z) for m in ["if", "markers"] for e, z in d[m].items()]
+        for m, e, z in embs:
+            for i, (key, pr) in enumerate(prompts.items()):
+                prb.append({"split": split, "fold": fold, "text_enc": t, "modality": m, "enc": e, "concept": key,
+                            **probe((z @ zp[i]).numpy(), tiles, key, pr["label"])})
+
+R, P = pd.DataFrame(ret), pd.DataFrame(prb)
+R.to_csv(os.path.join(OUT, "fm_retrieval.csv"), index=False)
+P.to_csv(os.path.join(OUT, "fm_text_probes.csv"), index=False)
+pd.set_option("display.width", 250); pd.set_option("display.max_rows", 500)
+R["pair2"] = R.pair.map(lambda p: "<->".join(sorted(p.split("->"))))
+g = R.groupby(["kind", "pair2", "enc", "split", "fold"])[["R@1_g100", "pair_acc"]].mean().reset_index()
+print(g.groupby(["kind", "pair2", "enc", "split"])[["R@1_g100", "pair_acc"]].agg(["mean", "std"]).round(3).to_string())
+print(P.groupby(["text_enc", "modality", "enc", "split"])[["auroc", "spearman"]].mean().round(3).unstack("split").to_string())

experiments/19_pairsets_prep.py

@@ -0,0 +1,162 @@
+"""Unified pool of ALL H&E image-text pairs in PathGen-1.6M/ (every text-image source of PAIR_SUMMARY.md; all TCGA
+patches). Images are de-duplicated by patch key <wsi>_<x>_<y> (the same patch recurs across sources); one image can
+carry several texts. Text types:
+  caption  : PathGen-1.6M captions, figure_caption_summary / summary_testing / subsample_summary, slide_report desc,
+             sample_data/output txt, output_subsample txt, path_reason_data caption
+  summary  : path_summary_train(.json, _5w.json) one-paragraph trainee summaries
+  qa_open  : open-ended Q/A -> assistant answers joined (PathGen-Instruct OPEN, output_subsample json,
+             qa_pathinstruct_withreason*.csv answers, path_reason_data question+answer)
+  qa_close : multiple-choice Q/A -> question stem + text of the correct option (PathGen-Instruct CLOSE,
+             qa_instruct_merged_onlyqaforrl, output_reason json, multiple-choice csv rows)
+Outputs
+  cache/allimg_meta.csv : unique images  (idx, key, wsi_id, patient, image_path)            -> image features
+  cache/alltxt_meta.csv : unique (image, text) pairs (idx, img_idx, caption, ttype, source)   -> text features
+(the `caption` column holds the text so the extractors' --set option works unchanged)
+Not used: pathor1_output*.csv (hashed non-TCGA image names, images absent).
+"""
+import ast
+import glob
+import json
+import os
+import re
+
+import pandas as pd
+
+from common import CACHE, ROOT
+
+PG = os.path.join(ROOT, "..", "PathGen-1.6M")
+img, txt = {}, []          # key -> image path ; list of (key, text, ttype, source)
+KEY = re.compile(r"(TCGA-[A-Z0-9-]+\.[A-Za-z0-9-]+)_(\d+)_(\d+)")
+
+
+def key_of(s):
+    m = KEY.search(os.path.basename(str(s)))
+    return f"{m.group(1)}_{m.group(2)}_{m.group(3)}" if m else None
+
+
+def add_img(key, path):
+    if key and key not in img and os.path.exists(path):
+        img[key] = path
+
+
+def add_txt(key, text, ttype, source):
+    text = re.sub(r"\s+", " ", str(text).replace("<image>", " ")).strip()
+    if key and text and text.lower() != "nan":
+        txt.append((key, text, ttype, source))
+
+
+def close_text(question, answer):
+    """multiple-choice -> 'question stem + correct option text'"""
+    q = question.replace("<image>", " ")
+    opts = dict(re.findall(r"(?:^|\n)\s*\(?([A-F])[\).:]\s*([^\n]+)", q))
+    stem = re.split(r"\n\s*\(?A[\).:]", q)[0].strip()
+    a = answer.strip()
+    letter = re.match(r"^\(?([A-F])\)?(?:[\).:\s]|$)", a)
+    if letter and letter.group(1) in opts:
+        a = opts[letter.group(1)]
+    elif letter and len(a) <= 3:
+        return None
+    else:
+        a = re.sub(r"^\(?[A-F]\)\s*", "", a)
+    return f"{stem} {a.strip()}"
+
+
+def conv_text(conv):
+    """returns (ttype, text) for a LLaVA-style conversation list"""
+    gpt = [c["value"].strip() for c in conv if c["from"] == "gpt"]
+    if len(gpt) == 1 and re.fullmatch(r"\(?[A-F]\)?\.?", gpt[0]):
+        human = [c["value"] for c in conv if c["from"] == "human"][0]
+        t = close_text(human, gpt[0])
+        return ("qa_close", t) if t else (None, None)
+    return "qa_open", " ".join(gpt)
+
+
+# 1. PathGen-1.6M captions (Data/images, 672 px)
+for x in json.load(open(os.path.join(PG, "Data", "PathGen-1.6M.json"))):
+    k = f"{x['wsi_id']}_{x['position'][0]}_{x['position'][1]}"
+    add_img(k, os.path.join(PG, "Data", "images", k + ".png"))
+    add_txt(k, x["caption"], "caption", "pathgen1.6m")
+# 2. PathGen-Instruct (OPEN + CLOSE)
+idx = pd.read_csv(os.path.join(PG, "PathGen-Instruct.json.image_index.csv")).drop_duplicates("key").set_index("key")
+for x in json.load(open(os.path.join(PG, "PathGen-Instruct.json"))):
+    k = f"{x['wsi_id']}_{x['position'][0]}_{x['position'][1]}"
+    if k in idx.index:
+        add_img(k, os.path.join(PG, idx.loc[k, "local_image"]))
+    tt, t = conv_text(x["conversations"])
+    if tt:
+        add_txt(k, t, tt, "pathgen-instruct")
+# 3. qa_instruct_merged_onlyqaforrl
+for x in json.load(open(os.path.join(PG, "qa_instruct_merged_onlyqaforrl.json"))):
+    k = x.get("path_list") or f"{x['wsi_id']}_{x['position'][0]}_{x['position'][1]}"
+    for c in glob.glob(os.path.join(PG, "qa_instruct_merged_images", k + ".*")):
+        if c.endswith((".png", ".jpg")):
+            add_img(k, c)
+    tt, t = conv_text(x["conversations"])
+    if tt:
+        add_txt(k, t, tt, "qa_instruct_merged")
+# 4. summaries
+for f in ["path_summary_train.json", "path_summary_train_5w.json"]:
+    for x in json.load(open(os.path.join(PG, f))):
+        k = key_of(x["images"][0])
+        add_img(k, os.path.join(PG, x["images"][0]))
+        add_txt(k, [m["content"] for m in x["messages"] if m["role"] == "assistant"][0], "summary", f.replace(".json", ""))
+# 5. output_reason (png + json conversation), output_subsample (png + txt caption / json conversation), sample_data
+for d, src in [("output_reason", "output_reason"), ("output_subsample", "output_subsample"), ("sample_data/output", "sample_data")]:
+    for pth in glob.glob(os.path.join(PG, d, "*.png")):
+        k = key_of(pth)
+        add_img(k, pth)
+        base = pth[:-4]
+        if os.path.exists(base + ".txt"):
+            add_txt(k, open(base + ".txt").read(), "caption", src)
+        if os.path.exists(base + ".json"):
+            tt, t = conv_text(json.load(open(base + ".json")))
+            if tt:
+                add_txt(k, t, tt, src)
+# 6. qa_pathinstruct_withreason*.csv (images in output_reason/)
+for f in sorted(glob.glob(os.path.join(PG, "qa_pathinstruct_withreason*.csv"))):
+    for _, r in pd.read_csv(f).iterrows():
+        k = key_of(r.get("path_list", r.get("imagep")))
+        add_img(k, os.path.join(PG, str(r["imagep"]).lstrip("./")))
+        t = close_text(r["question"], str(r["answer"])) if re.search(r"\n\s*\(?A[\).:]", str(r["question"])) else r["answer"]
+        if t:
+            add_txt(k, t, "qa_close" if t is not r["answer"] else "qa_open", "qa_pathinstruct_withreason")
+# 7. path_reason_data (caption + diagnostic Q/A)
+for f in sorted(glob.glob(os.path.join(PG, "path_reason_data_3600*.csv"))):
+    for _, r in pd.read_csv(f).iterrows():
+        k = key_of(r["patch_id"])
+        for c in [os.path.join(PG, str(r["patch_id"]).lstrip("./")), os.path.join(PG, "output_subsample", os.path.basename(str(r["patch_id"])))]:
+            add_img(k, c)
+        add_txt(k, r["caption"], "caption", "path_reason_data")
+        add_txt(k, f"{r['question']} {r['answer']}", "qa_open", "path_reason_data")
+# 8. figure_caption_summary (summary_testing_data / subsample_summary_sentence), slide_report
+for f in ["summary_testing_data.csv", "subsample_summary_sentence.csv"]:
+    for _, r in pd.read_csv(os.path.join(PG, f)).iterrows():
+        x, y = ast.literal_eval(r["position"])
+        k = f"{r['wsi_id']}_{x}_{y}"
+        add_img(k, os.path.join(PG, "figure_caption_summary", k + ".png"))
+        add_txt(k, r["caption"], "caption", "figure_caption_summary")
+for f in glob.glob(os.path.join(PG, "slide_report", "*", "*_desc.csv")):
+    k = key_of(f.replace("_desc.csv", ""))
+    add_img(k, f.replace("_desc.csv", ".png"))
+    for _, r in pd.read_csv(f).iterrows():
+        add_txt(k, r["caption"], "caption", "slide_report")
+
+T = pd.DataFrame(txt, columns=["key", "caption", "ttype", "source"])
+n_raw = len(T)
+T = T[T.key.isin(img)]
+n_noimg = n_raw - len(T)
+T = T.groupby(["key", "caption"], as_index=False).agg(ttype=("ttype", "first"), source=("source", lambda s: "|".join(sorted(set(s)))))
+I = pd.DataFrame({"key": sorted(T.key.unique())})
+I["image_path"] = I.key.map(img)
+I["wsi_id"] = I.key.str.rsplit("_", n=2).str[0]
+I["patient"] = I.wsi_id.str[:12]
+I["caption"] = ""
+T["img_idx"] = T.key.map(dict(zip(I.key, I.index)))
+I.to_csv(os.path.join(CACHE, "allimg_meta.csv"), index_label="idx")
+T[["img_idx", "key", "caption", "ttype", "source"]].to_csv(os.path.join(CACHE, "alltxt_meta.csv"), index_label="idx")
+print(f"texts: {n_raw} raw, {n_noimg} without an image on disk, {len(T)} unique (image, text) pairs")
+print(f"images: {len(I)} unique patches, {I.wsi_id.nunique()} slides, {I.patient.nunique()} patients")
+print(T.ttype.value_counts().to_string())
+print(T.source.str.split("|").explode().value_counts().to_string())
+print("texts per image:", T.groupby("img_idx").size().describe().round(2).to_dict())
+print("words per text by type:", T.assign(w=T.caption.str.split().str.len()).groupby("ttype").w.median().to_dict())

experiments/20_pairing.py

@@ -0,0 +1,153 @@
+"""Text <-> H&E image pairing benchmark on ALL real image-text pairs (unified pool from 19_pairsets_prep.py:
+6144 TCGA patches, 6480 pairs; text types caption / summary / qa_open / qa_close).
+
+Protocol: patient-level 5-fold CV over the pool (inside each training fold 10% of patients are held out for early
+stopping). Training uses every (image, text) pair of the training patients; an image with several texts counts all of
+them as positives (multi-positive InfoNCE). Every image encoder x text encoder combination in three modes:
+  zeroshot  : the CLIP model's own image and text towers, no training (only when image enc == text enc)
+  text_head : frozen image features = anchor; trainable MLP head on the frozen text features (ImageBind recipe)
+  two_heads : trainable MLP heads on both frozen towers into a shared 512-d space (CLIP-style projection)
+Evaluation on the held-out fold, per text type and for "all" (one text per image): tie-aware R@1/5/10, MRR,
+Hungarian one-to-one pairing accuracy over the whole gallery, R@1 on random 100-pair galleries; both directions.
+Caveat: PathGen-CLIP-L was trained on PathGen-1.6M (most of this pool) -> its numbers are not held-out.
+Outputs results/pairing_long.csv, results/pairing_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 sklearn.model_selection import GroupKFold
+
+from common import CACHE, ROOT, metrics
+
+p = argparse.ArgumentParser()
+p.add_argument("--img", nargs="+", default=["uni2", "phikonv2", "plip", "biomedclip", "imagebind", "pathgenclip"])
+p.add_argument("--txt", nargs="+", default=["plip", "biomedclip", "imagebind", "pathgenclip"])
+p.add_argument("--epochs", type=int, default=60)
+p.add_argument("--bs", type=int, default=256)
+p.add_argument("--lr", type=float, default=3e-4)
+p.add_argument("--tag", default="")
+args = p.parse_args()
+dev = "cuda"
+OUT = os.path.join(ROOT, "results")
+I = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv"))
+T = pd.read_csv(os.path.join(CACHE, "alltxt_meta.csv"))
+TTYPES = ["caption", "summary", "qa_open", "qa_close"]
+
+
+def load(kind, enc):
+    f = os.path.join(CACHE, f"all{kind}_{enc}.npy")
+    return torch.from_numpy(np.load(f).astype(np.float32)).to(dev) if os.path.exists(f) else None
+
+
+IMG = {e: load("img_he", e) for e in args.img}                       # (n_img, 8, D)
+TXT = {e: (lambda a: a[:, 0] if a is not None else None)(load("txt_text", e)) for e in args.txt}  # (n_txt, D)
+t_img = torch.as_tensor(T.img_idx.values, device=dev)
+
+
+def mlp(din, dout, hid=512):  # same head as 12_train_text_pathgen.py / 09_bench.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 eval_pairs(t_ids, rng):
+    """one text per image (random) -> aligned (img_ids, txt_ids) without duplicate images"""
+    sub = T.loc[t_ids]
+    pick = sub.groupby("img_idx").apply(lambda g: g.index[rng.integers(len(g))], include_groups=False).values
+    return T.img_idx.values[pick], pick
+
+
+def train(ie, te, mode, tr_t, va_t, seed=0):
+    torch.manual_seed(seed); rng = np.random.default_rng(seed)
+    X, Tx = IMG[ie], TXT[te]
+    th = mlp(Tx.shape[1], 512 if mode == "two_heads" else X.shape[-1]).to(dev)
+    ih = mlp(X.shape[-1], 512).to(dev) if mode == "two_heads" else nn.Identity()
+    logit = nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32, device=dev))
+    opt = torch.optim.AdamW(list(th.parameters()) + list(ih.parameters()) + [logit], lr=args.lr, weight_decay=0.05)
+    steps = args.epochs * max(1, len(tr_t) // args.bs)
+    sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=args.lr, total_steps=steps, pct_start=0.05)
+    va_i, va_tt = eval_pairs(va_t, np.random.default_rng(0))
+
+    def enc_img(ids, view=None):
+        return F.normalize(ih(X[ids, 0 if view is None else view]), dim=-1)
+
+    def enc_txt(ids):
+        return F.normalize(th(Tx[ids]), dim=-1)
+
+    best, state = -1, None
+    for ep in range(args.epochs):
+        th.train(); ih.train()
+        perm = rng.permutation(tr_t)
+        for s in range(0, len(perm) - args.bs + 1, args.bs):
+            b = torch.as_tensor(perm[s:s + args.bs], device=dev)
+            ii = t_img[b]
+            zi = enc_img(ii, torch.randint(0, 8, (len(b),), device=dev))
+            zt = enc_txt(b)
+            lg = logit.exp().clamp(max=100) * zi @ zt.T
+            pos = (ii[:, None] == ii[None, :]).float()           # same image -> positive (multi-positive InfoNCE)
+            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(); ih.eval()
+        with torch.no_grad():
+            zi, zt = enc_img(torch.as_tensor(va_i, device=dev)), enc_txt(torch.as_tensor(va_tt, device=dev))
+            v = (r1(zi, zt) + r1(zt, zi)) / 2
+        if v > best:
+            best = v
+            state = ({k: x.clone() for k, x in th.state_dict().items()}, {k: x.clone() for k, x in ih.state_dict().items()})
+    th.load_state_dict(state[0]); ih.load_state_dict(state[1]); th.eval(); ih.eval()
+    return (lambda i_ids, t_ids: (enc_img(torch.as_tensor(i_ids, device=dev)).cpu(), enc_txt(torch.as_tensor(t_ids, device=dev)).cpu())), best
+
+
+folds = list(GroupKFold(5).split(I, groups=I.patient))
+rows = []
+for k, (trv_img, te_img) in enumerate(folds):
+    pats = np.random.default_rng(k).permutation(sorted(I.patient.iloc[trv_img].unique()))
+    va_p = set(pats[: max(1, len(pats) // 10)])
+    img_split = pd.Series("test", index=I.index)
+    img_split.iloc[trv_img] = np.where(I.patient.iloc[trv_img].isin(va_p), "val", "train")
+    tsplit = img_split.loc[T.img_idx].values
+    tr_t, va_t, te_t = (T.index[tsplit == s].values for s in ["train", "val", "test"])
+    # evaluation sets: per text type and 'all', one text per image
+    ev = {"all": eval_pairs(te_t, np.random.default_rng(k))}
+    for tt in TTYPES:
+        ids = te_t[T.ttype.values[te_t] == tt]
+        if len(ids) >= 20:
+            ev[tt] = eval_pairs(ids, np.random.default_rng(k))
+    for ie in args.img:
+        for te in args.txt:
+            if IMG[ie] is None or TXT[te] is None:
+                continue
+            for mode in (["zeroshot"] if ie == te else []) + ["text_head", "two_heads"]:
+                if mode == "zeroshot":
+                    f = lambda i_ids, t_ids: (F.normalize(IMG[ie][torch.as_tensor(i_ids, device=dev), 0], dim=-1).cpu(),
+                                              F.normalize(TXT[te][torch.as_tensor(t_ids, device=dev)], dim=-1).cpu())
+                    val = np.nan
+                else:
+                    f, val = train(ie, te, mode, tr_t, va_t)
+                for name, (i_ids, t_ids) in ev.items():
+                    with torch.no_grad():
+                        zi, zt = f(i_ids, t_ids)
+                    for q, g, d in [(zi, zt, "img->text"), (zt, zi, "text->img")]:
+                        rows.append({"eval": name, "fold": k, "img_enc": ie, "txt_enc": te, "mode": mode, "dir": d,
+                                     "val_r1": val, **metrics(q, g)})
+            print(f"fold {k} {ie:11s} {te:11s} done (n_train_pairs={len(tr_t)}, n_test_img={len(ev['all'][0])})", flush=True)
+    pd.DataFrame(rows).to_csv(os.path.join(OUT, f"pairing_long{args.tag}.csv"), index=False)
+
+R = pd.DataFrame(rows)
+cols = ["R@1", "R@5", "R@10", "pair_acc", "R@1_g100"]
+S = R.groupby(["eval", "mode", "img_enc", "txt_enc", "fold"])[cols + ["N"]].mean().groupby(["eval", "mode", "img_enc", "txt_enc"]).agg(["mean", "std"])
+S.round(4).to_csv(os.path.join(OUT, f"pairing_summary{args.tag}.csv"))
+pd.set_option("display.width", 250); pd.set_option("display.max_rows", 500)
+for e in ["all"] + TTYPES:
+    print(f"\n=== eval: {e} (mean over both directions and 5 folds) ===")
+    print(R[R["eval"] == e].groupby(["mode", "img_enc", "txt_enc"])[cols + ["N"]].mean().round(3).sort_values("R@1_g100", ascending=False).to_string())

experiments/21_hest_prep.py

@@ -0,0 +1,124 @@
+"""Gene-expression modality: HEST-1k human Visium spots (H&E patch + whole-transcriptome spot expression).
+
+Input  ../HEST/visium_humandisease_adata.pickle            list of 421 AnnData (1% spot subsample per sample)
+       ../HEST/UNI/visium_hest_univ2_imagefeature.pickle    list of 421 (n_spot, 1536) UNI2-h spot features (aligned)
+       ../HEST/prov-gigapath-conchinside/visium_hest_gigapath_imagefeature.pickle   (GigaPath, aligned)
+       ../HEST/HEST_v1_1_0.csv                              sample metadata; list order = Visium/Homo sapiens rows
+Kept   samples whose panel covers the common whole-transcriptome genes (targeted ~1k-gene panels dropped).
+Groups for CV: study + patient (patient missing -> whole study), i.e. no patient leaks across folds.
+Outputs cache/hest/
+  meta.csv            one row per spot (sample id, organ, oncotree, group, total counts)
+  counts.npz          raw counts (n_spot, n_common_genes) sparse + gene names
+  hvg.npy             log1p(CP10k) of the top-2000 HVGs (n_spot, 2000) float32      -> MLP gene encoder input
+  he_uni2.npy, he_gigapath.npy   (n_spot, 1, 1536) float16                            -> frozen H&E anchor
+"""
+import os
+import pickle
+import weakref
+
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+import torch
+from anndata._core import file_backing
+
+from common import CACHE, ROOT
+
+HEST = os.path.join(ROOT, "..", "HEST")
+OUT = os.path.join(CACHE, "hest")
+os.makedirs(OUT, exist_ok=True)
+
+
+def _ss(self, state):  # old anndata pickles store '_adata' instead of '_adata_ref'
+    self.__dict__ = state.copy()
+    ref = state.get("_adata_ref", state.get("_adata"))
+    self.__dict__["_adata_ref"] = weakref.ref(ref) if ref is not None else (lambda: None)
+
+
+file_backing.AnnDataFileManager.__setstate__ = _ss
+A = pickle.load(open(os.path.join(HEST, "visium_humandisease_adata.pickle"), "rb"))
+FU = pickle.load(open(os.path.join(HEST, "UNI", "visium_hest_univ2_imagefeature.pickle"), "rb"))
+FG = pickle.load(open(os.path.join(HEST, "prov-gigapath-conchinside", "visium_hest_gigapath_imagefeature.pickle"), "rb"))
+H = pd.read_csv(os.path.join(HEST, "HEST_v1_1_0.csv"))
+H = H[(H.st_technology == "Visium") & (H.species == "Homo sapiens")].reset_index(drop=True)
+assert len(H) == len(A) == len(FU) == len(FG)
+
+# --- check the list order against the metadata (full-res size from thumbnail / scalefactor, oncotree code)
+ok = []
+for a, (_, h) in zip(A, H.iterrows()):
+    st = a.uns["spatial"]["ST"]
+    sh = np.array(st["images"]["downscaled_fullres"].shape[:2]) / st["scalefactors"]["tissue_downscaled_fullres_scalef"]
+    size_ok = np.allclose(sh, [h.fullres_px_height, h.fullres_px_width], rtol=0.02)
+    onc = a.obs.oncotree_code.iloc[0]
+    onc_ok = (pd.isna(onc) and pd.isna(h.oncotree_code)) or onc == h.oncotree_code
+    ok.append(size_ok and onc_ok)
+print(f"metadata order check: {sum(ok)}/{len(ok)} samples match (image size + oncotree)")
+assert np.mean(ok) > 0.95
+
+# --- harmonise gene names to HGNC symbols (samples use symbols, Ensembl ids + SYMBOL column, or 'GRCh38__' prefixes)
+e2s = {}
+for a in A:
+    if "gene_ids" in a.var and not str(a.var_names[0]).startswith(("ENSG", "GRCh")):
+        e2s.update(dict(zip(a.var.gene_ids.astype(str), a.var_names)))
+
+
+def symbols(a):
+    if "SYMBOL" in a.var:
+        return a.var.SYMBOL.astype(str).tolist()
+    out = []
+    for x in map(str, a.var_names):
+        x = x.split("__")[-1] if x.startswith("GRCh") else x
+        out.append(e2s.get(x.split(".")[0], x) if x.startswith("ENSG") else x)
+    return out
+
+
+for a in A:
+    a.var_names = symbols(a)
+
+# --- common gene set over whole-transcriptome samples
+wt = [i for i, a in enumerate(A) if a.n_vars >= 17000]
+common = set(A[wt[0]].var_names)
+for i in wt[1:]:
+    common &= set(A[i].var_names)
+common = sorted(g for g in common if not g.startswith(("MT-", "RPL", "RPS")))
+print(f"whole-transcriptome samples: {len(wt)}/{len(A)}; common genes (no MT/ribo): {len(common)}")
+
+rows, Xs, fu, fg = [], [], [], []
+for i in wt:
+    a, h = A[i], H.iloc[i]
+    a = a[:, ~a.var_names.duplicated()]
+    X = a[:, common].X
+    Xs.append(sp.csr_matrix(X, dtype=np.float32))
+    fu.append(FU[i].numpy()); fg.append(FG[i].numpy())
+    study = str(h.dataset_title)
+    grp = f"{study}|{h.patient}" if pd.notna(h.patient) else study
+    for b in a.obs_names:
+        rows.append({"sample": h.id, "spot": b, "organ": h.organ, "oncotree": h.oncotree_code,
+                     "disease": h.disease_state, "study": study, "patient": h.patient, "group": grp})
+M = pd.DataFrame(rows)
+X = sp.vstack(Xs).tocsr()
+M["total_counts"] = np.asarray(X.sum(1)).ravel()
+keep = (M.total_counts >= 200).values
+M, X = M[keep].reset_index(drop=True), X[keep]
+fu, fg = np.concatenate(fu)[keep], np.concatenate(fg)[keep]
+M["n_genes"] = np.asarray((X > 0).sum(1)).ravel()
+
+# log1p CP10k, HVG = top-2000 by variance of the log-normalised values (genes expressed in >=1% spots)
+L = sp.diags(1e4 / M.total_counts.values) @ X
+L.data = np.log1p(L.data)
+frac = np.asarray((X > 0).mean(0)).ravel()
+mu = np.asarray(L.mean(0)).ravel()
+var = np.asarray(L.multiply(L).mean(0)).ravel() - mu ** 2
+var[frac < 0.01] = -1
+hvg = np.sort(np.argsort(-var)[:2000])
+np.save(os.path.join(OUT, "hvg.npy"), L[:, hvg].toarray().astype(np.float32))
+sp.save_npz(os.path.join(OUT, "counts.npz"), X)
+pd.Series(common).to_csv(os.path.join(OUT, "genes.csv"), index=False, header=["gene"])
+pd.Series(np.array(common)[hvg]).to_csv(os.path.join(OUT, "hvg_genes.csv"), index=False, header=["gene"])
+np.save(os.path.join(OUT, "he_uni2.npy"), fu.astype(np.float16)[:, None])
+np.save(os.path.join(OUT, "he_gigapath.npy"), fg.astype(np.float16)[:, None])
+M.to_csv(os.path.join(OUT, "meta.csv"), index_label="idx")
+print(f"spots kept: {len(M)} (dropped {int((~keep).sum())} with <200 counts); samples {M['sample'].nunique()}, "
+      f"groups {M.group.nunique()}, studies {M.study.nunique()}")
+print("organs:", M.organ.value_counts().to_dict())
+print("median counts / genes per spot:", M.total_counts.median(), M.n_genes.median())

experiments/22_extract_gene_fm.py

@@ -0,0 +1,126 @@
+"""Frozen single-cell foundation-model embeddings of HEST Visium spot expression (each spot = one pseudo-cell).
+
+    python 22_extract_gene_fm.py scgpt | geneformer
+
+scgpt      : wanglab/scGPT-human (same re-implementation / strict load as 15_extract_marker_fm.py); input = expressed
+             genes among the 2000 HVGs that are in the scGPT vocab, raw counts -> log1p CP10k -> per-spot quantile
+             binning (51 bins); at most 1200 genes (random subset, as scgpt.tasks.embed_data max_length=1200);
+             <cls> output, L2-normalised.
+geneformer : ctheodoris/Geneformer V2-104M; input = rank-value encoding of all common genes (count / total * 1e4 /
+             gc104M gene median, descending, nonzero), truncated to 4096 tokens incl. <cls>/<eos>;
+             penultimate layer, mean over gene tokens.
+Output cache/hest/gene_<name>.npy (n_spot, 1, D) float16
+"""
+import importlib
+import os
+import pickle
+import sys
+import time
+
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+import torch
+
+from common import CACHE
+
+M15 = importlib.import_module("15_extract_marker_fm")
+OUT = os.path.join(CACHE, "hest")
+name = sys.argv[1]
+dev = "cuda"
+X = sp.load_npz(os.path.join(OUT, "counts.npz")).tocsr()
+genes = pd.read_csv(os.path.join(OUT, "genes.csv")).gene.values
+tot = np.asarray(X.sum(1)).ravel()
+t0 = time.time()
+
+
+def batches(seqs, bs):
+    for i in range(0, len(seqs), bs):
+        yield seqs[i:i + bs]
+
+
+if name == "scgpt":
+    import json
+    from huggingface_hub import hf_hub_download
+    args = json.load(open(hf_hub_download(M15.SCGPT_REPO, "args.json")))
+    vocab = json.load(open(hf_hub_download(M15.SCGPT_REPO, "vocab.json")))
+    sd = torch.load(hf_hub_download(M15.SCGPT_REPO, "best_model.pt"), map_location="cpu", weights_only=True)
+    sd = {k.replace("self_attn.Wqkv.", "self_attn.in_proj_"): v for k, v in sd.items()}
+    PAD, CLS = vocab[args["pad_token"]], vocab["<cls>"]
+    model = M15.ScGPT(len(vocab), args["embsize"], args["nheads"], args["d_hid"], args["nlayers"], PAD)
+    model.load_state_dict(sd, strict=True)
+    model.eval().to(dev)
+    hvg = pd.read_csv(os.path.join(OUT, "hvg_genes.csv")).gene.values
+    gi = np.array([i for i, g in enumerate(genes) if g in set(hvg) and g in vocab])
+    gid = np.array([vocab[g] for g in genes[gi]])
+    L = sp.diags(1e4 / tot) @ X[:, gi]
+    L.data = np.log1p(L.data)
+    L = L.tocsr()
+    rng = np.random.default_rng(0)
+    seqs = []
+    for r in range(L.shape[0]):
+        row = L.getrow(r)
+        nz, v = row.indices, row.data
+        if len(nz) > 1199:
+            keep = np.sort(rng.choice(len(nz), 1199, replace=False))
+            nz, v = nz[keep], v[keep]
+        vals = M15.scgpt_binning(v, args["n_bins"], rng).astype(np.float32)
+        seqs.append((np.r_[CLS, gid[nz]], np.r_[np.float32(args["pad_value"]), vals]))
+    print(f"scGPT: {len(gi)} HVGs in vocab; genes/spot median {np.median([len(s[0]) - 1 for s in seqs])}", flush=True)
+    embs = []
+    with torch.no_grad():
+        for b in batches(seqs, 32):
+            Lm = max(len(g) for g, _ in b)
+            g = torch.full((len(b), Lm), PAD, dtype=torch.long)
+            v = torch.full((len(b), Lm), float(args["pad_value"]))
+            for j, (gg, vv) in enumerate(b):
+                g[j, :len(gg)], v[j, :len(vv)] = torch.from_numpy(gg), torch.from_numpy(vv)
+            g, v = g.to(dev), v.to(dev)
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                h = model.encode(g, v, g.eq(PAD))
+            embs.append(h[:, 0].float().cpu())
+    E = torch.cat(embs).numpy()
+    E /= np.linalg.norm(E, axis=1, keepdims=True)
+else:
+    from transformers import BertModel
+    tokd = pickle.load(open(M15.hf_file(M15.GF_REPO, "geneformer/token_dictionary_gc104M.pkl"), "rb"))
+    medd = pickle.load(open(M15.hf_file(M15.GF_REPO, "geneformer/gene_median_dictionary_gc104M.pkl"), "rb"))
+    n2id = pickle.load(open(M15.hf_file(M15.GF_REPO, "geneformer/gene_name_id_dict_gc104M.pkl"), "rb"))
+    mdir = os.path.dirname(M15.hf_file(M15.GF_REPO, "Geneformer-V2-104M/model.safetensors"))
+    M15.hf_file(M15.GF_REPO, "Geneformer-V2-104M/config.json")
+    gi = np.array([i for i, g in enumerate(genes) if g in n2id and n2id[g] in tokd and n2id[g] in medd])
+    tok = np.array([tokd[n2id[g]] for g in genes[gi]])
+    med = np.array([medd[n2id[g]] for g in genes[gi]])
+    CLS, EOS, PAD = tokd["<cls>"], tokd["<eos>"], tokd["<pad>"]
+    model = BertModel.from_pretrained(mdir, add_pooling_layer=False, attn_implementation="sdpa").eval().to(dev)
+    Xg = (sp.diags(1e4 / tot) @ X[:, gi]).tocsr()
+    seqs = []
+    for r in range(Xg.shape[0]):
+        row = Xg.getrow(r)
+        nz, v = row.indices, row.data / med[row.indices]
+        order = nz[np.argsort(-v, kind="stable")][:4094]
+        seqs.append([CLS] + tok[order].tolist() + [EOS])
+    print(f"Geneformer: {len(gi)} genes in dictionary; tokens/spot median {np.median([len(s) for s in seqs])}", flush=True)
+    order_len = np.argsort([len(s) for s in seqs])        # length-sorted batches (less padding), restored below
+    embs = np.zeros((len(seqs), model.config.hidden_size), np.float32)
+    with torch.no_grad():
+        for k in range(0, len(seqs), 8):
+            idx = order_len[k:k + 8]
+            b = [seqs[i] for i in idx]
+            Lm = max(map(len, b))
+            ids = torch.full((len(b), Lm), PAD, dtype=torch.long)
+            for j, s in enumerate(b):
+                ids[j, :len(s)] = torch.tensor(s)
+            ids = ids.to(dev)
+            att = (ids != PAD).long()
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                h = model(input_ids=ids, attention_mask=att, output_hidden_states=True).hidden_states[-2].float()
+            gm = att.clone().float()
+            gm[:, 0] = 0
+            gm[torch.arange(len(b)), att.sum(1) - 1] = 0
+            embs[idx] = ((h * gm[..., None]).sum(1) / gm.sum(1, keepdim=True).clamp(min=1)).cpu().numpy()
+            if k % 800 == 0:
+                print(f"  {k}/{len(seqs)} {time.time() - t0:.0f}s", flush=True)
+    E = embs
+np.save(os.path.join(OUT, f"gene_{name}.npy"), E.astype(np.float16)[:, None])
+print(f"[{name}] saved {E.shape}  finite={np.isfinite(E).all()}  {time.time() - t0:.0f}s")

experiments/23_bind_gene.py

@@ -0,0 +1,106 @@
+"""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):  # same head as 09_bench / 12_ / 20_
+    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())

experiments/24_eval_gene_emergent.py

@@ -0,0 +1,160 @@
+"""Emergent alignment of the gene-expression modality (HEST, 23_) with protein (ORION mIF, 17_/09_) and text
+(PathGen, 12_). All heads were trained separately against the same frozen UNI2 H&E space; gene never co-occurs with
+mIF or text in training, and no dataset measures RNA + protein + text on the same tissue, so 1:1 pairing accuracy
+cannot be computed for these pairs. Instead:
+
+(A) protein <-> gene, cross-dataset neighbour transfer
+    protein->gene: each ORION test tile (mIF embedding) retrieves its k=20 nearest HEST spots by gene embedding;
+                   predicted RNA level = mean log1p-CP10k of the matching gene over those spots;
+                   score = Spearman(predicted RNA, measured protein = marker+ cells / nuclei) over tiles (>=20 nuclei).
+    gene->protein: each HEST spot (gene embedding) retrieves its k=20 nearest ORION tiles by mIF embedding;
+                   Spearman(predicted protein, measured RNA of the spot).
+    references   : he->gene (ORION UNI2 vs HEST gene head; trained modality pair, cross-dataset),
+                   he->he (UNI2 vs UNI2, image similarity only), and the same with labels shuffled (chance ~0).
+    galleries    : all HEST spots, and bowel-only spots (same organ as ORION colorectal cancer).
+(B) text <-> gene concept probes: prompts.json through the PathGen-trained text head, scored against HEST test-fold
+    gene embeddings (and UNI2 H&E of the same spots); label = RNA of the concept gene; Spearman + AUROC (top quartile).
+Outputs results/gene_emergent_protein.csv, results/gene_emergent_text.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.functional as F
+from scipy.stats import spearmanr
+from sklearn.metrics import roc_auc_score
+
+from common import CACHE, ROOT, RUNS, load_all_tiles
+
+OUT = os.path.join(ROOT, "results")
+H = os.path.join(CACHE, "hest")
+K = 20
+M = pd.read_csv(os.path.join(H, "meta.csv"))
+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)}
+
+
+def rna(gs):
+    return np.asarray(L[:, [gpos[g] for g in gs]].mean(1)).ravel()
+
+
+PROT2RNA = {"CD45": ["PTPRC"], "Pan-CK": ["KRT8", "EPCAM"], "E-cadherin": ["CDH1"], "SMA": ["ACTA2"],
+            "Ki67": ["MKI67"], "CD31": ["PECAM1"], "CD68": ["CD68"], "CD163": ["CD163"], "CD3e": ["CD3E"],
+            "CD4": ["CD4"], "CD8a": ["CD8A"], "FOXP3": ["FOXP3"], "CD20": ["MS4A1"], "PD-L1": ["CD274"],
+            "PD-1": ["PDCD1"]}
+RNA = {m: rna(g) for m, g in PROT2RNA.items()}
+
+# ---------------- HEST: gene embeddings of every spot from the fold whose head did not train on it
+hest = {}
+for f in glob.glob(os.path.join(RUNS, "gene", "uni2_*_fold*", "test.pt")):
+    enc = os.path.basename(os.path.dirname(f)).split("_", 1)[1].rsplit("_fold", 1)[0]
+    d = torch.load(f, weights_only=False)
+    z = hest.setdefault(enc, {"gene": torch.zeros(len(M), d["gene"].shape[1]), "he": torch.zeros(len(M), d["he"].shape[1]),
+                              "fold": np.full(len(M), -1)})
+    z["gene"][d["idx"]], z["he"][d["idx"]] = d["gene"], d["he"]
+    z["fold"][d["idx"]] = int(f.split("_fold")[1].split(os.sep)[0])
+hest = {e: z for e, z in hest.items() if (z["fold"] >= 0).all()}
+print("gene encoders:", list(hest))
+
+# ---------------- ORION: patient 5-fold test embeddings (mIF encoders) in the same UNI2 space
+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
+orion = {}
+for root in ["bench", "bench_fm"]:
+    for f in glob.glob(os.path.join(RUNS, root, "patient_5fold_*_uni2_*", "test_embeddings.pt")):
+        a = json.load(open(os.path.join(os.path.dirname(f), "train.json")))["args"]
+        E = torch.load(f, weights_only=False)
+        o = orion.setdefault(a["if_enc"], {"if": torch.zeros(len(df), E["z"]["if"].shape[1]), "he": torch.zeros(len(df), E["z"]["he"].shape[1]),
+                                           "seen": np.zeros(len(df), bool)})
+        t = np.asarray(E["tile_id"])
+        o["if"][t], o["he"][t], o["seen"][t] = E["z"]["if"], E["z"]["he"], True
+orion = {e: o for e, o in orion.items() if o["seen"].all()}
+print("mIF encoders:", list(orion))
+tissue = nn_ >= 20
+PROT = {m: df[f"{m}_count"].values / np.maximum(nn_, 1) for m in PROT2RNA}
+
+
+def knn_mean(q, g, vals, k=K, bs=1024):
+    out = np.zeros(len(q))
+    for i in range(0, len(q), bs):
+        s = q[i:i + bs] @ g.T
+        idx = s.topk(k, dim=1).indices.numpy()
+        out[i:i + bs] = vals[idx].mean(1)
+    return out
+
+
+rows = []
+rng = np.random.default_rng(0)
+for genc, hz in hest.items():
+    for gal in ["all", "bowel"]:
+        gm = np.ones(len(M), bool) if gal == "all" else (M.organ == "Bowel").values
+        for ienc, oz in orion.items():
+            queries = [("protein->gene (emergent)", oz["if"], hz["gene"]), ("he->gene (cross-dataset)", oz["he"], hz["gene"]),
+                       ("he->he (image only)", oz["he"], hz["he"])]
+            for kind, q, g in queries:
+                for m in PROT2RNA:
+                    pred = knn_mean(q[tissue], g[gm], RNA[m][gm])
+                    y = PROT[m][tissue]
+                    rows.append({"gene_enc": genc, "if_enc": ienc, "gallery": gal, "kind": kind, "marker": m,
+                                 "spearman": spearmanr(pred, y).correlation,
+                                 "spearman_shuffled": spearmanr(pred, rng.permutation(y)).correlation})
+            # reverse: each HEST spot -> nearest ORION tiles
+            for kind, q, g in [("gene->protein (emergent)", hz["gene"][gm], oz["if"][tissue]), ("gene->he (cross-dataset)", hz["gene"][gm], oz["he"][tissue])]:
+                for m in PROT2RNA:
+                    pred = knn_mean(q, g, PROT[m][tissue])
+                    y = RNA[m][gm]
+                    rows.append({"gene_enc": genc, "if_enc": ienc, "gallery": gal, "kind": kind, "marker": m,
+                                 "spearman": spearmanr(pred, y).correlation,
+                                 "spearman_shuffled": spearmanr(pred, rng.permutation(y)).correlation})
+    print("done", genc, flush=True)
+P = pd.DataFrame(rows)
+P.to_csv(os.path.join(OUT, "gene_emergent_protein.csv"), index=False)
+
+# ---------------- (B) text <-> gene concept probes
+prompts = json.load(open(os.path.join(ROOT, "prompts.json")))
+CONCEPT_RNA = {"immune": ["PTPRC"], "tumor_epith": ["KRT8", "EPCAM"], "stroma": ["ACTA2"], "proliferation": ["MKI67"],
+               "vessels": ["PECAM1"], "macrophage": ["CD68"]}
+TXT = sorted(os.path.basename(d).split("_", 1)[1] for d in glob.glob(os.path.join(RUNS, "text_pathgen", "uni2_*"))
+             if os.path.exists(os.path.join(CACHE, f"prompts_text_{os.path.basename(d).split('_', 1)[1]}.npy")))
+
+
+def load_head(enc, din, dout):
+    h = torch.nn.Sequential(torch.nn.Linear(din, 512), torch.nn.GELU(), torch.nn.Dropout(0.1), torch.nn.Linear(512, 512),
+                            torch.nn.GELU(), torch.nn.LayerNorm(512), torch.nn.Linear(512, dout))
+    h.load_state_dict(torch.load(os.path.join(RUNS, "text_pathgen", f"uni2_{enc}", "head.pt"), map_location="cpu"))
+    return h.eval()
+
+
+trow = []
+for t in TXT:
+    prm = torch.from_numpy(np.load(os.path.join(CACHE, f"prompts_text_{t}.npy"))[:, 0].astype(np.float32))
+    with torch.no_grad():
+        zp = F.normalize(load_head(t, prm.shape[1], 1536)(prm), dim=-1)
+    keys = list(prompts)
+    for genc, hz in hest.items():
+        embs = [("he (UNI2)", hz["he"]), (f"gene ({genc})", hz["gene"])] if genc == list(hest)[0] else [(f"gene ({genc})", hz["gene"])]
+        for mod, z in embs:
+            for c, gs in CONCEPT_RNA.items():
+                s = (z @ zp[keys.index(c)]).numpy()
+                for gal in ["all", "bowel"]:
+                    gm = np.ones(len(M), bool) if gal == "all" else (M.organ == "Bowel").values
+                    for k in range(5):
+                        sel = gm & (hz["fold"] == k)
+                        y = rna(gs)[sel]
+                        lab = y >= np.quantile(y, 0.75)
+                        trow.append({"text_enc": t, "modality": mod, "concept": c, "gallery": gal, "fold": k,
+                                     "spearman": spearmanr(s[sel], y).correlation,
+                                     "auroc": roc_auc_score(lab, s[sel]) if 0 < lab.sum() < len(lab) else np.nan})
+T = pd.DataFrame(trow)
+T.to_csv(os.path.join(OUT, "gene_emergent_text.csv"), index=False)
+pd.set_option("display.width", 250); pd.set_option("display.max_rows", 500)
+print(P.groupby(["gallery", "kind", "if_enc", "gene_enc"])[["spearman", "spearman_shuffled"]].mean().round(3).to_string())
+print(P[P.gallery == "bowel"].groupby(["kind", "marker"]).spearman.mean().unstack("kind").round(2).to_string())
+print(T.groupby(["gallery", "text_enc", "modality"])[["auroc", "spearman"]].mean().round(3).to_string())

experiments/25_gigatime_infer.py

@@ -0,0 +1,361 @@
+"""GigaTIME (Providence/Microsoft, Cell 2025) virtual mIF inference.
+
+Model : prov-gigatime/GigaTIME (HF, gated, Apache-2.0) -- UNet++ CNN, H&E RGB -> 23 sigmoid maps
+        (21 proteins + 2 background channels TRITC/Cy5, which are dropped here).
+Code  : https://github.com/prov-gigatime/GigaTIME (architecture re-implemented below, verbatim
+        layer layout of scripts/archs.py, so nothing from the downloaded repo is imported).
+
+Preprocessing (as in the official notebooks): RGB/255, ImageNet mean/std normalisation, no stain norm.
+Optional bilinear upsampling by --scale before inference (the model was trained on native-resolution
+Providence tiles; ORION/PathGen tiles are ~0.5 um/px, so --scale 2 brings them to ~0.25 um/px = 40x).
+The net is fully convolutional (4 poolings): images are reflect-padded to a multiple of 16 and run
+whole; images larger than --max_full px are run as 512-px tiles with a 64-px overlap margin
+(only the tile centres are kept), which avoids the seams of the official non-overlapping 256 windows.
+Outputs are sigmoid probabilities (the official code thresholds at 0.5; we keep the probabilities).
+
+Usage (HF token only via environment, never stored):
+  HF_TOKEN=... .venv/bin/python 25_gigatime_infer.py --set orion --scale 2
+  .venv/bin/python 25_gigatime_infer.py --set orion --scale 1 --subset 200 --tag scaletest_s1 --no_maps
+  .venv/bin/python 25_gigatime_infer.py --set tcga_text --scale 2
+  .venv/bin/python 25_gigatime_infer.py --set ge_patches --scale 2
+
+Outputs (cache/gigatime/):
+  <set>_maps.npy   uint8 (N, 21, R, R)  probability*255, area-resampled to R=--map_res (default 128)
+  <set>_means.npy  float32 (N, 21)      per-tile channel mean probability (at inference resolution,
+                                         padding excluded)
+  <set>_frac.npy   float32 (N, 21)      per-tile fraction of pixels with prob > 0.5
+  <set>_meta.csv   row order, source path, status, sizes
+  channels.json
+  orion: orion_validation.csv/json (shared-channel agreement with the real ORION mIF)
+"""
+import argparse
+import glob
+import json
+import os
+import queue
+import sys
+import threading
+import time
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image, PngImagePlugin
+
+PngImagePlugin.MAX_TEXT_CHUNK = 64 * 1024 * 1024
+
+ROOT = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, ROOT)
+from common import CACHE, ORION, load_all_tiles  # noqa: E402
+
+OUT = os.path.join(CACHE, "gigatime")
+WDIR = os.path.join(ROOT, "ext", "gigatime", "weights", "GigaTIME")
+PATHGEN = os.path.join(ROOT, "..", "PathGen-1.6M")
+GE_DIRS = ["brain_image", "brain_train_image", "brain_trainnew_image", "gliomas_test_image_new",
+           "gliomas_train_image", "idc_image", "idc_train_image", "saromas_test_image_new",
+           "saromas_train_image_new", "skin_image", "skin_train_image"]
+
+ALL_CHANNELS = ['DAPI', 'TRITC', 'Cy5', 'PD-1', 'CD14', 'CD4', 'T-bet', 'CD34', 'CD68', 'CD16', 'CD11c',
+                'CD138', 'CD20', 'CD3', 'CD8', 'PD-L1', 'CK', 'Ki67', 'Tryptase', 'Actin-D', 'Caspase3-D',
+                'PHH3-B', 'Transgelin']
+BACKGROUND = ['TRITC', 'Cy5']
+KEEP = [i for i, c in enumerate(ALL_CHANNELS) if c not in BACKGROUND]
+CHANNELS = [ALL_CHANNELS[i] for i in KEEP]
+# GigaTIME channel -> ORION channel (ORION if/*.tiff channel order == *_prop column order)
+SHARED = {'DAPI': 'Hoechst', 'PD-1': 'PD-1', 'CD4': 'CD4', 'CD68': 'CD68', 'CD20': 'CD20', 'CD3': 'CD3e',
+          'CD8': 'CD8a', 'PD-L1': 'PD-L1', 'CK': 'Pan-CK', 'Ki67': 'Ki67',
+          # tentative (smooth-muscle markers vs SMA)
+          'Actin-D': 'SMA', 'Transgelin': 'SMA'}
+MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
+STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
+
+
+# ---------------- architecture (same layout as GigaTIME scripts/archs.py, UNet++) ----------------
+class VGGBlock(nn.Module):
+    def __init__(self, i, m, o):
+        super().__init__()
+        self.relu = nn.ReLU(inplace=True)
+        self.conv1, self.bn1 = nn.Conv2d(i, m, 3, padding=1), nn.BatchNorm2d(m)
+        self.conv2, self.bn2 = nn.Conv2d(m, o, 3, padding=1), nn.BatchNorm2d(o)
+
+    def forward(self, x):
+        return self.relu(self.bn2(self.conv2(self.relu(self.bn1(self.conv1(x))))))
+
+
+class GigaTIME(nn.Module):
+    def __init__(self, num_classes=23, input_channels=3):
+        super().__init__()
+        f = [32, 64, 128, 256, 512]
+        self.pool = nn.MaxPool2d(2, 2)
+        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        self.conv0_0 = VGGBlock(input_channels, f[0], f[0])
+        self.conv1_0 = VGGBlock(f[0], f[1], f[1])
+        self.conv2_0 = VGGBlock(f[1], f[2], f[2])
+        self.conv3_0 = VGGBlock(f[2], f[3], f[3])
+        self.conv4_0 = VGGBlock(f[3], f[4], f[4])
+        self.conv0_1 = VGGBlock(f[0] + f[1], f[0], f[0])
+        self.conv1_1 = VGGBlock(f[1] + f[2], f[1], f[1])
+        self.conv2_1 = VGGBlock(f[2] + f[3], f[2], f[2])
+        self.conv3_1 = VGGBlock(f[3] + f[4], f[3], f[3])
+        self.conv0_2 = VGGBlock(f[0] * 2 + f[1], f[0], f[0])
+        self.conv1_2 = VGGBlock(f[1] * 2 + f[2], f[1], f[1])
+        self.conv2_2 = VGGBlock(f[2] * 2 + f[3], f[2], f[2])
+        self.conv0_3 = VGGBlock(f[0] * 3 + f[1], f[0], f[0])
+        self.conv1_3 = VGGBlock(f[1] * 3 + f[2], f[1], f[1])
+        self.conv0_4 = VGGBlock(f[0] * 4 + f[1], f[0], f[0])
+        self.final = nn.Conv2d(f[0], num_classes, kernel_size=1)
+
+    def forward(self, x):
+        up, pool = self.up, self.pool
+        x0_0 = self.conv0_0(x)
+        x1_0 = self.conv1_0(pool(x0_0))
+        x0_1 = self.conv0_1(torch.cat([x0_0, up(x1_0)], 1))
+        x2_0 = self.conv2_0(pool(x1_0))
+        x1_1 = self.conv1_1(torch.cat([x1_0, up(x2_0)], 1))
+        x0_2 = self.conv0_2(torch.cat([x0_0, x0_1, up(x1_1)], 1))
+        x3_0 = self.conv3_0(pool(x2_0))
+        x2_1 = self.conv2_1(torch.cat([x2_0, up(x3_0)], 1))
+        x1_2 = self.conv1_2(torch.cat([x1_0, x1_1, up(x2_1)], 1))
+        x0_3 = self.conv0_3(torch.cat([x0_0, x0_1, x0_2, up(x1_2)], 1))
+        x4_0 = self.conv4_0(pool(x3_0))
+        x3_1 = self.conv3_1(torch.cat([x3_0, up(x4_0)], 1))
+        x2_2 = self.conv2_2(torch.cat([x2_0, x2_1, up(x3_1)], 1))
+        x1_3 = self.conv1_3(torch.cat([x1_0, x1_1, x1_2, up(x2_2)], 1))
+        x0_4 = self.conv0_4(torch.cat([x0_0, x0_1, x0_2, x0_3, up(x1_3)], 1))
+        return self.final(x0_4)
+
+
+def load_model(device):
+    wp = os.path.join(WDIR, "model.pth")
+    if not os.path.exists(wp):  # token is read from the HF_TOKEN env var by huggingface_hub
+        from huggingface_hub import snapshot_download
+        snapshot_download("prov-gigatime/GigaTIME", local_dir=WDIR)
+    sd = torch.load(wp, map_location="cpu", weights_only=True)
+    m = GigaTIME(23, 3)
+    m.load_state_dict(sd, strict=True)
+    return m.eval().to(device)
+
+
+# ---------------- inference ----------------
+def _pad_reflect(x, top, bottom, left, right):
+    # reflect padding must be < dim; fall back to replicate for tiny images
+    mode = "reflect" if max(top, bottom) < x.shape[2] and max(left, right) < x.shape[3] else "replicate"
+    return F.pad(x, (left, right, top, bottom), mode=mode)
+
+
+@torch.no_grad()
+def infer(model, x, max_full=768, tile=512, margin=64, tile_bs=4):
+    """x: (1,3,H,W) normalised float on device -> logits (1,23,H,W) float32."""
+    _, _, H, W = x.shape
+    with torch.autocast("cuda", dtype=torch.float16):
+        if max(H, W) <= max_full:
+            xp = _pad_reflect(x, 0, (-H) % 16, 0, (-W) % 16)
+            return model(xp)[:, :, :H, :W].float()
+        step = tile - 2 * margin
+        ny, nx = -(-H // step), -(-W // step)
+        xp = _pad_reflect(x, margin, ny * step + margin - H, margin, nx * step + margin - W)
+        out = torch.zeros(1, 23, ny * step, nx * step, device=x.device)
+        pos = [(iy, ix) for iy in range(ny) for ix in range(nx)]
+        for k in range(0, len(pos), tile_bs):
+            chunk = pos[k:k + tile_bs]
+            tiles = torch.cat([xp[:, :, iy * step:iy * step + tile, ix * step:ix * step + tile] for iy, ix in chunk])
+            o = model(tiles)[:, :, margin:margin + step, margin:margin + step].float()
+            for j, (iy, ix) in enumerate(chunk):
+                out[:, :, iy * step:(iy + 1) * step, ix * step:(ix + 1) * step] = o[j]
+        return out[:, :, :H, :W]
+
+
+def read_rgb(path):
+    with Image.open(path) as im:
+        return np.asarray(im.convert("RGB"))
+
+
+def prefetch(paths, extra=None, qsize=16):
+    q = queue.Queue(qsize)
+
+    def work():
+        for i, p in enumerate(paths):
+            try:
+                img = read_rgb(p)
+                ex = extra(i) if extra else None
+                q.put((i, img, ex, None))
+            except Exception as e:  # unreadable -> logged as failed
+                q.put((i, None, None, f"{type(e).__name__}: {str(e)[:200]}"))
+        q.put(None)
+
+    threading.Thread(target=work, daemon=True).start()
+    while True:
+        it = q.get()
+        if it is None:
+            return
+        yield it
+
+
+# ---------------- datasets ----------------
+def get_set(name, subset, seed=0):
+    if name == "orion":
+        df = load_all_tiles()
+        meta = pd.DataFrame({"tile_id": df.tile_id, "source_path": [os.path.normpath(os.path.join(ORION, p)) for p in df.image_path],
+                             "if_path": [os.path.normpath(os.path.join(ORION, p)) for p in df.target_path],
+                             "orion_split": df.orion_split})
+        prop = df[[c + "_prop" for c in sorted(set(SHARED.values()), key=list(SHARED.values()).index)]]
+    elif name == "tcga_text":
+        d = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv"))
+        meta = pd.DataFrame({"idx": d["idx"], "key": d["key"], "source_path": d.image_path.map(os.path.normpath)})
+        prop = None
+    elif name == "ge_patches":
+        files = []
+        for g in GE_DIRS:
+            files += [(g, f) for f in sorted(glob.glob(os.path.join(PATHGEN, g, "*"))) if os.path.isfile(f) and f.lower().endswith((".png", ".jpg", ".jpeg", ".tif", ".tiff"))]
+        meta = pd.DataFrame({"group": [g for g, _ in files], "source_path": [os.path.normpath(f) for _, f in files]})
+        prop = None
+    else:
+        raise ValueError(name)
+    if subset:
+        sel = np.sort(np.random.default_rng(seed).choice(len(meta), subset, replace=False))
+        meta = meta.iloc[sel].reset_index(drop=True)
+        prop = prop.iloc[sel].reset_index(drop=True) if prop is not None else None
+    return meta.reset_index(drop=True), prop
+
+
+def pearson_rows(a, b):
+    """row-wise Pearson for (n, p) arrays; NaN where a row is constant."""
+    a = a - a.mean(1, keepdims=True)
+    b = b - b.mean(1, keepdims=True)
+    den = np.sqrt((a * a).sum(1) * (b * b).sum(1))
+    with np.errstate(invalid="ignore", divide="ignore"):
+        return np.where(den > 0, (a * b).sum(1) / den, np.nan)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--set", required=True, choices=["orion", "tcga_text", "ge_patches"])
+    ap.add_argument("--scale", type=float, default=2.0)
+    ap.add_argument("--subset", type=int, default=0)
+    ap.add_argument("--tag", default="")
+    ap.add_argument("--no_maps", action="store_true")
+    ap.add_argument("--map_res", type=int, default=128)
+    ap.add_argument("--pix_tiles", type=int, default=400, help="ORION tiles used for pixel-level Pearson")
+    ap.add_argument("--gpu_frac", type=float, default=0.2, help="cap on this process' share of GPU memory")
+    a = ap.parse_args()
+
+    os.makedirs(OUT, exist_ok=True)
+    dev = "cuda"
+    torch.cuda.set_per_process_memory_fraction(a.gpu_frac)
+    torch.backends.cudnn.benchmark = False
+    model = load_model(dev)
+    mean, std = MEAN.to(dev), STD.to(dev)
+    tag = a.tag or a.set
+    meta, prop = get_set(a.set, a.subset)
+    N, C, R = len(meta), len(KEEP), a.map_res
+    print(f"[{tag}] N={N} scale={a.scale} channels={C}", flush=True)
+
+    with open(os.path.join(OUT, "channels.json"), "w") as f:
+        json.dump({"channels": CHANNELS, "model_output_channels": ALL_CHANNELS, "dropped_background": BACKGROUND,
+                   "orion_shared": SHARED, "value": "sigmoid probability (maps: uint8 = round(prob*255))",
+                   "source": "huggingface.co/prov-gigatime/GigaTIME"}, f, indent=1)
+
+    maps = None if a.no_maps else np.lib.format.open_memmap(os.path.join(OUT, f"{tag}_maps.npy"), "w+", np.uint8, (N, C, R, R))
+    means = np.full((N, C), np.nan, np.float32)
+    frac = np.full((N, C), np.nan, np.float32)
+    status, hw, inhw = ["failed: not processed"] * N, [(0, 0)] * N, [(0, 0)] * N
+
+    is_orion = a.set == "orion"
+    if is_orion:
+        import tifffile
+        orion_ch = ["Hoechst", "CD31", "CD45", "CD68", "CD4", "FOXP3", "CD8a", "CD45RO", "CD20", "PD-L1", "CD3e",
+                    "CD163", "E-cadherin", "PD-1", "Ki67", "Pan-CK", "SMA"]
+        pairs = [(g, CHANNELS.index(g), o, orion_ch.index(o)) for g, o in SHARED.items()]
+        real_mean = np.full((N, len(orion_ch)), np.nan, np.float32)
+        pix_sel = set(np.linspace(0, N - 1, min(a.pix_tiles, N)).astype(int).tolist())
+        pix_r = {g: [] for g, *_ in pairs}
+        box_r = {g: [] for g, *_ in pairs}
+        pooled = {g: ([], []) for g, *_ in pairs}
+        rng = np.random.default_rng(0)
+        extra = lambda i: tifffile.imread(meta.if_path[i])  # (333,333,17) uint8
+    else:
+        extra = None
+
+    t0 = time.time()
+    for i, img, ifa, err in prefetch(list(meta.source_path), extra):
+        if err is not None:
+            status[i] = "failed: " + err
+            print(f"  FAIL {i} {meta.source_path[i]} {err}", flush=True)
+            continue
+        try:
+            H, W = img.shape[:2]
+            x = torch.from_numpy(img).to(dev).permute(2, 0, 1)[None].float() / 255.0
+            if a.scale != 1:
+                x = F.interpolate(x, scale_factor=a.scale, mode="bilinear", align_corners=False).clamp(0, 1)
+            x = (x - mean) / std
+            p = torch.sigmoid(infer(model, x))[:, KEEP]  # (1,C,h,w)
+            means[i] = p.mean((0, 2, 3)).cpu().numpy()
+            frac[i] = (p > 0.5).float().mean((0, 2, 3)).cpu().numpy()
+            if maps is not None:
+                maps[i] = (F.interpolate(p, size=(R, R), mode="area")[0] * 255).round().clamp(0, 255).byte().cpu().numpy()
+            status[i], hw[i], inhw[i] = "ok", (H, W), tuple(p.shape[2:])
+            if is_orion:
+                real_mean[i] = ifa.reshape(-1, ifa.shape[2]).mean(0)
+                if i in pix_sel:
+                    pr = F.interpolate(p, size=(H, W), mode="area")[0].cpu().numpy()  # back onto the IF grid
+                    for g, gi, o, oi in pairs:
+                        a1, b1 = pr[gi].ravel()[None], ifa[..., oi].astype(np.float32).ravel()[None]
+                        pix_r[g].append(pearson_rows(a1, b1)[0])
+                        bh, bw = H // 8 * 8, W // 8 * 8  # 8x8-pixel boxes (~4 um), as in the official eval
+                        a8 = pr[gi][:bh, :bw].reshape(bh // 8, 8, bw // 8, 8).mean((1, 3)).ravel()[None]
+                        b8 = ifa[:bh, :bw, oi].astype(np.float32).reshape(bh // 8, 8, bw // 8, 8).mean((1, 3)).ravel()[None]
+                        box_r[g].append(pearson_rows(a8, b8)[0])
+                        s = rng.choice(a1.shape[1], 2000, replace=False)
+                        pooled[g][0].append(a1[0, s]); pooled[g][1].append(b1[0, s])
+        except Exception as e:
+            status[i] = f"failed: {type(e).__name__}: {str(e)[:200]}"
+            print(f"  FAIL {i} {meta.source_path[i]} {status[i]}", flush=True)
+        if (i + 1) % 500 == 0:
+            print(f"  {i + 1}/{N}  {time.time() - t0:.0f}s  gpu_max={torch.cuda.max_memory_allocated() / 2**30:.2f}GB", flush=True)
+
+    if maps is not None:
+        maps.flush()
+    meta["status"] = status
+    meta["orig_h"], meta["orig_w"] = [h for h, _ in hw], [w for _, w in hw]
+    meta["infer_h"], meta["infer_w"] = [h for h, _ in inhw], [w for _, w in inhw]
+    meta["scale"] = a.scale
+    meta.insert(0, "row", np.arange(N))
+    meta.to_csv(os.path.join(OUT, f"{tag}_meta.csv"), index=False)
+    np.save(os.path.join(OUT, f"{tag}_means.npy"), means)
+    np.save(os.path.join(OUT, f"{tag}_frac.npy"), frac)
+    nok = int((meta.status == "ok").sum())
+    print(f"[{tag}] done ok={nok} failed={N - nok} time={time.time() - t0:.0f}s "
+          f"gpu_max={torch.cuda.max_memory_allocated() / 2**30:.2f}GB", flush=True)
+
+    if is_orion:
+        from scipy.stats import spearmanr, pearsonr
+        ok = (meta.status == "ok").values
+        rows = []
+        for g, gi, o, oi in pairs:
+            r = dict(gigatime=g, orion=o,
+                     tile_mean_spearman=spearmanr(means[ok, gi], real_mean[ok, oi])[0],
+                     tile_frac_spearman=spearmanr(frac[ok, gi], real_mean[ok, oi])[0],
+                     tile_mean_vs_prop_spearman=spearmanr(means[ok, gi], prop[o + "_prop"].values[ok])[0],
+                     pixel_pearson_tile_median=float(np.nanmedian(pix_r[g])),
+                     pixel_pearson_tile_mean=float(np.nanmean(pix_r[g])),
+                     pixel_pearson_pooled=pearsonr(np.concatenate(pooled[g][0]), np.concatenate(pooled[g][1]))[0],
+                     box8_pearson_tile_mean=float(np.nanmean(box_r[g])),
+                     n_tiles=int(ok.sum()), n_pix_tiles=len(pix_r[g]))
+            rows.append(r)
+        v = pd.DataFrame(rows)
+        v.to_csv(os.path.join(OUT, f"{tag}_validation.csv"), index=False)
+        core = v[~v.gigatime.isin(["Actin-D", "Transgelin"])]
+        summ = {"scale": a.scale, "N": int(ok.sum()),
+                "mean_over_10_core_channels": core.drop(columns=["gigatime", "orion"]).mean().to_dict()}
+        with open(os.path.join(OUT, f"{tag}_validation.json"), "w") as f:
+            json.dump(summ, f, indent=1, default=float)
+        np.save(os.path.join(OUT, f"{tag}_real_if_means.npy"), real_mean)
+        pd.set_option("display.width", 250)
+        print(v.round(3).to_string(index=False))
+        print(json.dumps(summ, indent=1, default=float))
+
+
+if __name__ == "__main__":
+    main()