| """v3c PANDA-MLP training. hybrid sampler: guaranteed per-class + natural-freq fill."""
|
| 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 torch
|
| import 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/pan_skin/harmonized/corpus.h5ad"))
|
| OUT = Path(str(PANDA_ROOT / "checkpoints/pan_skin"))
|
| V2_CK = Path(str(PANDA_ROOT / "checkpoints/pan_skin_v2_kept/panda_final.pt"))
|
| OUT.mkdir(parents=True, exist_ok=True)
|
|
|
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
| BALANCE_MIX = 0.5
|
| GUARANTEED_PER_CLASS = 6
|
| NATURAL_SLOTS = 96
|
|
|
|
|
| class CorpusDataset(Dataset):
|
| def __init__(self, X, y, d, mhf, logc):
|
| self.X = X.astype(np.float32); self.y = y.astype(np.int64)
|
| self.d = d.astype(np.int64); self.mhf = mhf.astype(np.float32)
|
| self.logc = 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 = np.asarray(y); self.d = 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}
|
| self.class_counts = np.bincount(self.y, minlength=len(self.classes))
|
| p = self.class_counts / self.class_counts.sum()
|
| self.natural_p = p
|
| 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 > 0:
|
| pick = self.rng.choice(idx, size=take, replace=(len(idx) < take))
|
| batch.extend(pick.tolist())
|
| n_extra = NATURAL_SLOTS
|
| for _ in range(n_extra):
|
| c_pick = self.rng.choice(self.classes, p=self.natural_p)
|
| idx = self.by_cls[c_pick]
|
| batch.append(int(self.rng.choice(idx)))
|
| yield batch
|
| def __len__(self): return self.n_batches
|
|
|
|
|
| def rankme(z):
|
| with torch.no_grad():
|
| _, s, _ = torch.svd(z - z.mean(0, keepdim=True))
|
| p = s / s.sum().clamp_min(1e-12)
|
| H = -(p * (p + 1e-12).log()).sum()
|
| return float(H.exp().item())
|
|
|
|
|
| def main():
|
| print(f"[train] device: {DEVICE}", flush=True)
|
| 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())
|
| cls_ix = {c: i for i, c in enumerate(classes)}
|
| ds_ix = {d: i for i, d in enumerate(datasets)}
|
| y = np.array([cls_ix[c] for c in a.obs["canonical_label"].astype(str)])
|
| d = np.array([ds_ix[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
|
| if "total_counts" in a.obs.columns:
|
| counts = a.obs["total_counts"].astype(float).values
|
| else:
|
| counts = np.asarray(a.X.sum(axis=1)).ravel()
|
| logc = np.log10(counts + 1); logc = (logc - logc.mean()) / (logc.std() + 1e-6)
|
|
|
| with open(OUT / "label_encoding.json", "w") as f:
|
| json.dump({"classes": classes, "datasets": datasets}, f, indent=2)
|
|
|
| counts_per = np.bincount(y, minlength=len(classes))
|
| print(f"[train] shape {a.shape}, n_classes {len(classes)}, n_datasets {len(datasets)}",
|
| flush=True)
|
| print(f"[train] class counts: {dict(zip(classes, counts_per.tolist()))}", flush=True)
|
|
|
| inv_sqrt = 1.0 / np.sqrt(counts_per + 1)
|
| inv_sqrt = inv_sqrt / inv_sqrt.mean()
|
| class_w = BALANCE_MIX * inv_sqrt + (1 - BALANCE_MIX) * np.ones_like(inv_sqrt)
|
| class_w = torch.tensor(class_w, dtype=torch.float32).to(DEVICE)
|
|
|
| ds = CorpusDataset(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)
|
|
|
| if V2_CK.exists():
|
| try:
|
| ck = torch.load(V2_CK, map_location=DEVICE, weights_only=False)
|
| if "classes" in ck:
|
|
|
| idx_v2 = {c: i for i, c in enumerate(ck["classes"])}
|
| for c in classes:
|
| if c in idx_v2:
|
| model.prototypes[cls_ix[c]].copy_(
|
| torch.tensor(ck["prototypes"][idx_v2[c]]).to(DEVICE))
|
| print(f"[warmstart] copied {sum(1 for c in classes if c in idx_v2)}/{len(classes)}"
|
| f" prototypes from v2", flush=True)
|
| except Exception as exc:
|
| print(f"[warmstart] skipped ({exc})")
|
|
|
| opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
|
| stage_epochs = [15, 25, 40, 40]
|
|
|
| def train_one_epoch(stage):
|
| L = {"supcon": [], "vic": [], "ce": [], "proto": [], "dom": [], "depth": [], "hsic": []}
|
| 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_supcon = 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_supcon + 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
|
| L["proto"].append(float(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
|
| L["dom"].append(float(L_d)); L["depth"].append(float(L_dep)); L["hsic"].append(float(L_h))
|
| L["supcon"].append(float(L_supcon)); L["vic"].append(float(L_vic)); L["ce"].append(float(L_ce))
|
| 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)
|
| return {k: float(np.mean(v)) if v else 0.0 for k, v in L.items()}
|
|
|
| global_ep = 0
|
| for stage in range(4):
|
| n = stage_epochs[stage]
|
| print(f"\n=== stage {stage} ({n} epochs) ===", flush=True)
|
| for e in range(n):
|
| t0 = time.time()
|
| m = train_one_epoch(stage)
|
| global_ep += 1
|
| if global_ep % 3 == 0 or e == n - 1:
|
| X_b, *_ = next(iter(loader))
|
| with torch.no_grad():
|
| z_b = model(X_b.to(DEVICE),
|
| torch.zeros(len(X_b), 2, device=DEVICE))["z"]
|
| rk = rankme(z_b)
|
| print(f"[s{stage}][ep {global_ep}] sup={m['supcon']:.3f} vic={m['vic']:.3f} "
|
| f"ce={m['ce']:.3f} pro={m['proto']:.3f} dom={m['dom']:.3f} "
|
| f"depth={m['depth']:.3f} hsic={m['hsic']:.3f} | "
|
| f"RankMe={rk:.1f} 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"\n[done] saved {OUT}/panda_final.pt", flush=True)
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
