"""sulic-in-panda held-out 5-fold cv. Test A: binary facs (placode vs epi). Test C: 4-way placode subtype.""" from __future__ import annotations from pathlib import Path import warnings, json, sys, time warnings.filterwarnings("ignore") import numpy as np import anndata as ad import scanpy as sc import scipy.sparse as sp import torch import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader from sklearn.decomposition import PCA from sklearn.model_selection import StratifiedKFold from sklearn.preprocessing import StandardScaler from sklearn.metrics import roc_auc_score, accuracy_score 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, prototype_infonce ) SULIC_H5AD = Path(str(PANDA_ROOT / "data/processed/sulic/adata_sulic_clustered.h5ad")) OUT = Path(str(PANDA_ROOT / "scripts/sulic")) OUT.mkdir(parents=True, exist_ok=True) DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") N_FOLDS = 5 class SulicDataset(Dataset): def __init__(self, X, y): self.X = X.astype(np.float32); self.y = y.astype(np.int64) 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.zeros(1, dtype=torch.int64), # single-dataset torch.zeros(2, dtype=torch.float32)) # no aux class PxKSampler: def __init__(self, y, P=None, K=16, n_batches=80, seed=0): self.y = np.asarray(y) self.classes = np.unique(self.y) self.P = P or len(self.classes) self.K = K self.n_batches = n_batches self.rng = np.random.default_rng(seed) self.by_cls = {c: np.where(self.y == c)[0] for c in self.classes} def __iter__(self): for _ in range(self.n_batches): classes_p = self.rng.choice(self.classes, size=min(self.P, len(self.classes)), replace=False) batch = [] for c in classes_p: idx = self.by_cls[c] take = self.K pick = self.rng.choice(idx, size=take, replace=(len(idx) < take)) batch.extend(pick.tolist()) yield batch def __len__(self): return self.n_batches def prepare_pca(a, n_pca=50): if a.raw is not None: a = a.raw.to_adata() sc.pp.normalize_total(a, target_sum=1e4) sc.pp.log1p(a) sc.pp.highly_variable_genes(a, n_top_genes=2000, flavor="seurat", subset=False) a = a[:, a.var["highly_variable"]].copy() X = a.X.toarray() if sp.issparse(a.X) else a.X scaler = StandardScaler().fit(X) Xz = np.clip(scaler.transform(X), -10, 10) pca = PCA(n_components=n_pca, random_state=42).fit(Xz) Xp = pca.transform(Xz).astype(np.float32) return a, Xp def train_fold(Xp, y, classes, tr, te, fold_id, ensemble_seeds=5): K = len(classes) all_probs = [] for seed in range(ensemble_seeds): torch.manual_seed(fold_id * 100 + seed) np.random.seed(fold_id * 100 + seed) torch.cuda.empty_cache() model = PANDAEncoder(n_pca=Xp.shape[1], n_classes=K, n_datasets=1).to(DEVICE) opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4) ds = SulicDataset(Xp[tr], y[tr]) sampler = PxKSampler(y[tr], K=16, n_batches=80, seed=seed) loader = DataLoader(ds, batch_sampler=sampler, num_workers=0) for stage, ne in enumerate([15, 20, 25]): for e in range(ne): for X_b, y_b, _, aux_b in loader: X_b, y_b = X_b.to(DEVICE), y_b.to(DEVICE) aux_b = aux_b.to(DEVICE) out = model(X_b, aux_b, lam_dann=0.0) L_sup = supcon_loss(out["z"], y_b) L_vic = vicreg_loss(out["z"]) L_ce = F.cross_entropy(out["logits"], y_b, label_smoothing=0.05) total = L_sup + 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 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() with torch.no_grad(): Xt = torch.from_numpy(Xp[te]).to(DEVICE) aux = torch.zeros(len(te), 2, device=DEVICE) out = model(Xt, aux, lam_dann=0.0) cos = out["z"] @ model.prototypes.T probs = torch.softmax(cos / 0.07, dim=1).cpu().numpy() all_probs.append(probs) ensemble_probs = np.mean(all_probs, axis=0) pred = ensemble_probs.argmax(axis=1) yte = y[te] return pred, ensemble_probs, yte def evaluate_test(name, X, y_bin_or_multi, class_list): print(f"\n=== {name} ===", flush=True) skf = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=42) aurocs, accs = [], [] for fold, (tr, te) in enumerate(skf.split(X, y_bin_or_multi)): t0 = time.time() pred, probs, yte = train_fold(X, y_bin_or_multi, class_list, tr, te, fold) acc = accuracy_score(yte, pred) if len(class_list) == 2: auc = roc_auc_score(yte, probs[:, 1]) else: try: auc = roc_auc_score(np.eye(len(class_list))[yte], probs, average="macro", multi_class="ovr") except Exception: auc = float("nan") print(f"[{name} fold {fold}] acc={acc:.4f} AUROC={auc:.4f} " f"wall={time.time()-t0:.0f}s", flush=True) aurocs.append(auc); accs.append(acc) print(f"[{name}] MEAN AUROC = {np.mean(aurocs):.4f} +- {np.std(aurocs):.4f}", flush=True) print(f"[{name}] MEAN ACC = {np.mean(accs):.4f} +- {np.std(accs):.4f}", flush=True) return {"aurocs": aurocs, "accs": accs, "mean_auroc": float(np.mean(aurocs)), "std_auroc": float(np.std(aurocs)), "mean_acc": float(np.mean(accs)), "std_acc": float(np.std(accs))} def main(): print(f"[sulic-panda] loading {SULIC_H5AD}", flush=True) a = ad.read_h5ad(SULIC_H5AD) print(f"[sulic-panda] shape {a.shape}, samples: {a.obs['sample'].value_counts().to_dict()}", flush=True) a_p, Xp = prepare_pca(a, n_pca=50) print(f"[sulic-panda] Xp {Xp.shape}", flush=True) y_A = (a.obs["sample"].isin(["Placode1", "Placode2"])).astype(int).values print(f"[sulic-panda] Test A class balance: {np.bincount(y_A).tolist()}", flush=True) resA = evaluate_test("TestA", Xp, y_A, ["Epithelium", "Placode"]) if "placode_enriched" in a.obs.columns: mask_p = (a.obs["placode_enriched"] == 1).values sub = a[mask_p].copy() if "paper_subtype" not in sub.obs.columns: # fallback: kmeans on placode-cell Xp gives 4 pseudo-subtypes print("[sulic-panda] paper_subtype missing — deriving 4-way clustering on Xp", flush=True) from sklearn.cluster import KMeans Xp_sub = Xp[mask_p] km = KMeans(n_clusters=4, random_state=42, n_init=10).fit(Xp_sub) paper_subtype = np.array([f"PlacodeK{i}" for i in km.labels_]) else: paper_subtype = sub.obs["paper_subtype"].astype(str).values cls = sorted(np.unique(paper_subtype)) y_C = np.array([cls.index(v) for v in paper_subtype], dtype=np.int64) Xp_C = Xp[mask_p] print(f"[sulic-panda] Test C n={len(y_C)}, classes={cls}, " f"counts={np.bincount(y_C).tolist()}", flush=True) resC = evaluate_test("TestC", Xp_C, y_C, cls) else: resC = None # save result = {"testA": resA, "testC": resC} with open(OUT / "sulic_panda_heldout_results.json", "w") as f: json.dump(result, f, indent=2) print(f"\n[sulic-panda] wrote {OUT}/sulic_panda_heldout_results.json", flush=True) if __name__ == "__main__": main()