| """variant A — semi-supervised panda on dingwall using data s3 marker panels with score+margin gate."""
|
| from __future__ import annotations
|
| from pathlib import Path
|
| import warnings, json, sys, time
|
| warnings.filterwarnings("ignore")
|
|
|
| import numpy as np
|
| import pandas as pd
|
| import anndata as ad
|
| import scanpy as sc
|
| import scipy.sparse as sp
|
| from scipy.stats import fisher_exact
|
| 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.model import (
|
| PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, subcenter_angular_infonce
|
| )
|
|
|
| ROOT = Path(str(PANDA_ROOT))
|
| RAW_H5 = ROOT / "data/raw/GSE220977_combined.h5ad"
|
| DERM_MARKERS = ROOT / "data/external_labels/dingwall_supp/biorxiv_media-3.xlsx"
|
| OUT_DIR = ROOT / "discovery/pan_skin/marker"
|
| CK_DIR = ROOT / "checkpoints/pan_skin_dingwall_variantA"
|
|
|
|
|
| CKO_GSMS = {"GSM6833482", "GSM6833483"}
|
| WT_GSMS = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"}
|
|
|
| TOP_N = 30
|
| SCORE_MIN = 0.10
|
| MARGIN_MIN = 0.05
|
| N_HVG = 2000
|
| N_PCA = 40
|
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
|
| GUARANTEED_PER_CLASS = 6
|
| NATURAL_SLOTS = 96
|
| STAGE_EPOCHS = [15, 25, 40, 40]
|
| BALANCE_MIX = 0.5
|
|
|
|
|
|
|
|
|
| def load_derm_panels(top_n: int = TOP_N) -> dict[int, list[str]]:
|
| df = pd.read_excel(DERM_MARKERS)
|
| df = df.sort_values(["cluster", "avg_log2FC"], ascending=[True, False])
|
| return {int(c): df[df["cluster"] == c].head(top_n)["gene"].tolist()
|
| for c in sorted(df["cluster"].unique())}
|
|
|
|
|
| def load_dingwall_dermal() -> ad.AnnData:
|
| a = ad.read_h5ad(RAW_H5)
|
| sample = a.obs["sample"].astype(str)
|
| a.obs["genotype"] = np.where(sample.isin(list(CKO_GSMS)), "En1-cKO",
|
| np.where(sample.isin(list(WT_GSMS)), "WT", "other"))
|
| a = a[a.obs["genotype"].isin(["WT", "En1-cKO"])].copy()
|
| return a
|
|
|
|
|
| def preprocess_paper_style(a: ad.AnnData) -> ad.AnnData:
|
| """lognormalize + hvg(2000) + pca(40) + harmony per-sample, matches dingwall STAR methods."""
|
| sc.pp.filter_genes(a, min_cells=10)
|
| sc.pp.normalize_total(a, target_sum=1e4)
|
| sc.pp.log1p(a)
|
| sc.pp.highly_variable_genes(a, n_top_genes=N_HVG, flavor="seurat", batch_key="sample")
|
| sc.pp.scale(a, max_value=10, zero_center=False)
|
| sc.tl.pca(a, n_comps=N_PCA, use_highly_variable=True, zero_center=False)
|
| try:
|
| import harmonypy as hm
|
| sc.external.pp.harmony_integrate(a, key="sample", basis="X_pca",
|
| adjusted_basis="X_pca_harmony", max_iter_harmony=20)
|
| a.obsm["X_train"] = a.obsm["X_pca_harmony"]
|
| except Exception as exc:
|
| print(f"[preprocess] harmony skipped ({exc}); using raw PCA", flush=True)
|
| a.obsm["X_train"] = a.obsm["X_pca"]
|
| return a
|
|
|
|
|
|
|
|
|
| def score_and_gate(a: ad.AnnData, panels: dict[int, list[str]],
|
| score_min: float = SCORE_MIN,
|
| margin_min: float = MARGIN_MIN) -> ad.AnnData:
|
| """score cells on 12 derm panels; accept label if best>score_min and margin>margin_min."""
|
| for cl, genes in panels.items():
|
| present = [g for g in genes if g in a.var_names]
|
| if not present:
|
| a.obs[f"derm{cl}_score"] = 0.0
|
| else:
|
| sc.tl.score_genes(a, gene_list=present, score_name=f"derm{cl}_score",
|
| random_state=0, use_raw=False)
|
| cols = [f"derm{cl}_score" for cl in sorted(panels)]
|
| S = a.obs[cols].values
|
| top1_ix = S.argmax(axis=1)
|
| top1 = S[np.arange(len(S)), top1_ix]
|
| S_copy = S.copy(); S_copy[np.arange(len(S)), top1_ix] = -np.inf
|
| top2 = S_copy.max(axis=1)
|
| margin = top1 - top2
|
| accept = (top1 > score_min) & (margin > margin_min)
|
|
|
| ids = np.array([int(cols[i].replace("derm", "").replace("_score", "")) for i in top1_ix])
|
| a.obs["derm_pseudo"] = ids
|
| a.obs["derm_pseudo_top1"] = top1
|
| a.obs["derm_pseudo_margin"] = margin
|
| a.obs["derm_pseudo_accept"] = accept
|
| return a
|
|
|
|
|
|
|
|
|
| class CorpusDataset(Dataset):
|
| def __init__(self, X, y, d, aux):
|
| self.X = X.astype(np.float32); self.y = y.astype(np.int64)
|
| self.d = d.astype(np.int64); self.aux = aux.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.from_numpy(self.aux[i]))
|
|
|
|
|
| class HybridSampler:
|
| def __init__(self, y, n_batches=100, seed=0):
|
| self.y = np.asarray(y); self.n_batches = n_batches
|
| self.rng = np.random.default_rng(seed)
|
| self.classes = np.unique(self.y)
|
| self.by_cls = {int(c): np.where(self.y == c)[0] for c in self.classes}
|
| counts = np.bincount(self.y, minlength=int(self.classes.max()) + 1).astype(float)
|
| self.natural_p = counts / counts.sum()
|
| def __iter__(self):
|
| for _ in range(self.n_batches):
|
| batch = []
|
| for c in self.classes:
|
| idx = self.by_cls[int(c)]
|
| take = min(GUARANTEED_PER_CLASS, len(idx))
|
| if take > 0:
|
| batch.extend(self.rng.choice(idx, size=take, replace=(len(idx) < take)).tolist())
|
| for _ in range(NATURAL_SLOTS):
|
| c = self.rng.choice(len(self.natural_p), p=self.natural_p)
|
| idx = self.by_cls.get(int(c), self.by_cls[int(self.classes[0])])
|
| batch.append(int(self.rng.choice(idx)))
|
| yield batch
|
| def __len__(self): return self.n_batches
|
|
|
|
|
| def train_panda(X_tr, y_tr, d_tr, aux_tr, n_classes, n_datasets, ck_out: Path):
|
| ck_out.mkdir(parents=True, exist_ok=True)
|
| counts = np.bincount(y_tr, minlength=n_classes)
|
| inv_sqrt = 1.0 / np.sqrt(counts + 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, device=DEVICE)
|
|
|
| ds = CorpusDataset(X_tr, y_tr, d_tr, aux_tr)
|
| loader = DataLoader(ds, batch_sampler=HybridSampler(y_tr, n_batches=100), num_workers=0)
|
|
|
| model = PANDAEncoder(variant="pca", n_pca=X_tr.shape[1], n_classes=n_classes,
|
| n_datasets=n_datasets).to(DEVICE)
|
| opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
|
|
|
| for stage, n_ep in enumerate(STAGE_EPOCHS):
|
| print(f"[panda-A] stage {stage} ({n_ep} epochs)", flush=True)
|
| for e in range(n_ep):
|
| t0 = time.time(); losses = []
|
| for X_b, y_b, d_b, aux_b in loader:
|
| X_b = X_b.to(DEVICE); y_b = y_b.to(DEVICE); d_b = d_b.to(DEVICE); aux_b = aux_b.to(DEVICE)
|
| 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()
|
| total = total + 0.6 * subcenter_angular_infonce(out["z"], y_b, proto_ref)
|
| if stage >= 2:
|
| total = total + F.cross_entropy(out["dom"], d_b)
|
| total = total + 0.3 * F.mse_loss(out["depth"].squeeze(1), aux_b[:, 1])
|
| total = total + 0.05 * hsic_biased(out["repr"], aux_b[:, 1:2])
|
| 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)
|
| losses.append(float(total.item()))
|
| if (e + 1) % 5 == 0:
|
| print(f" ep {e+1}/{n_ep} loss={np.mean(losses):.3f} dt={time.time()-t0:.1f}s", flush=True)
|
| torch.save({"model": model.state_dict()}, ck_out / f"panda_stage{stage}.pt")
|
|
|
| torch.save({"model": model.state_dict(),
|
| "prototypes": model.prototypes.detach().cpu().numpy()},
|
| ck_out / "panda_final.pt")
|
| return model
|
|
|
|
|
| @torch.no_grad()
|
| def infer(model, X, aux):
|
| model.eval()
|
| Xt = torch.from_numpy(X.astype(np.float32)).to(DEVICE)
|
| at = torch.from_numpy(aux.astype(np.float32)).to(DEVICE)
|
| B = 4096; preds = []; confs = []
|
| for i in range(0, len(Xt), B):
|
| out = model(Xt[i:i+B], at[i:i+B])
|
| p = F.softmax(out["logits"], dim=1)
|
| preds.append(p.argmax(dim=1).cpu().numpy())
|
| confs.append(p.max(dim=1).values.cpu().numpy())
|
| return np.concatenate(preds), np.concatenate(confs)
|
|
|
|
|
|
|
|
|
| def report_depletion(labels: np.ndarray, genotype: np.ndarray, n_classes: int) -> pd.DataFrame:
|
| n_wt = int((genotype == "WT").sum()); n_cko = int((genotype == "En1-cKO").sum())
|
| base = n_cko / max(n_wt + n_cko, 1)
|
| rows = []
|
| for c in range(n_classes):
|
| m = labels == c
|
| w = int(((genotype == "WT") & m).sum()); k = int(((genotype == "En1-cKO") & m).sum())
|
| if w + k == 0: continue
|
| try:
|
| odds, p = fisher_exact([[w, n_wt - w], [k, n_cko - k]], alternative="two-sided")
|
| except ValueError:
|
| odds, p = 1.0, 1.0
|
| rows.append({"derm_id": c, "n": w + k, "n_WT": w, "n_cKO": k,
|
| "cko_frac": k / (w + k), "baseline": base,
|
| "odds_ratio": float(odds), "fisher_p": float(p)})
|
| return pd.DataFrame(rows).sort_values("cko_frac")
|
|
|
|
|
| def main():
|
| OUT_DIR.mkdir(parents=True, exist_ok=True); CK_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
| print("[A] load panels + dermal Dingwall", flush=True)
|
| panels = load_derm_panels()
|
| a = load_dingwall_dermal()
|
|
|
| pred_csv = ROOT / "discovery/pan_skin/marker/dingwall_predictions.csv"
|
| if pred_csv.exists():
|
| pred = pd.read_csv(pred_csv)
|
| pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"]))
|
| a.obs["v3_label"] = [pred_map.get(c, "unknown") for c in a.obs_names.astype(str)]
|
| a = a[np.isin(a.obs["v3_label"], ["fibroblast-papillary", "fibroblast-reticular"])].copy()
|
| print(f"[A] restricted to PANDA-v3 fibroblasts: n={a.n_obs}", flush=True)
|
|
|
| print("[A] paper-style preprocess", flush=True)
|
| a = preprocess_paper_style(a)
|
|
|
| print("[A] score + gate pseudo-labels", flush=True)
|
| a = score_and_gate(a, panels)
|
| n_acc = int(a.obs["derm_pseudo_accept"].sum())
|
| print(f"[A] pseudo-label acceptance: {n_acc}/{a.n_obs} ({100*n_acc/a.n_obs:.1f}%)", flush=True)
|
|
|
|
|
| train_mask = a.obs["derm_pseudo_accept"].values.astype(bool)
|
| X_all = np.asarray(a.obsm["X_train"])
|
| y_all = a.obs["derm_pseudo"].astype(int).values
|
| sample_ix = {s: i for i, s in enumerate(sorted(a.obs["sample"].astype(str).unique()))}
|
| d_all = np.array([sample_ix[s] for s in a.obs["sample"].astype(str)])
|
| aux_all = np.stack([np.zeros(a.n_obs, dtype=np.float32),
|
| np.log10(np.asarray(a.X.sum(axis=1)).ravel() + 1)], axis=1)
|
| aux_all[:, 1] = (aux_all[:, 1] - aux_all[:, 1].mean()) / (aux_all[:, 1].std() + 1e-6)
|
|
|
| classes = sorted(np.unique(y_all[train_mask]).tolist())
|
| if len(classes) < 2:
|
| print("[A] not enough classes accepted; abort", flush=True); return
|
| cls_ix = {c: i for i, c in enumerate(classes)}
|
| y_all_ix = np.array([cls_ix.get(int(c), -1) for c in y_all])
|
| y_tr = y_all_ix[train_mask]
|
| X_tr = X_all[train_mask]; d_tr = d_all[train_mask]; aux_tr = aux_all[train_mask]
|
|
|
| print(f"[A] train n={train_mask.sum()} on {len(classes)} classes: {classes}", flush=True)
|
| model = train_panda(X_tr, y_tr, d_tr, aux_tr, n_classes=len(classes),
|
| n_datasets=len(sample_ix), ck_out=CK_DIR)
|
|
|
|
|
| infer_mask = ~train_mask
|
| preds_ix, confs = infer(model, X_all[infer_mask], aux_all[infer_mask])
|
| preds_derm = np.array([classes[p] for p in preds_ix])
|
|
|
|
|
| final = np.where(train_mask, y_all,
|
| np.concatenate([y_all[train_mask].astype(int) * 0 - 1,
|
| preds_derm.astype(int)])[:a.n_obs] if False else 0)
|
|
|
| final = y_all.astype(int).copy()
|
| final[infer_mask] = preds_derm.astype(int)
|
|
|
| df = pd.DataFrame({
|
| "cell_id": a.obs_names.astype(str).values,
|
| "sample": a.obs["sample"].astype(str).values,
|
| "genotype": a.obs["genotype"].astype(str).values,
|
| "pseudo_derm": y_all,
|
| "pseudo_accept": train_mask,
|
| "final_derm": final,
|
| })
|
| df.to_csv(OUT_DIR / "102_variantA_predictions.csv", index=False)
|
|
|
| dep = report_depletion(final, a.obs["genotype"].values, n_classes=12)
|
| dep.to_csv(OUT_DIR / "102_variantA_depletion.csv", index=False)
|
| summary = {
|
| "variant": "A_semi_supervised_S3_scoring",
|
| "score_min": SCORE_MIN, "margin_min": MARGIN_MIN, "top_n": TOP_N,
|
| "n_total": int(a.n_obs), "n_train_pseudo": int(train_mask.sum()),
|
| "classes_trained": classes,
|
| "derm10": dep[dep["derm_id"] == 10].to_dict("records"),
|
| "derm2": dep[dep["derm_id"] == 2].to_dict("records"),
|
| "derm9": dep[dep["derm_id"] == 9].to_dict("records"),
|
| "all": dep.to_dict("records"),
|
| }
|
| (OUT_DIR / "102_variantA_summary.json").write_text(json.dumps(summary, indent=2, default=str))
|
| print(f"[A] done -> {OUT_DIR}/102_variantA_*", flush=True)
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
