| """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),
|
| torch.zeros(2, dtype=torch.float32))
|
|
|
|
|
| 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:
|
|
|
| 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
|
|
|
|
|
| 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()
|
|
|
