"""train panda on the canonical paper-labeled corpus for one system + variant.""" from __future__ import annotations import argparse, sys, json, pickle, warnings, numpy as np, pandas as pd, torch, torch.nn.functional as F from pathlib import Path import anndata as ad, scanpy as sc, scipy.sparse as sp, yaml warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 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 import ( PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, subcenter_angular_infonce, prototype_repulsion, ) ROOT = Path(str(PANDA_ROOT)) DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") def load_corpus(system): """canonical loader; alias for load_corpus_v3 after finalize_rename.""" return load_corpus_v3(system) def load_corpus_v3(system): # kept for backward-compat with older scripts that import load_corpus_v3 p = ROOT / f"data/corpus/{system}/harmonized/corpus.h5ad" stats = np.load(ROOT / f"data/corpus/{system}/harmonized/corpus_stats.npz", allow_pickle=True) pca = pickle.load(open(ROOT / f"data/corpus/{system}/harmonized/pca_basis.pkl", "rb")) a = ad.read_h5ad(p) hvgs = [str(g) for g in stats["shared_hvgs"]] return a, hvgs, stats["mean"], stats["std"], pca def get_marker_gene_list(system): y = yaml.safe_load(open(ROOT / "panda/markers.yaml")) return y[system] def prepare_batches(adata, hvgs, mu, sig, pca, marker_genes=None, variant="pca", legacy_double_norm=True): # corpus.h5ad X is already normalize_total+log1p'd by the corpus builders, and the # published checkpoints were trained with a second normalize/log1p applied on top # (against per-gene stats computed from singly-normalized data). legacy_double_norm=True # reproduces that behaviour bit-for-bit; pass False to train on the corpus values the # stats were actually computed from. Do not mix: a checkpoint must be evaluated under # the same setting it was trained with. hvg2i = {g: i for i, g in enumerate(hvgs)} common = [g for g in adata.var_names.astype(str) if g in hvg2i] a_c = adata[:, common].copy() if legacy_double_norm: sc.pp.normalize_total(a_c, target_sum=1e4); sc.pp.log1p(a_c) X_ = a_c.X.toarray().astype(np.float32) if sp.issparse(a_c.X) else a_c.X.astype(np.float32) Xf = np.zeros((adata.n_obs, len(hvgs)), dtype=np.float32) cols = np.array([hvg2i[g] for g in common]) Xf[:, cols] = X_ Xz = np.clip((Xf - mu.astype(np.float32)) / sig.astype(np.float32), -10, 10) Xpca = pca.transform(Xz).astype(np.float32) Xmark = None if variant == "marker" and marker_genes: mvals = np.zeros((adata.n_obs, len(marker_genes)), dtype=np.float32) for j, g in enumerate(marker_genes): if g in adata.var_names: col = adata[:, g].X if sp.issparse(col): col = col.toarray() mvals[:, j] = col.flatten().astype(np.float32) mmu = mvals.mean(axis=0, keepdims=True); msig = mvals.std(axis=0, keepdims=True) + 1e-6 Xmark = np.clip((mvals - mmu) / msig, -5, 5).astype(np.float32) # expose the training-corpus marker stats so the checkpoint can carry them; # inference must reuse these rather than refit on the target dataset prepare_batches.last_marker_stats = (mmu, msig) labels = adata.obs["canonical_label"].astype(str).values classes = sorted(set(labels)) y = np.array([classes.index(l) for l in labels], dtype=np.int64) datasets = sorted(set(adata.obs["dataset"].astype(str).values)) y_dset = np.array([datasets.index(d) for d in adata.obs["dataset"].astype(str).values], dtype=np.int64) counts = np.asarray(adata.X.sum(axis=1)).ravel() log10cz = ((np.log10(counts + 1) - np.log10(counts + 1).mean()) / (np.log10(counts + 1).std() + 1e-6)).astype(np.float32) return Xpca, Xmark, y, classes, y_dset, datasets, log10cz def train(system, variant, epochs=8, batch=256, lr=1e-3): a, hvgs, mu, sig, pca = load_corpus_v3(system) marker_genes = get_marker_gene_list(system) if variant == "marker" else [] Xpca, Xmark, y, classes, y_dset, datasets, log10cz = prepare_batches( a, hvgs, mu, sig, pca, marker_genes, variant ) print(f"[train] {system}/{variant} n={a.n_obs} K={len(classes)} datasets={len(datasets)}", flush=True) print(f"[train] classes: {classes}", flush=True) n_markers = Xmark.shape[1] if Xmark is not None else 0 model = PANDAEncoder(variant=variant, n_pca=50, n_markers=n_markers, n_classes=len(classes), n_sub=3, n_datasets=len(datasets), dropout=0.2).to(DEVICE) opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4) rng = np.random.default_rng(0) for epoch in range(epochs): stage = 0 if epoch < 1 else 1 if epoch < 3 else 2 if epoch < 6 else 3 for g in opt.param_groups: g["lr"] = lr * (0.5 if epoch >= epochs - 1 else 1.0) perm = rng.permutation(a.n_obs) losses = [] for bstart in range(0, a.n_obs, batch): idx = perm[bstart:bstart+batch] x = torch.from_numpy(Xpca[idx]).to(DEVICE) xm = torch.from_numpy(Xmark[idx]).to(DEVICE) if Xmark is not None else None yy = torch.from_numpy(y[idx]).to(DEVICE) yd = torch.from_numpy(y_dset[idx]).to(DEVICE) dd = torch.from_numpy(log10cz[idx]).float().to(DEVICE).unsqueeze(1) aux = torch.zeros(len(idx), 2, device=DEVICE) lam = 0.1 if stage >= 2 else 0.0 out = model(x, aux, x_markers=xm, lam_dann=lam) z = out["z"] L = supcon_loss(z, yy, 0.1) + 1.0 * vicreg_loss(z) + 0.4 * F.cross_entropy(out["logits"], yy) if stage >= 1: L = L + 0.6 * subcenter_angular_infonce(z, yy, model.prototypes.detach().clone(), margin=0.15, temperature=0.07) if stage >= 2: L = L + F.cross_entropy(out["dom"], yd) + 0.3 * F.mse_loss(out["depth"], dd) + 0.05 * hsic_biased(out["repr"], dd) # NOTE: earlier revisions added 0.5 * prototype_repulsion(model.prototypes.detach()) # at stage 3. prototypes is a gradient-free EMA buffer and the tensor was detached, # so the term contributed exactly zero gradient — it only inflated the printed loss. # Removed from the objective (behaviour-preserving); prototype_repulsion() remains # in panda.model as an analysis metric. opt.zero_grad(); L.backward(); opt.step() if stage >= 1: with torch.no_grad(): model.update_prototypes(z.detach(), yy) losses.append(float(L)) print(f"[train {system}/{variant}] epoch {epoch}/{epochs} stage={stage} loss={np.mean(losses):.4f}", flush=True) # save to the same path every consumer (zero_shot, run_all_zero_shot, extract_prototypes) # loads from. Earlier revisions wrote to checkpoints/{system}_v3/ while consumers read # checkpoints/{system}/ — a retrain silently never propagated. ck_dir = ROOT / f"checkpoints/{system}/{variant}" ck_dir.mkdir(parents=True, exist_ok=True) mstats = getattr(prepare_batches, "last_marker_stats", None) if variant == "marker" else None torch.save({"model": model.state_dict(), "classes": classes, "datasets": datasets, "marker_genes": marker_genes if variant == "marker" else [], "marker_mu": mstats[0] if mstats else None, "marker_sig": mstats[1] if mstats else None, "legacy_double_norm": True, "prototypes": model.prototypes.detach().cpu().numpy()}, ck_dir / "panda_final.pt") print(f"[save] {ck_dir}/panda_final.pt", flush=True) if __name__ == "__main__": ap = argparse.ArgumentParser() ap.add_argument("system", choices=["pan_skin", "hematopoiesis", "pancreas"]) ap.add_argument("--variant", choices=["pca", "marker"], required=True) ap.add_argument("--epochs", type=int, default=8) ap.add_argument("--batch", type=int, default=256) ap.add_argument("--lr", type=float, default=1e-3) args = ap.parse_args() train(args.system, args.variant, args.epochs, args.batch, args.lr)