| """End-to-end check of embed.py from RAW inputs on examples/ (64 held-out fold-0 test items per dataset; needs the |
| frozen encoders, see setup_encoders.sh; GPU). |
| |
| python -I example_check.py [models/ej_ad_mv_fold0] |
| |
| 1. raw input -> encoder features, compared with the cached features the model was trained / tested on |
| (examples/ref_feats.npz): per-item cosine. Expected ~1.000; Tahoe-x1 / SE-600M ~0.99 (their official input |
| pipelines sample genes at random); VirTues lower here because mIF standardisation statistics are computed over the |
| tiles given in one call (64 tiles here vs the whole ORION set for the cache). |
| 2. raw input -> 1536-d embeddings -> retrieval inside the 64 items (R@1, both directions; chance 1/64), next to the |
| same retrieval from the cached features. |
| Output: examples/example_check.csv; exit code 1 if a deterministic encoder does not reproduce its cached features. |
| """ |
| import os |
| import sys |
| import time |
|
|
| import numpy as np |
| import pandas as pd |
| import scipy.sparse as sp |
| import torch |
|
|
| ROOT = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.insert(0, ROOT) |
| from embed import RNA_E, TXT_E, Embedder |
|
|
| MD = sys.argv[1] if len(sys.argv) > 1 else os.path.join(ROOT, "models", "ej_ad_mv_fold0") |
| EX = os.path.join(ROOT, "examples") |
| E = Embedder(MD) |
| R = dict(np.load(os.path.join(EX, "ref_feats.npz"))) |
| rows, t0 = [], time.time() |
| MIN_COS = {"uni2": .999, "medvtok": .999, "keep": .999, "biomedclip": .99, "qwen3-0.6b": .999, "hvg": .999, |
| "aidocell": .999, "tahoex1": .97, "se600m": .97, "virtues": 0.0} |
|
|
|
|
| def cos(a, b): |
| a, b = torch.as_tensor(np.asarray(a, np.float32)), torch.as_tensor(np.asarray(b, np.float32)) |
| return torch.nn.functional.cosine_similarity(a, b, dim=-1).numpy() |
|
|
|
|
| def feat(ds, enc, raw): |
| c = cos(raw, R[f"{ds}_{enc}"]) |
| rows.append({"check": "features", "data": ds, "encoder": enc, "n": len(c), "cos_median": float(np.median(c)), |
| "cos_min": float(c.min()), "ok": bool(np.median(c) >= MIN_COS[enc])}) |
| print(f"[{time.time() - t0:4.0f}s] {ds:5s} {enc:11s} cos median {np.median(c):.4f} min {c.min():.4f}" |
| f"{'' if rows[-1]['ok'] else ' <-- FAIL'}", flush=True) |
|
|
|
|
| def r1(a, b): |
| s = a @ b.T; d = s.diag() |
| return float((((s > d[:, None]).sum(1) == 0).float().mean() + ((s.T > d[:, None]).sum(1) == 0).float().mean()) / 2) |
|
|
|
|
| |
| U, M = {}, {} |
| U["orion"], M["orion"] = E.he_feats(sorted(os.path.join(EX, "orion_he", f) for f in os.listdir(os.path.join(EX, "orion_he")))) |
| U["hest"], M["hest"] = E.he_feats(np.load(os.path.join(EX, "hest_he.npy"))) |
| C = pd.read_csv(os.path.join(EX, "tcga_captions.csv")) |
| U["tcga"], M["tcga"] = E.he_feats([os.path.join(EX, "tcga_he", f) for f in C.file]) |
| for ds in U: |
| feat(ds, "uni2", U[ds]); feat(ds, "medvtok", M[ds]) |
| E.release() |
| |
| Ft = E.text_feats(C.caption.fillna("").astype(str).tolist()) |
| for e in TXT_E: |
| feat("tcga", e, Ft[e]) |
| E.release() |
| |
| genes = pd.read_csv(os.path.join(EX, "hest_genes.csv")).gene.values |
| Fr = E.rna_feats(sp.load_npz(os.path.join(EX, "hest_counts.npz")), genes) |
| for e in RNA_E: |
| feat("hest", e, Fr[e]) |
| E.release() |
| |
| m = np.load(os.path.join(EX, "orion_mif.npz"), allow_pickle=True) |
| Fm = E.mif_feats(m["tiles"], m["slide"]) |
| feat("orion", "virtues", Fm) |
| E.release() |
|
|
| |
| pairs = {"H&E<->mIF (ORION)": ((E.he_from_feats(U["orion"], M["orion"]), E.mif_from_feats(Fm)), |
| (E.he_from_feats(R["orion_uni2"], R["orion_medvtok"]), E.mif_from_feats(R["orion_virtues"]))), |
| "H&E<->RNA (HEST)": ((E.he_from_feats(U["hest"], M["hest"]), E.rna_from_feats(Fr)), |
| (E.he_from_feats(R["hest_uni2"], R["hest_medvtok"]), |
| E.rna_from_feats({e: R[f"hest_{e}"] for e in RNA_E}))), |
| "H&E<->text (TCGA)": ((E.he_from_feats(U["tcga"], M["tcga"]), E.text_from_feats(Ft)), |
| (E.he_from_feats(R["tcga_uni2"], R["tcga_medvtok"]), |
| E.text_from_feats({e: R[f"tcga_{e}"] for e in TXT_E})))} |
| for p, ((a, b), (ac, bc)) in pairs.items(): |
| rows.append({"check": "retrieval", "data": p, "n": len(a), "R@1_raw": r1(a, b), "R@1_cached": r1(ac, bc), |
| "chance": 1 / len(a)}) |
| print(f"{p:18s} R@1 within {len(a)}: raw inputs {rows[-1]['R@1_raw']:.3f} | cached features " |
| f"{rows[-1]['R@1_cached']:.3f} (chance {1 / len(a):.3f})", flush=True) |
| D = pd.DataFrame(rows) |
| D.to_csv(os.path.join(EX, "example_check.csv"), index=False) |
| bad = D[(D.check == "features") & ~D.ok.astype(bool)] |
| print(f"done [{time.time() - t0:.0f}s]:", "all encoders reproduce the cached features" if bad.empty else |
| f"FAILED {bad[['data', 'encoder']].values.tolist()}") |
| sys.exit(0 if bad.empty else 1) |
|
|
