| """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() | |