"""5-fold CV: retrain PANDA per fold, eval by prototype-cosine on held-out 20%.""" from __future__ import annotations from pathlib import Path import warnings, json, sys, time warnings.filterwarnings("ignore") import numpy as np import pandas as pd import anndata as ad import torch import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader from sklearn.model_selection import StratifiedKFold from sklearn.metrics import (accuracy_score, f1_score, roc_auc_score, classification_report) import os as _os from pathlib import Path as _Path PANDA_ROOT = _Path(_os.environ.get("PANDA_ROOT", str(_Path(__file__).resolve().parents[2]))) sys.path.insert(0, str(PANDA_ROOT)) from panda.pan_skin.model import ( PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, prototype_infonce ) CORPUS = Path(str(PANDA_ROOT / "data/corpus/pancreas/harmonized/corpus.h5ad")) OUT = Path(str(PANDA_ROOT / "discovery/pancreas/marker")) DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") N_FOLDS = 5 GUARANTEED_PER_CLASS = 6 NATURAL_SLOTS = 96 class CorpusDataset(Dataset): def __init__(self, X, y, d, mhf, logc): self.X = X.astype(np.float32); self.y = y.astype(np.int64) self.d = d.astype(np.int64); self.mhf = mhf.astype(np.float32) self.logc = logc.astype(np.float32) def __len__(self): return self.X.shape[0] def __getitem__(self, i): return (torch.from_numpy(self.X[i]), torch.tensor(self.y[i]), torch.tensor(self.d[i]), torch.tensor([self.mhf[i], self.logc[i]], dtype=torch.float32)) class HybridSampler: def __init__(self, y, d, n_batches=100, seed=0): self.y = np.asarray(y); self.d = np.asarray(d) self.n_batches = n_batches self.rng = np.random.default_rng(seed) self.classes = np.unique(self.y) self.by_cls = {c: np.where(self.y == c)[0] for c in self.classes} counts = np.bincount(self.y, minlength=int(self.classes.max())+1) self.p = counts[self.classes] / counts[self.classes].sum() def __iter__(self): for _ in range(self.n_batches): batch = [] for c in self.classes: idx = self.by_cls[c] take = min(GUARANTEED_PER_CLASS, len(idx)) if take: pick = self.rng.choice(idx, size=take, replace=(len(idx) < take)) batch.extend(pick.tolist()) for _ in range(NATURAL_SLOTS): c_pick = self.rng.choice(self.classes, p=self.p) batch.append(int(self.rng.choice(self.by_cls[c_pick]))) yield batch def __len__(self): return self.n_batches def train_one_fold(X, y, d, mhf, logc, classes, datasets, tr, te, fold_id, log_prefix): torch.cuda.empty_cache() Xtr, ytr, dtr, mtr, ltr = X[tr], y[tr], d[tr], mhf[tr], logc[tr] ds = CorpusDataset(Xtr, ytr, dtr, mtr, ltr) sampler = HybridSampler(ytr, dtr, n_batches=100, seed=fold_id) loader = DataLoader(ds, batch_sampler=sampler, num_workers=0) counts_per = np.bincount(ytr, minlength=len(classes)) inv_sqrt = 1.0 / np.sqrt(counts_per + 1) inv_sqrt = inv_sqrt / inv_sqrt.mean() class_w_np = 0.5 * inv_sqrt + 0.5 * np.ones_like(inv_sqrt) class_w = torch.tensor(class_w_np, dtype=torch.float32).to(DEVICE) model = PANDAEncoder(n_pca=X.shape[1], n_classes=len(classes), n_datasets=len(datasets)).to(DEVICE) opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4) stage_epochs = [15, 25, 40, 40] for stage in range(4): for e in range(stage_epochs[stage]): if e % 10 == 0: print(f"{log_prefix} s{stage} ep {e}/{stage_epochs[stage]}", flush=True) for X_b, y_b, d_b, aux_b in loader: X_b, y_b, d_b, aux_b = X_b.to(DEVICE), y_b.to(DEVICE), d_b.to(DEVICE), aux_b.to(DEVICE) if stage >= 2: jitter = torch.empty_like(aux_b[:, 1:2]).uniform_(-2, 0) aux_b = aux_b.clone(); aux_b[:, 1:2] = aux_b[:, 1:2] + jitter lam = 1.0 if stage >= 2 else 0.0 out = model(X_b, aux_b, lam_dann=lam) L_supcon = supcon_loss(out["z"], y_b) L_vic = vicreg_loss(out["z"]) L_ce = F.cross_entropy(out["logits"], y_b, weight=class_w, label_smoothing=0.05) total = L_supcon + 1.0 * L_vic + 0.4 * L_ce if stage >= 1: proto_ref = model.prototypes.detach().clone() L_p = prototype_infonce(out["z"], y_b, proto_ref) total = total + 0.6 * L_p if stage >= 2: L_d = F.cross_entropy(out["dom"], d_b) L_dep = F.mse_loss(out["depth"].squeeze(1), aux_b[:, 1]) L_h = hsic_biased(out["repr"], aux_b[:, 1:2]) total = total + L_d + 0.3 * L_dep + 0.05 * L_h opt.zero_grad(); total.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0) opt.step() if stage >= 1: model.update_prototypes(out["z"].detach(), y_b) model.eval() Xte, yte = X[te], y[te] with torch.no_grad(): Xt = torch.from_numpy(Xte.astype(np.float32)).to(DEVICE) aux = torch.zeros(len(te), 2, device=DEVICE) out = model(Xt, aux, lam_dann=0.0) z = out["z"] cos = z @ model.prototypes.T pred = cos.argmax(dim=1).cpu().numpy() probs = torch.softmax(cos / 0.07, dim=1).cpu().numpy() acc = accuracy_score(yte, pred) f1 = f1_score(yte, pred, average="macro", zero_division=0) try: auc = roc_auc_score(np.eye(len(classes))[yte], probs, average="macro", multi_class="ovr") except Exception: auc = float("nan") print(f"{log_prefix} acc={acc:.4f} macro_f1={f1:.4f} macro_auc={auc:.4f}", flush=True) return acc, f1, auc, pred, yte def main(): a = ad.read_h5ad(CORPUS) keep = (a.obs["canonical_label"].astype(str) != "UNK").values a = a[keep].copy() classes = sorted(a.obs["canonical_label"].astype(str).unique()) datasets = sorted(a.obs["dataset"].astype(str).unique()) c2i = {c: i for i, c in enumerate(classes)} d2i = {d: i for i, d in enumerate(datasets)} X = np.asarray(a.obsm["X_pca"]) y = np.array([c2i[c] for c in a.obs["canonical_label"].astype(str)]) d = np.array([d2i[dd] for dd in a.obs["dataset"].astype(str)]) mhf = a.obs.get("missing_hvg_frac", np.zeros(len(a))).astype(np.float32).values counts = a.obs["total_counts"].astype(float).values if "total_counts" in a.obs.columns \ else np.asarray(a.X.sum(axis=1)).ravel() logc = np.log10(counts + 1); logc = (logc - logc.mean()) / (logc.std() + 1e-6) print(f"[cv] corpus: {a.shape} n_classes={len(classes)} n_datasets={len(datasets)}", flush=True) print(f"[cv] class counts: {dict(zip(classes, np.bincount(y, minlength=len(classes)).tolist()))}", flush=True) skf = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=42) accs, f1s, aucs = [], [], [] all_preds = [] partial_path = OUT / "61_pancreas_heldout_5fold_partial.json" for fold, (tr, te) in enumerate(skf.split(X, y)): t0 = time.time() try: acc, f1, auc, pred, yte = train_one_fold( X, y, d, mhf, logc, classes, datasets, tr, te, fold_id=fold, log_prefix=f"[fold {fold}]") except Exception as exc: import traceback; traceback.print_exc() print(f"[fold {fold}] FAILED: {exc}", flush=True) continue print(f"[fold {fold}] wall={time.time()-t0:.0f}s", flush=True) accs.append(acc); f1s.append(f1); aucs.append(auc) all_preds.append({"fold": fold, "test_idx": te.tolist(), "pred": pred.tolist(), "true": yte.tolist()}) with open(partial_path, "w") as f: json.dump({"folds_done": fold + 1, "accs": accs, "f1s": f1s, "aucs": aucs}, f) print(f"\n[cv] 5-FOLD ACC: {np.mean(accs):.4f} +- {np.std(accs):.4f}") print(f"[cv] 5-FOLD F1: {np.mean(f1s):.4f} +- {np.std(f1s):.4f}") print(f"[cv] 5-FOLD AUC: {np.mean(aucs):.4f} +- {np.std(aucs):.4f}") all_y_true = np.concatenate([np.array(p["true"]) for p in all_preds]) all_y_pred = np.concatenate([np.array(p["pred"]) for p in all_preds]) print("\n[cv] Concatenated held-out classification report:") rep = classification_report(all_y_true, all_y_pred, target_names=classes, digits=3, zero_division=0, output_dict=True) print(classification_report(all_y_true, all_y_pred, target_names=classes, digits=3, zero_division=0)) result = { "mean_acc": float(np.mean(accs)), "std_acc": float(np.std(accs)), "mean_f1": float(np.mean(f1s)), "std_f1": float(np.std(f1s)), "mean_auc": float(np.mean(aucs)), "std_auc": float(np.std(aucs)), "per_fold_acc": accs, "per_fold_f1": f1s, "per_fold_auc": aucs, "per_class_report": rep, "n_folds": N_FOLDS, "n_classes": len(classes), "protocol": "StratifiedKFold(5) retrain from scratch per fold; " "eval by prototype-cosine argmax on held-out 20%.", } (OUT / "61_pancreas_heldout_5fold_cv.json").write_text(json.dumps(result, indent=2)) print(f"\n[cv] wrote {OUT}/61_pancreas_heldout_5fold_cv.json") if __name__ == "__main__": main()