"""zero-shot HSC PANDA on dahlin 2018: WT vs Kit W41/W41 class enrichment.""" from __future__ import annotations from pathlib import Path import warnings, json, sys, pickle warnings.filterwarnings("ignore") import numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp import torch from scipy import stats 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 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") CKPT = Path(str(PANDA_ROOT / "checkpoints/hematopoiesis")) HARM = Path(str(PANDA_ROOT / "data/corpus/hematopoiesis/harmonized")) DAHLIN = Path(str(PANDA_ROOT / "data/corpus/hematopoiesis/held_out_unlabeled/dahlin_extract")) OUT = Path(str(PANDA_ROOT / "discovery/hematopoiesis/marker")) GENOTYPE_MAP = { "SIGAB1": "WT", "SIGAC1": "WT", "SIGAD1": "WT", "SIGAF1": "WT", "SIGAG1": "WT", "SIGAH1": "WT", "SIGAG8": "Kit_W41", "SIGAH8": "Kit_W41", } # Dahlin sorted TWO different FACS gates (GEO source_name, GSE107727): # SIGAB1/C1/D1 -> "LSK_sample_1..3" (WT, LSK gate) # SIGAF1/G1/H1 -> "Lin_negative_cKit_positive_sample_4..6" (WT, LK gate) # SIGAG8/H8 -> "Lin_negative_cKit_positive_sample_39/40" (W41, LK gate) # LSK is a subset of LK enriched for immature progenitors, so the two WT # groups differ enormously (WT-LSK is ~84% MPP / ~1% erythroid; WT-LK is # ~36% / ~35%). Pooling them into one WT baseline against LK-only mutants # makes the sorting gate, not the genotype, the dominant contrast: the pure # WT-LK-vs-WT-LSK effect is log2fc +4.79 erythroid / +4.22 myeloid, i.e. # LARGER than the Kit-W41 effect it was being attributed to. # Dahlin themselves compare W41 LK vs WT LK. We now do the same. GATE_MAP = { "SIGAB1": "LSK", "SIGAC1": "LSK", "SIGAD1": "LSK", "SIGAF1": "LK", "SIGAG1": "LK", "SIGAH1": "LK", "SIGAG8": "LK", "SIGAH8": "LK", } GATE_MATCHED = "LK" # restrict the genotype contrast to this gate def load_dahlin_all(): print("[dahlin] loading 8 samples …", flush=True) parts = [] for f in sorted(DAHLIN.glob("*.txt.gz")): gsm = f.name.split("_")[0] sample = f.name.split("_")[1].split(".")[0] genotype = GENOTYPE_MAP.get(sample, "unknown") print(f" {sample} ({genotype})", flush=True) df = pd.read_csv(f, sep="\t", compression="gzip", index_col=0) X = sp.csr_matrix(df.values.T.astype(np.float32)) obs = pd.DataFrame(index=[f"{sample}_{bc}" for bc in df.columns.astype(str)]) obs["sample"] = sample obs["genotype"] = genotype obs["gate"] = GATE_MAP.get(sample, "unknown") obs["dataset"] = "dahlin_GSE107727" var = pd.DataFrame(index=df.index.astype(str)) var["ensmusg"] = var.index.values a = ad.AnnData(X=X, obs=obs, var=var) parts.append(a) return ad.concat(parts, join="outer", label="_batch") def convert_ensembl_to_symbol(a): import mygene mg = mygene.MyGeneInfo() ids = a.var_names.astype(str).tolist() print(f"[dahlin] querying {len(ids)} ENSMUSG IDs …", flush=True) res = mg.querymany(ids, scopes="ensembl.gene", fields="symbol", species="mouse", verbose=False) id2sym = {r["query"]: r["symbol"] for r in res if "symbol" in r} syms = pd.Series(a.var_names.astype(str)).map(id2sym).values keep = pd.notna(syms) print(f"[dahlin] mapped {int(keep.sum())}/{len(a.var_names)} genes", flush=True) a = a[:, keep].copy() a.var_names = syms[keep] a.var_names_make_unique() return a def main(): a = load_dahlin_all() print(f"[dahlin] concat shape: {a.shape}", flush=True) a = convert_ensembl_to_symbol(a) print(f"[dahlin] after symbol conversion: {a.shape}", flush=True) ck = torch.load(CKPT / "panda_final.pt", map_location=DEVICE, weights_only=False) classes = ck["classes"]; datasets = ck["datasets"] model = PANDAEncoder(n_pca=50, n_classes=len(classes), n_datasets=len(datasets)).to(DEVICE).eval() model.load_state_dict(ck["model"]) stats_ = np.load(HARM / "corpus_stats.npz", allow_pickle=True) shared_hvgs = [str(g) for g in stats_["shared_hvgs"]] mu, sig = stats_["mean"], stats_["std"] with open(HARM / "pca_basis.pkl", "rb") as f: pca = pickle.load(f) G = len(shared_hvgs) hvg2i = {g: i for i, g in enumerate(shared_hvgs)} common = [g for g in a.var_names.astype(str) if g in hvg2i] frac = len(common) / G print(f"[proj] {len(common)}/{G} HVGs present ({frac:.1%})", flush=True) a_c = a[:, common].copy() 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((a.n_obs, G), dtype=np.float32) cols = [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) all_z = [] with torch.no_grad(): for i in range(0, a.n_obs, 4096): xb = torch.from_numpy(Xpca[i:i+4096]).to(DEVICE) aux = torch.zeros(len(xb), 2, device=DEVICE) all_z.append(model(xb, aux, lam_dann=0.0)["z"].cpu().numpy()) Z = np.concatenate(all_z, axis=0) protos = ck["prototypes"] protos = protos / (np.linalg.norm(protos, axis=1, keepdims=True) + 1e-8) cos = Z @ protos.T pred = np.array([classes[i] for i in cos.argmax(axis=1)], dtype=object) conf = cos.max(axis=1) a.obs["pred_label"] = pred a.obs["pred_conf"] = conf.astype(np.float32) print(f"\n[dahlin] predicted class distribution overall:") print(a.obs["pred_label"].value_counts()) print(f"\n[dahlin] per genotype:") xt = pd.crosstab(a.obs["pred_label"], a.obs["genotype"], normalize="columns") print(xt.round(4)) xt.to_csv(OUT / "66_dahlin_class_per_genotype.csv") def enrichment(mask_a, mask_b, label_a, label_b): rows = [] n_a = int(mask_a.sum()); n_b = int(mask_b.sum()) for c in classes: n_c_a = int(((pred == c) & mask_a).sum()) n_c_b = int(((pred == c) & mask_b).sum()) odds, p = stats.fisher_exact(np.array( [[n_c_a, n_a - n_c_a], [n_c_b, n_b - n_c_b]])) f_a = (n_c_a + 1) / (n_a + 2); f_b = (n_c_b + 1) / (n_b + 2) rows.append({"class": c, f"n_{label_a}": n_c_a, f"n_{label_b}": n_c_b, f"pct_{label_a}": round(100 * n_c_a / max(n_a, 1), 3), f"pct_{label_b}": round(100 * n_c_b / max(n_b, 1), 3), "log2_fold": round(np.log2(f_a / f_b), 3), "fisher_p": p}) return pd.DataFrame(rows).sort_values("log2_fold"), n_a, n_b gate = a.obs["gate"].values geno = a.obs["genotype"].values # PRIMARY, gate-matched: W41 LK vs WT LK (what Dahlin themselves compare) print(f"\n[dahlin] PRIMARY gate-matched Fisher enrichment " f"(Kit_W41 {GATE_MATCHED} vs WT {GATE_MATCHED}):") df, n_kit, n_wt = enrichment((geno == "Kit_W41") & (gate == GATE_MATCHED), (geno == "WT") & (gate == GATE_MATCHED), "Kit_W41", "WT") print(f" Kit_W41 n={n_kit}, WT n={n_wt} (both {GATE_MATCHED} gate)") print(df.to_string(index=False)) df.to_csv(OUT / "66_dahlin_enrichment.csv", index=False) # NEGATIVE CONTROL: WT LK vs WT LSK — same genotype, gate only. Quantifies # how much of any "genotype" signal is really the sorting gate. print("\n[dahlin] NEGATIVE CONTROL (WT LK vs WT LSK, genotype held constant):") dfg, n_lk, n_lsk = enrichment((geno == "WT") & (gate == "LK"), (geno == "WT") & (gate == "LSK"), "WT_LK", "WT_LSK") print(f" WT_LK n={n_lk}, WT_LSK n={n_lsk}") print(dfg.to_string(index=False)) dfg.to_csv(OUT / "66_dahlin_gate_negative_control.csv", index=False) # legacy pooled contrast, kept for provenance only dfp, _, _ = enrichment(geno == "Kit_W41", geno == "WT", "Kit_W41", "WT") dfp.to_csv(OUT / "66_dahlin_enrichment_pooled_CONFOUNDED.csv", index=False) a.obs.to_csv(OUT / "66_dahlin_predictions.csv") print(f"\n[dahlin] complete. Outputs in {OUT}/") if __name__ == "__main__": main()