"""5-fold stratified CV, 3 systems x 2 variants. 5 epochs/fold (shorter than train_panda).""" from __future__ import annotations import argparse, sys, json, pickle, warnings, numpy as np, pandas as pd, torch, torch.nn.functional as F from pathlib import Path import anndata as ad, scanpy as sc, scipy.sparse as sp, yaml warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 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 import (PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, subcenter_angular_infonce, prototype_repulsion) ROOT = Path(str(PANDA_ROOT)) DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") def load_and_prepare(system, variant): """canonical corpus + features (PCA, optional marker channel).""" a = ad.read_h5ad(ROOT / f"data/corpus/{system}/harmonized/corpus.h5ad") stats = np.load(ROOT / f"data/corpus/{system}/harmonized/corpus_stats.npz", allow_pickle=True) pca = pickle.load(open(ROOT / f"data/corpus/{system}/harmonized/pca_basis.pkl", "rb")) hvgs = [str(g) for g in stats["shared_hvgs"]] hvg2i = {g: i for i, g in enumerate(hvgs)} common = [g for g in a.var_names.astype(str) if g in hvg2i] a_c = a[:, common].copy() sc.pp.normalize_total(a_c, target_sum=1e4); sc.pp.log1p(a_c) X = a_c.X.toarray().astype(np.float32) if sp.issparse(a_c.X) else a_c.X.astype(np.float32) Xf = np.zeros((a.n_obs, len(hvgs)), dtype=np.float32) Xf[:, np.array([hvg2i[g] for g in common])] = X Xz = np.clip((Xf - stats["mean"].astype(np.float32)) / stats["std"].astype(np.float32), -10, 10) Xpca = pca.transform(Xz).astype(np.float32) Xmark = None; marker_genes = [] if variant == "marker": marker_genes = yaml.safe_load(open(ROOT / "panda/markers.yaml"))[system] mv = np.zeros((a.n_obs, len(marker_genes)), dtype=np.float32) for j, g in enumerate(marker_genes): if g in a.var_names: col = a[:, g].X if sp.issparse(col): col = col.toarray() mv[:, j] = col.flatten().astype(np.float32) mmu = mv.mean(axis=0, keepdims=True); msig = mv.std(axis=0, keepdims=True) + 1e-6 Xmark = np.clip((mv - mmu) / msig, -5, 5).astype(np.float32) labels = a.obs["canonical_label"].astype(str).values classes = sorted(set(labels)) y = np.array([classes.index(l) for l in labels], dtype=np.int64) datasets = sorted(set(a.obs["dataset"].astype(str).values)) y_dset = np.array([datasets.index(d) for d in a.obs["dataset"].astype(str).values], dtype=np.int64) return Xpca, Xmark, y, classes, y_dset, datasets, marker_genes def train_fold(Xpca, Xmark, y, y_dset, classes, variant, tr_ix, epochs=5, batch=256, lr=1e-3, seed=0): n_classes = len(classes) n_datasets = int(max(y_dset[tr_ix].max() + 1, 1)) n_markers = Xmark.shape[1] if Xmark is not None else 0 torch.manual_seed(seed); np.random.seed(seed) model = PANDAEncoder(variant=variant, n_pca=50, n_markers=n_markers, n_classes=n_classes, n_sub=3, n_datasets=n_datasets, dropout=0.2).to(DEVICE) opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4) Xt = Xpca[tr_ix]; yt = y[tr_ix]; ydt = y_dset[tr_ix] Xmt = Xmark[tr_ix] if Xmark is not None else None rng = np.random.default_rng(seed) n = len(tr_ix) for epoch in range(epochs): stage = 0 if epoch < 1 else 1 if epoch < 3 else 2 perm = rng.permutation(n) for bstart in range(0, n, batch): idx = perm[bstart:bstart+batch] x = torch.from_numpy(Xt[idx]).to(DEVICE) xm = torch.from_numpy(Xmt[idx]).to(DEVICE) if Xmt is not None else None yy = torch.from_numpy(yt[idx]).to(DEVICE) yd = torch.from_numpy(ydt[idx]).to(DEVICE) aux = torch.zeros(len(idx), 2, device=DEVICE) lam = 0.1 if stage >= 2 else 0.0 out = model(x, aux, x_markers=xm, lam_dann=lam) z = out["z"] L = supcon_loss(z, yy, 0.1) + 1.0 * vicreg_loss(z) + 0.4 * F.cross_entropy(out["logits"], yy) if stage >= 1: L = L + 0.6 * subcenter_angular_infonce(z, yy, model.prototypes.detach().clone(), margin=0.15, temperature=0.07) if stage >= 2: L = L + F.cross_entropy(out["dom"], yd) opt.zero_grad(); L.backward(); opt.step() if stage >= 1: with torch.no_grad(): model.update_prototypes(z.detach(), yy) return model def evaluate(model, Xpca, Xmark, y, val_ix, classes): model.eval() preds, probs = [], [] Xv = Xpca[val_ix]; Xmv = Xmark[val_ix] if Xmark is not None else None with torch.no_grad(): for i in range(0, len(val_ix), 2048): xb = torch.from_numpy(Xv[i:i+2048]).to(DEVICE) xmb = torch.from_numpy(Xmv[i:i+2048]).to(DEVICE) if Xmv is not None else None aux = torch.zeros(len(xb), 2, device=DEVICE) out = model(xb, aux, x_markers=xmb, lam_dann=0.0) z = out["z"] mc = model.max_sub_cos(z) preds.append(mc.argmax(dim=1).cpu().numpy()) probs.append(F.softmax(mc / 0.07, dim=1).cpu().numpy()) preds = np.concatenate(preds); probs = np.concatenate(probs) yv = y[val_ix] acc = accuracy_score(yv, preds) f1 = f1_score(yv, preds, average="macro", zero_division=0) try: auc = roc_auc_score(np.eye(len(classes))[yv], probs, average="macro", multi_class="ovr") except Exception: auc = float("nan") rep = classification_report(yv, preds, labels=list(range(len(classes))), target_names=classes, output_dict=True, zero_division=0) return acc, f1, auc, rep def cv(system, variant, folds=5, epochs=5, seed=0): print(f"\n=== CV {system}/{variant} ({folds}-fold, {epochs} epochs) ===", flush=True) Xpca, Xmark, y, classes, y_dset, datasets, _ = load_and_prepare(system, variant) print(f"[cv] n={len(y):,} K={len(classes)}", flush=True) skf = StratifiedKFold(n_splits=folds, shuffle=True, random_state=seed) accs, f1s, aucs = [], [], [] last_rep = None for fold, (tr, va) in enumerate(skf.split(np.zeros(len(y)), y)): model = train_fold(Xpca, Xmark, y, y_dset, classes, variant, tr, epochs=epochs, seed=seed * 100 + fold) acc, f1, auc, rep = evaluate(model, Xpca, Xmark, y, va, classes) accs.append(acc); f1s.append(f1); aucs.append(auc) last_rep = rep print(f"[fold {fold+1}] acc={acc:.4f} F1={f1:.4f} AUC={auc:.4f}", flush=True) result = { "system": system, "variant": variant, "folds": folds, "epochs": epochs, "seed": seed, "n_cells": int(len(y)), "n_classes": len(classes), "per_class_report_note": "per_class_report is from the LAST fold only, not aggregated", "per_fold_acc": accs, "per_fold_f1": f1s, "per_fold_auc": aucs, "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.nanmean(aucs)), "std_auc": float(np.nanstd(aucs)), "per_class_report": last_rep, } out_dir = ROOT / f"discovery/{system}/{variant}" out_dir.mkdir(parents=True, exist_ok=True) # seed 0 is the canonical file; other seeds get a suffix so true seed-replicates # (same corpus, same script, same curriculum) can be compared side by side fname = "cv_5fold.json" if seed == 0 else f"cv_5fold_seed{seed}.json" (out_dir / fname).write_text(json.dumps(result, indent=2, default=str)) print(f"[cv] mean acc={result['mean_acc']:.4f}±{result['std_acc']:.4f} " f"F1={result['mean_f1']:.4f} AUC={result['mean_auc']:.4f}", flush=True) if __name__ == "__main__": ap = argparse.ArgumentParser() ap.add_argument("--systems", nargs="*", default=["pan_skin", "hematopoiesis", "pancreas"]) ap.add_argument("--variants", nargs="*", default=["pca", "marker"]) ap.add_argument("--folds", type=int, default=5) ap.add_argument("--epochs", type=int, default=5) ap.add_argument("--seed", type=int, default=0, help="fold-assignment + model-init seed; non-zero seeds write cv_5fold_seed{N}.json") args = ap.parse_args() for s in args.systems: for v in args.variants: cv(s, v, args.folds, args.epochs, seed=args.seed)