"""train panda on pan-hsc corpus.""" from __future__ import annotations from pathlib import Path import warnings, json, sys, time warnings.filterwarnings("ignore") import numpy as np, anndata as ad, torch, torch.nn.functional as F from torch.utils.data import Dataset, DataLoader 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/hematopoiesis/harmonized/corpus.h5ad")) OUT = Path(str(PANDA_ROOT / "checkpoints/hematopoiesis")) OUT.mkdir(parents=True, exist_ok=True) DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") GUARANTEED_PER_CLASS = 6 NATURAL_SLOTS = 96 class Ds(Dataset): def __init__(self, X, y, d, mhf, logc): self.X, self.y, self.d = X.astype(np.float32), y.astype(np.int64), d.astype(np.int64) self.mhf, self.logc = mhf.astype(np.float32), 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, self.d = np.asarray(y), 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 = self.rng.choice(self.classes, p=self.p) batch.append(int(self.rng.choice(self.by_cls[c]))) yield batch def __len__(self): return self.n_batches 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)} 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)]) X = np.asarray(a.obsm["X_pca"]) 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) counts_per = np.bincount(y, minlength=len(classes)) print(f"[train] {a.shape}, n_classes={len(classes)}, " f"class_counts={dict(zip(classes, counts_per.tolist()))}", flush=True) with open(OUT / "label_encoding.json", "w") as f: json.dump({"classes": classes, "datasets": datasets}, f, indent=2) inv_sqrt = 1.0 / np.sqrt(counts_per + 1); inv_sqrt = inv_sqrt / inv_sqrt.mean() class_w = torch.tensor(0.5 * inv_sqrt + 0.5 * np.ones_like(inv_sqrt), dtype=torch.float32).to(DEVICE) ds = Ds(X, y, d, mhf, logc) sampler = HybridSampler(y, d, n_batches=100) loader = DataLoader(ds, batch_sampler=sampler, num_workers=0) 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): n_ep = stage_epochs[stage] print(f"\n=== stage {stage} ({n_ep}) ===", flush=True) for e in range(n_ep): t0 = time.time() 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_sup = 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_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 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) if e % 5 == 0: print(f"[s{stage}][ep {e}] dt={time.time()-t0:.1f}s", flush=True) torch.save({"model": model.state_dict(), "classes": classes, "datasets": datasets}, OUT / f"panda_stage{stage}.pt") torch.save({"model": model.state_dict(), "classes": classes, "datasets": datasets, "prototypes": model.prototypes.detach().cpu().numpy()}, OUT / "panda_final.pt") np.save(OUT / "prototypes.npy", model.prototypes.detach().cpu().numpy()) print(f"[done] saved {OUT}/panda_final.pt", flush=True) if __name__ == "__main__": main()