| """5-fold CV on pan-skin corpus. retrain 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/pan_skin/harmonized/corpus.h5ad")) | |
| OUT = Path(str(PANDA_ROOT / "discovery/pan_skin/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 / "60_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 / "60_heldout_5fold_cv.json").write_text(json.dumps(result, indent=2)) | |
| print(f"\n[cv] wrote {OUT}/60_heldout_5fold_cv.json") | |
| if __name__ == "__main__": | |
| main() | |