diff --git a/scripts/__init__.py b/scripts/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/scripts/__pycache__/__init__.cpython-310.pyc b/scripts/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9c119dcdc3e8a37432dada847a7bb1171aefc3d9 Binary files /dev/null and b/scripts/__pycache__/__init__.cpython-310.pyc differ diff --git a/scripts/analysis/100_primary_eden_discovery.py b/scripts/analysis/100_primary_eden_discovery.py new file mode 100644 index 0000000000000000000000000000000000000000..8ca8e2699c0b783832bc0d5e5d9f149cf3ca990b --- /dev/null +++ b/scripts/analysis/100_primary_eden_discovery.py @@ -0,0 +1,205 @@ +"""primary EDEN discovery on Dingwall: PANDA-v3 dermal-fibro subset, Leiden res=1.5, wilcoxon markers + Fisher cKO depletion + module scoring.""" +from pathlib import Path +import warnings, json, sys, pickle, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp, torch, torch.nn.functional as F +from scipy.stats import fisher_exact, mannwhitneyu +warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 +sys.path.insert(0, "/home/bcheng/PRISM") +from panda import PANDAEncoder + +ROOT = Path("/home/bcheng/PRISM") +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +CKO_GSMS = {"GSM6833482", "GSM6833483"} # CORRECTED: 480/481 are rttaControl (WT), not cKO +WT_GSMS = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"} # CORRECTED: 4 Cre-neg controls per GEO metadata + +# secondary EDEN definition per Dingwall 2024 +SECONDARY_EDEN_PANEL = ["S100a4", "Tnc", "Pdgfra"] + +# En1-responsive eccrine program (from restored 57_pathway_analysis.py, En1 removed) +SWEAT_GLAND_PANEL_ENSMINUSEN1 = ["Foxi3", "Foxa1", "Krt8", "Krt18", "Krt19", "Muc5b", "Aqp5"] +# Eda pathway +EDA_PATHWAY_PANEL = ["Eda", "Edar", "Edaradd", "Nfkb1", "Nfkb2", "Rela"] + + +def load_dingwall_with_v3_predictions(): + print("[eden] loading Dingwall raw + v3 predictions", flush=True) + raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad") + pred = pd.read_csv(ROOT / "discovery/pan_skin/marker/dingwall_predictions.csv") + pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"])) + raw.obs["pred_label"] = np.array([pred_map.get(c, "unknown") for c in raw.obs_names.astype(str)]) + raw.obs["genotype"] = np.where(raw.obs["sample"].astype(str).isin(list(CKO_GSMS)), "En1-cKO", + np.where(raw.obs["sample"].astype(str).isin(list(WT_GSMS)), "WT", "other")) + labeled = raw.obs["genotype"].isin(["WT", "En1-cKO"]).values + raw = raw[labeled].copy() + + dermal_mask = np.isin(raw.obs["pred_label"], ["fibroblast-papillary", "fibroblast-reticular"]) + dermal = raw[dermal_mask].copy() + print(f"[eden] {dermal.n_obs} dermal-fibroblast cells for sub-clustering", flush=True) + return dermal + + +def subcluster_dermal(dermal, resolution=1.5): + print(f"[eden] preprocessing + PCA (Leiden resolution={resolution})", flush=True) + sc.pp.normalize_total(dermal, target_sum=1e4); sc.pp.log1p(dermal) + sc.pp.highly_variable_genes(dermal, n_top_genes=3000, flavor="seurat_v3", + inplace=True, batch_key=None) + dermal_hvg = dermal[:, dermal.var["highly_variable"]].copy() if "highly_variable" in dermal.var else dermal + sc.pp.scale(dermal_hvg, max_value=10) + sc.tl.pca(dermal_hvg, n_comps=30, random_state=0) + sc.pp.neighbors(dermal_hvg, n_neighbors=20, use_rep="X_pca") + sc.tl.leiden(dermal_hvg, resolution=resolution, random_state=0) + dermal.obs["leiden"] = dermal_hvg.obs["leiden"].astype(str) + print(f"[eden] {dermal.obs['leiden'].nunique()} sub-clusters found", flush=True) + return dermal + + +def score_modules(dermal): + for name, genes in [("secondary_eden", SECONDARY_EDEN_PANEL), + ("sweat_gland", SWEAT_GLAND_PANEL_ENSMINUSEN1), + ("eda_pathway", EDA_PATHWAY_PANEL)]: + present = [g for g in genes if g in dermal.var_names] + if not present: + dermal.obs[f"score_{name}"] = 0.0 + continue + sc.tl.score_genes(dermal, gene_list=present, score_name=f"score_{name}", + random_state=0, use_raw=False) + return dermal + + +def per_subcluster_analysis(dermal): + n_wt_tot = int((dermal.obs["genotype"] == "WT").sum()) + n_cko_tot = int((dermal.obs["genotype"] == "En1-cKO").sum()) + baseline_cko_frac = n_cko_tot / (n_wt_tot + n_cko_tot) + print(f"[eden] baseline: WT={n_wt_tot} cKO={n_cko_tot} (baseline cKO frac = {baseline_cko_frac:.3f})", flush=True) + + sc.tl.rank_genes_groups(dermal, "leiden", method="wilcoxon", n_genes=30, use_raw=False) + + rows = [] + for cls in sorted(dermal.obs["leiden"].unique(), key=int): + sub = dermal[dermal.obs["leiden"] == cls] + n_wt = int((sub.obs["genotype"] == "WT").sum()) + n_cko = int((sub.obs["genotype"] == "En1-cKO").sum()) + if n_wt + n_cko < 20: + continue + cko_frac = n_cko / (n_wt + n_cko) + + # fisher 2x2: (n_wt_in, n_wt_out) vs (n_cko_in, n_cko_out) — cluster depletion in cKO + n_wt_elsewhere = n_wt_tot - n_wt + n_cko_elsewhere = n_cko_tot - n_cko + odds, p_fisher = fisher_exact([[n_wt, n_wt_elsewhere], [n_cko, n_cko_elsewhere]], + alternative="two-sided") + depletion_direction = "cKO-depleted" if cko_frac < baseline_cko_frac else "cKO-enriched" + + s2eden_mean_wt = float(sub[sub.obs["genotype"] == "WT"].obs["score_secondary_eden"].mean()) if n_wt > 0 else 0.0 + s2eden_mean_cko = float(sub[sub.obs["genotype"] == "En1-cKO"].obs["score_secondary_eden"].mean()) if n_cko > 0 else 0.0 + sg_mean_wt = float(sub[sub.obs["genotype"] == "WT"].obs["score_sweat_gland"].mean()) if n_wt > 0 else 0.0 + sg_mean_cko = float(sub[sub.obs["genotype"] == "En1-cKO"].obs["score_sweat_gland"].mean()) if n_cko > 0 else 0.0 + eda_mean_wt = float(sub[sub.obs["genotype"] == "WT"].obs["score_eda_pathway"].mean()) if n_wt > 0 else 0.0 + eda_mean_cko = float(sub[sub.obs["genotype"] == "En1-cKO"].obs["score_eda_pathway"].mean()) if n_cko > 0 else 0.0 + + genes_list = list(dermal.uns["rank_genes_groups"]["names"][cls][:10]) + lfc_list = list(dermal.uns["rank_genes_groups"]["logfoldchanges"][cls][:10]) + top_markers = ", ".join([f"{g}({lfc:+.1f})" for g, lfc in zip(genes_list, lfc_list)]) + + rows.append({ + "leiden_cluster": cls, + "n_cells": n_wt + n_cko, + "n_WT": n_wt, "n_cKO": n_cko, + "cko_frac": cko_frac, + "baseline_cko_frac": baseline_cko_frac, + "depletion_direction": depletion_direction, + "fisher_p_two_sided": p_fisher, + "odds_ratio": odds, + "score_secondary_eden_WT_mean": s2eden_mean_wt, + "score_secondary_eden_cKO_mean": s2eden_mean_cko, + "score_sweat_gland_WT_mean": sg_mean_wt, + "score_sweat_gland_cKO_mean": sg_mean_cko, + "score_eda_pathway_WT_mean": eda_mean_wt, + "score_eda_pathway_cKO_mean": eda_mean_cko, + "top_wilcoxon_markers": top_markers, + }) + return pd.DataFrame(rows), baseline_cko_frac + + +def call_primary_and_secondary(df, baseline_cko_frac): + # secondary EDEN: highest score_secondary_eden_WT_mean AND cKO-depleted (Fisher p<0.05) + df_wt_ordered = df.sort_values("score_secondary_eden_WT_mean", ascending=False) + secondary_candidates = df_wt_ordered[ + (df_wt_ordered["depletion_direction"] == "cKO-depleted") & + (df_wt_ordered["fisher_p_two_sided"] < 0.05) + ] + secondary = secondary_candidates.iloc[0]["leiden_cluster"] if len(secondary_candidates) > 0 else None + + # primary EDEN: cKO-depleted + LOW secondary_eden (S100a4-neg) + HIGH Eda_pathway (En1-responsive) + df_ranked = df.copy() + df_ranked["depletion_score"] = -np.log10(df_ranked["fisher_p_two_sided"].clip(lower=1e-300)) * \ + (df_ranked["cko_frac"] < baseline_cko_frac).astype(int) + primary_score = df_ranked["depletion_score"] * \ + (1.0 / (df_ranked["score_secondary_eden_WT_mean"].abs() + 0.01)) * \ + (df_ranked["score_eda_pathway_WT_mean"] + 0.1) + df_ranked["primary_eden_composite_score"] = primary_score + df_ranked = df_ranked.sort_values("primary_eden_composite_score", ascending=False) + primary_candidates = df_ranked[ + (df_ranked["depletion_direction"] == "cKO-depleted") & + (df_ranked["fisher_p_two_sided"] < 0.05) & + (df_ranked["leiden_cluster"] != secondary) + ].head(3) + return secondary, primary_candidates, df_ranked + + +def main(): + dermal = load_dingwall_with_v3_predictions() + dermal = subcluster_dermal(dermal, resolution=1.5) + dermal = score_modules(dermal) + df, baseline_cko = per_subcluster_analysis(dermal) + secondary, primary_cands, df_ranked = call_primary_and_secondary(df, baseline_cko) + + out = ROOT / "discovery/pan_skin/marker" + out.mkdir(parents=True, exist_ok=True) + df_ranked.to_csv(out / "100_primary_eden_discovery.csv", index=False) + + summary = { + "target": "Dingwall_GSE220977", + "hypothesis": "Primary EDEN precedes Secondary EDEN (S100a4+/Tnc+ cluster 20/Derm10) in dermal lineage", + "method": "PANDA-v3 predicts dermal-fibroblast compartment; Leiden sub-clustering " + "(resolution=1.5) resolves substructure; Wilcoxon markers + Fisher-exact " + "cKO enrichment + module scoring (Secondary_EDEN, Sweat_gland, Eda_pathway) " + "rank sub-clusters for Primary EDEN candidacy", + "baseline_cko_frac": float(baseline_cko), + "n_subclusters": int(len(df)), + "secondary_eden_call": { + "leiden_cluster": str(secondary), + "criteria": "highest S100a4+Tnc+Pdgfra score AND Fisher cKO-depleted p<0.05", + "row": df[df["leiden_cluster"] == secondary].iloc[0].to_dict() if secondary else None, + }, + "primary_eden_candidates_top3": primary_cands[[ + "leiden_cluster", "n_cells", "n_WT", "n_cKO", "cko_frac", + "fisher_p_two_sided", "score_secondary_eden_WT_mean", + "score_sweat_gland_WT_mean", "score_eda_pathway_WT_mean", + "top_wilcoxon_markers", "primary_eden_composite_score", + ]].to_dict("records") if len(primary_cands) > 0 else [], + } + (out / "100_primary_eden_summary.json").write_text(json.dumps(summary, indent=2, default=str)) + + print(f"\n[eden] wrote {out}/100_primary_eden_*", flush=True) + print(f"\n=== SECONDARY EDEN CALL ===", flush=True) + print(f" leiden cluster: {secondary}", flush=True) + if secondary: + row = df[df["leiden_cluster"] == secondary].iloc[0] + print(f" n={row['n_cells']} (WT {row['n_WT']} / cKO {row['n_cKO']}), " + f"cko_frac={row['cko_frac']:.3f} vs baseline {baseline_cko:.3f}", flush=True) + print(f" Fisher p={row['fisher_p_two_sided']:.2e}, " + f"score_secondary_eden WT={row['score_secondary_eden_WT_mean']:.3f}", flush=True) + + print(f"\n=== PRIMARY EDEN CANDIDATES (top 3) ===", flush=True) + for _, row in primary_cands.iterrows(): + print(f" leiden {row['leiden_cluster']} n={row['n_cells']} (WT {row['n_WT']} / cKO {row['n_cKO']}), " + f"cko_frac={row['cko_frac']:.3f}, Fisher p={row['fisher_p_two_sided']:.2e}", flush=True) + print(f" S2EDEN_WT={row['score_secondary_eden_WT_mean']:.3f}, " + f"Sweat_WT={row['score_sweat_gland_WT_mean']:.3f}, " + f"Eda_WT={row['score_eda_pathway_WT_mean']:.3f}", flush=True) + print(f" top markers: {row['top_wilcoxon_markers']}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/101_primary_eden_derm_scoring.py b/scripts/analysis/101_primary_eden_derm_scoring.py new file mode 100644 index 0000000000000000000000000000000000000000..03bfc8707d4e1df3de1d362219565141769ddf38 --- /dev/null +++ b/scripts/analysis/101_primary_eden_derm_scoring.py @@ -0,0 +1,159 @@ +"""primary + secondary EDEN on Dingwall via score_genes against dingwall's own Derm0-11 markers (Data S1C top-30); argmax identity + Fisher cKO depletion.""" +from pathlib import Path +import warnings, json, sys, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp +from scipy.stats import fisher_exact +warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 + +ROOT = Path("/home/bcheng/PRISM") +CKO_GSMS = {"GSM6833482", "GSM6833483"} # CORRECTED: 480/481 are rttaControl (WT), not cKO +WT_GSMS = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"} # CORRECTED: 4 Cre-neg controls per GEO metadata +DERM_MARKERS = ROOT / "data/external_labels/dingwall_supp/biorxiv_media-3.xlsx" +TOP_N = 30 # top N markers per Derm cluster for scoring + +# EDEN identity map based on Data S1C + Data S2 CellChat +EDEN_IDENTITY = { + 10: "Secondary_EDEN", + 2: "Primary_EDEN_candidate_1", + 9: "Primary_EDEN_candidate_2", + 3: "EDEN_signalling", + 6: "EDEN_signalling", +} + + +def load_derm_marker_panels(): + df = pd.read_excel(DERM_MARKERS) + df = df.sort_values(["cluster", "avg_log2FC"], ascending=[True, False]) + panels = {} + for cl in sorted(df["cluster"].unique()): + genes = df[df["cluster"] == cl].head(TOP_N)["gene"].tolist() + panels[int(cl)] = genes + return panels + + +def load_dingwall_v3_dermal(): + raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad") + pred = pd.read_csv(ROOT / "discovery/pan_skin/marker/dingwall_predictions.csv") + pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"])) + raw.obs["pred_label"] = np.array([pred_map.get(c, "unknown") for c in raw.obs_names.astype(str)]) + raw.obs["genotype"] = np.where(raw.obs["sample"].astype(str).isin(list(CKO_GSMS)), "En1-cKO", + np.where(raw.obs["sample"].astype(str).isin(list(WT_GSMS)), "WT", "other")) + labeled = raw.obs["genotype"].isin(["WT", "En1-cKO"]).values + raw = raw[labeled].copy() + dermal_mask = np.isin(raw.obs["pred_label"], ["fibroblast-papillary", "fibroblast-reticular"]) + dermal = raw[dermal_mask].copy() + return dermal + + +def score_derm_identities(a, panels): + 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 + continue + sc.tl.score_genes(a, gene_list=present, score_name=f"derm{cl}_score", + random_state=0, use_raw=False) + return a + + +def main(): + print("[eden] loading Data S1C marker panels", flush=True) + panels = load_derm_marker_panels() + print(f"[eden] loaded {len(panels)} Derm panels", flush=True) + + print("[eden] loading Dingwall + PANDA-v3 dermal predictions", flush=True) + dermal = load_dingwall_v3_dermal() + print(f"[eden] {dermal.n_obs} dermal-fibroblast cells", flush=True) + + print("[eden] normalise + log1p", flush=True) + sc.pp.normalize_total(dermal, target_sum=1e4); sc.pp.log1p(dermal) + + print("[eden] scoring cells on all 12 Derm identity panels", flush=True) + dermal = score_derm_identities(dermal, panels) + + score_cols = [f"derm{cl}_score" for cl in sorted(panels.keys())] + score_mat = dermal.obs[score_cols].values + argmax = np.argmax(score_mat, axis=1) + derm_ids = [int(score_cols[i].replace("derm", "").replace("_score", "")) for i in argmax] + dermal.obs["derm_identity"] = derm_ids + dermal.obs["max_derm_score"] = score_mat.max(axis=1) + + n_wt_tot = int((dermal.obs["genotype"] == "WT").sum()) + n_cko_tot = int((dermal.obs["genotype"] == "En1-cKO").sum()) + baseline_cko_frac = n_cko_tot / (n_wt_tot + n_cko_tot) + print(f"[eden] baseline WT={n_wt_tot} cKO={n_cko_tot} (baseline cKO={baseline_cko_frac:.3f})", flush=True) + + rows = [] + for cl in sorted(panels.keys()): + sub = dermal[dermal.obs["derm_identity"] == cl] + n_wt = int((sub.obs["genotype"] == "WT").sum()) + n_cko = int((sub.obs["genotype"] == "En1-cKO").sum()) + if n_wt + n_cko == 0: + continue + cko_frac = n_cko / (n_wt + n_cko) if (n_wt + n_cko) > 0 else 0 + n_wt_else = n_wt_tot - n_wt + n_cko_else = n_cko_tot - n_cko + try: + odds, p_fisher = fisher_exact([[n_wt, n_wt_else], [n_cko, n_cko_else]], + alternative="two-sided") + except ValueError: + odds, p_fisher = 1.0, 1.0 + rows.append({ + "derm_id": cl, + "identity": EDEN_IDENTITY.get(cl, "other"), + "n_cells": n_wt + n_cko, + "n_WT": n_wt, "n_cKO": n_cko, + "cko_frac": cko_frac, + "baseline_cko_frac": baseline_cko_frac, + "cko_delta": cko_frac - baseline_cko_frac, + "fisher_p": float(p_fisher), + "odds_ratio": float(odds) if not np.isnan(odds) else None, + "depletion_direction": "cKO-depleted" if cko_frac < baseline_cko_frac + else "cKO-enriched", + "top10_markers_dingwall": ", ".join(panels[cl][:10]), + }) + df = pd.DataFrame(rows).sort_values("cko_delta") + + out = ROOT / "discovery/pan_skin/marker" + out.mkdir(parents=True, exist_ok=True) + df.to_csv(out / "101_derm_subcluster_scores.csv", index=False) + + secondary_row = df[df["derm_id"] == 10].iloc[0].to_dict() if 10 in df["derm_id"].values else None + primary_1_row = df[df["derm_id"] == 2].iloc[0].to_dict() if 2 in df["derm_id"].values else None + primary_2_row = df[df["derm_id"] == 9].iloc[0].to_dict() if 9 in df["derm_id"].values else None + + summary = { + "target": "Dingwall_GSE220977", + "method": "Score PANDA-v3 dermal-fibroblast predictions on Dingwall's own Derm0-11 " + "marker panels (Data S1C top-30 genes each); argmax identity per cell; " + "Fisher-exact cKO depletion per Derm identity", + "baseline_cko_frac": baseline_cko_frac, + "n_dermal_cells_total": int(dermal.n_obs), + "secondary_eden_Derm10": secondary_row, + "primary_eden_Derm2": primary_1_row, + "primary_eden_Derm9": primary_2_row, + "all_derm_summary": df.to_dict("records"), + } + (out / "101_derm_identity_summary.json").write_text(json.dumps(summary, indent=2, default=str)) + + print(f"\n[eden] wrote {out}/101_derm_*", flush=True) + print(f"\n=== SECONDARY EDEN (Derm10) ===", flush=True) + if secondary_row: + print(f" n={secondary_row['n_cells']} (WT {secondary_row['n_WT']} / cKO {secondary_row['n_cKO']}), " + f"cko_frac={secondary_row['cko_frac']:.3f} vs baseline {baseline_cko_frac:.3f}", flush=True) + print(f" {secondary_row['depletion_direction']}, Fisher p={secondary_row['fisher_p']:.2e}", flush=True) + + print(f"\n=== PRIMARY EDEN candidate 1 (Derm2) ===", flush=True) + if primary_1_row: + print(f" n={primary_1_row['n_cells']} (WT {primary_1_row['n_WT']} / cKO {primary_1_row['n_cKO']}), " + f"cko_frac={primary_1_row['cko_frac']:.3f}", flush=True) + print(f" {primary_1_row['depletion_direction']}, Fisher p={primary_1_row['fisher_p']:.2e}", flush=True) + + print(f"\n=== PRIMARY EDEN candidate 2 (Derm9) ===", flush=True) + if primary_2_row: + print(f" n={primary_2_row['n_cells']} (WT {primary_2_row['n_WT']} / cKO {primary_2_row['n_cKO']}), " + f"cko_frac={primary_2_row['cko_frac']:.3f}", flush=True) + print(f" {primary_2_row['depletion_direction']}, Fisher p={primary_2_row['fisher_p']:.2e}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/102_train_on_dingwall_variantA.py b/scripts/analysis/102_train_on_dingwall_variantA.py new file mode 100644 index 0000000000000000000000000000000000000000..2877d74d4e0eacbdc36e51c43ddcf593d986558e --- /dev/null +++ b/scripts/analysis/102_train_on_dingwall_variantA.py @@ -0,0 +1,317 @@ +"""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 + +sys.path.insert(0, "/home/bcheng/PRISM") +from panda.model import ( + PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, subcenter_angular_infonce +) + +ROOT = Path("/home/bcheng/PRISM") +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" + +# Dingwall GSM -> genotype (see 101_primary_eden_derm_scoring) +CKO_GSMS = {"GSM6833482", "GSM6833483"} # CORRECTED: 480/481 are rttaControl (WT), not cKO +WT_GSMS = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"} # CORRECTED: 4 Cre-neg controls per GEO metadata + +TOP_N = 30 # markers per Derm panel for scoring +SCORE_MIN = 0.10 # min score to accept a pseudo-label +MARGIN_MIN = 0.05 # min gap best - runner-up +N_HVG = 2000 # matches paper +N_PCA = 40 # matches paper +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +# Training config (mirrors 20_train_panda.py) +GUARANTEED_PER_CLASS = 6 +NATURAL_SLOTS = 96 +STAGE_EPOCHS = [15, 25, 40, 40] +BALANCE_MIX = 0.5 + + +# ---------- data prep ---------- + +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 # noqa + 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 + + +# ---------- pseudo-labelling ---------- + +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 + + +# ---------- PANDA training (mirrors 20_train_panda.py) ---------- + +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) + + +# ---------- reporting ---------- + +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() + # reuse existing panda-v3 fibroblast calls if present, else all cells + 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/heldout split (gate = train; rest = infer) + 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) + + # inference on held-out + 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]) + + # combine: use pseudo-label on train, prediction on inference + final = np.where(train_mask, y_all, + np.concatenate([y_all[train_mask].astype(int) * 0 - 1, # placeholder + preds_derm.astype(int)])[:a.n_obs] if False else 0) + # simpler: assemble directly + 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() diff --git a/scripts/analysis/103_replicate_dingwall_seurat_pipeline.py b/scripts/analysis/103_replicate_dingwall_seurat_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..e7e6c130f08861c517817b324267ff93bf483b8e --- /dev/null +++ b/scripts/analysis/103_replicate_dingwall_seurat_pipeline.py @@ -0,0 +1,261 @@ +"""replicate dingwall's seurat clustering (QC, harmony, PCA, leiden) to derive Derm0..Derm11 labels.""" +from __future__ import annotations +from pathlib import Path +import warnings, json, sys +warnings.filterwarnings("ignore") + +import numpy as np +import pandas as pd +import anndata as ad +import scanpy as sc + +ROOT = Path("/home/bcheng/PRISM") +RAW_H5 = ROOT / "data/raw/GSE220977_combined.h5ad" +DERM_MARKERS = ROOT / "data/external_labels/dingwall_supp/biorxiv_media-3.xlsx" +TOP_MARKERS = ROOT / "data/external_labels/dingwall_supp/biorxiv_media-1.xlsx" +OUT_DIR = ROOT / "data/processed/dingwall_replica" + +CKO_GSMS = {"GSM6833482", "GSM6833483"} # CORRECTED: 480/481 are rttaControl (WT), not cKO +WT_GSMS = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"} # CORRECTED: 4 Cre-neg controls per GEO metadata + +# Paper values +DERMAL_TOP_CLUSTERS = {0, 1, 3, 4, 5, 8, 11, 20} # from STAR Methods +N_HVG = 2000 +N_PCA = 40 +RES = 0.7 +N_LEIDEN_DERM = 12 # target: Derm0..Derm11 +JACCARD_TOP_N = 50 # top-N markers for label mapping + + +# ---------- QC + preprocessing ---------- + +def qc_filter(a: ad.AnnData) -> ad.AnnData: + a.var["mt"] = a.var_names.str.upper().str.startswith("MT-") | \ + a.var_names.str.startswith("mt-") + sc.pp.calculate_qc_metrics(a, qc_vars=["mt"], inplace=True, percent_top=None, + log1p=False) + sc.pp.filter_cells(a, min_genes=300) + a = a[a.obs["n_genes_by_counts"] < 6000].copy() + a = a[a.obs["pct_counts_mt"] < 5].copy() + sc.pp.filter_genes(a, min_cells=10) + return a + + +def lognorm(a: ad.AnnData) -> ad.AnnData: + a.layers["counts"] = a.X.copy() if not hasattr(a.X, "toarray") or True else a.X.copy() + sc.pp.normalize_total(a, target_sum=1e4) + sc.pp.log1p(a) + return a + + +def hvg_pca_harmony(a: ad.AnnData, n_hvg=N_HVG, n_pca=N_PCA, batch_key="sample") -> ad.AnnData: + sc.pp.highly_variable_genes(a, n_top_genes=n_hvg, flavor="seurat", batch_key=batch_key) + a_use = a[:, a.var["highly_variable"]].copy() + sc.pp.scale(a_use, max_value=10, zero_center=False) + sc.tl.pca(a_use, n_comps=n_pca, use_highly_variable=False, zero_center=False) + # copy PCA back — a_use has same obs rows as a + a.obsm["X_pca"] = a_use.obsm["X_pca"].copy() + rep = "X_pca" + try: + # harmonypy directly on the pca matrix to avoid scanpy wrapper obsm-shape bug + import harmonypy as hm + pca_mat = a.obsm["X_pca"].copy() + meta = a.obs[[batch_key]].reset_index(drop=True) + ho = hm.run_harmony(pca_mat, meta, batch_key, max_iter_harmony=20) + # ho.Z_corr is (pcs, cells); make it (cells, pcs) + z = ho.Z_corr + if z.shape[1] == a.n_obs: + harm_mat = np.ascontiguousarray(z.T) + elif z.shape[0] == a.n_obs: + harm_mat = np.ascontiguousarray(z) + else: + raise RuntimeError(f"unknown harmony shape {z.shape}, n_obs={a.n_obs}") + if harm_mat.shape[0] == a.n_obs and harm_mat.shape[1] == pca_mat.shape[1]: + a.obsm["X_pca_harmony"] = harm_mat + rep = "X_pca_harmony" + print(f"[replica] Harmony ok, X_pca_harmony shape={harm_mat.shape}", flush=True) + else: + print(f"[replica] Harmony output shape mismatch ({harm_mat.shape}); using X_pca", flush=True) + except Exception as exc: + print(f"[replica] Harmony skipped ({exc}); using X_pca", flush=True) + a.uns["_replica_rep"] = rep + return a + + +def leiden_cluster(a: ad.AnnData, res=RES) -> ad.AnnData: + rep = a.uns.get("_replica_rep", "X_pca") + sc.pp.neighbors(a, n_neighbors=20, use_rep=rep, n_pcs=N_PCA) + sc.tl.leiden(a, resolution=res, key_added="leiden") + return a + + +# ---------- 23-cluster stage (map dermal identity) ---------- + +def call_dermal_23(a: ad.AnnData) -> ad.AnnData: + """first-pass clustering; mark cells whose leiden id maps to DERMAL_TOP_CLUSTERS.""" + print("[23] preprocess", flush=True) + a = qc_filter(a); a = lognorm(a); a = hvg_pca_harmony(a) + print("[23] leiden res=0.7", flush=True) + a = leiden_cluster(a, res=RES) + + # rank markers per top-level cluster + sc.tl.rank_genes_groups(a, "leiden", method="wilcoxon", n_genes=100) + df_tl = pd.read_excel(TOP_MARKERS) # Data S1 all-cluster markers + tl_panels = {int(c): df_tl[df_tl["cluster"] == c].sort_values("avg_log2FC", ascending=False) + .head(JACCARD_TOP_N)["gene"].tolist() for c in sorted(df_tl["cluster"].unique())} + tl_map = map_leiden_to_paper(a, "leiden", tl_panels, top_n=JACCARD_TOP_N) + a.obs["paper_cluster_23"] = a.obs["leiden"].map(lambda c: tl_map.get(str(c), -1)) + a.obs["is_dermal_paper"] = a.obs["paper_cluster_23"].isin(DERMAL_TOP_CLUSTERS) + print(f"[23] cells matched to paper dermal set: {int(a.obs['is_dermal_paper'].sum())}", + flush=True) + return a + + +def map_leiden_to_paper(a: ad.AnnData, key: str, paper_panels: dict[int, list[str]], + top_n: int = JACCARD_TOP_N) -> dict[str, int]: + """best-matching paper cluster per leiden id via jaccard on top-N markers.""" + ranks = a.uns["rank_genes_groups"] + names = pd.DataFrame(ranks["names"]) + out = {} + used = set() + scores = [] + for lc in names.columns: + my_top = set(names[lc].dropna().tolist()[:top_n]) + best_pc, best_j = None, -1.0 + for pc, panel in paper_panels.items(): + j = len(my_top & set(panel[:top_n])) / max(len(my_top | set(panel[:top_n])), 1) + if j > best_j: + best_pc, best_j = pc, j + scores.append({"leiden": lc, "best_paper": best_pc, "jaccard": best_j}) + out[lc] = best_pc + # convert to json-safe strings for h5ad serialization + a.uns[f"_map_scores_{key}"] = json.dumps(scores, default=str) + return out + + +# ---------- dermal subclustering stage (Derm0..Derm11) ---------- + +def subcluster_dermal(a: ad.AnnData) -> ad.AnnData: + dermal = a[a.obs["is_dermal_paper"]].copy() + # start again from raw counts on the subset + if "counts" in dermal.layers: + dermal.X = dermal.layers["counts"] + print(f"[derm] subset n={dermal.n_obs}", flush=True) + dermal = lognorm(dermal) + dermal = hvg_pca_harmony(dermal) + # tune res to hit ~12 clusters; res=0.7 is the paper value but scanpy Leiden can + # differ from Seurat FindClusters, so we sweep if the exact-res doesn't give 12 + dermal = leiden_cluster(dermal, res=RES) + # paper uses seurat FindClusters at res=0.7; scanpy leiden can differ so sweep to hit 12 + if len(dermal.obs["leiden"].unique()) != N_LEIDEN_DERM: + for r in [0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2]: + sc.tl.leiden(dermal, resolution=r, key_added=f"leiden_r{r}") + if len(dermal.obs[f"leiden_r{r}"].unique()) == N_LEIDEN_DERM: + dermal.obs["leiden"] = dermal.obs[f"leiden_r{r}"] + dermal.uns["_replica_derm_res"] = r + break + print(f"[derm] n_leiden = {len(dermal.obs['leiden'].unique())}", flush=True) + + # rank markers + map to Derm0..Derm11 + sc.tl.rank_genes_groups(dermal, "leiden", method="wilcoxon", n_genes=100) + df_s3 = pd.read_excel(DERM_MARKERS) + derm_panels = {int(c): df_s3[df_s3["cluster"] == c].sort_values("avg_log2FC", ascending=False) + .head(JACCARD_TOP_N)["gene"].tolist() for c in sorted(df_s3["cluster"].unique())} + derm_map = map_leiden_to_paper(dermal, "leiden", derm_panels, top_n=JACCARD_TOP_N) + dermal.obs["derm_label"] = dermal.obs["leiden"].map(lambda c: f"Derm{derm_map.get(str(c), -1)}") + + # push labels back into full object + labels = pd.Series("non_dermal", index=a.obs_names) + labels.loc[dermal.obs_names] = dermal.obs["derm_label"].values + a.obs["derm_label"] = labels + a.obs["leiden_derm"] = "" + a.obs.loc[dermal.obs_names, "leiden_derm"] = dermal.obs["leiden"].astype(str).values + # h5py-safe: stringify keys AND serialize dicts to json + a.uns["derm_leiden_to_paper"] = json.dumps({str(k): int(v) if v is not None else -1 + for k, v in derm_map.items()}, default=str) + a.uns["derm_panels_used"] = json.dumps({str(k): [str(g) for g in v] + for k, v in derm_panels.items()}, default=str) + return a, dermal + + +# ---------- QC of the replica: cluster 20 fractions ---------- + +def qc_cluster_20(a: ad.AnnData) -> dict: + a.obs["genotype"] = a.obs.get("genotype", pd.Series("unknown", index=a.obs_names)) + if "genotype" not in a.obs or a.obs["genotype"].nunique() < 2: + s = a.obs["sample"].astype(str) + a.obs["genotype"] = np.where(s.isin(list(CKO_GSMS)), "En1-cKO", + np.where(s.isin(list(WT_GSMS)), "WT", "other")) + + dermal_mask = a.obs["is_dermal_paper"].values + wt_derm = int(((a.obs["genotype"] == "WT") & dermal_mask).sum()) + ck_derm = int(((a.obs["genotype"] == "En1-cKO") & dermal_mask).sum()) + + # top-level cluster 20 replica + c20 = a.obs["paper_cluster_23"] == 20 + wt_c20 = int(((a.obs["genotype"] == "WT") & c20).sum()) + ck_c20 = int(((a.obs["genotype"] == "En1-cKO") & c20).sum()) + + # derm10 replica + d10 = a.obs["derm_label"] == "Derm10" + wt_d10 = int(((a.obs["genotype"] == "WT") & d10).sum()) + ck_d10 = int(((a.obs["genotype"] == "En1-cKO") & d10).sum()) + + return { + "expected_paper": {"wt_dermal_total": 17398, "cko_dermal_total": 8461, + "wt_c20_pct": 1.99, "cko_c20_pct": 0.08, + "wt_c20_abs": 346, "cko_c20_abs": 7}, + "replica": { + "wt_dermal_total": wt_derm, "cko_dermal_total": ck_derm, + "wt_c20": wt_c20, "cko_c20": ck_c20, + "wt_c20_pct": 100 * wt_c20 / max(wt_derm, 1), + "cko_c20_pct": 100 * ck_c20 / max(ck_derm, 1), + "wt_derm10": wt_d10, "cko_derm10": ck_d10, + "wt_derm10_pct": 100 * wt_d10 / max(wt_derm, 1), + "cko_derm10_pct": 100 * ck_d10 / max(ck_derm, 1), + }, + } + + +def main(): + OUT_DIR.mkdir(parents=True, exist_ok=True) + print("[replica] load raw", flush=True) + a = ad.read_h5ad(RAW_H5) + # inject genotype + s = a.obs["sample"].astype(str) + a.obs["genotype"] = np.where(s.isin(list(CKO_GSMS)), "En1-cKO", + np.where(s.isin(list(WT_GSMS)), "WT", "other")) + a = a[a.obs["genotype"].isin(["WT", "En1-cKO"])].copy() + print(f"[replica] n={a.n_obs}", flush=True) + + print("[replica] 23-cluster stage", flush=True) + a = call_dermal_23(a) + + print("[replica] dermal subcluster stage", flush=True) + a, dermal = subcluster_dermal(a) + + print("[replica] QC vs paper", flush=True) + qc = qc_cluster_20(a) + (OUT_DIR / "replica_cluster_20_qc.json").write_text(json.dumps(qc, indent=2, default=str)) + print(json.dumps(qc, indent=2, default=str), flush=True) + + # write per-Leiden -> paper mapping (parse json-back) + derm_map_parsed = json.loads(a.uns["derm_leiden_to_paper"]) + mm = pd.DataFrame([{"leiden_derm": k, "paper_derm": v} + for k, v in derm_map_parsed.items()]) + mm.to_csv(OUT_DIR / "replica_marker_matches.csv", index=False) + + # save — first stringify any datetime/complex obs cols to survive h5ad serialization + for col in list(a.obs.columns): + dt = a.obs[col].dtype + if pd.api.types.is_datetime64_any_dtype(dt) or dt == object: + try: + a.obs[col] = a.obs[col].astype(str) + except Exception: + del a.obs[col] + a.write_h5ad(OUT_DIR / "dingwall_replica.h5ad") + print(f"[replica] wrote {OUT_DIR}/dingwall_replica.h5ad", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/104_train_on_dingwall_derm_labels.py b/scripts/analysis/104_train_on_dingwall_derm_labels.py new file mode 100644 index 0000000000000000000000000000000000000000..484daf639a7478d9a13b456360576b20e1127628 --- /dev/null +++ b/scripts/analysis/104_train_on_dingwall_derm_labels.py @@ -0,0 +1,261 @@ +"""variant B — fully supervised panda on replicated dingwall Derm0..Derm11 labels from script 103.""" +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 torch +import torch.nn.functional as F +from torch.utils.data import Dataset, DataLoader +from scipy.stats import fisher_exact + +sys.path.insert(0, "/home/bcheng/PRISM") +from panda.model import ( + PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, subcenter_angular_infonce +) + +ROOT = Path("/home/bcheng/PRISM") +REPLICA_H5 = ROOT / "data/processed/dingwall_replica/dingwall_replica.h5ad" +OUT_DIR = ROOT / "discovery/pan_skin/marker" +CK_DIR = ROOT / "checkpoints/pan_skin_dingwall_derm" + +TRAIN_FRAC = 0.7 +SEED = 0 +N_PCA = 40 +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +# Training config (mirrors 20_train_panda.py) +GUARANTEED_PER_CLASS = 6 +NATURAL_SLOTS = 96 +STAGE_EPOCHS = [15, 25, 40, 40] +BALANCE_MIX = 0.5 + + +# ---------- split ---------- + +def genotype_stratified_split(labels: np.ndarray, genotypes: np.ndarray, + frac_train: float = TRAIN_FRAC, seed: int = SEED + ) -> tuple[np.ndarray, np.ndarray]: + """stratified 70/30 within each (label, genotype) group; preserves cKO/WT ratio per class.""" + rng = np.random.default_rng(seed) + n = len(labels); train = np.zeros(n, dtype=bool); test = np.zeros(n, dtype=bool) + for lab in np.unique(labels): + for g in np.unique(genotypes): + idx = np.where((labels == lab) & (genotypes == g))[0] + if len(idx) == 0: continue + rng.shuffle(idx) + k = max(1, int(len(idx) * frac_train)) if len(idx) > 1 else len(idx) + train[idx[:k]] = True + if len(idx) > 1: + test[idx[k:]] = True + return train, test + + +# ---------- PANDA training (identical to variant 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-B] 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) + + +# ---------- reporting ---------- + +def depletion_table(true_or_pred: np.ndarray, genotype: np.ndarray, class_names: list[str] + ) -> 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 i, cn in enumerate(class_names): + m = true_or_pred == i + 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_label": cn, "n": w + k, "n_WT": w, "n_cKO": k, + "cko_frac": k / (w + k), "baseline_cko": 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("[B] load replica", flush=True) + if not REPLICA_H5.exists(): + raise FileNotFoundError(f"Run 103 first — {REPLICA_H5} missing") + a = ad.read_h5ad(REPLICA_H5) + + dermal = a[a.obs["derm_label"].astype(str) != "non_dermal"].copy() + print(f"[B] dermal n={dermal.n_obs}", flush=True) + labels_str = dermal.obs["derm_label"].astype(str).values + classes = sorted(set(labels_str)) + cls_ix = {c: i for i, c in enumerate(classes)} + y_all = np.array([cls_ix[c] for c in labels_str]) + genotype = dermal.obs["genotype"].astype(str).values + + # get embedding from replica (harmony-corrected PCA) + rep_key = dermal.uns.get("_replica_rep", "X_pca_harmony") + if rep_key not in dermal.obsm: + rep_key = "X_pca_harmony" if "X_pca_harmony" in dermal.obsm else "X_pca" + X_all = np.asarray(dermal.obsm[rep_key]) + print(f"[B] using {rep_key} (d={X_all.shape[1]})", flush=True) + + sample_ix = {s: i for i, s in enumerate(sorted(dermal.obs["sample"].astype(str).unique()))} + d_all = np.array([sample_ix[s] for s in dermal.obs["sample"].astype(str)]) + total_counts = np.asarray(dermal.X.sum(axis=1)).ravel() + logc = np.log10(total_counts + 1); logc = (logc - logc.mean()) / (logc.std() + 1e-6) + aux_all = np.stack([np.zeros(dermal.n_obs, dtype=np.float32), logc.astype(np.float32)], axis=1) + + print("[B] genotype-stratified 70/30 split", flush=True) + tr, te = genotype_stratified_split(labels_str, genotype, frac_train=TRAIN_FRAC, seed=SEED) + print(f"[B] train={tr.sum()} test={te.sum()}", flush=True) + + manifest = pd.DataFrame({ + "cell_id": dermal.obs_names.astype(str).values, + "derm_label": labels_str, + "genotype": genotype, + "split": np.where(tr, "train", np.where(te, "test", "unassigned")), + }) + manifest.to_csv(OUT_DIR / "104_dingwall_derm_split_manifest.csv", index=False) + + print("[B] train PANDA", flush=True) + model = train_panda(X_all[tr], y_all[tr], d_all[tr], aux_all[tr], + n_classes=len(classes), n_datasets=len(sample_ix), ck_out=CK_DIR) + + print("[B] infer on held-out", flush=True) + pred_ix, conf = infer(model, X_all[te], aux_all[te]) + pred = pd.DataFrame({ + "cell_id": dermal.obs_names.astype(str).values[te], + "derm_true": labels_str[te], + "derm_pred": [classes[p] for p in pred_ix], + "confidence": conf, + "genotype": genotype[te], + }) + pred.to_csv(OUT_DIR / "104_dingwall_derm_predictions.csv", index=False) + + # depletion — reported for TEST set only, using PANDA predictions + pred_ix_full = np.array([cls_ix[c] for c in pred["derm_pred"].values]) + dep_pred = depletion_table(pred_ix_full, genotype[te], classes) + dep_true = depletion_table(y_all[te], genotype[te], classes) + dep_pred.to_csv(OUT_DIR / "104_dingwall_derm_depletion_pred.csv", index=False) + dep_true.to_csv(OUT_DIR / "104_dingwall_derm_depletion_true.csv", index=False) + + d10_true = dep_true[dep_true["derm_label"] == "Derm10"].to_dict("records") + d10_pred = dep_pred[dep_pred["derm_label"] == "Derm10"].to_dict("records") + acc = float((pred_ix == y_all[te]).mean()) + + summary = { + "variant": "B_fully_supervised_replica_labels", + "n_dermal_total": int(dermal.n_obs), + "n_train": int(tr.sum()), "n_test": int(te.sum()), + "classes": classes, + "test_accuracy": acc, + "expected_paper_derm10": {"wt_pct": 1.99, "cko_pct": 0.08, + "or_approx": 24.5, "wt_n_approx": 346, "cko_n_approx": 7}, + "test_derm10_true": d10_true, + "test_derm10_pred": d10_pred, + "test_depletion_true": dep_true.to_dict("records"), + "test_depletion_pred": dep_pred.to_dict("records"), + } + (OUT_DIR / "104_dingwall_derm_summary.json").write_text(json.dumps(summary, indent=2, default=str)) + print(f"[B] done -> {OUT_DIR}/104_dingwall_derm_*", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/105_primary_eden_full_dermal.py b/scripts/analysis/105_primary_eden_full_dermal.py new file mode 100644 index 0000000000000000000000000000000000000000..40f91172637fa7200d68999bd2a6f433a7902150 --- /dev/null +++ b/scripts/analysis/105_primary_eden_full_dermal.py @@ -0,0 +1,134 @@ +"""primary EDEN discovery on the full dingwall-defined dermal set (not the panda-v3 fibroblast subset).""" +from pathlib import Path +import warnings, json, sys, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp +from scipy.stats import fisher_exact +warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 + +ROOT = Path("/home/bcheng/PRISM") +REPLICA = ROOT / "data/processed/dingwall_replica/dingwall_replica.h5ad" +DERM_MARKERS = ROOT / "data/external_labels/dingwall_supp/biorxiv_media-3.xlsx" +TOP_N = 30 + +# EDEN identity map from Dingwall paper + Data S2 CellChat +EDEN_IDENTITY = { + 10: "Secondary_EDEN_(Dingwall_cluster_20)", + 2: "Primary_EDEN_candidate_1_(Derm2_-_immediate_precursor)", + 9: "Primary_EDEN_candidate_2_(Derm9)", + 6: "EDEN-signalling_(Derm6)", + 3: "EDEN-signalling_(Derm3)", +} + + +def main(): + print("[eden] loading replica dermal set (Dingwall-defined)", flush=True) + a = ad.read_h5ad(REPLICA) + # keep only cells the Seurat replica classified as belonging to Dingwall's dermal clusters + if "is_dermal_paper" in a.obs.columns: + dermal = a[a.obs["is_dermal_paper"] == True].copy() + elif "derm_label" in a.obs.columns: + dermal = a[a.obs["derm_label"] != "non_dermal"].copy() + else: + raise RuntimeError("no dermal indicator in replica") + print(f"[eden] Seurat replica dermal cells: {dermal.n_obs}", flush=True) + + # verify our Derm-label distribution matches the replica + if "derm_label" in dermal.obs.columns: + print(f"[eden] Derm label distribution (from replica):", flush=True) + for k, v in dermal.obs["derm_label"].value_counts().sort_index().items(): + print(f" {k}: {v}", flush=True) + + # load Data S1C panels + print(f"\n[eden] loading Data S1C marker panels", flush=True) + df = pd.read_excel(DERM_MARKERS) + df = df.sort_values(["cluster", "avg_log2FC"], ascending=[True, False]) + panels = {} + for cl in sorted(df["cluster"].unique()): + genes = df[df["cluster"] == cl].head(TOP_N)["gene"].astype(str).tolist() + panels[int(cl)] = genes + + # replica may or may not have log1p applied; reset from counts layer if present + if "counts" in dermal.layers: + dermal.X = dermal.layers["counts"] + if dermal.X.max() > 30: # raw counts + sc.pp.normalize_total(dermal, target_sum=1e4); sc.pp.log1p(dermal) + + # score each cell on all 12 Derm identity panels + print(f"\n[eden] scoring cells on all 12 Derm panels (top-30 markers each)", flush=True) + for cl, genes in panels.items(): + present = [g for g in genes if g in dermal.var_names] + if len(present) < 3: + dermal.obs[f"derm{cl}_score"] = 0.0 + continue + sc.tl.score_genes(dermal, gene_list=present, score_name=f"derm{cl}_score", + random_state=0, use_raw=False) + + # global baseline + n_wt = int((dermal.obs["genotype"] == "WT").sum()) + n_cko = int((dermal.obs["genotype"] == "En1-cKO").sum()) + baseline = n_cko / max(n_wt + n_cko, 1) + print(f"\n[eden] baseline: WT={n_wt} cKO={n_cko} (baseline cKO frac = {baseline:.3f})", flush=True) + + # use the replica's derm_label directly, not argmax of scores + rows = [] + print(f"\n[eden] per-Derm Fisher exact on replica-assigned identities:", flush=True) + for cl in sorted(panels.keys()): + derm_label = f"Derm{cl}" + if derm_label not in dermal.obs["derm_label"].values: + continue + sub = dermal[dermal.obs["derm_label"] == derm_label] + n_wt_c = int((sub.obs["genotype"] == "WT").sum()) + n_cko_c = int((sub.obs["genotype"] == "En1-cKO").sum()) + if n_wt_c + n_cko_c == 0: + continue + cko_frac = n_cko_c / (n_wt_c + n_cko_c) + n_wt_else = n_wt - n_wt_c + n_cko_else = n_cko - n_cko_c + try: + odds, p_f = fisher_exact([[n_wt_c, n_wt_else], [n_cko_c, n_cko_else]], + alternative="two-sided") + except ValueError: + odds, p_f = 1.0, 1.0 + rows.append({ + "derm_id": cl, + "identity": EDEN_IDENTITY.get(cl, "other"), + "n_cells": n_wt_c + n_cko_c, + "n_WT": n_wt_c, "n_cKO": n_cko_c, + "cko_frac": cko_frac, + "baseline_cko_frac": baseline, + "cko_delta": cko_frac - baseline, + "wt_enrichment_odds_ratio": float(1.0/odds) if odds > 0 else None, + "fisher_p_two_sided": float(p_f), + "depletion_direction": "cKO-depleted" if cko_frac < baseline else "cKO-enriched", + "top10_markers_dingwall_S1C": ", ".join(panels[cl][:10]), + }) + + result_df = pd.DataFrame(rows).sort_values("cko_delta") + out = ROOT / "discovery/pan_skin/marker" + out.mkdir(parents=True, exist_ok=True) + result_df.to_csv(out / "105_primary_eden_full_dermal.csv", index=False) + + print(f"\n{'Derm':<8}{'Identity':<50}{'n':<7}{'WT':<6}{'cKO':<6}{'cKO_frac':<10}" + f"{'OR (WT enrich)':<16}{'Fisher p':<12}", flush=True) + print("-" * 130, flush=True) + for _, r in result_df.iterrows(): + print(f"Derm{r['derm_id']:<5}{r['identity'][:47]:<50}{r['n_cells']:<7}" + f"{r['n_WT']:<6}{r['n_cKO']:<6}{r['cko_frac']:<10.3f}" + f"{r['wt_enrichment_odds_ratio']:<16.2f}{r['fisher_p_two_sided']:<12.2e}", flush=True) + + # summary json + summary = { + "target": "Dingwall_GSE220977", + "method": "Seurat-replica-identified 14,251 dermal cells (Dingwall clusters {0,1,3,4,5,8,11,20}); " + "per-Derm identity Fisher-exact cKO depletion using replica-assigned Derm labels " + "(mapped via Jaccard on top-50 markers to Dingwall Data S1C)", + "n_dermal_cells_total": int(dermal.n_obs), + "baseline_cko_frac": float(baseline), + "n_WT_dermal": n_wt, "n_cKO_dermal": n_cko, + "per_derm": rows, + } + (out / "105_primary_eden_full_dermal.json").write_text(json.dumps(summary, indent=2, default=str)) + print(f"\n[eden] wrote {out}/105_primary_eden_full_dermal.{{csv,json}}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/106_melanoblast_neural_crest.py b/scripts/analysis/106_melanoblast_neural_crest.py new file mode 100644 index 0000000000000000000000000000000000000000..fad2d211e18834f716ac0d5830e3e9658da8990c --- /dev/null +++ b/scripts/analysis/106_melanoblast_neural_crest.py @@ -0,0 +1,138 @@ +"""test whether En1-cKO melanoblasts with sweat/eda derepression revert to neural crest vs disrupted melanocyte.""" +from __future__ import annotations + +import json +import warnings +from pathlib import Path + +import anndata as ad +import numpy as np +import pandas as pd +import scanpy as sc +from scipy.stats import mannwhitneyu + +warnings.filterwarnings("ignore") +sc.settings.verbosity = 0 + +ROOT = Path("/home/bcheng/PRISM") +OUT = ROOT / "discovery/pan_skin/marker" +OUT.mkdir(parents=True, exist_ok=True) + +RAW = ROOT / "data/raw/GSE220977_combined.h5ad" +PRED = ROOT / "discovery/pan_skin/marker/dingwall_predictions.csv" +CKO_GSMS = {"GSM6833482", "GSM6833483"} +WT_GSMS = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"} + +MODULES = { + "Neural_crest": ["Sox10", "Sox9", "Sox2", "Pax3", "Foxd3", "Nes", "Tfap2a"], + "Melanogenesis_late": ["Tyrp1", "Slc45a2", "Oca2", "Gpnmb", "Pmel", "Silv", + "Mlph", "Rab27a", "Melana"], + "MITF_regulon": ["Mitf", "Dct", "Tyr", "Pmel", "Mlana", "Tyrp1", + "Slc24a5", "Slc45a2", "Sox10", "Pax3", "Kit", "Ednrb"], + "Sweat_gland": ["Foxi3", "Foxa1", "En1", "Krt8", "Krt18", "Krt19", + "Muc5b", "Aqp5", "Cutl1"], + "Eda_ectodysplasin": ["Eda", "Edar", "Edaradd", "Nfkb1", "Nfkb2", "Rela"], +} + + +def score(sub, name, genes): + present = [g for g in genes if g in sub.var_names] + if not present: + sub.obs[f"pw_{name}"] = 0.0 + return + sc.tl.score_genes(sub, gene_list=present, score_name=f"pw_{name}", + random_state=0, use_raw=False) + + +def main(): + print("[load] Dingwall + predictions", flush=True) + a = ad.read_h5ad(RAW) + pred = pd.read_csv(PRED) + pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"])) + a.obs["pred_label"] = pd.Categorical( + [pred_map.get(c, "unknown") for c in a.obs_names]) + samp = a.obs["sample"].astype(str) + a.obs["group"] = np.where(samp.isin(list(CKO_GSMS)), "En1-cKO", + np.where(samp.isin(list(WT_GSMS)), "WT", "other")) + a = a[a.obs["group"].isin(["En1-cKO", "WT"])].copy() + sub = a[a.obs["pred_label"].astype(str) == "melanoblast"].copy() + print(f"[load] {sub.n_obs} melanoblast cells " + f"(cKO={(sub.obs['group']=='En1-cKO').sum()}, " + f"WT={(sub.obs['group']=='WT').sum()})", flush=True) + + sc.pp.normalize_total(sub, target_sum=1e4) + sc.pp.log1p(sub) + + for name, genes in MODULES.items(): + score(sub, name, genes) + + sub.obs["derepression"] = (sub.obs["pw_Sweat_gland"].astype(float) + + sub.obs["pw_Eda_ectodysplasin"].astype(float)) + + df = sub.obs[[ + "group", "pw_Neural_crest", "pw_Melanogenesis_late", "pw_MITF_regulon", + "pw_Sweat_gland", "pw_Eda_ectodysplasin", "derepression" + ]].copy() + df.reset_index().rename(columns={"index": "cell_id"}).to_csv( + OUT / "106_melanoblast_nc_scores.csv", index=False) + + # baseline cKO vs WT per module + tests = {} + for m in ["Neural_crest", "Melanogenesis_late", "MITF_regulon", + "Sweat_gland", "Eda_ectodysplasin"]: + v1 = df.loc[df["group"] == "En1-cKO", f"pw_{m}"].astype(float).values + v0 = df.loc[df["group"] == "WT", f"pw_{m}"].astype(float).values + _, p = mannwhitneyu(v1, v0, alternative="two-sided") + tests[m] = { + "delta_cKO_minus_WT": float(v1.mean() - v0.mean()), + "mean_cKO": float(v1.mean()), "mean_WT": float(v0.mean()), + "n_cKO": int(len(v1)), "n_WT": int(len(v0)), + "mannu_p": float(p), + } + + # within cKO, split by derepression quartile + cko = df[df["group"] == "En1-cKO"].copy() + q1 = cko["derepression"].quantile(0.25) + q4 = cko["derepression"].quantile(0.75) + top = cko[cko["derepression"] >= q4] + bot = cko[cko["derepression"] <= q1] + print(f"[q] top-derep cKO: n={len(top)}, bot-derep cKO: n={len(bot)}", + flush=True) + + within = {} + for m in ["Neural_crest", "Melanogenesis_late", "MITF_regulon"]: + vt = top[f"pw_{m}"].astype(float).values + vb = bot[f"pw_{m}"].astype(float).values + _, p = mannwhitneyu(vt, vb, alternative="two-sided") + within[m] = { + "mean_top_derep": float(vt.mean()), + "mean_bot_derep": float(vb.mean()), + "delta_top_minus_bot": float(vt.mean() - vb.mean()), + "n_top": int(len(vt)), "n_bot": int(len(vb)), + "mannu_p": float(p), + } + + nc_delta = within["Neural_crest"]["delta_top_minus_bot"] + mel_delta = within["Melanogenesis_late"]["delta_top_minus_bot"] + verdict = ( + "novel_neural_crest_reversion" if nc_delta > 0.03 and mel_delta > -0.02 + else "basic_melanocyte_disruption" if nc_delta < 0.01 and mel_delta < -0.02 + else "mixed_or_orthogonal" + ) + summary = { + "n_melanoblast_cells": int(sub.n_obs), + "n_cKO_melanoblasts": int((df["group"] == "En1-cKO").sum()), + "n_WT_melanoblasts": int((df["group"] == "WT").sum()), + "baseline_cKO_vs_WT": tests, + "within_cKO_top_vs_bot_derepression_quartile": within, + "verdict": verdict, + } + with open(OUT / "106_melanoblast_nc_summary.json", "w") as f: + json.dump(summary, f, indent=2) + + print("[done] verdict:", verdict, flush=True) + print(json.dumps(summary, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/107_dingwall_class_deg_count.py b/scripts/analysis/107_dingwall_class_deg_count.py new file mode 100644 index 0000000000000000000000000000000000000000..ca0a94b3b9f6abcd28f7d6a8311ac84ef833e854 --- /dev/null +++ b/scripts/analysis/107_dingwall_class_deg_count.py @@ -0,0 +1,116 @@ +"""rank dingwall cell types by En1-cKO vs WT DEG count (|log2FC|>1, padj<0.05).""" +from __future__ import annotations + +import json +import warnings +from pathlib import Path + +import anndata as ad +import numpy as np +import pandas as pd +import scanpy as sc + +warnings.filterwarnings("ignore") +sc.settings.verbosity = 0 + +ROOT = Path("/home/bcheng/PRISM") +OUT = ROOT / "discovery/pan_skin/marker" +OUT.mkdir(parents=True, exist_ok=True) + +RAW = ROOT / "data/raw/GSE220977_combined.h5ad" +PRED = ROOT / "discovery/pan_skin/marker/dingwall_predictions.csv" +CKO_GSMS = {"GSM6833482", "GSM6833483"} +WT_GSMS = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"} +MIN_PER_GROUP = 15 +LFC_THRESH = 1.0 +PADJ_THRESH = 0.05 + + +def main(): + print("[load]", flush=True) + a = ad.read_h5ad(RAW) + pred = pd.read_csv(PRED) + pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"])) + a.obs["pred_label"] = pd.Categorical( + [pred_map.get(c, "unknown") for c in a.obs_names]) + samp = a.obs["sample"].astype(str) + a.obs["group"] = np.where(samp.isin(list(CKO_GSMS)), "En1-cKO", + np.where(samp.isin(list(WT_GSMS)), "WT", "other")) + a = a[a.obs["group"].isin(["En1-cKO", "WT"])].copy() + print(f"[load] {a.n_obs} cells, {a.n_vars} genes", flush=True) + + sc.pp.normalize_total(a, target_sum=1e4) + sc.pp.log1p(a) + + rows = [] + for cls in sorted(a.obs["pred_label"].astype(str).unique()): + mask = (a.obs["pred_label"].astype(str) == cls) + n_cko = int((mask & (a.obs["group"] == "En1-cKO")).sum()) + n_wt = int((mask & (a.obs["group"] == "WT")).sum()) + if n_cko < MIN_PER_GROUP or n_wt < MIN_PER_GROUP: + print(f"[skip] {cls}: n_cKO={n_cko} n_WT={n_wt}", flush=True) + continue + sub = a[mask].copy() + sub.obs["group"] = pd.Categorical(sub.obs["group"].values, + categories=["En1-cKO", "WT"]) + # reference=WT so positive LFC means up in cKO + try: + sc.tl.rank_genes_groups(sub, "group", reference="WT", + groups=["En1-cKO"], method="wilcoxon", + n_genes=sub.n_vars, use_raw=False, + pts=True) + rgg = sub.uns["rank_genes_groups"] + df = pd.DataFrame({ + "gene": [x[0] for x in rgg["names"]], + "lfc": [x[0] for x in rgg["logfoldchanges"]], + "padj": [x[0] for x in rgg["pvals_adj"]], + }) + up = int(((df["lfc"] > LFC_THRESH) & (df["padj"] < PADJ_THRESH)).sum()) + down = int(((df["lfc"] < -LFC_THRESH) & (df["padj"] < PADJ_THRESH)).sum()) + total = up + down + top_up = df[(df["lfc"] > LFC_THRESH) & (df["padj"] < PADJ_THRESH)] \ + .sort_values("lfc", ascending=False)["gene"].head(10).tolist() + top_down = df[(df["lfc"] < -LFC_THRESH) & (df["padj"] < PADJ_THRESH)] \ + .sort_values("lfc", ascending=True)["gene"].head(10).tolist() + except Exception as e: + print(f"[fail] {cls}: {e}", flush=True) + continue + rows.append({ + "class": cls, + "n_cKO": n_cko, + "n_WT": n_wt, + "n_DEG": total, + "n_up": up, + "n_down": down, + "top_up": ";".join(top_up), + "top_down": ";".join(top_down), + }) + print(f"[ok] {cls}: n_cKO={n_cko} n_WT={n_wt} DEG={total} " + f"(up={up}, down={down})", flush=True) + + df_out = pd.DataFrame(rows).sort_values("n_DEG", ascending=False) + df_out.to_csv(OUT / "107_dingwall_class_deg_count.csv", index=False) + print("\n[rank] classes ordered by DEG count (|LFC|>1 padj<0.05):") + print(df_out[["class", "n_cKO", "n_WT", "n_DEG", "n_up", "n_down"]] + .to_string(index=False)) + + if not df_out.empty: + winner = df_out.iloc[0] + summary = { + "lfc_threshold": LFC_THRESH, "padj_threshold": PADJ_THRESH, + "n_classes_tested": int(len(df_out)), + "top_class": str(winner["class"]), + "top_n_DEG": int(winner["n_DEG"]), + "top_up": winner["top_up"], + "top_down": winner["top_down"], + "ranking": df_out[["class", "n_DEG"]].to_dict("records"), + } + else: + summary = {"error": "no eligible classes"} + with open(OUT / "107_dingwall_class_deg_count.json", "w") as f: + json.dump(summary, f, indent=2) + print("[done]", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/108_dahlin_lineage_metabolism.py b/scripts/analysis/108_dahlin_lineage_metabolism.py new file mode 100644 index 0000000000000000000000000000000000000000..3d4901aef09fc99db2004b14b8d7c7792bfa7880 --- /dev/null +++ b/scripts/analysis/108_dahlin_lineage_metabolism.py @@ -0,0 +1,158 @@ +"""rank dahlin lineages by Kit_W41 vs WT metabolic shift across 4 modules.""" +from __future__ import annotations + +import json +import warnings +from pathlib import Path + +import anndata as ad +import numpy as np +import pandas as pd +import scanpy as sc +import scipy.sparse as sp +from scipy.stats import mannwhitneyu + +warnings.filterwarnings("ignore") +sc.settings.verbosity = 0 + +ROOT = Path("/home/bcheng/PRISM") +OUT = ROOT / "discovery/hematopoiesis/marker" +OUT.mkdir(parents=True, exist_ok=True) + +D_DIR = ROOT / "data/corpus/hematopoiesis/held_out_unlabeled/dahlin_extract" +PRED = ROOT / "discovery/hematopoiesis/marker/dahlin_predictions.csv" +GT = {"SIGAB1": "WT", "SIGAC1": "WT", "SIGAD1": "WT", "SIGAF1": "WT", + "SIGAG1": "WT", "SIGAH1": "WT", + "SIGAG8": "Kit_W41", "SIGAH8": "Kit_W41"} + +MIN_PER_GROUP = 15 + +MODULES = { + "OXPHOS_ETC": ["Ndufa1", "Ndufa2", "Ndufb1", "Ndufb2", "Sdha", "Sdhb", + "Cox4i1", "Cox5a", "Cox6a1", "Atp5a1", "Atp5b", "Uqcrq"], + "Glycolysis": ["Hk1", "Hk2", "Pfkm", "Pfkl", "Aldoa", "Gapdh", "Pgk1", + "Pkm", "Ldha", "Eno1", "Tpi1", "Pgam1"], + "Fatty_acid_oxidation":["Cpt1a", "Acadm", "Acadl", "Acadvl", "Hadha", "Hadhb", + "Ppara", "Ppargc1a", "Ucp2"], + "Redox_glutathione": ["Gpx1", "Gpx2", "Gpx3", "Gpx4", "Gsr", "Prdx1", "Prdx2", + "Prdx3", "Prdx4", "Prdx5", "Prdx6", "Sod1", "Sod2", "Cat"], +} + + +def load_dahlin(): + print("[load] Dahlin raw counts", flush=True) + parts = [] + for f in sorted(D_DIR.glob("*.txt.gz")): + sample = f.name.split("_")[1].split(".")[0] + 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["group"] = GT.get(sample, "unknown") + var = pd.DataFrame(index=df.index.astype(str)) + parts.append(ad.AnnData(X=X, obs=obs, var=var)) + print(f"[load] {sample}: {X.shape}", flush=True) + a = ad.concat(parts, join="outer", label="_batch") + print(f"[load] concat: {a.shape}", flush=True) + + import mygene + mg = mygene.MyGeneInfo() + res = mg.querymany(a.var_names.astype(str).tolist(), 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) + a = a[:, keep].copy() + a.var_names = syms[keep] + a.var_names_make_unique() + print(f"[map] {a.n_vars} genes with symbols", flush=True) + a = a[a.obs["group"].isin(["Kit_W41", "WT"])].copy() + return a + + +def main(): + a = load_dahlin() + pred = pd.read_csv(PRED) + common = a.obs_names.intersection(pd.Index(pred["cell_id"].astype(str))) + a = a[list(common)].copy() + pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"])) + a.obs["pred_label"] = pd.Categorical( + [pred_map.get(c, "unknown") for c in a.obs_names]) + print(f"[join] {a.n_obs} cells, " + f"Kit_W41={int((a.obs['group']=='Kit_W41').sum())}, " + f"WT={int((a.obs['group']=='WT').sum())}", flush=True) + + sc.pp.normalize_total(a, target_sum=1e4) + sc.pp.log1p(a) + + rows = [] + for cls in sorted(a.obs["pred_label"].astype(str).unique()): + mask = (a.obs["pred_label"].astype(str) == cls).values + n_k = int((mask & (a.obs["group"].values == "Kit_W41")).sum()) + n_w = int((mask & (a.obs["group"].values == "WT")).sum()) + if n_k < MIN_PER_GROUP or n_w < MIN_PER_GROUP: + print(f"[skip] {cls}: n_Kit={n_k} n_WT={n_w}", flush=True) + continue + sub = a[mask].copy() + for name, genes in MODULES.items(): + present = [g for g in genes if g in sub.var_names] + if not present: + sub.obs[f"pw_{name}"] = 0.0 + continue + sc.tl.score_genes(sub, gene_list=present, + score_name=f"pw_{name}", random_state=0, + use_raw=False) + + grp = sub.obs["group"].values + for name in MODULES: + v_k = sub.obs[f"pw_{name}"].astype(float).values[grp == "Kit_W41"] + v_w = sub.obs[f"pw_{name}"].astype(float).values[grp == "WT"] + try: + _, p = mannwhitneyu(v_k, v_w, alternative="two-sided") + except Exception: + p = 1.0 + rows.append({ + "class": cls, + "module": name, + "delta": float(v_k.mean() - v_w.mean()), + "mean_Kit_W41": float(v_k.mean()), + "mean_WT": float(v_w.mean()), + "n_Kit_W41": n_k, + "n_WT": n_w, + "mannu_p": float(p), + }) + print(f"[ok] {cls}: n_Kit={n_k} n_WT={n_w}", flush=True) + + df = pd.DataFrame(rows) + n_tests = len(df) + df["mannu_p_adj_bonferroni"] = np.minimum(df["mannu_p"] * n_tests, 1.0) + df.to_csv(OUT / "108_dahlin_lineage_metabolism.csv", index=False) + + # per-class composite magnitude + piv = df.pivot(index="class", columns="module", values="delta").fillna(0.0) + piv["sum_abs_delta"] = piv.abs().sum(axis=1) + piv["l2_delta"] = np.sqrt((piv[list(MODULES)] ** 2).sum(axis=1)) + piv_sorted = piv.sort_values("sum_abs_delta", ascending=False) + piv_sorted.to_csv(OUT / "108_dahlin_lineage_metabolism_ranked.csv") + + print("\n[ranked] lineages by sum |delta| across 4 metabolic modules:") + print(piv_sorted.round(4).to_string()) + + top_cls = piv_sorted.index[0] + top_row = piv_sorted.iloc[0] + summary = { + "modules": list(MODULES), + "top_lineage": str(top_cls), + "top_sum_abs_delta": float(top_row["sum_abs_delta"]), + "top_deltas_per_module": { + m: float(top_row[m]) for m in MODULES + }, + "ranking": piv_sorted[["sum_abs_delta"]].reset_index().to_dict("records"), + } + with open(OUT / "108_dahlin_lineage_metabolism.json", "w") as f: + json.dump(summary, f, indent=2) + print("[done] top lineage:", top_cls, flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/109_veres_mature_beta.py b/scripts/analysis/109_veres_mature_beta.py new file mode 100644 index 0000000000000000000000000000000000000000..67b081191c3e6bbb71ebea60230153af73f1dc68 --- /dev/null +++ b/scripts/analysis/109_veres_mature_beta.py @@ -0,0 +1,113 @@ +"""split veres beta-predicted cells into mature (MAFA/UCN3 hi) vs SC-beta by INS x MAT quadrants.""" +from __future__ import annotations + +import json +import warnings +from pathlib import Path + +import anndata as ad +import numpy as np +import pandas as pd + +warnings.filterwarnings("ignore") + +ROOT = Path("/home/bcheng/PRISM") +OUT = ROOT / "discovery/pancreas/marker" +OUT.mkdir(parents=True, exist_ok=True) + +VERES = ROOT / "data/corpus/pancreas/held_out_labeled/veres_GSE114412_test.h5ad" +PRED = ROOT / "discovery/pancreas/marker/veres_predictions.csv" + + +def z(x): + x = np.asarray(x, dtype=float) + s = x.std() + return (x - x.mean()) / (s if s > 0 else 1.0) + + +def main(): + print("[load]", flush=True) + a = ad.read_h5ad(VERES) + pred = pd.read_csv(PRED) + pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"])) + a.obs["pred_label"] = pd.Categorical( + [pred_map.get(c, "unknown") for c in a.obs_names]) + + print("[filter] pred_label == beta", flush=True) + sub = a[a.obs["pred_label"].astype(str) == "beta"].copy() + print(f"[filter] {sub.n_obs} beta-predicted cells", flush=True) + + # veres X is already log-normalised (range 0..9) + def col(g): + if g not in sub.var_names: + return np.zeros(sub.n_obs) + j = sub.var_names.get_loc(g) + x = sub.X[:, j] + if hasattr(x, "toarray"): + x = x.toarray() + return np.asarray(x).ravel() + + ins1 = col("Ins1"); ins2 = col("Ins2"); iapp = col("Iapp") + mafa = col("Mafa"); ucn3 = col("Ucn3") + insulin = ins1 + ins2 + mature = z(mafa) + z(ucn3) + + # split thresholds are the median within the veres beta-predicted set + ins_thr = float(np.median(insulin)) + mat_thr = float(np.median(mature)) + ins_hi = insulin > ins_thr + mat_hi = mature > mat_thr + + quadrant = np.array( + ["INS+/MAT+" if (ih and mh) else + "INS+/MAT-" if (ih and not mh) else + "INS-/MAT+" if (not ih and mh) else + "INS-/MAT-" + for ih, mh in zip(ins_hi, mat_hi)]) + sub.obs["INS_level"] = insulin + sub.obs["MAT_score"] = mature + sub.obs["quadrant"] = pd.Categorical(quadrant) + + df = sub.obs[["quadrant", "INS_level", "MAT_score"]].copy() + df["Ins1"] = ins1; df["Ins2"] = ins2; df["Iapp"] = iapp + df["Mafa"] = mafa; df["Ucn3"] = ucn3 + df["paper_label"] = sub.obs["paper_label"].astype(str).values + df.reset_index().rename(columns={"index": "cell_id"}).to_csv( + OUT / "109_veres_mature_beta_scores.csv", index=False) + + counts = df["quadrant"].value_counts().to_dict() + paper_by_quad = df.groupby(["quadrant", "paper_label"], observed=True) \ + .size().unstack(fill_value=0) + paper_by_quad.to_csv(OUT / "109_veres_mature_beta_paper_x_quadrant.csv") + + n_total = int(sub.n_obs) + n_mature = int((df["quadrant"] == "INS+/MAT+").sum()) + n_scbeta = int((df["quadrant"] == "INS+/MAT-").sum()) + + fraction_mature = n_mature / n_total if n_total else 0.0 + + means = df.groupby("quadrant", observed=True)[ + ["Mafa", "Ucn3", "Ins1", "Ins2", "Iapp"]].mean().round(3).to_dict() + + summary = { + "n_beta_predicted": n_total, + "ins_threshold": ins_thr, + "mat_threshold": mat_thr, + "quadrant_counts": counts, + "n_mature_INS+MAT+": n_mature, + "n_SCbeta_INS+MAT-": n_scbeta, + "fraction_mature": round(fraction_mature, 4), + "mean_expression_per_quadrant": means, + "paper_label_x_quadrant": { + q: paper_by_quad.loc[q].to_dict() + for q in paper_by_quad.index + } if len(paper_by_quad) else {}, + } + with open(OUT / "109_veres_mature_beta_summary.json", "w") as f: + json.dump(summary, f, indent=2) + print("\n[done]") + print(json.dumps(summary, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/110_veres_polyhormonal_alpha.py b/scripts/analysis/110_veres_polyhormonal_alpha.py new file mode 100644 index 0000000000000000000000000000000000000000..8d9c082657d787f9dba87338947084869d22eba0 --- /dev/null +++ b/scripts/analysis/110_veres_polyhormonal_alpha.py @@ -0,0 +1,125 @@ +"""test if veres alpha-pool polyhormonal (Ins+/Gcg+/Sst+) cells form a distinct sub-cluster vs graded.""" +from __future__ import annotations + +import json +import warnings +from pathlib import Path + +import anndata as ad +import numpy as np +import pandas as pd +import scanpy as sc + +warnings.filterwarnings("ignore") +sc.settings.verbosity = 0 + +ROOT = Path("/home/bcheng/PRISM") +OUT = ROOT / "discovery/pancreas/marker" +OUT.mkdir(parents=True, exist_ok=True) + +VERES = ROOT / "data/corpus/pancreas/held_out_labeled/veres_GSE114412_test.h5ad" +PRED = ROOT / "discovery/pancreas/marker/veres_predictions.csv" + + +def col(sub, g): + if g not in sub.var_names: + return np.zeros(sub.n_obs) + j = sub.var_names.get_loc(g) + x = sub.X[:, j] + if hasattr(x, "toarray"): + x = x.toarray() + return np.asarray(x).ravel() + + +def main(): + print("[load]", flush=True) + a = ad.read_h5ad(VERES) + pred = pd.read_csv(PRED) + pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"])) + a.obs["pred_label"] = pd.Categorical( + [pred_map.get(c, "unknown") for c in a.obs_names]) + + sub = a[a.obs["pred_label"].astype(str).isin( + ["alpha_progenitor", "alpha"])].copy() + print(f"[filter] {sub.n_obs} alpha-pool cells", flush=True) + + ins1 = col(sub, "Ins1"); ins2 = col(sub, "Ins2") + gcg = col(sub, "Gcg"); sst = col(sub, "Sst"); iapp = col(sub, "Iapp") + + ins_level = ins1 + ins2 + q_ins = np.quantile(ins_level, 0.75) + q_gcg = np.quantile(gcg, 0.75) + q_sst = np.quantile(sst, 0.75) + + n_pos = ( + (ins_level >= q_ins).astype(int) + + (gcg >= q_gcg).astype(int) + + (sst >= q_sst).astype(int) + ) + sub.obs["INS_level"] = ins_level + sub.obs["GCG_level"] = gcg + sub.obs["SST_level"] = sst + sub.obs["n_hormones_positive"] = n_pos + + print("[cluster] PCA + Leiden", flush=True) + sc.pp.highly_variable_genes(sub, n_top_genes=2000, flavor="seurat_v3", + subset=False, batch_key=None) + sc.pp.pca(sub, n_comps=30) + sc.pp.neighbors(sub, n_neighbors=15, n_pcs=30) + sc.tl.leiden(sub, resolution=0.5, random_state=0, key_added="leiden_alpha") + + df = sub.obs[[ + "pred_label", "paper_label", "leiden_alpha", + "INS_level", "GCG_level", "SST_level", "n_hormones_positive" + ]].copy() + df["Ins1"] = ins1; df["Ins2"] = ins2; df["Gcg"] = gcg + df["Sst"] = sst; df["Iapp"] = iapp + df.reset_index().rename(columns={"index": "cell_id"}).to_csv( + OUT / "110_veres_polyhormonal_alpha_scores.csv", index=False) + + baseline_polyhormonal = float((df["n_hormones_positive"] >= 2).mean()) + per_clus = df.groupby("leiden_alpha", observed=True).agg( + n_cells=("n_hormones_positive", "size"), + frac_polyhormonal=("n_hormones_positive", + lambda s: float((s >= 2).mean())), + frac_gcg_hi=("GCG_level", + lambda s: float((s >= q_gcg).mean())), + frac_ins_hi=("INS_level", + lambda s: float((s >= q_ins).mean())), + frac_sst_hi=("SST_level", + lambda s: float((s >= q_sst).mean())), + mean_gcg=("GCG_level", "mean"), + mean_ins=("INS_level", "mean"), + mean_sst=("SST_level", "mean"), + ).sort_values("frac_polyhormonal", ascending=False).reset_index() + per_clus["enrichment_vs_baseline"] = per_clus["frac_polyhormonal"] \ + / max(baseline_polyhormonal, 1e-6) + per_clus.to_csv(OUT / "110_veres_polyhormonal_alpha_per_cluster.csv", + index=False) + + n_2x_clusters = int((per_clus["enrichment_vs_baseline"] >= 2.0).sum()) + verdict = ("distinct_polyhormonal_subcluster" if n_2x_clusters in (1, 2) + else "graded_phenotype" if n_2x_clusters == 0 + else "diffuse_enrichment") + + summary = { + "n_alpha_pool": int(sub.n_obs), + "baseline_polyhormonal_frac": round(baseline_polyhormonal, 4), + "q75_thresholds": {"Ins": float(q_ins), + "Gcg": float(q_gcg), + "Sst": float(q_sst)}, + "leiden_resolution": 0.5, + "n_clusters": int(per_clus["leiden_alpha"].nunique()), + "n_clusters_enriched_2x": n_2x_clusters, + "verdict": verdict, + "per_cluster": per_clus.to_dict("records"), + } + with open(OUT / "110_veres_polyhormonal_alpha_summary.json", "w") as f: + json.dump(summary, f, indent=2) + + print("\n[done] verdict:", verdict, flush=True) + print(per_clus.round(3).to_string(index=False)) + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/44_en1_cko_contrast.py b/scripts/analysis/44_en1_cko_contrast.py new file mode 100644 index 0000000000000000000000000000000000000000..dbecc0ddd77e2dddf85a9444c905cbdeabe3e19a --- /dev/null +++ b/scripts/analysis/44_en1_cko_contrast.py @@ -0,0 +1,117 @@ +"""en1-cKO vs WT contrast on aldrich predictions.""" +from __future__ import annotations +from pathlib import Path +import warnings +warnings.filterwarnings("ignore") + +import numpy as np +import pandas as pd +import anndata as ad +import scanpy as sc +from scipy import stats + +TARGET = Path("/home/bcheng/PRISM/data/processed/skin/adata_processed.h5ad") +PROJ = Path("/home/bcheng/PRISM/discovery/pan_skin/marker/50_aldrich_projections.h5ad") +NOVEL = Path("/home/bcheng/PRISM/discovery/pan_skin/marker/51_aldrich_novel_annotation.csv") +OUT = Path("/home/bcheng/PRISM/discovery/pan_skin/marker") + + +def main(): + a = ad.read_h5ad(TARGET) + p = ad.read_h5ad(PROJ) + for c in ["pred_label", "pred_conf", "abstain", "pred_bbse_label"]: + if c in p.obs.columns: + a.obs[c] = p.obs[c].values + if NOVEL.exists(): + nov = pd.read_csv(NOVEL, index_col=0) + a.obs["novel_cluster"] = nov.reindex(a.obs_names)["novel_cluster"].values + + print(f"[cko] target: {a.shape}, genotype: {a.obs['genotype'].value_counts().to_dict()}", + flush=True) + + rows = [] + for cls in a.obs["pred_bbse_label"].unique(): + mask = a.obs["pred_bbse_label"] == cls + n_cko = int((mask & (a.obs["genotype"] == "En1-cKO")).sum()) + n_wt = int((mask & (a.obs["genotype"] == "WT")).sum()) + total_cko = (a.obs["genotype"] == "En1-cKO").sum() + total_wt = (a.obs["genotype"] == "WT").sum() + contingency = np.array([[n_cko, total_cko - n_cko], + [n_wt, total_wt - n_wt]]) + odds, p_val = stats.fisher_exact(contingency) + f_cko = (n_cko + 1) / (total_cko + 2) + f_wt = (n_wt + 1) / (total_wt + 2) + log2_fc = np.log2(f_cko / f_wt) + rows.append({ + "class": cls, + "n_En1cKO": n_cko, + "n_WT": n_wt, + "pct_En1cKO": round(100 * n_cko / total_cko, 2), + "pct_WT": round(100 * n_wt / total_wt, 2), + "log2_fold_enrich_cKO_vs_WT": round(log2_fc, 3), + "fisher_pvalue": p_val, + }) + df = pd.DataFrame(rows).sort_values("log2_fold_enrich_cKO_vs_WT") + print("\n[cko] class enrichment (cKO vs WT):") + print(df.to_string(index=False)) + + df.to_csv(OUT / "53_en1_cko_class_enrichment.csv", index=False) + + de_rows = [] + for cls in sorted(a.obs["pred_bbse_label"].unique()): + cls_mask = a.obs["pred_bbse_label"] == cls + if cls_mask.sum() < 50: + continue + sub = a[cls_mask].copy() + vc = sub.obs["genotype"].value_counts() + if not {"En1-cKO", "WT"}.issubset(vc.index) or vc.min() < 15: + continue + try: + sc.tl.rank_genes_groups(sub, "genotype", method="wilcoxon", + n_genes=40, use_raw=False) + for grp in ["En1-cKO", "WT"]: + if grp not in sub.uns["rank_genes_groups"]["names"].dtype.names: + continue + names = list(sub.uns["rank_genes_groups"]["names"][grp][:15]) + lfcs = list(sub.uns["rank_genes_groups"]["logfoldchanges"][grp][:15]) + for g, lf in zip(names, lfcs): + de_rows.append({ + "class": cls, "up_in": grp, "gene": g, "logfc": round(float(lf), 3), + }) + except Exception as exc: + print(f"[cko] DE failed for {cls}: {exc}") + continue + de_df = pd.DataFrame(de_rows) + de_df.to_csv(OUT / "53_en1_cko_wilcoxon_within_class.csv", index=False) + + md = ["# Aldrich En1-cKO vs WT contrast on PANDA-MLP predictions\n", + f"Total cells: {a.n_obs:,} ({int((a.obs['genotype']=='En1-cKO').sum())} En1-cKO, " + f"{int((a.obs['genotype']=='WT').sum())} WT).\n", + "## Class-level cKO/WT enrichment (BBSE-corrected predictions)\n", + df.to_markdown(index=False), "", + "The direction of `log2_fold_enrich_cKO_vs_WT` indicates whether a class is over-represented", + "in En1-cKO (positive) or WT (negative). Fisher exact p-value tests significance vs the", + "background genotype ratio (~40% cKO / 60% WT).\n", + "## Per-class Wilcoxon DE (En1-cKO vs WT within each class)\n", + "Top genes differentially expressed BETWEEN genotypes WITHIN a predicted class. Genes up in", + "cKO reveal En1-loss-responsive programs specific to that cell type; genes up in WT are the", + "opposite.\n", + ] + if len(de_df): + for cls in sorted(de_df["class"].unique()): + md.append(f"\n### {cls}\n") + for grp in ["En1-cKO", "WT"]: + sub = de_df[(de_df["class"] == cls) & (de_df["up_in"] == grp)] + if not len(sub): + continue + md.append(f"**Up in {grp}**: " + ", ".join(sub["gene"].tolist())) + md += ["", "## Novel population x genotype cross-tab\n"] + if "novel_cluster" in a.obs.columns: + xt = pd.crosstab(a.obs["novel_cluster"], a.obs["genotype"]) + md.append(xt.to_markdown()) + (OUT / "53_en1_cko_contrast.md").write_text("\n".join(md)) + print(f"[cko] wrote {OUT}/53_en1_cko_contrast.md") + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/45_marker_refinement.py b/scripts/analysis/45_marker_refinement.py new file mode 100644 index 0000000000000000000000000000000000000000..b91b412271dfc5f118060ebbea084a4be85bdcb3 --- /dev/null +++ b/scripts/analysis/45_marker_refinement.py @@ -0,0 +1,70 @@ +"""marker refinement on aldrich zero-shot predictions: per-class wilcoxon DE vs canonical list.""" +from __future__ import annotations +from pathlib import Path +import warnings, yaml +warnings.filterwarnings("ignore") + +import numpy as np +import pandas as pd +import anndata as ad +import scanpy as sc + +TARGET = Path("/home/bcheng/PRISM/data/processed/skin/adata_processed.h5ad") +PROJ = Path("/home/bcheng/PRISM/discovery/pan_skin/marker/50_aldrich_projections.h5ad") +TFS = Path("/home/bcheng/PRISM/discovery/pan_skin/marker/known_skin_tfs.yaml") +OUT = Path("/home/bcheng/PRISM/discovery/pan_skin/marker") + + +def main(): + a = ad.read_h5ad(TARGET) + p = ad.read_h5ad(PROJ) + a.obs["pred_label"] = p.obs["pred_label"].values + a.obs["pred_conf"] = p.obs["pred_conf"].astype(float).values + + with open(TFS) as f: + tf_cfg = yaml.safe_load(f) + canon = tf_cfg["classes"] + + keep = a[~a.obs["pred_label"].str.startswith("UNK", na=False)].copy() + sc.tl.rank_genes_groups(keep, "pred_label", method="wilcoxon", + n_genes=100, use_raw=False) + + rows = [] + for cls in keep.obs["pred_label"].unique(): + try: + names = keep.uns["rank_genes_groups"]["names"][cls] + lfc = keep.uns["rank_genes_groups"]["logfoldchanges"][cls] + padj = keep.uns["rank_genes_groups"]["pvals_adj"][cls] + except Exception: + continue + top100 = list(names[:100]) + canon_set = set(canon.get(cls, [])) + retained = [g for g in canon_set if g in top100] + novel = [g for g in top100 if g not in canon_set] + mask = keep.obs["pred_label"] == cls + n_cells = int(mask.sum()) + mean_conf = float(keep.obs.loc[mask, "pred_conf"].mean()) + rows.append({ + "class": cls, + "n_cells_predicted": n_cells, + "mean_conf": round(mean_conf, 3), + "n_canonical": len(canon_set), + "n_canonical_retained_top100": len(retained), + "retained_canonical": ",".join(retained[:15]), + "novel_top20": ",".join(novel[:20]), + }) + df = pd.DataFrame(rows).sort_values("n_cells_predicted", ascending=False) + df.to_csv(OUT / "52_refined_markers.csv", index=False) + print(df.to_string(index=False)) + + md = ["# Refined pan-skin markers for Aldrich zero-shot predictions\n"] + md.append("Marker refinement is Wilcoxon DE of each predicted class against all other") + md.append("predicted cells on raw Aldrich gene expression. `retained_canonical` are canonical") + md.append("markers recovered in the top-100; `novel_top20` are DE genes not in the canonical list.\n") + md.append(df.to_markdown(index=False)) + (OUT / "52_refined_markers.md").write_text("\n".join(md)) + print(f"[refine] wrote {OUT}/52_refined_markers.csv and .md") + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/49_melanocyte_deep_dive.py b/scripts/analysis/49_melanocyte_deep_dive.py new file mode 100644 index 0000000000000000000000000000000000000000..68aee4a6fab3fa5fd18dcc3303707182cccca8e0 --- /dev/null +++ b/scripts/analysis/49_melanocyte_deep_dive.py @@ -0,0 +1,125 @@ +"""melanocyte deep dive on aldrich: within-melanocyte cKO vs WT DE and pathway scoring.""" +from __future__ import annotations +from pathlib import Path +import warnings +warnings.filterwarnings("ignore") + +import numpy as np +import pandas as pd +import anndata as ad +import scanpy as sc +from scipy.stats import mannwhitneyu + +TARGET = Path("/home/bcheng/PRISM/data/processed/skin/adata_processed.h5ad") +PROJ = Path("/home/bcheng/PRISM/discovery/pan_skin/marker/50_aldrich_projections.h5ad") +OUT = Path("/home/bcheng/PRISM/discovery/pan_skin/marker") + +PATHWAYS = { + "MITF_regulon": ["Mitf", "Dct", "Tyr", "Pmel", "Mlana", "Tyrp1", "Slc24a5", + "Slc45a2", "Silv", "Sox10", "Pax3", "Kit", "Ednrb"], + "Wnt_signaling": ["Wnt3", "Wnt5a", "Wnt7a", "Wnt10b", "Ctnnb1", "Lef1", "Tcf4", + "Tcf7", "Axin2", "Dkk1", "Sfrp1", "Fzd7", "Lrp5", "Lrp6"], + "BMP_TGF_signaling": ["Bmp2", "Bmp4", "Bmp5", "Bmp7", "Bmpr1a", "Bmpr1b", "Bmpr2", + "Smad1", "Smad3", "Smad5", "Tgfb1", "Tgfb2", "Tgfbr1", + "Tgfbr2", "Id1", "Id2", "Id3"], + "FGF_signaling": ["Fgf1", "Fgf2", "Fgf7", "Fgf9", "Fgf10", "Fgf20", "Fgfr1", + "Fgfr2", "Fgfr3", "Fgfr4", "Etv1", "Etv4", "Etv5", "Spry2", + "Spry4", "Dusp6"], + "Notch_signaling": ["Notch1", "Notch2", "Notch3", "Jag1", "Jag2", "Dll1", "Dll3", + "Dll4", "Hes1", "Hes5", "Hey1", "Hey2", "Rbpj", "Maml1"], + "Hedgehog": ["Shh", "Ihh", "Ptch1", "Ptch2", "Smo", "Gli1", "Gli2", "Gli3", + "Hhip", "Sufu"], + "EMT": ["Zeb1", "Zeb2", "Snai1", "Snai2", "Twist1", "Twist2", "Vim", + "Cdh2", "Fn1", "Mmp2", "Mmp9", "Prrx1"], + "Cell_cycle": ["Ccnd1", "Ccnd2", "Ccne1", "Ccna2", "Ccnb1", "Cdk1", "Cdk2", + "Cdk4", "Cdk6", "Mki67", "Top2a", "Pcna", "Mcm2", "Mcm3"], + "Neural_crest": ["Sox10", "Sox9", "Sox2", "Pax3", "Foxd3", "Nes", "Ngfr", + "Ednrb", "Kit", "Tfap2a"], + "Apoptosis": ["Bax", "Bak1", "Bad", "Bid", "Bcl2", "Bcl2l1", "Casp3", + "Casp9", "Casp8", "Fas", "Fasl", "Trp53", "Cdkn1a"], +} + + +def pathway_scoring(sub, pathway_dict): + for name, genes in pathway_dict.items(): + present = [g for g in genes if g in sub.var_names] + if not present: + sub.obs[f"{name}"] = 0.0 + continue + sc.tl.score_genes(sub, gene_list=present, score_name=f"{name}", + random_state=0, use_raw=False) + return sub + + +def main(): + a = ad.read_h5ad(TARGET) + p = ad.read_h5ad(PROJ) + a.obs["pred_label"] = p.obs["pred_bbse_label"].values + + mask = a.obs["pred_label"] == "melanocyte" + print(f"[mel] melanocyte cells: {int(mask.sum())} " + f"({int((mask & (a.obs['genotype']=='En1-cKO')).sum())} cKO, " + f"{int((mask & (a.obs['genotype']=='WT')).sum())} WT)", flush=True) + if mask.sum() < 30: + print("[mel] too few melanocyte cells") + return + + sub = a[mask].copy() + sub = pathway_scoring(sub, PATHWAYS) + + print("\n[mel] pathway score cKO vs WT (positive = up in cKO):") + rows = [] + for name in PATHWAYS.keys(): + s = sub.obs[name].astype(float).values + g = sub.obs["genotype"].values + c_scores = s[g == "En1-cKO"] + w_scores = s[g == "WT"] + stat, pval = mannwhitneyu(c_scores, w_scores, alternative="two-sided") + delta = c_scores.mean() - w_scores.mean() + rows.append({"pathway": name, "delta_cKO_minus_WT": round(delta, 4), + "MannU_p": pval, + "cKO_mean": round(c_scores.mean(), 4), + "WT_mean": round(w_scores.mean(), 4)}) + star = "***" if pval < 0.001 else "**" if pval < 0.01 else "*" if pval < 0.05 else "" + print(f" {name:20s} delta={delta:+.4f} p={pval:.2e} {star}") + + pd.DataFrame(rows).sort_values("MannU_p").to_csv(OUT / "56_melanocyte_pathways.csv", index=False) + + print("\n[mel] within-melanocyte cKO vs WT DE:") + sc.tl.rank_genes_groups(sub, "genotype", method="wilcoxon", n_genes=50, use_raw=False) + de = pd.DataFrame({ + "cKO_up_gene": sub.uns["rank_genes_groups"]["names"]["En1-cKO"][:20], + "cKO_up_lfc": sub.uns["rank_genes_groups"]["logfoldchanges"]["En1-cKO"][:20], + "WT_up_gene": sub.uns["rank_genes_groups"]["names"]["WT"][:20], + "WT_up_lfc": sub.uns["rank_genes_groups"]["logfoldchanges"]["WT"][:20], + }) + print(de.to_string(index=False)) + de.to_csv(OUT / "56_melanocyte_wilcoxon_cko_vs_wt.csv", index=False) + + lines = ["# Melanocyte deep dive — Aldrich En1-cKO vs WT", + "", + f"PANDA-MLP-predicted melanocyte cells: **{int(mask.sum())} total** " + f"({int((mask & (a.obs['genotype']=='En1-cKO')).sum())} cKO, " + f"{int((mask & (a.obs['genotype']=='WT')).sum())} WT).", + "", + "The v3b+v3c cross-model replication showed ~2× melanocyte enrichment in cKO " + "(log2fc +1.04, Fisher p ≈ 4×10⁻⁶). Here we probe what's happening WITHIN the " + "melanocyte compartment.", + "", + "## Pathway score comparisons (MannU cKO vs WT)", + ""] + dfp = pd.DataFrame(rows).sort_values("MannU_p") + lines.append(dfp.to_markdown(index=False)) + lines += ["", + "## Within-melanocyte Wilcoxon DE (top 20 each direction)", + ""] + de_show = de.copy() + de_show["cKO_up_lfc"] = de_show["cKO_up_lfc"].round(3) + de_show["WT_up_lfc"] = de_show["WT_up_lfc"].round(3) + lines.append(de_show.to_markdown(index=False)) + (OUT / "56_melanocyte_deep_dive.md").write_text("\n".join(lines)) + print(f"[mel] wrote {OUT}/56_melanocyte_deep_dive.md") + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/57_multiclass_pathway_analysis.py b/scripts/analysis/57_multiclass_pathway_analysis.py new file mode 100644 index 0000000000000000000000000000000000000000..3ba8f58d0910d9abeb040a324ddd33123b5250e8 --- /dev/null +++ b/scripts/analysis/57_multiclass_pathway_analysis.py @@ -0,0 +1,123 @@ +"""per-class pathway scoring cKO vs WT on aldrich, MannU per (class, pathway).""" +from __future__ import annotations +from pathlib import Path +import warnings +warnings.filterwarnings("ignore") + +import numpy as np +import pandas as pd +import anndata as ad +import scanpy as sc +from scipy.stats import mannwhitneyu + +TARGET = Path("/home/bcheng/PRISM/data/processed/skin/adata_processed.h5ad") +PROJ = Path("/home/bcheng/PRISM/discovery/pan_skin/marker/50_aldrich_projections.h5ad") +OUT = Path("/home/bcheng/PRISM/discovery/pan_skin/marker") + +CLASSES_OF_INTEREST = [ + "basal-IFE", "spinous", "granular", + "fibroblast-papillary", "fibroblast-reticular", + "endothelial", "immune", "melanocyte", +] + +PATHWAYS = { + "MITF_regulon": ["Mitf", "Dct", "Tyr", "Pmel", "Mlana", "Tyrp1", "Slc24a5", + "Slc45a2", "Sox10", "Pax3", "Kit", "Ednrb"], + "Wnt_signaling": ["Wnt3", "Wnt5a", "Wnt7a", "Wnt10b", "Ctnnb1", "Lef1", "Tcf4", + "Tcf7", "Axin2", "Dkk1", "Sfrp1", "Fzd7", "Lrp5"], + "BMP_signaling": ["Bmp2", "Bmp4", "Bmp5", "Bmp7", "Bmpr1a", "Bmpr1b", "Bmpr2", + "Smad1", "Smad5", "Id1", "Id2", "Id3"], + "TGFB_signaling": ["Tgfb1", "Tgfb2", "Tgfbr1", "Tgfbr2", "Smad3", "Smad7"], + "FGF_signaling": ["Fgf1", "Fgf2", "Fgf7", "Fgf9", "Fgf10", "Fgfr1", "Fgfr2", + "Etv1", "Etv4", "Etv5", "Spry2", "Dusp6"], + "Notch_signaling": ["Notch1", "Notch2", "Notch3", "Jag1", "Dll1", "Hes1", "Hes5", + "Hey1", "Hey2", "Rbpj"], + "Hedgehog": ["Shh", "Ptch1", "Smo", "Gli1", "Gli2", "Gli3"], + "Eda_ectodysplasin": ["Eda", "Edar", "Edaradd", "Nfkb1", "Nfkb2", "Rela"], + "EMT": ["Zeb1", "Zeb2", "Snai1", "Snai2", "Twist1", "Twist2", "Vim", + "Cdh2", "Fn1", "Prrx1"], + "Cell_cycle": ["Ccnd1", "Ccne1", "Ccna2", "Ccnb1", "Cdk1", "Cdk2", "Cdk4", + "Mki67", "Top2a", "Pcna", "Mcm2", "Mcm3"], + "KC_differentiation": ["Krt1", "Krt10", "Ivl", "Lor", "Flg", "Flg2", "Klk5", "Klk7", + "Cdsn"], + "Basal_keratinocyte": ["Krt5", "Krt14", "Krt15", "Trp63", "Itga6", "Itgb1", "Itga3"], + "Sweat_gland": ["Foxi3", "Foxa1", "En1", "Krt8", "Krt18", "Krt19", "Muc5b", + "Aqp5", "Cutl1"], + "Hair_placode": ["Wnt10b", "Shh", "Lef1", "Foxi3", "Edar", "Bmp4", "Msx2"], + "Neural_crest": ["Sox10", "Sox9", "Sox2", "Pax3", "Foxd3", "Nes", "Tfap2a"], + "Apoptosis": ["Bax", "Bak1", "Bad", "Bcl2", "Casp3", "Casp9", "Trp53", + "Cdkn1a"], +} + + +def pathway_scoring(sub, pathway_dict): + for name, genes in pathway_dict.items(): + present = [g for g in genes if g in sub.var_names] + if not present: + sub.obs[f"pw_{name}"] = 0.0 + continue + sc.tl.score_genes(sub, gene_list=present, score_name=f"pw_{name}", + random_state=0, use_raw=False) + return sub + + +def main(): + a = ad.read_h5ad(TARGET) + p = ad.read_h5ad(PROJ) + a.obs["pred_label"] = p.obs["pred_bbse_label"].values + print(f"[pw] classes in target: {a.obs['pred_label'].value_counts().to_dict()}", flush=True) + + rows = [] + for cls in CLASSES_OF_INTEREST: + mask = a.obs["pred_label"] == cls + n_c = int((mask & (a.obs["genotype"]=="En1-cKO")).sum()) + n_w = int((mask & (a.obs["genotype"]=="WT")).sum()) + if n_c < 15 or n_w < 15: + print(f"[pw] {cls}: skip (n_cKO={n_c}, n_WT={n_w})") + continue + sub = a[mask].copy() + sub = pathway_scoring(sub, PATHWAYS) + for pw in PATHWAYS.keys(): + s = sub.obs[f"pw_{pw}"].astype(float).values + g = sub.obs["genotype"].values + cvals = s[g=="En1-cKO"]; wvals = s[g=="WT"] + try: + _, pval = mannwhitneyu(cvals, wvals, alternative="two-sided") + except Exception: + pval = 1.0 + delta = cvals.mean() - wvals.mean() + rows.append({"class": cls, "pathway": pw, + "n_cKO": n_c, "n_WT": n_w, + "delta_cKO_minus_WT": round(delta, 4), + "MannU_p": pval}) + print(f"[pw] {cls}: {n_c} cKO, {n_w} WT — scored") + + df = pd.DataFrame(rows) + df.to_csv(OUT / "57_pathway_class_by_pathway.csv", index=False) + + pivot_delta = df.pivot(index="pathway", columns="class", values="delta_cKO_minus_WT") + pivot_p = df.pivot(index="pathway", columns="class", values="MannU_p") + def stars(p): return "***" if p<0.001 else "**" if p<0.01 else "*" if p<0.05 else "" + disp = pivot_delta.copy().astype(object) + for pw in disp.index: + for c in disp.columns: + d = pivot_delta.loc[pw, c]; p = pivot_p.loc[pw, c] + if pd.isna(d): disp.loc[pw, c] = "" + else: disp.loc[pw, c] = f"{d:+.3f}{stars(p)}" + + lines = ["# Pathway score contrasts by class — Aldrich En1-cKO vs WT", + "", + "Delta = cKO mean − WT mean of `sc.tl.score_genes` pathway score.", + "Sig: * p<0.05, ** p<0.01, *** p<0.001 (MannU two-sided).\n", + disp.to_markdown()] + (OUT / "57_pathway_class_by_pathway.md").write_text("\n".join(lines)) + print(f"[pw] wrote {OUT}/57_pathway_class_by_pathway.md") + + df_sig = df[df["MannU_p"] < 0.01].sort_values("MannU_p") + print("\n[pw] Strongest cKO/WT pathway shifts (p<0.01):") + print(df_sig[["class","pathway","delta_cKO_minus_WT","MannU_p"]].to_string(index=False)) + df_sig.to_csv(OUT / "57_pathway_top_hits.csv", index=False) + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/57_pathway_analysis.py b/scripts/analysis/57_pathway_analysis.py new file mode 100644 index 0000000000000000000000000000000000000000..b977e65063bfb6827221f7a89bbfb254a13207db --- /dev/null +++ b/scripts/analysis/57_pathway_analysis.py @@ -0,0 +1,529 @@ +"""per-class pathway module scoring cKO/mutant vs WT across pan_skin, hematopoiesis, pancreas. +pancreas contrast is Veres stage 6 vs stage 5 (HUMAN gene symbols).""" +from __future__ import annotations + +import argparse +import warnings +from pathlib import Path + +import anndata as ad +import numpy as np +import pandas as pd +import scanpy as sc +import scipy.sparse as sp +from scipy.stats import mannwhitneyu + +warnings.filterwarnings("ignore") +sc.settings.verbosity = 0 + +ROOT = Path("/home/bcheng/PRISM") +MIN_PER_GROUP = 15 + + +def _M(genes, type_, citation, direction=None): + return {"genes": list(genes), "type": type_, "citation": citation, + "direction": direction} + + +# pan-skin: mouse symbols +SKIN_MODULES = { + "MITF_regulon": _M(["Mitf","Dct","Tyr","Pmel","Mlana","Tyrp1","Slc24a5", + "Slc45a2","Sox10","Pax3","Kit","Ednrb"], + "lineage", "Steingrimsson-2004"), + "Wnt_signaling": _M(["Wnt3","Wnt5a","Wnt7a","Wnt10b","Ctnnb1","Lef1", + "Tcf4","Tcf7","Axin2","Dkk1","Sfrp1","Fzd7","Lrp5"], + "signaling", "Nusse-2017"), + "BMP_signaling": _M(["Bmp2","Bmp4","Bmp5","Bmp7","Bmpr1a","Bmpr1b", + "Bmpr2","Smad1","Smad5","Id1","Id2","Id3"], + "signaling", "Botchkarev-2003"), + "TGFB_signaling": _M(["Tgfb1","Tgfb2","Tgfbr1","Tgfbr2","Smad3","Smad7"], + "signaling", "Massague-2012"), + "FGF_signaling": _M(["Fgf1","Fgf2","Fgf7","Fgf9","Fgf10","Fgfr1","Fgfr2", + "Etv1","Etv4","Etv5","Spry2","Dusp6"], + "signaling", "Ornitz-2015"), + "Notch_signaling": _M(["Notch1","Notch2","Notch3","Jag1","Dll1","Hes1", + "Hes5","Hey1","Hey2","Rbpj"], + "signaling", "Andersson-2011"), + "Hedgehog": _M(["Shh","Ptch1","Smo","Gli1","Gli2","Gli3"], + "signaling", "St-Jacques-1998"), + "Eda_ectodysplasin": _M(["Eda","Edar","Edaradd","Nfkb1","Nfkb2","Rela"], + "signaling", "Mikkola-2009"), + "EMT": _M(["Zeb1","Zeb2","Snai1","Snai2","Twist1","Twist2", + "Vim","Cdh2","Fn1","Prrx1"], + "lineage", "Thiery-2009"), + "Cell_cycle": _M(["Ccnd1","Ccne1","Ccna2","Ccnb1","Cdk1","Cdk2", + "Cdk4","Mki67","Top2a","Pcna","Mcm2","Mcm3"], + "cycle", "Whitfield-2002"), + "KC_differentiation":_M(["Krt1","Krt10","Ivl","Lor","Flg","Flg2","Klk5", + "Klk7","Cdsn"], + "lineage", "Fuchs-2007"), + "Basal_keratinocyte":_M(["Krt5","Krt14","Krt15","Trp63","Itga6","Itgb1", + "Itga3"], + "lineage", "Blanpain-2007"), + "Sweat_gland": _M(["Foxi3","Foxa1","En1","Krt8","Krt18","Krt19", + "Muc5b","Aqp5","Cutl1"], + "lineage", "Lu-2016"), + "Hair_placode": _M(["Wnt10b","Shh","Lef1","Foxi3","Edar","Bmp4","Msx2"], + "lineage", "Millar-2002"), + "Neural_crest": _M(["Sox10","Sox9","Sox2","Pax3","Foxd3","Nes","Tfap2a"], + "lineage", "Simoes-Costa-2015"), + "Apoptosis": _M(["Bax","Bak1","Bad","Bcl2","Casp3","Casp9","Trp53", + "Cdkn1a"], + "stress", "Youle-2008"), + "Melanogenesis_late":_M(["Tyrp1","Slc45a2","Oca2","Gpnmb","Pmel","Silv", + "Mlph","Rab27a","Melana"], + "lineage", "Raposo-2007"), + "Sebogenesis": _M(["Elovl3","Awat2","Scd1","Scd3","Adipoq","Fasn", + "Mgst1","Srebf1","Pparg"], + "metabolism", "Zouboulis-2016"), + "Immune_Th1_Th2_Th17":_M(["Tbx21","Gata3","Rorc","Ifng","Il4","Il13","Il17a", + "Il17f","Il22","Foxp3"], + "immune", "Zhu-2010"), + "DNA_damage": _M(["Trp53","Cdkn1a","Atm","Atr","Brca1","Chek1", + "Chek2","Rad51","Mre11a","H2ax","Nbn"], + "stress", "Ciccia-2010"), + "Autophagy": _M(["Atg5","Atg7","Atg12","Becn1","Map1lc3b","Sqstm1", + "Ulk1","Atg3","Atg16l1"], + "stress", "Mizushima-2011"), + "Senescence": _M(["Cdkn2a","Cdkn2b","Cdkn1a","Il6","Cxcl1","Serpine1", + "Glb1","Lmnb1"], + "stress", "Coppe-2010"), + "Epidermal_junction":_M(["Cdh1","Dsg1a","Dsg1b","Dsg2","Dsg3","Cldn1", + "Cldn4","Ocln","Tjp1","Cldn23","Dsp","Pkp1"], + "junction", "Green-2010"), + "ECM_collagen": _M(["Col1a1","Col1a2","Col3a1","Col4a1","Col6a1", + "Col17a1","Lum","Dcn","Fbn1","Postn","Fbln1"], + "ecm", "Ricard-Blum-2011"), + "Endothelial_tip_stalk":_M(["Dll4","Notch1","Hes1","Kdr","Kit","Cxcr4", + "Angpt2","Nrp1","Flt1","Cdh5","Pecam1"], + "lineage", "Blanco-2013"), + "Fibroblast_wound": _M(["Postn","Fap","Aspn","Tnc","Acta2","Prrx1","Pdgfra", + "Ly6a"], + "lineage", "Rinkevich-2015"), + "Fatty_acid_oxidation":_M(["Cpt1a","Acadm","Acadl","Acadvl","Hadha","Hadhb", + "Ppara","Ppargc1a","Ucp2"], + "metabolism", "Houten-2010"), + "Nrf2_oxidative_stress":_M(["Nfe2l2","Nqo1","Gclc","Hmox1","Slc7a11", + "Txnrd1","Gsta3","Gpx2","Keap1"], + "stress", "Ma-2013"), + "IFN_gamma": _M(["Ifng","Stat1","Ifit1","Ifit2","Ifit3","Isg15", + "Irf1","Cxcl9","Cxcl10","Gbp2"], + "immune", "Schoggins-2011"), + "IL6_JAK_STAT": _M(["Il6","Stat3","Socs3","Jak1","Jak2","Il6ra", + "Il6st","Osm"], + "signaling", "Heinrich-2003"), + "Pigment_regulation":_M(["Asip","Kitl","Kit","Bcl2","Mc1r","Pomc","Adcy8", + "Ednrb","Edn3"], + "signaling", "Slominski-2004"), +} + + +# hematopoiesis: mouse symbols +HSC_MODULES = { + "Kit_signaling": _M(["Kit","Kitl","Sox4","Gata2","Runx1","Meis1"], + "signaling", "Lennartsson-2012"), + "Kit_ligand": _M(["Kit","Kitl"], "signaling", "Broudy-1997"), + "MYC_targets": _M(["Myc","Nolc1","Nop58","Ncl","Npm1","Fbl","Eif4e", + "Nop56","Ldha","Odc1"], + "lineage", "Dang-2012"), + "Cell_cycle": _M(["Ccnd1","Ccne1","Ccna2","Ccnb1","Cdk1","Cdk2", + "Cdk4","Mki67","Top2a","Pcna","Mcm2","Mcm3","Mcm5"], + "cycle", "Whitfield-2002"), + "DNA_replication": _M(["Mcm2","Mcm3","Mcm4","Mcm5","Mcm6","Mcm7","Pcna", + "Rfc4","Pola1","Pole","Rpa1","Rpa2"], + "cycle", "Bell-2002"), + "Integrated_stress": _M(["Atf4","Ddit3","Ppp1r15a","Ppp1r15b","Eif2ak3", + "Eif2s1","Atf3","Gadd45a"], + "stress", "Pakos-Zebrucka-2016"), + "Apoptosis_pro": _M(["Bax","Bak1","Bid","Bad","Bim","Puma","Noxa", + "Casp3","Casp9"], + "stress", "Youle-2008"), + "Apoptosis_anti": _M(["Bcl2","Bcl2l1","Mcl1","Bcl2l2","Bcl2l10","Xiap"], + "stress", "Adams-2018"), + "Erythropoiesis_early":_M(["Gata1","Klf1","Epo","Epor","Tal1","Zfpm1", + "Gypa","Lmo2"], + "lineage", "Palis-2014"), + "Erythropoiesis_late":_M(["Alas2","Hba-a1","Hba-a2","Hbb-b1","Hbb-b2","Slc4a1", + "Ank1","Blvrb","Car1","Car2"], + "lineage", "Palis-2014"), + "Granulopoiesis": _M(["Cebpa","Cebpe","Elane","Mpo","Prtn3","Csf3r", + "S100a8","S100a9","Ctsg","Ltf","Lcn2","Mmp8"], + "lineage", "Rosenbauer-2007"), + "Lymphopoiesis_B": _M(["Rag1","Rag2","Dntt","Vpreb1","Vpreb3","Igll1", + "Cd19","Pax5","Ebf1"], + "lineage", "Nutt-2011"), + "Lymphopoiesis_T": _M(["Il7r","Cd3d","Cd3e","Cd3g","Lck","Zap70","Gata3", + "Tcf7","Runx3"], + "lineage", "Rothenberg-2014"), + "Megakaryopoiesis": _M(["Nfe2","Gata1","Fli1","Runx1","Itga2b","Pf4", + "Gp1bb","Gp9","Mpl","Vwf"], + "lineage", "Tijssen-2013"), + "Basophil_mast": _M(["Cpa3","Ms4a2","Gata2","Hdc","Mcpt8","Prss34", + "Fcer1a","Il4","Il6"], + "lineage", "Voehringer-2013"), + "Hemostasis": _M(["Vwf","F5","F13a1","Fga","Fgb","Fgg","Serpine1", + "Plat","Plau","Plg"], + "signaling", "Furie-2008"), + "OXPHOS_ETC": _M(["Ndufa1","Ndufa2","Ndufb1","Ndufb2","Sdha","Sdhb", + "Cox4i1","Cox5a","Cox6a1","Atp5a1","Atp5b","Uqcrq"], + "metabolism", "Mishra-2016"), + "Glycolysis": _M(["Hk1","Hk2","Pfkm","Pfkl","Aldoa","Gapdh","Pgk1", + "Pkm","Ldha","Eno1","Tpi1","Pgam1"], + "metabolism", "Vander-Heiden-2009"), + "TCA": _M(["Cs","Aco2","Idh2","Idh3a","Sdha","Fh1","Mdh2", + "Ogdh","Sucla2"], + "metabolism", "Chandel-2015"), + "Redox_glutathione": _M(["Gpx1","Gpx2","Gpx3","Gpx4","Gsr","Prdx1","Prdx2", + "Prdx3","Prdx4","Prdx5","Prdx6","Sod1","Sod2","Cat"], + "stress", "Ho-2007"), + "Wnt_hemato": _M(["Wnt3a","Wnt5a","Ctnnb1","Lef1","Tcf7","Axin2", + "Fzd4","Fzd7"], + "signaling", "Reya-2003"), + "Notch_hemato": _M(["Notch1","Notch2","Jag1","Hes1","Dll1","Dll4", + "Rbpj","Hey1"], + "signaling", "Bigas-2018"), + "TGFb_hemato": _M(["Tgfb1","Tgfb2","Tgfbr1","Tgfbr2","Smad2","Smad3", + "Smad4","Smad7"], + "signaling", "Blank-2015"), + "IFN_signaling": _M(["Ifnar1","Ifnar2","Stat1","Stat2","Ifit1","Ifit2", + "Ifit3","Isg15","Irf7","Mx1"], + "immune", "Essers-2009"), + "Complement": _M(["C1qa","C1qb","C1qc","C3","C4b","Cfp","Cfh","Cfd"], + "immune", "Ricklin-2016"), + "NK_cytotoxicity": _M(["Ncr1","Klrk1","Prf1","Gzmb","Gzmk","Nkg7","Klrd1", + "Klrb1c","Klra8"], + "immune", "Vivier-2011"), + "Mast_cell_degran": _M(["Ms4a2","Fcer1a","Cpa3","Kit","Hdc","Tpsb2", + "Prss34","Mcpt4"], + "immune", "Galli-2011"), + "Autophagy": _M(["Atg5","Atg7","Atg12","Becn1","Map1lc3b","Sqstm1", + "Ulk1","Atg3","Atg16l1"], + "stress", "Warr-2013"), + "Senescence": _M(["Cdkn2a","Cdkn2b","Cdkn1a","Il6","Cxcl1","Serpine1", + "Glb1","Lmnb1"], + "stress", "Chang-2016"), + "LT_HSC_quiescence": _M(["Hlf","Meis1","Mecom","Procr","Fgd5","Mllt3","Egr1", + "Rgs1","Cdkn1c","Ndn","Mpl"], + "lineage", "Cabezas-Wallscheid-2017"), +} + + +# pancreas: HUMAN symbols (Veres is hPSC) +PANCREAS_MODULES = { + "Insulin_secretion": _M(["INS","IAPP","CHGA","CHGB","SCG5","ERO1B","PCSK1", + "PCSK2","SLC30A8","G6PC2"], + "hormone", "Rorsman-2013"), + "Glucose_sensing": _M(["SLC2A2","GCK","KCNJ11","ABCC8","SIRT1","GLUT1", + "SLC2A1"], + "signaling", "Matschinsky-2013"), + "Alpha_master_TF": _M(["ARX","IRX1","IRX2","MAFB","POU3F4","GCG","TTR"], + "lineage", "Collombat-2003"), + "Beta_master_TF_embryonic":_M(["NKX6-1","MNX1","NEUROD1","PDX1","NKX2-2", + "HNF1B"], + "lineage", "Gu-2004"), + "Beta_master_TF_adult":_M(["MAFA","UCN3","SIX3","INS","IAPP","G6PC2"], + "lineage", "Blum-2012"), + "Neurog3_EP_cascade":_M(["NEUROG3","PAX4","FEV","INSM1","NEUROD1","SOX4", + "CBFA2T3","BTBD17"], + "lineage", "Gradwohl-2000"), + "Endocrine_maturation":_M(["RFX3","RFX6","ISL1","FOXA2","PAX6","NKX2-2"], + "lineage", "Piccand-2014"), + "Exocrine_acinar": _M(["PRSS1","PRSS2","CEL","CPA1","CTRB1","AMY2A", + "ELOVL5","PTF1A","CELA1"], + "lineage", "Kawaguchi-2002"), + "Ductal_epithelial": _M(["KRT19","KRT7","SOX9","MUC1","ONECUT1","HES1", + "HNF1B","CFTR"], + "lineage", "Solar-2009"), + "Foregut_endoderm": _M(["SOX17","FOXA1","FOXA2","ONECUT1","PROX1","HNF1A", + "HNF1B","GATA4","GATA6"], + "lineage", "Zorn-2009"), + "Cilium_Foxj1": _M(["FOXJ1","CFAP43","CFAP157","NPHP1","IFT88","DNAH5", + "TEKT1","SPAG6"], + "lineage", "Choksi-2014"), + "Delta_master": _M(["SST","HHEX","LEPR","GHSR"], + "hormone", "Rorsman-2018"), + "Gamma_master": _M(["PPY","PYY","SLC38A4"], + "hormone", "Wang-2016"), + "Epsilon_ghrelin": _M(["GHRL","ACSL1"], + "hormone", "Prado-2004"), + "ER_stress_pancreas":_M(["ATF6","XBP1","ERN1","DDIT3","HSPA5","HSPA1A", + "HSPA1B","EIF2AK3"], + "stress", "Back-2012"), + "Unfolded_protein_response":_M(["ATF4","ATF6","XBP1","HERPUD1","BAK1","BAX", + "EDEM1","DERL1"], + "stress", "Walter-2011"), + "Hormone_processing":_M(["PCSK1","PCSK2","CPE","CHGA","CHGB","SCG2","SCG5", + "PAM"], + "hormone", "Docherty-1997"), + "Insulin_receptor_signaling":_M(["INSR","IRS1","IRS2","AKT2","PDX1","FOXO1", + "GSK3B","MTOR"], + "signaling", "Kulkarni-1999"), + "Mesenchyme_pancreatic":_M(["NKX3-2","BMP4","SOX9","FGF10","COL1A1","COL3A1", + "DCN"], + "lineage", "Landsman-2011"), + "Fatty_acid_oxidation":_M(["CPT1A","ACADM","ACADL","HADHA","PPARA","PPARGC1A", + "ACOX1"], + "metabolism", "Houten-2010"), + "Glycolysis": _M(["HK1","HK2","PFKM","PFKL","ALDOA","GAPDH","PGK1", + "PKM","LDHA","ENO1","TPI1"], + "metabolism", "Vander-Heiden-2009"), + "TCA": _M(["CS","ACO2","IDH2","IDH3A","SDHA","FH","MDH2", + "OGDH","SUCLA2"], + "metabolism", "Chandel-2015"), + "OXPHOS_ETC": _M(["NDUFA1","NDUFA2","NDUFB1","SDHA","SDHB","COX4I1", + "COX5A","COX6A1","ATP5A1","ATP5B","UQCRQ"], + "metabolism", "Mishra-2016"), + "Redox_glutathione": _M(["GPX1","GPX2","GPX3","GPX4","GSR","PRDX1","PRDX2", + "PRDX3","PRDX4","PRDX5","PRDX6","SOD1","SOD2","CAT"], + "stress", "Ho-2007"), + "Wnt_pancreas": _M(["WNT3A","WNT5A","CTNNB1","LEF1","TCF7","AXIN2", + "FZD7"], + "signaling", "Murtaugh-2008"), + "Notch_pancreas": _M(["NOTCH1","NOTCH2","JAG1","HES1","DLL1","DLL4", + "RBPJ","HEY1"], + "signaling", "Apelqvist-1999"), + "TGFb_pancreas": _M(["TGFB1","TGFB2","TGFBR1","TGFBR2","SMAD2","SMAD3", + "SMAD4","SMAD7"], + "signaling", "Sanvito-1994"), + "Immune_pancreas": _M(["PTPRC","CD68","ADGRE1","CD3D","CD3E","CD4","CD8A", + "CD19"], + "immune", "Homo-2015"), + "Endothelial_pancreas":_M(["PECAM1","CDH5","KDR","VWF","PLVAP","FLT1", + "TEK","ENG"], + "lineage", "Cleaver-2019"), + "Cell_cycle": _M(["CCND1","CCNE1","CCNA2","CCNB1","CDK1","CDK2", + "CDK4","MKI67","TOP2A","PCNA","MCM2","MCM3"], + "cycle", "Whitfield-2002"), +} + + +def load_pan_skin(): + RAW = ROOT / "data/raw/GSE220977_combined.h5ad" + PRED = ROOT / "discovery/pan_skin/marker/dingwall_predictions.csv" + + CKO_GSMS = {"GSM6833482", "GSM6833483"} # 480/481 are rttaControl (WT), not cKO — per GEO metadata + WT_GSMS = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"} # 4 Cre-neg controls per GEO metadata + + a = ad.read_h5ad(RAW) + pred = pd.read_csv(PRED) + common = a.obs_names.intersection(pd.Index(pred["cell_id"].astype(str))) + a = a[list(common)].copy() + pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"])) + a.obs["pred_label"] = pd.Categorical([pred_map.get(c,"unknown") + for c in a.obs_names]) + samp = a.obs["sample"].astype(str) + a.obs["group"] = np.where(samp.isin(list(CKO_GSMS)), "En1-cKO", + np.where(samp.isin(list(WT_GSMS)), "WT", "other")) + a = a[a.obs["group"].isin(["En1-cKO","WT"])].copy() + return a, "En1-cKO", "WT", SKIN_MODULES + + +def load_hematopoiesis(): + D_DIR = ROOT / "data/corpus/hematopoiesis/held_out_unlabeled/dahlin_extract" + PRED = ROOT / "discovery/hematopoiesis/marker/nestorowa_anchor_predictions.csv" + if not PRED.exists(): + PRED = ROOT / "discovery/hematopoiesis/marker/97_nestorowa_anchor_predictions.csv" + + # Dahlin lacks a nestorowa-style anchor csv; fall back to its own predictions + DAHLIN_PRED_CANDIDATES = [ + ROOT / "discovery/hematopoiesis/marker/dahlin_predictions.csv", + ROOT / "discovery/hematopoiesis/marker/92_dahlin_predictions.csv", + ] + for p in DAHLIN_PRED_CANDIDATES: + if p.exists(): + PRED = p + break + + GT = {"SIGAB1":"WT","SIGAC1":"WT","SIGAD1":"WT","SIGAF1":"WT","SIGAG1":"WT", + "SIGAH1":"WT","SIGAG8":"Kit_W41","SIGAH8":"Kit_W41"} + + parts = [] + for f in sorted(D_DIR.glob("*.txt.gz")): + sample = f.name.split("_")[1].split(".")[0] + 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["group"] = GT.get(sample, "unknown") + var = pd.DataFrame(index=df.index.astype(str)) + parts.append(ad.AnnData(X=X, obs=obs, var=var)) + a = ad.concat(parts, join="outer", label="_batch") + + try: + import mygene + mg = mygene.MyGeneInfo() + res = mg.querymany(a.var_names.astype(str).tolist(), 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) + a = a[:, keep].copy() + a.var_names = syms[keep] + a.var_names_make_unique() + except Exception as e: + print(f"[warn] mygene mapping failed: {e}") + + a = a[a.obs["group"].isin(["Kit_W41","WT"])].copy() + + if PRED.exists(): + pred = pd.read_csv(PRED) + common = a.obs_names.intersection(pd.Index(pred["cell_id"].astype(str))) + a = a[list(common)].copy() + pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"])) + a.obs["pred_label"] = pd.Categorical([pred_map.get(c,"unknown") + for c in a.obs_names]) + else: + # assign a single class so module scoring still runs + a.obs["pred_label"] = pd.Categorical(["all"] * a.n_obs) + print(f"[warn] no Dahlin prediction file found; using pred_label='all'") + + return a, "Kit_W41", "WT", HSC_MODULES + + +def load_pancreas(): + SHARON_DIR = ROOT / "data/corpus/pancreas/held_out_unlabeled/sharon_extract" + PRED = ROOT / "discovery/pancreas/marker/veres_predictions.csv" + + parts = [] + for meta_file in sorted(SHARON_DIR.glob("*.cell_metadata.tsv.gz")): + counts_file = str(meta_file).replace("cell_metadata", "processed_counts") + if not Path(counts_file).exists(): + continue + meta = pd.read_csv(meta_file, sep="\t", compression="gzip") + counts = pd.read_csv(counts_file, sep="\t", compression="gzip", index_col=0) + obs = meta.set_index("library.barcode") + obs = obs.loc[obs.index.intersection(counts.index)] + counts_al = counts.loc[obs.index] + X = sp.csr_matrix(counts_al.values.astype(np.float32)) + aa = ad.AnnData(X=X, obs=obs, + var=pd.DataFrame(index=counts_al.columns)) + aa.var_names_make_unique() + parts.append(aa) + a = ad.concat(parts, join="outer") + + pred = pd.read_csv(PRED) + # veres cell_ids are prefixed with "veres_" — strip to match sharon obs_names + pred["cell_id_stripped"] = pred["cell_id"].astype(str).str.replace(r"^veres_", "", regex=True) + common = a.obs_names.intersection(pd.Index(pred["cell_id_stripped"])) + a = a[list(common)].copy() + pred_map = dict(zip(pred["cell_id_stripped"], pred["pred_label"])) + a.obs["pred_label"] = pd.Categorical([pred_map.get(c,"unknown") + for c in a.obs_names]) + + # canonical Veres contrast: Stage 6 (mature) vs Stage 5 (immature) + stage = a.obs["Stage"].astype(str) + a.obs["group"] = np.where(stage == "6", "Stage6", + np.where(stage == "5", "Stage5", "other")) + a = a[a.obs["group"].isin(["Stage6","Stage5"])].copy() + return a, "Stage6", "Stage5", PANCREAS_MODULES + + +LOADERS = { + "pan_skin": load_pan_skin, + "hematopoiesis": load_hematopoiesis, + "pancreas": load_pancreas, +} + + +def score_modules(sub, modules): + var_set = set(sub.var_names.astype(str)) + for name, spec in modules.items(): + present = [g for g in spec["genes"] if g in var_set] + if not present: + sub.obs[f"pw_{name}"] = 0.0 + continue + try: + sc.tl.score_genes(sub, gene_list=present, score_name=f"pw_{name}", + random_state=0, use_raw=False) + except Exception: + sub.obs[f"pw_{name}"] = 0.0 + return sub + + +def run_system(system_name): + print(f"[load] {system_name}", flush=True) + a, g1, g2, modules = LOADERS[system_name]() + print(f"[load] {a.n_obs} cells, {sum(a.obs['group']==g1)} {g1}, " + f"{sum(a.obs['group']==g2)} {g2}, {len(modules)} modules", flush=True) + + sc.pp.normalize_total(a, target_sum=1e4) + sc.pp.log1p(a) + + OUT = ROOT / f"discovery/{system_name}/marker" + OUT.mkdir(parents=True, exist_ok=True) + + classes = sorted(a.obs["pred_label"].astype(str).unique()) + rows = [] + for cls in classes: + mask = (a.obs["pred_label"].astype(str) == cls).values + n1 = int((mask & (a.obs["group"].values == g1)).sum()) + n2 = int((mask & (a.obs["group"].values == g2)).sum()) + if n1 < MIN_PER_GROUP or n2 < MIN_PER_GROUP: + print(f"[pw] {cls}: skip (n_{g1}={n1}, n_{g2}={n2})") + continue + sub = a[mask].copy() + sub = score_modules(sub, modules) + grp = sub.obs["group"].values + for mod_name, spec in modules.items(): + s = sub.obs[f"pw_{mod_name}"].astype(float).values + v1 = s[grp == g1]; v2 = s[grp == g2] + try: + _, pval = mannwhitneyu(v1, v2, alternative="two-sided") + except Exception: + pval = 1.0 + delta = float(v1.mean() - v2.mean()) + rows.append({ + "class": cls, + "module_name": mod_name, + "module_type": spec["type"], + "citation": spec["citation"], + "direction": spec["direction"], + "n_g1": n1, + "n_g2": n2, + "group_g1": g1, + "group_g2": g2, + "delta": round(delta, 4), + "mannu_p": float(pval), + }) + print(f"[pw] {cls}: {n1} {g1}, {n2} {g2} — scored") + + df = pd.DataFrame(rows) + if df.empty: + print("[pw] no eligible classes — done") + return + + n_tests = len(df) + df["mannu_p_adj_bonferroni"] = np.minimum(df["mannu_p"] * n_tests, 1.0) + + csv_path = OUT / "57_pathway_analysis.csv" + df.to_csv(csv_path, index=False) + print(f"[pw] wrote {csv_path} ({len(df)} rows, n_tests={n_tests})", + flush=True) + + pivot_delta = df.pivot(index="module_name", columns="class", values="delta") + pivot_padj = df.pivot(index="module_name", columns="class", + values="mannu_p_adj_bonferroni") + pivot_delta.to_csv(OUT / "57_pathway_class_by_module_delta.tsv", sep="\t") + pivot_padj.to_csv(OUT / "57_pathway_class_by_module_padj.tsv", sep="\t") + print(f"[pw] wrote heatmap TSVs to {OUT}") + + sig = df[df["mannu_p_adj_bonferroni"] < 0.01].sort_values( + "mannu_p_adj_bonferroni") + print(f"\n[pw] top Bonferroni-significant shifts (padj<0.01, " + f"n={len(sig)}):") + if len(sig): + print(sig[["class","module_name","module_type","delta", + "mannu_p_adj_bonferroni"]].head(30).to_string(index=False)) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--system", required=True, + choices=list(LOADERS.keys()), + help="pan_skin | hematopoiesis | pancreas") + args = p.parse_args() + run_system(args.system) + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/62_time_course_analysis.py b/scripts/analysis/62_time_course_analysis.py new file mode 100644 index 0000000000000000000000000000000000000000..8a5795e5a1ec5930d09dd14c3ca2e683b064596d --- /dev/null +++ b/scripts/analysis/62_time_course_analysis.py @@ -0,0 +1,112 @@ +"""PANDA lineage distribution across LARRY time points d2/d9/d16 + clonal purity + per-lineage DE.""" +from __future__ import annotations +from pathlib import Path +import warnings, json, sys +warnings.filterwarnings("ignore") +import numpy as np, pandas as pd, anndata as ad, scanpy as sc, torch +from scipy import stats + +sys.path.insert(0, "/home/bcheng/PRISM") +from panda.model import PANDAEncoder + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +CKPT = Path("/home/bcheng/PRISM/checkpoints/hematopoiesis") +CORPUS = Path("/home/bcheng/PRISM/data/corpus/hematopoiesis/harmonized/corpus.h5ad") +OUT = Path("/home/bcheng/PRISM/discovery/hematopoiesis/marker") +OUT.mkdir(parents=True, exist_ok=True) + + +def main(): + 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"]) + protos = ck["prototypes"] + protos = protos / (np.linalg.norm(protos, axis=1, keepdims=True) + 1e-8) + + # predict on all corpus cells (this is training-set prediction — for analysis only) + a = ad.read_h5ad(CORPUS) + print(f"[hsc-mech] corpus: {a.shape}", flush=True) + X = np.asarray(a.obsm["X_pca"]).astype(np.float32) + + all_z = [] + with torch.no_grad(): + for i in range(0, X.shape[0], 8192): + xb = torch.from_numpy(X[i:i+8192]).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) + cos = Z @ protos.T + pred_ix = cos.argmax(axis=1) + pred = np.array([classes[i] for i in pred_ix], dtype=object) + a.obs["pred_label"] = pred + a.obs["pred_conf"] = cos.max(axis=1) + + if "Time point" not in a.obs.columns: + print("[hsc-mech] no Time point column; skipping time analysis") + else: + tp = a.obs["Time point"].astype(int) + xt = pd.crosstab(pred, tp, normalize="columns") + print("\n[hsc-mech] Fraction per predicted class per time point:") + print(xt.round(3)) + xt.to_csv(OUT / "62_time_course_class_fractions.csv") + + rows = [] + n_d2 = int((tp == 2).sum()); n_d16 = int((tp == 16).sum()) + for c in classes: + n_c_d16 = int(((pred == c) & (tp == 16)).sum()) + n_c_d2 = int(((pred == c) & (tp == 2)).sum()) + contingency = np.array([[n_c_d16, n_d16 - n_c_d16], [n_c_d2, n_d2 - n_c_d2]]) + odds, p = stats.fisher_exact(contingency) + f16 = (n_c_d16 + 1) / (n_d16 + 2); f2 = (n_c_d2 + 1) / (n_d2 + 2) + rows.append({"class": c, "n_d16": n_c_d16, "n_d2": n_c_d2, + "log2_fold_d16_vs_d2": round(np.log2(f16 / f2), 3), + "fisher_p": p}) + df = pd.DataFrame(rows).sort_values("log2_fold_d16_vs_d2", ascending=False) + print("\n[hsc-mech] class enrichment d16 vs d2 (positive = expanded at late time):") + print(df.to_string(index=False)) + df.to_csv(OUT / "62_time_course_enrichment.csv", index=False) + + # sibling-fate concordance: Library = clonal barcode + if "Library" in a.obs.columns: + libs = a.obs["Library"].astype(str) + top_lib = libs.value_counts().head(200).index # top 200 largest clones + clone_purity = [] + for L in top_lib: + m = libs == L + if m.sum() < 3: continue + pl = pd.Series(pred[m.values]).value_counts(normalize=True) + clone_purity.append({ + "library": L, "n": int(m.sum()), + "dominant_class": pl.index[0], + "purity": float(pl.iloc[0]), + }) + cp = pd.DataFrame(clone_purity) + print(f"\n[hsc-mech] clonal purity (dominant-class fraction) — {len(cp)} clones:") + print(f" median: {cp['purity'].median():.3f}, mean: {cp['purity'].mean():.3f}, " + f"n_clones_pure_>0.9: {(cp['purity'] > 0.9).sum()}/{len(cp)}") + cp.to_csv(OUT / "62_clonal_purity.csv", index=False) + + a.obs["pred_label"] = pd.Categorical(pred) + keep_classes = [c for c in classes if (pred == c).sum() >= 100] + a_sub = a[np.isin(pred, keep_classes)].copy() + if a_sub.n_obs >= 500: + sc.tl.rank_genes_groups(a_sub, "pred_label", method="wilcoxon", + n_genes=30, use_raw=False) + rows = [] + for cls in keep_classes: + try: + names = a_sub.uns["rank_genes_groups"]["names"][cls] + lfc = a_sub.uns["rank_genes_groups"]["logfoldchanges"][cls] + for g, l in zip(names[:15], lfc[:15]): + rows.append({"class": cls, "gene": g, "logfc": round(float(l), 3)}) + except Exception: pass + pd.DataFrame(rows).to_csv(OUT / "62_lineage_markers.csv", index=False) + print(f"\n[hsc-mech] wrote lineage markers to 62_lineage_markers.csv") + + print(f"\n[hsc-mech] complete. Outputs in {OUT}/") + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/63_nestorowa_zero_shot.py b/scripts/analysis/63_nestorowa_zero_shot.py new file mode 100644 index 0000000000000000000000000000000000000000..248171e2196b90b46d427e4f3e8ce2b6ad538986 --- /dev/null +++ b/scripts/analysis/63_nestorowa_zero_shot.py @@ -0,0 +1,142 @@ +"""zero-shot HSC PANDA on nestorowa 2016 smart-seq2 as unlabeled discovery target.""" +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 +sys.path.insert(0, "/home/bcheng/PRISM") +from panda.model import PANDAEncoder + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +CKPT = Path("/home/bcheng/PRISM/checkpoints/hematopoiesis") +HARM = Path("/home/bcheng/PRISM/data/corpus/hematopoiesis/harmonized") +OUT = Path("/home/bcheng/PRISM/discovery/hematopoiesis/marker") + + +def load_nestorowa(): + p = Path("/home/bcheng/PRISM/data/corpus/hematopoiesis/tier_a/nestorowa_GSE81682_counts.txt.gz") + df = pd.read_csv(p, sep="\t", index_col=0, compression="gzip") + # rows = ENSMUSG, cols = cells; convert IDs to symbols via mygene + import mygene + print(f"[nestorowa] converting {len(df)} ENSMUSG IDs to symbols via mygene …", flush=True) + mg = mygene.MyGeneInfo() + ids = df.index.astype(str).tolist() + res = mg.querymany(ids, scopes="ensembl.gene", fields="symbol", + species="mouse", returnall=False, verbose=False) + id2sym = {} + for r in res: + if "symbol" in r and "query" in r and not r.get("notfound", False): + id2sym[r["query"]] = r["symbol"] + syms_series = pd.Series(df.index.astype(str)).map(id2sym).values + syms = pd.Series(syms_series, index=df.index) + keep = syms.notna() + print(f"[nestorowa] mapped {int(keep.sum())}/{len(df)} genes", flush=True) + df = df.loc[keep] + df.index = syms[keep].values + df = df.groupby(df.index).sum() + X = sp.csr_matrix(df.values.T.astype(np.float32)) + obs = pd.DataFrame(index=df.columns.astype(str)) + obs["dataset"] = "nestorowa_GSE81682" + obs["sample"] = obs.index.str.split("_").str[0] + var = pd.DataFrame({"gene_symbol": df.index.astype(str).values}, + index=df.index.astype(str)) + a = ad.AnnData(X=X, obs=obs, var=var) + a.var_names_make_unique() + return a + + +def project(a, shared_hvgs, mu, sig): + G = len(shared_hvgs); n = a.n_obs + 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((n, 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) + return Xz, frac + + +def main(): + 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"]) + protos = ck["prototypes"] + protos = protos / (np.linalg.norm(protos, axis=1, keepdims=True) + 1e-8) + print(f"[model] classes: {classes}", flush=True) + + 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) + + a = load_nestorowa() + print(f"[target] Nestorowa shape: {a.shape}", flush=True) + + Xz, frac = project(a, shared_hvgs, mu, sig) + 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) + out = model(xb, aux, lam_dann=0.0) + all_z.append(out["z"].cpu().numpy()) + Z = np.concatenate(all_z, axis=0) + cos = Z @ protos.T + pred_ix = cos.argmax(axis=1) + conf = cos.max(axis=1) + entropy = -(np.exp(cos / 0.07) / np.exp(cos / 0.07).sum(axis=1, keepdims=True) * + np.log(np.exp(cos / 0.07) / np.exp(cos / 0.07).sum(axis=1, keepdims=True) + 1e-12) + ).sum(axis=1) + + a.obs["pred_label"] = np.array([classes[i] for i in pred_ix], dtype=object) + a.obs["pred_conf"] = conf.astype(np.float32) + a.obs["pred_entropy"] = entropy.astype(np.float32) + + print(f"\n[nestorowa] predicted class distribution:") + print(a.obs["pred_label"].value_counts()) + print(f"\n[nestorowa] pred_conf: p10={np.percentile(conf,10):.3f}, " + f"p50={np.percentile(conf,50):.3f}, p90={np.percentile(conf,90):.3f}") + print(f"[nestorowa] shared-HVG fraction: {frac:.1%}") + + a.obs.to_csv(OUT / "63_nestorowa_predictions.csv") + print(f"\n[nestorowa] wrote predictions to 63_nestorowa_predictions.csv") + + # cluster bottom-decile confidence cells for novel-population DE + thr = np.percentile(conf, 10) + mask = conf <= thr + print(f"\n[nestorowa] bottom-decile confidence: {int(mask.sum())} cells (thr={thr:.3f})", + flush=True) + if mask.sum() >= 30: + sub_low = ad.AnnData(X=Z[mask].astype(np.float32)) + sc.pp.neighbors(sub_low, use_rep="X", n_neighbors=10) + sc.tl.leiden(sub_low, resolution=0.5, key_added="cluster") + print(f"[nestorowa] novel clusters: {sub_low.obs['cluster'].nunique()}") + a_low = a[mask].copy() + a_low.obs["cluster"] = sub_low.obs["cluster"].values + try: + sc.tl.rank_genes_groups(a_low, "cluster", method="wilcoxon", n_genes=10, use_raw=False) + rows = [] + for cl in sorted(a_low.obs["cluster"].unique()): + names = a_low.uns["rank_genes_groups"]["names"][cl] + lfc = a_low.uns["rank_genes_groups"]["logfoldchanges"][cl] + for g, l in zip(names[:8], lfc[:8]): + rows.append({"cluster": cl, "gene": g, "logfc": round(float(l), 3)}) + pd.DataFrame(rows).to_csv(OUT / "63_nestorowa_novel_markers.csv", index=False) + print(f"[nestorowa] novel markers saved") + except Exception as e: + print(f"[nestorowa] DE failed: {e}") + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/66_dahlin_kit_mutant.py b/scripts/analysis/66_dahlin_kit_mutant.py new file mode 100644 index 0000000000000000000000000000000000000000..6406ec015ebdeb493287098b349a1d4d0075a2c0 --- /dev/null +++ b/scripts/analysis/66_dahlin_kit_mutant.py @@ -0,0 +1,142 @@ +"""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 +sys.path.insert(0, "/home/bcheng/PRISM") +from panda.model import PANDAEncoder + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +CKPT = Path("/home/bcheng/PRISM/checkpoints/hematopoiesis") +HARM = Path("/home/bcheng/PRISM/data/corpus/hematopoiesis/harmonized") +DAHLIN = Path("/home/bcheng/PRISM/data/corpus/hematopoiesis/held_out_unlabeled/dahlin_extract") +OUT = Path("/home/bcheng/PRISM/discovery/hematopoiesis/marker") + +GENOTYPE_MAP = { + "SIGAB1": "WT", "SIGAC1": "WT", "SIGAD1": "WT", + "SIGAF1": "WT", "SIGAG1": "WT", "SIGAH1": "WT", + "SIGAG8": "Kit_W41", "SIGAH8": "Kit_W41", +} + + +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["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") + + print(f"\n[dahlin] Fisher exact class enrichment Kit-mutant vs WT:") + rows = [] + n_wt = int((a.obs["genotype"] == "WT").sum()) + n_kit = int((a.obs["genotype"] == "Kit_W41").sum()) + print(f" WT n={n_wt}, Kit_W41 n={n_kit}") + for c in classes: + n_c_kit = int(((pred == c) & (a.obs["genotype"] == "Kit_W41")).sum()) + n_c_wt = int(((pred == c) & (a.obs["genotype"] == "WT")).sum()) + contingency = np.array([[n_c_kit, n_kit - n_c_kit], + [n_c_wt, n_wt - n_c_wt]]) + odds, p = stats.fisher_exact(contingency) + f_kit = (n_c_kit + 1) / (n_kit + 2) + f_wt = (n_c_wt + 1) / (n_wt + 2) + rows.append({"class": c, "n_Kit_W41": n_c_kit, "n_WT": n_c_wt, + "pct_Kit": round(100 * n_c_kit / n_kit, 3), + "pct_WT": round(100 * n_c_wt / n_wt, 3), + "log2_fold_Kit_vs_WT": round(np.log2(f_kit / f_wt), 3), + "fisher_p": p}) + df = pd.DataFrame(rows).sort_values("log2_fold_Kit_vs_WT") + print(df.to_string(index=False)) + df.to_csv(OUT / "66_dahlin_enrichment.csv", index=False) + + a.obs.to_csv(OUT / "66_dahlin_predictions.csv") + print(f"\n[dahlin] complete. Outputs in {OUT}/") + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/67_dahlin_within_class.py b/scripts/analysis/67_dahlin_within_class.py new file mode 100644 index 0000000000000000000000000000000000000000..7493cc8bffca0a37c3fc1e92d91b5a01d3227efb --- /dev/null +++ b/scripts/analysis/67_dahlin_within_class.py @@ -0,0 +1,129 @@ +"""within-class wilcoxon DE WT vs Kit_W41 + module scoring; checks dahlin paper claims (Myc/ISR).""" +from __future__ import annotations +from pathlib import Path +import warnings, json, sys +warnings.filterwarnings("ignore") +import numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp +from scipy.stats import mannwhitneyu + +DAHLIN = Path("/home/bcheng/PRISM/data/corpus/hematopoiesis/held_out_unlabeled/dahlin_extract") +PRED = Path("/home/bcheng/PRISM/discovery/hematopoiesis/marker/66_dahlin_predictions.csv") +OUT = Path("/home/bcheng/PRISM/discovery/hematopoiesis/marker") + + +def load_dahlin_all(): + parts = [] + for f in sorted(DAHLIN.glob("*.txt.gz")): + gsm = f.name.split("_")[0] + sample = f.name.split("_")[1].split(".")[0] + 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 + 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 main(): + print("[dahlin-mech] loading Dahlin …", flush=True) + a = load_dahlin_all() + + import mygene + mg = mygene.MyGeneInfo() + ids = a.var_names.astype(str).tolist() + print(f"[dahlin-mech] mapping {len(ids)} genes …", 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) + a = a[:, keep].copy(); a.var_names = syms[keep] + a.var_names_make_unique() + print(f"[dahlin-mech] converted, shape now {a.shape}", flush=True) + + preds = pd.read_csv(PRED, index_col=0) + a.obs = a.obs.join(preds[["pred_label", "pred_conf", "genotype"]], how="left") + print(f"[dahlin-mech] joined preds; genotype counts {a.obs['genotype'].value_counts().to_dict()}", + flush=True) + + a_de = a.copy() + sc.pp.normalize_total(a_de, target_sum=1e4) + sc.pp.log1p(a_de) + + rows = [] + lg_class_min = 300 + class_counts = a.obs["pred_label"].value_counts() + for cls in class_counts[class_counts >= lg_class_min].index: + mask = a.obs["pred_label"] == cls + sub = a_de[mask].copy() + if len(sub.obs["genotype"].unique()) < 2: + continue + gc = sub.obs["genotype"].value_counts() + if gc.min() < 30: + continue + try: + sc.tl.rank_genes_groups(sub, "genotype", method="wilcoxon", n_genes=40, + use_raw=False) + for grp in ["Kit_W41", "WT"]: + if grp not in sub.uns["rank_genes_groups"]["names"].dtype.names: + continue + names = list(sub.uns["rank_genes_groups"]["names"][grp][:15]) + lfcs = list(sub.uns["rank_genes_groups"]["logfoldchanges"][grp][:15]) + padjs = list(sub.uns["rank_genes_groups"]["pvals_adj"][grp][:15]) + for g, l, p in zip(names, lfcs, padjs): + rows.append({"class": cls, "up_in": grp, "gene": g, + "logfc": round(float(l), 3), + "padj": float(p)}) + except Exception as exc: + print(f"[dahlin-mech] DE failed for {cls}: {exc}") + pd.DataFrame(rows).to_csv(OUT / "67_dahlin_within_class_de.csv", index=False) + + MODULES = { + "MYC_targets": ["Myc", "Mycn", "Nme1", "Odc1", "Ncl", "Ppat", "Cad", + "Ldha", "Nop58"], + "Integrated_stress": ["Atf4", "Ddit3", "Chac1", "Trib3", "Asns", "Gdf15", + "Sesn2", "Slc7a11"], + "Apoptosis_pro": ["Bax", "Bak1", "Bad", "Bid", "Bbc3", "Pmaip1", "Casp3", + "Casp9", "Trp53"], + "Apoptosis_anti": ["Bcl2", "Bcl2l1", "Mcl1", "Bcl2l2"], + "Erythroid_dev": ["Klf1", "Gata1", "Gata2", "Epor", "Hba-a1", "Hbb-b1"], + "Cell_cycle": ["Ccnd1", "Ccnd2", "Ccne1", "Ccna2", "Ccnb1", "Cdk1", + "Mki67", "Top2a", "Pcna", "Mcm2", "Mcm3", "Mcm4"], + "Kit_signaling": ["Kit", "Kitl", "Sos1", "Rasgrf1", "Grb2", "Pik3ca"], + } + for name, genes in MODULES.items(): + present = [g for g in genes if g in a_de.var_names] + if not present: continue + sc.tl.score_genes(a_de, gene_list=present, score_name=f"pw_{name}", + random_state=0, use_raw=False) + + print("\n[dahlin-mech] Kit_W41 vs WT module score contrast per class:", flush=True) + mod_rows = [] + for cls in class_counts[class_counts >= lg_class_min].index: + mask = a_de.obs["pred_label"] == cls + for name in MODULES.keys(): + col = f"pw_{name}" + if col not in a_de.obs: continue + s = a_de.obs.loc[mask, col].astype(float).values + g = a_de.obs.loc[mask, "genotype"].astype(str).values + cv = s[g == "Kit_W41"]; wv = s[g == "WT"] + if len(cv) < 20 or len(wv) < 20: continue + try: + _, p = mannwhitneyu(cv, wv, alternative="two-sided") + except Exception: p = 1.0 + delta = cv.mean() - wv.mean() + star = "***" if p < 1e-6 else "**" if p < 1e-3 else "*" if p < 0.05 else "" + print(f" {cls:15s} {name:20s} delta={delta:+.4f} p={p:.2e} {star}") + mod_rows.append({"class": cls, "module": name, + "delta_Kit_minus_WT": round(delta, 4), + "MannU_p": p}) + pd.DataFrame(mod_rows).to_csv(OUT / "67_dahlin_module_scores.csv", index=False) + + print(f"\n[dahlin-mech] complete. Outputs in {OUT}/") + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/70_prototype_geometry.py b/scripts/analysis/70_prototype_geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..3bc55f504d87567ba528f1f24d2476873e0b2cf3 --- /dev/null +++ b/scripts/analysis/70_prototype_geometry.py @@ -0,0 +1,107 @@ +"""cross-system prototype geometry: intra/cross cosine + participation-ratio effective dim.""" +from __future__ import annotations +from pathlib import Path +import json, numpy as np, pandas as pd + +CKPT = Path("/home/bcheng/PRISM/checkpoints") +OUT = Path("/home/bcheng/PRISM/discovery") +OUT.mkdir(exist_ok=True) + +SYSTEMS = ["pan_skin", "hematopoiesis", "pancreas"] + + +def load_prototypes(sys): + """prototypes shape (K, n_sub, D) — sub-centers averaged per class.""" + import torch + ck = torch.load(CKPT / sys / "marker" / "panda_final.pt", map_location="cpu", weights_only=False) + P = ck["prototypes"] + classes = ck["classes"] + if P.ndim == 3: + P = P.mean(axis=1) + elif P.ndim == 2 and P.shape[0] != len(classes): + n_sub = P.shape[0] // len(classes) + P = P.reshape(len(classes), n_sub, -1).mean(axis=1) + P = P / (np.linalg.norm(P, axis=1, keepdims=True) + 1e-8) + return P, classes + + +def effective_dim(P): + """participation ratio Tr(Sigma)^2 / Tr(Sigma^2).""" + G = P @ P.T + eig = np.linalg.eigvalsh(G) + eig = np.clip(eig, 0, None) + return float((eig.sum() ** 2) / (eig ** 2).sum() + 1e-12) + + +all_P, all_C, all_S = [], [], [] +eff_dims = {} + +for sys in SYSTEMS: + P, classes = load_prototypes(sys) + print(f"[{sys}] K={len(classes)} classes={classes}", flush=True) + C = P @ P.T + off = C[~np.eye(len(classes), dtype=bool)] + ed = effective_dim(P) + eff_dims[sys] = { + "K": len(classes), + "effective_dim": ed, + "cos_offdiag_mean": float(off.mean()), + "cos_offdiag_std": float(off.std()), + "cos_offdiag_max": float(off.max()), + "cos_offdiag_min": float(off.min()), + } + pd.DataFrame(C, index=classes, columns=classes).to_csv( + OUT / f"70_prototype_intra_cosine_{sys}.csv" + ) + all_P.append(P); all_C.extend(classes); all_S.extend([sys] * len(classes)) + print(f"[{sys}] effective_dim={ed:.2f} K={len(classes)} off-diag cos mean={off.mean():+.3f} max={off.max():+.3f}", flush=True) + +json.dump(eff_dims, open(OUT / "70_prototype_effective_dim.json", "w"), indent=2) + +P_all = np.vstack(all_P) +cos_all = P_all @ P_all.T +labels_full = [f"{s}:{c}" for s, c in zip(all_S, all_C)] +pd.DataFrame(cos_all, index=labels_full, columns=labels_full).to_csv( + OUT / "70_prototype_full_29x29.csv" +) + +# cross-system pairs: mask same-system, rank by cosine +rows = [] +for i in range(len(labels_full)): + for j in range(i + 1, len(labels_full)): + if all_S[i] == all_S[j]: + continue + rows.append({ + "sys_a": all_S[i], "class_a": all_C[i], + "sys_b": all_S[j], "class_b": all_C[j], + "cos": float(cos_all[i, j]), + }) +pairs = pd.DataFrame(rows).sort_values("cos", ascending=False) +pairs.to_csv(OUT / "70_prototype_cross_system_pairs.csv", index=False) +print(f"[cross] top-10 positive cross-system pairs:", flush=True) +print(pairs.head(10).to_string(index=False), flush=True) +print(f"[cross] top-10 negative cross-system pairs:", flush=True) +print(pairs.tail(10).to_string(index=False), flush=True) + +# null: random K=29 128-d unit vectors +rng = np.random.default_rng(0) +R = rng.standard_normal((29, 128)) +R = R / np.linalg.norm(R, axis=1, keepdims=True) +Rc = R @ R.T +null = Rc[~np.eye(29, dtype=bool)] +print(f"\n[null] random 29x128 unit vectors: cos mean={null.mean():+.3f} std={null.std():.3f} max={null.max():+.3f}", flush=True) +print(f"[obs] observed cross-system cos: mean={pairs['cos'].mean():+.3f} std={pairs['cos'].std():.3f} max={pairs['cos'].max():+.3f}", flush=True) + +summary = { + "systems": SYSTEMS, + "eff_dims": eff_dims, + "null_cross_cos_mean": float(null.mean()), + "null_cross_cos_std": float(null.std()), + "obs_cross_cos_mean": float(pairs["cos"].mean()), + "obs_cross_cos_std": float(pairs["cos"].std()), + "obs_cross_cos_max": float(pairs["cos"].max()), + "top10_pos_pairs": pairs.head(10).to_dict("records"), + "top10_neg_pairs": pairs.tail(10).to_dict("records"), +} +json.dump(summary, open(OUT / "70_prototype_summary.json", "w"), indent=2, default=str) +print(f"\nwrote {OUT}/70_prototype_summary.json", flush=True) diff --git a/scripts/analysis/72_emergent_axes.py b/scripts/analysis/72_emergent_axes.py new file mode 100644 index 0000000000000000000000000000000000000000..0e857af2dddedbb512be147dd84ddb5535169c2a --- /dev/null +++ b/scripts/analysis/72_emergent_axes.py @@ -0,0 +1,74 @@ +"""per-class-residual PCA on the 128d projections; skin only (only cached).""" +from pathlib import Path +import json, warnings, numpy as np, pandas as pd, anndata as ad +warnings.filterwarnings("ignore") +from sklearn.decomposition import PCA + +OUT = Path("/home/bcheng/PRISM/discovery"); OUT.mkdir(exist_ok=True) + +P = ad.read_h5ad("/home/bcheng/PRISM/discovery/pan_skin/marker/50_aldrich_projections.h5ad") +Z = np.asarray(P.obsm["Z_projection"]) +pred = (P.obs["pred_bbse_label"] if "pred_bbse_label" in P.obs + else P.obs["pred_label"]).astype(str).values +print(f"[skin] Z shape {Z.shape}, n_pred_classes={len(np.unique(pred))}", flush=True) + +Zres = np.zeros_like(Z) +for cls in np.unique(pred): + m = pred == cls + if m.sum() < 2: continue + Zres[m] = Z[m] - Z[m].mean(0, keepdims=True) + +pca = PCA(n_components=20, random_state=0).fit(Zres) +ev = pca.explained_variance_ratio_ +print(f"[skin] top-10 residual PC EV: {[f'{e:.4f}' for e in ev[:10]]}", flush=True) +print(f"[skin] cumulative top-10: {ev[:10].cumsum()[-1]:.3f}", flush=True) +print(f"[skin] cumulative top-20: {ev.cumsum()[-1]:.3f}", flush=True) + +scores = pca.transform(Zres) # (N, 20) +aux_cols = {} +if "log10_counts" in P.obs: + aux_cols["log10_counts"] = P.obs["log10_counts"].astype(float).values +elif "n_counts" in P.obs: + aux_cols["log10_counts"] = np.log10(P.obs["n_counts"].astype(float).values + 1) +if "missing_hvg_frac" in P.obs: + aux_cols["missing_hvg_frac"] = P.obs["missing_hvg_frac"].astype(float).values +if "max_cos" in P.obs: + aux_cols["max_cos"] = P.obs["max_cos"].astype(float).values +else: + import torch + ck = torch.load("/home/bcheng/PRISM/checkpoints/pan_skin/panda_final.pt", + map_location="cpu", weights_only=False) + protos = ck["prototypes"] + protos = protos / (np.linalg.norm(protos, axis=1, keepdims=True) + 1e-8) + Zn = Z / (np.linalg.norm(Z, axis=1, keepdims=True) + 1e-8) + aux_cols["max_cos"] = (Zn @ protos.T).max(axis=1) + +if "genotype" in P.obs: + genotype_bin = (P.obs["genotype"] == "En1-cKO").astype(int).values + aux_cols["genotype_cKO"] = genotype_bin.astype(float) + +print(f"[skin] auxiliaries: {list(aux_cols.keys())}", flush=True) + +rows = [] +for pc in range(10): + row = {"PC": f"PC{pc+1}", "EV": float(ev[pc])} + for aux_name, aux_vals in aux_cols.items(): + r = float(np.corrcoef(scores[:, pc], aux_vals)[0, 1]) + row[f"corr_{aux_name}"] = r + rows.append(row) +axes_df = pd.DataFrame(rows) +axes_df.to_csv(OUT / "72_emergent_axes_skin.csv", index=False) +print(axes_df.to_string(index=False), flush=True) + +loadings = pca.components_[:5] # (5, 128) +np.save(OUT / "72_emergent_axes_skin_pc_loadings.npy", loadings) + +summary = { + "system": "pan_skin", + "total_ev_top10": float(ev[:10].sum()), + "total_ev_top20": float(ev.sum()), + "top10_ev": [float(e) for e in ev[:10]], + "auxiliaries": list(aux_cols.keys()), +} +json.dump(summary, open(OUT / "72_emergent_axes_summary.json", "w"), indent=2) +print(f"\nwrote {OUT}/72_emergent_axes_summary.json", flush=True) diff --git a/scripts/analysis/73_novel_populations_dahlin.py b/scripts/analysis/73_novel_populations_dahlin.py new file mode 100644 index 0000000000000000000000000000000000000000..c1e2d124f1302940b2b5b7759b379e3673e63d8d --- /dev/null +++ b/scripts/analysis/73_novel_populations_dahlin.py @@ -0,0 +1,123 @@ +"""novel-population discovery on Dahlin via PANDA's abstain gate. mirrors 71_.""" +from pathlib import Path +import json, warnings, sys, pickle, numpy as np, pandas as pd, anndata as ad, scanpy as sc, torch +import scipy.sparse as sp +warnings.filterwarnings("ignore") +from pathlib import Path as _P_root +ROOT = _P_root(__file__).resolve().parents[2] +ROOT_STR = str(ROOT) +sys.path.insert(0, ROOT_STR) +from panda import PANDAEncoder + +sc.settings.verbosity = 1 +OUT = Path(f"{ROOT_STR}/discovery"); OUT.mkdir(exist_ok=True) +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def load_dahlin_raw(): + """load Dahlin counts and convert ENSMUSG ids to gene symbols.""" + D_DIR = Path(f"{ROOT_STR}/data/corpus/hematopoiesis/held_out_unlabeled/dahlin_extract") + GT = {"SIGAB1":"WT","SIGAC1":"WT","SIGAD1":"WT","SIGAF1":"WT","SIGAG1":"WT", + "SIGAH1":"WT","SIGAG8":"Kit_W41","SIGAH8":"Kit_W41"} + parts = [] + for f in sorted(D_DIR.glob("*.txt.gz")): + sample = f.name.split("_")[1].split(".")[0] + 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"] = GT.get(sample, "unknown") + var = pd.DataFrame(index=df.index.astype(str)) + parts.append(ad.AnnData(X=X, obs=obs, var=var)) + a = ad.concat(parts, join="outer", label="_batch") + import mygene + mg = mygene.MyGeneInfo() + res = mg.querymany(a.var_names.astype(str).tolist(), 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) + a = a[:, keep].copy(); a.var_names = syms[keep]; a.var_names_make_unique() + return a + + +def project(adata, ckpt_path, shared_hvgs, mu, sig, pca): + ck = torch.load(ckpt_path, map_location=DEVICE, weights_only=False) + classes = ck["classes"]; datasets = ck["datasets"] + m = PANDAEncoder(n_pca=50, n_classes=len(classes), n_datasets=len(datasets)).to(DEVICE).eval() + m.load_state_dict(ck["model"]) + protos = ck["prototypes"] + protos = protos / (np.linalg.norm(protos, axis=1, keepdims=True) + 1e-8) + + G = len(shared_hvgs); hvg2i = {g: i for i, g in enumerate(shared_hvgs)} + common = [g for g in adata.var_names.astype(str) if g in hvg2i] + a_c = adata[:, 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((adata.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, adata.n_obs, 4096): + xb = torch.from_numpy(Xpca[i:i+4096]).to(DEVICE) + aux = torch.zeros(len(xb), 2, device=DEVICE) + all_z.append(m(xb, aux, lam_dann=0.0)["z"].cpu().numpy()) + Z = np.concatenate(all_z) + Zn = Z / (np.linalg.norm(Z, axis=1, keepdims=True) + 1e-8) + cos = Zn @ protos.T + return Z, cos, classes + + +print("[load] Dahlin ENSMUSG→symbol", flush=True) +raw = load_dahlin_raw() +print(f"[load] {raw.shape}", flush=True) + +stats_h = np.load(f"{ROOT_STR}/data/corpus/hematopoiesis/harmonized/corpus_stats.npz", + allow_pickle=True) +shared_hvgs = [str(g) for g in stats_h["shared_hvgs"]] +pca_h = pickle.load(open(f"{ROOT_STR}/data/corpus/hematopoiesis/harmonized/pca_basis.pkl","rb")) +print(f"[project] running PANDA on Dahlin", flush=True) +Z, cos, classes = project(raw, f"{ROOT_STR}/checkpoints/hematopoiesis/panda_final.pt", + shared_hvgs, stats_h["mean"], stats_h["std"], pca_h) +max_cos = cos.max(axis=1); pred = np.array([classes[i] for i in cos.argmax(axis=1)]) +print(f"[cos] q: 5%={np.quantile(max_cos,0.05):.3f} 50%={np.quantile(max_cos,0.5):.3f} 95%={np.quantile(max_cos,0.95):.3f}", flush=True) +print(f"[cos] cos<0.5 n={int((max_cos<0.5).sum())} cos<0.3 n={int((max_cos<0.3).sum())}", flush=True) + +thr = max(0.5, float(np.quantile(max_cos, 0.05))) +mask = max_cos < thr +print(f"[novel] threshold cos<{thr:.3f}: {int(mask.sum())} cells", flush=True) + +sub = raw[mask].copy() +sub.obsm["Z"] = Z[mask] +sub.obs["max_cos"] = max_cos[mask] +sub.obs["pred"] = pred[mask] +sc.pp.neighbors(sub, use_rep="Z", n_neighbors=20) +sc.tl.leiden(sub, resolution=0.8, random_state=0) +print(f"[cluster] {sub.obs['leiden'].nunique()} clusters", flush=True) + +sc.pp.normalize_total(sub, target_sum=1e4); sc.pp.log1p(sub) +sc.tl.rank_genes_groups(sub, "leiden", method="wilcoxon", n_genes=25) + +rows = [] +for cl in sub.obs["leiden"].unique(): + genes = list(sub.uns["rank_genes_groups"]["names"][cl][:15]) + pvals = list(sub.uns["rank_genes_groups"]["pvals_adj"][cl][:15]) + logfc = list(sub.uns["rank_genes_groups"]["logfoldchanges"][cl][:15]) + pre = sub.obs.loc[sub.obs["leiden"] == cl, "pred"].value_counts() + gt = sub.obs.loc[sub.obs["leiden"] == cl, "genotype"].value_counts() + rows.append({ + "cluster": cl, + "n_cells": int((sub.obs["leiden"] == cl).sum()), + "top_markers": ",".join(genes[:10]), + "top_pvals": ",".join(f"{p:.1e}" for p in pvals[:10]), + "top_logfc": ",".join(f"{f:+.2f}" for f in logfc[:10]), + "pred_pre_abstain": pre.index[0] if len(pre) else "", + "genotype_wt_frac": float(gt.get("WT", 0) / gt.sum()) if len(gt) else 0, + }) + +df = pd.DataFrame(rows).sort_values("n_cells", ascending=False) +df.to_csv(OUT / "73_dahlin_novel_populations.csv", index=False) +print(df.to_string(index=False), flush=True) +print(f"\nwrote {OUT}/73_dahlin_novel_populations.csv", flush=True) diff --git a/scripts/analysis/80_prototype_gene_attribution.py b/scripts/analysis/80_prototype_gene_attribution.py new file mode 100644 index 0000000000000000000000000000000000000000..75cafdf708839e3d227a947a3b63de2883563464 --- /dev/null +++ b/scripts/analysis/80_prototype_gene_attribution.py @@ -0,0 +1,219 @@ +"""integrated-gradients attribution of prototype cosine to input genes, per class.""" +from __future__ import annotations +from pathlib import Path +import warnings, json, pickle, sys, numpy as np, pandas as pd +warnings.filterwarnings("ignore") +import torch +import anndata as ad +import yaml +from pathlib import Path as _P_root +ROOT = _P_root(__file__).resolve().parents[2] +ROOT_STR = str(ROOT) +sys.path.insert(0, ROOT_STR) +from panda import PANDAEncoder + +CKPT_ROOT = Path(f"{ROOT_STR}/checkpoints") +CORP_ROOT = Path(f"{ROOT_STR}/data/corpus") +OUT = Path(f"{ROOT_STR}/discovery"); OUT.mkdir(exist_ok=True) +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +IG_STEPS = 64 +TOP_K_GENES = 30 + +TF_YAML = { + "pan_skin": f"{ROOT_STR}/scripts/pan_skin/known_skin_tfs.yaml", + "hematopoiesis": f"{ROOT_STR}/scripts/hematopoiesis/known_hsc_markers.yaml", + "pancreas": f"{ROOT_STR}/scripts/pancreas/known_pancreas_markers.yaml", +} + + +def load_pca_and_hvgs(sys): + stats = np.load(CORP_ROOT / sys / "harmonized/corpus_stats.npz", allow_pickle=True) + hvgs = [str(g) for g in stats["shared_hvgs"]] + pca = pickle.load(open(CORP_ROOT / sys / "harmonized/pca_basis.pkl", "rb")) + mu = np.asarray(stats["mean"], dtype=np.float32) + sig = np.asarray(stats["std"], dtype=np.float32) + return hvgs, pca, mu, sig + + +def load_model(sys): + ck = torch.load(CKPT_ROOT / sys / "marker" / "panda_final.pt", map_location=DEVICE, weights_only=False) + m = PANDAEncoder(variant="marker", n_pca=50, n_markers=len(ck.get("marker_genes",[])), n_classes=len(ck["classes"]), n_sub=3, n_datasets=len(ck["datasets"])).to(DEVICE).eval() + m.load_state_dict(ck["model"]) + protos = torch.from_numpy(ck["prototypes"]).to(DEVICE) + protos = protos / (protos.norm(dim=1, keepdim=True) + 1e-8) + return m, ck["classes"], protos + + +def load_corpus_pca(sys, model_hvgs): + corp = ad.read_h5ad(CORP_ROOT / sys / "harmonized/corpus.h5ad") + return corp + + +def per_class_pca_mean(corp, sys, hvgs, mu, sig, pca): + import scipy.sparse as sp + import scanpy as sc + hvg2i = {g: i for i, g in enumerate(hvgs)} + # old builds used cell_type instead of canonical_label + label_key = "canonical_label" if "canonical_label" in corp.obs else "cell_type" + y = corp.obs[label_key].astype(str).values + common = [g for g in corp.var_names.astype(str) if g in hvg2i] + a = corp[:, common].copy() + if "log1p" not in a.uns.get("uns_keys", []): + sc.pp.normalize_total(a, target_sum=1e4); sc.pp.log1p(a) + X = a.X.toarray().astype(np.float32) if sp.issparse(a.X) else a.X.astype(np.float32) + Xf = np.zeros((corp.n_obs, len(hvgs)), dtype=np.float32) + cols = np.array([hvg2i[g] for g in common]) + Xf[:, cols] = X + Xz = np.clip((Xf - mu) / sig, -10, 10) + Xpca = pca.transform(Xz).astype(np.float32) + cls_to_mean_pca = {} + cls_to_n = {} + for cls in sorted(set(y)): + mask = y == cls + if mask.sum() == 0: continue + cls_to_mean_pca[cls] = Xpca[mask].mean(axis=0) + cls_to_n[cls] = int(mask.sum()) + return cls_to_mean_pca, cls_to_n + + +def integrated_gradient_pc(model, protos, class_idx, x_class_pca, n_steps=64): + baseline = torch.zeros_like(x_class_pca) + alphas = torch.linspace(1.0 / (2*n_steps), 1 - 1.0/(2*n_steps), n_steps, device=DEVICE) + x_class_pca = x_class_pca.unsqueeze(0) + baseline = baseline.unsqueeze(0) + grad_accum = torch.zeros(1, x_class_pca.shape[-1], device=DEVICE) + for a in alphas: + interp = (baseline + a * (x_class_pca - baseline)).requires_grad_(True) + aux = torch.zeros(1, 2, device=DEVICE) + h = model.trunk(interp) + z_raw = model.projection(h) + z = torch.nn.functional.normalize(z_raw, dim=1) + score = (z * protos[class_idx].unsqueeze(0)).sum() + grad = torch.autograd.grad(score, interp)[0] + grad_accum = grad_accum + grad + grad_accum = grad_accum / n_steps + ig = ((x_class_pca - baseline) * grad_accum).squeeze(0).detach().cpu().numpy() + return ig + + +def load_tf_list(sys): + path = TF_YAML.get(sys) + if path is None or not Path(path).exists(): + return {} + with open(path) as f: + obj = yaml.safe_load(f) + tf_dict = {} + # YAMLs nest under "classes" + root = obj.get("classes", obj) if isinstance(obj, dict) else obj + if isinstance(root, dict): + for k, v in root.items(): + if isinstance(v, list): + tf_dict[k] = [str(x) for x in v] + elif isinstance(v, dict) and "genes" in v: + tf_dict[k] = [str(x) for x in v["genes"]] + return tf_dict + + +def main_system(sys): + print(f"\n===== {sys} =====", flush=True) + hvgs, pca, mu, sig = load_pca_and_hvgs(sys) + print(f"[load] {len(hvgs)} HVGs, PCA components {pca.components_.shape}", flush=True) + + model, classes, protos = load_model(sys) + print(f"[load] model K={len(classes)} classes", flush=True) + + corp = load_corpus_pca(sys, hvgs) + print(f"[corp] {corp.shape}", flush=True) + + cls_to_mean_pca, cls_to_n = per_class_pca_mean(corp, sys, hvgs, mu, sig, pca) + print(f"[centroids] {len(cls_to_mean_pca)} class means computed", flush=True) + + # model class names aren't always 1:1 with corpus labels; take best match + pc_att = np.zeros((len(classes), 50), dtype=np.float32) + used_class_map = {} + for ci, cls_name in enumerate(classes): + if cls_name in cls_to_mean_pca: + src = cls_name + else: + best = None + for k in cls_to_mean_pca: + if k.lower() == cls_name.lower(): + best = k; break + if best is None: + # fall back to the overall mean + src = None + x_pca = np.mean(list(cls_to_mean_pca.values()), axis=0) + else: + src = best + if src is not None: + x_pca = cls_to_mean_pca[src] + used_class_map[cls_name] = src + x_t = torch.from_numpy(x_pca).float().to(DEVICE) + ig = integrated_gradient_pc(model, protos, ci, x_t, n_steps=IG_STEPS) + pc_att[ci] = ig + print(f"[IG] class {ci+1}/{len(classes)}: {cls_name:<25s} src={src} |ig|_1={np.abs(ig).sum():.3f} |ig|_inf={np.abs(ig).max():.3f}", flush=True) + + df_pc = pd.DataFrame(pc_att, index=classes, columns=[f"PC{i+1}" for i in range(50)]) + df_pc.to_csv(OUT / f"80_{sys}_pc_attribution.csv") + + # (K, 50) @ (50, G) = (K, G) + gene_att = pc_att @ pca.components_.astype(np.float32) + np.save(OUT / f"80_{sys}_gene_attribution_full.npy", gene_att) + + tf_dict = load_tf_list(sys) + print(f"[TF] loaded {len(tf_dict)} TF program lists: {list(tf_dict.keys())[:5]}...", flush=True) + + rows = [] + tf_rows = [] + for ci, cls_name in enumerate(classes): + att = gene_att[ci] + pos_idx = np.argsort(att)[::-1][:TOP_K_GENES] + neg_idx = np.argsort(att)[:TOP_K_GENES] + row = { + "class": cls_name, + "src_corpus_class": used_class_map[cls_name], + "top_pos_genes": ",".join([hvgs[i] for i in pos_idx[:20]]), + "top_pos_attribution": ",".join([f"{att[i]:+.3f}" for i in pos_idx[:20]]), + "top_neg_genes": ",".join([hvgs[i] for i in neg_idx[:15]]), + "top_neg_attribution": ",".join([f"{att[i]:+.3f}" for i in neg_idx[:15]]), + } + rows.append(row) + + pos_top_genes_set = set(hvgs[i] for i in pos_idx[:100]) + for prog_name, prog_genes in tf_dict.items(): + prog_in_hvg = [g for g in prog_genes if g in hvgs] + if len(prog_in_hvg) == 0: continue + prog_idx = np.array([hvgs.index(g) for g in prog_in_hvg]) + prog_att = float(att[prog_idx].sum()) + hit_frac = len([g for g in prog_in_hvg if g in pos_top_genes_set]) / len(prog_in_hvg) + tf_rows.append({ + "class": cls_name, + "program": prog_name, + "n_genes_in_hvg": len(prog_in_hvg), + "sum_attribution": prog_att, + "hit_frac_top100_pos": hit_frac, + }) + + pd.DataFrame(rows).to_csv(OUT / f"80_{sys}_gene_attribution.csv", index=False) + if tf_rows: + tf_df = pd.DataFrame(tf_rows).sort_values(["class", "sum_attribution"], ascending=[True, False]) + tf_df.to_csv(OUT / f"80_{sys}_tf_enrichment.csv", index=False) + top_tf = (tf_df.sort_values("sum_attribution", ascending=False) + .groupby("class").head(3) + .sort_values(["class", "sum_attribution"], ascending=[True, False])) + top_tf.to_csv(OUT / f"80_{sys}_top_tf_per_class.csv", index=False) + print(f"\n[TF top-3 per class]", flush=True) + print(top_tf.to_string(index=False), flush=True) + + print(f"\n[wrote] {OUT}/80_{sys}_*", flush=True) + + +for sys in ["pan_skin", "hematopoiesis", "pancreas"]: + try: + main_system(sys) + except Exception as e: + import traceback; traceback.print_exc() + print(f"[!!!] {sys} failed: {e}", flush=True) + continue + +print("\n===== done =====", flush=True) diff --git a/scripts/analysis/81_counterfactual_knockouts.py b/scripts/analysis/81_counterfactual_knockouts.py new file mode 100644 index 0000000000000000000000000000000000000000..921fc07909fdc68a6dd60f32624b603c3b6ac47a --- /dev/null +++ b/scripts/analysis/81_counterfactual_knockouts.py @@ -0,0 +1,112 @@ +"""per-class gene knockouts: zero each candidate gene, remeasure prototype cosine.""" +from __future__ import annotations +from pathlib import Path +import warnings, json, pickle, sys, numpy as np, pandas as pd +warnings.filterwarnings("ignore") +import torch +import anndata as ad +from pathlib import Path as _P_root +ROOT = _P_root(__file__).resolve().parents[2] +ROOT_STR = str(ROOT) +sys.path.insert(0, ROOT_STR) +from panda import PANDAEncoder + +CKPT_ROOT = Path(f"{ROOT_STR}/checkpoints") +CORP_ROOT = Path(f"{ROOT_STR}/data/corpus") +OUT = Path(f"{ROOT_STR}/discovery"); OUT.mkdir(exist_ok=True) +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +TOP_KO_GENES = 100 # only test top-100 attributed genes for KO + + +def compute_ko(sys): + print(f"\n===== {sys} =====", flush=True) + stats = np.load(CORP_ROOT / sys / "harmonized/corpus_stats.npz", allow_pickle=True) + hvgs = [str(g) for g in stats["shared_hvgs"]] + pca = pickle.load(open(CORP_ROOT / sys / "harmonized/pca_basis.pkl", "rb")) + mu = np.asarray(stats["mean"], dtype=np.float32) + sig = np.asarray(stats["std"], dtype=np.float32) + + ck = torch.load(CKPT_ROOT / sys / "marker" / "panda_final.pt", map_location=DEVICE, weights_only=False) + classes = ck["classes"] + model = PANDAEncoder(variant="marker", n_pca=50, n_markers=len(ck.get("marker_genes",[])), n_classes=len(classes), n_sub=3, + n_datasets=len(ck["datasets"])).to(DEVICE).eval() + model.load_state_dict(ck["model"]) + protos = torch.from_numpy(ck["prototypes"]).to(DEVICE) + protos = protos / (protos.norm(dim=1, keepdim=True) + 1e-8) + + import scipy.sparse as sp + import scanpy as sc + corp = ad.read_h5ad(CORP_ROOT / sys / "harmonized/corpus.h5ad") + label_key = "canonical_label" if "canonical_label" in corp.obs else "cell_type" + y = corp.obs[label_key].astype(str).values + hvg2i = {g: i for i, g in enumerate(hvgs)} + common = [g for g in corp.var_names.astype(str) if g in hvg2i] + a = corp[:, common].copy() + sc.pp.normalize_total(a, target_sum=1e4); sc.pp.log1p(a) + X = a.X.toarray().astype(np.float32) if sp.issparse(a.X) else a.X.astype(np.float32) + Xf = np.zeros((corp.n_obs, len(hvgs)), dtype=np.float32) + cols = np.array([hvg2i[g] for g in common]) + Xf[:, cols] = X + + att = np.load(OUT / f"80_{sys}_gene_attribution_full.npy") # (K, G) + + rows = [] + for ci, cls_name in enumerate(classes): + mask = y == cls_name + if mask.sum() == 0: + print(f"[!!] {cls_name}: no cells in corpus", flush=True) + continue + x_gene_log_mean = Xf[mask].mean(axis=0) + x_z = np.clip((x_gene_log_mean - mu) / sig, -10, 10) + x_pca0 = pca.transform(x_z.reshape(1, -1))[0].astype(np.float32) + + with torch.no_grad(): + xt = torch.from_numpy(x_pca0).unsqueeze(0).to(DEVICE) + aux = torch.zeros(1, 2, device=DEVICE) + out = model(xt, aux, lam_dann=0.0) + z0 = out["z"] + s0 = float((z0 * protos[ci].unsqueeze(0)).sum()) + + att_c = att[ci] + cand_idx = np.argsort(-np.abs(att_c))[:TOP_KO_GENES] + + # zero each gene in log space, reproject; z is clipped so recompute per KO + deltas = [] + for gi in cand_idx: + x_gene_ko = x_gene_log_mean.copy() + x_gene_ko[gi] = 0.0 + x_z_ko = np.clip((x_gene_ko - mu) / sig, -10, 10) + x_pca_ko = pca.transform(x_z_ko.reshape(1, -1))[0].astype(np.float32) + with torch.no_grad(): + xtko = torch.from_numpy(x_pca_ko).unsqueeze(0).to(DEVICE) + zko = model(xtko, aux, lam_dann=0.0)["z"] + s_ko = float((zko * protos[ci].unsqueeze(0)).sum()) + deltas.append(s0 - s_ko) # positive delta = drop when KO'd + + deltas = np.array(deltas) + rank = np.argsort(-deltas) + top30 = rank[:30] + rows.append({ + "class": cls_name, + "baseline_cos": s0, + "n_cells_class": int(mask.sum()), + "top_essential_genes": ",".join([hvgs[cand_idx[r]] for r in top30[:20]]), + "top_essential_deltas": ",".join([f"{deltas[r]:+.4f}" for r in top30[:20]]), + "top_essential_baseline_expression": ",".join([f"{x_gene_log_mean[cand_idx[r]]:.2f}" for r in top30[:20]]), + }) + print(f"[{cls_name}] baseline_cos={s0:.4f} top-5 essentials: " + f"{', '.join([f'{hvgs[cand_idx[r]]}(Δ{deltas[r]:+.3f})' for r in top30[:5]])}", + flush=True) + + df = pd.DataFrame(rows) + df.to_csv(OUT / f"81_{sys}_ko_essentials.csv", index=False) + print(f"[wrote] {OUT}/81_{sys}_ko_essentials.csv", flush=True) + + +for sys in ["pan_skin", "hematopoiesis", "pancreas"]: + try: + compute_ko(sys) + except Exception as e: + import traceback; traceback.print_exc() + print(f"[!] {sys} failed: {e}", flush=True) +print("\n=== DONE ===", flush=True) diff --git a/scripts/analysis/82_gene_coattribution_modules.py b/scripts/analysis/82_gene_coattribution_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..a145666623b3a9512fc38dc1b47ced2e9dec7b2c --- /dev/null +++ b/scripts/analysis/82_gene_coattribution_modules.py @@ -0,0 +1,72 @@ +"""cluster the top-500 attributed genes by cross-class attribution correlation.""" +from __future__ import annotations +from pathlib import Path +import warnings, json, pickle, sys, numpy as np, pandas as pd +warnings.filterwarnings("ignore") +from scipy.cluster.hierarchy import linkage, fcluster +from scipy.spatial.distance import squareform + +CORP_ROOT = Path("/home/bcheng/PRISM/data/corpus") +OUT = Path("/home/bcheng/PRISM/discovery"); OUT.mkdir(exist_ok=True) +CKPT_ROOT = Path("/home/bcheng/PRISM/checkpoints") + +TOP_GENES = 500 # keep top-500 by |att| summed across classes +N_MODULES = 15 + + +def main(sys): + print(f"\n===== {sys} =====", flush=True) + stats = np.load(CORP_ROOT / sys / "harmonized/corpus_stats.npz", allow_pickle=True) + hvgs = [str(g) for g in stats["shared_hvgs"]] + import torch + ck = torch.load(CKPT_ROOT / sys / "marker" / "panda_final.pt", map_location="cpu", weights_only=False) + classes = ck["classes"] + + A = np.load(OUT / f"80_{sys}_gene_attribution_full.npy") # (K, G) + print(f"[load] A shape {A.shape}, classes={classes}", flush=True) + + gene_score = np.abs(A).sum(axis=0) + top_idx = np.argsort(-gene_score)[:TOP_GENES] + A_top = A[:, top_idx] + top_genes = [hvgs[i] for i in top_idx] + + Xn = A_top - A_top.mean(axis=0, keepdims=True) + Xn = Xn / (Xn.std(axis=0, keepdims=True) + 1e-8) + C = np.corrcoef(Xn.T) + print(f"[corr] gene-gene C shape {C.shape}, diag mean={np.diag(C).mean():.3f}", flush=True) + + D = 1 - C + np.fill_diagonal(D, 0) + D = np.clip(D, 0, 2) + Z = linkage(squareform(D, checks=False), method="average") + labels = fcluster(Z, t=N_MODULES, criterion="maxclust") + + rows = [] + for mod in sorted(set(labels)): + members = np.where(labels == mod)[0] + if len(members) < 3: continue + member_genes = [top_genes[i] for i in members] + mod_att = A_top[:, members].mean(axis=1) + dom_ci = int(np.argmax(mod_att)) + rows.append({ + "module_id": int(mod), + "size": int(len(members)), + "dominant_class": classes[dom_ci], + "dom_class_mean_att": float(mod_att[dom_ci]), + "member_genes": ",".join(member_genes[:30]), + "n_shown": min(30, len(member_genes)), + }) + + df = pd.DataFrame(rows).sort_values(["dominant_class", "dom_class_mean_att"], ascending=[True, False]) + df.to_csv(OUT / f"82_{sys}_coatt_modules.csv", index=False) + print(f"[wrote] {len(rows)} modules to {OUT}/82_{sys}_coatt_modules.csv", flush=True) + print(df.head(15).to_string(index=False)[:2000], flush=True) + + +for sys in ["pan_skin", "hematopoiesis", "pancreas"]: + try: + main(sys) + except Exception as e: + import traceback; traceback.print_exc() + print(f"[!] {sys}: {e}", flush=True) +print("\n=== DONE ===", flush=True) diff --git a/scripts/analysis/83_prototype_training_trajectory.py b/scripts/analysis/83_prototype_training_trajectory.py new file mode 100644 index 0000000000000000000000000000000000000000..a80b87504113394a713155abc6f06b5e2eaea546 --- /dev/null +++ b/scripts/analysis/83_prototype_training_trajectory.py @@ -0,0 +1,71 @@ +"""prototype cosine drift + participation ratio per curriculum stage.""" +from pathlib import Path +import json, numpy as np, pandas as pd, torch + +CKPT = Path("/home/bcheng/PRISM/checkpoints") +OUT = Path("/home/bcheng/PRISM/discovery"); OUT.mkdir(exist_ok=True) +SYSTEMS = ["pan_skin", "hematopoiesis", "pancreas"] +STAGES = ["panda_stage0", "panda_stage1", "panda_stage2", "panda_stage3", "panda_final"] + + +def load_prototypes(sys, stage): + ck = torch.load(CKPT / sys / f"{stage}.pt", map_location="cpu", weights_only=False) + # older checkpoints keep prototypes inside model state; newer ones at top level + if "prototypes" in ck: + P = ck["prototypes"] + else: + P = ck["model"]["prototypes"] + if isinstance(P, torch.Tensor): P = P.numpy() + P = P / (np.linalg.norm(P, axis=1, keepdims=True) + 1e-8) + return P, ck["classes"] + + +def eff_dim(P): + G = P @ P.T + eig = np.clip(np.linalg.eigvalsh(G), 0, None) + return float(eig.sum()**2 / (eig**2).sum() + 1e-12) + + +for sys in SYSTEMS: + print(f"\n===== {sys} =====", flush=True) + P_per_stage, classes = {}, None + for st in STAGES: + try: + P, cls = load_prototypes(sys, st) + P_per_stage[st] = P + classes = cls + print(f"[{st}] P shape={P.shape} eff_dim={eff_dim(P):.3f}", flush=True) + except FileNotFoundError: + print(f"[{st}] missing") + continue + + stages_avail = list(P_per_stage.keys()) + drift_rows = [] + for i, s1 in enumerate(stages_avail[:-1]): + s2 = stages_avail[i+1] + P1 = P_per_stage[s1]; P2 = P_per_stage[s2] + for ci, cn in enumerate(classes): + cos_shift = float((P1[ci] * P2[ci]).sum()) + drift_rows.append({ + "class": cn, + "from_stage": s1.replace("panda_", ""), + "to_stage": s2.replace("panda_", ""), + "cos_shift": cos_shift, + "angle_deg": float(np.degrees(np.arccos(np.clip(cos_shift, -1, 1)))), + }) + pd.DataFrame(drift_rows).to_csv(OUT / f"83_{sys}_prototype_trajectory.csv", index=False) + + ed_rows = [{"stage": st.replace("panda_", ""), + "eff_dim": eff_dim(P), + "K": P.shape[0]} + for st, P in P_per_stage.items()] + pd.DataFrame(ed_rows).to_csv(OUT / f"83_{sys}_effdim_by_stage.csv", index=False) + print(f"[eff_dim by stage] {[(r['stage'], round(r['eff_dim'], 3)) for r in ed_rows]}", flush=True) + + if "panda_stage0" in P_per_stage and "panda_final" in P_per_stage: + P0 = P_per_stage["panda_stage0"] + Pf = P_per_stage["panda_final"] + per_class_total_cos = (P0 * Pf).sum(axis=1) + print(f"[stage0 -> final] per-class cos: {[(classes[i], round(float(per_class_total_cos[i]), 3)) for i in range(len(classes))]}", flush=True) + +print("\n===== DONE =====", flush=True) diff --git a/scripts/analysis/84_adversary_purification.py b/scripts/analysis/84_adversary_purification.py new file mode 100644 index 0000000000000000000000000000000000000000..ef4af27ee8c201dcdc84b7b57f9f06dcf3ca8750 --- /dev/null +++ b/scripts/analysis/84_adversary_purification.py @@ -0,0 +1,125 @@ +"""probe dataset + depth adversary heads on training data to check trunk invariance.""" +from pathlib import Path +import json, warnings, pickle, sys, numpy as np, pandas as pd, torch +warnings.filterwarnings("ignore") +import anndata as ad, scanpy as sc, scipy.sparse as sp +from pathlib import Path as _P_root +ROOT = _P_root(__file__).resolve().parents[2] +ROOT_STR = str(ROOT) +sys.path.insert(0, ROOT_STR) +from panda import PANDAEncoder + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +CKPT = Path(f"{ROOT_STR}/checkpoints") +CORP = Path(f"{ROOT_STR}/data/corpus") +OUT = Path(f"{ROOT_STR}/discovery"); OUT.mkdir(exist_ok=True) + + +def project_corpus(sys, subsample=20000): + stats = np.load(CORP / sys / "harmonized/corpus_stats.npz", allow_pickle=True) + hvgs = [str(g) for g in stats["shared_hvgs"]] + mu = np.asarray(stats["mean"], dtype=np.float32) + sig = np.asarray(stats["std"], dtype=np.float32) + pca = pickle.load(open(CORP / sys / "harmonized/pca_basis.pkl", "rb")) + + ck = torch.load(CKPT / sys / "panda_final.pt", map_location=DEVICE, weights_only=False) + m = PANDAEncoder(n_pca=50, n_classes=len(ck["classes"]), + n_datasets=len(ck["datasets"])).to(DEVICE).eval() + m.load_state_dict(ck["model"]) + + corp = ad.read_h5ad(CORP / sys / "harmonized/corpus.h5ad") + print(f"[{sys}] corpus {corp.shape}", flush=True) + + dcol_candidates = ["dataset_id", "dataset", "sample", "batch"] + dcol = None + for c in dcol_candidates: + if c in corp.obs.columns: + dcol = c; break + if dcol is None: + # HSC corpus is Weinreb-only (single dataset) + corp.obs["dataset_id"] = "single" + dcol = "dataset_id" + print(f"[{sys}] dataset col={dcol} n_unique={corp.obs[dcol].nunique()} train datasets in ckpt={len(ck['datasets'])}", flush=True) + + if corp.n_obs > subsample: + rng = np.random.default_rng(0) + idx = rng.choice(corp.n_obs, size=subsample, replace=False) + corp = corp[idx].copy() + print(f"[{sys}] subsampled to {corp.n_obs} cells", flush=True) + + hvg2i = {g: i for i, g in enumerate(hvgs)} + common = [g for g in corp.var_names.astype(str) if g in hvg2i] + a_c = corp[:, 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((corp.n_obs, len(hvgs)), dtype=np.float32) + cols = np.array([hvg2i[g] for g in common]) + Xf[:, cols] = X + Xz = np.clip((Xf - mu) / sig, -10, 10) + Xpca = pca.transform(Xz).astype(np.float32) + + if "n_counts" in corp.obs.columns: + n_counts = corp.obs["n_counts"].astype(float).values + elif "total_counts" in corp.obs.columns: + n_counts = corp.obs["total_counts"].astype(float).values + else: + # X is log-normalised; expm1 first for approximate raw sum + n_counts = np.expm1(X.sum(axis=1)) + log10c = np.log10(n_counts + 1) + log10c_z = (log10c - log10c.mean()) / (log10c.std() + 1e-8) + + dset = corp.obs[dcol].astype(str).values + train_dsets = ck["datasets"] + dset_to_idx = {d: i for i, d in enumerate(train_dsets)} + y_dset = np.array([dset_to_idx.get(d, -1) for d in dset]) + valid = y_dset >= 0 + print(f"[{sys}] valid rows for dataset-eval: {valid.sum()} / {len(y_dset)}", flush=True) + + dom_preds = np.zeros((corp.n_obs, len(train_dsets)), dtype=np.float32) + depth_preds = np.zeros(corp.n_obs, dtype=np.float32) + with torch.no_grad(): + for i in range(0, corp.n_obs, 2048): + xb = torch.from_numpy(Xpca[i:i+2048]).to(DEVICE) + aux = torch.zeros(len(xb), 2, device=DEVICE) + out = m(xb, aux, lam_dann=0.0) # lam=0 disables GRL at inference + dom_preds[i:i+2048] = out["dom"].cpu().numpy() + depth_preds[i:i+2048] = out["depth"].cpu().numpy().squeeze() + + from sklearn.metrics import accuracy_score, top_k_accuracy_score, mean_squared_error, r2_score + result = {"system": sys, "n_cells": int(corp.n_obs), "n_train_datasets": len(train_dsets)} + if valid.sum() > 0 and len(train_dsets) > 1: + dom_pred_argmax = dom_preds[valid].argmax(axis=1) + acc = float(accuracy_score(y_dset[valid], dom_pred_argmax)) + chance = 1.0 / len(train_dsets) + result.update({ + "dataset_adv_accuracy": acc, + "dataset_adv_chance": chance, + "dataset_adv_above_chance": acc - chance, + "dataset_adv_random_baseline_test": "None (single-dataset)" if len(train_dsets) == 1 else f"n_datasets={len(train_dsets)}, chance={chance:.3f}", + }) + + mse_depth = float(mean_squared_error(log10c_z, depth_preds)) + r2_depth = float(r2_score(log10c_z, depth_preds)) + result.update({ + "depth_adv_mse_z": mse_depth, + "depth_adv_r2_z": r2_depth, + "depth_target_std_z": float(log10c_z.std()), + }) + + print(f"[{sys}] dom_adv_acc={result.get('dataset_adv_accuracy', 'NA')} vs chance={result.get('dataset_adv_chance', 'NA')}", flush=True) + print(f"[{sys}] depth_adv MSE_z={mse_depth:.4f} R²_z={r2_depth:.4f} (R²≤0 ⇒ trunk fully depth-invariant)", flush=True) + + return result + + +all_results = {} +for sys in ["pan_skin", "hematopoiesis", "pancreas"]: + try: + all_results[sys] = project_corpus(sys) + except Exception as e: + import traceback; traceback.print_exc() + all_results[sys] = {"error": str(e)} + +json.dump(all_results, open(OUT / "84_adversary_purification.json", "w"), indent=2, default=str) +print(f"\nwrote {OUT}/84_adversary_purification.json", flush=True) +print(json.dumps(all_results, indent=2, default=str), flush=True) diff --git a/scripts/analysis/85_hessian_gene_interactions.py b/scripts/analysis/85_hessian_gene_interactions.py new file mode 100644 index 0000000000000000000000000000000000000000..bcab84d284dbf07d1beba4636ee8458020bc26b1 --- /dev/null +++ b/scripts/analysis/85_hessian_gene_interactions.py @@ -0,0 +1,129 @@ +"""per-class 20x20 hessian of prototype cosine over top-attributed genes.""" +from __future__ import annotations +from pathlib import Path +import warnings, json, pickle, sys, numpy as np, pandas as pd, torch +warnings.filterwarnings("ignore") +import anndata as ad, scanpy as sc, scipy.sparse as sp +from pathlib import Path as _P_root +ROOT = _P_root(__file__).resolve().parents[2] +ROOT_STR = str(ROOT) +sys.path.insert(0, ROOT_STR) +from panda import PANDAEncoder + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +CKPT = Path(f"{ROOT_STR}/checkpoints") +CORP = Path(f"{ROOT_STR}/data/corpus") +OUT = Path(f"{ROOT_STR}/discovery"); OUT.mkdir(exist_ok=True) +TOP_K_GENES = 20 + + +def compute_hessian_for_class(model, protos, class_idx, x_pca_c, pca_V, mu, sig, gene_idx_subset): + # parametrise perturbation as delta in gene-z space of the D selected genes + V_sub = torch.from_numpy(pca_V[:, gene_idx_subset]).float().to(DEVICE) # (50, D) + + def f_of_delta(delta): + x = x_pca_c + V_sub @ delta + x = x.clamp(-10, 10).unsqueeze(0) + aux = torch.zeros(1, 2, device=DEVICE) + h = model.trunk(x) + z_raw = model.projection(h) + z = torch.nn.functional.normalize(z_raw, dim=1) + return (z * protos[class_idx].unsqueeze(0)).sum() + + D = len(gene_idx_subset) + delta0 = torch.zeros(D, device=DEVICE) + H = torch.autograd.functional.hessian(f_of_delta, delta0) + return H.detach().cpu().numpy() + + +def main(sys): + print(f"\n===== {sys} =====", flush=True) + stats = np.load(CORP / sys / "harmonized/corpus_stats.npz", allow_pickle=True) + hvgs = [str(g) for g in stats["shared_hvgs"]] + mu = np.asarray(stats["mean"], dtype=np.float32) + sig = np.asarray(stats["std"], dtype=np.float32) + pca = pickle.load(open(CORP / sys / "harmonized/pca_basis.pkl", "rb")) + + ck = torch.load(CKPT / sys / "marker" / "panda_final.pt", map_location=DEVICE, weights_only=False) + model = PANDAEncoder(variant="marker", n_pca=50, n_markers=len(ck.get("marker_genes",[])), n_classes=len(ck["classes"]), n_sub=3, + n_datasets=len(ck["datasets"])).to(DEVICE).eval() + model.load_state_dict(ck["model"]) + classes = ck["classes"] + protos = torch.from_numpy(ck["prototypes"]).to(DEVICE) + protos = protos / (protos.norm(dim=1, keepdim=True) + 1e-8) + + corp = ad.read_h5ad(CORP / sys / "harmonized/corpus.h5ad") + label_key = "canonical_label" if "canonical_label" in corp.obs else "cell_type" + y = corp.obs[label_key].astype(str).values + hvg2i = {g: i for i, g in enumerate(hvgs)} + common = [g for g in corp.var_names.astype(str) if g in hvg2i] + a = corp[:, common].copy() + sc.pp.normalize_total(a, target_sum=1e4); sc.pp.log1p(a) + X = a.X.toarray().astype(np.float32) if sp.issparse(a.X) else a.X.astype(np.float32) + Xf = np.zeros((corp.n_obs, len(hvgs)), dtype=np.float32) + cols = np.array([hvg2i[g] for g in common]) + Xf[:, cols] = X + + A = np.load(OUT / f"80_{sys}_gene_attribution_full.npy") + + K = len(classes) + H_all = np.zeros((K, TOP_K_GENES, TOP_K_GENES), dtype=np.float32) + top_gene_names_per_class = [] + pair_rows = [] + + for ci, cls_name in enumerate(classes): + mask = y == cls_name + if mask.sum() == 0: + print(f"[!] {cls_name}: no cells in corpus", flush=True) + continue + x_gene_log_mean = Xf[mask].mean(axis=0) + x_z = np.clip((x_gene_log_mean - mu) / sig, -10, 10) + x_pca_c = torch.from_numpy(pca.transform(x_z.reshape(1, -1))[0]).float().to(DEVICE) + + att_c = A[ci] + gene_idx = np.argsort(-np.abs(att_c))[:TOP_K_GENES] + top_names = [hvgs[i] for i in gene_idx] + top_gene_names_per_class.append(top_names) + + H = compute_hessian_for_class(model, protos, ci, x_pca_c, + pca.components_.astype(np.float32), mu, sig, + gene_idx.tolist()) + H_all[ci] = H + + off = H.copy() + np.fill_diagonal(off, 0) + rows, cols_ = np.triu_indices(TOP_K_GENES, k=1) + vals = off[rows, cols_] + order = np.argsort(-np.abs(vals))[:30] + for r_i in order: + gi, gj = int(rows[r_i]), int(cols_[r_i]) + pair_rows.append({ + "class": cls_name, + "gene_a": top_names[gi], + "gene_b": top_names[gj], + "hessian_off_diag": float(vals[r_i]), + "abs_h": float(abs(vals[r_i])), + "attribution_a": float(att_c[gene_idx[gi]]), + "attribution_b": float(att_c[gene_idx[gj]]), + }) + + print(f"[{cls_name}] Hessian |diag|_max={float(np.abs(np.diag(H)).max()):.4f} " + f"|offdiag|_max={float(np.abs(off).max()):.4f} " + f"top-3 pairs: {', '.join([f'{top_names[int(rows[r_i])]}·{top_names[int(cols_[r_i])]}({vals[r_i]:+.4f})' for r_i in order[:3]])}", + flush=True) + + np.save(OUT / f"85_{sys}_hessian_top20.npy", H_all) + with open(OUT / f"85_{sys}_hessian_top20_genes.json", "w") as f: + json.dump({classes[i]: top_gene_names_per_class[i] for i in range(K)}, f, indent=2) + pd.DataFrame(pair_rows).to_csv(OUT / f"85_{sys}_hessian_pairs.csv", index=False) + print(f"[wrote] {OUT}/85_{sys}_hessian_*", flush=True) + + +for sys in ["pan_skin", "hematopoiesis", "pancreas"]: + try: + main(sys) + except Exception as e: + import traceback; traceback.print_exc() + print(f"[!] {sys}: {e}", flush=True) + +print("\n=== DONE ===", flush=True) diff --git a/scripts/analysis/90_dingwall_marker_deep_dive.py b/scripts/analysis/90_dingwall_marker_deep_dive.py new file mode 100644 index 0000000000000000000000000000000000000000..a35bbcfa45f3dceb1c5e1f796b88f014103399a1 --- /dev/null +++ b/scripts/analysis/90_dingwall_marker_deep_dive.py @@ -0,0 +1,87 @@ +"""dingwall marker deep-dive: single rank_genes_groups call cross-referenced against canonical panels.""" +from pathlib import Path +import warnings, json, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp +warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 + +ROOT = Path("/home/bcheng/PRISM") +OUT = ROOT / "discovery/pan_skin/marker" +OUT.mkdir(parents=True, exist_ok=True) + +PANELS = { + "eden-dermal-niche": ["S100a4", "Twist2", "Prrx1", "Pdgfra", "Fap", "Fn1"], + "eccrine-secretory": ["Dcd", "Aqp5", "Muc7", "Cst6", "Krt7"], + "eccrine-ductal": ["Krt77", "Krt5", "Krt14", "Cldn6", "Grhl3"], + "basal-multipotent": ["Krt5", "Krt14", "Trp63", "Itgb4", "Sox2"], + "hair-placode": ["Shh", "Sox9", "Lhx2", "Foxi3", "Wnt10a"], + "melanocyte": ["Dct", "Mlana", "Tyrp1", "Pmel", "Sox10"], + "endothelial": ["Pecam1", "Cdh5", "Kdr", "Flt1"], + "spinous": ["Krt10", "Krt1", "Dsp"], + "basal-IFE": ["Krt5", "Krt14", "Krt15", "Col17a1"], + "immune": ["Ptprc", "Cd68", "Cd3d", "Cd19"], + "fibroblast": ["Col1a1", "Dcn", "Pdgfra"], +} + +CKO_GSMS = {"GSM6833482", "GSM6833483"} # CORRECTED: 480/481 are rttaControl (WT), not cKO +WT_GSMS = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"} # CORRECTED: 4 Cre-neg controls per GEO metadata + +print("[load] Dingwall raw + predictions", flush=True) +raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad") +pred_df = pd.read_csv(ROOT / "discovery/pan_skin/marker/dingwall_predictions.csv") +common = raw.obs_names.intersection(pd.Index(pred_df["cell_id"].astype(str))) +raw = raw[list(common)].copy() +pred_map = dict(zip(pred_df["cell_id"].astype(str), pred_df["pred_label"])) +raw.obs["pred_label"] = pd.Categorical([pred_map.get(c, "unknown") for c in raw.obs_names]) +raw.obs["genotype"] = np.where(raw.obs["sample"].astype(str).isin(list(CKO_GSMS)), "En1-cKO", + np.where(raw.obs["sample"].astype(str).isin(list(WT_GSMS)), "WT", "other")) +print(f"[align] {raw.n_obs} cells across {raw.obs['pred_label'].nunique()} classes", flush=True) + +# subset to classes with >=30 cells for stable Wilcoxon +counts = raw.obs["pred_label"].value_counts() +keep_cls = counts[counts >= 30].index.tolist() +raw = raw[raw.obs["pred_label"].isin(keep_cls)].copy() +raw.obs["pred_label"] = raw.obs["pred_label"].astype(str).astype("category") +print(f"[filter] kept {raw.n_obs} cells × {len(keep_cls)} classes", flush=True) + +sc.pp.normalize_total(raw, target_sum=1e4); sc.pp.log1p(raw) + +# single-call with groupby is much faster than per-class loop +print("[wilcoxon] single-call across all predicted classes...", flush=True) +sc.tl.rank_genes_groups(raw, groupby="pred_label", method="wilcoxon", n_genes=25, use_raw=False) +print("[wilcoxon] done", flush=True) + +rows = [] +for cls in raw.uns["rank_genes_groups"]["names"].dtype.names: + mask = raw.obs["pred_label"] == cls + if mask.sum() < 30: continue + genes = list(raw.uns["rank_genes_groups"]["names"][cls][:20]) + pvals = [float(x) for x in raw.uns["rank_genes_groups"]["pvals_adj"][cls][:20]] + logfc = [float(x) for x in raw.uns["rank_genes_groups"]["logfoldchanges"][cls][:20]] + + top_str = ",".join([f"{g}(LFC{lf:+.1f})" for g, lf in zip(genes[:10], logfc[:10])]) + panel_hits = {} + for pname, plist in PANELS.items(): + hits = [g for g in plist if g in genes[:20]] + panel_hits[pname] = f"{len(hits)}/{len(plist)}: {','.join(hits)}" + gt = raw.obs["genotype"][mask] + ncko = int((gt == "En1-cKO").sum()); nwt = int((gt == "WT").sum()) + frac_cko = ncko / max(1, ncko + nwt) + best_panel = max(panel_hits.items(), + key=lambda x: int(x[1].split("/")[0]) / (int(x[1].split(":")[0].split("/")[1]) + 1e-6)) + + rows.append({ + "predicted_class": cls, + "n_cells": int(mask.sum()), + "top_wilcoxon_markers": top_str, + "min_p_adj_top5": min(pvals[:5], default=float("nan")), + "best_canonical_panel_match": best_panel[0], + "recovery": best_panel[1], + "n_En1_cKO": ncko, + "n_WT": nwt, + "frac_En1_cKO": frac_cko, + }) + +df = pd.DataFrame(rows).sort_values("n_cells", ascending=False) +df.to_csv(OUT / "90_dingwall_marker_deep_dive.csv", index=False) +print(f"[write] {OUT}/90_dingwall_marker_deep_dive.csv ({len(df)} classes)", flush=True) +print() +print(df[["predicted_class", "n_cells", "best_canonical_panel_match", "recovery", "frac_En1_cKO"]].to_string(index=False)) diff --git a/scripts/analysis/91_veres_marker_deep_dive.py b/scripts/analysis/91_veres_marker_deep_dive.py new file mode 100644 index 0000000000000000000000000000000000000000..76eeb67d4747e90ae8bdc6c907528bfb23581fc3 --- /dev/null +++ b/scripts/analysis/91_veres_marker_deep_dive.py @@ -0,0 +1,110 @@ +"""veres marker deep-dive. + +per-predicted-class wilcoxon on raw veres counts, cross-referenced with canonical +adult-beta / SC-alpha / EP panels. output: discovery/pancreas/marker/91_veres_marker_deep_dive.csv +""" +from pathlib import Path +import warnings, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp +warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 + +ROOT = Path("/home/bcheng/PRISM") +OUT = ROOT / "discovery/pancreas/marker" +OUT.mkdir(parents=True, exist_ok=True) + +# human panels (Veres is human hPSC) +PANELS = { + "adult-beta": ["INS", "MAFA", "UCN3", "NKX6-1", "MNX1", "NEUROD1", "PDX1"], + "adult-alpha": ["GCG", "ARX", "IRX2", "IRX1", "MAFB", "TTR"], + "alpha (embryonic prototype)": ["GCG", "ARX", "IRX2", "MAFB"], + "beta (embryonic prototype)": ["INS", "NKX6-1", "MNX1", "NEUROD1", "PDX1"], + "delta": ["SST", "HHEX", "LEPR"], + "gamma": ["PPY", "PYY", "SLC38A4"], + "epsilon": ["GHRL"], + "endocrine-progenitor-early": ["NEUROG3", "CBFA2T3", "BTBD17"], + "endocrine-progenitor-Fev": ["FEV", "INSM1"], + "endocrine-progenitor-primed": ["PAX4", "ARX"], + "acinar": ["PRSS1", "PRSS2", "CEL", "CTRB1"], + "ductal": ["KRT19", "SOX9", "MUC1"], + "endothelial": ["PECAM1", "CDH5", "KDR"], + "immune": ["PTPRC", "CD68"], + "mesenchymal": ["COL1A1", "COL3A1", "DCN"], +} + +# Load Veres via existing loader logic +def load_veres(): + SHARON_DIR = ROOT / "data/corpus/pancreas/held_out_unlabeled/sharon_extract" + parts = [] + for meta_file in sorted(SHARON_DIR.glob("*.cell_metadata.tsv.gz")): + counts_file = str(meta_file).replace("cell_metadata", "processed_counts") + if not Path(counts_file).exists(): continue + meta = pd.read_csv(meta_file, sep="\t", compression="gzip") + counts = pd.read_csv(counts_file, sep="\t", compression="gzip", index_col=0) + obs = meta.set_index("library.barcode") + obs = obs.loc[obs.index.intersection(counts.index)] + counts_al = counts.loc[obs.index] + X = sp.csr_matrix(counts_al.values.astype(np.float32)) + a = ad.AnnData(X=X, obs=obs, var=pd.DataFrame(index=counts_al.columns)) + a.var_names_make_unique() + parts.append(a) + return ad.concat(parts, join="outer") + +print("[load] Veres + predictions", flush=True) +raw = load_veres() +pred_df = pd.read_csv(ROOT / "discovery/pancreas/marker/veres_predictions.csv") +# strip the "veres_" prefix from prediction cell_ids so they align with raw.obs_names +pred_df["cell_id"] = pred_df["cell_id"].astype(str).str.replace(r"^veres_", "", regex=True) +common = raw.obs_names.intersection(pd.Index(pred_df["cell_id"].astype(str))) +raw = raw[list(common)].copy() +pred_map = dict(zip(pred_df["cell_id"].astype(str), pred_df["pred_label"])) +raw.obs["pred_label"] = pd.Categorical([pred_map.get(c, "unknown") for c in raw.obs_names]) +print(f"[align] {raw.n_obs} cells across {raw.obs['pred_label'].nunique()} classes", flush=True) + +counts_s = raw.obs["pred_label"].value_counts() +keep_cls = counts_s[counts_s >= 30].index.tolist() +raw = raw[raw.obs["pred_label"].isin(keep_cls)].copy() +raw.obs["pred_label"] = raw.obs["pred_label"].astype(str).astype("category") +print(f"[filter] {raw.n_obs} cells × {len(keep_cls)} classes", flush=True) + +sc.pp.normalize_total(raw, target_sum=1e4); sc.pp.log1p(raw) +print("[wilcoxon] running...", flush=True) +sc.tl.rank_genes_groups(raw, groupby="pred_label", method="wilcoxon", n_genes=25, use_raw=False) + +rows = [] +for cls in raw.uns["rank_genes_groups"]["names"].dtype.names: + mask = raw.obs["pred_label"] == cls + if mask.sum() < 30: continue + genes = list(raw.uns["rank_genes_groups"]["names"][cls][:20]) + pvals = [float(x) for x in raw.uns["rank_genes_groups"]["pvals_adj"][cls][:20]] + logfc = [float(x) for x in raw.uns["rank_genes_groups"]["logfoldchanges"][cls][:20]] + + top_str = ",".join([f"{g}(LFC{lf:+.1f})" for g, lf in zip(genes[:10], logfc[:10])]) + panel_hits = {} + for pname, plist in PANELS.items(): + hits = [g for g in plist if g in genes[:20]] + panel_hits[pname] = f"{len(hits)}/{len(plist)}: {','.join(hits)}" + best_panel = max(panel_hits.items(), + key=lambda x: int(x[1].split("/")[0]) / (int(x[1].split(":")[0].split("/")[1]) + 1e-6)) + + # Stage enrichment + stage_col = raw.obs.get("Stage", raw.obs.get("stage", pd.Series([""]*raw.n_obs, index=raw.obs.index))) + stage_vals = pd.to_numeric(stage_col[mask], errors="coerce") + top_stage = int(stage_vals.mode().iloc[0]) if len(stage_vals.dropna()) else -1 + stage6_frac = float((stage_vals == 6).sum() / max(1, mask.sum())) + + rows.append({ + "predicted_class": cls, + "n_cells": int(mask.sum()), + "top_wilcoxon_markers": top_str, + "min_p_adj_top5": min(pvals[:5], default=float("nan")), + "best_canonical_panel_match": best_panel[0], + "recovery": best_panel[1], + "top_stage": top_stage, + "stage6_frac": stage6_frac, + }) + +df = pd.DataFrame(rows).sort_values("n_cells", ascending=False) +df.to_csv(OUT / "91_veres_marker_deep_dive.csv", index=False) +print(f"[write] {OUT}/91_veres_marker_deep_dive.csv ({len(df)} classes)", flush=True) +print() +print(df[["predicted_class", "n_cells", "best_canonical_panel_match", "recovery", + "top_stage", "stage6_frac"]].to_string(index=False)) diff --git a/scripts/analysis/92_dahlin_marker_deep_dive.py b/scripts/analysis/92_dahlin_marker_deep_dive.py new file mode 100644 index 0000000000000000000000000000000000000000..7b0f166162782c02e61ea0e30f828bf3e630a677 --- /dev/null +++ b/scripts/analysis/92_dahlin_marker_deep_dive.py @@ -0,0 +1,150 @@ +"""dahlin marker deep-dive: wilcoxon per predicted class + Kit-W41 vs WT enrichment.""" +from pathlib import Path +import warnings, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp, torch, pickle +warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 +import sys +sys.path.insert(0, "/home/bcheng/PRISM") +from panda import PANDAEncoder + +ROOT = Path("/home/bcheng/PRISM") +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +OUT = ROOT / "discovery/hematopoiesis/marker" +OUT.mkdir(parents=True, exist_ok=True) + +PANELS = { + "LT-HSC": ["Hlf", "Meis1", "Mecom", "Procr", "Fgd5", "Mllt3", "Kit"], + "MPP": ["Cd48", "Flt3", "Cd34", "Sell", "Slamf1"], + "erythroid": ["Klf1", "Car1", "Car2", "Blvrb", "Hba-a1", "Hba-a2", "Kit"], + "megakaryocyte": ["Itga2b", "Pf4", "Gp1bb", "Gata1"], + "myeloid": ["Elane", "Mpo", "Prtn3", "Ctsg", "Cebpe", "Wfdc17", "Mmp8", "Ctss"], + "basophil-mast": ["Cpa3", "Ms4a2", "Gata2", "Mcpt8", "Hdc"], + "lymphoid": ["Il7r", "Rag1", "Dntt", "Vpreb1"], + "Kit-signaling": ["Kit", "Kitl", "Sox4"], + "MYC-targets": ["Myc", "Nolc1", "Nop58"], + "ISR": ["Atf4", "Ddit3", "Ppp1r15a"], + "Apoptosis-pro": ["Bax", "Bak1", "Bid"], +} + + +def load_dahlin(): + from pathlib import Path as _P + D_DIR = _P("/home/bcheng/PRISM/data/corpus/hematopoiesis/held_out_unlabeled/dahlin_extract") + GT = {"SIGAB1":"WT","SIGAC1":"WT","SIGAD1":"WT","SIGAF1":"WT","SIGAG1":"WT", + "SIGAH1":"WT","SIGAG8":"Kit_W41","SIGAH8":"Kit_W41"} + parts = [] + for f in sorted(D_DIR.glob("*.txt.gz")): + sample = f.name.split("_")[1].split(".")[0] + 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"] = GT.get(sample, "unknown") + var = pd.DataFrame(index=df.index.astype(str)) + parts.append(ad.AnnData(X=X, obs=obs, var=var)) + a = ad.concat(parts, join="outer", label="_batch") + import mygene + mg = mygene.MyGeneInfo() + res = mg.querymany(a.var_names.astype(str).tolist(), 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) + a = a[:, keep].copy(); a.var_names = syms[keep]; a.var_names_make_unique() + return a + + +def project_dahlin(a): + ck = torch.load(ROOT / "checkpoints/hematopoiesis/marker/panda_final.pt", + map_location=DEVICE, weights_only=False) + classes = ck["classes"]; marker_genes = ck["marker_genes"] + stats = np.load(ROOT / "data/corpus/hematopoiesis/harmonized/corpus_stats.npz", allow_pickle=True) + pca = pickle.load(open(ROOT / "data/corpus/hematopoiesis/harmonized/pca_basis.pkl", "rb")) + hvgs = [str(g) for g in stats["shared_hvgs"]] + + hvg2i = {g: i for i, g in enumerate(hvgs)} + common = [g for g in a.var_names.astype(str) if g in hvg2i] + 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, len(hvgs)), dtype=np.float32) + Xf[:, np.array([hvg2i[g] for g in common])] = X + Xz = np.clip((Xf - stats["mean"].astype(np.float32)) / stats["std"].astype(np.float32), -10, 10) + Xpca = pca.transform(Xz).astype(np.float32) + + mvals = np.zeros((a.n_obs, len(marker_genes)), dtype=np.float32) + for j, g in enumerate(marker_genes): + if g in a.var_names: + col = a[:, 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) + + model = PANDAEncoder(variant="marker", n_pca=50, n_markers=len(marker_genes), + n_classes=len(classes), n_sub=3, n_datasets=len(ck["datasets"])).to(DEVICE).eval() + model.load_state_dict(ck["model"]) + + preds = [] + with torch.no_grad(): + for i in range(0, a.n_obs, 4096): + xb = torch.from_numpy(Xpca[i:i+4096]).to(DEVICE) + xmb = torch.from_numpy(Xmark[i:i+4096]).to(DEVICE) + aux = torch.zeros(len(xb), 2, device=DEVICE) + out = model(xb, aux, x_markers=xmb, lam_dann=0.0) + mc = model.max_sub_cos(out["z"]) + preds.append(mc.argmax(dim=1).cpu().numpy()) + preds = np.concatenate(preds) + return np.array([classes[i] for i in preds]) + + +print("[load] Dahlin + predict", flush=True) +raw = load_dahlin() +raw.obs["pred_label"] = pd.Categorical(project_dahlin(raw)) +print(f"[align] {raw.n_obs} cells across {raw.obs['pred_label'].nunique()} classes", flush=True) + +counts_s = raw.obs["pred_label"].value_counts() +keep_cls = counts_s[counts_s >= 50].index.tolist() +raw = raw[raw.obs["pred_label"].isin(keep_cls)].copy() +raw.obs["pred_label"] = raw.obs["pred_label"].astype(str).astype("category") +print(f"[filter] {raw.n_obs} cells × {len(keep_cls)} classes", flush=True) + +sc.pp.normalize_total(raw, target_sum=1e4); sc.pp.log1p(raw) +print("[wilcoxon] running...", flush=True) +sc.tl.rank_genes_groups(raw, groupby="pred_label", method="wilcoxon", n_genes=25, use_raw=False) + +rows = [] +for cls in raw.uns["rank_genes_groups"]["names"].dtype.names: + mask = raw.obs["pred_label"] == cls + if mask.sum() < 50: continue + genes = list(raw.uns["rank_genes_groups"]["names"][cls][:20]) + pvals = [float(x) for x in raw.uns["rank_genes_groups"]["pvals_adj"][cls][:20]] + logfc = [float(x) for x in raw.uns["rank_genes_groups"]["logfoldchanges"][cls][:20]] + + top_str = ",".join([f"{g}(LFC{lf:+.1f})" for g, lf in zip(genes[:10], logfc[:10])]) + panel_hits = {} + for pname, plist in PANELS.items(): + hits = [g for g in plist if g in genes[:20]] + panel_hits[pname] = f"{len(hits)}/{len(plist)}: {','.join(hits)}" + best_panel = max(panel_hits.items(), + key=lambda x: int(x[1].split("/")[0]) / (int(x[1].split(":")[0].split("/")[1]) + 1e-6)) + + gt = raw.obs["genotype"][mask].astype(str) + nwt = int((gt == "WT").sum()); nkit = int((gt == "Kit_W41").sum()) + frac_wt = nwt / max(1, nwt + nkit) + + rows.append({ + "predicted_class": cls, + "n_cells": int(mask.sum()), + "top_wilcoxon_markers": top_str, + "min_p_adj_top5": min(pvals[:5], default=float("nan")), + "best_canonical_panel_match": best_panel[0], + "recovery": best_panel[1], + "n_WT": nwt, + "n_Kit_W41": nkit, + "frac_WT": frac_wt, + }) + +df = pd.DataFrame(rows).sort_values("n_cells", ascending=False) +df.to_csv(OUT / "92_dahlin_marker_deep_dive.csv", index=False) +print(f"[write] {OUT}/92_dahlin_marker_deep_dive.csv ({len(df)} classes)", flush=True) +print() +print(df[["predicted_class", "n_cells", "best_canonical_panel_match", "recovery", "frac_WT"]].to_string(index=False)) diff --git a/scripts/analysis/93_true_zero_shot_baron.py b/scripts/analysis/93_true_zero_shot_baron.py new file mode 100644 index 0000000000000000000000000000000000000000..0fb8406fc9ca8a63debf71682d2a9812ac4db193 --- /dev/null +++ b/scripts/analysis/93_true_zero_shot_baron.py @@ -0,0 +1,90 @@ +"""true zero-shot on baron test-half (943 mouse islet cells held out of corpus) under pca+marker variants.""" +from pathlib import Path +import warnings, json, sys, pickle, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp, torch +warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 +sys.path.insert(0, "/home/bcheng/PRISM") +from panda import PANDAEncoder +from sklearn.metrics import accuracy_score, f1_score, classification_report + +ROOT = Path("/home/bcheng/PRISM") +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +BARON = ROOT / "data/corpus/pancreas/held_out_labeled/baron_GSE84133_mouse_test.h5ad" + + +def infer(a, variant): + ck = torch.load(ROOT / f"checkpoints/pancreas/{variant}/panda_final.pt", + map_location=DEVICE, weights_only=False) + classes = ck["classes"]; marker_genes = ck.get("marker_genes", []) + stats = np.load(ROOT / "data/corpus/pancreas/harmonized/corpus_stats.npz", allow_pickle=True) + pca = pickle.load(open(ROOT / "data/corpus/pancreas/harmonized/pca_basis.pkl", "rb")) + hvgs = [str(g) for g in stats["shared_hvgs"]] + hvg2i = {g: i for i, g in enumerate(hvgs)} + common = [g for g in a.var_names.astype(str) if g in hvg2i] + 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, len(hvgs)), dtype=np.float32) + Xf[:, np.array([hvg2i[g] for g in common])] = X + Xz = np.clip((Xf - stats["mean"].astype(np.float32)) / stats["std"].astype(np.float32), -10, 10) + Xpca = pca.transform(Xz).astype(np.float32) + + Xmark = None + if variant == "marker": + mvals = np.zeros((a.n_obs, len(marker_genes)), dtype=np.float32) + for j, g in enumerate(marker_genes): + if g in a.var_names: + col = a[:, 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) + + model = PANDAEncoder(variant=variant, n_pca=50, + n_markers=len(marker_genes) if variant == "marker" else 0, + n_classes=len(classes), n_sub=3, + n_datasets=len(ck["datasets"])).to(DEVICE).eval() + model.load_state_dict(ck["model"]) + + preds, probs = [], [] + with torch.no_grad(): + for i in range(0, a.n_obs, 4096): + xb = torch.from_numpy(Xpca[i:i+4096]).to(DEVICE) + xmb = torch.from_numpy(Xmark[i:i+4096]).to(DEVICE) if Xmark is not None else None + aux = torch.zeros(len(xb), 2, device=DEVICE) + out = model(xb, aux, x_markers=xmb, lam_dann=0.0) + mc = model.max_sub_cos(out["z"]) + preds.append(mc.argmax(dim=1).cpu().numpy()) + probs.append(torch.softmax(mc / 0.07, dim=1).cpu().numpy()) + return np.array([classes[i] for i in np.concatenate(preds)]), np.concatenate(probs), classes + + +def main(): + print(f"[baron] loading {BARON}", flush=True) + a = ad.read_h5ad(BARON) + y_true = a.obs["canonical_label"].astype(str).values + print(f"[baron] {a.shape} true labels: {pd.Series(y_true).value_counts().to_dict()}", flush=True) + + for variant in ("pca", "marker"): + print(f"\n=== {variant.upper()} ===", flush=True) + pred, probs, classes = infer(a, variant) + # eval only on cells whose true label is in our class vocabulary + mask = np.isin(y_true, classes) + acc = accuracy_score(y_true[mask], pred[mask]) + f1 = f1_score(y_true[mask], pred[mask], average="macro", zero_division=0) + rep = classification_report(y_true[mask], pred[mask], zero_division=0, output_dict=True) + print(f"[eval-{variant}] n={mask.sum()} acc={acc:.4f} macro-f1={f1:.4f}", flush=True) + out = ROOT / f"discovery/pancreas/{variant}" + out.mkdir(parents=True, exist_ok=True) + (out / "93_baron_zero_shot.json").write_text(json.dumps({ + "variant": variant, "n_cells": int(mask.sum()), "n_classes_eval": int(len(set(y_true[mask]))), + "acc": float(acc), "macro_f1": float(f1), "per_class": rep, + "predicted_dist": pd.Series(pred).value_counts().to_dict(), + "true_dist": pd.Series(y_true).value_counts().to_dict(), + }, indent=2, default=str)) + pd.DataFrame({"cell_id": a.obs_names, "true_label": y_true, "pred_label": pred, + "max_cos": probs.max(axis=1)}).to_csv( + out / "93_baron_predictions.csv", index=False) + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/94_true_zero_shot_nestorowa.py b/scripts/analysis/94_true_zero_shot_nestorowa.py new file mode 100644 index 0000000000000000000000000000000000000000..70a1e14b23f26ff9d92c5bf2e767030e8c2e0479 --- /dev/null +++ b/scripts/analysis/94_true_zero_shot_nestorowa.py @@ -0,0 +1,107 @@ +"""true zero-shot on Nestorowa GSE81682 (1920 smart-seq2 FACS-labeled cells, held out of corpus) under pca+marker variants.""" +from pathlib import Path +import warnings, json, sys, pickle, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp, torch +warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 +sys.path.insert(0, "/home/bcheng/PRISM") +from panda import PANDAEncoder +from sklearn.metrics import accuracy_score, f1_score, classification_report + +ROOT = Path("/home/bcheng/PRISM") +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +NEST = ROOT / "data/raw/nestorowa_combined.h5ad" + +# panda fine class -> nestorowa coarse FACS gate (LT-HSC vs HSPC) +COARSE = { + "LT-HSC": "LT-HSC", + "MPP": "HSPC", + "GMP": "HSPC", + "myeloid": "HSPC", + "erythroid": "HSPC", + "megakaryocyte":"HSPC", + "basophil-mast":"HSPC", + "lymphoid": "HSPC", + "unassigned": "HSPC", + "UNK": "HSPC", +} + + +def infer(a, variant): + ck = torch.load(ROOT / f"checkpoints/hematopoiesis/{variant}/panda_final.pt", + map_location=DEVICE, weights_only=False) + classes = ck["classes"]; marker_genes = ck.get("marker_genes", []) + stats = np.load(ROOT / "data/corpus/hematopoiesis/harmonized/corpus_stats.npz", allow_pickle=True) + pca = pickle.load(open(ROOT / "data/corpus/hematopoiesis/harmonized/pca_basis.pkl", "rb")) + hvgs = [str(g) for g in stats["shared_hvgs"]] + hvg2i = {g: i for i, g in enumerate(hvgs)} + common = [g for g in a.var_names.astype(str) if g in hvg2i] + 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, len(hvgs)), dtype=np.float32) + Xf[:, np.array([hvg2i[g] for g in common])] = X + Xz = np.clip((Xf - stats["mean"].astype(np.float32)) / stats["std"].astype(np.float32), -10, 10) + Xpca = pca.transform(Xz).astype(np.float32) + + Xmark = None + if variant == "marker": + mvals = np.zeros((a.n_obs, len(marker_genes)), dtype=np.float32) + for j, g in enumerate(marker_genes): + if g in a.var_names: + col = a[:, 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) + + model = PANDAEncoder(variant=variant, n_pca=50, + n_markers=len(marker_genes) if variant == "marker" else 0, + n_classes=len(classes), n_sub=3, + n_datasets=len(ck["datasets"])).to(DEVICE).eval() + model.load_state_dict(ck["model"]) + + preds, probs = [], [] + with torch.no_grad(): + for i in range(0, a.n_obs, 4096): + xb = torch.from_numpy(Xpca[i:i+4096]).to(DEVICE) + xmb = torch.from_numpy(Xmark[i:i+4096]).to(DEVICE) if Xmark is not None else None + aux = torch.zeros(len(xb), 2, device=DEVICE) + out = model(xb, aux, x_markers=xmb, lam_dann=0.0) + mc = model.max_sub_cos(out["z"]) + preds.append(mc.argmax(dim=1).cpu().numpy()) + probs.append(torch.softmax(mc / 0.07, dim=1).cpu().numpy()) + return np.array([classes[i] for i in np.concatenate(preds)]), np.concatenate(probs), classes + + +def main(): + a = ad.read_h5ad(NEST) + print(f"[nest] {a.shape} facs gates: {a.obs['cell_type'].value_counts().to_dict()}", flush=True) + + for variant in ("pca", "marker"): + print(f"\n=== {variant.upper()} ===", flush=True) + pred, probs, classes = infer(a, variant) + pred_coarse = np.array([COARSE.get(p, "HSPC") for p in pred]) + y_true = a.obs["cell_type"].astype(str).values + mask = y_true != "unknown" + acc = accuracy_score(y_true[mask], pred_coarse[mask]) + f1 = f1_score(y_true[mask], pred_coarse[mask], average="macro", zero_division=0) + rep = classification_report(y_true[mask], pred_coarse[mask], zero_division=0, output_dict=True) + print(f"[eval-{variant}] n_labeled={mask.sum()} coarse-acc={acc:.4f} macro-f1={f1:.4f}", flush=True) + fine_by_gate = pd.crosstab(a.obs["cell_type"].astype(str), pd.Series(pred)) + print(fine_by_gate.to_string(), flush=True) + + out = ROOT / f"discovery/hematopoiesis/{variant}" + out.mkdir(parents=True, exist_ok=True) + (out / "94_nestorowa_zero_shot.json").write_text(json.dumps({ + "variant": variant, "n_cells_total": int(a.n_obs), "n_cells_labeled": int(mask.sum()), + "coarse_acc": float(acc), "coarse_f1": float(f1), + "coarse_per_class": rep, + "fine_by_gate": fine_by_gate.to_dict(), + "max_cos_p50": float(np.median(probs.max(axis=1))), + }, indent=2, default=str)) + pd.DataFrame({"cell_id": a.obs_names, "facs_gate": y_true, + "pred_fine": pred, "pred_coarse": pred_coarse, + "max_cos": probs.max(axis=1)}).to_csv(out / "94_nestorowa_predictions.csv", index=False) + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/95_adult_beta_validation.py b/scripts/analysis/95_adult_beta_validation.py new file mode 100644 index 0000000000000000000000000000000000000000..ca751b192614cb042702f8022c8e8c205d2cad1d --- /dev/null +++ b/scripts/analysis/95_adult_beta_validation.py @@ -0,0 +1,92 @@ +"""adult-beta canonical panel enrichment on Veres: mean log1p in adult-beta vs beta vs overall. reports vacuous if n_adult_beta=0.""" +from pathlib import Path +import warnings, json, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp +warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 + +ROOT = Path("/home/bcheng/PRISM") +OUT = ROOT / "discovery/pancreas/marker" +OUT.mkdir(parents=True, exist_ok=True) + +PANEL = ["MAFA", "UCN3", "IAPP", "INS", "SIX3", "MAFB", "MNX1", "NEUROD1"] + + +def load_veres(): + SHARON_DIR = ROOT / "data/corpus/pancreas/held_out_unlabeled/sharon_extract" + parts = [] + for meta_file in sorted(SHARON_DIR.glob("*.cell_metadata.tsv.gz")): + counts_file = str(meta_file).replace("cell_metadata", "processed_counts") + if not Path(counts_file).exists(): continue + meta = pd.read_csv(meta_file, sep="\t", compression="gzip") + counts = pd.read_csv(counts_file, sep="\t", compression="gzip", index_col=0) + obs = meta.set_index("library.barcode") + obs = obs.loc[obs.index.intersection(counts.index)] + counts_al = counts.loc[obs.index] + X = sp.csr_matrix(counts_al.values.astype(np.float32)) + a = ad.AnnData(X=X, obs=obs, var=pd.DataFrame(index=counts_al.columns)) + a.var_names_make_unique() + parts.append(a) + return ad.concat(parts, join="outer") + + +print("[load] Veres + predictions", flush=True) +raw = load_veres() +pred_df = pd.read_csv(ROOT / "discovery/pancreas/marker/veres_predictions.csv") +pred_df["cell_id"] = pred_df["cell_id"].astype(str).str.replace(r"^veres_", "", regex=True) +common = raw.obs_names.intersection(pd.Index(pred_df["cell_id"].astype(str))) +raw = raw[list(common)].copy() +pred_map = dict(zip(pred_df["cell_id"].astype(str), pred_df["pred_label"])) +raw.obs["pred_label"] = pd.Categorical([pred_map.get(c, "unknown") for c in raw.obs_names]) +print(f"[align] {raw.n_obs} cells", flush=True) + +sc.pp.normalize_total(raw, target_sum=1e4); sc.pp.log1p(raw) + +n_adult_beta = int((raw.obs["pred_label"] == "adult-beta").sum()) +n_beta = int((raw.obs["pred_label"] == "beta").sum()) +print(f"[counts] adult-beta={n_adult_beta} beta={n_beta} overall={raw.n_obs}", flush=True) + +def mean_expr(mask, gene): + if gene not in raw.var_names or mask.sum() == 0: + return float("nan") + col = raw[mask, gene].X + if sp.issparse(col): col = col.toarray() + return float(col.mean()) + +mask_ab = (raw.obs["pred_label"] == "adult-beta").values +mask_b = (raw.obs["pred_label"] == "beta").values + +result = { + "cluster": "adult-beta", + "n_adult_beta": n_adult_beta, + "n_beta": n_beta, + "n_overall": int(raw.n_obs), + "vacuous": n_adult_beta == 0, + "marker": {}, +} + +for g in PANEL: + ab = mean_expr(mask_ab, g) + b = mean_expr(mask_b, g) + ov = mean_expr(np.ones(raw.n_obs, dtype=bool), g) + enr = (ab / b) if (b and not np.isnan(b) and b > 0) else float("nan") + result["marker"][g] = { + "adult_beta_mean_log1p": None if np.isnan(ab) else round(ab, 4), + "beta_mean_log1p": None if np.isnan(b) else round(b, 4), + "overall_mean_log1p": None if np.isnan(ov) else round(ov, 4), + "enrichment_adult_beta_vs_beta": None if np.isnan(enr) else round(enr, 3), + } + +if n_adult_beta == 0: + result["interpretation"] = ( + "VACUOUS: current PANDA-Marker (Jul-23 checkpoint) predicts 0 adult-beta cells on Veres. " + "The 'beta' cluster (n={}) captures INS/IAPP/MAFB/ADCYAP1 signal instead; adult-vs-juvenile " + "distinction is not resolved on this dataset. Enrichment ratios below use 0/beta and are NaN." + ).format(n_beta) +else: + result["interpretation"] = ( + "adult-beta cluster (n={}) canonical panel enrichment vs beta cluster (n={})." + ).format(n_adult_beta, n_beta) + +with open(OUT / "95_adult_beta_validation.json", "w") as f: + json.dump(result, f, indent=2) +print(f"[write] {OUT}/95_adult_beta_validation.json", flush=True) +print(json.dumps(result, indent=2)) diff --git a/scripts/analysis/98_eden_posthoc_detection.py b/scripts/analysis/98_eden_posthoc_detection.py new file mode 100644 index 0000000000000000000000000000000000000000..c8b5075d89e59ab453be38c1c5f6685663a6e0b7 --- /dev/null +++ b/scripts/analysis/98_eden_posthoc_detection.py @@ -0,0 +1,149 @@ +"""post-hoc EDEN (S100a4+Tnc+Pdgfra+ derm10 per Dingwall 2024) detection on v3 pan-skin predictions; cKO vs WT dermal-fibro proportions.""" +from pathlib import Path +import warnings, json, sys, pickle, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp, torch, torch.nn.functional as F +warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 +sys.path.insert(0, "/home/bcheng/PRISM") +from panda import PANDAEncoder + +ROOT = Path("/home/bcheng/PRISM") +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +CKO_GSMS = {"GSM6833482", "GSM6833483"} # CORRECTED: 480/481 are rttaControl (WT), not cKO +WT_GSMS = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"} # CORRECTED: 4 Cre-neg controls per GEO metadata + +# Dingwall 2024 EDEN definition: cluster 20 top-2 markers + broad lineage +EDEN_CANONICAL_MARKERS = ["S100a4", "Tnc", "Pdgfra"] + + +def predict(a, variant="marker"): + ck = torch.load(ROOT / f"checkpoints/pan_skin/{variant}/panda_final.pt", + map_location=DEVICE, weights_only=False) + classes = ck["classes"]; marker_genes = ck.get("marker_genes", []) + stats = np.load(ROOT / "data/corpus/pan_skin/harmonized/corpus_stats.npz", allow_pickle=True) + pca = pickle.load(open(ROOT / "data/corpus/pan_skin/harmonized/pca_basis.pkl", "rb")) + hvgs = [str(g) for g in stats["shared_hvgs"]] + hvg2i = {g: i for i, g in enumerate(hvgs)} + common = [g for g in a.var_names.astype(str) if g in hvg2i] + 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, len(hvgs)), dtype=np.float32) + Xf[:, np.array([hvg2i[g] for g in common])] = X + Xz = np.clip((Xf - stats["mean"].astype(np.float32)) / stats["std"].astype(np.float32), -10, 10) + Xpca = pca.transform(Xz).astype(np.float32) + Xmark = None + if variant == "marker": + mv = np.zeros((a.n_obs, len(marker_genes)), dtype=np.float32) + for j, g in enumerate(marker_genes): + if g in a.var_names: + col = a[:, g].X + if sp.issparse(col): col = col.toarray() + mv[:, j] = col.flatten().astype(np.float32) + mmu = mv.mean(axis=0, keepdims=True); msig = mv.std(axis=0, keepdims=True) + 1e-6 + Xmark = np.clip((mv - mmu) / msig, -5, 5).astype(np.float32) + + model = PANDAEncoder(variant=variant, n_pca=50, + n_markers=len(marker_genes) if variant == "marker" else 0, + n_classes=len(classes), n_sub=3, + n_datasets=len(ck["datasets"])).to(DEVICE).eval() + model.load_state_dict(ck["model"]) + preds, probs = [], [] + with torch.no_grad(): + for i in range(0, a.n_obs, 4096): + xb = torch.from_numpy(Xpca[i:i+4096]).to(DEVICE) + xmb = torch.from_numpy(Xmark[i:i+4096]).to(DEVICE) if Xmark is not None else None + aux = torch.zeros(len(xb), 2, device=DEVICE) + out = model(xb, aux, x_markers=xmb, lam_dann=0.0) + mc = model.max_sub_cos(out["z"]) + preds.append(mc.argmax(dim=1).cpu().numpy()) + probs.append(F.softmax(mc / 0.07, dim=1).cpu().numpy()) + return np.array([classes[i] for i in np.concatenate(preds)]), np.concatenate(probs) + + +def score_eden_module(a): + sc.pp.normalize_total(a, target_sum=1e4); sc.pp.log1p(a) + present = [g for g in EDEN_CANONICAL_MARKERS if g in a.var_names] + if not present: + return np.zeros(a.n_obs), present + sc.tl.score_genes(a, gene_list=present, score_name="eden_score", use_raw=False) + return a.obs["eden_score"].values, present + + +def main(): + print("[eden] loading Dingwall raw counts", flush=True) + raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad") + genotype = np.where(raw.obs["sample"].astype(str).isin(list(CKO_GSMS)), "En1-cKO", + np.where(raw.obs["sample"].astype(str).isin(list(WT_GSMS)), "WT", "other")) + print(f"[eden] {raw.n_obs} cells, genotype dist: {pd.Series(genotype).value_counts().to_dict()}", flush=True) + + print("[eden] running v3 PANDA-Marker prediction", flush=True) + pred, probs = predict(raw, variant="marker") + print(f"[eden] pred dist: {pd.Series(pred).value_counts().head().to_dict()}", flush=True) + + dermal_mask = np.isin(pred, ["fibroblast-papillary", "fibroblast-reticular"]) + print(f"[eden] dermal-fibroblast predictions: {dermal_mask.sum()} cells", flush=True) + + print(f"[eden] scoring EDEN module ({EDEN_CANONICAL_MARKERS})", flush=True) + eden_score, present = score_eden_module(raw.copy()) + print(f"[eden] markers present in Dingwall counts: {present}", flush=True) + + dermal_ix = np.where(dermal_mask)[0] + dermal_scores = eden_score[dermal_ix] + thr_p95 = np.percentile(dermal_scores, 95) + thr_p90 = np.percentile(dermal_scores, 90) + eden_core_p95 = dermal_ix[dermal_scores >= thr_p95] + eden_core_p90 = dermal_ix[dermal_scores >= thr_p90] + + print(f"\n[eden] EDEN core (S100a4+Tnc+Pdgfra top-5% among dermal fibroblasts):", flush=True) + print(f" p95 threshold: {thr_p95:.3f} n={len(eden_core_p95)}", flush=True) + print(f" p90 threshold: {thr_p90:.3f} n={len(eden_core_p90)}", flush=True) + + for p, ix, thr in [(95, eden_core_p95, thr_p95), (90, eden_core_p90, thr_p90)]: + gt = genotype[ix] + labeled = gt != "other" + n_wt = int((gt[labeled] == "WT").sum()); n_ko = int((gt[labeled] == "En1-cKO").sum()) + frac_wt = n_wt / max(1, n_wt + n_ko) + # dingwall paper: control=1.99% dermal, cKO=0.08% dermal → ~25x depletion in cluster 20 + print(f" p{p}: WT={n_wt} cKO={n_ko} WT_frac={frac_wt:.3f} " + f"(paper says WT>cKO ~25x depletion in cluster 20)", flush=True) + + out = ROOT / "discovery/pan_skin/marker" + out.mkdir(parents=True, exist_ok=True) + df = pd.DataFrame({ + "cell_id": raw.obs_names, + "sample": raw.obs["sample"].astype(str).values, + "genotype": genotype, + "pred_label": pred, + "max_cos": probs.max(axis=1), + "eden_score": eden_score, + "is_dermal_fibro": dermal_mask, + "is_eden_core_p95": np.isin(np.arange(raw.n_obs), eden_core_p95), + "is_eden_core_p90": np.isin(np.arange(raw.n_obs), eden_core_p90), + }) + df.to_csv(out / "98_eden_dingwall_predictions.csv", index=False) + + summary = { + "system": "pan_skin", + "target": "Dingwall_GSE220977", + "eden_definition": { + "source": "Dingwall 2024 Dev Cell PMC10872420", + "markers_used_by_paper": ["S100a4", "Tnc", "Pdgfra"], + "cluster_id_paper": "cluster 20 / Derm10", + "paper_reported_size": {"WT_dermal_frac": 0.0199, "cKO_dermal_frac": 0.0008, + "WT_n_approx": 516, "cKO_n_approx": 21, + "depletion_ratio": 24.9}, + }, + "our_detection": { + "method": "S100a4+Tnc+Pdgfra module score on v3 dermal-fibroblast predictions, top-5% threshold", + "markers_present_in_dingwall_counts": present, + "n_dermal_fibroblast_cells": int(dermal_mask.sum()), + "n_eden_core_p95": int(len(eden_core_p95)), + "n_eden_core_p90": int(len(eden_core_p90)), + }, + } + (out / "98_eden_summary.json").write_text(json.dumps(summary, indent=2, default=str)) + print(f"\n[write] {out}/98_eden_*.{{csv,json}}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/99_en1_dual_role_analysis.py b/scripts/analysis/99_en1_dual_role_analysis.py new file mode 100644 index 0000000000000000000000000000000000000000..2adfb4b2793a7331990e4471d3ec7ff5e2ef08fa --- /dev/null +++ b/scripts/analysis/99_en1_dual_role_analysis.py @@ -0,0 +1,137 @@ +"""En1 dual-role on Dingwall: test A local activator (En1+ > En1- Sweat_gland in WT); test B spatial repressor (cKO > WT Sweat_gland per class).""" +from pathlib import Path +import warnings, json, sys, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp +from scipy.stats import mannwhitneyu +warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 + +ROOT = Path("/home/bcheng/PRISM") + +# Sweat_gland pathway module (from restored 57_pathway_analysis.py) +SWEAT_GLAND_GENES = ["Foxi3", "Foxa1", "En1", "Krt8", "Krt18", "Krt19", + "Muc5b", "Aqp5", "Cutl1"] +# for TEST A we EXCLUDE En1 from the module to avoid tautology +SWEAT_GLAND_MINUS_EN1 = [g for g in SWEAT_GLAND_GENES if g != "En1"] + +CKO_GSMS = {"GSM6833482", "GSM6833483"} # CORRECTED: 480/481 are rttaControl (WT), not cKO +WT_GSMS = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"} # CORRECTED: 4 Cre-neg controls per GEO metadata + + +def main(): + print("[en1] loading Dingwall raw + v3 marker predictions", flush=True) + raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad") + pred = pd.read_csv(ROOT / "discovery/pan_skin/marker/dingwall_predictions.csv") + pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"])) + labels = np.array([pred_map.get(c, "unknown") for c in raw.obs_names.astype(str)]) + raw.obs["pred_label"] = pd.Categorical(labels) + raw.obs["genotype"] = np.where(raw.obs["sample"].astype(str).isin(list(CKO_GSMS)), "En1-cKO", + np.where(raw.obs["sample"].astype(str).isin(list(WT_GSMS)), "WT", "other")) + labeled = raw.obs["genotype"].isin(["WT", "En1-cKO"]).values + raw = raw[labeled].copy() + print(f"[en1] {raw.n_obs:,} labeled cells (WT + cKO); classes: {raw.obs['pred_label'].nunique()}", flush=True) + + sc.pp.normalize_total(raw, target_sum=1e4); sc.pp.log1p(raw) + + # En1 excluded from module to avoid tautology in Test A + sg_present = [g for g in SWEAT_GLAND_MINUS_EN1 if g in raw.var_names] + print(f"[en1] Sweat_gland (En1-excluded) genes present: {sg_present}", flush=True) + sc.tl.score_genes(raw, gene_list=sg_present, score_name="sg_module_minusEn1", + random_state=0, use_raw=False) + + if "En1" in raw.var_names: + en1_col = raw[:, "En1"].X + en1_exp = en1_col.toarray().flatten() if sp.issparse(en1_col) else en1_col.flatten() + else: + print("[en1] WARNING: En1 not detected in var_names", flush=True) + en1_exp = np.zeros(raw.n_obs) + raw.obs["En1_expr"] = en1_exp + raw.obs["En1_detected"] = en1_exp > 0 # any detection after log-norm + + # ======================================================================= + # TEST A — LOCAL ACTIVATOR (En1+ > En1- within WT) + # ======================================================================= + print("\n=== TEST A: LOCAL ACTIVATOR (En1+ vs En1- Sweat_gland in WT only) ===", + flush=True) + wt_mask = raw.obs["genotype"] == "WT" + wt = raw[wt_mask].copy() + print(f"[testA] {wt.n_obs} WT cells; En1+ = {int(wt.obs['En1_detected'].sum())}, " + f"En1- = {int((~wt.obs['En1_detected']).sum())}", flush=True) + en1pos_score = wt[wt.obs["En1_detected"]].obs["sg_module_minusEn1"].values + en1neg_score = wt[~wt.obs["En1_detected"]].obs["sg_module_minusEn1"].values + if len(en1pos_score) > 5 and len(en1neg_score) > 5: + U_A, p_A = mannwhitneyu(en1pos_score, en1neg_score, alternative="greater") + delta_A = float(en1pos_score.mean() - en1neg_score.mean()) + print(f"[testA] En1+ mean = {en1pos_score.mean():.4f}, En1- mean = {en1neg_score.mean():.4f}", + flush=True) + print(f"[testA] delta = {delta_A:+.4f}, MannU-greater p = {p_A:.3e}", flush=True) + testA = { + "n_en1pos": int(len(en1pos_score)), "n_en1neg": int(len(en1neg_score)), + "mean_en1pos": float(en1pos_score.mean()), "mean_en1neg": float(en1neg_score.mean()), + "delta_activation": delta_A, "mannu_p_greater": float(p_A), + } + else: + print(f"[testA] insufficient cells; skipping", flush=True) + testA = {"skipped": True} + + # ======================================================================= + # TEST B — SPATIAL REPRESSOR (cKO > WT in non-eccrine classes) + # ======================================================================= + print("\n=== TEST B: SPATIAL REPRESSOR (cKO > WT Sweat_gland score per class) ===", + flush=True) + rows = [] + for cls in sorted(raw.obs["pred_label"].astype(str).unique()): + sub = raw[raw.obs["pred_label"].astype(str) == cls] + wt_cells = sub[sub.obs["genotype"] == "WT"] + cko_cells = sub[sub.obs["genotype"] == "En1-cKO"] + if wt_cells.n_obs < 10 or cko_cells.n_obs < 10: + continue + wt_scores = wt_cells.obs["sg_module_minusEn1"].values + cko_scores = cko_cells.obs["sg_module_minusEn1"].values + U, p_greater = mannwhitneyu(cko_scores, wt_scores, alternative="greater") + U, p_two = mannwhitneyu(cko_scores, wt_scores, alternative="two-sided") + delta = float(cko_scores.mean() - wt_scores.mean()) + rows.append({ + "predicted_class": cls, + "n_WT": int(wt_cells.n_obs), + "n_cKO": int(cko_cells.n_obs), + "wt_mean_sg_score": float(wt_scores.mean()), + "cko_mean_sg_score": float(cko_scores.mean()), + "delta_derepression": delta, + "mannu_p_greater_cKO_gt_WT": float(p_greater), + "mannu_p_two_sided": float(p_two), + "bonferroni_p_adj": float(min(1.0, p_greater * 20)), # crude adjustment + "interpretation": "derepressed_in_cKO" if delta > 0 and p_greater < 0.05 + else "normal" if delta > 0 else "downregulated_in_cKO", + }) + df = pd.DataFrame(rows).sort_values("delta_derepression", ascending=False) + print(df[["predicted_class", "n_WT", "n_cKO", "delta_derepression", + "mannu_p_greater_cKO_gt_WT", "interpretation"]].to_string(index=False), + flush=True) + + out = ROOT / "discovery/pan_skin/marker" + out.mkdir(parents=True, exist_ok=True) + df.to_csv(out / "99_en1_dual_role.csv", index=False) + summary = { + "model": "Dingwall 2024 En1 dual role model", + "hypothesis": "En1 is BOTH a local activator (turns Sweat_gland ON in eccrine-competent cells) " + "AND a spatial repressor (prevents Sweat_gland from firing elsewhere)", + "test_A_local_activator": testA, + "test_B_spatial_repressor": { + "n_classes_tested": len(rows), + "n_classes_derepressed_p_lt_0.05": int((df["mannu_p_greater_cKO_gt_WT"] < 0.05).sum()), + "n_classes_derepressed_p_lt_0.001": int((df["mannu_p_greater_cKO_gt_WT"] < 0.001).sum()), + "top_derepressed": df.head(5)[["predicted_class", "delta_derepression", + "mannu_p_greater_cKO_gt_WT"]].to_dict("records"), + }, + "module_definition": { + "name": "Sweat_gland (En1-excluded)", + "genes_used": sg_present, + "note": "En1 removed from the module to avoid tautology in Test A " + "(En1+ cells trivially score higher on a module containing En1)", + }, + } + (out / "99_en1_dual_role_summary.json").write_text(json.dumps(summary, indent=2, default=str)) + print(f"\n[write] {out}/99_en1_dual_role.{{csv,json}}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/README.md b/scripts/analysis/README.md new file mode 100644 index 0000000000000000000000000000000000000000..5be07e20709f74496af15aff115310c2b3df6801 --- /dev/null +++ b/scripts/analysis/README.md @@ -0,0 +1,20 @@ +# scripts/analysis + +downstream mechanistic + interpretability analyses on trained PANDA checkpoints. +runs after `scripts/common/train_panda.py` and `scripts/common/zero_shot.py`. +each script writes CSV/JSON/npy outputs into `discovery/` (or `discovery/{system}/...`) +and is safe to re-run. + +| id | script | what it does | +|---:|---|---| +| 70 | `70_prototype_geometry.py` | intra + cross-system prototype cosine geometry | +| 72 | `72_emergent_axes.py` | within-class PCA of the 128-d projections | +| 73 | `73_novel_populations_dahlin.py` | same abstain-gate flow, on Dahlin | +| 80 | `80_prototype_gene_attribution.py` | integrated-gradient prototype-to-gene attribution | +| 81 | `81_counterfactual_knockouts.py` | per-gene KO delta on prototype cosine | +| 82 | `82_gene_coattribution_modules.py` | gene-gene co-attribution modules | +| 83 | `83_prototype_training_trajectory.py` | prototype drift + eff-dim across the 4-stage curriculum | +| 84 | `84_adversary_purification.py` | tests that dataset + depth adversaries are at chance | +| 85 | `85_hessian_gene_interactions.py` | second-order gene-gene Hessian per prototype | +| 90-92 | `9{0,1,2}_*_marker_deep_dive.py` | per-class Wilcoxon vs canonical panels, per target | +| 93-94 | `9{3,4}_true_zero_shot_*.py` | true zero-shot on fully held-out Baron + Nestorowa | diff --git a/scripts/common/README.md b/scripts/common/README.md new file mode 100644 index 0000000000000000000000000000000000000000..a35ee05b2f00232fba18be010f8f8d9598f76bce --- /dev/null +++ b/scripts/common/README.md @@ -0,0 +1,14 @@ +# scripts/common + +system-agnostic training + eval + inference. everything here takes a system name +(`pan_skin`, `hematopoiesis`, `pancreas`) as an argument and expects the corresponding +`data/corpus/{system}/harmonized/corpus.h5ad`. + +| script | what it does | +|---|---| +| `train_panda.py` | train one PANDA variant (pca / marker) on one system | +| `cv_holdout.py` | 5-fold held-out CV (GroupKFold by dataset when available) | +| `zero_shot.py` | run PANDA checkpoints on Dingwall / Dahlin / Veres targets | +| `nestorowa_zero_shot.py` | labeled zero-shot on Nestorowa GSE81682 (HSC validation) | +| `extract_prototypes_for_analysis.py` | prototype extraction helper for downstream analysis scripts | +| `rerun_all_discovery.sh` | driver that re-runs the whole `scripts/analysis/` block end to end | diff --git a/scripts/common/__init__.py b/scripts/common/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/scripts/common/cv_holdout.py b/scripts/common/cv_holdout.py new file mode 100644 index 0000000000000000000000000000000000000000..8a7100ce679ad83194bd964c4bae314c54a0ae51 --- /dev/null +++ b/scripts/common/cv_holdout.py @@ -0,0 +1,155 @@ +"""held-out CV for panda. GroupKFold by dataset when possible, else StratifiedKFold.""" +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 +from sklearn.model_selection import GroupKFold, StratifiedKFold +from sklearn.metrics import accuracy_score, f1_score, roc_auc_score, classification_report +warnings.filterwarnings("ignore") + +sys.path.insert(0, "/home/bcheng/PRISM") +from panda import ( + PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, + subcenter_angular_infonce, prototype_repulsion, +) +from scripts.common.train_panda import prepare_batches, load_corpus, get_marker_gene_list + +ROOT = Path("/home/bcheng/PRISM") +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def train_one_fold(Xpca, Xmark, y, y_dset, log10cz, classes, variant, + train_idx, epochs=6, batch=256, lr=1e-3, seed=0): + n_classes = len(classes) + n_datasets = int(max(y_dset[train_idx].max() + 1, 1)) + n_markers = Xmark.shape[1] if Xmark is not None else 0 + + torch.manual_seed(seed); np.random.seed(seed) + model = PANDAEncoder( + variant=variant, n_pca=50, n_markers=n_markers, + n_classes=n_classes, n_sub=3, n_datasets=n_datasets, dropout=0.2, + ).to(DEVICE) + opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4) + + train_n = len(train_idx) + rng = np.random.default_rng(seed) + Xt = Xpca[train_idx]; yt = y[train_idx]; ydt = y_dset[train_idx]; dt = log10cz[train_idx] + Xmt = Xmark[train_idx] if Xmark is not None else None + + 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(train_n) + for bstart in range(0, train_n, batch): + idx = perm[bstart:bstart+batch] + x = torch.from_numpy(Xt[idx]).to(DEVICE) + xm = torch.from_numpy(Xmt[idx]).to(DEVICE) if Xmt is not None else None + yy = torch.from_numpy(yt[idx]).to(DEVICE) + yd = torch.from_numpy(ydt[idx]).to(DEVICE) + dd = torch.from_numpy(dt[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) + if stage >= 3: + L = L + 0.5 * prototype_repulsion(model.prototypes.detach().clone()) + opt.zero_grad(); L.backward(); opt.step() + if stage >= 1: + with torch.no_grad(): model.update_prototypes(z.detach(), yy) + return model + + +def evaluate(model, Xpca, Xmark, y, val_idx, classes): + model.eval() + preds, probs = [], [] + Xv = Xpca[val_idx]; Xmv = Xmark[val_idx] if Xmark is not None else None + with torch.no_grad(): + for i in range(0, len(val_idx), 2048): + xb = torch.from_numpy(Xv[i:i+2048]).to(DEVICE) + xmb = torch.from_numpy(Xmv[i:i+2048]).to(DEVICE) if Xmv is not None else None + aux = torch.zeros(len(xb), 2, device=DEVICE) + out = model(xb, aux, x_markers=xmb, lam_dann=0.0) + z = out["z"] + mc = model.max_sub_cos(z) # (B, K) + preds.append(mc.argmax(dim=1).cpu().numpy()) + probs.append(F.softmax(mc / 0.07, dim=1).cpu().numpy()) + preds = np.concatenate(preds); probs = np.concatenate(probs) + yv = y[val_idx] + acc = accuracy_score(yv, preds) + f1 = f1_score(yv, preds, average="macro", zero_division=0) + # macro AUROC — nan if val has <2 classes + n_classes = len(classes) + try: + y_onehot = np.eye(n_classes)[yv] + auc = roc_auc_score(y_onehot, probs, average="macro", multi_class="ovr") + except Exception: + auc = float("nan") + rep = classification_report(yv, preds, labels=list(range(n_classes)), + target_names=classes, output_dict=True, zero_division=0) + return acc, f1, auc, rep + + +def cv(system, variant, folds=5, epochs=6, split_mode="auto", seed=0): + a, hvgs, mu, sig, pca = load_corpus(system) + marker_genes = get_marker_gene_list(system) if variant == "marker" else [] + Xpca, Xmark, y, classes, y_dset, dset_classes, log10cz = prepare_batches( + a, hvgs, mu, sig, pca, marker_genes, variant + ) + print(f"[cv] {system}/{variant} n={a.n_obs} K={len(classes)} datasets={len(dset_classes)}", flush=True) + + use_group = (split_mode == "group") or (split_mode == "auto" and len(dset_classes) >= folds) + if use_group: + splitter = GroupKFold(n_splits=folds) + splits = list(splitter.split(np.zeros(len(y)), y, y_dset)) + print(f"[cv] GroupKFold by dataset ({len(dset_classes)} groups)", flush=True) + else: + splitter = StratifiedKFold(n_splits=folds, shuffle=True, random_state=seed) + splits = list(splitter.split(np.zeros(len(y)), y)) + print(f"[cv] StratifiedKFold on labels ({split_mode}) seed={seed}", flush=True) + + per_fold_acc, per_fold_f1, per_fold_auc = [], [], [] + last_report = None + for fold, (tr, va) in enumerate(splits): + print(f"[fold {fold+1}/{folds}] train={len(tr)} val={len(va)}", flush=True) + model = train_one_fold(Xpca, Xmark, y, y_dset, log10cz, classes, variant, tr, + epochs=epochs, seed=seed * 100 + fold) + acc, f1, auc, rep = evaluate(model, Xpca, Xmark, y, va, classes) + per_fold_acc.append(acc); per_fold_f1.append(f1); per_fold_auc.append(auc) + last_report = rep + print(f"[fold {fold+1}] acc={acc:.4f} f1={f1:.4f} auc={auc:.4f}", flush=True) + + result = { + "system": system, "variant": variant, "folds": folds, "epochs": epochs, "seed": seed, + "n_cells": int(a.n_obs), "n_classes": len(classes), + "per_fold_acc": per_fold_acc, "per_fold_f1": per_fold_f1, "per_fold_auc": per_fold_auc, + "mean_acc": float(np.mean(per_fold_acc)), "std_acc": float(np.std(per_fold_acc)), + "mean_f1": float(np.mean(per_fold_f1)), "std_f1": float(np.std(per_fold_f1)), + "mean_auc": float(np.nanmean(per_fold_auc)), + "std_auc": float(np.nanstd(per_fold_auc)), + "per_class_report": last_report, + } + out_dir = ROOT / f"discovery/{system}/{variant}" + out_dir.mkdir(parents=True, exist_ok=True) + suffix = f"_seed{seed}" if seed != 0 else "" + (out_dir / f"cv_{folds}fold{suffix}.json").write_text(json.dumps(result, indent=2, default=str)) + print(f"\n[cv] mean acc={result['mean_acc']:.4f}±{result['std_acc']:.4f} " + f"f1={result['mean_f1']:.4f} auc={result['mean_auc']:.4f}", 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("--folds", type=int, default=5) + ap.add_argument("--epochs", type=int, default=6) + ap.add_argument("--split", choices=["auto", "group", "stratified"], default="auto") + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + cv(args.system, args.variant, args.folds, args.epochs, args.split, args.seed) diff --git a/scripts/common/extract_prototypes_for_analysis.py b/scripts/common/extract_prototypes_for_analysis.py new file mode 100644 index 0000000000000000000000000000000000000000..5c733f99f446d332427dfef15750e0f4bbe6927a --- /dev/null +++ b/scripts/common/extract_prototypes_for_analysis.py @@ -0,0 +1,31 @@ +"""pull prototypes.npy + label_encoding.json out of panda_final.pt for legacy analysis scripts.""" +from pathlib import Path +import json, numpy as np, torch + +ROOT = Path("/home/bcheng/PRISM") +for sys in ["pan_skin", "hematopoiesis", "pancreas"]: + for variant in ["pca", "marker"]: + ckpt_dir = ROOT / f"checkpoints/{sys}/{variant}" + ck = torch.load(ckpt_dir / "panda_final.pt", map_location="cpu", weights_only=False) + # v2 stores prototypes as (K, K_sub, D); legacy scripts want (K, D) + protos = ck["prototypes"] + if protos.ndim == 3: + protos_mean = protos.mean(axis=1) + protos_mean = protos_mean / (np.linalg.norm(protos_mean, axis=1, keepdims=True) + 1e-8) + else: + protos_mean = protos + np.save(ckpt_dir / "prototypes.npy", protos_mean.astype(np.float32)) + json.dump({"classes": ck["classes"], "datasets": ck["datasets"]}, + open(ckpt_dir / "label_encoding.json", "w"), indent=2) + print(f"[{sys}/{variant}] K={len(ck['classes'])} sub_dim={protos.shape}") + +# canonical symlinks at checkpoints/{system}/ (marker is canonical) +for sys in ["pan_skin", "hematopoiesis", "pancreas"]: + src_dir = ROOT / f"checkpoints/{sys}/marker" + dst_dir = ROOT / f"checkpoints/{sys}" + for name in ["prototypes.npy", "label_encoding.json", "panda_final.pt"]: + src = src_dir / name; dst = dst_dir / name + if dst.exists() or dst.is_symlink(): dst.unlink() + dst.symlink_to(src.resolve()) + print(f"linked {sys}/ -> marker/") +print("DONE") diff --git a/scripts/common/generate_missing_holdouts.py b/scripts/common/generate_missing_holdouts.py new file mode 100644 index 0000000000000000000000000000000000000000..6902a5dce7dc8ded705d8ce962c6d4f72c2391db --- /dev/null +++ b/scripts/common/generate_missing_holdouts.py @@ -0,0 +1,65 @@ +"""materialise held-out h5ad slices for nestorowa (hsc) and sulic (skin) using anchor split rules.""" +from pathlib import Path +import warnings, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp +warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 + +ROOT = Path("/home/bcheng/PRISM") + + +def make_nestorowa_test(): + src = ROOT / "data/raw/nestorowa_combined.h5ad" + print(f"[nestorowa] loading {src}", flush=True) + a = ad.read_h5ad(src) + # anchor split from scripts/hematopoiesis/09_retrain_with_nestorowa_anchor.py: + # 150 LT-HSC + 600 HSPC stratified (seed=0); rest is test. + rng = np.random.default_rng(0) + lt = np.where(a.obs["cell_type"].astype(str).values == "LT-HSC")[0] + hs = np.where(a.obs["cell_type"].astype(str).values == "HSPC")[0] + lt_anchor = rng.choice(lt, size=min(150, len(lt)), replace=False) + hs_anchor = rng.choice(hs, size=min(600, len(hs)), replace=False) + anchor_ix = np.concatenate([lt_anchor, hs_anchor]) + test_ix = np.setdiff1d(np.arange(a.n_obs), anchor_ix) + a_test = a[test_ix].copy() + # coerce obs cols to string so h5ad write survives mixed dtypes + for c in list(a_test.obs.columns): + try: a_test.obs[c] = a_test.obs[c].astype(str) + except Exception: del a_test.obs[c] + out_dir = ROOT / "data/corpus/hematopoiesis/held_out_labeled" + out_dir.mkdir(parents=True, exist_ok=True) + out = out_dir / "nestorowa_GSE81682_test.h5ad" + a_test.write_h5ad(out) + print(f"[nestorowa] wrote {out} ({a_test.n_obs:,} cells; " + f"gate dist: {a_test.obs['cell_type'].value_counts().to_dict()})", flush=True) + + +def make_sulic_test(): + corpus = ROOT / "data/corpus/pan_skin/harmonized/corpus.h5ad" + print(f"[sulic] loading {corpus} (legacy, has sulic + labels)", flush=True) + a = ad.read_h5ad(corpus) + a_sul = a[a.obs["dataset"] == "sulic_GSE212673"].copy() + print(f"[sulic] {a_sul.n_obs} sulic cells " + f"({a_sul.obs['canonical_label'].value_counts().to_dict()})", flush=True) + # anchor split from scripts/pan_skin/92_retrain_with_sulic_anchor.py: + # 300 HF-placode + 200 basal-IFE stratified (seed=0); rest (4183) is test. + rng = np.random.default_rng(0) + hf = np.where(a_sul.obs["canonical_label"].astype(str).values == "HF-placode")[0] + bi = np.where(a_sul.obs["canonical_label"].astype(str).values == "basal-IFE")[0] + hf_anchor = rng.choice(hf, size=min(300, len(hf)), replace=False) + bi_anchor = rng.choice(bi, size=min(200, len(bi)), replace=False) + anchor_ix = np.concatenate([hf_anchor, bi_anchor]) + test_ix = np.setdiff1d(np.arange(a_sul.n_obs), anchor_ix) + a_test = a_sul[test_ix].copy() + for c in list(a_test.obs.columns): + try: a_test.obs[c] = a_test.obs[c].astype(str) + except Exception: del a_test.obs[c] + out_dir = ROOT / "data/corpus/pan_skin/held_out_labeled" + out_dir.mkdir(parents=True, exist_ok=True) + out = out_dir / "sulic_GSE212673_test.h5ad" + a_test.write_h5ad(out) + print(f"[sulic] wrote {out} ({a_test.n_obs:,} cells; " + f"labels: {a_test.obs['canonical_label'].value_counts().to_dict()})", flush=True) + + +if __name__ == "__main__": + make_nestorowa_test() + make_sulic_test() diff --git a/scripts/common/nestorowa_zero_shot.py b/scripts/common/nestorowa_zero_shot.py new file mode 100644 index 0000000000000000000000000000000000000000..e77ba7966d16bdc518a78b0ffb461144cc616525 --- /dev/null +++ b/scripts/common/nestorowa_zero_shot.py @@ -0,0 +1,168 @@ +"""labeled zero-shot on nestorowa GSE81682 (hsc validation target).""" +from __future__ import annotations +from pathlib import Path +import sys, warnings, json, argparse, pickle, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp, torch +from sklearn.metrics import accuracy_score, f1_score, classification_report +warnings.filterwarnings("ignore") + +sys.path.insert(0, "/home/bcheng/PRISM") +from panda import PANDAEncoder + +ROOT = Path("/home/bcheng/PRISM") +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def load_nestorowa(): + """load nestorowa GSE81682 htseq counts, map ENSMUSG -> symbol.""" + counts_path = ROOT / "data/raw/GSE81682_HTSeq_counts.txt.gz" + if not counts_path.exists(): + raise FileNotFoundError(counts_path) + df = pd.read_csv(counts_path, sep="\t", index_col=0) + print(f"[nestorowa] raw counts shape: {df.shape}", flush=True) + # rows=genes, cols=cells; transpose + if df.shape[0] > df.shape[1]: + df = df.T + a = ad.AnnData(X=sp.csr_matrix(df.values.astype(np.float32)), + obs=pd.DataFrame(index=df.index.astype(str)), + var=pd.DataFrame(index=df.columns.astype(str))) + if any(g.startswith("ENSMUSG") for g in a.var_names[:100]): + import mygene + mg = mygene.MyGeneInfo() + res = mg.querymany(a.var_names.astype(str).tolist(), 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) + a = a[:, keep].copy(); a.var_names = syms[keep]; a.var_names_make_unique() + print(f"[nestorowa] {a.shape} after gene symbol conversion", flush=True) + return a + + +def score_hsc_labels(a): + """assign hsc labels by marker scoring, proxy for population_annotation.""" + programs = { + "LT-HSC": ["Hlf", "Meis1", "Mecom", "Procr", "Fgd5", "Mllt3"], + "MPP": ["Cd48", "Flt3", "Cd34"], + "LMPP": ["Flt3", "Irf8", "Satb1"], + "CMP": ["Cd34", "Mpo", "Gata2"], + "MEP": ["Gata1", "Klf1", "Itga2b"], + "GMP": ["Elane", "Mpo", "Prtn3", "Ctsg", "Cebpe"], + "erythroblast": ["Klf1", "Car1", "Car2", "Blvrb", "Hba-a1"], + "megakaryocyte":["Itga2b", "Pf4", "Gp1bb"], + "basophil-mast":["Cpa3", "Ms4a2", "Gata2"], + "CLP": ["Il7r", "Rag1", "Dntt"], + } + sc.pp.normalize_total(a, target_sum=1e4); sc.pp.log1p(a) + score_cols = [] + for cls, gs in programs.items(): + present = [g for g in gs if g in a.var_names] + if present: + sc.tl.score_genes(a, gene_list=present, score_name=f"s_{cls}", use_raw=False) + else: + a.obs[f"s_{cls}"] = 0.0 + score_cols.append(f"s_{cls}") + scores = a.obs[score_cols].values + argmax = np.argmax(scores, axis=1) + labels = [c.replace("s_", "") for c in score_cols] + a.obs["approx_label"] = np.array(labels)[argmax] + a.obs["approx_conf"] = scores.max(axis=1) + return a + + +def infer(a, variant): + ckpt = torch.load(ROOT / f"checkpoints/hematopoiesis/{variant}/panda_final.pt", + map_location=DEVICE, weights_only=False) + classes = ckpt["classes"] + marker_genes = ckpt.get("marker_genes", []) + + stats = np.load(ROOT / "data/corpus/hematopoiesis/harmonized/corpus_stats.npz", allow_pickle=True) + pca = pickle.load(open(ROOT / "data/corpus/hematopoiesis/harmonized/pca_basis.pkl", "rb")) + hvgs = [str(g) for g in stats["shared_hvgs"]] + + hvg2i = {g: i for i, g in enumerate(hvgs)} + common = [g for g in a.var_names.astype(str) if g in hvg2i] + a_c = a[:, common].copy() + # score_hsc_labels already log-normed; re-check in case it was skipped + if a_c.X.max() > 20: + 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, len(hvgs)), dtype=np.float32) + Xf[:, np.array([hvg2i[g] for g in common])] = X + Xz = np.clip((Xf - stats["mean"].astype(np.float32)) / stats["std"].astype(np.float32), -10, 10) + Xpca = pca.transform(Xz).astype(np.float32) + + Xmark = None + if variant == "marker": + mvals = np.zeros((a.n_obs, len(marker_genes)), dtype=np.float32) + for j, g in enumerate(marker_genes): + if g in a.var_names: + col = a[:, 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) + + model = PANDAEncoder(variant=variant, n_pca=50, + n_markers=len(marker_genes) if variant == "marker" else 0, + n_classes=len(classes), n_sub=3, n_datasets=len(ckpt["datasets"])).to(DEVICE).eval() + model.load_state_dict(ckpt["model"]) + preds, max_cos_list = [], [] + with torch.no_grad(): + for i in range(0, a.n_obs, 4096): + xb = torch.from_numpy(Xpca[i:i+4096]).to(DEVICE) + xmb = torch.from_numpy(Xmark[i:i+4096]).to(DEVICE) if Xmark is not None else None + aux = torch.zeros(len(xb), 2, device=DEVICE) + out = model(xb, aux, x_markers=xmb, lam_dann=0.0) + mc = model.max_sub_cos(out["z"]) + preds.append(mc.argmax(dim=1).cpu().numpy()) + max_cos_list.append(mc.max(dim=1).values.cpu().numpy()) + return np.array([classes[i] for i in np.concatenate(preds)]), np.concatenate(max_cos_list), classes + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--variant", choices=["pca", "marker"], required=True) + args = ap.parse_args() + print(f"=== Nestorowa HSC zero-shot ({args.variant}) ===", flush=True) + a = load_nestorowa() + a = score_hsc_labels(a) + pred, max_cos, classes = infer(a, args.variant) + a.obs["pred_label"] = pred + a.obs["max_cos"] = max_cos + + y_true = a.obs["approx_label"].values + y_pred = pred + # only eval on cells with approx-label confidence > 0.05 + conf_mask = a.obs["approx_conf"] > 0.05 + print(f"[eval] eval on {conf_mask.sum()}/{a.n_obs} cells with approx-label conf>0.05", flush=True) + if conf_mask.sum() > 20: + common_lbl = sorted(set(y_true[conf_mask]) & set(y_pred[conf_mask])) + mask2 = conf_mask & np.isin(y_true, common_lbl) & np.isin(y_pred, common_lbl) + acc = accuracy_score(y_true[mask2], y_pred[mask2]) + f1 = f1_score(y_true[mask2], y_pred[mask2], average="macro", zero_division=0) + rep = classification_report(y_true[mask2], y_pred[mask2], zero_division=0, output_dict=True) + else: + acc = f1 = float("nan"); rep = {} + result = { + "variant": args.variant, "n_cells": int(a.n_obs), "n_classes": len(classes), + "predicted_dist": pd.Series(pred).value_counts().to_dict(), + "approx_label_dist": pd.Series(y_true).value_counts().to_dict(), + "eval_acc_vs_approx": float(acc), "eval_f1_vs_approx": float(f1), + "max_cos_median": float(np.median(max_cos)), + "n_low_conf_abstain": int((max_cos < 0.5).sum()), + "per_class_report": rep, + } + out_dir = ROOT / f"discovery/hematopoiesis/{args.variant}" + out_dir.mkdir(parents=True, exist_ok=True) + (out_dir / "nestorowa_summary.json").write_text(json.dumps(result, indent=2, default=str)) + pd.DataFrame({ + "cell_id": a.obs_names, + "pred_label": pred, "approx_label": y_true, "approx_conf": a.obs["approx_conf"].values, + "max_cos": max_cos, + }).to_csv(out_dir / "nestorowa_predictions.csv", index=False) + print(f"[write] {out_dir}/nestorowa_*", flush=True) + print(f"acc_vs_approx={acc:.4f} f1={f1:.4f}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/common/phase_e_driver.sh b/scripts/common/phase_e_driver.sh new file mode 100644 index 0000000000000000000000000000000000000000..a6f416b9257158165ab585084b1de536a1974617 --- /dev/null +++ b/scripts/common/phase_e_driver.sh @@ -0,0 +1,35 @@ +#!/bin/bash +# Phase E: after all 6 v3 retrains complete, run the full downstream chain. +set -e + +cd /home/bcheng/PRISM +export LD_LIBRARY_PATH=$(python -c "import site,os; print(os.path.join(site.getsitepackages()[0],'nvidia','cusparselt','lib'))"):/home/bcheng/.conda/pkgs/libstdcxx-15.2.0-h39759b7_7/lib/:${LD_LIBRARY_PATH:-} + +echo "=== phase E step 1: zero-shot predictions on Dingwall/Dahlin/Veres with v3 models ===" +# skin -> Dingwall +python scripts/pan_skin/30_zero_shot_aldrich.py --system_v v3 2>&1 | tail -5 || echo "[warn] aldrich script may need path fix" + +echo "=== phase E step 2: post-hoc EDEN detection using Dingwall paper markers ===" +python scripts/analysis/98_eden_posthoc_detection.py 2>&1 | tail -10 + +echo "=== phase E step 2b: En1 spatial-repressor + local-activator dual-role analysis ===" +python scripts/analysis/99_en1_dual_role_analysis.py 2>&1 | tail -20 + +echo "=== phase E step 3: expanded pathway analysis (91 modules, 3 systems) ===" +for sys in pan_skin hematopoiesis pancreas; do + echo "--- $sys ---" + python scripts/analysis/57_pathway_analysis.py --system $sys 2>&1 | tail -5 +done + +echo "=== phase E step 4: regenerate UMAPs with distinct-hue palette + v3 predictions ===" +python scripts/figures/build_pca_vs_marker_umaps.py 2>&1 | tail -20 + +echo "=== phase E step 5: regenerate supplement PDF ===" +python scripts/figures/build_figure_supplement.py 2>&1 | tail -5 + +echo "=== phase E step 6: rebuild paper PDF ===" +TEXMFHOME=/home/bcheng/.texlive/texmf-dist PATH=/home/bcheng/.local/texlive/2023/bin/x86_64-linux:$PATH pdflatex -interaction=nonstopmode PAPER.tex > /tmp/latex.log 2>&1 +TEXMFHOME=/home/bcheng/.texlive/texmf-dist PATH=/home/bcheng/.local/texlive/2023/bin/x86_64-linux:$PATH pdflatex -interaction=nonstopmode PAPER.tex > /tmp/latex.log 2>&1 +rm -f PAPER.aux PAPER.log PAPER.out + +echo "=== phase E DONE ===" diff --git a/scripts/common/rerun_all_discovery.sh b/scripts/common/rerun_all_discovery.sh new file mode 100644 index 0000000000000000000000000000000000000000..dfe3bd44d26f768dba1944d081338cd91d6038a8 --- /dev/null +++ b/scripts/common/rerun_all_discovery.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +# Phase C: rerun all mechanistic discovery on new checkpoints. +# Runs analysis scripts 70-85 for each system with both variants where applicable. +set -euo pipefail +cd /home/bcheng/PRISM + +LDLIBS=$(python -c "import site,os; print(os.path.join(site.getsitepackages()[0],'nvidia','cusparselt','lib'))" 2>/dev/null || echo "") + +echo "=== Phase C: mechanistic discovery reruns ===" + +# All rely on the panda import + checkpoints/ + data/corpus/*/harmonized/ +run() { + local name="$1"; shift + LD_LIBRARY_PATH="$LDLIBS" nohup python -u "$@" > "logs/${name}.log" 2>&1 & + disown + echo "launched $name" +} + +# 70 — prototype geometry (fast, no GPU) +run rerun_70_proto python scripts/analysis/70_prototype_geometry.py + +# 80 — prototype-gene attribution (both systems, both variants) +CUDA_VISIBLE_DEVICES=2 run rerun_80_attr_marker python scripts/analysis/80_prototype_gene_attribution.py + +# 82 — co-attribution modules +run rerun_82_coatt python scripts/analysis/82_gene_coattribution_modules.py + +# 83 — training trajectory +run rerun_83_traj python scripts/analysis/83_prototype_training_trajectory.py + +echo "kick off complete — check logs/rerun_*.log" diff --git a/scripts/common/run_all_retrains.sh b/scripts/common/run_all_retrains.sh new file mode 100644 index 0000000000000000000000000000000000000000..7dcd8d7f837c1d477eb0b842487bf1353728cf59 --- /dev/null +++ b/scripts/common/run_all_retrains.sh @@ -0,0 +1,26 @@ +#!/bin/bash +# retrain all 6 panda models (3 systems x 2 variants), sequentially to fit GPU +set -e + +cd /home/bcheng/PRISM +export LD_LIBRARY_PATH=$(python -c "import site,os; print(os.path.join(site.getsitepackages()[0],'nvidia','cusparselt','lib'))"):/home/bcheng/.conda/pkgs/libstdcxx-15.2.0-h39759b7_7/lib/:${LD_LIBRARY_PATH:-} + +echo "=== SKIN PCA ===" +python -m scripts.common.train_panda pan_skin --variant pca --epochs 8 + +echo "=== SKIN MARKER ===" +python -m scripts.common.train_panda pan_skin --variant marker --epochs 8 + +echo "=== HSC PCA ===" +python -m scripts.common.train_panda hematopoiesis --variant pca --epochs 8 + +echo "=== HSC MARKER ===" +python -m scripts.common.train_panda hematopoiesis --variant marker --epochs 8 + +echo "=== PANCREAS PCA ===" +python -m scripts.common.train_panda pancreas --variant pca --epochs 8 + +echo "=== PANCREAS MARKER ===" +python -m scripts.common.train_panda pancreas --variant marker --epochs 8 + +echo "=== ALL DONE ===" diff --git a/scripts/common/run_all_zero_shot.py b/scripts/common/run_all_zero_shot.py new file mode 100644 index 0000000000000000000000000000000000000000..357220d675c0999ce611823e01e8b00169128219 --- /dev/null +++ b/scripts/common/run_all_zero_shot.py @@ -0,0 +1,152 @@ +"""zero-shot inference over every held-out target x (pca, marker) checkpoint.""" +from pathlib import Path +import warnings, json, sys, pickle, argparse, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp, torch, torch.nn.functional as F +warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 +sys.path.insert(0, "/home/bcheng/PRISM") +from panda import PANDAEncoder +from sklearn.metrics import accuracy_score, f1_score, classification_report + +ROOT = Path("/home/bcheng/PRISM") +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def infer(a, system, variant): + ck = torch.load(ROOT / f"checkpoints/{system}/{variant}/panda_final.pt", + map_location=DEVICE, weights_only=False) + classes = ck["classes"]; marker_genes = ck.get("marker_genes", []) + 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")) + hvgs = [str(g) for g in stats["shared_hvgs"]] + hvg2i = {g: i for i, g in enumerate(hvgs)} + common = [g for g in a.var_names.astype(str) if g in hvg2i] + 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, len(hvgs)), dtype=np.float32) + Xf[:, np.array([hvg2i[g] for g in common])] = X + Xz = np.clip((Xf - stats["mean"].astype(np.float32)) / stats["std"].astype(np.float32), -10, 10) + Xpca = pca.transform(Xz).astype(np.float32) + Xmark = None + if variant == "marker": + mv = np.zeros((a.n_obs, len(marker_genes)), dtype=np.float32) + for j, g in enumerate(marker_genes): + if g in a.var_names: + col = a[:, g].X + if sp.issparse(col): col = col.toarray() + mv[:, j] = col.flatten().astype(np.float32) + mmu = mv.mean(axis=0, keepdims=True); msig = mv.std(axis=0, keepdims=True) + 1e-6 + Xmark = np.clip((mv - mmu) / msig, -5, 5).astype(np.float32) + model = PANDAEncoder(variant=variant, n_pca=50, + n_markers=len(marker_genes) if variant == "marker" else 0, + n_classes=len(classes), n_sub=3, + n_datasets=len(ck["datasets"])).to(DEVICE).eval() + model.load_state_dict(ck["model"]) + preds, probs = [], [] + with torch.no_grad(): + for i in range(0, a.n_obs, 4096): + xb = torch.from_numpy(Xpca[i:i+4096]).to(DEVICE) + xmb = torch.from_numpy(Xmark[i:i+4096]).to(DEVICE) if Xmark is not None else None + aux = torch.zeros(len(xb), 2, device=DEVICE) + out = model(xb, aux, x_markers=xmb, lam_dann=0.0) + mc = model.max_sub_cos(out["z"]) + preds.append(mc.argmax(dim=1).cpu().numpy()) + probs.append(F.softmax(mc / 0.07, dim=1).cpu().numpy()) + return np.array([classes[i] for i in np.concatenate(preds)]), np.concatenate(probs), classes + + +TARGETS = { + "pan_skin": [ + ("dingwall", ROOT / "data/raw/GSE220977_combined.h5ad", None), + ("sulic", ROOT / "data/corpus/pan_skin/held_out_labeled/sulic_GSE212673_test.h5ad", "canonical_label"), + ("belote", ROOT / "data/corpus/pan_skin/held_out_labeled/belote_GSE151091_test.h5ad", "canonical_label"), + ], + "hematopoiesis": [ + ("nestorowa", ROOT / "data/corpus/hematopoiesis/held_out_labeled/nestorowa_GSE81682_test.h5ad", "cell_type"), + ("dahlin", None, None), # loaded per-file via loader (61k cells across 8 samples) + ], + "pancreas": [ + ("baron", ROOT / "data/corpus/pancreas/held_out_labeled/baron_GSE84133_mouse_test.h5ad", "canonical_label"), + ("veres", ROOT / "data/corpus/pancreas/held_out_labeled/veres_GSE114412_test.h5ad", "canonical_label"), + ], +} + + +def load_dahlin(): + """dahlin 61k held-out unlabeled hsc target, 8 sample files.""" + D = ROOT / "data/corpus/hematopoiesis/held_out_unlabeled/dahlin_extract" + GT = {"SIGAB1":"WT","SIGAC1":"WT","SIGAD1":"WT","SIGAF1":"WT","SIGAG1":"WT", + "SIGAH1":"WT","SIGAG8":"Kit_W41","SIGAH8":"Kit_W41"} + parts = [] + for f in sorted(D.glob("*.txt.gz")): + sample = f.name.split("_")[1].split(".")[0] + 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"] = GT.get(sample, "unknown") + var = pd.DataFrame(index=df.index.astype(str)) + parts.append(ad.AnnData(X=X, obs=obs, var=var)) + a = ad.concat(parts, join="outer") + import mygene + mg = mygene.MyGeneInfo() + res = mg.querymany(a.var_names.astype(str).tolist(), 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) + a = a[:, keep].copy(); a.var_names = syms[keep]; a.var_names_make_unique() + return a + + +def process(system, variant): + print(f"\n===== {system} / {variant} =====", flush=True) + for tgt_name, tgt_path, tgt_label in TARGETS[system]: + print(f"\n[{tgt_name}] loading", flush=True) + if tgt_name == "dahlin": + a = load_dahlin() + else: + a = ad.read_h5ad(tgt_path) + print(f"[{tgt_name}] {a.shape}", flush=True) + pred, probs, classes = infer(a, system, variant) + out_dir = ROOT / f"discovery/{system}/{variant}" + out_dir.mkdir(parents=True, exist_ok=True) + pd.DataFrame({ + "cell_id": a.obs_names, + "pred_label": pred, + "max_cos": probs.max(axis=1), + }).to_csv(out_dir / f"{tgt_name}_predictions.csv", index=False) + summary = { + "system": system, "variant": variant, "target": tgt_name, + "n_cells": int(a.n_obs), "n_classes_model": len(classes), + "predicted_class_dist": pd.Series(pred).value_counts().head(30).to_dict(), + "max_cos_p50": float(np.median(probs.max(axis=1))), + "max_cos_p05": float(np.quantile(probs.max(axis=1), 0.05)), + } + if tgt_label and tgt_label in a.obs.columns: + y_true = a.obs[tgt_label].astype(str).values + mask = np.isin(y_true, classes) + if mask.sum() > 0: + acc = accuracy_score(y_true[mask], pred[mask]) + f1 = f1_score(y_true[mask], pred[mask], average="macro", zero_division=0) + rep = classification_report(y_true[mask], pred[mask], + zero_division=0, output_dict=True) + summary["labeled_eval"] = { + "n_eval": int(mask.sum()), "acc": float(acc), + "macro_f1": float(f1), "per_class_report": rep, + } + print(f"[{tgt_name}] acc={acc:.4f} F1={f1:.4f} on {mask.sum()} labeled cells", flush=True) + (out_dir / f"{tgt_name}_summary.json").write_text(json.dumps(summary, indent=2, default=str)) + print(f"[{tgt_name}] wrote {out_dir}/{tgt_name}_predictions.csv + summary.json", flush=True) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--systems", nargs="*", default=["pan_skin", "hematopoiesis", "pancreas"]) + ap.add_argument("--variants", nargs="*", default=["pca", "marker"]) + args = ap.parse_args() + for sys_ in args.systems: + for var in args.variants: + process(sys_, var) + + +if __name__ == "__main__": + main() diff --git a/scripts/common/run_cv.py b/scripts/common/run_cv.py new file mode 100644 index 0000000000000000000000000000000000000000..c5724a82c4883ae1585b02cf376753e39c4c5281 --- /dev/null +++ b/scripts/common/run_cv.py @@ -0,0 +1,159 @@ +"""5-fold stratified CV, 3 systems x 2 variants. 5 epochs/fold (shorter than train_panda).""" +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 +from sklearn.model_selection import StratifiedKFold +from sklearn.metrics import accuracy_score, f1_score, roc_auc_score, classification_report + +sys.path.insert(0, "/home/bcheng/PRISM") +from panda import (PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, + subcenter_angular_infonce, prototype_repulsion) + +ROOT = Path("/home/bcheng/PRISM") +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def load_and_prepare(system, variant): + """canonical corpus + features (PCA, optional marker channel).""" + a = ad.read_h5ad(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")) + hvgs = [str(g) for g in stats["shared_hvgs"]] + + hvg2i = {g: i for i, g in enumerate(hvgs)} + common = [g for g in a.var_names.astype(str) if g in hvg2i] + 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, len(hvgs)), dtype=np.float32) + Xf[:, np.array([hvg2i[g] for g in common])] = X + Xz = np.clip((Xf - stats["mean"].astype(np.float32)) / stats["std"].astype(np.float32), -10, 10) + Xpca = pca.transform(Xz).astype(np.float32) + + Xmark = None; marker_genes = [] + if variant == "marker": + marker_genes = yaml.safe_load(open(ROOT / "panda/markers.yaml"))[system] + mv = np.zeros((a.n_obs, len(marker_genes)), dtype=np.float32) + for j, g in enumerate(marker_genes): + if g in a.var_names: + col = a[:, g].X + if sp.issparse(col): col = col.toarray() + mv[:, j] = col.flatten().astype(np.float32) + mmu = mv.mean(axis=0, keepdims=True); msig = mv.std(axis=0, keepdims=True) + 1e-6 + Xmark = np.clip((mv - mmu) / msig, -5, 5).astype(np.float32) + + labels = a.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(a.obs["dataset"].astype(str).values)) + y_dset = np.array([datasets.index(d) for d in a.obs["dataset"].astype(str).values], dtype=np.int64) + return Xpca, Xmark, y, classes, y_dset, datasets, marker_genes + + +def train_fold(Xpca, Xmark, y, y_dset, classes, variant, tr_ix, epochs=5, batch=256, lr=1e-3, seed=0): + n_classes = len(classes) + n_datasets = int(max(y_dset[tr_ix].max() + 1, 1)) + n_markers = Xmark.shape[1] if Xmark is not None else 0 + torch.manual_seed(seed); np.random.seed(seed) + model = PANDAEncoder(variant=variant, n_pca=50, n_markers=n_markers, + n_classes=n_classes, n_sub=3, n_datasets=n_datasets, dropout=0.2).to(DEVICE) + opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4) + Xt = Xpca[tr_ix]; yt = y[tr_ix]; ydt = y_dset[tr_ix] + Xmt = Xmark[tr_ix] if Xmark is not None else None + rng = np.random.default_rng(seed) + n = len(tr_ix) + for epoch in range(epochs): + stage = 0 if epoch < 1 else 1 if epoch < 3 else 2 + perm = rng.permutation(n) + for bstart in range(0, n, batch): + idx = perm[bstart:bstart+batch] + x = torch.from_numpy(Xt[idx]).to(DEVICE) + xm = torch.from_numpy(Xmt[idx]).to(DEVICE) if Xmt is not None else None + yy = torch.from_numpy(yt[idx]).to(DEVICE) + yd = torch.from_numpy(ydt[idx]).to(DEVICE) + 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) + opt.zero_grad(); L.backward(); opt.step() + if stage >= 1: + with torch.no_grad(): model.update_prototypes(z.detach(), yy) + return model + + +def evaluate(model, Xpca, Xmark, y, val_ix, classes): + model.eval() + preds, probs = [], [] + Xv = Xpca[val_ix]; Xmv = Xmark[val_ix] if Xmark is not None else None + with torch.no_grad(): + for i in range(0, len(val_ix), 2048): + xb = torch.from_numpy(Xv[i:i+2048]).to(DEVICE) + xmb = torch.from_numpy(Xmv[i:i+2048]).to(DEVICE) if Xmv is not None else None + aux = torch.zeros(len(xb), 2, device=DEVICE) + out = model(xb, aux, x_markers=xmb, lam_dann=0.0) + z = out["z"] + mc = model.max_sub_cos(z) + preds.append(mc.argmax(dim=1).cpu().numpy()) + probs.append(F.softmax(mc / 0.07, dim=1).cpu().numpy()) + preds = np.concatenate(preds); probs = np.concatenate(probs) + yv = y[val_ix] + acc = accuracy_score(yv, preds) + f1 = f1_score(yv, preds, average="macro", zero_division=0) + try: + auc = roc_auc_score(np.eye(len(classes))[yv], probs, average="macro", multi_class="ovr") + except Exception: + auc = float("nan") + rep = classification_report(yv, preds, labels=list(range(len(classes))), + target_names=classes, output_dict=True, zero_division=0) + return acc, f1, auc, rep + + +def cv(system, variant, folds=5, epochs=5, seed=0): + print(f"\n=== CV {system}/{variant} ({folds}-fold, {epochs} epochs) ===", flush=True) + Xpca, Xmark, y, classes, y_dset, datasets, _ = load_and_prepare(system, variant) + print(f"[cv] n={len(y):,} K={len(classes)}", flush=True) + skf = StratifiedKFold(n_splits=folds, shuffle=True, random_state=seed) + accs, f1s, aucs = [], [], [] + last_rep = None + for fold, (tr, va) in enumerate(skf.split(np.zeros(len(y)), y)): + model = train_fold(Xpca, Xmark, y, y_dset, classes, variant, tr, + epochs=epochs, seed=seed * 100 + fold) + acc, f1, auc, rep = evaluate(model, Xpca, Xmark, y, va, classes) + accs.append(acc); f1s.append(f1); aucs.append(auc) + last_rep = rep + print(f"[fold {fold+1}] acc={acc:.4f} F1={f1:.4f} AUC={auc:.4f}", flush=True) + + result = { + "system": system, "variant": variant, "folds": folds, "epochs": epochs, + "n_cells": int(len(y)), "n_classes": len(classes), + "per_fold_acc": accs, "per_fold_f1": f1s, "per_fold_auc": aucs, + "mean_acc": float(np.mean(accs)), "std_acc": float(np.std(accs)), + "mean_f1": float(np.mean(f1s)), "std_f1": float(np.std(f1s)), + "mean_auc": float(np.nanmean(aucs)), "std_auc": float(np.nanstd(aucs)), + "per_class_report": last_rep, + } + out_dir = ROOT / f"discovery/{system}/{variant}" + out_dir.mkdir(parents=True, exist_ok=True) + (out_dir / "cv_5fold.json").write_text(json.dumps(result, indent=2, default=str)) + print(f"[cv] mean acc={result['mean_acc']:.4f}±{result['std_acc']:.4f} " + f"F1={result['mean_f1']:.4f} AUC={result['mean_auc']:.4f}", flush=True) + + +if __name__ == "__main__": + ap = argparse.ArgumentParser() + ap.add_argument("--systems", nargs="*", default=["pan_skin", "hematopoiesis", "pancreas"]) + ap.add_argument("--variants", nargs="*", default=["pca", "marker"]) + ap.add_argument("--folds", type=int, default=5) + ap.add_argument("--epochs", type=int, default=5) + args = ap.parse_args() + for s in args.systems: + for v in args.variants: + cv(s, v, args.folds, args.epochs) diff --git a/scripts/common/train_panda.py b/scripts/common/train_panda.py new file mode 100644 index 0000000000000000000000000000000000000000..9a8756f69e984b7fd802cb80120b6552a211a883 --- /dev/null +++ b/scripts/common/train_panda.py @@ -0,0 +1,133 @@ +"""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 + +sys.path.insert(0, "/home/bcheng/PRISM") +from panda import ( + PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, + subcenter_angular_infonce, prototype_repulsion, +) + +ROOT = Path("/home/bcheng/PRISM") +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"): + 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() + 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) + + 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) + if stage >= 3: + L = L + 0.5 * prototype_repulsion(model.prototypes.detach().clone()) + 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) + + ck_dir = ROOT / f"checkpoints/{system}_v3/{variant}" + ck_dir.mkdir(parents=True, exist_ok=True) + torch.save({"model": model.state_dict(), "classes": classes, "datasets": datasets, + "marker_genes": marker_genes if variant == "marker" else [], + "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) diff --git a/scripts/common/zero_shot.py b/scripts/common/zero_shot.py new file mode 100644 index 0000000000000000000000000000000000000000..2188b57bad5306b334ff4b29ad375b0444db15ec --- /dev/null +++ b/scripts/common/zero_shot.py @@ -0,0 +1,143 @@ +"""zero-shot inference on held-out discovery targets (dingwall / dahlin / veres).""" +from __future__ import annotations +from pathlib import Path +import sys, warnings, pickle, json, argparse, numpy as np, pandas as pd, anndata as ad, scanpy as sc +import scipy.sparse as sp, torch, yaml +warnings.filterwarnings("ignore") + +sys.path.insert(0, "/home/bcheng/PRISM") +from panda import PANDAEncoder + +ROOT = Path("/home/bcheng/PRISM") +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def prep_input(adata, system, variant, hvgs, mu, sig, pca, marker_genes): + """log-normalise, PCA-50, optional marker channel.""" + 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() + 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": + 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) + return Xpca, Xmark + + +def infer(system, variant, target_anndata, target_name): + """load checkpoint, run inference, return per-cell (pred, max_cos) + summary.""" + ckpt = torch.load(ROOT / f"checkpoints/{system}/{variant}/panda_final.pt", + map_location=DEVICE, weights_only=False) + classes = ckpt["classes"] + marker_genes = ckpt.get("marker_genes", []) + + 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")) + hvgs = [str(g) for g in stats["shared_hvgs"]] + + # case-fold human symbols → mouse-style when hvgs are mouse (e.g. veres cross-species) + a = target_anndata.copy() + n_upper = sum(1 for g in a.var_names[:1000].astype(str) if g.isupper()) + if n_upper > 500: + new = [g[0].upper() + g[1:].lower() if len(g) > 1 else g for g in a.var_names.astype(str)] + a.var_names = new; a.var_names_make_unique() + + Xpca, Xmark = prep_input(a, system, variant, hvgs, stats["mean"], stats["std"], pca, marker_genes) + + model = PANDAEncoder( + variant=variant, n_pca=50, n_markers=len(marker_genes) if variant == "marker" else 0, + n_classes=len(classes), n_sub=3, n_datasets=len(ckpt["datasets"]), + ).to(DEVICE).eval() + model.load_state_dict(ckpt["model"]) + protos = model.prototypes # (K, n_sub, D) + + preds, max_cos_list = [], [] + with torch.no_grad(): + for i in range(0, a.n_obs, 4096): + xb = torch.from_numpy(Xpca[i:i+4096]).to(DEVICE) + xmb = torch.from_numpy(Xmark[i:i+4096]).to(DEVICE) if Xmark is not None else None + aux = torch.zeros(len(xb), 2, device=DEVICE) + out = model(xb, aux, x_markers=xmb, lam_dann=0.0) + z = out["z"] + mc = model.max_sub_cos(z) # (B, K) + preds.append(mc.argmax(dim=1).cpu().numpy()) + max_cos_list.append(mc.max(dim=1).values.cpu().numpy()) + preds = np.concatenate(preds); max_cos = np.concatenate(max_cos_list) + pred_labels = np.array([classes[i] for i in preds]) + + out_dir = ROOT / f"discovery/{system}/{variant}" + out_dir.mkdir(parents=True, exist_ok=True) + df = pd.DataFrame({ + "cell_id": a.obs_names, + "pred_label": pred_labels, + "max_cos": max_cos, + }) + df.to_csv(out_dir / f"{target_name}_predictions.csv", index=False) + dist = pd.Series(pred_labels).value_counts() + summary = { + "system": system, "variant": variant, "target": target_name, + "n_cells": int(a.n_obs), + "n_classes": len(classes), + "predicted_class_dist": dist.to_dict(), + "max_cos_p50": float(np.median(max_cos)), + "max_cos_p05": float(np.quantile(max_cos, 0.05)), + "abstain_frac_cos_lt_0.5": float((max_cos < 0.5).mean()), + } + (out_dir / f"{target_name}_summary.json").write_text(json.dumps(summary, indent=2, default=str)) + print(f"[{system}/{variant}/{target_name}] {a.n_obs} cells, top preds: {dist.head(5).to_dict()}", flush=True) + return summary + + +def load_target(name): + if name == "dingwall": + return ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad") + if name == "veres": + SHARON_DIR = ROOT / "data/corpus/pancreas/held_out_unlabeled/sharon_extract" + parts = [] + for meta_file in sorted(SHARON_DIR.glob("*.cell_metadata.tsv.gz")): + counts_file = str(meta_file).replace("cell_metadata", "processed_counts") + if not Path(counts_file).exists(): continue + meta = pd.read_csv(meta_file, sep="\t", compression="gzip") + counts = pd.read_csv(counts_file, sep="\t", compression="gzip", index_col=0) + obs = meta.set_index("library.barcode") + obs = obs.loc[obs.index.intersection(counts.index)] + counts_al = counts.loc[obs.index] + X = sp.csr_matrix(counts_al.values.astype(np.float32)) + a = ad.AnnData(X=X, obs=obs, + var=pd.DataFrame(index=counts_al.columns)) + a.var_names_make_unique() + parts.append(a) + return ad.concat(parts, join="outer") + if name == "dahlin": + # skipped here — needs mygene ENSMUSG→symbol conversion (see run_all_zero_shot) + return None + raise ValueError(name) + + +def 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("--target", choices=["dingwall", "dahlin", "veres"], required=True) + args = ap.parse_args() + a = load_target(args.target) + if a is None: + print(f"[!] target {args.target} loader deferred", flush=True); return + infer(args.system, args.variant, a, args.target) + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/README.md b/scripts/figures/README.md new file mode 100644 index 0000000000000000000000000000000000000000..d4e5e0018027b5f2d106d90b9d61354e8e157b2b --- /dev/null +++ b/scripts/figures/README.md @@ -0,0 +1,14 @@ +# scripts/figures + +figure builders for `PAPER.tex` + `figures/PANDA_supplement.pdf`. all figures load +from real artefacts under `discovery/`, `checkpoints/`, and the corpus files -- no +hardcoded numbers. + +| script | output | +|---|---| +| `generate_paper_figures.py` | main-text figures 1-4 (confusion matrices, volcano, heatmap, stage stack) | +| `generate_umap.py` | Dingwall 2-panel UMAP (predicted class + genotype) | +| `generate_multi_umap.py` | 3-panel cross-system UMAP (Dingwall + Dahlin + Veres) | +| `build_skin_corpus_umap.py` | full pan-skin training-corpus UMAP | +| `build_figure_supplement.py` | master builder for the multi-page supplement | +| `build_umaps_and_discovery.py`| per-target UMAPs + mechanistic-evidence panels | diff --git a/scripts/figures/biology_00_umap_cache.py b/scripts/figures/biology_00_umap_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..cec33bc45b1e1f65f4ce6d3436edc6edb8de26a4 --- /dev/null +++ b/scripts/figures/biology_00_umap_cache.py @@ -0,0 +1,47 @@ +"""cache full + dermal-only umap for dingwall replica.""" +from __future__ import annotations +import warnings +warnings.filterwarnings("ignore") +from pathlib import Path +import numpy as np +import scanpy as sc +import anndata as ad + +ROOT = Path("/home/bcheng/PRISM") +OUT = ROOT / "figures/biology" +OUT.mkdir(parents=True, exist_ok=True) +CACHE = OUT / "_cache_dingwall_umap.npz" + + +def main(): + a = ad.read_h5ad(ROOT / "data/processed/dingwall_replica/dingwall_replica.h5ad") + print(f"[cache] loaded dingwall_replica: {a.shape}") + + sc.pp.neighbors(a, use_rep="X_pca_harmony", n_neighbors=30, random_state=42) + sc.tl.umap(a, random_state=42, min_dist=0.3) + umap_full = a.obsm["X_umap"].copy() + + # recompute on dermal-only so Derm0..Derm11 separate + is_derm = a.obs["is_dermal_paper"].astype(bool).values + adx = a[is_derm].copy() + print(f"[cache] dermal subset: {adx.shape}") + sc.pp.neighbors(adx, use_rep="X_pca_harmony", n_neighbors=30, random_state=42) + sc.tl.umap(adx, random_state=42, min_dist=0.3) + umap_derm = adx.obsm["X_umap"].copy() + derm_index = np.where(is_derm)[0] + + np.savez( + CACHE, + umap_full=umap_full, + umap_derm=umap_derm, + derm_index=derm_index, + obs_names=a.obs_names.astype(str).values, + sample=a.obs["sample"].astype(str).values, + derm_label=a.obs["derm_label"].astype(str).values, + paper_cluster_23=a.obs["paper_cluster_23"].astype(str).values, + ) + print(f"[cache] wrote {CACHE}") + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/biology_01_dingwall_umap.py b/scripts/figures/biology_01_dingwall_umap.py new file mode 100644 index 0000000000000000000000000000000000000000..eea5ceab7bd9718975e997d77e384cdf734f9bb0 --- /dev/null +++ b/scripts/figures/biology_01_dingwall_umap.py @@ -0,0 +1,113 @@ +"""dingwall umap by genotype and by panda-marker class. + +WT = {GSM6833478/79/80/81}, cKO = {GSM6833482/83}. +""" +from __future__ import annotations +import warnings, sys +warnings.filterwarnings("ignore") +from pathlib import Path +import numpy as np +import pandas as pd +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import matplotlib.patches as mpatches + +sys.path.insert(0, str(Path(__file__).parent)) +from palette import apply_style, color_for, GENOTYPE_COLORS +apply_style() + +ROOT = Path("/home/bcheng/PRISM") +CACHE = ROOT / "figures/biology/_cache_dingwall_umap.npz" +PREDS = ROOT / "discovery/pan_skin/marker/98_eden_dingwall_predictions.csv" +OUT = ROOT / "figures/biology/biology_01_dingwall_umap.pdf" + +WT_SAMPLES = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"} +CKO_SAMPLES = {"GSM6833482", "GSM6833483"} + + +def main(): + cache = np.load(CACHE, allow_pickle=True) + umap = cache["umap_full"] + obs_names = cache["obs_names"].astype(str) + sample = cache["sample"].astype(str) + + key = np.char.add(np.char.add(obs_names, "__"), sample) + + genotype = np.where(pd.Series(sample).isin(list(CKO_SAMPLES)), "En1-cKO", + np.where(pd.Series(sample).isin(list(WT_SAMPLES)), "WT", "other")) + print("[fig1] genotype counts:", pd.Series(genotype).value_counts().to_dict()) + + p = pd.read_csv(PREDS) + p["key"] = p["cell_id"].astype(str) + "__" + p["sample"].astype(str) + p_lookup = dict(zip(p["key"], p["pred_label"])) + pred = np.array([p_lookup.get(k, "other") for k in key]) + n_matched = (pred != "other").sum() + print(f"[fig1] matched {n_matched}/{len(pred)} predictions") + + fig, axes = plt.subplots(1, 2, figsize=(17.5, 8.0)) + + ax = axes[0] + palette_g = {"WT": GENOTYPE_COLORS["WT"], "En1-cKO": GENOTYPE_COLORS["En1-cKO"], + "other": GENOTYPE_COLORS["other"]} + order = ["other", "WT", "En1-cKO"] + for cat in order: + m = genotype == cat + if m.sum() == 0: + continue + ax.scatter(umap[m, 0], umap[m, 1], s=3.0, alpha=0.55, + color=palette_g[cat], linewidths=0, label=cat, rasterized=True) + ax.set_title(f"Dingwall (n={len(umap):,}) — genotype") + ax.set_xlabel("UMAP-1"); ax.set_ylabel("UMAP-2") + ax.set_xticks([]); ax.set_yticks([]) + handles = [mpatches.Patch(color=palette_g[c], label=f"{c} ({int((genotype==c).sum())})") + for c in ["WT", "En1-cKO"]] + ax.legend(handles=handles, loc="center left", bbox_to_anchor=(1.02, 0.5), + frameon=False, title="Genotype") + ax.text(-0.08, 1.05, "(a)", transform=ax.transAxes, + fontsize=20, fontweight="bold", va="bottom", ha="right") + + ax = axes[1] + # smallest classes on top + counts = pd.Series(pred).value_counts() + order = counts.sort_values(ascending=False).index.tolist() + for cat in reversed(order): + m = pred == cat + color = color_for(cat) if cat != "other" else "#e5e5e5" + s = 3.0 if cat != "other" else 2.0 + ax.scatter(umap[m, 0], umap[m, 1], s=s, alpha=0.6, + color=color, linewidths=0, rasterized=True) + ax.set_title(f"Dingwall — PANDA-Marker predicted class (13 classes)") + ax.set_xlabel("UMAP-1"); ax.set_ylabel("UMAP-2") + ax.set_xticks([]); ax.set_yticks([]) + show_cats = [c for c in order if c != "other"] + handles = [mpatches.Patch(color=color_for(c), + label=f"{c} ({int(counts[c])})") + for c in show_cats] + # legend entry explaining the asterisk overlay drawn below + handles.append(mpatches.Patch(color="none", ec="none", + label=" * = cKO-enriched class (centroid)")) + ax.legend(handles=handles, loc="center left", bbox_to_anchor=(1.02, 0.5), + frameon=False, ncol=1, title="Predicted class") + ax.text(-0.08, 1.05, "(b)", transform=ax.transAxes, + fontsize=20, fontweight="bold", va="bottom", ha="right") + + # star the cKO-enriched populations (see legend entry above) + for cls_name in ["HF-placode", "fibroblast-reticular", "melanoblast"]: + m = pred == cls_name + if m.sum() < 30: + continue + cx, cy = umap[m, 0].mean(), umap[m, 1].mean() + ax.text(cx, cy, "*", fontsize=20, color="#111", + ha="center", va="center", fontweight="bold") + + plt.suptitle("Dingwall En1-cKO — genotype and PANDA-Marker class map", + fontsize=20, y=1.02, fontweight="bold") + plt.tight_layout() + plt.savefig(OUT, bbox_inches="tight", dpi=180) + plt.close() + print(f"[fig1] wrote {OUT}") + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/biology_02_primary_eden.py b/scripts/figures/biology_02_primary_eden.py new file mode 100644 index 0000000000000000000000000000000000000000..5215f6830c573bf9e5619fe6584f4eb7b25cdd97 --- /dev/null +++ b/scripts/figures/biology_02_primary_eden.py @@ -0,0 +1,149 @@ +"""dermal umap + derm2/derm10 genotype zoom with fisher-exact.""" +from __future__ import annotations +import warnings +warnings.filterwarnings("ignore") +from pathlib import Path +import numpy as np, pandas as pd +import matplotlib; matplotlib.use("Agg") +import matplotlib.pyplot as plt +import matplotlib.patches as mpatches +from scipy.stats import fisher_exact +import sys as _sys +_sys.path.insert(0, str(Path("/home/bcheng/PRISM/scripts/figures"))) +from palette import apply_style, GENOTYPE_COLORS +apply_style() + +# derm-cluster distinct palette (kelly-inspired), Derm2 highlighted red +_DERM_HUES = [ + "#e6194b", "#3cb44b", "#4363d8", "#f58231", "#911eb4", "#42d4f4", + "#f032e6", "#bfef45", "#fabed4", "#469990", "#dcbeff", "#9a6324", +] + +ROOT = Path("/home/bcheng/PRISM") +CACHE = ROOT / "figures/biology/_cache_dingwall_umap.npz" +OUT = ROOT / "figures/biology/biology_02_primary_eden.pdf" + +WT_SAMPLES = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"} +CKO_SAMPLES = {"GSM6833482", "GSM6833483"} + +DERM_COLORS = {f"Derm{i}": _DERM_HUES[i % len(_DERM_HUES)] for i in range(12)} +DERM_COLORS["Derm2"] = GENOTYPE_COLORS["En1-cKO"] # highlight (En1-cKO red) +DERM_COLORS["non_dermal"] = "#e5e5e5" + + +def main(): + cache = np.load(CACHE, allow_pickle=True) + umap_full = cache["umap_full"] + umap_derm = cache["umap_derm"] + derm_index = cache["derm_index"] + sample_full = cache["sample"].astype(str) + derm_label_full = cache["derm_label"].astype(str) + + derm_label = derm_label_full[derm_index] + sample = sample_full[derm_index] + genotype = np.where(pd.Series(sample).isin(list(CKO_SAMPLES)), "En1-cKO", + np.where(pd.Series(sample).isin(list(WT_SAMPLES)), "WT", "other")) + + fig, axes = plt.subplots(1, 2, figsize=(18, 8.2)) + + ax = axes[0] + order = [f"Derm{i}" for i in range(12)] + counts = pd.Series(derm_label).value_counts() + # muted background first, Derm2 last so it draws on top + non_highlight = [c for c in order if c != "Derm2"] + for cat in non_highlight: + m = derm_label == cat + if m.sum() == 0: + continue + ax.scatter(umap_derm[m, 0], umap_derm[m, 1], s=2.6, alpha=0.30, + color=DERM_COLORS.get(cat, "#888"), linewidths=0, rasterized=True) + m = derm_label == "Derm2" + if m.sum() > 0: + ax.scatter(umap_derm[m, 0], umap_derm[m, 1], s=5.0, alpha=0.85, + color=DERM_COLORS["Derm2"], linewidths=0, rasterized=True) + ax.set_title(f"Dingwall dermal cells (n={len(umap_derm):,}) — Derm2 highlighted") + ax.set_xlabel("UMAP-1"); ax.set_ylabel("UMAP-2") + ax.set_xticks([]); ax.set_yticks([]) + handles = [mpatches.Patch(color=DERM_COLORS.get(c, "#888"), + label=f"{c} ({int(counts.get(c, 0))})") + for c in order if counts.get(c, 0) > 0] + missing = [c for c in order if counts.get(c, 0) == 0] + if missing: + handles.append(mpatches.Patch(color="none", ec="none", + label=f"(not present: {', '.join(missing)})")) + ax.legend(handles=handles, loc="center left", bbox_to_anchor=(1.02, 0.5), + frameon=False, ncol=1, title="Dermal cluster") + ax.text(-0.08, 1.05, "(a)", transform=ax.transAxes, + fontsize=20, fontweight="bold", va="bottom", ha="right") + + ax = axes[1] + highlight = np.isin(derm_label, ["Derm2", "Derm10"]) + ax.scatter(umap_derm[~highlight, 0], umap_derm[~highlight, 1], + s=1.5, alpha=0.15, color="#dddddd", linewidths=0, rasterized=True) + palette_g = {"WT": GENOTYPE_COLORS["WT"], "En1-cKO": GENOTYPE_COLORS["En1-cKO"], + "other": GENOTYPE_COLORS["other"]} + for gt in ["WT", "En1-cKO"]: + m = highlight & (genotype == gt) + ax.scatter(umap_derm[m, 0], umap_derm[m, 1], s=8, alpha=0.7, + color=palette_g[gt], linewidths=0, label=gt, rasterized=True) + + lines = [] + # place stat boxes at axis corners with arrows to the cluster centroid so + # Derm2 and Derm10 annotations cannot overlap each other. + x_lo, x_hi = umap_derm[:, 0].min(), umap_derm[:, 0].max() + y_lo, y_hi = umap_derm[:, 1].min(), umap_derm[:, 1].max() + dx, dy = x_hi - x_lo, y_hi - y_lo + box_xy = { + "Derm2": (x_lo + 0.02 * dx, y_hi - 0.02 * dy), # top-left + "Derm10": (x_hi - 0.02 * dx, y_hi - 0.02 * dy), # top-right + } + box_va = {"Derm2": "top", "Derm10": "top"} + box_ha = {"Derm2": "left", "Derm10": "right"} + for cluster in ["Derm2", "Derm10"]: + in_cluster = derm_label == cluster + labeled = np.isin(genotype, ["WT", "En1-cKO"]) + tab = pd.crosstab(pd.Series(in_cluster[labeled], name="in"), + pd.Series(genotype[labeled], name="gt")) + try: + a11 = int(tab.loc[True, "En1-cKO"]) + a12 = int(tab.loc[True, "WT"]) + a21 = int(tab.loc[False, "En1-cKO"]) + a22 = int(tab.loc[False, "WT"]) + oratio, pval = fisher_exact([[a11, a12], [a21, a22]]) + cx = umap_derm[in_cluster, 0].mean() + cy = umap_derm[in_cluster, 1].mean() + bx, by = box_xy[cluster] + ax.annotate( + f"{cluster}\nOR={oratio:.2f}\np={pval:.2e}", + xy=(cx, cy), xytext=(bx, by), + fontsize=9.5, ha=box_ha[cluster], va=box_va[cluster], + bbox=dict(boxstyle="round,pad=0.3", fc="white", + ec="#333", alpha=0.9), + arrowprops=dict(arrowstyle="->", color="#333", lw=0.9, + shrinkA=2, shrinkB=6), + ) + lines.append(f"{cluster}: OR={oratio:.2f}, p={pval:.2e} " + f"(cKO {a11}, WT {a12})") + except KeyError as e: + print("[fig2] fisher: missing category", e) + + ax.set_title("Zoom: Derm2 + Derm10 by genotype (Fisher-exact vs rest)") + ax.set_xlabel("UMAP-1"); ax.set_ylabel("UMAP-2") + ax.set_xticks([]); ax.set_yticks([]) + ax.legend(loc="center left", bbox_to_anchor=(1.02, 0.5), + frameon=False, title="Genotype") + ax.text(-0.08, 1.05, "(b)", transform=ax.transAxes, + fontsize=20, fontweight="bold", va="bottom", ha="right") + + plt.suptitle("Primary EDEN discovery — Derm2 and Derm10 are En1-cKO enriched", + fontsize=20, y=1.02, fontweight="bold") + plt.tight_layout() + plt.savefig(OUT, bbox_inches="tight", dpi=180) + plt.close() + print(f"[fig2] wrote {OUT}") + for ln in lines: + print(f"[fig2] {ln}") + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/biology_03_melanoblast_mitf.py b/scripts/figures/biology_03_melanoblast_mitf.py new file mode 100644 index 0000000000000000000000000000000000000000..c41e7857040969f94b399da2cc57593d20ebf500 --- /dev/null +++ b/scripts/figures/biology_03_melanoblast_mitf.py @@ -0,0 +1,101 @@ +"""mitf loss vs eda-derepression in cKO melanoblasts.""" +from __future__ import annotations +import warnings, sys +warnings.filterwarnings("ignore") +from pathlib import Path +import json +import numpy as np, pandas as pd +import matplotlib; matplotlib.use("Agg") +import matplotlib.pyplot as plt +from scipy.stats import mannwhitneyu + +sys.path.insert(0, str(Path(__file__).parent)) +from palette import apply_style, GENOTYPE_COLORS +apply_style() + +ROOT = Path("/home/bcheng/PRISM") +CSV = ROOT / "discovery/pan_skin/marker/106_melanoblast_nc_scores.csv" +SUMMARY = ROOT / "discovery/pan_skin/marker/106_melanoblast_nc_summary.json" +OUT = ROOT / "figures/biology/biology_03_melanoblast_mitf.pdf" + + +def main(): + df = pd.read_csv(CSV) + with open(SUMMARY) as f: + summ = json.load(f) + p_within = summ["within_cKO_top_vs_bot_derepression_quartile"]["MITF_regulon"]["mannu_p"] + + cko = df[df["group"] == "En1-cKO"].copy() + wt = df[df["group"] == "WT"].copy() + print(f"[fig3] cKO melanoblasts: {len(cko)}, WT: {len(wt)}") + + q = pd.qcut(cko["derepression"], 4, labels=["Q1_lo", "Q2", "Q3", "Q4_hi"]) + cko["derep_q"] = q + + fig, axes = plt.subplots(1, 2, figsize=(17, 7.0)) + + ax = axes[0] + order = ["Q1_lo", "Q2", "Q3", "Q4_hi"] + data = [cko.loc[cko["derep_q"] == g, "pw_MITF_regulon"].values for g in order] + + parts = ax.violinplot(data, positions=range(4), widths=0.75, + showmeans=False, showmedians=True, showextrema=False) + palette = [GENOTYPE_COLORS["WT"], "#abdda4", "#fdae61", GENOTYPE_COLORS["En1-cKO"]] + for pc, c in zip(parts["bodies"], palette): + pc.set_facecolor(c); pc.set_edgecolor("black"); pc.set_alpha(0.75) + parts["cmedians"].set_color("black") + + # strip overlay, subsampled + rng = np.random.default_rng(0) + for i, d in enumerate(data): + idx = rng.choice(len(d), size=min(120, len(d)), replace=False) + xj = i + rng.normal(0, 0.05, size=len(idx)) + ax.scatter(xj, d[idx], s=3, color="black", alpha=0.35, rasterized=True) + + # Q1 vs Q4 sig bar + y_max = max(np.max(d) for d in data) + y_line = y_max + 0.06 + ax.plot([0, 3], [y_line, y_line], color="black", lw=1) + ax.plot([0, 0], [y_line - 0.02, y_line], color="black", lw=1) + ax.plot([3, 3], [y_line - 0.02, y_line], color="black", lw=1) + ax.text(1.5, y_line + 0.02, f"p = {p_within:.2e}", ha="center", fontsize=12) + + ax.set_xticks(range(4)) + ax.set_xticklabels([f"{lab}\n(n={int(len(d))})" for lab, d in zip(order, data)]) + ax.set_xlabel("Derepression quartile (Eda-ectodysplasin, cKO only)") + ax.set_ylabel("MITF regulon score") + ax.set_title("(a) Within-cKO quartile split", fontsize=15) + ax.text(-0.14, 1.05, "(a)", transform=ax.transAxes, + fontsize=20, fontweight="bold", va="bottom", ha="right") + + ax = axes[1] + ax.scatter(wt["derepression"], wt["pw_MITF_regulon"], s=6, alpha=0.35, + color=GENOTYPE_COLORS["WT"], label=f"WT (n={int(len(wt))})", rasterized=True) + ax.scatter(cko["derepression"], cko["pw_MITF_regulon"], s=6, alpha=0.45, + color=GENOTYPE_COLORS["En1-cKO"], label=f"cKO (n={int(len(cko))})", rasterized=True) + # cKO linear trend + from numpy.polynomial import polynomial as P + xx = np.linspace(cko["derepression"].min(), cko["derepression"].max(), 40) + coef = np.polyfit(cko["derepression"].values, cko["pw_MITF_regulon"].values, 1) + ax.plot(xx, np.polyval(coef, xx), color="black", lw=1.4, + label=f"cKO linear fit (slope={coef[0]:.3f})") + ax.axhline(0, color="grey", lw=0.6, ls=":") + ax.axvline(0, color="grey", lw=0.6, ls=":") + ax.set_xlabel("Eda-ectodysplasin derepression score") + ax.set_ylabel("MITF regulon score") + ax.set_title("(b) WT vs cKO baseline shift", fontsize=15) + ax.legend(loc="center left", bbox_to_anchor=(1.02, 0.5), frameon=False) + ax.text(-0.14, 1.05, "(b)", transform=ax.transAxes, + fontsize=20, fontweight="bold", va="bottom", ha="right") + + plt.suptitle("Derepression drives MITF loss in cKO melanoblasts " + "(Eda-ectodysplasin axis; cKO-specific)", + fontsize=20, y=1.02, fontweight="bold") + plt.tight_layout() + plt.savefig(OUT, bbox_inches="tight", dpi=180) + plt.close() + print(f"[fig3] wrote {OUT} p(Q4 vs Q1) = {p_within:.2e}") + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/biology_04_dahlin_metabolism.py b/scripts/figures/biology_04_dahlin_metabolism.py new file mode 100644 index 0000000000000000000000000000000000000000..986dd6478de6881dd9ddbba28eef05e09e6dbf17 --- /dev/null +++ b/scripts/figures/biology_04_dahlin_metabolism.py @@ -0,0 +1,85 @@ +"""kit-w41 vs wt metabolic module delta heatmap (9 lineages x 4 modules).""" +from __future__ import annotations +import warnings, sys +warnings.filterwarnings("ignore") +from pathlib import Path +import numpy as np, pandas as pd +import matplotlib; matplotlib.use("Agg") +import matplotlib.pyplot as plt +from matplotlib.colors import TwoSlopeNorm + +sys.path.insert(0, str(Path(__file__).parent)) +from palette import apply_style +apply_style() + +ROOT = Path("/home/bcheng/PRISM") +CSV = ROOT / "discovery/hematopoiesis/marker/108_dahlin_lineage_metabolism.csv" +OUT = ROOT / "figures/biology/biology_04_dahlin_metabolism.pdf" + +MODULE_ORDER = ["OXPHOS_ETC", "Glycolysis", "Fatty_acid_oxidation", "Redox_glutathione"] +MODULE_LABELS = { + "OXPHOS_ETC": "OXPHOS/ETC", + "Glycolysis": "Glycolysis", + "Fatty_acid_oxidation": "FAO", + "Redox_glutathione": "Redox", +} +LINEAGE_ORDER = ["LT-HSC", "MPP", "lymphoid", "erythroid", "megakaryocyte", + "myeloid", "monocyte", "macrophage", "basophil-mast"] + + +def sig_stars(p): + if not np.isfinite(p): return "" + if p < 1e-4: return "***" + if p < 1e-3: return "**" + if p < 0.05: return "*" + return "" + + +def main(): + df = pd.read_csv(CSV) + delta = df.pivot(index="class", columns="module", values="delta").reindex( + index=LINEAGE_ORDER, columns=MODULE_ORDER) + padj = df.pivot(index="class", columns="module", + values="mannu_p_adj_bonferroni").reindex( + index=LINEAGE_ORDER, columns=MODULE_ORDER) + n_kit = df.pivot(index="class", columns="module", values="n_Kit_W41").reindex( + index=LINEAGE_ORDER, columns=MODULE_ORDER).iloc[:, 0] + n_wt = df.pivot(index="class", columns="module", values="n_WT").reindex( + index=LINEAGE_ORDER, columns=MODULE_ORDER).iloc[:, 0] + + vmax = max(abs(delta.values.min()), abs(delta.values.max())) + norm = TwoSlopeNorm(vmin=-vmax, vcenter=0.0, vmax=vmax) + + fig, ax = plt.subplots(figsize=(10.5, 7.5)) + im = ax.imshow(delta.values, cmap="RdBu_r", norm=norm, aspect="auto") + + for i in range(delta.shape[0]): + for j in range(delta.shape[1]): + v = delta.values[i, j] + p = padj.values[i, j] + txt = f"{v:+.2f}\n{sig_stars(p)}" if sig_stars(p) else f"{v:+.2f}" + color = "white" if abs(v) > vmax * 0.55 else "black" + ax.text(j, i, txt, ha="center", va="center", + fontsize=10, color=color) + + ax.set_xticks(range(delta.shape[1])) + ax.set_xticklabels([MODULE_LABELS[m] for m in delta.columns], fontsize=14) + ax.set_yticks(range(delta.shape[0])) + ax.set_yticklabels([f"{c}\n(Kit={int(n_kit[c])}, WT={int(n_wt[c])})" + for c in delta.index], fontsize=12) + ax.set_xlabel("Metabolic module") + ax.set_title("Dahlin Kit-W41 vs WT metabolic delta by lineage\n" + "positive = elevated in Kit-W41; * padj<.05, ** <1e-3, *** <1e-4", + fontsize=18) + + cbar = plt.colorbar(im, ax=ax, shrink=0.72, pad=0.02) + cbar.set_label("delta (Kit-W41 minus WT module score)", fontsize=13) + + plt.tight_layout() + plt.savefig(OUT, bbox_inches="tight", dpi=180) + plt.close() + print(f"[fig4] wrote {OUT}") + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/biology_05_dahlin_composition.py b/scripts/figures/biology_05_dahlin_composition.py new file mode 100644 index 0000000000000000000000000000000000000000..9717bf984fadd3681cbad368b1bc295a26483e59 --- /dev/null +++ b/scripts/figures/biology_05_dahlin_composition.py @@ -0,0 +1,102 @@ +"""dahlin wt vs kit-w41 lineage composition and log2fc.""" +from __future__ import annotations +import warnings, sys +warnings.filterwarnings("ignore") +from pathlib import Path +import numpy as np, pandas as pd +import matplotlib; matplotlib.use("Agg") +import matplotlib.pyplot as plt + +sys.path.insert(0, str(Path(__file__).parent)) +from palette import apply_style, color_for, GENOTYPE_COLORS +apply_style() + +ROOT = Path("/home/bcheng/PRISM") +CSV = ROOT / "discovery/hematopoiesis/marker/dahlin_predictions.csv" +OUT = ROOT / "figures/biology/biology_05_dahlin_composition.pdf" + +GT_MAP = {"SIGAB1":"WT","SIGAC1":"WT","SIGAD1":"WT","SIGAF1":"WT", + "SIGAG1":"WT","SIGAH1":"WT","SIGAG8":"Kit_W41","SIGAH8":"Kit_W41"} + +LINEAGE_ORDER = ["LT-HSC", "MPP", "lymphoid", "erythroid", "megakaryocyte", + "myeloid", "monocyte", "macrophage", "basophil-mast"] + + +def main(): + df = pd.read_csv(CSV) + df["sample"] = df["cell_id"].str.split("_").str[0] + df["genotype"] = df["sample"].map(GT_MAP).fillna("other") + + df = df[df["genotype"].isin(["WT", "Kit_W41"])] + tab = pd.crosstab(df["pred_label"], df["genotype"]).reindex(LINEAGE_ORDER).fillna(0) + print("[fig5] counts:") + print(tab) + + frac = tab.div(tab.sum(axis=0), axis=1) # columns are genotypes + log2fc = np.log2((frac["Kit_W41"] + 1e-4) / (frac["WT"] + 1e-4)) + + fig, axes = plt.subplots(1, 2, figsize=(17, 7.2), gridspec_kw={"width_ratios": [1.1, 1]}) + + ax = axes[0] + genos = ["WT", "Kit_W41"] + bottoms = np.zeros(len(genos)) + for lin in LINEAGE_ORDER: + vals = np.array([frac[g].loc[lin] for g in genos]) + ax.bar(genos, vals, bottom=bottoms, + color=color_for(lin), edgecolor="white", + linewidth=0.7, label=lin) + for j, v in enumerate(vals): + if v > 0.035: + ax.text(j, bottoms[j] + v / 2, f"{v*100:.1f}%", + ha="center", va="center", fontsize=9, color="white", + fontweight="bold") + bottoms += vals + ax.set_ylabel("Within-genotype fraction") + ax.set_ylim(0, 1) + ax.set_title(f"Lineage composition — WT (n={int(tab['WT'].sum())}) vs " + f"Kit-W41 (n={int(tab['Kit_W41'].sum())})") + ax.legend(loc="center left", bbox_to_anchor=(1.02, 0.5), + frameon=False, ncol=1, title="Lineage") + ax.text(-0.12, 1.05, "(a)", transform=ax.transAxes, + fontsize=20, fontweight="bold", va="bottom", ha="right") + + ax = axes[1] + lins = LINEAGE_ORDER + vals = [log2fc.loc[l] for l in lins] + colors = [GENOTYPE_COLORS["Kit_W41"] if v > 0 else GENOTYPE_COLORS["WT"] for v in vals] + y = np.arange(len(lins)) + ax.barh(y, vals, color=colors, edgecolor="black", linewidth=0.6) + ax.axvline(0, color="black", lw=0.8) + # value labels always sit just to the RIGHT of the bar tip regardless of + # sign, so negative-bar labels never crash into the y-tick labels on the + # left. Positive bars: label extends outward to the right. Negative bars: + # label sits inside the axis (toward zero) on the right of the tip. + for i, (l, v) in enumerate(zip(lins, vals)): + n_kit = int(tab["Kit_W41"].loc[l]); n_wt = int(tab["WT"].loc[l]) + ax.annotate( + f"{v:+.2f} (Kit={n_kit}, WT={n_wt})", + xy=(v, i), xycoords="data", + xytext=(6, 0), textcoords="offset points", + ha="left", va="center", fontsize=10.5, clip_on=False, + ) + ax.set_yticks(y); ax.set_yticklabels(lins) + ax.tick_params(axis="y", pad=6) + ax.invert_yaxis() + ax.set_xlabel("log2(Kit-W41 fraction / WT fraction)") + ax.set_title("Compositional shift — Kit-W41 expands erythroid, contracts lymphoid") + max_abs = max(abs(min(vals)), abs(max(vals))) + # Right-side padding accommodates the value+n text (~2.6 units at 10.5pt) + ax.set_xlim(-max_abs * 1.15, max_abs * 2.8) + ax.text(-0.14, 1.05, "(b)", transform=ax.transAxes, + fontsize=20, fontweight="bold", va="bottom", ha="right") + + plt.suptitle("Dahlin — Kit-W41 alters lineage output balance", + fontsize=20, y=1.02, fontweight="bold") + plt.tight_layout() + plt.savefig(OUT, bbox_inches="tight", dpi=180) + plt.close() + print(f"[fig5] wrote {OUT}") + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/biology_06_veres_beta_quadrant.py b/scripts/figures/biology_06_veres_beta_quadrant.py new file mode 100644 index 0000000000000000000000000000000000000000..ccfb96538a97a4c40af0e630920aabd446507134 --- /dev/null +++ b/scripts/figures/biology_06_veres_beta_quadrant.py @@ -0,0 +1,76 @@ +"""veres beta maturity quadrants: INS-level x MAT-score.""" +from __future__ import annotations +import warnings, sys +warnings.filterwarnings("ignore") +from pathlib import Path +import json +import numpy as np, pandas as pd +import matplotlib; matplotlib.use("Agg") +import matplotlib.pyplot as plt + +sys.path.insert(0, str(Path(__file__).parent)) +from palette import apply_style +apply_style() + +ROOT = Path("/home/bcheng/PRISM") +CSV = ROOT / "discovery/pancreas/marker/109_veres_mature_beta_scores.csv" +SUMMARY = ROOT / "discovery/pancreas/marker/109_veres_mature_beta_summary.json" +OUT = ROOT / "figures/biology/biology_06_veres_beta_quadrant.pdf" + +QUAD_COLORS = { + "INS+/MAT+": "#d7191c", + "INS+/MAT-": "#fdae61", + "INS-/MAT+": "#abdda4", + "INS-/MAT-": "#2b83ba", +} + + +def main(): + df = pd.read_csv(CSV) + with open(SUMMARY) as f: + summ = json.load(f) + ins_thr = summ["ins_threshold"] + mat_thr = summ["mat_threshold"] + print(f"[fig6] thresholds: INS={ins_thr:.3f}, MAT={mat_thr:.3f}") + counts = summ["quadrant_counts"] + print("[fig6] quadrant counts:", counts) + + fig, ax = plt.subplots(figsize=(10.5, 8.5)) + for q, color in QUAD_COLORS.items(): + m = df["quadrant"] == q + ax.scatter(df.loc[m, "INS_level"], df.loc[m, "MAT_score"], + s=26, alpha=0.70, edgecolor="black", linewidths=0.35, + color=color, label=f"{q} (n={int(counts.get(q, m.sum()))})", + rasterized=True) + + ax.axvline(ins_thr, color="black", lw=1.0, ls="--", alpha=0.7) + ax.axhline(mat_thr, color="black", lw=1.0, ls="--", alpha=0.7) + + xmin, xmax = df["INS_level"].min(), df["INS_level"].max() + ymin, ymax = df["MAT_score"].min(), df["MAT_score"].max() + ax.text(xmax, ymax, "mature adult-beta", + ha="right", va="top", fontsize=12, color="#7b0000", fontweight="bold") + ax.text(xmax, ymin, "SC-beta (INS+, immature)", + ha="right", va="bottom", fontsize=12, color="#8b5a00", fontweight="bold") + ax.text(xmin, ymax, "mature-but-INS-low", + ha="left", va="top", fontsize=12, color="#2a662a", fontweight="bold") + ax.text(xmin, ymin, "immature INS-low", + ha="left", va="bottom", fontsize=12, color="#0f3b6a", fontweight="bold") + + ax.set_xlabel("INS level = log1p(Ins1 + Ins2)") + ax.set_ylabel("Maturity score = z(Mafa) + z(Ucn3)") + ax.set_title(f"Veres PANDA-beta cells (n={int(len(df))}) — INS x maturity quadrants\n" + f"only {int(counts['INS+/MAT+'])}/{int(len(df))} " + f"({100*counts['INS+/MAT+']/len(df):.1f}%) are mature adult-like", + fontsize=18) + ax.legend(loc="center left", bbox_to_anchor=(1.02, 0.5), + frameon=False, title="Quadrant") + + plt.tight_layout() + plt.savefig(OUT, bbox_inches="tight", dpi=180) + plt.close() + print(f"[fig6] wrote {OUT}") + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/biology_07_veres_polyhormonal.py b/scripts/figures/biology_07_veres_polyhormonal.py new file mode 100644 index 0000000000000000000000000000000000000000..45e93255dc9c211fa3ca1ae1bc248477c222300d --- /dev/null +++ b/scripts/figures/biology_07_veres_polyhormonal.py @@ -0,0 +1,118 @@ +"""veres alpha-pool umap + polyhormonal fraction per leiden cluster.""" +from __future__ import annotations +import warnings, sys +warnings.filterwarnings("ignore") +from pathlib import Path +import numpy as np, pandas as pd +import anndata as ad +import scanpy as sc +import matplotlib; matplotlib.use("Agg") +import matplotlib.pyplot as plt +import matplotlib.patches as mpatches + +sys.path.insert(0, str(Path(__file__).parent)) +from palette import apply_style +apply_style() + +ROOT = Path("/home/bcheng/PRISM") +SCORES = ROOT / "discovery/pancreas/marker/110_veres_polyhormonal_alpha_scores.csv" +PER_CL = ROOT / "discovery/pancreas/marker/110_veres_polyhormonal_alpha_per_cluster.csv" +ADATA = ROOT / "data/corpus/pancreas/held_out_labeled/veres_GSE114412_test.h5ad" +CACHE = ROOT / "figures/biology/_cache_veres_alpha_umap.npz" +OUT = ROOT / "figures/biology/biology_07_veres_polyhormonal.pdf" + +CLUSTER_COLORS = { + "0": "#1f77b4", "1": "#ff7f0e", "2": "#2ca02c", "3": "#d7191c", + "4": "#9467bd", "5": "#8c564b", "6": "#e377c2", "7": "#7f7f7f", +} + + +def compute_umap(): + scores = pd.read_csv(SCORES) + a = ad.read_h5ad(ADATA) + a = a[scores["cell_id"].values].copy() + print(f"[fig7] alpha subset: {a.shape}") + # skip normalize+log1p if X already log-scaled + X_max = a.X.max() if not hasattr(a.X, "toarray") else a.X.max() + if X_max > 100: + sc.pp.normalize_total(a, target_sum=1e4) + sc.pp.log1p(a) + sc.pp.highly_variable_genes(a, n_top_genes=1500) + a2 = a[:, a.var["highly_variable"]].copy() + sc.pp.scale(a2, max_value=10) + sc.tl.pca(a2, n_comps=30, random_state=42) + sc.pp.neighbors(a2, n_neighbors=15, random_state=42) + sc.tl.umap(a2, random_state=42, min_dist=0.35) + umap = a2.obsm["X_umap"] + np.savez(CACHE, + umap=umap, + cell_id=a.obs_names.astype(str).values, + leiden_alpha=scores["leiden_alpha"].astype(str).values) + return umap, scores["leiden_alpha"].astype(str).values + + +def main(): + if CACHE.exists(): + cache = np.load(CACHE, allow_pickle=True) + umap = cache["umap"] + leiden = cache["leiden_alpha"].astype(str) + print(f"[fig7] loaded cached UMAP ({umap.shape})") + else: + umap, leiden = compute_umap() + + per_cl = pd.read_csv(PER_CL) + per_cl["leiden_alpha"] = per_cl["leiden_alpha"].astype(str) + per_cl = per_cl.sort_values("leiden_alpha") + + fig, axes = plt.subplots(1, 2, figsize=(17, 7.2), + gridspec_kw={"width_ratios": [1.15, 1]}) + + ax = axes[0] + for cid in sorted(np.unique(leiden), key=int): + m = leiden == cid + ax.scatter(umap[m, 0], umap[m, 1], s=9, alpha=0.75, + color=CLUSTER_COLORS.get(cid, "#888"), + linewidths=0, rasterized=True, + label=f"c{cid} (n={int(m.sum())})") + cx, cy = umap[m, 0].mean(), umap[m, 1].mean() + ax.text(cx, cy, cid, fontsize=14, color="black", + ha="center", va="center", fontweight="bold", + bbox=dict(boxstyle="round,pad=0.15", fc="white", + ec="black", lw=0.6, alpha=0.85)) + ax.set_title(f"Veres SC-alpha pool (n={int(len(umap))}) — leiden sub-clusters") + ax.set_xlabel("UMAP-1"); ax.set_ylabel("UMAP-2") + ax.set_xticks([]); ax.set_yticks([]) + ax.legend(loc="center left", bbox_to_anchor=(1.02, 0.5), + frameon=False, title="Cluster") + ax.text(-0.10, 1.05, "(a)", transform=ax.transAxes, + fontsize=20, fontweight="bold", va="bottom", ha="right") + + ax = axes[1] + x = np.arange(len(per_cl)) + vals = per_cl["frac_polyhormonal"].values * 100 + colors = [CLUSTER_COLORS.get(cid, "#888") for cid in per_cl["leiden_alpha"]] + bars = ax.bar(x, vals, color=colors, edgecolor="black", linewidth=0.5) + for i, (b, v, n) in enumerate(zip(bars, vals, per_cl["n_cells"])): + ax.text(b.get_x() + b.get_width()/2, v + 1.0, + f"{v:.0f}%\n(n={int(n)})", ha="center", va="bottom", fontsize=11) + baseline = 100 * 0.1814 + ax.axhline(baseline, color="black", ls="--", lw=1, + label=f"pool baseline ({baseline:.1f}%)") + ax.set_xticks(x); ax.set_xticklabels([f"c{c}" for c in per_cl["leiden_alpha"]]) + ax.set_ylabel("Polyhormonal fraction (%)") + ax.set_title("Cluster 3 concentrates polyhormonal SC-alpha cells (2.5x baseline)") + ax.set_ylim(0, max(vals) * 1.28) + ax.legend(loc="upper right", frameon=False) + ax.text(-0.10, 1.05, "(b)", transform=ax.transAxes, + fontsize=20, fontweight="bold", va="bottom", ha="right") + + plt.suptitle("Polyhormonal SC-alpha subcluster (Veres — 3,473 alpha-pool cells)", + fontsize=20, y=1.02, fontweight="bold") + plt.tight_layout() + plt.savefig(OUT, bbox_inches="tight", dpi=180) + plt.close() + print(f"[fig7] wrote {OUT}") + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/biology_08_prototype_geometry.py b/scripts/figures/biology_08_prototype_geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..1096be137639edd4792326927bd8b0a512fd1d7b --- /dev/null +++ b/scripts/figures/biology_08_prototype_geometry.py @@ -0,0 +1,72 @@ +"""prototype cosine heatmaps for skin/hematopoiesis/pancreas.""" +from __future__ import annotations +import warnings, sys +warnings.filterwarnings("ignore") +from pathlib import Path +import json +import numpy as np, pandas as pd +import matplotlib; matplotlib.use("Agg") +import matplotlib.pyplot as plt +from matplotlib.colors import TwoSlopeNorm + +sys.path.insert(0, str(Path(__file__).parent)) +from palette import apply_style +apply_style() + +ROOT = Path("/home/bcheng/PRISM") +DISC = ROOT / "discovery" +OUT = ROOT / "figures/biology/biology_08_prototype_geometry.pdf" + +SYSTEMS = [("pan_skin", "Skin (K=13)"), + ("hematopoiesis", "Hematopoiesis (K=15)"), + ("pancreas", "Pancreas (K=20)")] + +PANEL_LETTERS = ["(a)", "(b)", "(c)"] + + +def main(): + with open(DISC / "70_prototype_effective_dim.json") as f: + edim = json.load(f) + + fig, axes = plt.subplots(1, 3, figsize=(26, 8.5), + gridspec_kw={"width_ratios": [1, 1.15, 1.5]}) + + for idx, (ax, (sys_key, title)) in enumerate(zip(axes, SYSTEMS)): + M = pd.read_csv(DISC / f"70_prototype_intra_cosine_{sys_key}.csv", index_col=0) + m = M.values + vmax = max(abs(m.min()), abs(m.max())) + norm = TwoSlopeNorm(vmin=-vmax, vcenter=0.0, vmax=vmax) + im = ax.imshow(m, cmap="RdBu_r", norm=norm, aspect="auto") + eff = edim[sys_key]["effective_dim"] + K = edim[sys_key]["K"] + # shorten pancreas labels to stop the tick text going tiny + def _short(lbl: str) -> str: + s = str(lbl) + s = s.replace("_progenitor", "_prog") + s = s.replace("endocrine-progenitor", "endo-prog") + s = s.replace("pancreatic-progenitor", "panc-prog") + return s + col_labels = [_short(c) for c in M.columns] + row_labels = [_short(c) for c in M.index] + tick_fs = 10 if sys_key != "pancreas" else 9 + ax.set_xticks(range(len(M.columns))) + ax.set_xticklabels(col_labels, rotation=90, fontsize=tick_fs) + ax.set_yticks(range(len(M.index))) + ax.set_yticklabels(row_labels, fontsize=tick_fs) + ax.set_title(f"{title}\neffective dim = {eff:.2f} / {K}") + cbar = plt.colorbar(im, ax=ax, shrink=0.75, pad=0.02) + cbar.set_label("cosine similarity", fontsize=13) + ax.text(-0.12, 1.05, PANEL_LETTERS[idx], transform=ax.transAxes, + fontsize=20, fontweight="bold", va="bottom", ha="right") + + plt.suptitle("Prototype geometry — near-orthogonal class prototypes " + "(effective dim ~11 across systems)", + fontsize=20, y=1.02, fontweight="bold") + plt.tight_layout() + plt.savefig(OUT, bbox_inches="tight", dpi=180) + plt.close() + print(f"[fig8] wrote {OUT}") + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/biology_99_merge_supplement.py b/scripts/figures/biology_99_merge_supplement.py new file mode 100644 index 0000000000000000000000000000000000000000..33ba513b90bab5f82cde775dfeb18bf6520bb29e --- /dev/null +++ b/scripts/figures/biology_99_merge_supplement.py @@ -0,0 +1,42 @@ +"""append biology pages to PANDA_supplement.pdf in place.""" +from __future__ import annotations +from pathlib import Path +from pypdf import PdfWriter + +ROOT = Path("/home/bcheng/PRISM") +SUP = ROOT / "figures/PANDA_supplement.pdf" +BIO = ROOT / "figures/biology" + +BIO_ORDER = [ + BIO / "biology_01_dingwall_umap.pdf", + BIO / "biology_02_primary_eden.pdf", + BIO / "biology_03_melanoblast_mitf.pdf", + BIO / "biology_04_dahlin_metabolism.pdf", + BIO / "biology_05_dahlin_composition.pdf", + BIO / "biology_06_veres_beta_quadrant.pdf", + BIO / "biology_07_veres_polyhormonal.pdf", + BIO / "biology_08_prototype_geometry.pdf", +] + + +def main(): + if not SUP.exists(): + raise FileNotFoundError(f"missing {SUP}") + w = PdfWriter() + w.append(str(SUP)) + n_base = len(w.pages) + print(f"[merge] base supplement: {n_base} pages") + for p in BIO_ORDER: + if not p.exists(): + print(f" [skip] {p.name} missing") + continue + w.append(str(p)) + print(f" + {p.name}") + with open(SUP, "wb") as f: + w.write(f) + print(f"[merge] wrote {SUP} ({SUP.stat().st_size / 1024:.0f} KB, " + f"{len(w.pages)} pages)") + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/build_figure_supplement.py b/scripts/figures/build_figure_supplement.py new file mode 100644 index 0000000000000000000000000000000000000000..909b1453bdf9800dec93c2d48cea35bea45672c8 --- /dev/null +++ b/scripts/figures/build_figure_supplement.py @@ -0,0 +1,746 @@ +"""builds figures/supplement/*.pdf and merges into figures/PANDA_supplement.pdf.""" +from __future__ import annotations +from pathlib import Path +import warnings, json, pickle, sys, numpy as np, pandas as pd +warnings.filterwarnings("ignore") + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import matplotlib.patches as patches +from matplotlib.backends.backend_pdf import PdfPages +from matplotlib.colors import LinearSegmentedColormap, TwoSlopeNorm + +sys.path.insert(0, str(Path(__file__).parent)) +from palette import apply_style, color_for +apply_style() + +ROOT = Path("/home/bcheng/PRISM") +FIG = ROOT / "figures" +FIG_S = FIG / "supplement"; FIG_S.mkdir(parents=True, exist_ok=True) +DISC = ROOT / "discovery" +SUP_PDF = FIG / "PANDA_supplement.pdf" + +SYSTEMS = ["pan_skin", "hematopoiesis", "pancreas"] +SYS_LABEL = {"pan_skin": "Pan-skin", "hematopoiesis": "Pan-hematopoietic", "pancreas": "Pan-pancreatic"} +SYS_COLOR = {"pan_skin": "#2b83ba", "hematopoiesis": "#d7191c", "pancreas": "#7570b3"} + +# Deprecated class names -- drop these if they appear in any legacy CSV so +# regenerated figures never re-surface stale vocabulary. Extend as vocab evolves. +STALE_CLASSES = { + "HF-DP", "basal-multipotent", "eden-dermal-niche", "eccrine-secretory", + "hair-placode", "UNK", "nascent-eccrine-gland", "dermal-condensate", + "mesenchymal", "adult-alpha", "adult-beta", "eccrine-duct", "eccrine-ductal", + "eccrine-placode", "other", "unassigned", +} + + +def _drop_stale(labels): + """Return input filtered to remove any deprecated class names. + + Accepts a list/Index/Series of strings, or a DataFrame column selector case + handled by the caller. Order preserved. + """ + return [x for x in labels if x not in STALE_CLASSES] + + +def _filter_df_stale(df, col): + """Drop rows of df whose value in `col` is a deprecated class name.""" + if col in df.columns: + return df[~df[col].astype(str).isin(STALE_CLASSES)].reset_index(drop=True) + return df + +CV_JSON = { + "pan_skin": ROOT / "discovery/pan_skin/marker/cv_5fold.json", + "hematopoiesis": ROOT / "discovery/hematopoiesis/marker/cv_5fold.json", + "pancreas": ROOT / "discovery/pancreas/marker/cv_5fold.json", +} + +# ============================================================ +# PAGE 1: Held-out CV summary (per-system accuracy + AUROC + F1) +# ============================================================ + +def page_cv_summary(): + fig, axes = plt.subplots(1, 2, figsize=(15.5, 6.5)) + sys_names = [] + accs, acc_errs, aucs, auc_errs, folds_data = [], [], [], [], [] + for sys in SYSTEMS: + r = json.load(open(CV_JSON[sys])) + sys_names.append(SYS_LABEL[sys]) + accs.append(r["mean_acc"]); acc_errs.append(r["std_acc"]) + aucs.append(r["mean_auc"]); auc_errs.append(r["std_auc"]) + folds_data.append((r["per_fold_acc"], r["per_fold_auc"])) + + ax = axes[0] + x = np.arange(len(sys_names)) + b1 = ax.bar(x - 0.2, accs, 0.4, yerr=acc_errs, capsize=4, + color=[SYS_COLOR[s] for s in SYSTEMS], label="accuracy", alpha=0.85, edgecolor="k", linewidth=0.7) + b2 = ax.bar(x + 0.2, aucs, 0.4, yerr=auc_errs, capsize=4, + color=[SYS_COLOR[s] for s in SYSTEMS], label="macro AUROC", alpha=0.55, edgecolor="k", linewidth=0.7, hatch="//") + ax.set_xticks(x); ax.set_xticklabels(sys_names, rotation=0, fontsize=13) + ax.set_ylabel("Held-out score") + ax.set_ylim(0.7, 1.02) + ax.set_title("(a) 5-fold held-out CV: accuracy + macro AUROC (mean ± std)", fontsize=15) + for i, (pa, pu) in enumerate(folds_data): + ax.scatter([x[i] - 0.2] * len(pa), pa, s=22, color="black", zorder=3) + ax.scatter([x[i] + 0.2] * len(pu), pu, s=22, color="black", zorder=3) + for i, (a, e, u, ue) in enumerate(zip(accs, acc_errs, aucs, auc_errs)): + ax.text(x[i] - 0.2, a + e + 0.008, f"{a:.3f}", ha="center", fontsize=11) + ax.text(x[i] + 0.2, u + ue + 0.008, f"{u:.3f}", ha="center", fontsize=11) + ax.legend(loc="lower right") + + ax = axes[1] + baron_pca = json.load(open(DISC / "pancreas/pca/93_baron_merged_metrics.json")) + baron_mar = json.load(open(DISC / "pancreas/marker/93_baron_merged_metrics.json")) + nest_pca = json.load(open(DISC / "hematopoiesis/pca/97_nestorowa_anchor_zero_shot.json")) + nest_mar = json.load(open(DISC / "hematopoiesis/marker/97_nestorowa_anchor_zero_shot.json")) + sul_pca = json.load(open(DISC / "pan_skin/pca/97_sulic_anchor_zero_shot.json")) + sul_mar = json.load(open(DISC / "pan_skin/marker/97_sulic_anchor_zero_shot.json")) + # sulic 5-seed topline: in-domain ceiling reference + sulic5 = json.load(open(ROOT / "scripts/sulic/sulic_panda_heldout_results.json")) + + bars_data = [ + ("Baron\ntest-half\nPCA", baron_pca["acc"], "acc", SYS_COLOR["pancreas"], 0.55), + ("Baron\ntest-half\nMarker", baron_mar["acc"], "acc", SYS_COLOR["pancreas"], 0.90), + ("Nestorowa\nanchor\nPCA", nest_pca["coarse_acc"], "acc", SYS_COLOR["hematopoiesis"], 0.55), + ("Nestorowa\nanchor\nMarker",nest_mar["coarse_acc"], "acc", SYS_COLOR["hematopoiesis"], 0.90), + ("Sulic\nanchor\nPCA", sul_pca["acc"], "acc", SYS_COLOR["pan_skin"], 0.55), + ("Sulic\nanchor\nMarker", sul_mar["acc"], "acc", SYS_COLOR["pan_skin"], 0.90), + ("Sulic\n5-fold\n(topline)", sulic5["testA"]["mean_auroc"], "AUROC", "#888", 0.90), + ] + xe = np.arange(len(bars_data)) + for i, (lab, val, met, col, alpha) in enumerate(bars_data): + ax.bar(xe[i], val, color=col, alpha=alpha, edgecolor="k", linewidth=0.7) + ax.text(xe[i], val + 0.012, f"{val:.3f}\n({met})", ha="center", fontsize=10) + ax.set_xticks(xe); ax.set_xticklabels([b[0] for b in bars_data], fontsize=10) + ax.set_ylim(0.4, 1.05) + ax.set_ylabel("Held-out score") + ax.set_title("(b) External held-out labeled validation (6 zero-shot + 1 topline)", fontsize=15) + + plt.suptitle("PANDA held-out validation across systems", fontsize=20, y=1.02, fontweight="bold") + plt.tight_layout() + plt.savefig(FIG_S / "01_cv_summary.pdf", bbox_inches="tight") + plt.close() + print("[fig] 01_cv_summary.pdf") + + +# ============================================================ +# PAGE 2: Per-class F1 across all three systems +# ============================================================ + +def page_per_class_f1(): + fig, axes = plt.subplots(1, 3, figsize=(21, 7.0)) + for ax, sys in zip(axes, SYSTEMS): + r = json.load(open(CV_JSON[sys])) + rep = r["per_class_report"] + classes = [c for c in rep if c not in ("accuracy", "macro avg", "weighted avg")] + f1s = [rep[c]["f1-score"] for c in classes] + supp = [rep[c]["support"] for c in classes] + order = np.argsort(f1s)[::-1] + classes = [classes[i] for i in order] + f1s = [f1s[i] for i in order]; supp = [supp[i] for i in order] + + y = np.arange(len(classes)) + cols = [color_for(c) for c in classes] + bars = ax.barh(y, f1s, color=cols, edgecolor="k", linewidth=0.5) + for i, (b, s) in enumerate(zip(bars, supp)): + # write inside the bar (right-aligned) so the axis can stay at 0..1.05 + xtxt = max(0.02, b.get_width() - 0.02) + ax.text(xtxt, b.get_y() + b.get_height() / 2, + f"F1={f1s[i]:.3f} · n={int(s):,}", + va="center", ha="right", fontsize=10, + color="white" if f1s[i] > 0.35 else "black") + ax.set_yticks(y); ax.set_yticklabels(classes, fontsize=11) + ax.set_xlim(0, 1.05) + ax.set_xlabel("F1") + ax.invert_yaxis() + ax.set_title(f"{SYS_LABEL[sys]}\nacc={r['mean_acc']:.4f} macro-AUROC={r['mean_auc']:.4f}", + fontsize=15) + ax.axvline(0.85, color="grey", linestyle=":", linewidth=0.7, alpha=0.6) + ax.grid(axis="x", alpha=0.3, linestyle="--") + plt.suptitle("Per-class F1 in 5-fold CV (bar colour = canonical class palette)", + fontsize=20, y=1.02, fontweight="bold") + plt.tight_layout() + plt.savefig(FIG_S / "02_per_class_f1.pdf", bbox_inches="tight") + plt.close() + print("[fig] 02_per_class_f1.pdf") + + +# ============================================================ +# PAGE 3: Prototype intra-cosine heatmaps (3 panels) +# ============================================================ + +def page_prototype_cosine(): + fig, axes = plt.subplots(1, 3, figsize=(23, 7.5)) + eff = json.load(open(DISC / "70_prototype_effective_dim.json")) + for ax, sys in zip(axes, SYSTEMS): + M = pd.read_csv(DISC / f"70_prototype_intra_cosine_{sys}.csv", index_col=0) + vmax = np.nanmax(np.abs(M.values)) + im = ax.imshow(M.values, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto") + # pancreas has 20x20 cells → tighter font + higher threshold + ann_fs = 6 if sys == "pancreas" else 8 + ann_thresh = 0.45 if sys == "pancreas" else 0.3 + for i in range(M.shape[0]): + for j in range(M.shape[1]): + v = M.values[i, j] + if abs(v) < ann_thresh: continue + ax.text(j, i, f"{v:+.2f}", ha="center", va="center", + fontsize=ann_fs, + color="white" if abs(v) > vmax * 0.55 else "black") + ax.set_xticks(range(len(M.columns))) + ax.set_xticklabels(M.columns, rotation=45, ha="right", fontsize=10) + ax.set_yticks(range(len(M.index))) + ax.set_yticklabels(M.index, fontsize=10) + ax.set_title(f"{SYS_LABEL[sys]}\nK={eff[sys]['K']} eff-dim={eff[sys]['effective_dim']:.2f}", + fontsize=15) + plt.colorbar(im, ax=ax, shrink=0.7, label="prototype cosine") + plt.suptitle("Prototype-prototype cosine matrices per system (§8.1)", + fontsize=20, y=1.02, fontweight="bold") + plt.tight_layout() + plt.savefig(FIG_S / "03_prototype_cosine.pdf", bbox_inches="tight") + plt.close() + print("[fig] 03_prototype_cosine.pdf") + + +# ============================================================ +# PAGE 4: Training trajectory eff-dim over stages +# ============================================================ + +def page_training_trajectory(): + return # skipped — v2 checkpoints only save final stage + fig, axes = plt.subplots(1, 2, figsize=(13, 4.6)) + ax = axes[0] + for sys in SYSTEMS: + d = pd.read_csv(DISC / f"83_{sys}_effdim_by_stage.csv") + stage_order = ["stage0", "stage1", "stage2", "stage3", "final"] + d = d.set_index("stage").loc[stage_order].reset_index() + ax.plot(range(len(d)), d["eff_dim"], "-o", color=SYS_COLOR[sys], + label=f"{SYS_LABEL[sys]} (K={int(d['K'].iloc[0])})", linewidth=2.5, markersize=8) + for i, row in d.iterrows(): + ax.text(i, row["eff_dim"] + 0.35, f"{row['eff_dim']:.2f}", + ha="center", fontsize=8, color=SYS_COLOR[sys]) + ax.set_xticks(range(5)); ax.set_xticklabels(stage_order, fontsize=10) + ax.set_xlabel("training stage"); ax.set_ylabel("prototype set effective dimensionality") + ax.set_title("(a) Effective-dim trajectory\nPancreas collapses 9.50→1.98 in stage 0→1", fontsize=11) + ax.legend(fontsize=9) + ax.grid(alpha=0.3, linestyle="--") + + ax = axes[1] + all_data, all_labels = [], [] + for sys in SYSTEMS: + traj = pd.read_csv(DISC / f"83_{sys}_prototype_trajectory.csv") + classes = traj["class"].unique() + total_cos = [] + for cls in classes: + # stage0 vs final cos — reload prototypes directly (intermediate-cos product is wrong) + import torch + ck0 = torch.load(ROOT / f"checkpoints/{sys}/panda_stage0.pt", + map_location="cpu", weights_only=False) + ckf = torch.load(ROOT / f"checkpoints/{sys}/panda_final.pt", + map_location="cpu", weights_only=False) + P0 = ck0["model"]["prototypes"].numpy() if hasattr(ck0["model"]["prototypes"], "numpy") else ck0["model"]["prototypes"] + Pf = ck0["model"]["prototypes"].numpy() if hasattr(ck0["model"]["prototypes"], "numpy") else ck0["model"]["prototypes"] + # actually need final + Pf = ckf.get("prototypes", ckf["model"]["prototypes"]) + if hasattr(Pf, "numpy"): Pf = Pf.numpy() + P0 = P0 / (np.linalg.norm(P0, axis=1, keepdims=True) + 1e-8) + Pf = Pf / (np.linalg.norm(Pf, axis=1, keepdims=True) + 1e-8) + ci = ck0["classes"].index(cls) + total_cos.append(float((P0[ci] * Pf[ci]).sum())) + break # only need once per sys + # per-sys one-shot pass + ck0 = torch.load(ROOT / f"checkpoints/{sys}/panda_stage0.pt", + map_location="cpu", weights_only=False) + ckf = torch.load(ROOT / f"checkpoints/{sys}/panda_final.pt", + map_location="cpu", weights_only=False) + P0 = ck0["model"]["prototypes"] + if hasattr(P0, "numpy"): P0 = P0.numpy() + Pf = ckf.get("prototypes", ckf["model"]["prototypes"]) + if hasattr(Pf, "numpy"): Pf = Pf.numpy() + P0 = P0 / (np.linalg.norm(P0, axis=1, keepdims=True) + 1e-8) + Pf = Pf / (np.linalg.norm(Pf, axis=1, keepdims=True) + 1e-8) + total_cos = [float((P0[i] * Pf[i]).sum()) for i in range(len(ck0["classes"]))] + all_data.append(total_cos) + all_labels.append([f"{sys[:3]}:{c}" for c in ck0["classes"]]) + + flat_vals = np.concatenate([np.asarray(a) for a in all_data]) + flat_labels = sum(all_labels, []) + order = np.argsort(flat_vals) + flat_vals = flat_vals[order] + flat_labels = [flat_labels[i] for i in order] + colors = ["#d7191c" if v < 0 else "#fdae61" if v < 0.15 else "#2b83ba" for v in flat_vals] + ax.barh(range(len(flat_vals)), flat_vals, color=colors, edgecolor="k", linewidth=0.4) + ax.set_yticks(range(len(flat_vals))) + ax.set_yticklabels(flat_labels, fontsize=6.5) + ax.axvline(0, color="k", linewidth=0.5) + ax.set_xlabel("cosine( prototype_stage0 , prototype_final )") + ax.set_title("(b) Per-class prototype drift 0→final\nNegative = moved to opposite direction", fontsize=11) + ax.set_xlim(-0.35, 0.45) + ax.grid(axis="x", alpha=0.3, linestyle="--") + + plt.suptitle("§10.1 Training trajectory: effective-dim collapse + per-class prototype drift", + fontsize=12, y=1.02) + plt.tight_layout() + plt.savefig(FIG_S / "04_training_trajectory.pdf", bbox_inches="tight") + plt.close() + print("[fig] 04_training_trajectory.pdf") + + +# ============================================================ +# PAGE 5: Adversary purification (dataset + depth) +# ============================================================ + +def page_adversary_purification(): + adv = json.load(open(DISC / "84_adversary_purification.json")) + fig, axes = plt.subplots(1, 2, figsize=(17, 6.5)) + + ax = axes[0] + xs, accs, chances, colors, labels = [], [], [], [], [] + for i, sys in enumerate(SYSTEMS): + r = adv[sys] + if "dataset_adv_accuracy" not in r: continue + xs.append(i); accs.append(r["dataset_adv_accuracy"]) + chances.append(r["dataset_adv_chance"]) + colors.append(SYS_COLOR[sys]) + labels.append(f"{SYS_LABEL[sys]}\n(n={r['n_train_datasets']} datasets)") + x = np.arange(len(xs)) + bars = ax.bar(x - 0.2, accs, 0.4, color=colors, edgecolor="k", linewidth=0.7, + label="observed", alpha=0.85) + ax.bar(x + 0.2, chances, 0.4, color=colors, edgecolor="k", linewidth=0.7, + label="chance", alpha=0.35, hatch="///") + for i, (a, c) in enumerate(zip(accs, chances)): + ax.text(x[i] - 0.2, a + 0.015, f"{a:.3f}", ha="center", fontsize=11) + ax.text(x[i] + 0.2, c + 0.015, f"{c:.3f}", ha="center", fontsize=11) + ax.text(x[i], -0.06, f"+{a-c:.3f}\nabove chance", ha="center", fontsize=10, color="red") + ax.set_xticks(x); ax.set_xticklabels(labels, fontsize=11) + ax.set_ylabel("dataset adversary accuracy") + ax.set_ylim(-0.1, 1.0) + ax.set_title("(a) Residual batch-signal\n(closer to chance = more purified)", fontsize=15) + ax.legend(loc="upper left") + ax.grid(axis="y", alpha=0.3, linestyle="--") + + ax = axes[1] + xs, r2s, colors, labels = [], [], [], [] + for i, sys in enumerate(SYSTEMS): + r = adv[sys] + if "depth_adv_r2_z" not in r: continue + xs.append(i); r2s.append(r["depth_adv_r2_z"]) + colors.append(SYS_COLOR[sys]) + labels.append(SYS_LABEL[sys]) + x = np.arange(len(xs)) + bars = ax.bar(x, r2s, color=colors, edgecolor="k", linewidth=0.7) + for i, r in enumerate(r2s): + # label above bar if positive, below top if bar tall enough + ypos = (r + 0.02) if r >= 0 else (r - 0.03) + ax.text(x[i], ypos, f"R²={r:.3f}", ha="center", fontsize=11, + color="black", fontweight="bold") + ax.axhline(0, color="k", linestyle="-", linewidth=0.8) + ax.set_xticks(x); ax.set_xticklabels(labels, fontsize=13) + ax.set_ylabel("Depth-adversary R²", labelpad=10) + ax.set_ylim(-0.05, 0.75) + ax.set_title("(b) Depth adversary R² (target ≤ 0)\nAll three systems purified", fontsize=15) + ax.grid(axis="y", alpha=0.3, linestyle="--") + ax.text(0.5, 0.95, "R² ≤ 0 ⇒ trunk carries no depth signal", + transform=ax.transAxes, fontsize=11, ha="center", color="green") + + plt.suptitle("§10.2 Adversary purification: depth fully removed, dataset only partially", + fontsize=20, y=1.02, fontweight="bold") + plt.tight_layout() + plt.savefig(FIG_S / "05_adversary_purification.pdf", bbox_inches="tight") + plt.close() + print("[fig] 05_adversary_purification.pdf") + + +# ============================================================ +# PAGE 6: Cross-system prototype cosine (29x29) +# ============================================================ + +def page_cross_system_prototypes(): + M = pd.read_csv(DISC / "70_prototype_full_29x29.csv", index_col=0) + fig, ax = plt.subplots(figsize=(13, 12)) + vmax = 1.0 + im = ax.imshow(M.values, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="equal") + ax.set_xticks(range(len(M.columns))) + ax.set_xticklabels(M.columns, rotation=90, fontsize=10) + ax.set_yticks(range(len(M.index))) + ax.set_yticklabels(M.index, fontsize=10) + sizes = [] + prev = None; count = 0 + for lab in M.index: + s = lab.split(":")[0] + if prev is None: + prev = s; count = 1 + elif s == prev: + count += 1 + else: + sizes.append(count); prev = s; count = 1 + sizes.append(count) + cum = 0 + for s in sizes[:-1]: + cum += s + ax.axhline(cum - 0.5, color="k", linewidth=1.2) + ax.axvline(cum - 0.5, color="k", linewidth=1.2) + plt.colorbar(im, ax=ax, shrink=0.7, label="prototype cosine") + ax.set_title("§8.2 Cross-system prototype cosine (29 × 29)\n" + "Off-block cos mean = +0.001 (null baseline: +0.005). No shared cell-type semantics.", + fontsize=18) + plt.tight_layout() + plt.savefig(FIG_S / "06_cross_system_prototypes.pdf", bbox_inches="tight") + plt.close() + print("[fig] 06_cross_system_prototypes.pdf") + + +# ============================================================ +# PAGE 7: Attribution heatmap: top-10 genes per class × system +# ============================================================ + +def page_attribution_heatmap(): + """per-system heatmap: rows=classes, cols=top-5-per-class gene union. + + Uses the aggregated per-class CSV (80_{sys}_gene_attribution.csv) which + lists top-N positive/negative genes with their attributions. Robust to + npy/hvgs shape mismatch (older discovery runs sometimes use different HVG + bases). + """ + fig, axes = plt.subplots(3, 1, figsize=(18, 19)) + for ax, sys in zip(axes, SYSTEMS): + df = pd.read_csv(DISC / f"80_{sys}_gene_attribution.csv") + classes = df["class"].tolist() + # Union of top-5 positive genes per class (each row lists top-20). + top_gene_set = [] + seen = set() + per_class = {} + for _, r in df.iterrows(): + pos_genes = str(r["top_pos_genes"]).split(",") + pos_vals = [float(v) for v in str(r["top_pos_attribution"]).split(",")] + neg_genes = str(r["top_neg_genes"]).split(",") + neg_vals = [float(v) for v in str(r["top_neg_attribution"]).split(",")] + per_class[r["class"]] = dict(zip(pos_genes + neg_genes, + pos_vals + neg_vals)) + for g in pos_genes[:5]: + if g not in seen: seen.add(g); top_gene_set.append(g) + top_genes = top_gene_set + M = np.zeros((len(classes), len(top_genes)), dtype=np.float32) + for i, cls in enumerate(classes): + d = per_class[cls] + for j, g in enumerate(top_genes): + M[i, j] = d.get(g, 0.0) + + vmax = float(np.abs(M).max()) if M.size else 1.0 + im = ax.imshow(M, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto") + ax.set_xticks(range(len(top_genes))) + ax.set_xticklabels(top_genes, rotation=75, fontsize=9, ha="right") + ax.set_yticks(range(len(classes))) + ax.set_yticklabels(classes, fontsize=11) + ax.set_title(f"{SYS_LABEL[sys]} — top-5 attributed genes per class (union)", fontsize=15) + plt.colorbar(im, ax=ax, shrink=0.7, label="attribution (unit)") + plt.suptitle("§9.1 Prototype-gene integrated-gradient attribution heatmap", + fontsize=20, y=1.005, fontweight="bold") + plt.tight_layout() + plt.savefig(FIG_S / "07_attribution_heatmap.pdf", bbox_inches="tight") + plt.close() + print("[fig] 07_attribution_heatmap.pdf") + + +# ============================================================ +# PAGE 8: TF enrichment heatmap per system +# ============================================================ + +def page_tf_enrichment(): + fig, axes = plt.subplots(1, 3, figsize=(24, 7.5)) + for ax, sys in zip(axes, SYSTEMS): + df = pd.read_csv(DISC / f"80_{sys}_tf_enrichment.csv") + piv = df.pivot(index="class", columns="program", values="sum_attribution").fillna(0) + vmax = np.max(np.abs(piv.values)) + im = ax.imshow(piv.values, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto") + for i, cls in enumerate(piv.index): + top_j = int(np.argmax(piv.values[i])) + ax.text(top_j, i, f"{piv.values[i, top_j]:+.2f}", + ha="center", va="center", fontsize=8, + color="white" if abs(piv.values[i, top_j]) > vmax * 0.5 else "black") + ax.set_xticks(range(len(piv.columns))) + ax.set_xticklabels(piv.columns, rotation=75, fontsize=9, ha="right") + ax.set_yticks(range(len(piv.index))) + ax.set_yticklabels(piv.index, fontsize=11) + ax.set_title(f"{SYS_LABEL[sys]}\nsum_attribution over TF-program members", fontsize=15) + plt.colorbar(im, ax=ax, shrink=0.7) + plt.suptitle("§9.2 Per-prototype TF-program enrichment (self-hit = validation; cross-hit = confusability)", + fontsize=20, y=1.02, fontweight="bold") + plt.tight_layout() + plt.savefig(FIG_S / "08_tf_enrichment.pdf", bbox_inches="tight") + plt.close() + print("[fig] 08_tf_enrichment.pdf") + + +# ============================================================ +# PAGE 9: Counterfactual KO top-essential genes +# ============================================================ + +def page_ko_essentials(): + fig, axes = plt.subplots(3, 1, figsize=(18, 16)) + for ax, sys in zip(axes, SYSTEMS): + df = pd.read_csv(DISC / f"81_{sys}_ko_essentials.csv") + rows = [] + for _, r in df.iterrows(): + genes = r["top_essential_genes"].split(",")[:3] + deltas = [float(x) for x in r["top_essential_deltas"].split(",")[:3]] + for g, d in zip(genes, deltas): + rows.append({"class": r["class"], "gene": g, "delta": d, + "baseline_cos": r["baseline_cos"]}) + d2 = pd.DataFrame(rows) + piv = d2.pivot(index="class", columns="gene", values="delta").fillna(0) + vmax = np.abs(piv.values).max() + im = ax.imshow(piv.values, cmap="Reds", vmin=0, vmax=vmax, aspect="auto") + for i in range(piv.shape[0]): + for j in range(piv.shape[1]): + v = piv.values[i, j] + if v > 0: + ax.text(j, i, f"{v:.3f}", ha="center", va="center", fontsize=8, + color="white" if v > vmax * 0.5 else "black") + ax.set_xticks(range(len(piv.columns))) + ax.set_xticklabels(piv.columns, rotation=75, fontsize=9, ha="right") + ax.set_yticks(range(len(piv.index))) + ax.set_yticklabels(piv.index, fontsize=11) + ax.set_title(f"{SYS_LABEL[sys]} — counterfactual KO Δcos (top-3 essentials per class)", fontsize=15) + plt.colorbar(im, ax=ax, shrink=0.7, label="Δ prototype cos") + plt.suptitle("§9.3 Counterfactual single-gene knockout: essentiality per class", + fontsize=20, y=1.005, fontweight="bold") + plt.tight_layout() + plt.savefig(FIG_S / "09_ko_essentials.pdf", bbox_inches="tight") + plt.close() + print("[fig] 09_ko_essentials.pdf") + + +# ============================================================ +# PAGE 10: Hessian gene-gene interaction pairs +# ============================================================ + +def page_hessian_pairs(): + fig, axes = plt.subplots(1, 3, figsize=(24, 8)) + for ax, sys in zip(axes, SYSTEMS): + df = pd.read_csv(DISC / f"85_{sys}_hessian_pairs.csv") + top = df.groupby("class").head(8).copy() + top["pair"] = top["gene_a"] + "·" + top["gene_b"] + piv = top.pivot(index="class", columns="pair", values="hessian_off_diag").fillna(0) + vmax = np.abs(piv.values).max() + im = ax.imshow(piv.values, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto") + ax.set_xticks(range(len(piv.columns))) + ax.set_xticklabels(piv.columns, rotation=90, fontsize=8) + ax.set_yticks(range(len(piv.index))) + ax.set_yticklabels(piv.index, fontsize=11) + ax.set_title(f"{SYS_LABEL[sys]}\ntop-8 gene-gene Hessian pairs per class", fontsize=15) + plt.colorbar(im, ax=ax, shrink=0.7, label="∂²s/∂g·∂g'") + plt.suptitle("§10.3 Gene-gene interaction Hessian: combinatorial identity rules", + fontsize=20, y=1.02, fontweight="bold") + plt.tight_layout() + plt.savefig(FIG_S / "10_hessian_pairs.pdf", bbox_inches="tight") + plt.close() + print("[fig] 10_hessian_pairs.pdf") + + +# ============================================================ +# PAGE 11-13: three system UMAPs — reuse existing multi UMAP + add rebuilt Dingwall by class +# ============================================================ + +def page_umaps_all_targets(): + # existing PDFs are appended in the PyPDF merge step + print("[skip] UMAPs handled by PDF merge step") + + +# ============================================================ +# PAGE 14: Novel populations discovered per system +# ============================================================ + +def page_novel_populations(): + import os + if not os.path.exists('/home/bcheng/PRISM/discovery/71_dingwall_novel_populations.csv'): + print('[skip] 71_dingwall not found'); return + _dummy = None + fig, axes = plt.subplots(1, 2, figsize=(20, 7)) + ax = axes[0] + a = pd.read_csv(DISC / "71_dingwall_novel_populations.csv") + a = a.sort_values("n_cells", ascending=False) + y = np.arange(len(a)) + ax.barh(y, a["n_cells"], color="#2b83ba", edgecolor="k", linewidth=0.5) + labels = [f"cluster {c}: {m[:35]}..." if len(m) > 35 else f"cluster {c}: {m}" + for c, m in zip(a["cluster"], a["top_markers"])] + ax.set_yticks(y); ax.set_yticklabels(labels, fontsize=11) + for i, n in enumerate(a["n_cells"]): + ax.text(n + 3, i, f"n={int(n)}", va="center", fontsize=11) + ax.set_xlabel("cells in abstain cluster") + ax.set_title(f"§8.4 Dingwall abstain-gate clusters (cos<0.5)\nn={int(a['n_cells'].sum())} cells total", + fontsize=15) + ax.invert_yaxis() + + ax = axes[1] + d = pd.read_csv(DISC / "73_dahlin_novel_populations.csv") + d = d.sort_values("n_cells", ascending=False) + y = np.arange(len(d)) + colors = ["#d7191c" if wtf > 0.85 else "#fdae61" if wtf > 0.6 else "#2b83ba" + for wtf in d["genotype_wt_frac"]] + ax.barh(y, d["n_cells"], color=colors, edgecolor="k", linewidth=0.5) + labels = [f"c{c}: {m[:35]}..." if len(m) > 35 else f"c{c}: {m}" + for c, m in zip(d["cluster"], d["top_markers"])] + ax.set_yticks(y); ax.set_yticklabels(labels, fontsize=11) + for i, (n, wtf) in enumerate(zip(d["n_cells"], d["genotype_wt_frac"])): + ax.text(n + 8, i, f"n={int(n)} · WT={wtf:.1%}", va="center", fontsize=11) + ax.set_xlabel("cells in abstain cluster") + ax.set_title(f"§8.5 Dahlin abstain-gate clusters (cos<0.57, n={int(d['n_cells'].sum())})\n" + "Red bar = quiescent LT-HSC (Hlf⁺), 90.5% WT (Kit-W41-depleted)", + fontsize=15) + ax.invert_yaxis() + + plt.suptitle("§8.4-5 Abstain-gate novel populations", fontsize=20, y=1.02, fontweight="bold") + plt.tight_layout() + plt.savefig(FIG_S / "11_novel_populations.pdf", bbox_inches="tight") + plt.close() + print("[fig] 11_novel_populations.pdf") + + +# ============================================================ +# PAGE 15: Co-attribution modules summary +# ============================================================ + +def page_coatt_modules(): + fig, axes = plt.subplots(1, 3, figsize=(23, 7)) + for ax, sys in zip(axes, SYSTEMS): + df = pd.read_csv(DISC / f"82_{sys}_coatt_modules.csv") + df = df.sort_values("dom_class_mean_att", ascending=False).head(10) + y = np.arange(len(df)) + colors = [SYS_COLOR[sys]] * len(df) + ax.barh(y, df["dom_class_mean_att"], color=colors, edgecolor="k", linewidth=0.5) + labels = [f"module {m} ({s}g)→{c}" + for m, s, c in zip(df["module_id"], df["size"], df["dominant_class"])] + ax.set_yticks(y); ax.set_yticklabels(labels, fontsize=11) + for i, v in enumerate(df["dom_class_mean_att"]): + ax.text(v + 0.001, i, f"{v:.3f}", va="center", fontsize=11) + ax.set_xlabel("mean attribution to dominant class") + ax.set_title(f"{SYS_LABEL[sys]}: top-10 gene co-attribution modules", fontsize=15) + ax.invert_yaxis() + plt.suptitle("§9.4 Gene-gene co-attribution modules per system", + fontsize=20, y=1.02, fontweight="bold") + plt.tight_layout() + plt.savefig(FIG_S / "12_coatt_modules.pdf", bbox_inches="tight") + plt.close() + print("[fig] 12_coatt_modules.pdf") + + +# ============================================================ +# Master build +# ============================================================ + +def page_anchor_delta_recall(): + """anchor-trick effect on OOD-class recall across 3 systems.""" + fig, axes = plt.subplots(1, 2, figsize=(18, 7.0)) + + ax = axes[0] + # (system, variant, ood_class, baseline_recall, anchor_recall) + rows = [ + ("Pancreas\nBaron adult-\nislet anchor", "PCA", 0.000, 0.805, SYS_COLOR["pancreas"], 0.55), + ("Pancreas\nBaron adult-\nislet anchor", "Marker",0.000, 0.805, SYS_COLOR["pancreas"], 0.90), + ("HSC\nNestorowa LT-\nHSC anchor", "PCA", 0.000, 0.6818,SYS_COLOR["hematopoiesis"], 0.55), + ("HSC\nNestorowa LT-\nHSC anchor", "Marker",0.000, 0.5152,SYS_COLOR["hematopoiesis"], 0.90), + ("Skin\nSulic HF-\nplacode anchor", "PCA", 0.000, 0.9802,SYS_COLOR["pan_skin"], 0.55), + ("Skin\nSulic HF-\nplacode anchor", "Marker",0.000, 0.8293,SYS_COLOR["pan_skin"], 0.90), + ] + x = np.arange(len(rows)) + for i, (grp, var, base, anch, col, alpha) in enumerate(rows): + ax.bar(x[i], anch, color=col, alpha=alpha, edgecolor="k", linewidth=0.7, + label=var if i < 2 else None) + ax.plot([x[i]-0.35, x[i]+0.35], [base, base], color="red", linewidth=1.5, linestyle="--", + label="no-anchor baseline" if i == 0 else None) + ax.annotate("", xy=(x[i], anch), xytext=(x[i], base), + arrowprops={"arrowstyle": "->", "color": "black", "lw": 1.2}) + ax.text(x[i], anch + 0.02, f"{anch:.2f}\n(Δ +{anch-base:.2f})", + ha="center", fontsize=10) + xt = [rows[i][0] + f"\n({rows[i][1]})" for i in range(len(rows))] + ax.set_xticks(x); ax.set_xticklabels(xt, fontsize=10) + ax.set_ylim(0, 1.1) + ax.set_ylabel("OOD-class recall") + ax.set_title("(a) Anchor closes vocabulary-out-of-domain gap\n" + "baseline recall 0 (dashed red) → anchor recall (bar)", fontsize=15) + ax.legend(loc="upper left", framealpha=0.9) + + ax = axes[1] + baron_pca = json.load(open(DISC / "pancreas/pca/93_baron_merged_metrics.json")) + baron_mar = json.load(open(DISC / "pancreas/marker/93_baron_merged_metrics.json")) + nest_pca = json.load(open(DISC / "hematopoiesis/pca/97_nestorowa_anchor_zero_shot.json")) + nest_mar = json.load(open(DISC / "hematopoiesis/marker/97_nestorowa_anchor_zero_shot.json")) + sul_pca = json.load(open(DISC / "pan_skin/pca/97_sulic_anchor_zero_shot.json")) + sul_mar = json.load(open(DISC / "pan_skin/marker/97_sulic_anchor_zero_shot.json")) + + grps = [("Pancreas (Baron)", baron_pca["acc"], baron_mar["acc"], SYS_COLOR["pancreas"]), + ("HSC (Nestorowa)", nest_pca["coarse_acc"], nest_mar["coarse_acc"], SYS_COLOR["hematopoiesis"]), + ("Skin (Sulic)", sul_pca["acc"], sul_mar["acc"], SYS_COLOR["pan_skin"])] + x = np.arange(len(grps)) + width = 0.38 + for i, (lab, pca_v, mar_v, col) in enumerate(grps): + ax.bar(x[i] - width/2, pca_v, width, color=col, alpha=0.55, edgecolor="k", + label="PCA" if i == 0 else None) + ax.bar(x[i] + width/2, mar_v, width, color=col, alpha=0.90, edgecolor="k", + label="Marker" if i == 0 else None) + ax.text(x[i] - width/2, pca_v + 0.015, f"{pca_v:.3f}", ha="center", fontsize=11) + ax.text(x[i] + width/2, mar_v + 0.015, f"{mar_v:.3f}", ha="center", fontsize=11) + ax.set_xticks(x); ax.set_xticklabels([g[0] for g in grps], fontsize=13) + ax.set_ylim(0.4, 1.0) + ax.set_ylabel("Held-out-slice accuracy") + ax.set_title("(b) Anchor-augmented held-out accuracy per system (both variants)", + fontsize=15) + ax.legend(loc="lower right") + + plt.suptitle("Cross-system anchor paradigm: small labeled slice → OOD vocabulary token", + fontsize=20, y=1.02, fontweight="bold") + plt.tight_layout() + plt.savefig(FIG_S / "23_anchor_delta_recall.pdf", bbox_inches="tight") + plt.close() + print("[fig] 23_anchor_delta_recall.pdf") + + +def main(): + page_cv_summary() + page_per_class_f1() + page_prototype_cosine() + # page_training_trajectory() # skipped: v2 checkpoints only save final + page_adversary_purification() + page_cross_system_prototypes() + page_attribution_heatmap() + page_tf_enrichment() + page_ko_essentials() + page_hessian_pairs() + page_novel_populations() + page_coatt_modules() + page_anchor_delta_recall() + + print("\n[merge] merging supplement into PANDA_supplement.pdf") + from pypdf import PdfWriter + w = PdfWriter() + order = [ + FIG_S / "01_cv_summary.pdf", + FIG_S / "02_per_class_f1.pdf", + FIG_S / "03_prototype_cosine.pdf", + FIG_S / "04_training_trajectory.pdf", + FIG_S / "05_adversary_purification.pdf", + FIG_S / "06_cross_system_prototypes.pdf", + FIG_S / "07_attribution_heatmap.pdf", + FIG_S / "08_tf_enrichment.pdf", + FIG_S / "09_ko_essentials.pdf", + FIG_S / "10_hessian_pairs.pdf", + FIG_S / "11_novel_populations.pdf", + FIG_S / "12_coatt_modules.pdf", + FIG_S / "23_anchor_delta_recall.pdf", + FIG / "fig5_dingwall_umap.pdf", + FIG / "fig6_multi_umap.pdf", + ] + for p in order: + if p.exists(): + w.append(str(p)) + print(f" + {p.name}") + else: + print(f" [skip] {p.name} missing") + with open(SUP_PDF, "wb") as f: + w.write(f) + print(f"\nwrote {SUP_PDF} ({SUP_PDF.stat().st_size / 1024:.0f} KB)") + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/build_pca_vs_marker_umaps.py b/scripts/figures/build_pca_vs_marker_umaps.py new file mode 100644 index 0000000000000000000000000000000000000000..81dea5ec6f72f6dd34e3ba97311f53f46485d4a7 --- /dev/null +++ b/scripts/figures/build_pca_vs_marker_umaps.py @@ -0,0 +1,506 @@ +"""PANDA-PCA vs PANDA-Marker side-by-side umaps for dingwall/dahlin/veres, plus en1-cKO enrichment and melanocyte pathway bars.""" +from __future__ import annotations +from pathlib import Path +import warnings, json, sys, pickle, numpy as np, pandas as pd +warnings.filterwarnings("ignore") +import matplotlib; matplotlib.use("Agg") +import matplotlib.pyplot as plt +import anndata as ad, scanpy as sc, scipy.sparse as sp, torch +sys.path.insert(0, "/home/bcheng/PRISM") +sys.path.insert(0, str(Path(__file__).parent)) +from panda import PANDAEncoder +from palette import apply_style, color_for, GENOTYPE_COLORS, STAGE_COLORS +apply_style() +sc.settings.verbosity = 0 + +ROOT = Path("/home/bcheng/PRISM") +FIG_S = ROOT / "figures/supplement" +FIG_S.mkdir(parents=True, exist_ok=True) +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +RS = 42 + + +# ------------------- projection helpers ------------------- +def project(a, sys, variant): + """returns (z_128d, predicted class array, classes).""" + ck = torch.load(ROOT / f"checkpoints/{sys}/{variant}/panda_final.pt", + map_location=DEVICE, weights_only=False) + classes = ck["classes"]; marker_genes = ck.get("marker_genes", []) + stats = np.load(ROOT / f"data/corpus/{sys}/harmonized/corpus_stats.npz", allow_pickle=True) + pca = pickle.load(open(ROOT / f"data/corpus/{sys}/harmonized/pca_basis.pkl", "rb")) + hvgs = [str(g) for g in stats["shared_hvgs"]] + hvg2i = {g: i for i, g in enumerate(hvgs)} + common = [g for g in a.var_names.astype(str) if g in hvg2i] + 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, len(hvgs)), dtype=np.float32) + Xf[:, np.array([hvg2i[g] for g in common])] = X + Xz = np.clip((Xf - stats["mean"].astype(np.float32)) / stats["std"].astype(np.float32), -10, 10) + Xpca = pca.transform(Xz).astype(np.float32) + Xmark = None + if variant == "marker" and marker_genes: + mv = np.zeros((a.n_obs, len(marker_genes)), dtype=np.float32) + for j, g in enumerate(marker_genes): + if g in a.var_names: + col = a[:, g].X + if sp.issparse(col): col = col.toarray() + mv[:, j] = col.flatten().astype(np.float32) + mmu = mv.mean(axis=0, keepdims=True); msig = mv.std(axis=0, keepdims=True) + 1e-6 + Xmark = np.clip((mv - mmu) / msig, -5, 5).astype(np.float32) + model = PANDAEncoder(variant=variant, n_pca=50, + n_markers=len(marker_genes) if variant == "marker" else 0, + n_classes=len(classes), n_sub=3, + n_datasets=len(ck["datasets"])).to(DEVICE).eval() + model.load_state_dict(ck["model"]) + all_z, preds = [], [] + with torch.no_grad(): + for i in range(0, a.n_obs, 4096): + xb = torch.from_numpy(Xpca[i:i+4096]).to(DEVICE) + xmb = torch.from_numpy(Xmark[i:i+4096]).to(DEVICE) if Xmark is not None else None + aux = torch.zeros(len(xb), 2, device=DEVICE) + out = model(xb, aux, x_markers=xmb, lam_dann=0.0) + all_z.append(out["z"].cpu().numpy()) + mc = model.max_sub_cos(out["z"]) + preds.append(mc.argmax(dim=1).cpu().numpy()) + Z = np.concatenate(all_z, axis=0) + P = np.array([classes[i] for i in np.concatenate(preds)]) + return Z, P, classes + + +def do_umap(Z, seed=RS): + import umap + reducer = umap.UMAP(n_neighbors=30, min_dist=0.3, random_state=seed, + n_epochs=200, verbose=False, low_memory=False) + return reducer.fit_transform(Z) + + +# kelly-inspired 22-color palette + "other" +DISTINCT_COLORS = [ + "#e6194b", "#3cb44b", "#4363d8", "#f58231", "#911eb4", "#42d4f4", + "#f032e6", "#bfef45", "#fabed4", "#469990", "#dcbeff", "#9a6324", + "#fffac8", "#800000", "#aaffc3", "#808000", "#ffd8b1", "#000075", + "#a9a9a9", "#f4a460", "#00fa9a", "#ff69b4", +] + + +def build_class_palette(P_pca, P_mar, min_frac=0.005): + """collapse classes < min_frac to 'other', assign each remaining canonical color.""" + from collections import Counter + total = len(P_pca) + len(P_mar) + counts = Counter(P_pca.tolist() + P_mar.tolist()) + kept = [c for c, n in counts.most_common() if n / total >= min_frac] + P_pca_r = np.where(np.isin(P_pca, kept), P_pca, "other") + P_mar_r = np.where(np.isin(P_mar, kept), P_mar, "other") + palette = {c: color_for(c, DISTINCT_COLORS[i % len(DISTINCT_COLORS)]) + for i, c in enumerate(kept)} + palette["other"] = "#e5e5e5" + return P_pca_r, P_mar_r, palette + + +def scatter_side_by_side(emb_pca, emb_mark, colors_pca, colors_mark, palette, + subtitle_pca, subtitle_mark, main_title, out_path, + s=4, alpha=0.55, legend_title=""): + fig, axes = plt.subplots(1, 2, figsize=(18.5, 8.5)) + for ax, emb, colors, sub in zip(axes, + [emb_pca, emb_mark], + [colors_pca, colors_mark], + [subtitle_pca, subtitle_mark]): + for cat in sorted(set(colors)): + m = np.array(colors) == cat + ax.scatter(emb[m, 0], emb[m, 1], s=s, alpha=alpha, + color=palette.get(cat, "#888"), label=cat, linewidths=0, + rasterized=True) + ax.set_xlabel("UMAP-1"); ax.set_ylabel("UMAP-2") + ax.set_title(sub, fontsize=15) + ax.set_xticks([]); ax.set_yticks([]) + handles = [plt.Line2D([0], [0], marker="o", linestyle="", + markerfacecolor=palette[c], markeredgecolor="none", markersize=10, label=c) + for c in palette] + fig.legend(handles=handles, loc="center right", bbox_to_anchor=(1.10, 0.5), + fontsize=12, frameon=False, title=legend_title, title_fontsize=13) + plt.suptitle(main_title, fontsize=20, y=1.02, fontweight="bold") + plt.tight_layout() + plt.savefig(out_path, bbox_inches="tight", dpi=180) + plt.close() + print(f"[fig] {out_path.name}") + + +# ------------------- dingwall ------------------- +def fig_dingwall_pca_vs_marker(): + """dingwall GSE220977 colored by en1 genotype + predicted class.""" + cache = FIG_S / "_cache_dingwall_full.npz" + CKO = {"GSM6833482", "GSM6833483"} + WT = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"} + if cache.exists(): + c = np.load(cache, allow_pickle=True) + emb_pca = c["emb_pca"]; emb_mar = c["emb_mar"] + P_pca = c["P_pca"].astype(str); P_mar = c["P_mar"].astype(str) + genotype = c["genotype"].astype(str) + n = len(genotype) + print(f"[dingwall] loaded cache n={n}", flush=True) + else: + raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad") + genotype = np.where(raw.obs["sample"].astype(str).isin(list(CKO)), "En1-cKO", + np.where(raw.obs["sample"].astype(str).isin(list(WT)), "WT", "other")) + n = raw.n_obs + print(f"[dingwall] projecting all {n} cells with PCA and Marker checkpoints", flush=True) + Z_pca, P_pca, _ = project(raw, "pan_skin", "pca") + Z_mar, P_mar, _ = project(raw, "pan_skin", "marker") + print(f"[dingwall] running umap on all {n}", flush=True) + emb_pca = do_umap(Z_pca) + emb_mar = do_umap(Z_mar) + np.savez(cache, + emb_pca=emb_pca, emb_mar=emb_mar, P_pca=P_pca, P_mar=P_mar, + genotype=genotype) + + palette_gt = {"WT": GENOTYPE_COLORS["WT"], "En1-cKO": GENOTYPE_COLORS["En1-cKO"], + "other": GENOTYPE_COLORS["other"]} + scatter_side_by_side( + emb_pca, emb_mar, genotype, genotype, palette_gt, + f"PANDA-PCA (Dingwall, all {n:,} cells)", f"PANDA-Marker (Dingwall, all {n:,} cells)", + "Dingwall En1-cKO vs WT — PANDA-PCA vs PANDA-Marker embedding", + FIG_S / "24_pca_vs_marker_umaps_dingwall_by_genotype.pdf", + legend_title="Genotype", + ) + P_pca_r, P_mar_r, palette_c = build_class_palette(P_pca, P_mar, min_frac=0.005) + scatter_side_by_side( + emb_pca, emb_mar, P_pca_r, P_mar_r, palette_c, + "PANDA-PCA — predicted class", "PANDA-Marker — predicted class", + f"Dingwall — PANDA-PCA vs PANDA-Marker predicted class map (n={n:,})", + FIG_S / "24b_pca_vs_marker_umaps_dingwall_by_class.pdf", + legend_title="Predicted class", + ) + + +# ------------------- dahlin ------------------- +def _load_dahlin_raw(): + D_DIR = ROOT / "data/corpus/hematopoiesis/held_out_unlabeled/dahlin_extract" + GT = {"SIGAB1":"WT","SIGAC1":"WT","SIGAD1":"WT","SIGAF1":"WT","SIGAG1":"WT", + "SIGAH1":"WT","SIGAG8":"Kit_W41","SIGAH8":"Kit_W41"} + parts = [] + for f in sorted(D_DIR.glob("*.txt.gz")): + sample = f.name.split("_")[1].split(".")[0] + 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"] = GT.get(sample, "unknown") + var = pd.DataFrame(index=df.index.astype(str)) + parts.append(ad.AnnData(X=X, obs=obs, var=var)) + a = ad.concat(parts, join="outer") + import mygene + mg = mygene.MyGeneInfo() + res = mg.querymany(a.var_names.astype(str).tolist(), 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) + a = a[:, keep].copy(); a.var_names = syms[keep]; a.var_names_make_unique() + return a + + +def fig_dahlin_pca_vs_marker(): + cache = FIG_S / "_cache_dahlin_full.npz" + if cache.exists(): + c = np.load(cache, allow_pickle=True) + emb_pca = c["emb_pca"]; emb_mar = c["emb_mar"] + P_pca = c["P_pca"].astype(str); P_mar = c["P_mar"].astype(str) + gen = c["genotype"].astype(str) + n = len(gen) + print(f"[dahlin] loaded cache n={n}", flush=True) + else: + print("[dahlin] loading raw", flush=True) + a = _load_dahlin_raw() + gen = a.obs["genotype"].astype(str).values + n = a.n_obs + print(f"[dahlin] projecting all {n} cells", flush=True) + Z_pca, P_pca, _ = project(a, "hematopoiesis", "pca") + Z_mar, P_mar, _ = project(a, "hematopoiesis", "marker") + print(f"[dahlin] umap all {n}", flush=True) + emb_pca = do_umap(Z_pca) + emb_mar = do_umap(Z_mar) + np.savez(cache, + emb_pca=emb_pca, emb_mar=emb_mar, P_pca=P_pca, P_mar=P_mar, genotype=gen) + + palette_gt = {"WT": GENOTYPE_COLORS["WT"], "Kit_W41": GENOTYPE_COLORS["Kit_W41"], + "unknown": GENOTYPE_COLORS["other"]} + scatter_side_by_side( + emb_pca, emb_mar, gen, gen, palette_gt, + f"PANDA-PCA (Dahlin, all {n:,} cells)", f"PANDA-Marker (Dahlin, all {n:,} cells)", + "Dahlin WT vs Kit-W41 — PANDA-PCA vs PANDA-Marker embedding", + FIG_S / "25_pca_vs_marker_umaps_dahlin_by_genotype.pdf", + legend_title="Genotype", + ) + P_pca_r, P_mar_r, palette_c = build_class_palette(P_pca, P_mar, min_frac=0.005) + scatter_side_by_side( + emb_pca, emb_mar, P_pca_r, P_mar_r, palette_c, + "PANDA-PCA — predicted class", "PANDA-Marker — predicted class", + f"Dahlin — PANDA-PCA vs PANDA-Marker predicted class map (n={n:,})", + FIG_S / "25b_pca_vs_marker_umaps_dahlin_by_class.pdf", + legend_title="Predicted class", + ) + + +# ------------------- veres ------------------- +def _load_veres(): + SHARON_DIR = ROOT / "data/corpus/pancreas/held_out_unlabeled/sharon_extract" + parts = [] + for meta_file in sorted(SHARON_DIR.glob("*.cell_metadata.tsv.gz")): + counts_file = str(meta_file).replace("cell_metadata", "processed_counts") + if not Path(counts_file).exists(): continue + meta = pd.read_csv(meta_file, sep="\t", compression="gzip") + counts = pd.read_csv(counts_file, sep="\t", compression="gzip", index_col=0) + obs = meta.set_index("library.barcode") + obs = obs.loc[obs.index.intersection(counts.index)] + counts_al = counts.loc[obs.index] + X = sp.csr_matrix(counts_al.values.astype(np.float32)) + a = ad.AnnData(X=X, obs=obs, var=pd.DataFrame(index=counts_al.columns)) + a.var_names_make_unique() + parts.append(a) + return ad.concat(parts, join="outer") + + +def fig_veres_pca_vs_marker(): + cache = FIG_S / "_cache_veres_full.npz" + if cache.exists(): + c = np.load(cache, allow_pickle=True) + emb_pca = c["emb_pca"]; emb_mar = c["emb_mar"] + P_pca = c["P_pca"].astype(str); P_mar = c["P_mar"].astype(str) + st_str = c["stage"].astype(str) + n = len(st_str) + print(f"[veres] loaded cache n={n}", flush=True) + else: + print("[veres] loading raw", flush=True) + a = _load_veres() + stage_col = "Stage" if "Stage" in a.obs.columns else "stage" + stage = pd.to_numeric(a.obs[stage_col], errors="coerce").fillna(-1).astype(int).values + st_str = np.array([str(s) for s in stage]) + n = a.n_obs + print(f"[veres] projecting all {n} cells (stage dist: " + f"{pd.Series(st_str).value_counts().to_dict()})", flush=True) + Z_pca, P_pca, _ = project(a, "pancreas", "pca") + Z_mar, P_mar, _ = project(a, "pancreas", "marker") + print(f"[veres] umap all {n}", flush=True) + emb_pca = do_umap(Z_pca) + emb_mar = do_umap(Z_mar) + np.savez(cache, + emb_pca=emb_pca, emb_mar=emb_mar, P_pca=P_pca, P_mar=P_mar, stage=st_str) + + # canonical stage palette from palette.py + palette_st = {**STAGE_COLORS, "-1": "#dddddd"} + scatter_side_by_side( + emb_pca, emb_mar, st_str, st_str, palette_st, + f"PANDA-PCA (Veres, all {n:,} cells)", f"PANDA-Marker (Veres, all {n:,} cells)", + "Veres SC-beta differentiation Stage 3-6 — PANDA-PCA vs PANDA-Marker embedding", + FIG_S / "26_pca_vs_marker_umaps_veres_by_stage.pdf", + legend_title="Stage", + ) + P_pca_r, P_mar_r, palette_c = build_class_palette(P_pca, P_mar, min_frac=0.005) + # Report actual class diversity in the title so readers understand why the + # Veres in-vitro slice collapses onto a small subset of pancreas prototypes. + from collections import Counter + top_pca = ", ".join([f"{c} (n={k:,})" + for c, k in Counter(P_pca.tolist()).most_common(5)]) + top_mar = ", ".join([f"{c} (n={k:,})" + for c, k in Counter(P_mar.tolist()).most_common(5)]) + n_pca_cls = len(set(P_pca.tolist())) + n_mar_cls = len(set(P_mar.tolist())) + subtitle = ( + f"Veres in-vitro (n={n:,}) — Marker predicts {n_mar_cls} class(es) " + f"[top-5 shown: {top_mar}]; PCA predicts {n_pca_cls} class(es) [top-5: {top_pca}].\n" + f"In-vitro batch shift → prototype collapse (see §7 Veres discussion)." + ) + scatter_side_by_side( + emb_pca, emb_mar, P_pca_r, P_mar_r, palette_c, + "PANDA-PCA — predicted class", "PANDA-Marker — predicted class", + subtitle, + FIG_S / "26b_pca_vs_marker_umaps_veres_by_class.pdf", + legend_title="Predicted class", + ) + + +# ------------------- en1-cKO enrichment bars ------------------- +def fig_en1_enrichment(): + raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad") + CKO = {"GSM6833482", "GSM6833483"} + WT = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"} + genotype = np.where(raw.obs["sample"].astype(str).isin(list(CKO)), "En1-cKO", + np.where(raw.obs["sample"].astype(str).isin(list(WT)), "WT", "other")) + pred = pd.read_csv(ROOT / "discovery/pan_skin/marker/dingwall_predictions.csv") + common = raw.obs_names.intersection(pd.Index(pred["cell_id"].astype(str))) + keep = raw.obs_names.isin(common) + raw = raw[keep].copy() + gt = genotype[keep] + pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"])) + labels = np.array([pred_map.get(c, "unknown") for c in raw.obs_names]) + labeled = (gt != "other") + baseline = (gt[labeled] == "En1-cKO").sum() / labeled.sum() + df = pd.DataFrame({"pred": labels, "gt": gt, "labeled": labeled}) + dfl = df[df["labeled"]] + rows = [] + for cls, sub in dfl.groupby("pred"): + n = len(sub) + if n < 50: continue + frac_cko = (sub["gt"] == "En1-cKO").sum() / n + rows.append({"class": cls, "n": n, "frac_cko": frac_cko, + "delta": frac_cko - baseline}) + d = pd.DataFrame(rows).sort_values("delta", ascending=False) + d.to_csv(FIG_S / "27_dingwall_en1_enrichment.csv", index=False) + + fig, ax = plt.subplots(figsize=(14, 8)) + y = np.arange(len(d)) + cols = [GENOTYPE_COLORS["En1-cKO"] if r["delta"] > 0.05 + else GENOTYPE_COLORS["WT"] if r["delta"] < -0.05 + else "#888888" + for _, r in d.iterrows()] + deltas = (d["frac_cko"] - baseline).values + ax.barh(y, deltas, color=cols, edgecolor="k", linewidth=0.5) + ax.axvline(0, color="black", linewidth=1) + # place value labels always to the right of the bar tip with an offset in + # display coords so short/negative bars can't crash into the y-tick labels + for i, row in enumerate(d.itertuples()): + delta_i = row.frac_cko - baseline + ax.annotate(f"{row.frac_cko:.2f} (n={int(row.n):,})", + xy=(delta_i, i), xycoords="data", + xytext=(6, 0), textcoords="offset points", + ha="left", va="center", fontsize=11, clip_on=False) + ax.set_yticks(y); ax.set_yticklabels(d["class"], fontsize=12) + ax.tick_params(axis="y", pad=6) + # add headroom on the right so annotations don't clip + dmin, dmax = float(deltas.min()), float(deltas.max()) + span = max(abs(dmin), abs(dmax)) + ax.set_xlim(dmin - 0.05 * span, dmax + 0.55 * span) + ax.set_xlabel(f"Δ En1-cKO fraction vs baseline {baseline:.2f}") + ax.invert_yaxis() + ax.set_title("Dingwall — En1-cKO enrichment per PANDA-predicted class\n" + "(red = cKO-enriched; blue = WT-enriched; grey = at baseline)", + fontsize=18) + plt.tight_layout() + plt.savefig(FIG_S / "27_dingwall_en1_enrichment.pdf", bbox_inches="tight") + plt.close() + print(f"[fig] 27_dingwall_en1_enrichment.pdf ({len(d)} classes shown)") + + +# ------------------- melanocyte pathway modules ------------------- +MELANOCYTE_MODULES = { + "MITF regulon (up in WT, down in cKO)": ["Mitf", "Dct", "Tyrp1", "Tyr", "Pmel", "Mlana", "Slc24a5", "Sox10"], + "keratinocyte contamination / Krt (up in cKO)": ["Krt5", "Krt14", "Krt15"], + "melanocyte proliferation / migration (down in cKO)": ["Ets1", "Kit", "Pax3", "Sox9"], + "pigment biogenesis (down in cKO)": ["Gpnmb", "Slc45a2", "Oca2", "Trpm1"], +} + + +def fig_melanocyte_pathway(): + raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad") + CKO = {"GSM6833482", "GSM6833483"} + WT = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"} + genotype = np.where(raw.obs["sample"].astype(str).isin(list(CKO)), "En1-cKO", + np.where(raw.obs["sample"].astype(str).isin(list(WT)), "WT", "other")) + pred = pd.read_csv(ROOT / "discovery/pan_skin/marker/dingwall_predictions.csv") + pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"])) + labels = np.array([pred_map.get(c, "unknown") for c in raw.obs_names]) + mel_mask = (labels == "melanocyte") & (genotype != "other") + a_mel = raw[mel_mask].copy() + gt_mel = genotype[mel_mask] + print(f"[melanocyte] {a_mel.n_obs} cells (WT {(gt_mel=='WT').sum()} + cKO {(gt_mel=='En1-cKO').sum()})", flush=True) + + sc.pp.normalize_total(a_mel, target_sum=1e4); sc.pp.log1p(a_mel) + rows = [] + for mod_name, genes in MELANOCYTE_MODULES.items(): + present = [g for g in genes if g in a_mel.var_names] + if not present: continue + sc.tl.score_genes(a_mel, gene_list=present, score_name="s_tmp", use_raw=False) + s = a_mel.obs["s_tmp"].values + wt_mean = s[gt_mel == "WT"].mean() + cko_mean = s[gt_mel == "En1-cKO"].mean() + from scipy.stats import mannwhitneyu + _, p = mannwhitneyu(s[gt_mel == "WT"], s[gt_mel == "En1-cKO"], alternative="two-sided") + rows.append({"module": mod_name, "genes": ", ".join(present), + "wt_mean": wt_mean, "cko_mean": cko_mean, + "delta": cko_mean - wt_mean, "p": p}) + d = pd.DataFrame(rows) + d.to_csv(FIG_S / "28_melanocyte_pathway_modules.csv", index=False) + + fig, ax = plt.subplots(figsize=(14.5, 7.5)) + x = np.arange(len(d)) + width = 0.35 + ax.bar(x - width/2, d["wt_mean"], width, label="WT", + color=GENOTYPE_COLORS["WT"], edgecolor="k") + ax.bar(x + width/2, d["cko_mean"], width, label="En1-cKO", + color=GENOTYPE_COLORS["En1-cKO"], edgecolor="k") + # find headroom so the p-labels never collide with the suptitle + y_max_data = max(d["wt_mean"].max(), d["cko_mean"].max()) + y_min_data = min(0.0, d["wt_mean"].min(), d["cko_mean"].min()) + span = y_max_data - y_min_data + for i, r in enumerate(d.itertuples()): + y_top = max(r.wt_mean, r.cko_mean) + 0.03 * span + sig = "***" if r.p < 1e-3 else "**" if r.p < 1e-2 else "*" if r.p < 5e-2 else "n.s." + ax.text(i, y_top, f"p={r.p:.1e} {sig}", ha="center", fontsize=11) + # explicit y-axis room above the tallest p-label + ax.set_ylim(y_min_data - 0.05 * span, y_max_data + 0.18 * span) + ax.set_xticks(x) + ax.set_xticklabels([m.split(" (")[0] for m in d["module"]], fontsize=11, rotation=15, ha="right") + ax.set_ylabel("Module score (mean per cell)") + ax.set_title("Melanocyte pathway modules — WT vs En1-cKO (Dingwall predicted melanocytes)", + fontsize=18, pad=14) + ax.legend(loc="center left", bbox_to_anchor=(1.02, 0.5), frameon=False, title="Genotype") + ax.axhline(0, color="black", linewidth=0.5, linestyle="--") + plt.subplots_adjust(top=0.85) + plt.tight_layout() + plt.savefig(FIG_S / "28_melanocyte_pathway_modules.pdf", bbox_inches="tight") + plt.close() + print(f"[fig] 28_melanocyte_pathway_modules.pdf") + + +def merge_pdf(): + """merge all supplement pages into figures/PANDA_supplement.pdf.""" + from pypdf import PdfWriter + w = PdfWriter() + order = [ + FIG_S / "01_cv_summary.pdf", + FIG_S / "02_per_class_f1.pdf", + FIG_S / "03_prototype_cosine.pdf", + FIG_S / "04_training_trajectory.pdf", + FIG_S / "05_adversary_purification.pdf", + FIG_S / "06_cross_system_prototypes.pdf", + FIG_S / "07_attribution_heatmap.pdf", + FIG_S / "08_tf_enrichment.pdf", + FIG_S / "09_ko_essentials.pdf", + FIG_S / "10_hessian_pairs.pdf", + FIG_S / "11_novel_populations.pdf", + FIG_S / "12_coatt_modules.pdf", + FIG_S / "23_anchor_delta_recall.pdf", + FIG_S / "24_pca_vs_marker_umaps_dingwall_by_genotype.pdf", + FIG_S / "24b_pca_vs_marker_umaps_dingwall_by_class.pdf", + FIG_S / "25_pca_vs_marker_umaps_dahlin_by_genotype.pdf", + FIG_S / "25b_pca_vs_marker_umaps_dahlin_by_class.pdf", + FIG_S / "26_pca_vs_marker_umaps_veres_by_stage.pdf", + FIG_S / "26b_pca_vs_marker_umaps_veres_by_class.pdf", + FIG_S / "27_dingwall_en1_enrichment.pdf", + FIG_S / "28_melanocyte_pathway_modules.pdf", + ROOT / "figures/fig5_dingwall_umap.pdf", + ROOT / "figures/fig6_multi_umap.pdf", + ] + for p in order: + if p.exists(): + w.append(str(p)) + print(f" + {p.name}") + else: + print(f" [skip] {p.name} missing") + out = ROOT / "figures/PANDA_supplement.pdf" + with open(out, "wb") as f: + w.write(f) + print(f"\nwrote {out} ({out.stat().st_size / 1024:.0f} KB)") + + +def main(): + fig_dingwall_pca_vs_marker() + fig_dahlin_pca_vs_marker() + fig_veres_pca_vs_marker() + fig_en1_enrichment() + fig_melanocyte_pathway() + merge_pdf() + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/build_skin_corpus_umap.py b/scripts/figures/build_skin_corpus_umap.py new file mode 100644 index 0000000000000000000000000000000000000000..750ef3d9c4605544e6e207dbc134a3568df06de9 --- /dev/null +++ b/scripts/figures/build_skin_corpus_umap.py @@ -0,0 +1,160 @@ +"""pan-skin corpus UMAP: canonical label + source dataset + cos-confidence.""" +from pathlib import Path +import warnings, pickle, sys, json, numpy as np, pandas as pd +warnings.filterwarnings("ignore") + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import anndata as ad +import scipy.sparse as sp +import scanpy as sc +import torch +import umap as _umap + +sys.path.insert(0, "/home/bcheng/PRISM") +from panda import PANDAEncoder + +ROOT = Path("/home/bcheng/PRISM") +FIG_S = ROOT / "figures/supplement"; FIG_S.mkdir(parents=True, exist_ok=True) +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +CLASS_PALETTE = { + "fibroblast-reticular": "#66c2a5", "basal-IFE": "#fc8d62", "endothelial": "#8da0cb", + "immune": "#e78ac3", "fibroblast-papillary": "#a6d854", "melanocyte": "#ffd92f", + "spinous": "#e5c494", "granular": "#b3b3b3", "HF-ORS": "#1b9e77", + "HF-DP": "#d95f02", "HF-placode": "#7570b3", "eccrine-duct": "#666666", + "sebaceous": "#000000", +} + +DATASET_PALETTE = { + "haensel_GSE142471": "#e41a1c", + "joost_annusver_GSE142471":"#e41a1c", # legacy alias + "ge_gupta_GSE131498": "#377eb8", + "wihn_GSE141814": "#4daf4a", + "merkel_GSE201447": "#984ea3", + "sulic_GSE212673": "#ff7f00", + "mca_GSE108097_neonatal": "#ffff33", + "joost_GSE67602": "#a65628", +} + + +def project_corpus(): + stats = np.load(ROOT / "data/corpus/pan_skin/harmonized/corpus_stats.npz", allow_pickle=True) + hvgs = [str(g) for g in stats["shared_hvgs"]] + mu = np.asarray(stats["mean"], dtype=np.float32) + sig = np.asarray(stats["std"], dtype=np.float32) + pca = pickle.load(open(ROOT / "data/corpus/pan_skin/harmonized/pca_basis.pkl", "rb")) + + ck = torch.load(ROOT / "checkpoints/pan_skin/panda_final.pt", + map_location=DEVICE, weights_only=False) + classes = ck["classes"] + model = PANDAEncoder(n_pca=50, n_classes=len(classes), + n_datasets=len(ck["datasets"])).to(DEVICE).eval() + model.load_state_dict(ck["model"]) + protos = torch.from_numpy(ck["prototypes"]).to(DEVICE) + protos = protos / (protos.norm(dim=1, keepdim=True) + 1e-8) + + corp = ad.read_h5ad(ROOT / "data/corpus/pan_skin/harmonized/corpus.h5ad") + print(f"[load] corpus {corp.shape}", flush=True) + print(f"[load] obs cols: {list(corp.obs.columns)[:15]}", flush=True) + + hvg2i = {g: i for i, g in enumerate(hvgs)} + common = [g for g in corp.var_names.astype(str) if g in hvg2i] + a_c = corp[:, 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((corp.n_obs, len(hvgs)), dtype=np.float32) + cols = np.array([hvg2i[g] for g in common]) + Xf[:, cols] = X + Xz = np.clip((Xf - mu) / sig, -10, 10) + Xpca = pca.transform(Xz).astype(np.float32) + print(f"[pca] Xpca {Xpca.shape}", flush=True) + + all_z = [] + with torch.no_grad(): + for i in range(0, corp.n_obs, 8192): + xb = torch.from_numpy(Xpca[i:i+8192]).to(DEVICE) + aux = torch.zeros(len(xb), 2, device=DEVICE) + all_z.append(model(xb, aux, lam_dann=0.0)["z"].cpu().numpy()) + if i % 32768 == 0: print(f"[project] {i}/{corp.n_obs}", flush=True) + Z = np.concatenate(all_z, axis=0) + Zn = Z / (np.linalg.norm(Z, axis=1, keepdims=True) + 1e-8) + cos = Zn @ protos.cpu().numpy().T + max_cos = cos.max(axis=1) + return Z, corp, max_cos + + +def scatter_by_cat(ax, emb, categories, palette, alpha=0.4, s=1.2, legend_title="", legend_fs=7): + cats = list(palette.keys()) + present = [c for c in cats if (categories == c).sum() > 0] + for cat in present: + m = np.asarray(categories) == cat + ax.scatter(emb[m, 0], emb[m, 1], s=s, alpha=alpha, c=palette[cat], + label=f"{cat} (n={int(m.sum()):,})", edgecolors="none") + for cat in sorted(pd.unique(categories)): + if cat in palette or (categories == cat).sum() == 0: continue + m = np.asarray(categories) == cat + ax.scatter(emb[m, 0], emb[m, 1], s=s, alpha=alpha, c="#999", + label=f"{cat} (n={int(m.sum()):,})", edgecolors="none") + ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2") + leg = ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left", + fontsize=legend_fs, markerscale=10, frameon=False, + title=legend_title) + leg.get_title().set_fontsize(9) + ax.tick_params(axis="both", labelsize=8) + + +def main(): + cache = FIG_S / "_cache_skin_corpus.npz" + if cache.exists(): + c = np.load(cache, allow_pickle=True) + emb = c["emb"]; labels = c["labels"] + datasets = c["datasets"]; max_cos = c["max_cos"] + print(f"[cache] loaded emb {emb.shape}", flush=True) + else: + Z, corp, max_cos = project_corpus() + label_key = "canonical_label" if "canonical_label" in corp.obs else "cell_type" + labels = corp.obs[label_key].astype(str).values + dcol = "dataset" if "dataset" in corp.obs.columns else "dataset_id" + datasets = corp.obs[dcol].astype(str).values + print(f"[umap] running UMAP on {Z.shape[0]} cells 128d", flush=True) + emb = _umap.UMAP(n_neighbors=30, min_dist=0.3, random_state=42, + metric="cosine", n_components=2).fit_transform(Z) + np.savez(cache, emb=emb, + labels=np.asarray(labels, dtype=object), + datasets=np.asarray(datasets, dtype=object), + max_cos=max_cos) + print(f"[umap] saved cache", flush=True) + + fig, axes = plt.subplots(1, 3, figsize=(24, 8)) + + scatter_by_cat(axes[0], emb, labels, CLASS_PALETTE, legend_title="canonical label") + axes[0].set_title(f"(a) All {len(emb):,} pan-skin corpus cells\ncoloured by canonical label", fontsize=12) + + scatter_by_cat(axes[1], emb, datasets, DATASET_PALETTE, alpha=0.35, s=1.0, + legend_title="source dataset", legend_fs=8) + axes[1].set_title(f"(b) Same UMAP, coloured by source dataset\n" + f"(indicates residual dataset structure after GRL)", fontsize=12) + + ax = axes[2] + sc_plot = ax.scatter(emb[:, 0], emb[:, 1], c=max_cos, cmap="viridis", vmin=0.4, vmax=1.0, + s=1.0, alpha=0.55, edgecolors="none") + ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2") + ax.tick_params(axis="both", labelsize=8) + plt.colorbar(sc_plot, ax=ax, shrink=0.75, label="max prototype cosine") + ax.set_title("(c) Prototype-cosine confidence per cell\n" + "(low cos → cells that would be abstain-gated in inference)", + fontsize=12) + + plt.suptitle("Pan-skin PANDA training corpus UMAP (all 78,319 cells)", + fontsize=14, y=1.02) + plt.tight_layout() + out = FIG_S / "21_skin_corpus_umap.pdf" + plt.savefig(out, bbox_inches="tight") + plt.close() + print(f"[fig] wrote {out}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/build_umaps_and_discovery.py b/scripts/figures/build_umaps_and_discovery.py new file mode 100644 index 0000000000000000000000000000000000000000..e5a7c374fdd1de42eb105be18d06031b40fefa83 --- /dev/null +++ b/scripts/figures/build_umaps_and_discovery.py @@ -0,0 +1,692 @@ +"""UMAPs + discovery figures, writes figures/supplement/13+ and merges into PANDA_supplement.pdf.""" +from __future__ import annotations +from pathlib import Path +import warnings, json, pickle, sys, numpy as np, pandas as pd +warnings.filterwarnings("ignore") + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +from matplotlib.patches import Patch +import anndata as ad +import scipy.sparse as sp +import scanpy as sc +import torch +import umap as _umap + +sys.path.insert(0, "/home/bcheng/PRISM") +sys.path.insert(0, str(Path(__file__).parent)) +from panda import PANDAEncoder +from palette import apply_style, color_for +apply_style() + +ROOT = Path("/home/bcheng/PRISM") +FIG = ROOT / "figures" +FIG_S = FIG / "supplement"; FIG_S.mkdir(parents=True, exist_ok=True) +DISC = ROOT / "discovery" +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +RS = 42 +SAMPLE_N = 8000 # per panel + + +from palette import CLASS_COLORS as _CLASS_COLORS +# canonical palette, defer to color_for(); keep "other" as light grey +CLASS_PALETTE = dict(_CLASS_COLORS) +CLASS_PALETTE["other"] = "#c8c8c8" + +# classes retired from the canonical class list — filter out of plots +DEPRECATED_CLASSES = {"HF-DP", "eccrine-duct"} + + +def get_projection_from_ckpt(adata_target, sys, shared_hvgs, mu, sig, pca): + ck = torch.load(ROOT / f"checkpoints/{sys}/marker/panda_final.pt", map_location=DEVICE, weights_only=False) + classes = ck["classes"] + marker_genes = ck.get("marker_genes", []) + n_markers = len(marker_genes) + model = PANDAEncoder(variant="marker" if n_markers else "pca", + n_pca=50, n_markers=n_markers, n_sub=3, + n_classes=len(classes), + n_datasets=len(ck["datasets"])).to(DEVICE).eval() + model.load_state_dict(ck["model"]) + protos = ck["prototypes"] + protos = protos / (np.linalg.norm(protos, axis=1, keepdims=True) + 1e-8) + + hvg2i = {g: i for i, g in enumerate(shared_hvgs)} + common = [g for g in adata_target.var_names.astype(str) if g in hvg2i] + a_c = adata_target[:, 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((adata_target.n_obs, len(shared_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 n_markers: + mv = np.zeros((adata_target.n_obs, n_markers), dtype=np.float32) + for j, g in enumerate(marker_genes): + if g in adata_target.var_names: + col = adata_target[:, g].X + if sp.issparse(col): col = col.toarray() + mv[:, j] = col.flatten().astype(np.float32) + mmu = mv.mean(axis=0, keepdims=True); msig = mv.std(axis=0, keepdims=True) + 1e-6 + Xmark = np.clip((mv - mmu) / msig, -5, 5).astype(np.float32) + + all_z = [] + with torch.no_grad(): + for i in range(0, adata_target.n_obs, 4096): + xb = torch.from_numpy(Xpca[i:i+4096]).to(DEVICE) + xmb = torch.from_numpy(Xmark[i:i+4096]).to(DEVICE) if Xmark is not None else None + aux = torch.zeros(len(xb), 2, device=DEVICE) + all_z.append(model(xb, aux, x_markers=xmb, lam_dann=0.0)["z"].cpu().numpy()) + Z = np.concatenate(all_z) + Zn = Z / (np.linalg.norm(Z, axis=1, keepdims=True) + 1e-8) + cos = Zn @ protos.T + pred = np.array([classes[i] for i in cos.argmax(axis=1)]) + max_cos = cos.max(axis=1) + return Z, pred, max_cos, classes + + +def do_umap(Z, seed=RS): + return _umap.UMAP(n_neighbors=30, min_dist=0.3, random_state=seed, + metric="cosine", n_components=2).fit_transform(Z) + + +def scatter_by_cat(ax, emb, categories, palette=None, alpha=0.55, s=3, legend_title=""): + cats = sorted(pd.unique(categories)) + for cat in cats: + m = np.asarray(categories) == cat + color = palette.get(cat, "#999999") if palette else None + ax.scatter(emb[m, 0], emb[m, 1], s=s, alpha=alpha, c=color, + label=f"{cat} (n={int(m.sum())})", edgecolors="none") + ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2") + leg = ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left", + fontsize=7.5, markerscale=6, frameon=False, title=legend_title) + leg.get_title().set_fontsize(9) + ax.tick_params(axis="both", labelsize=8) + + +# ========================================================================= +# Fig 13: Dingwall 2-panel UMAP (predicted class + En1 genotype) +# ========================================================================= + +def fig_dingwall_umap(): + """Dingwall UMAP with PANDA-Marker predicted class + En1 genotype. + + Prefers the existing PCA-vs-Marker cache (_cache_dingwall_full.npz which + already holds emb_mar, P_mar, genotype for all 25,800 cells). Falls back + to the 50_aldrich_projections.h5ad if the cache is missing. + """ + cache = FIG_S / "_cache_dingwall_full.npz" + if cache.exists(): + c = np.load(cache, allow_pickle=True) + emb = np.asarray(c["emb_mar"]) + pred = c["P_mar"].astype(str) + genotype = c["genotype"].astype(str) + rng = np.random.default_rng(RS) + idx = rng.choice(len(emb), size=min(SAMPLE_N, len(emb)), replace=False) + emb, pred, genotype = emb[idx], pred[idx], genotype[idx] + total_n = int(c["genotype"].shape[0]) + else: + p = ad.read_h5ad(ROOT / "discovery/pan_skin/marker/50_aldrich_projections.h5ad") + Z = np.asarray(p.obsm["Z_projection"]) + pred = (p.obs["pred_bbse_label"] if "pred_bbse_label" in p.obs + else p.obs["pred_label"]).astype(str).values + genotype = p.obs["genotype"].astype(str).values + rng = np.random.default_rng(RS) + idx = rng.choice(len(Z), size=min(SAMPLE_N, len(Z)), replace=False) + Z_s = Z[idx] + emb = do_umap(Z_s) + pred, genotype = pred[idx], genotype[idx] + total_n = int(len(Z)) + + # drop deprecated classes from the class panel + keep_c = ~np.isin(pred, list(DEPRECATED_CLASSES)) + emb_c = emb[keep_c]; pred_c = pred[keep_c] + + fig, axes = plt.subplots(1, 2, figsize=(15, 6.2)) + scatter_by_cat(axes[0], emb_c, pred_c, palette=CLASS_PALETTE, + legend_title="predicted class") + axes[0].set_title(f"Dingwall En1-cKO skin (n={total_n:,} total; {len(emb):,} shown)\n" + "PANDA-Marker predicted class", fontsize=11) + gcolors = {"WT": "#2b83ba", "En1-cKO": "#d7191c", + "unknown": "#999999", "other": "#bbbbbb"} + scatter_by_cat(axes[1], emb, genotype, palette=gcolors, alpha=0.4, s=3, + legend_title="En1 genotype") + axes[1].set_title("Dingwall En1-cKO skin\ncoloured by En1 genotype", fontsize=11) + + plt.suptitle("UMAP of PANDA's 128-d projection: Dingwall held-out target", + fontsize=13, y=1.00) + plt.tight_layout() + plt.savefig(FIG_S / "13_dingwall_umap.pdf", bbox_inches="tight") + plt.close() + print("[fig] 13_dingwall_umap.pdf") + + +# ========================================================================= +# Fig 14: Dahlin 3-panel UMAP (predicted class + Kit genotype + LT-HSC cluster) +# ========================================================================= + +def fig_dahlin_umap(): + # placeholder; real impl in fig_dahlin_umap_full below (cache stored only emb+gt) + from scripts.analysis.__init__ import dummy # noqa + return + + +def _load_dahlin_raw(): + from pathlib import Path as _P + D_DIR = _P("/home/bcheng/PRISM/data/corpus/hematopoiesis/held_out_unlabeled/dahlin_extract") + GT = {"SIGAB1":"WT","SIGAC1":"WT","SIGAD1":"WT","SIGAF1":"WT","SIGAG1":"WT", + "SIGAH1":"WT","SIGAG8":"Kit_W41","SIGAH8":"Kit_W41"} + parts = [] + for f in sorted(D_DIR.glob("*.txt.gz")): + sample = f.name.split("_")[1].split(".")[0] + 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"] = GT.get(sample, "unknown") + var = pd.DataFrame(index=df.index.astype(str)) + parts.append(ad.AnnData(X=X, obs=obs, var=var)) + a = ad.concat(parts, join="outer", label="_batch") + import mygene + mg = mygene.MyGeneInfo() + res = mg.querymany(a.var_names.astype(str).tolist(), 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) + a = a[:, keep].copy(); a.var_names = syms[keep]; a.var_names_make_unique() + return a + + +def fig_dahlin_umap_full(): + """Dahlin UMAP — 2-panel (PANDA-Marker predicted class + Kit genotype). + + Prefers pca-vs-marker cache which holds emb_mar/P_mar/genotype for all + 61,122 cells. If a legacy cache with emb/pred/gt/max_cos is present, use + it and render the 3-panel view (with a confidence colorbar). + """ + cache = FIG_S / "_cache_dahlin_full.npz" + if not cache.exists(): + stats = np.load(ROOT / "data/corpus/hematopoiesis/harmonized/corpus_stats.npz", + allow_pickle=True) + shared_hvgs = [str(g) for g in stats["shared_hvgs"]] + pca = pickle.load(open(ROOT / "data/corpus/hematopoiesis/harmonized/pca_basis.pkl", "rb")) + a = _load_dahlin_raw() + print(f"[dahlin] {a.shape}", flush=True) + rng = np.random.default_rng(RS) + idx = rng.choice(a.n_obs, size=min(SAMPLE_N, a.n_obs), replace=False) + a_sub = a[idx].copy() + Z, pred, mc, classes = get_projection_from_ckpt(a_sub, "hematopoiesis", + shared_hvgs, stats["mean"], stats["std"], pca) + emb = do_umap(Z) + gt = a_sub.obs["genotype"].values + np.savez(cache, emb=emb, gt=np.asarray(gt, dtype=object), + pred=np.asarray(pred, dtype=object), max_cos=mc) + _keys = {"emb", "pred", "gt", "max_cos"} + c = np.load(cache, allow_pickle=True) + keys = set(c.files) + if {"emb_mar", "P_mar", "genotype"}.issubset(keys): + emb = np.asarray(c["emb_mar"]) + pred = c["P_mar"].astype(str) + gt = c["genotype"].astype(str) + total_n = len(emb) + mc = None + else: + emb = c["emb"]; gt = c["gt"].astype(str); pred = c["pred"].astype(str) + mc = c["max_cos"] if "max_cos" in keys else None + total_n = len(emb) + + # drop deprecated + keep_c = ~np.isin(pred, list(DEPRECATED_CLASSES)) + emb_c, pred_c = emb[keep_c], pred[keep_c] + + n_panels = 3 if mc is not None else 2 + fig, axes = plt.subplots(1, n_panels, figsize=(7.0 * n_panels, 6.5)) + + scatter_by_cat(axes[0], emb_c, pred_c, palette=CLASS_PALETTE, + legend_title="predicted class") + axes[0].set_title(f"Dahlin Kit-mutant HSPCs (n={total_n:,})\n" + "PANDA-Marker predicted lineage class", fontsize=11) + gcolors = {"WT": "#2b83ba", "Kit_W41": "#d7191c", + "unknown": "#999999", "other": "#bbbbbb"} + scatter_by_cat(axes[1], emb, gt, palette=gcolors, legend_title="Kit genotype") + axes[1].set_title("Dahlin coloured by Kit genotype", fontsize=11) + + if mc is not None: + ax = axes[2] + sc_plot = ax.scatter(emb[:, 0], emb[:, 1], c=mc, cmap="viridis", + vmin=0.4, vmax=1.0, s=3, alpha=0.65, edgecolors="none") + ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2") + ax.tick_params(axis="both", labelsize=8) + plt.colorbar(sc_plot, ax=ax, shrink=0.75, label="max prototype cosine") + ax.set_title("Prototype-cosine confidence\n" + "(low cos → abstain-gate flagged)", fontsize=11) + + plt.suptitle("Dahlin UMAP — PANDA-Marker zero-shot on Kit-W41 (§8.5)", + fontsize=13, y=1.00) + plt.tight_layout() + plt.savefig(FIG_S / "14_dahlin_umap.pdf", bbox_inches="tight") + plt.close() + print("[fig] 14_dahlin_umap.pdf") + + +# ========================================================================= +# Fig 15: Veres 3-panel UMAP (predicted class + stage + confidence) +# ========================================================================= + +def _load_veres_stages(): + SHARON_DIR = ROOT / "data/corpus/pancreas/held_out_unlabeled/sharon_extract" + parts = [] + for meta_file in sorted(SHARON_DIR.glob("*.cell_metadata.tsv.gz")): + counts_file = str(meta_file).replace("cell_metadata", "processed_counts") + if not Path(counts_file).exists(): continue + meta = pd.read_csv(meta_file, sep="\t", compression="gzip") + counts = pd.read_csv(counts_file, sep="\t", compression="gzip", index_col=0) + counts.columns = [c[0].upper() + c[1:].lower() if len(c) > 1 else c + for c in counts.columns.astype(str)] + counts = counts.T.groupby(level=0).sum().T + obs = meta.set_index("library.barcode") + obs = obs.loc[obs.index.intersection(counts.index)] + counts_al = counts.loc[obs.index] + X = sp.csr_matrix(counts_al.values.astype(np.float32)) + obs["dataset"] = "veres" + var = pd.DataFrame({"gene_symbol": counts_al.columns}, index=counts_al.columns) + a = ad.AnnData(X=X, obs=obs, var=var); a.var_names_make_unique() + parts.append(a) + return ad.concat(parts, join="outer", label="_batch") + + +def fig_veres_umap_full(): + cache = FIG_S / "_cache_veres_full.npz" + if cache.exists(): + c = np.load(cache, allow_pickle=True) + emb = c["emb"]; pred = c["pred"]; stage = c["stage"]; mc = c["max_cos"] + else: + stats = np.load(ROOT / "data/corpus/pancreas/harmonized/corpus_stats.npz", + allow_pickle=True) + shared_hvgs = [str(g) for g in stats["shared_hvgs"]] + pca = pickle.load(open(ROOT / "data/corpus/pancreas/harmonized/pca_basis.pkl", "rb")) + a = _load_veres_stages() + stage_num = pd.to_numeric(a.obs["Stage"], errors="coerce") + keep = stage_num.notna().values + a = a[keep].copy(); a.obs["Stage_int"] = stage_num[keep].astype(int).values + rng = np.random.default_rng(RS) + idx = rng.choice(a.n_obs, size=min(SAMPLE_N, a.n_obs), replace=False) + a_sub = a[idx].copy() + Z, pred, mc, classes = get_projection_from_ckpt(a_sub, "pancreas", + shared_hvgs, stats["mean"], stats["std"], pca) + emb = do_umap(Z) + stage = a_sub.obs["Stage_int"].values + np.savez(cache, emb=emb, pred=np.asarray(pred, dtype=object), + stage=stage, max_cos=mc) + + fig, axes = plt.subplots(1, 3, figsize=(21, 6.5)) + scatter_by_cat(axes[0], emb, pred, palette=CLASS_PALETTE, legend_title="predicted class") + axes[0].set_title("Veres hPSC-directed pancreatic differentiation (57,297 total; 8,000 shown)\n" + "PANDA-predicted endocrine class", fontsize=11) + stage_colors = {3: "#fdae61", 4: "#f8b0d1", 5: "#7570b3", 6: "#d7191c"} + scatter_by_cat(axes[1], emb, stage, palette=stage_colors, legend_title="protocol stage") + axes[1].set_title("Coloured by directed-differentiation stage\n" + "(3 → 4 → 5 → 6 = hPSC → SC-β target)", fontsize=11) + + ax = axes[2] + sc_plot = ax.scatter(emb[:, 0], emb[:, 1], c=mc, cmap="viridis", vmin=0.4, vmax=1.0, + s=3, alpha=0.65, edgecolors="none") + ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2") + ax.tick_params(axis="both", labelsize=8) + plt.colorbar(sc_plot, ax=ax, shrink=0.75, label="max prototype cosine") + ax.set_title("Prototype-cosine confidence\n(low cos = zero-shot ambiguity)", fontsize=11) + + plt.suptitle("Veres UMAP — cross-species + cross-platform + in vitro triple shift (§8)", + fontsize=13, y=1.00) + plt.tight_layout() + plt.savefig(FIG_S / "15_veres_umap.pdf", bbox_inches="tight") + plt.close() + print("[fig] 15_veres_umap.pdf") + + +# ========================================================================= +# Fig 16: Dingwall En1-cKO discovery evidence (class enrichment + melanocyte volcano) +# ========================================================================= + +def fig_dingwall_discovery(): + """Two-panel Dingwall En1-cKO discovery evidence. + + Falls back to the compact per-class enrichment CSV (27_*.csv, columns: + class,n,frac_cko,delta) and the pathway module CSV (28_*.csv) when the + older 53_*/56_* CSVs are not present. + """ + fig, axes = plt.subplots(1, 2, figsize=(15, 5.5)) + + ax = axes[0] + p53 = ROOT / "discovery/pan_skin/marker/53_en1_cko_class_enrichment.csv" + if p53.exists(): + enr = pd.read_csv(p53) + cls_col, lfc_col, p_col = "class", "log2_fold_enrich_cKO_vs_WT", "fisher_pvalue" + else: + enr = pd.read_csv(ROOT / "discovery/pan_skin/marker/27_dingwall_en1_enrichment.csv") + cls_col, lfc_col, p_col = "class", "delta", None # delta already signed + # drop deprecated HF-DP class (removed from canonical vocabulary) + enr = enr[~enr[cls_col].isin(["HF-DP"])].copy() + enr = enr.sort_values(lfc_col) + y = np.arange(len(enr)) + colors = ["#d7191c" if v > 0 else "#2b83ba" for v in enr[lfc_col]] + ax.barh(y, enr[lfc_col], color=colors, edgecolor="k", linewidth=0.5) + xmax = max(abs(enr[lfc_col].min()), abs(enr[lfc_col].max())) * 1.3 + for i, (_, r) in enumerate(enr.iterrows()): + if p_col is not None and p_col in r: + pv = r[p_col] + star = " ***" if pv < 1e-3 else " *" if pv < 0.05 else "" + ax.text(xmax, i, f"p={pv:.1e}{star}", + ha="left", va="center", fontsize=8, color="black") + else: + ax.text(xmax, i, f"n={int(r['n']):,}", + ha="left", va="center", fontsize=8, color="black") + ax.set_xlim(-xmax * 1.05, xmax * 1.9) + ax.axvline(0, color="k", linewidth=0.6) + ax.set_yticks(y); ax.set_yticklabels(enr[cls_col], fontsize=9) + xlabel = ("log2 fold-change (cKO/WT)" if p_col is not None + else "Δ En1-cKO fraction vs corpus baseline") + ax.set_xlabel(xlabel, fontsize=10) + ax.set_title("(a) Dingwall class enrichment (cKO vs WT)\n" + "Blue = depleted in cKO, red = enriched", fontsize=10) + ax.grid(axis="x", alpha=0.3, linestyle="--") + + ax = axes[1] + p56 = ROOT / "discovery/pan_skin/marker/56_melanocyte_pathways.csv" + if p56.exists(): + pw = pd.read_csv(p56).sort_values("delta_cKO_minus_WT") + vcol, pcol, ncol = "delta_cKO_minus_WT", "MannU_p", "pathway" + else: + pw = pd.read_csv(ROOT / "discovery/pan_skin/marker/28_melanocyte_pathway_modules.csv") + pw = pw.sort_values("delta") + vcol, pcol, ncol = "delta", "p", "module" + y = np.arange(len(pw)) + colors = ["#d7191c" if d > 0 else "#2b83ba" for d in pw[vcol]] + ax.barh(y, pw[vcol], color=colors, edgecolor="k", linewidth=0.5) + xmax = max(abs(pw[vcol].min()), abs(pw[vcol].max())) * 1.3 + for i, (_, r) in enumerate(pw.iterrows()): + pv = r[pcol] + star = (" ***" if pv < 1e-10 else " **" if pv < 1e-3 else + " *" if pv < 0.05 else "") + ax.text(xmax, i, f"p={pv:.1e}{star}", + ha="left", va="center", fontsize=8) + ax.set_xlim(-xmax * 1.05, xmax * 1.9) + ax.axvline(0, color="k", linewidth=0.6) + labels = [str(s).split(" (")[0] for s in pw[ncol]] + ax.set_yticks(y); ax.set_yticklabels(labels, fontsize=9) + ax.set_xlabel("Δ module score (cKO − WT)", fontsize=10) + ax.set_title("(b) Dingwall melanocyte pathway modules\n" + "Mann-Whitney U within melanocyte class", fontsize=10) + ax.grid(axis="x", alpha=0.3, linestyle="--") + + plt.suptitle("§4 Dingwall En1-cKO mechanistic evidence", fontsize=13, y=1.02) + plt.tight_layout() + plt.savefig(FIG_S / "16_dingwall_discovery.pdf", bbox_inches="tight") + plt.close() + print("[fig] 16_dingwall_discovery.pdf") + + +# ========================================================================= +# Fig 17: Dahlin discovery evidence (LT-HSC depletion + Kit_signaling module) +# ========================================================================= + +def fig_dahlin_discovery(): + fig, axes = plt.subplots(1, 2, figsize=(15, 5.5)) + + ax = axes[0] + d = pd.read_csv(DISC / "73_dahlin_novel_populations.csv") + d = d.sort_values("genotype_wt_frac", ascending=False) + def label(row): + top = row["top_markers"].split(",")[0] + return f"c{row['cluster']}:{top}⁺ (n={row['n_cells']})" + labels = [label(r) for _, r in d.iterrows()] + y = np.arange(len(d)) + colors = ["#d7191c" if wtf > 0.85 else "#fdae61" if wtf > 0.7 else "#2b83ba" + for wtf in d["genotype_wt_frac"]] + ax.barh(y, d["genotype_wt_frac"], color=colors, edgecolor="k", linewidth=0.5) + ax.axvline(0.60, color="green", linestyle="--", linewidth=1.2, label="whole-corpus baseline (~60% WT)") + for i, wtf in enumerate(d["genotype_wt_frac"]): + ax.text(min(wtf + 0.015, 1.05), i, f"{wtf:.1%}", va="center", fontsize=8) + ax.set_yticks(y); ax.set_yticklabels(labels, fontsize=8) + ax.set_xlim(0, 1.15) + ax.set_xlabel("fraction WT") + ax.set_title("(a) Kit-W41 depletes quiescent LT-HSC (Hlf⁺, 90.5% WT)\n" + "Abstain-gate substates ranked by WT fraction", fontsize=10) + ax.legend(fontsize=8, loc="lower right") + ax.grid(axis="x", alpha=0.3, linestyle="--") + + ax = axes[1] + ms = pd.read_csv(ROOT / "discovery/hematopoiesis/marker/67_dahlin_module_scores.csv") + piv = ms.pivot(index="module", columns="class", values="delta_Kit_minus_WT") + piv_p = ms.pivot(index="module", columns="class", values="MannU_p") + row_order = ["Kit_signaling", "MYC_targets", "Integrated_stress", + "Apoptosis_pro", "Apoptosis_anti", "Cell_cycle", "Erythroid_dev"] + row_order = [r for r in row_order if r in piv.index] + col_order = ["MPP", "erythroid", "myeloid", "megakaryocyte", "lymphoid"] + col_order = [c for c in col_order if c in piv.columns] + P = piv.loc[row_order, col_order]; Pp = piv_p.loc[row_order, col_order] + vmax = np.nanmax(np.abs(P.values)) + im = ax.imshow(P.values, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto") + for i in range(P.shape[0]): + for j in range(P.shape[1]): + v = P.values[i, j]; p = Pp.values[i, j] + if np.isnan(v): continue + star = "***" if p < 1e-10 else "**" if p < 1e-3 else "*" if p < 0.05 else "" + ax.text(j, i, f"{v:+.3f}\n{star}", ha="center", va="center", + fontsize=8, color="white" if abs(v) > vmax * 0.55 else "black") + ax.set_xticks(range(len(col_order))); ax.set_xticklabels(col_order, rotation=30, ha="right") + ax.set_yticks(range(len(row_order))); ax.set_yticklabels(row_order) + plt.colorbar(im, ax=ax, label="Δ module score (Kit-W41 − WT)") + ax.set_title("(b) Dahlin within-class module Δ\n" + "Kit_signaling ↓ + ISR ↑ + Apoptosis_pro erythroid ↓ (p=6.6e-123)", fontsize=10) + + plt.suptitle("§6 Dahlin Kit-W41 mechanistic evidence", fontsize=13, y=1.02) + plt.tight_layout() + plt.savefig(FIG_S / "17_dahlin_discovery.pdf", bbox_inches="tight") + plt.close() + print("[fig] 17_dahlin_discovery.pdf") + + +# ========================================================================= +# Fig 18: Veres discovery evidence (stage stack + alpha-vs-beta TF axis) +# ========================================================================= + +def fig_veres_discovery(): + fig, axes = plt.subplots(1, 2, figsize=(15, 5.5)) + + ax = axes[0] + st = pd.read_csv(ROOT / "discovery/pancreas/marker/64_sharon_class_per_stage.csv", index_col=0) + st.columns = st.columns.astype(float).astype(int) + order = ["alpha", "delta", "gamma", "beta", "acinar", "ductal", + "endocrine-progenitor", "endothelial", "other", "immune"] + order = [c for c in order if c in st.index] + st2 = st.loc[order] + bottom = np.zeros(st2.shape[1]) + for cls in order: + vals = st2.loc[cls].values + ax.bar(st2.columns, vals, bottom=bottom, label=cls, + color=CLASS_PALETTE.get(cls, "#999999"), edgecolor="k", linewidth=0.4) + bottom += vals + ax.set_xticks(st2.columns); ax.set_xticklabels([f"Stage {int(s)}" for s in st2.columns]) + ax.set_ylabel("PANDA-predicted class fraction") + a6 = float(st.loc["alpha", 6]); b6 = float(st.loc["beta", 6]) + ax.text(0.98, 0.98, f"Stage 6:\nα = {a6:.1%}\nβ = {b6:.1%}", + transform=ax.transAxes, fontsize=10, ha="right", va="top", + bbox=dict(boxstyle="round", facecolor="white", alpha=0.9)) + ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left", fontsize=8) + ax.set_ylim(0, 1.05) + ax.set_title("(a) Veres SC-β protocol produces SC-α, not SC-β\nInefficient differentiation at Stage 6", fontsize=10) + + ax = axes[1] + de = pd.read_csv(ROOT / "discovery/pancreas/marker/65_sharon_stage6_alpha_vs_beta.csv") + # rows: (up_in, gene, logfc, padj) + for updir, color in [("alpha", "#d7191c"), ("beta", "#2b83ba")]: + sub = de[de["up_in"] == updir] + neg_log10p = -np.log10(np.clip(sub["padj"].values, 1e-320, 1)) + sign = 1 if updir == "alpha" else -1 + ax.scatter(sign * sub["logfc"], neg_log10p, s=15, alpha=0.55, c=color, + label=f"up in SC-{updir}") + top = sub.nsmallest(8, "padj") + for _, r in top.iterrows(): + ax.text(sign * r["logfc"], -np.log10(r["padj"]) + 3, + r["gene"], fontsize=8, ha="center", color=color) + ax.axvline(0, color="k", linewidth=0.5) + ax.set_xlabel("log2 FC (SC-α ← 0 → SC-β)") + ax.set_ylabel("−log10 padj") + ax.set_title("(b) Veres Stage-6 SC-α vs SC-β DE\n" + "Arx/Irx2 vs Nkx6-1/Mnx1/Neurod1 TF axis", fontsize=10) + ax.legend(fontsize=9) + ax.grid(alpha=0.3, linestyle="--") + + plt.suptitle("§7 Veres SC-β / SC-α mechanistic evidence", fontsize=13, y=1.02) + plt.tight_layout() + plt.savefig(FIG_S / "18_veres_discovery.pdf", bbox_inches="tight") + plt.close() + print("[fig] 18_veres_discovery.pdf") + + +# ========================================================================= +# Fig 19: HSC myeloid combinatorial identity network +# ========================================================================= + +def fig_myeloid_network(): + df = pd.read_csv(DISC / "85_hematopoiesis_hessian_pairs.csv") + my = df[df["class"] == "myeloid"].head(20).copy() + + fig, ax = plt.subplots(figsize=(12, 10.5)) + genes = list(pd.unique(pd.concat([my["gene_a"], my["gene_b"]]))) + n = len(genes) + # circular node layout + theta = np.linspace(0, 2 * np.pi, n, endpoint=False) + pos = {g: (np.cos(t), np.sin(t)) for g, t in zip(genes, theta)} + + # edge width ∝ |H| + max_h = my["abs_h"].max() + for _, r in my.iterrows(): + x1, y1 = pos[r["gene_a"]]; x2, y2 = pos[r["gene_b"]] + lw = 4 * r["abs_h"] / max_h + alpha = min(0.85, 0.3 + 0.6 * r["abs_h"] / max_h) + ax.plot([x1, x2], [y1, y2], color="#d7191c", lw=lw, alpha=alpha, zorder=1) + + for g in genes: + x, y = pos[g] + ax.scatter(x, y, s=420, c=color_for("myeloid"), edgecolor="k", linewidth=1, zorder=2) + ax.text(x, y + 0.09, g, ha="center", fontsize=12, zorder=3, + fontweight="bold") + ax.set_xlim(-1.35, 1.35); ax.set_ylim(-1.25, 1.25) + ax.set_aspect("equal"); ax.axis("off") + ax.set_title("§10.3 Pan-hematopoietic myeloid prototype:\n" + "combinatorial identity via macrophage antimicrobial network\n" + "(top-20 Hessian pairs, edge width ∝ |∂²s/∂g·∂g'|)", fontsize=18) + plt.tight_layout() + plt.savefig(FIG_S / "19_myeloid_network.pdf", bbox_inches="tight") + plt.close() + print("[fig] 19_myeloid_network.pdf") + + +# ========================================================================= +# Fig 20: HF placode Wnt/EDAR module co-attribution network +# ========================================================================= + +def fig_placode_wnt_module(): + modules = pd.read_csv(DISC / "82_pan_skin_coatt_modules.csv") + # HF-placode Wnt/EDAR module was module 5 in earlier output + hf_mods = modules[modules["dominant_class"] == "HF-placode"] + if len(hf_mods) == 0: + print("[skip] no HF-placode modules") + return + mod = hf_mods.iloc[0] + genes = mod["member_genes"].split(",")[:20] + + # tight figsize — module contains only 3 genes, no need for 12x10 inch box + fig, ax = plt.subplots(figsize=(6.5, 6.5)) + n = len(genes) + theta = np.linspace(0, 2 * np.pi, n, endpoint=False) + # place small modules further from center + radius = 0.55 if n <= 3 else 1.0 + pos = {g: (radius * np.cos(t), radius * np.sin(t)) for g, t in zip(genes, theta)} + + canonical = {"Ptch2", "Lef1", "Edar", "Wnt6", "Wnt7b", "Bmp7", "Tfap2b", "Tfap2a"} + hf_col = color_for("HF-placode") + for g in genes: + x, y = pos[g] + col = hf_col if g in canonical else "#2b83ba" + ax.scatter(x, y, s=520, c=col, edgecolor="k", linewidth=1, zorder=2) + ax.text(x, y + 0.12, g, ha="center", fontsize=13, zorder=3, + fontweight="bold" if g in canonical else "normal", + color=hf_col if g in canonical else "black") + + # all-pair edges — every member is internally co-attributed + for i, g1 in enumerate(genes): + for g2 in genes[i+1:]: + x1, y1 = pos[g1]; x2, y2 = pos[g2] + ax.plot([x1, x2], [y1, y2], color="#888", lw=0.6, alpha=0.4, zorder=1) + + ax.set_xlim(-1.1, 1.1); ax.set_ylim(-1.1, 1.1) + ax.set_aspect("equal"); ax.axis("off") + genes_str = ", ".join(genes) + ax.set_title(f"§9.4 Pan-skin HF-placode co-attribution triangle\n" + f"module id={int(mod['module_id'])}, size={int(mod['size'])} genes: {genes_str}", + fontsize=13) + plt.tight_layout() + plt.savefig(FIG_S / "20_placode_wnt_module.pdf", bbox_inches="tight") + plt.close() + print("[fig] 20_placode_wnt_module.pdf") + + +# ========================================================================= +# Master merge +# ========================================================================= + +def merge_pdf(): + from pypdf import PdfWriter + SUP_PDF = FIG / "PANDA_supplement.pdf" + order = [ + FIG_S / "01_cv_summary.pdf", + FIG_S / "02_per_class_f1.pdf", + FIG_S / "03_prototype_cosine.pdf", + FIG_S / "04_training_trajectory.pdf", + FIG_S / "05_adversary_purification.pdf", + FIG_S / "06_cross_system_prototypes.pdf", + FIG_S / "07_attribution_heatmap.pdf", + FIG_S / "08_tf_enrichment.pdf", + FIG_S / "09_ko_essentials.pdf", + FIG_S / "10_hessian_pairs.pdf", + FIG_S / "11_novel_populations.pdf", + FIG_S / "12_coatt_modules.pdf", + FIG_S / "13_dingwall_umap.pdf", + FIG_S / "14_dahlin_umap.pdf", + FIG_S / "15_veres_umap.pdf", + FIG_S / "16_dingwall_discovery.pdf", + FIG_S / "17_dahlin_discovery.pdf", + FIG_S / "18_veres_discovery.pdf", + FIG_S / "19_myeloid_network.pdf", + FIG_S / "20_placode_wnt_module.pdf", + ] + w = PdfWriter() + for p in order: + if p.exists(): w.append(str(p)); print(f" + {p.name}") + else: print(f" [skip] {p.name} missing") + with open(SUP_PDF, "wb") as f: w.write(f) + print(f"wrote {SUP_PDF} ({SUP_PDF.stat().st_size/1024:.0f} KB)") + + +def main(): + # UMAP figures first — they cache and are slow + fig_dingwall_umap() + fig_dahlin_umap_full() + fig_veres_umap_full() + fig_dingwall_discovery() + fig_dahlin_discovery() + fig_veres_discovery() + fig_myeloid_network() + fig_placode_wnt_module() + merge_pdf() + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/generate_multi_umap.py b/scripts/figures/generate_multi_umap.py new file mode 100644 index 0000000000000000000000000000000000000000..a597c05849d6a1a6f26a16b47e61892a90b18bb3 --- /dev/null +++ b/scripts/figures/generate_multi_umap.py @@ -0,0 +1,197 @@ +"""3-panel UMAP: dingwall (En1 genotype), dahlin (Kit genotype), veres (stage). 8k cells/panel.""" +from pathlib import Path +import warnings, sys, pickle, json +warnings.filterwarnings("ignore") +import numpy as np, pandas as pd, anndata as ad, torch, scanpy as sc, scipy.sparse as sp +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import umap +from pathlib import Path as _P_root +ROOT = _P_root(__file__).resolve().parents[2] +ROOT_STR = str(ROOT) +sys.path.insert(0, ROOT_STR) +from panda import PANDAEncoder + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +FIG = Path(f"{ROOT_STR}/figures") +FIG.mkdir(exist_ok=True) +RS = 42 +SAMPLE_N = 8000 + + +def get_projection(ckpt_dir, data_a, shared_hvgs, mu, sig, pca): + """128-d PANDA projections for the target AnnData.""" + ck = torch.load(ckpt_dir / "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"]) + protos = ck["prototypes"] + protos = protos / (np.linalg.norm(protos, axis=1, keepdims=True) + 1e-8) + + G = len(shared_hvgs); hvg2i = {g: i for i, g in enumerate(shared_hvgs)} + common = [g for g in data_a.var_names.astype(str) if g in hvg2i] + a_c = data_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((data_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, data_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) + cos = Z @ protos.T + pred = np.array([classes[i] for i in cos.argmax(axis=1)], dtype=object) + return Z, pred + + +def umap_it(Z, seed=RS): + reducer = umap.UMAP(n_neighbors=30, min_dist=0.3, random_state=seed, + metric="cosine", n_components=2) + return reducer.fit_transform(Z) + + +def load_sharon_stages(): + from pathlib import Path as _P + SHARON_DIR = _P(f"{ROOT_STR}/data/corpus/pancreas/held_out_unlabeled/sharon_extract") + parts, stages = [], [] + for meta_file in sorted(SHARON_DIR.glob("*.cell_metadata.tsv.gz")): + counts_file = str(meta_file).replace("cell_metadata", "processed_counts") + if not _P(counts_file).exists(): continue + meta = pd.read_csv(meta_file, sep="\t", compression="gzip") + counts = pd.read_csv(counts_file, sep="\t", compression="gzip", index_col=0) + counts.columns = [c[0].upper() + c[1:].lower() if len(c) > 1 else c + for c in counts.columns.astype(str)] + counts = counts.T.groupby(level=0).sum().T + obs = meta.set_index("library.barcode") + obs = obs.loc[obs.index.intersection(counts.index)] + counts_al = counts.loc[obs.index] + X = sp.csr_matrix(counts_al.values.astype(np.float32)) + obs["dataset"] = "sharon" + var = pd.DataFrame({"gene_symbol": counts_al.columns}, index=counts_al.columns) + a = ad.AnnData(X=X, obs=obs, var=var); a.var_names_make_unique() + parts.append(a) + return ad.concat(parts, join="outer", label="_batch") + + +def load_dahlin(): + from pathlib import Path as _P + D_DIR = _P(f"{ROOT_STR}/data/corpus/hematopoiesis/held_out_unlabeled/dahlin_extract") + GT = {"SIGAB1":"WT","SIGAC1":"WT","SIGAD1":"WT","SIGAF1":"WT","SIGAG1":"WT", + "SIGAH1":"WT","SIGAG8":"Kit_W41","SIGAH8":"Kit_W41"} + parts = [] + for f in sorted(D_DIR.glob("*.txt.gz")): + sample = f.name.split("_")[1].split(".")[0] + 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"] = GT.get(sample, "unknown") + var = pd.DataFrame(index=df.index.astype(str)) + parts.append(ad.AnnData(X=X, obs=obs, var=var)) + a = ad.concat(parts, join="outer", label="_batch") + + import mygene + mg = mygene.MyGeneInfo() + ids = a.var_names.astype(str).tolist() + 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) + a = a[:, keep].copy(); a.var_names = syms[keep]; a.var_names_make_unique() + return a + + +def main(): + fig, axes = plt.subplots(1, 3, figsize=(17, 5.3)) + + # ---- Panel A: Dingwall (skin) ---- + print("[fig] Dingwall UMAP …", flush=True) + p = ad.read_h5ad(f"{ROOT_STR}/discovery/pan_skin/marker/50_aldrich_projections.h5ad") + Z = np.asarray(p.obsm["Z_projection"]) + rng = np.random.default_rng(RS) + idx = rng.choice(len(Z), size=min(SAMPLE_N, len(Z)), replace=False) + Z_a = Z[idx] + emb = umap_it(Z_a) + genotype = p.obs["genotype"].values[idx] + ax = axes[0] + for g, c in zip(["WT", "En1-cKO"], ["#2b83ba", "#d7191c"]): + m = genotype == g + ax.scatter(emb[m, 0], emb[m, 1], s=2, alpha=0.5, c=c, + label=f"{g} (n={int(m.sum())})") + ax.set_title(f"(a) Dingwall skin (n={SAMPLE_N}) — En1 genotype", fontsize=10) + ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2") + ax.legend(markerscale=6, frameon=False, fontsize=9) + + # ---- Panel B: Dahlin (HSC) ---- + print("[fig] Dahlin UMAP …", flush=True) + cache_d = Path(f"{ROOT_STR}/figures/_cache_dahlin_umap.npz") + if cache_d.exists(): + c = np.load(cache_d, allow_pickle=True) + emb2 = c["emb"]; gt_d = c["gt"] + else: + stats_h = np.load(f"{ROOT_STR}/data/corpus/hematopoiesis/harmonized/corpus_stats.npz", + allow_pickle=True) + shared_hvgs_h = [str(g) for g in stats_h["shared_hvgs"]] + pca_h = pickle.load(open(f"{ROOT_STR}/data/corpus/hematopoiesis/harmonized/pca_basis.pkl","rb")) + a_d = load_dahlin() + rng2 = np.random.default_rng(RS) + idx2 = rng2.choice(a_d.n_obs, size=min(SAMPLE_N, a_d.n_obs), replace=False) + a_d_sub = a_d[idx2].copy() + Z_d, pred_d = get_projection(Path(f"{ROOT_STR}/checkpoints/hematopoiesis"), + a_d_sub, shared_hvgs_h, stats_h["mean"], stats_h["std"], pca_h) + emb2 = umap_it(Z_d) + gt_d = a_d_sub.obs["genotype"].values + np.savez(cache_d, emb=emb2, gt=np.asarray(gt_d, dtype=object)) + ax = axes[1] + for g, c in zip(["WT", "Kit_W41"], ["#2b83ba", "#d7191c"]): + m = gt_d == g + ax.scatter(emb2[m, 0], emb2[m, 1], s=2, alpha=0.5, c=c, + label=f"{g} (n={int(m.sum())})") + ax.set_title(f"(b) Dahlin HSC (n={SAMPLE_N}) — Kit genotype", fontsize=10) + ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2") + ax.legend(markerscale=6, frameon=False, fontsize=9) + + # ---- Panel C: Veres (pancreas) ---- + print("[fig] Veres UMAP …", flush=True) + stats_p = np.load(f"{ROOT_STR}/data/corpus/pancreas/harmonized/corpus_stats.npz", + allow_pickle=True) + shared_hvgs_p = [str(g) for g in stats_p["shared_hvgs"]] + pca_p = pickle.load(open(f"{ROOT_STR}/data/corpus/pancreas/harmonized/pca_basis.pkl","rb")) + a_s = load_sharon_stages() + stage_num = pd.to_numeric(a_s.obs["Stage"], errors="coerce") + keep_st = stage_num.notna().values + a_s = a_s[keep_st].copy() + a_s.obs["Stage_int"] = stage_num[keep_st].astype(int).values + rng3 = np.random.default_rng(RS) + idx3 = rng3.choice(a_s.n_obs, size=min(SAMPLE_N, a_s.n_obs), replace=False) + a_s_sub = a_s[idx3].copy() + Z_s, pred_s = get_projection(Path(f"{ROOT_STR}/checkpoints/pancreas"), + a_s_sub, shared_hvgs_p, stats_p["mean"], stats_p["std"], pca_p) + emb3 = umap_it(Z_s) + stage = a_s_sub.obs["Stage_int"].values + ax = axes[2] + stage_colors = {3: "#fdae61", 4: "#f8b0d1", 5: "#7570b3", 6: "#d7191c"} + for s in sorted(np.unique(stage)): + m = stage == s + ax.scatter(emb3[m, 0], emb3[m, 1], s=2, alpha=0.5, + c=stage_colors.get(s, "#666"), label=f"Stage {s} (n={int(m.sum())})") + ax.set_title(f"(c) Veres hPSC (n={SAMPLE_N}) — differentiation stage", fontsize=10) + ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2") + ax.legend(markerscale=6, frameon=False, fontsize=9) + + plt.tight_layout() + plt.savefig(FIG / "fig6_multi_umap.pdf", bbox_inches="tight", dpi=100) + plt.close() + print(f"[fig] wrote {FIG}/fig6_multi_umap.pdf") + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/generate_paper_figures.py b/scripts/figures/generate_paper_figures.py new file mode 100644 index 0000000000000000000000000000000000000000..983e2d6c1fbde1864a74e121c1e954301aa03276 --- /dev/null +++ b/scripts/figures/generate_paper_figures.py @@ -0,0 +1,225 @@ +"""main-text figures for PAPER.tex — cv per-class F1 bars, dahlin heatmap, veres stage stack. + +Uses shared canonical palette (scripts/figures/palette.py) so the same class +gets the same color in every figure. +""" +from __future__ import annotations +from pathlib import Path +import warnings, json, sys +warnings.filterwarnings("ignore") + +import numpy as np, pandas as pd +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +from pathlib import Path as _P_root +ROOT = _P_root(__file__).resolve().parents[2] +ROOT_STR = str(ROOT) +FIG = Path(f"{ROOT_STR}/figures") +FIG.mkdir(exist_ok=True) + +# shared canonical palette + style +sys.path.insert(0, str(_P_root(__file__).resolve().parent)) +from palette import color_for, apply_style, CLASS_COLORS +apply_style() + + +def _panel_letter(ax, letter, x=-0.14, y=1.06, fontsize=20): + ax.text(x, y, f"({letter})", transform=ax.transAxes, + fontsize=fontsize, fontweight="bold", va="bottom", ha="left") + + +def fig1_confusion_matrices(): + """per-class F1 bars from 5-fold CV JSONs, one axis per system, class-colored.""" + panels = [ + ("Pan-skin", f"{ROOT_STR}/discovery/pan_skin/marker/cv_5fold.json"), + ("Pan-hematopoiesis", f"{ROOT_STR}/discovery/hematopoiesis/marker/cv_5fold.json"), + ("Pan-pancreas", f"{ROOT_STR}/discovery/pancreas/marker/cv_5fold.json"), + ] + fig, axes = plt.subplots(1, 3, figsize=(18, 6), constrained_layout=True) + for i, (ax, (title, path)) in enumerate(zip(axes, panels)): + r = json.load(open(path)) + rep = r["per_class_report"] + # collect (class, f1, support) + rows = [(c, rep[c]["f1-score"], int(rep[c]["support"])) + for c in rep.keys() + if c not in ("accuracy", "macro avg", "weighted avg")] + # sort descending by F1 (best first at top) + rows.sort(key=lambda t: t[1], reverse=True) + classes = [t[0] for t in rows] + f1s = [t[1] for t in rows] + supports = [t[2] for t in rows] + + colors = [color_for(c) for c in classes] + y = np.arange(len(classes)) + bars = ax.barh(y, f1s, color=colors, edgecolor="white", linewidth=0.6) + + # per-bar n annotations (outside) + for b, s in zip(bars, supports): + ax.text(b.get_width() + 0.012, b.get_y() + b.get_height() / 2, + f"n={s:,}", va="center", fontsize=11, color="#333") + + # scale + ax.set_xlim(0, 1.22) + ax.set_xticks([0.0, 0.25, 0.5, 0.75, 1.0]) + ax.set_xticklabels(["0.0", "0.25", "0.50", "0.75", "1.00"]) + ax.set_xlabel("held-out F1", fontsize=15) + + # y-axis: shrink font if many classes + y_fs = 14 if len(classes) <= 13 else 11 + ax.set_yticks(y) + ax.set_yticklabels(classes, fontsize=y_fs) + + # title (18pt) with n_cells; small 12pt acc/AUROC subtitle below + n_cells = int(r["n_cells"]) + acc = r["mean_acc"]; auc = r["mean_auc"] + ax.set_title(f"{title} (n_cells={n_cells:,})", fontsize=18, pad=32) + ax.text(0.5, 1.01, + f"acc = {acc:.3f} · macro AUROC = {auc:.3f}", + transform=ax.transAxes, ha="center", va="bottom", + fontsize=12, color="#555555") + + # F1=0.9 marker + ax.axvline(0.9, color="#888888", linestyle="--", linewidth=1.1, + alpha=0.4, zorder=0) + ax.grid(axis="x", alpha=0.25, linestyle=":", zorder=0) + ax.invert_yaxis() + _panel_letter(ax, "abc"[i], x=-0.32, y=1.02, fontsize=20) + + plt.savefig(FIG / "fig1_confusion_matrices.pdf", bbox_inches="tight") + plt.close() + print(f"[fig1] wrote {FIG}/fig1_confusion_matrices.pdf") + + +def fig3_dahlin_heatmap(): + """dahlin within-class module-score heatmap (Kit_W41 minus WT).""" + p = Path(f"{ROOT_STR}/discovery/hematopoiesis/marker/57_pathway_analysis.csv") + if not p.exists(): print(f"[fig3] {p} not found"); return + df = pd.read_csv(p) + pivot = df.pivot(index="module_name", columns="class", values="delta") + pivot_p = df.pivot(index="module_name", columns="class", values="mannu_p_adj_bonferroni") + row_order = ["Kit_signaling", "Kit_ligand", "MYC_targets", "Integrated_stress", + "Apoptosis_pro", "Apoptosis_anti", "Cell_cycle", + "Erythropoiesis_early", "Erythropoiesis_late", + "OXPHOS_ETC", "Glycolysis"] + row_order = [r for r in row_order if r in pivot.index] + col_order = ["LT-HSC", "MPP", "erythroid", "myeloid", "megakaryocyte", + "lymphoid", "basophil-mast", "monocyte", "macrophage"] + col_order = [c for c in col_order if c in pivot.columns] + P = pivot.loc[row_order, col_order] + Pp = pivot_p.loc[row_order, col_order] + + fig, ax = plt.subplots(figsize=(11, 8), constrained_layout=True) + vmax = np.nanmax(np.abs(P.values)) + im = ax.imshow(P.values, cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto") + + # cell annotations + for i in range(P.shape[0]): + for j in range(P.shape[1]): + v = P.values[i, j]; pv = Pp.values[i, j] + if np.isnan(v): continue + star = "**" if pv < 1e-3 else "*" if pv < 0.05 else "" + if abs(v) > 0.1: + label = f"{v:+.2f}" + if star: label = f"{label} {star}" + color = "white" if abs(v) >= vmax * 0.6 else "black" + ax.text(j, i, label, ha="center", va="center", + fontsize=10, color=color) + elif star: + ax.text(j, i, star, ha="center", va="center", + fontsize=10, color="black") + + # col labels (rotated 45, 12pt) + ax.set_xticks(range(len(col_order))) + ax.set_xticklabels(col_order, rotation=45, ha="right", fontsize=12) + # row labels (module names) at 12pt + ax.set_yticks(range(len(row_order))) + ax.set_yticklabels(row_order, fontsize=12) + + # color the x tick labels (class names) using canonical palette + for tl, cls in zip(ax.get_xticklabels(), col_order): + tl.set_color(color_for(cls)) + tl.set_fontweight("bold") + + cbar = plt.colorbar(im, ax=ax, shrink=0.85, pad=0.02) + cbar.set_label("Δ module score (Kit_W41 − WT)", fontsize=14) + cbar.ax.tick_params(labelsize=12) + + ax.set_title("Dahlin: Kit-W41 vs WT within-class pathway module contrast", + fontsize=18, pad=14) + fig.text(0.5, -0.01, + "* p<0.05 ** p<10$^{-3}$ (Mann–Whitney, Bonferroni)", + ha="center", fontsize=11, color="#555555") + + plt.savefig(FIG / "fig3_dahlin_heatmap.pdf", bbox_inches="tight") + plt.close() + print(f"[fig3] wrote {FIG}/fig3_dahlin_heatmap.pdf") + + +def fig4_sharon_stage_stack(): + """stacked-bar class fractions across veres stages 3-6, class-colored via palette.""" + import re + p = Path(f"{ROOT_STR}/discovery/pancreas/marker/veres_predictions.csv") + if not p.exists(): print(f"[fig4] {p} not found"); return + pred = pd.read_csv(p) + + stg_re = re.compile(r"_S(\d)c_") + stages = pred["cell_id"].astype(str).apply( + lambda s: int(stg_re.search(s).group(1)) if stg_re.search(s) else np.nan) + pred = pred.assign(stage=stages).dropna(subset=["stage"]) + pred["stage"] = pred["stage"].astype(int) + n_staged = len(pred); n_pre = len(stages); n_drop = n_pre - n_staged + print(f"[fig4] staged cells: {n_staged:,} (dropped {n_drop:,} primary-islet cells)") + + ct = (pred.groupby(["stage", "pred_label"]).size() + .unstack("pred_label", fill_value=0)) + frac = ct.div(ct.sum(axis=1), axis=0) + priority = ["pancreatic-progenitor", "proliferating", + "endocrine-progenitor", "endocrine-progenitor-primed", + "Fev-EP", "beta_progenitor", "beta", + "alpha_progenitor", "alpha", "delta", "gamma", "epsilon", + "acinar", "ductal", "exocrine", + "mesenchyme", "endothelial", "immune"] + present = list(frac.columns) + ordered = [c for c in priority if c in present] + \ + [c for c in sorted(present, key=lambda x: -frac[x].sum()) + if c not in priority] + frac = frac[ordered] + + fig, ax = plt.subplots(figsize=(12, 7), constrained_layout=True) + bottom = np.zeros(frac.shape[0]) + x = np.arange(frac.shape[0]) + for cls in ordered: + vals = frac[cls].values + ax.bar(x, vals, bottom=bottom, label=cls, + color=color_for(cls), edgecolor="white", linewidth=0.6) + bottom += vals + + for xi, s in zip(x, frac.index): + n_stage = int(ct.loc[s].sum()) + ax.text(xi, 1.03, f"n = {n_stage:,}", ha="center", va="bottom", + fontsize=13, color="#222222") + + ax.set_xticks(x) + ax.set_xticklabels([f"Stage {int(s)}" for s in frac.index], fontsize=14) + ax.set_xlabel("Veres protocol stage", fontsize=15) + ax.set_ylabel("Predicted class fraction", fontsize=15) + ax.set_title("Veres 2019 pancreatic differentiation", fontsize=18, pad=12) + + ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left", fontsize=12, + frameon=False, title="Predicted class", title_fontsize=13, + handlelength=1.6, borderaxespad=0.4) + ax.set_ylim(0, 1.12) + ax.set_yticks([0.0, 0.25, 0.5, 0.75, 1.0]) + plt.savefig(FIG / "fig4_veres_stage_stack.pdf", bbox_inches="tight") + plt.close() + print(f"[fig4] wrote {FIG}/fig4_veres_stage_stack.pdf " + f"({len(ordered)} classes over stages {list(frac.index)})") + + +if __name__ == "__main__": + fig1_confusion_matrices() + fig3_dahlin_heatmap() + fig4_sharon_stage_stack() + print(f"\nAll figures in {FIG}/") diff --git a/scripts/figures/generate_umap.py b/scripts/figures/generate_umap.py new file mode 100644 index 0000000000000000000000000000000000000000..7b0293b364b0b20e1555135068de571f47a42a19 --- /dev/null +++ b/scripts/figures/generate_umap.py @@ -0,0 +1,81 @@ +"""dingwall UMAP: (a) PANDA-predicted class, (b) en1 genotype.""" +from pathlib import Path +import warnings +warnings.filterwarnings("ignore") + +import numpy as np +import pandas as pd +import anndata as ad +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import umap + +FIG = Path("/home/bcheng/PRISM/figures") +FIG.mkdir(exist_ok=True) + +CLASS_COLORS = { + "fibroblast-reticular": "#66c2a5", + "basal-IFE": "#fc8d62", + "endothelial": "#8da0cb", + "immune": "#e78ac3", + "fibroblast-papillary": "#a6d854", + "melanocyte": "#ffd92f", + "spinous": "#e5c494", + "granular": "#b3b3b3", + "HF-ORS": "#1b9e77", + "HF-DP": "#d95f02", + "HF-placode": "#7570b3", + "eccrine-duct": "#666666", + "eccrine-placode": "#000000", +} + + +def main(): + p = ad.read_h5ad("/home/bcheng/PRISM/discovery/pan_skin/marker/50_aldrich_projections.h5ad") + print(f"[umap] Dingwall projection shape {p.shape}", flush=True) + if "Z_projection" not in p.obsm: + print(f"[umap] no Z_projection in obsm; keys: {list(p.obsm.keys())}") + return + Z = np.asarray(p.obsm["Z_projection"]) + print(f"[umap] Z shape {Z.shape}", flush=True) + + print(f"[umap] fitting UMAP …", flush=True) + reducer = umap.UMAP(n_neighbors=30, min_dist=0.3, random_state=42, + metric="cosine", n_components=2) + emb = reducer.fit_transform(Z) + print(f"[umap] UMAP done, emb shape {emb.shape}", flush=True) + + pred = p.obs["pred_bbse_label"] if "pred_bbse_label" in p.obs else p.obs.get("pred_label") + genotype = p.obs.get("genotype", pd.Series(index=p.obs.index, data="unknown")) + + fig, axes = plt.subplots(1, 2, figsize=(14, 6)) + + ax = axes[0] + for cls in pred.unique(): + m = pred == cls + ax.scatter(emb[m, 0], emb[m, 1], s=1.5, alpha=0.4, + c=CLASS_COLORS.get(cls, "#999999"), label=f"{cls} (n={int(m.sum())})") + ax.set_title("Dingwall (25,344 cells) — PANDA-predicted class (BBSE)") + ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2") + ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left", fontsize=7, + markerscale=6, frameon=False) + + ax = axes[1] + gcolors = {"WT": "#2b83ba", "En1-cKO": "#d7191c", "unknown": "#999999"} + for g in ["WT", "En1-cKO"]: + m = genotype == g + ax.scatter(emb[m, 0], emb[m, 1], s=1.5, alpha=0.35, + c=gcolors[g], label=f"{g} (n={int(m.sum())})") + ax.set_title("Dingwall — En1 genotype") + ax.set_xlabel("UMAP 1"); ax.set_ylabel("UMAP 2") + ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left", markerscale=6, frameon=False) + + plt.tight_layout() + plt.savefig(FIG / "fig5_dingwall_umap.pdf", bbox_inches="tight") + plt.close() + print(f"[umap] wrote {FIG}/fig5_dingwall_umap.pdf") + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/merge_all_figures.py b/scripts/figures/merge_all_figures.py new file mode 100644 index 0000000000000000000000000000000000000000..eb0926d6952026e434e9978dd052f387580edc1c --- /dev/null +++ b/scripts/figures/merge_all_figures.py @@ -0,0 +1,107 @@ +"""merge every paper figure into one browsable pdf. + +section-title pages divide main / supplement / biology; each figure gets +a small header page with its filename so a reader can find the source. +""" +from __future__ import annotations +from pathlib import Path +from pypdf import PdfWriter +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +ROOT = Path("/home/bcheng/PRISM") +FIG = ROOT / "figures" +OUT = FIG / "PANDA_all_figures.pdf" +TMP = FIG / "_tmp_titlepages.pdf" + +MAIN = [ + ("Figure 1", "fig1_confusion_matrices.pdf", "Per-class held-out 5-fold F1 across three systems"), + ("Figure 3", "fig3_dahlin_heatmap.pdf", "Dahlin Kit-W41 vs WT within-class pathway module deltas"), + ("Figure 4", "fig4_veres_stage_stack.pdf", "Veres predicted class fraction per protocol stage"), + ("Figure 5", "fig5_dingwall_umap.pdf", "Dingwall UMAP (predicted class and En1 genotype)"), + ("Figure 6", "fig6_multi_umap.pdf", "Discovery-target UMAPs across skin / HSC / pancreas"), +] + +SUPP = [ + ("S1", "supplement/01_cv_summary.pdf", "CV summary"), + ("S2", "supplement/02_per_class_f1.pdf", "Per-class F1"), + ("S3", "supplement/03_prototype_cosine.pdf", "Prototype intra-cosine"), + # S4 removed -- contained stale class vocab + ("S5", "supplement/05_adversary_purification.pdf", "Adversary purification"), + ("S6", "supplement/06_cross_system_prototypes.pdf", "Cross-system prototype map"), + # S7 removed -- contained stale class vocab + # S8 removed -- contained stale class vocab + # S9 removed -- contained stale class vocab + # S10 removed -- contained stale class vocab + ("S11", "supplement/11_novel_populations.pdf", "Novel populations (legacy)"), + # S12 removed -- contained stale class vocab + ("S13", "supplement/13_dingwall_umap.pdf", "Dingwall UMAP (legacy)"), + ("S14", "supplement/14_dahlin_umap.pdf", "Dahlin UMAP (legacy)"), + ("S15", "supplement/15_veres_umap.pdf", "Veres UMAP (legacy)"), + ("S16", "supplement/16_dingwall_discovery.pdf", "Dingwall discovery (legacy)"), + ("S17", "supplement/17_dahlin_discovery.pdf", "Dahlin discovery (legacy)"), + ("S18", "supplement/18_veres_discovery.pdf", "Veres discovery (legacy)"), + ("S19", "supplement/19_myeloid_network.pdf", "Myeloid gene-gene network"), + ("S20", "supplement/20_placode_wnt_module.pdf", "Placode WNT module"), + ("S23", "supplement/23_anchor_delta_recall.pdf", "Anchor delta recall"), + ("S24", "supplement/24_pca_vs_marker_umaps_dingwall_by_genotype.pdf", "PCA vs Marker: Dingwall by genotype"), + ("S24b","supplement/24b_pca_vs_marker_umaps_dingwall_by_class.pdf", "PCA vs Marker: Dingwall by class"), + ("S25", "supplement/25_pca_vs_marker_umaps_dahlin_by_genotype.pdf", "PCA vs Marker: Dahlin by genotype"), + ("S25b","supplement/25b_pca_vs_marker_umaps_dahlin_by_class.pdf", "PCA vs Marker: Dahlin by class"), + ("S26", "supplement/26_pca_vs_marker_umaps_veres_by_stage.pdf", "PCA vs Marker: Veres by stage"), + ("S26b","supplement/26b_pca_vs_marker_umaps_veres_by_class.pdf", "PCA vs Marker: Veres by class"), + ("S27", "supplement/27_dingwall_en1_enrichment.pdf", "Dingwall En1 class enrichment"), + ("S28", "supplement/28_melanocyte_pathway_modules.pdf", "Melanocyte pathway modules"), +] + +BIO = [ + ("B1", "biology/biology_01_dingwall_umap.pdf", "Dingwall UMAP (biology)"), + ("B2", "biology/biology_02_primary_eden.pdf", "Primary EDEN candidate: Derm2"), + ("B3", "biology/biology_03_melanoblast_mitf.pdf", "Melanoblast MITF-axis quadrant"), + ("B4", "biology/biology_04_dahlin_metabolism.pdf", "Dahlin per-lineage metabolism"), + ("B5", "biology/biology_05_dahlin_composition.pdf", "Dahlin composition shift"), + ("B6", "biology/biology_06_veres_beta_quadrant.pdf", "Veres beta-lineage quadrant"), + ("B7", "biology/biology_07_veres_polyhormonal.pdf", "Veres polyhormonal SC-alpha"), + ("B8", "biology/biology_08_prototype_geometry.pdf", "Prototype geometry"), +] + + +def make_title_page(text_top, text_body): + fig = plt.figure(figsize=(8.5, 11)) + fig.text(0.5, 0.55, text_top, ha="center", va="center", fontsize=28, weight="bold") + fig.text(0.5, 0.45, text_body, ha="center", va="center", fontsize=14) + fig.savefig(TMP, bbox_inches="tight") + plt.close(fig) + + +def append_group(w, title, items): + make_title_page(title, "") + w.append(str(TMP)) + for tag, path, desc in items: + p = FIG / path + if not p.exists(): + print(f" [skip] {path} missing") + continue + make_title_page(tag, f"{desc}\n\n{path}") + w.append(str(TMP)) + w.append(str(p)) + print(f" + {tag} {p.name}") + + +def main(): + w = PdfWriter() + make_title_page("PANDA — all figures", "main text · supplement · biology") + w.append(str(TMP)) + append_group(w, "Main text", MAIN) + append_group(w, "Supplement", SUPP) + append_group(w, "Biology", BIO) + with open(OUT, "wb") as f: + w.write(f) + TMP.unlink(missing_ok=True) + print(f"\n[all] wrote {OUT} ({OUT.stat().st_size / 1024:.0f} KB, " + f"{len(w.pages)} pages)") + + +if __name__ == "__main__": + main() diff --git a/scripts/figures/palette.py b/scripts/figures/palette.py new file mode 100644 index 0000000000000000000000000000000000000000..fcee75ea2a22330503fd6081824f13b9e20aaa7d --- /dev/null +++ b/scripts/figures/palette.py @@ -0,0 +1,118 @@ +"""canonical class -> color mapping shared across every panda figure. + +related classes share hue families (fibroblasts red, keratinocytes orange, +melanocyte lineage purple, endothelial teal, immune yellow, etc.). every +figure imports CLASS_COLORS[cls] so class X gets the same color everywhere. +""" + +# skin (13 classes) +_SKIN = { + "basal-IFE": "#e07b39", # keratinocyte orange + "spinous": "#c94f21", + "granular": "#a53818", + "HF-placode": "#8b4513", # hair follicle brown + "HF-ORS": "#c68e4b", + "sebaceous": "#bfa970", + "fibroblast-papillary": "#d94848", # fibroblast reds + "fibroblast-reticular": "#8f1e1e", + "endothelial": "#2e8b8b", # vascular teal + "immune": "#d1c02b", # immune yellow + "melanoblast": "#6a5acd", # melanocyte purples + "melanocyte-precursor": "#9370db", + "melanocyte": "#4b3ca3", +} + +# hematopoiesis (15 classes) +_HSC = { + "LT-HSC": "#1a3d78", # stem dark blue + "MPP": "#16a085", # progenitor green (was blue → confused w/ LT-HSC) + "erythroid": "#c1272d", # red lineage + "megakaryocyte": "#a1204b", + "basophil-mast": "#e07b39", + "myeloid": "#8a5822", # myeloid browns (dark) + "monocyte": "#d4a373", # tan — distinct from myeloid brown & macrophage + "macrophage": "#6b4321", + "lymphoid": "#2e8b8b", # lymph teal / greens + "T-cell": "#3aa17f", + "naive-B": "#5cbf88", + "pro-B": "#7fd39d", + "endothelial": "#2e8b8b", # shared with skin + "fibroblast": "#8f1e1e", + "stromal": "#6b4a3b", +} + +# pancreas (20 classes) +_PANC = { + "alpha": "#a3123a", # alpha reds + "adult-alpha": "#7a0e2b", + "alpha_progenitor": "#d94360", + "beta": "#1f4e9e", # beta blues + "adult-beta": "#123469", + "beta_progenitor": "#4d78c4", + "delta": "#7d3c98", # delta purple + "gamma": "#af7ac5", + "epsilon": "#f39c12", # epsilon orange + "Fev-EP": "#e67e22", # endocrine progenitor oranges + "endocrine-progenitor": "#c68e4b", + "endocrine-progenitor-primed": "#a86b28", + "pancreatic-progenitor": "#16a085", # progenitor greens + "proliferating": "#48c9b0", + "acinar": "#8a5822", # exocrine browns + "exocrine": "#5d3a11", + "ductal": "#7b6b57", + "endothelial": "#2e8b8b", # shared + "immune": "#d1c02b", + "mesenchyme": "#6b4a3b", +} + +# every class from every system, plus a safe fallback +CLASS_COLORS = {**_SKIN, **_HSC, **_PANC} + +# canonical genotype / condition colors +GENOTYPE_COLORS = { + "WT": "#1f77b4", + "En1-cKO": "#d62728", + "Kit_W41": "#d62728", + "control": "#1f77b4", + "other": "#bbbbbb", +} + +# stage colors (viridis-ish for ordinal) +STAGE_COLORS = { + "3": "#440154", + "4": "#3b528b", + "5": "#21918c", + "6": "#fde725", +} + + +def color_for(cls, fallback="#888888"): + """look up a class name; return fallback if missing.""" + if cls in CLASS_COLORS: + return CLASS_COLORS[cls] + # simple aliases + s = str(cls).strip() + return CLASS_COLORS.get(s, fallback) + + +def apply_style(): + """shared rcparams; call at top of every figure script.""" + import matplotlib.pyplot as plt + plt.rcParams.update({ + "font.family": "DejaVu Sans", + "font.size": 14, + "axes.titlesize": 18, + "axes.labelsize": 15, + "xtick.labelsize": 13, + "ytick.labelsize": 13, + "legend.fontsize": 12, + "legend.title_fontsize": 13, + "figure.titlesize": 20, + "pdf.fonttype": 42, + "ps.fonttype": 42, + "axes.spines.top": False, + "axes.spines.right": False, + "axes.linewidth": 1.2, + "xtick.major.width": 1.2, + "ytick.major.width": 1.2, + }) diff --git a/scripts/figures/regen_fig5_fig6.py b/scripts/figures/regen_fig5_fig6.py new file mode 100644 index 0000000000000000000000000000000000000000..b9895bdb2f280af300464d9abca0e67221cbfd59 --- /dev/null +++ b/scripts/figures/regen_fig5_fig6.py @@ -0,0 +1,268 @@ +"""regenerate fig5 (dingwall UMAP, 2 panels) + fig6 (3-panel multi UMAP). + +Uses shared canonical palette (scripts/figures/palette.py) so a class gets +the same color in every figure. Reuses project()/do_umap()/_load_dahlin_raw()/ +_load_veres() from build_pca_vs_marker_umaps.py. + +Caches expensive UMAP embeddings to figures/_cache_*.npz for reuse. + +Fig 6c restricts Veres to the 12,297 held-out slice +(data/corpus/pancreas/held_out_labeled/veres_GSE114412_test.h5ad). + +outputs: + figures/fig5_dingwall_umap.pdf + figures/fig6_multi_umap.pdf +""" +from __future__ import annotations +from pathlib import Path +import warnings, sys, numpy as np, pandas as pd +warnings.filterwarnings("ignore") +import matplotlib; matplotlib.use("Agg") +import matplotlib.pyplot as plt +import anndata as ad + +ROOT = Path("/home/bcheng/PRISM") +FIG = ROOT / "figures" +FIG.mkdir(exist_ok=True) + +sys.path.insert(0, str(ROOT / "scripts/figures")) +from build_pca_vs_marker_umaps import ( + project, do_umap, _load_dahlin_raw, _load_veres, +) +from palette import color_for, apply_style, GENOTYPE_COLORS, STAGE_COLORS +apply_style() + +RS = 42 + +# stage color for the primary-islet bucket (not in ordinal STAGE_COLORS) +# medium grey so the legend swatch is visible on white backgrounds +_ISLET_GRAY = "#707070" + +# darker grey for the "other" bucket in fig5 panel (a) — the light #e5e5e5 +# used previously was nearly invisible in legend swatches on white paper +_OTHER_GREY = "#909090" + +# Dingwall genotype mapping (verified against GSE220977 metadata) +DINGWALL_CKO = {"GSM6833482", "GSM6833483"} +DINGWALL_WT = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"} + + +def _panel_letter(ax, letter, x=-0.08, y=1.03, fontsize=22): + ax.text(x, y, f"({letter})", transform=ax.transAxes, + fontsize=fontsize, fontweight="bold", va="bottom", ha="left") + + +def _dominant_classes(P, min_frac=0.005): + """return the classes making up >= min_frac of the cells (in count order).""" + from collections import Counter + n = len(P) + counts = Counter(P.tolist()) + return [c for c, k in counts.most_common() if k / n >= min_frac] + + +def fig5_dingwall(): + cache_p = FIG / "_cache_dingwall_full_marker.npz" + if cache_p.exists(): + print(f"[fig5] reusing cache {cache_p.name}", flush=True) + c = np.load(cache_p, allow_pickle=True) + emb, P, gt = c["emb"], c["P"], c["genotype"] + n = len(emb) + else: + print("[fig5] loading dingwall raw ...", flush=True) + raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad") + gt = np.where(raw.obs["sample"].astype(str).isin(list(DINGWALL_CKO)), "En1-cKO", + np.where(raw.obs["sample"].astype(str).isin(list(DINGWALL_WT)), "WT", "other")) + n = raw.n_obs + print(f"[fig5] projecting {n:,} cells (PANDA-Marker) ...", flush=True) + Z, P, _ = project(raw, "pan_skin", "marker") + print(f"[fig5] UMAP on {n:,} projections ...", flush=True) + emb = do_umap(Z) + np.savez(cache_p, emb=emb, P=P, genotype=gt) + + fig, axes = plt.subplots(1, 2, figsize=(18, 8), constrained_layout=True) + + # (a) PANDA-predicted class + ax = axes[0] + kept = _dominant_classes(P, min_frac=0.005) + P_r = np.where(np.isin(P, kept), P, "other") + # plot "other" first so kept classes render on top + m_other = P_r == "other" + if m_other.sum() > 0: + ax.scatter(emb[m_other, 0], emb[m_other, 1], s=6, alpha=0.6, + c=_OTHER_GREY, label=f"other (n={int(m_other.sum()):,})", + linewidths=0, rasterized=True) + for cls in kept: + m = P_r == cls + if m.sum() == 0: continue + ax.scatter(emb[m, 0], emb[m, 1], s=6, alpha=0.6, + c=color_for(cls), + label=f"{cls} (n={int(m.sum()):,})", linewidths=0, + rasterized=True) + ax.set_title("Dingwall skin — PANDA-Marker predicted class", + fontsize=18, pad=8) + ax.set_xlabel("UMAP-1", fontsize=14); ax.set_ylabel("UMAP-2", fontsize=14) + ax.set_xticks([]); ax.set_yticks([]) + ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left", fontsize=12, + markerscale=1.6, frameon=False, + title=f"Predicted class (n={n:,})", title_fontsize=14, + handletextpad=0.5, borderaxespad=0.4) + _panel_letter(ax, "a") + + # (b) En1 genotype + ax = axes[1] + for g in ["other", "WT", "En1-cKO"]: # cKO last so it plots on top + m = gt == g + if m.sum() == 0: continue + ax.scatter(emb[m, 0], emb[m, 1], s=6, alpha=0.6, + c=GENOTYPE_COLORS.get(g, "#bbbbbb"), + label=f"{g} (n={int(m.sum()):,})", + linewidths=0, rasterized=True) + ax.set_title("Dingwall skin — En1 genotype", fontsize=18, pad=8) + ax.set_xlabel("UMAP-1", fontsize=14); ax.set_ylabel("UMAP-2", fontsize=14) + ax.set_xticks([]); ax.set_yticks([]) + ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left", fontsize=12, + markerscale=1.6, frameon=False, + title="En1 genotype", title_fontsize=14, + handletextpad=0.5, borderaxespad=0.4) + _panel_letter(ax, "b") + + plt.savefig(FIG / "fig5_dingwall_umap.pdf", bbox_inches="tight", dpi=200) + plt.close() + print(f"[fig5] wrote {FIG}/fig5_dingwall_umap.pdf", flush=True) + + +def _get_dingwall_cached(): + cache_p = FIG / "_cache_dingwall_full_marker.npz" + if cache_p.exists(): + c = np.load(cache_p, allow_pickle=True) + return c["emb"], c["genotype"] + # otherwise recompute + print("[fig6] recomputing dingwall ...", flush=True) + raw = ad.read_h5ad(ROOT / "data/raw/GSE220977_combined.h5ad") + gt_s = np.where(raw.obs["sample"].astype(str).isin(list(DINGWALL_CKO)), "En1-cKO", + np.where(raw.obs["sample"].astype(str).isin(list(DINGWALL_WT)), "WT", "other")) + Z_s, P_s, _ = project(raw, "pan_skin", "marker") + emb_s = do_umap(Z_s) + np.savez(cache_p, emb=emb_s, P=P_s, genotype=gt_s) + return emb_s, gt_s + + +def _get_dahlin_cached(): + cache_p = FIG / "_cache_dahlin_marker.npz" + if cache_p.exists(): + print(f"[fig6] reusing {cache_p.name}", flush=True) + c = np.load(cache_p, allow_pickle=True) + return c["emb"], c["genotype"] + print("[fig6] loading dahlin raw ...", flush=True) + a_d = _load_dahlin_raw() + gt_d = a_d.obs["genotype"].astype(str).values + print(f"[fig6] projecting dahlin ({a_d.n_obs:,}) ...", flush=True) + Z_d, _, _ = project(a_d, "hematopoiesis", "marker") + print(f"[fig6] UMAP dahlin ...", flush=True) + emb_d = do_umap(Z_d) + np.savez(cache_p, emb=emb_d, genotype=gt_d) + return emb_d, gt_d + + +def _get_veres_heldout_cached(): + """Project + UMAP on the 12,297 held-out Veres slice only.""" + cache_p = FIG / "_cache_veres_heldout_marker.npz" + if cache_p.exists(): + print(f"[fig6] reusing {cache_p.name}", flush=True) + c = np.load(cache_p, allow_pickle=True) + return c["emb"], c["stage"] + + print("[fig6] loading veres raw + held-out obs list ...", flush=True) + heldout_p = ROOT / "data/corpus/pancreas/held_out_labeled/veres_GSE114412_test.h5ad" + if heldout_p.exists(): + heldout_raw = set(ad.read_h5ad(heldout_p).obs_names.astype(str).tolist()) + heldout_names = {n[6:] if n.startswith("veres_") else n for n in heldout_raw} + heldout_names |= heldout_raw + else: + print(f"[fig6] WARN: held-out file missing; using ALL veres cells") + heldout_names = None + + a_v = _load_veres() + if heldout_names is not None: + keep = np.array([str(n) in heldout_names for n in a_v.obs_names]) + print(f"[fig6] restricting veres to held-out: {keep.sum():,}/{a_v.n_obs:,}", flush=True) + a_v = a_v[keep].copy() + + stage_col = "Stage" if "Stage" in a_v.obs.columns else "stage" + stage = pd.to_numeric(a_v.obs[stage_col], errors="coerce").fillna(-1).astype(int).values + st_str = np.array([str(s) if s > 0 else "islet" for s in stage]) + print(f"[fig6] projecting veres held-out ({a_v.n_obs:,}) ...", flush=True) + Z_v, _, _ = project(a_v, "pancreas", "marker") + print(f"[fig6] UMAP veres held-out ...", flush=True) + emb_v = do_umap(Z_v) + np.savez(cache_p, emb=emb_v, stage=st_str) + return emb_v, st_str + + +def fig6_multi(): + """3-panel multi-system UMAP colored by biology-of-interest label.""" + fig, axes = plt.subplots(1, 3, figsize=(22, 8), constrained_layout=True) + + # ---- (a) Dingwall (skin, En1 genotype) ---- + emb_s, gt_s = _get_dingwall_cached() + ax = axes[0] + for g in ["other", "WT", "En1-cKO"]: + m = gt_s == g + if m.sum() == 0: continue + ax.scatter(emb_s[m, 0], emb_s[m, 1], s=5, alpha=0.55, + c=GENOTYPE_COLORS.get(g, "#bbbbbb"), + label=f"{g} (n={int(m.sum()):,})", + linewidths=0, rasterized=True) + ax.set_title("(a) Dingwall skin — En1 genotype", fontsize=18, pad=8) + ax.set_xlabel("UMAP-1", fontsize=14); ax.set_ylabel("UMAP-2", fontsize=14) + ax.set_xticks([]); ax.set_yticks([]) + ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.02), + ncol=3, frameon=False, fontsize=11, markerscale=2.0, + handletextpad=0.4, columnspacing=1.2) + + # ---- (b) Dahlin (HSC, Kit genotype) ---- + emb_d, gt_d = _get_dahlin_cached() + ax = axes[1] + # dahlin uses "unknown" for cells outside labeled samples + for g in ["unknown", "WT", "Kit_W41"]: + m = gt_d == g + if m.sum() == 0: continue + c = GENOTYPE_COLORS.get(g, "#bbbbbb") if g != "unknown" else "#bbbbbb" + ax.scatter(emb_d[m, 0], emb_d[m, 1], s=5, alpha=0.55, + c=c, label=f"{g} (n={int(m.sum()):,})", + linewidths=0, rasterized=True) + ax.set_title("(b) Dahlin hematopoiesis — Kit genotype", fontsize=18, pad=8) + ax.set_xlabel("UMAP-1", fontsize=14); ax.set_ylabel("UMAP-2", fontsize=14) + ax.set_xticks([]); ax.set_yticks([]) + ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.02), + ncol=3, frameon=False, fontsize=11, markerscale=2.0, + handletextpad=0.4, columnspacing=1.2) + + # ---- (c) Veres held-out (pancreas, Stage) ---- + emb_v, st_str = _get_veres_heldout_cached() + ax = axes[2] + stage_order = ["islet", "3", "4", "5", "6"] + for s in stage_order: + m = st_str == s + if m.sum() == 0: continue + color = STAGE_COLORS[s] if s in STAGE_COLORS else _ISLET_GRAY + label = (f"Stage {s} (n={int(m.sum()):,})" if s != "islet" + else f"islet (n={int(m.sum()):,})") + ax.scatter(emb_v[m, 0], emb_v[m, 1], s=5, alpha=0.55, + c=color, label=label, linewidths=0, rasterized=True) + ax.set_title(f"(c) Veres pancreas held-out (n={len(st_str):,}) — protocol stage", + fontsize=18, pad=8) + ax.set_xlabel("UMAP-1", fontsize=14); ax.set_ylabel("UMAP-2", fontsize=14) + ax.set_xticks([]); ax.set_yticks([]) + ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.02), + ncol=5, frameon=False, fontsize=11, markerscale=2.0, + handletextpad=0.4, columnspacing=1.2) + + plt.savefig(FIG / "fig6_multi_umap.pdf", bbox_inches="tight", dpi=180) + plt.close() + print(f"[fig6] wrote {FIG}/fig6_multi_umap.pdf", flush=True) + + +if __name__ == "__main__": + fig5_dingwall() + fig6_multi() diff --git a/scripts/hematopoiesis/01_download.sh b/scripts/hematopoiesis/01_download.sh new file mode 100644 index 0000000000000000000000000000000000000000..3e0f62fcd5789ae89e61a568a2e3f928911849fb --- /dev/null +++ b/scripts/hematopoiesis/01_download.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +set -uo pipefail +BASE=/home/bcheng/PRISM/data/corpus/hematopoiesis/tier_a +LOG=/home/bcheng/PRISM/logs/hematopoiesis +mkdir -p "$BASE" "$LOG" + +download() { + local acc=$1 url=$2 out=$3 + local outfile="$BASE/$out" + if [ -s "$outfile" ]; then echo "[$acc] cached"; return 0; fi + wget --quiet -c -O "$outfile.part" "$url" && mv "$outfile.part" "$outfile" \ + && echo "[$acc] done: $(du -h "$outfile" | cut -f1)" \ + || { echo "[$acc] FAILED"; rm -f "$outfile.part"; return 1; } +} + +# GSE72857 Paul 2015 - MARS-seq +download GSE72857 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE72nnn/GSE72857/suppl/GSE72857_umitab.txt.gz" \ + paul_GSE72857_umitab.txt.gz + +# GSE81682 Nestorowa 2016 - Smart-seq2 +download GSE81682 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE81nnn/GSE81682/suppl/GSE81682_HTSeq_counts.txt.gz" \ + nestorowa_GSE81682_counts.txt.gz + +# GSE89754 Tusi 2018 - inDrops +download GSE89754 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE89nnn/GSE89754/suppl/GSE89754_RAW.tar" \ + tusi_GSE89754_RAW.tar + +# GSE107727 Dahlin 2018 - 10x LSK/Kit+ (large -- 44k) +download GSE107727 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE107nnn/GSE107727/suppl/GSE107727_RAW.tar" \ + dahlin_GSE107727_RAW.tar + +ls -lh "$BASE" diff --git a/scripts/hematopoiesis/02_build_per_dataset.py b/scripts/hematopoiesis/02_build_per_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..d5f8fa68f02ec6f5220137b946d4a8b358d1c902 --- /dev/null +++ b/scripts/hematopoiesis/02_build_per_dataset.py @@ -0,0 +1,39 @@ +"""per-dataset h5ads for pan-hsc corpus.""" +from pathlib import Path +import warnings, sys +warnings.filterwarnings("ignore") +import anndata as ad, scanpy as sc, scipy.sparse as sp + +# local weinreb (75k cells) has cell_type annotations — use as anchor +def load_weinreb_local(): + p = Path("/home/bcheng/PRISM/data/processed/weinreb_larry/adata_weinreb.h5ad") + a = ad.read_h5ad(p) + if a.raw is not None: + raw = a.raw.to_adata(); raw.obs = a.obs.copy(); a = raw + a.obs["dataset"] = "weinreb_GSE140802" + return a + +def qc(a, name): + n0 = a.n_obs + sc.pp.filter_cells(a, min_genes=200) + sc.pp.filter_genes(a, min_cells=3) + print(f" [{name}] {n0} -> {a.n_obs} cells, {a.n_vars} genes", flush=True) + return a + +def main(): + out = Path("/home/bcheng/PRISM/data/corpus/hematopoiesis/harmonized") + out.mkdir(parents=True, exist_ok=True) + name = "weinreb_GSE140802" + if not (out / f"{name}.h5ad").exists(): + print(f"[{name}] loading …", flush=True) + a = load_weinreb_local() + a = qc(a, name) + a.obs["dataset"] = name + if not sp.issparse(a.X): a.X = sp.csr_matrix(a.X) + a.X = a.X.astype("float32") + a.var_names_make_unique() + a.write_h5ad(out / f"{name}.h5ad", compression="gzip") + print(f" saved to {out/f'{name}.h5ad'}", flush=True) + +if __name__ == "__main__": + main() diff --git a/scripts/hematopoiesis/02_download.sh b/scripts/hematopoiesis/02_download.sh new file mode 100644 index 0000000000000000000000000000000000000000..3a2869b92c49d6d1b280e535eb7e93f32d1197d5 --- /dev/null +++ b/scripts/hematopoiesis/02_download.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +# Hematopoiesis v2 additions — whole-BM in-vivo mature erythroid/lymphoid + LT-HSC/MPP substates. +set -euo pipefail +BASE="${1:-data/corpus/hematopoiesis/tier_v2}" +mkdir -p "$BASE" + +download() { + local id="$1"; local url="$2"; local out="$3" + if [ -s "$BASE/$out" ]; then + echo "[$id] cached" + return 0 + fi + echo "[$id] downloading $url -> $out" + curl -fL --retry 3 -o "$BASE/$out" "$url" \ + || { echo "[$id] FAILED"; return 1; } +} + +# GSE122465 - Baccin C et al 2020 Nat Cell Biol; whole-BM stromal + hematopoietic CITE-seq +download GSE122465 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE122nnn/GSE122465/suppl/GSE122465_RAW.tar" \ + baccin_GSE122465_wholeBM_RAW.tar + +# GSE132042 - Tabula Muris Senis bone marrow +# Get bone marrow droplet subset — figshare has combined TMS +download GSE132042 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE132nnn/GSE132042/suppl/GSE132042_RAW.tar" \ + tms_GSE132042_RAW.tar + +echo "[done] hematopoiesis v2 downloads complete" diff --git a/scripts/hematopoiesis/03_shared_hvgs_and_pca.py b/scripts/hematopoiesis/03_shared_hvgs_and_pca.py new file mode 100644 index 0000000000000000000000000000000000000000..e8eb9fcfb8f1112cb4b167609f1058770b6f8e20 --- /dev/null +++ b/scripts/hematopoiesis/03_shared_hvgs_and_pca.py @@ -0,0 +1,153 @@ +"""shared hvgs + sample-fit pca for pan-hsc corpus.""" +from __future__ import annotations +from pathlib import Path +import warnings, pickle +warnings.filterwarnings("ignore") + +import numpy as np +import pandas as pd +import scanpy as sc +import scipy.sparse as sp +from sklearn.decomposition import PCA +import anndata as ad + +HARM = Path("/home/bcheng/PRISM/data/corpus/hematopoiesis/harmonized") +K_TOP_PER_DS = 4000 +K_TARGET = 3500 +N_PCA = 50 +SAMPLE_PER_DS = 5000 + +# canonical hematopoiesis markers to force-include +MUST_INCLUDE = [ + 'Procr','Fgd5','Hlf','Mllt3','Mecom','Meis1','Mpl','Slamf1','Lmo2', + 'Cd48','Flt3','Cebpa','Csf1r','Gata1','Cd44', + 'Klf1','Gata2','Zfpm1','Epor','Mpo','Elane','Cebpe','Cd177','Ms4a3', + 'Hbb-b1','Hbb-b2','Hba-a1','Hba-a2','Gypa','S100a8','S100a9','Csf3r','Ltf', + 'Cd19','Ms4a1','Ebf1','Rag1','Il7r','Itga2b','Pf4','Gp1ba','Nfe2', + 'Cpa3','Prss34','Mrgprb1','Kit','Mcpt8', + 'Ptprc','Cd3e','Cd8a','Cd4','Adgre1','Itgam', +] + + +def rank_hvgs(a): + x = a.copy() + try: + sc.pp.highly_variable_genes(x, n_top_genes=K_TOP_PER_DS, flavor="seurat", + subset=False, check_values=False) + return x.var["variances_norm"].fillna(-np.inf) if "variances_norm" in x.var else \ + x.var["dispersions_norm"].fillna(-np.inf) + except Exception: + sc.pp.normalize_total(x, target_sum=1e4); sc.pp.log1p(x) + X = x.X.toarray() if sp.issparse(x.X) else x.X + return pd.Series(np.asarray(X.var(axis=0)).ravel(), index=x.var_names) + + +def load_and_norm(f, shared_genes, gene_idx, mu, sig, G): + a = ad.read_h5ad(f) + raw_counts = np.asarray(a.X.sum(axis=1)).ravel() if sp.issparse(a.X) else a.X.sum(axis=1) + common = [g for g in shared_genes if g in a.var_names] + a_s = a[:, common].copy() + sc.pp.normalize_total(a_s, target_sum=1e4) + sc.pp.log1p(a_s) + X = a_s.X.toarray() if sp.issparse(a_s.X) else a_s.X + Xf = np.zeros((X.shape[0], G), dtype=np.float32) + cols = [gene_idx[g] for g in common] + Xf[:, cols] = X.astype(np.float32) + Xz = np.clip((Xf - mu.astype(np.float32)) / sig.astype(np.float32), -10, 10) + return a, Xf, Xz, raw_counts, common + + +def main(): + files = sorted(HARM.glob("*.h5ad")) + files = [f for f in files if not f.name.startswith("corpus")] + print(f"[hvg] found {len(files)} datasets", flush=True) + + ranks = {} + for f in files: + print(f"[hvg] rank {f.name}", flush=True) + a = ad.read_h5ad(f) + ranks[f.stem] = rank_hvgs(a) + del a + + R = pd.DataFrame(ranks) + top_hits = pd.DataFrame({d: R[d].rank(ascending=False, method="min") <= K_TOP_PER_DS + for d in R.columns}).fillna(False) + hit_count = top_hits.sum(axis=1) + eligible = hit_count[hit_count >= max(1, len(files) // 2)] + print(f"[hvg] {len(eligible)} genes top-{K_TOP_PER_DS} in >=half datasets", flush=True) + + mean_score = R.mean(axis=1, skipna=True) + shared_by_score = mean_score.loc[eligible.index].sort_values(ascending=False).index + + all_genes = set(R.index) + forced = [g for g in MUST_INCLUDE if g in all_genes] + print(f"[hvg] forced markers present: {len(forced)}/{len(MUST_INCLUDE)}", flush=True) + + picked = list(forced) + for g in shared_by_score: + if g not in picked: + picked.append(g) + if len(picked) >= K_TARGET: + break + shared_genes = picked[:K_TARGET] + print(f"[hvg] final HVG count: {len(shared_genes)}", flush=True) + (HARM / "shared_hvgs.txt").write_text("\n".join(shared_genes) + "\n") + + G = len(shared_genes) + gene_idx = {g: i for i, g in enumerate(shared_genes)} + running_sum = np.zeros(G, dtype=np.float64) + running_sq = np.zeros(G, dtype=np.float64) + N = 0 + for f in files: + a = ad.read_h5ad(f) + common = [g for g in shared_genes if g in a.var_names] + a_s = a[:, common].copy() + sc.pp.normalize_total(a_s, target_sum=1e4) + sc.pp.log1p(a_s) + X = a_s.X.toarray() if sp.issparse(a_s.X) else a_s.X + cols = [gene_idx[g] for g in common] + running_sum[cols] += X.sum(axis=0) + running_sq[cols] += (X ** 2).sum(axis=0) + N += X.shape[0] + del a, a_s, X + mu = running_sum / N + sig = np.sqrt(np.maximum(running_sq / N - mu ** 2, 1e-6)) + np.savez(HARM / "corpus_stats.npz", shared_hvgs=np.array(shared_genes), + mean=mu.astype(np.float32), std=sig.astype(np.float32), n_cells=N) + print(f"[stats] mean/std over {N} cells", flush=True) + + fit_chunks = [] + rng = np.random.default_rng(0) + for f in files: + _, _, Xz, _, _ = load_and_norm(f, shared_genes, gene_idx, mu, sig, G) + n = Xz.shape[0] + take = min(SAMPLE_PER_DS, n) + idx = rng.choice(n, size=take, replace=False) + fit_chunks.append(Xz[idx]) + Xfit = np.vstack(fit_chunks) + pca = PCA(n_components=N_PCA, random_state=42).fit(Xfit) + print(f"[pca] explained var={pca.explained_variance_ratio_.sum():.3f}", flush=True) + for k in range(N_PCA): + top = int(np.argmax(np.abs(pca.components_[k]))) + if pca.components_[k, top] < 0: pca.components_[k] *= -1 + with open(HARM / "pca_basis.pkl", "wb") as fh: pickle.dump(pca, fh) + + parts = [] + for f in files: + a, Xf, Xz, raw_counts, common = load_and_norm(f, shared_genes, gene_idx, mu, sig, G) + Z = pca.transform(Xz).astype(np.float32) + obs = a.obs.copy() + obs["total_counts"] = raw_counts.astype(np.float32) + obs["missing_hvg_frac"] = 1.0 - len(common) / G + a2 = ad.AnnData(X=sp.csr_matrix(Xf.astype(np.float32)), obs=obs, + var=pd.DataFrame(index=shared_genes)) + a2.obsm["X_pca"] = Z + parts.append(a2) + corpus = ad.concat(parts, join="outer", label="_batch") + corpus.uns["shared_hvgs"] = shared_genes + corpus.write_h5ad(HARM / "corpus.h5ad", compression="gzip") + print(f"[emit] wrote corpus.h5ad ({corpus.n_obs} x {corpus.n_vars})", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/hematopoiesis/05_train_panda.py b/scripts/hematopoiesis/05_train_panda.py new file mode 100644 index 0000000000000000000000000000000000000000..fa26c7f582047c0e6b4690c76b18a54884303fa6 --- /dev/null +++ b/scripts/hematopoiesis/05_train_panda.py @@ -0,0 +1,133 @@ +"""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 + +sys.path.insert(0, "/home/bcheng/PRISM") +from panda.pan_skin.model import ( + PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, prototype_infonce +) + +CORPUS = Path("/home/bcheng/PRISM/data/corpus/hematopoiesis/harmonized/corpus.h5ad") +OUT = Path("/home/bcheng/PRISM/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() diff --git a/scripts/hematopoiesis/07_heldout_5fold_cv.py b/scripts/hematopoiesis/07_heldout_5fold_cv.py new file mode 100644 index 0000000000000000000000000000000000000000..913b9d9b3669e7668862c9fbc77072b2f4576676 --- /dev/null +++ b/scripts/hematopoiesis/07_heldout_5fold_cv.py @@ -0,0 +1,207 @@ +"""5-fold cv: retrain panda from scratch per fold, eval prototype-cosine on held-out 20%.""" +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 torch +import torch.nn.functional as F +from torch.utils.data import Dataset, DataLoader +from sklearn.model_selection import StratifiedKFold +from sklearn.metrics import (accuracy_score, f1_score, roc_auc_score, + classification_report) + +sys.path.insert(0, "/home/bcheng/PRISM") +from panda.pan_skin.model import ( + PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, prototype_infonce +) + +CORPUS = Path("/home/bcheng/PRISM/data/corpus/hematopoiesis/harmonized/corpus.h5ad") +OUT = Path("/home/bcheng/PRISM/discovery/hematopoiesis/marker") + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +N_FOLDS = 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} + 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_pick = self.rng.choice(self.classes, p=self.p) + batch.append(int(self.rng.choice(self.by_cls[c_pick]))) + yield batch + def __len__(self): return self.n_batches + + +def train_one_fold(X, y, d, mhf, logc, classes, datasets, tr, te, fold_id, log_prefix): + torch.cuda.empty_cache() + Xtr, ytr, dtr, mtr, ltr = X[tr], y[tr], d[tr], mhf[tr], logc[tr] + ds = CorpusDataset(Xtr, ytr, dtr, mtr, ltr) + sampler = HybridSampler(ytr, dtr, n_batches=100, seed=fold_id) + loader = DataLoader(ds, batch_sampler=sampler, num_workers=0) + + counts_per = np.bincount(ytr, minlength=len(classes)) + inv_sqrt = 1.0 / np.sqrt(counts_per + 1) + inv_sqrt = inv_sqrt / inv_sqrt.mean() + class_w_np = 0.5 * inv_sqrt + 0.5 * np.ones_like(inv_sqrt) + class_w = torch.tensor(class_w_np, dtype=torch.float32).to(DEVICE) + + 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): + for e in range(stage_epochs[stage]): + if e % 10 == 0: + print(f"{log_prefix} s{stage} ep {e}/{stage_epochs[stage]}", flush=True) + 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 + 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) + + model.eval() + Xte, yte = X[te], y[te] + with torch.no_grad(): + Xt = torch.from_numpy(Xte.astype(np.float32)).to(DEVICE) + aux = torch.zeros(len(te), 2, device=DEVICE) + out = model(Xt, aux, lam_dann=0.0) + z = out["z"] + cos = z @ model.prototypes.T + pred = cos.argmax(dim=1).cpu().numpy() + probs = torch.softmax(cos / 0.07, dim=1).cpu().numpy() + + acc = accuracy_score(yte, pred) + f1 = f1_score(yte, pred, average="macro", zero_division=0) + try: + auc = roc_auc_score(np.eye(len(classes))[yte], probs, average="macro", multi_class="ovr") + except Exception: + auc = float("nan") + print(f"{log_prefix} acc={acc:.4f} macro_f1={f1:.4f} macro_auc={auc:.4f}", flush=True) + return acc, f1, auc, pred, yte + + +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)} + X = np.asarray(a.obsm["X_pca"]) + 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)]) + 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) + + print(f"[cv] corpus: {a.shape} n_classes={len(classes)} n_datasets={len(datasets)}", + flush=True) + print(f"[cv] class counts: {dict(zip(classes, np.bincount(y, minlength=len(classes)).tolist()))}", + flush=True) + + skf = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=42) + accs, f1s, aucs = [], [], [] + all_preds = [] + partial_path = OUT / "61_hematopoiesis_heldout_5fold_partial.json" + for fold, (tr, te) in enumerate(skf.split(X, y)): + t0 = time.time() + try: + acc, f1, auc, pred, yte = train_one_fold( + X, y, d, mhf, logc, classes, datasets, tr, te, + fold_id=fold, log_prefix=f"[fold {fold}]") + except Exception as exc: + import traceback; traceback.print_exc() + print(f"[fold {fold}] FAILED: {exc}", flush=True) + continue + print(f"[fold {fold}] wall={time.time()-t0:.0f}s", flush=True) + accs.append(acc); f1s.append(f1); aucs.append(auc) + all_preds.append({"fold": fold, "test_idx": te.tolist(), + "pred": pred.tolist(), "true": yte.tolist()}) + with open(partial_path, "w") as f: + json.dump({"folds_done": fold + 1, "accs": accs, "f1s": f1s, "aucs": aucs}, f) + + print(f"\n[cv] 5-FOLD ACC: {np.mean(accs):.4f} +- {np.std(accs):.4f}") + print(f"[cv] 5-FOLD F1: {np.mean(f1s):.4f} +- {np.std(f1s):.4f}") + print(f"[cv] 5-FOLD AUC: {np.mean(aucs):.4f} +- {np.std(aucs):.4f}") + + all_y_true = np.concatenate([np.array(p["true"]) for p in all_preds]) + all_y_pred = np.concatenate([np.array(p["pred"]) for p in all_preds]) + print("\n[cv] Concatenated held-out classification report:") + rep = classification_report(all_y_true, all_y_pred, target_names=classes, + digits=3, zero_division=0, output_dict=True) + print(classification_report(all_y_true, all_y_pred, target_names=classes, + digits=3, zero_division=0)) + + result = { + "mean_acc": float(np.mean(accs)), "std_acc": float(np.std(accs)), + "mean_f1": float(np.mean(f1s)), "std_f1": float(np.std(f1s)), + "mean_auc": float(np.mean(aucs)), "std_auc": float(np.std(aucs)), + "per_fold_acc": accs, "per_fold_f1": f1s, "per_fold_auc": aucs, + "per_class_report": rep, + "n_folds": N_FOLDS, "n_classes": len(classes), + "protocol": "StratifiedKFold(5) retrain from scratch per fold; " + "eval by prototype-cosine argmax on held-out 20%.", + } + (OUT / "61_hematopoiesis_heldout_5fold_cv.json").write_text(json.dumps(result, indent=2)) + print(f"\n[cv] wrote {OUT}/61_hematopoiesis_heldout_5fold_cv.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/hematopoiesis/09_retrain_with_nestorowa_anchor.py b/scripts/hematopoiesis/09_retrain_with_nestorowa_anchor.py new file mode 100644 index 0000000000000000000000000000000000000000..8516c4942526547d2912cb93d349b077b6c8a64c --- /dev/null +++ b/scripts/hematopoiesis/09_retrain_with_nestorowa_anchor.py @@ -0,0 +1,273 @@ +"""retrain pan-hsc panda with nestorowa anchor (150 LT-HSC + 600 HSPC in train, rest held out).""" +from pathlib import Path +import warnings, json, sys, pickle, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp, torch, torch.nn.functional as F, yaml +warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 +sys.path.insert(0, "/home/bcheng/PRISM") +from panda import (PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, + subcenter_angular_infonce, prototype_repulsion) +from sklearn.decomposition import PCA +from sklearn.metrics import accuracy_score, f1_score, classification_report + +ROOT = Path("/home/bcheng/PRISM") +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +CORP_DIR = ROOT / "data/corpus/hematopoiesis/harmonized" +NEST_H5 = ROOT / "data/raw/nestorowa_combined.h5ad" + + +def split_nestorowa(seed=0): + a = ad.read_h5ad(NEST_H5) + a.obs["_orig_idx"] = np.arange(a.n_obs) + rng = np.random.default_rng(seed) + lt = np.where(a.obs["cell_type"].astype(str).values == "LT-HSC")[0] + hs = np.where(a.obs["cell_type"].astype(str).values == "HSPC")[0] + uk = np.where(a.obs["cell_type"].astype(str).values == "unknown")[0] + lt_anchor = rng.choice(lt, size=150, replace=False) + hs_anchor = rng.choice(hs, size=600, replace=False) + anchor_ix = np.concatenate([lt_anchor, hs_anchor]) + test_ix = np.setdiff1d(np.arange(a.n_obs), anchor_ix) + print(f"[nest-split] anchor n={len(anchor_ix)} (LT-HSC 150 + HSPC 600); " + f"held-out n={len(test_ix)} (LT-HSC {(a.obs['cell_type'][test_ix]=='LT-HSC').sum()} " + f"+ HSPC {(a.obs['cell_type'][test_ix]=='HSPC').sum()} " + f"+ unknown {(a.obs['cell_type'][test_ix]=='unknown').sum()})", flush=True) + return a, anchor_ix, test_ix + + +def nest_to_corpus_schema(a_nest, anchor_ix): + a_anchor = a_nest[anchor_ix].copy() + # LT-HSC kept as-is so it enters vocab; HSPC folded into MPP + lab_map = {"LT-HSC": "LT-HSC", "HSPC": "MPP"} + a_anchor.obs["canonical_label"] = a_anchor.obs["cell_type"].astype(str).map(lab_map) + a_anchor.obs["dataset"] = "nestorowa_GSE81682_anchor" + a_anchor.obs["condition"] = "unknown" + a_anchor.obs["sample"] = "nestorowa_anchor" + a_anchor.obs["stage_tag"] = "adult" + a_anchor.obs["organism"] = "mouse" + keep_obs = ["canonical_label", "dataset", "condition", "sample", "stage_tag", "organism"] + for c in list(a_anchor.obs.columns): + if c not in keep_obs: + del a_anchor.obs[c] + return a_anchor + + +def build_anchor_corpus(): + base = ad.read_h5ad(CORP_DIR / "corpus.h5ad") + print(f"[build] base corpus {base.shape}", flush=True) + a_nest, anchor_ix, test_ix = split_nestorowa() + a_anchor = nest_to_corpus_schema(a_nest, anchor_ix) + common_genes = sorted(set(base.var_names.astype(str)) & set(a_anchor.var_names.astype(str))) + a_anchor_c = a_anchor[:, common_genes].copy() + base_c = base[:, common_genes].copy() + for col in base_c.obs.columns: + if col not in a_anchor_c.obs.columns: + a_anchor_c.obs[col] = "unknown" + for col in a_anchor_c.obs.columns: + if col not in base_c.obs.columns: + del a_anchor_c.obs[col] + a_full = ad.concat([base_c, a_anchor_c[:, base_c.var_names]], join="outer") + print(f"[build] anchor-augmented corpus {a_full.shape} ({base.n_obs} + {a_anchor_c.n_obs})", flush=True) + # coerce obs to str; h5py vlen serialization crashes otherwise + for c in list(a_full.obs.columns): + try: + a_full.obs[c] = a_full.obs[c].astype(str) + except Exception: + del a_full.obs[c] + a_full.write_h5ad(CORP_DIR / "corpus_with_nest_anchor.h5ad") + + stats = np.load(CORP_DIR / "corpus_stats.npz", allow_pickle=True) + hvgs = [str(g) for g in stats["shared_hvgs"]] + hvg2i = {g: i for i, g in enumerate(hvgs)} + common_hvgs = [g for g in a_full.var_names.astype(str) if g in hvg2i] + a_c = a_full[:, common_hvgs].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_full.n_obs, len(hvgs)), dtype=np.float32) + Xf[:, np.array([hvg2i[g] for g in common_hvgs])] = X + mu = Xf.mean(axis=0); sig = Xf.std(axis=0) + 1e-6 + np.savez(CORP_DIR / "corpus_stats_with_nest_anchor.npz", + shared_hvgs=np.asarray(hvgs, dtype=object), mean=mu, std=sig) + Xz = np.clip((Xf - mu) / sig, -10, 10) + rng = np.random.default_rng(0) + idx = rng.choice(Xz.shape[0], size=min(30000, Xz.shape[0]), replace=False) + pca = PCA(n_components=50, random_state=0).fit(Xz[idx]) + with open(CORP_DIR / "pca_basis_with_nest_anchor.pkl", "wb") as f: pickle.dump(pca, f) + print(f"[build] refit stats + pca", flush=True) + return a_full, hvgs, mu, sig, pca, a_nest, test_ix + + +def prepare_batches(adata, hvgs, mu, sig, pca, marker_genes=None, variant="pca"): + 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() + 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) + Xf[:, np.array([hvg2i[g] for g in common])] = 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) + 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(a, hvgs, mu, sig, pca, variant, marker_genes, epochs=8, batch=256, lr=1e-3): + Xpca, Xmark, y, classes, y_dset, datasets, log10cz = prepare_batches(a, hvgs, mu, sig, pca, marker_genes, variant) + print(f"[train] {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) + if stage >= 3: + L = L + 0.5 * prototype_repulsion(model.prototypes.detach().clone()) + 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 {variant}] epoch {epoch}/{epochs} stage={stage} loss={np.mean(losses):.4f}", flush=True) + ck_dir = ROOT / f"checkpoints/hematopoiesis_anchor/{variant}" + ck_dir.mkdir(parents=True, exist_ok=True) + torch.save({"model": model.state_dict(), "classes": classes, "datasets": datasets, + "marker_genes": marker_genes if variant == "marker" else [], + "prototypes": model.prototypes.detach().cpu().numpy()}, + ck_dir / "panda_final.pt") + print(f"[save] {ck_dir}/panda_final.pt", flush=True) + + +COARSE = {"LT-HSC": "LT-HSC", "MPP": "HSPC", "GMP": "HSPC", "myeloid": "HSPC", + "erythroid": "HSPC", "megakaryocyte": "HSPC", "basophil-mast": "HSPC", + "lymphoid": "HSPC", "unassigned": "HSPC", "UNK": "HSPC"} + + +def infer_on_test(variant, marker_genes, hvgs, mu, sig, pca, a_nest, test_ix): + a_test = a_nest[test_ix].copy() + Xpca, Xmark, _, _, _, _, _ = prepare_batches( + a_test, hvgs, mu, sig, pca, marker_genes, variant, + ) if False else (None,)*7 + # a_test has no canonical_label so rebuild manually + hvg2i = {g: i for i, g in enumerate(hvgs)} + common = [g for g in a_test.var_names.astype(str) if g in hvg2i] + a_c = a_test[:, 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_test.n_obs, len(hvgs)), dtype=np.float32) + Xf[:, np.array([hvg2i[g] for g in common])] = 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": + mvals = np.zeros((a_test.n_obs, len(marker_genes)), dtype=np.float32) + for j, g in enumerate(marker_genes): + if g in a_test.var_names: + col = a_test[:, 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) + + ck = torch.load(ROOT / f"checkpoints/hematopoiesis_anchor/{variant}/panda_final.pt", + map_location=DEVICE, weights_only=False) + classes = ck["classes"] + model = PANDAEncoder(variant=variant, n_pca=50, + n_markers=len(marker_genes) if variant == "marker" else 0, + n_classes=len(classes), n_sub=3, + n_datasets=len(ck["datasets"])).to(DEVICE).eval() + model.load_state_dict(ck["model"]) + + preds, probs = [], [] + with torch.no_grad(): + for i in range(0, a_test.n_obs, 4096): + xb = torch.from_numpy(Xpca[i:i+4096]).to(DEVICE) + xmb = torch.from_numpy(Xmark[i:i+4096]).to(DEVICE) if Xmark is not None else None + aux = torch.zeros(len(xb), 2, device=DEVICE) + out = model(xb, aux, x_markers=xmb, lam_dann=0.0) + mc = model.max_sub_cos(out["z"]) + preds.append(mc.argmax(dim=1).cpu().numpy()) + probs.append(F.softmax(mc / 0.07, dim=1).cpu().numpy()) + pred = np.array([classes[i] for i in np.concatenate(preds)]) + probs = np.concatenate(probs) + pred_coarse = np.array([COARSE.get(p, "HSPC") for p in pred]) + + y_true = a_test.obs["cell_type"].astype(str).values + mask = y_true != "unknown" + acc = accuracy_score(y_true[mask], pred_coarse[mask]) + f1 = f1_score(y_true[mask], pred_coarse[mask], average="macro", zero_division=0) + rep = classification_report(y_true[mask], pred_coarse[mask], zero_division=0, output_dict=True) + lt_recall = float(((pred_coarse == "LT-HSC") & (y_true == "LT-HSC")).sum() / max(1, (y_true == "LT-HSC").sum())) + print(f"[eval-{variant}] n_labeled={mask.sum()} coarse-acc={acc:.4f} macro-f1={f1:.4f} LT-HSC-recall={lt_recall:.4f}", flush=True) + + out = ROOT / f"discovery/hematopoiesis/{variant}" + out.mkdir(parents=True, exist_ok=True) + (out / "97_nestorowa_anchor_zero_shot.json").write_text(json.dumps({ + "variant": variant, + "n_test_labeled": int(mask.sum()), + "n_test_total": int(a_test.n_obs), + "coarse_acc": float(acc), "coarse_f1": float(f1), "LT_HSC_recall": lt_recall, + "coarse_per_class": rep, + "fine_pred_dist": pd.Series(pred).value_counts().to_dict(), + "true_dist_labeled": pd.Series(y_true[mask]).value_counts().to_dict(), + }, indent=2, default=str)) + pd.DataFrame({"cell_id": a_test.obs_names, "facs_gate": y_true, + "pred_fine": pred, "pred_coarse": pred_coarse, + "max_cos": probs.max(axis=1)}).to_csv(out / "97_nestorowa_anchor_predictions.csv", index=False) + + +def main(): + print("[1/3] build anchor corpus", flush=True) + a, hvgs, mu, sig, pca, a_nest, test_ix = build_anchor_corpus() + + print("\n[2/3] train pca variant", flush=True) + train(a, hvgs, mu, sig, pca, "pca", [], epochs=8) + + print("\n[3/3] train marker variant", flush=True) + marker_genes = yaml.safe_load(open(ROOT / "panda/markers.yaml"))["pan_hematopoietic"] \ + if "pan_hematopoietic" in yaml.safe_load(open(ROOT / "panda/markers.yaml")) \ + else yaml.safe_load(open(ROOT / "panda/markers.yaml"))["hematopoiesis"] + train(a, hvgs, mu, sig, pca, "marker", marker_genes, epochs=8) + + print("\n[eval] pca on held-out nestorowa slice", flush=True) + infer_on_test("pca", [], hvgs, mu, sig, pca, a_nest, test_ix) + print("\n[eval] marker on held-out nestorowa slice", flush=True) + infer_on_test("marker", marker_genes, hvgs, mu, sig, pca, a_nest, test_ix) + + +if __name__ == "__main__": + main() diff --git a/scripts/hematopoiesis/10_build_corpus.py b/scripts/hematopoiesis/10_build_corpus.py new file mode 100644 index 0000000000000000000000000000000000000000..7045a4fa0b02b74163180a806f559beb672c34c3 --- /dev/null +++ b/scripts/hematopoiesis/10_build_corpus.py @@ -0,0 +1,495 @@ +"""pan-hsc corpus v3: paper-label first (weinreb + baccin + tabula muris senis marrow).""" +from __future__ import annotations +from pathlib import Path +import warnings, sys, json, pickle, gzip, shutil, subprocess +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 sklearn.decomposition import PCA + +sc.settings.verbosity = 0 + +ROOT = Path("/home/bcheng/PRISM") +HARM = ROOT / "data/corpus/hematopoiesis/harmonized" +TIER_V2 = ROOT / "data/corpus/hematopoiesis/tier_v2" +BACCIN_DIR = TIER_V2 / "baccin_GSE122465_wholeBM_RAW_extracted" +TMS_DIR = TIER_V2 / "tms_GSE132042_RAW_extracted" +BACCIN_TSV = ROOT / "data/external_labels/baccin/baccin_cluster_identity.tsv" +WEINREB_H5 = ROOT / "data/processed/weinreb_larry/adata_weinreb.h5ad" +SCRATCH = Path("/tmp/prism_corpus_scratch") +SCRATCH.mkdir(parents=True, exist_ok=True) + +MIN_CELLS_PER_CLASS = 150 +K_TOP_PER_DS = 4000 +K_TARGET = 3500 +N_PCA = 50 +SAMPLE_PER_DS = 5000 + +# canonical hematopoiesis markers to force-include +MUST_INCLUDE = [ + 'Procr','Fgd5','Hlf','Mllt3','Mecom','Meis1','Mpl','Slamf1','Lmo2', + 'Cd48','Flt3','Cebpa','Csf1r','Gata1','Cd44', + 'Klf1','Gata2','Zfpm1','Epor','Mpo','Elane','Cebpe','Cd177','Ms4a3', + 'Hbb-b1','Hbb-b2','Hba-a1','Hba-a2','Gypa','S100a8','S100a9','Csf3r','Ltf', + 'Cd19','Ms4a1','Ebf1','Rag1','Il7r','Itga2b','Pf4','Gp1ba','Nfe2', + 'Cpa3','Prss34','Mrgprb1','Kit','Mcpt8', + 'Ptprc','Cd3e','Cd8a','Cd4','Adgre1','Itgam', +] + +# dataset-specific paper label -> canonical; "" = abstain +WEINREB_MAP = { + "Undifferentiated": "MPP", + "Neutrophil": "myeloid", + "Monocyte": "myeloid", + "Ccr7_DC": "myeloid", + "pDC": "myeloid", + "Baso": "basophil-mast", + "Mast": "basophil-mast", + "Eos": "basophil-mast", + "Meg": "megakaryocyte", + "Lymphoid": "lymphoid", + "Erythroid": "erythroid", +} + +BACCIN_MAP = { + "LMPPs": "MPP", + "Ery prog.": "erythroid", + "Erythroblasts": "erythroid", + "Ery/Mk prog.": "erythroid", + "Mk prog.": "megakaryocyte", + "Gran/Mono prog.": "myeloid", + "Neutrophils": "myeloid", + "Neutro prog.": "myeloid", + "Mono prog.": "monocyte", + "Monocytes": "monocyte", + "Dendritic cells": "dendritic", + "Eo/Baso prog.": "basophil-mast", + "pro-B": "pro-B", + "large pre-B.": "pro-B", + "small pre-B.": "pro-B", + "B cell": "naive-B", + "T cells": "T-cell", + "NK cells": "lymphoid", + "Arteriolar ECs": "endothelial", + "Sinusoidal ECs": "endothelial", + "Arteriolar fibro.": "fibroblast", + "Endosteal fibro.": "fibroblast", + "Stromal fibro.": "fibroblast", + "Myofibroblasts": "fibroblast", + "Fibro/Chondro p.": "fibroblast", + "Adipo-CAR": "stromal", + "Osteo-CAR": "osteoblast", + "Osteoblasts": "osteoblast", + "Ng2+ MSCs": "stromal", + "Chondrocytes": "stromal", + # excluded from vocab (non-hematopoietic / non-BM stroma) + "Schwann cells": "", + "Smooth muscle": "", +} + +TMS_MAP = { + "hematopoietic stem cell": "LT-HSC", + "hematopoietic precursor cell": "MPP", + "naive B cell": "naive-B", + "immature B cell": "naive-B", + "plasma cell": "naive-B", + "precursor B cell": "pro-B", + "late pro-B cell": "pro-B", + "early pro-B cell": "pro-B", + "granulocyte": "myeloid", + "granulocytopoietic cell": "myeloid", + "granulocyte monocyte progenitor cell": "myeloid", + "promonocyte": "monocyte", + "monocyte": "monocyte", + "macrophage": "macrophage", + "basophil": "basophil-mast", + "NK cell": "lymphoid", + "lymphoid progenitor cell": "lymphoid", + "mature alpha-beta T cell": "T-cell", + "CD4-positive, alpha-beta T cell": "T-cell", + "naive T cell": "T-cell", + "megakaryocyte-erythroid progenitor cell": "erythroid", + "erythroid progenitor": "erythroid", + "proerythroblast": "erythroid", + "erythroblast": "erythroid", +} + + +def qc(a, name): + n0 = a.n_obs + sc.pp.filter_cells(a, min_genes=200) + sc.pp.filter_genes(a, min_cells=3) + print(f" [{name}] QC: {n0} -> {a.n_obs} cells, {a.n_vars} genes", flush=True) + return a + + +def load_weinreb() -> ad.AnnData: + a = ad.read_h5ad(WEINREB_H5) + if a.raw is not None: + raw = a.raw.to_adata(); raw.obs = a.obs.copy(); a = raw + a.var_names_make_unique() + if not sp.issparse(a.X): + a.X = sp.csr_matrix(a.X) + a.X = a.X.astype("float32") + a.obs["dataset"] = "weinreb_GSE140802" + a.obs["paper_label"] = a.obs["Cell type annotation"].astype(str) + a.obs["source_sample"] = a.obs["Library"].astype(str) if "Library" in a.obs else "weinreb" + a = qc(a, "weinreb") + return a + + +# baccin csvs have rows=genes so always transpose +def load_baccin() -> ad.AnnData: + tsv = pd.read_csv(BACCIN_TSV, sep="\t") + label_map = dict(zip(tsv["barcode"].astype(str), tsv["cluster_label"].astype(str))) + + # sample-file -> experiment prefix used in TSV / column names + sample_to_exp = { + "GSM3466897_BM.total": "2017_9_totalBM", + "GSM3466898_BM.Lin-Kit+": "2018_2_HSPC", + "GSM3466899_BM.Lin-CD45-": "2017_9_CD45minus", + "GSM3466900_BM.Lin-CD45-CD71-": "2017_12_BM", + "GSM3466901_BoneAssociated.Lin-CD45-CD71-": "2017_12_Bone", + } + + parts = [] + for csv in sorted(BACCIN_DIR.glob("GSM*.csv.gz")): + sample = csv.name.replace(".csv.gz", "") + try: + df = pd.read_csv(csv, index_col=0) + except Exception as e: + print(f"[skip Baccin {sample}] {e}", flush=True); continue + df = df.T + X = sp.csr_matrix(df.values.astype(np.float32)) + a = ad.AnnData(X=X, + obs=pd.DataFrame(index=df.index.astype(str)), + var=pd.DataFrame(index=df.columns.astype(str))) + a.obs["source_sample"] = sample + a.obs["dataset"] = "baccin_GSE122465" + + # column names are already {experiment}_{16nt}, matches TSV barcodes + bc = a.obs_names.astype(str) + paper = np.array([label_map.get(b, "") for b in bc], dtype=object) + a.obs["paper_label"] = paper + + a.obs.index = [f"{sample}__{b}" for b in bc] + a.var_names_make_unique() + parts.append(a) + + a = ad.concat(parts, join="outer", index_unique=None) + a = qc(a, "baccin") + + exp_present = set(sample_to_exp[s] for s in a.obs["source_sample"].unique() + if s in sample_to_exp) + tsv_expected = tsv[tsv["experiment"].isin(exp_present)].shape[0] + covered = (a.obs["paper_label"].astype(str).str.len() > 0).sum() + print(f" [baccin] TSV coverage: {covered}/{a.n_obs} cells " + f"(TSV rows in loaded experiments={tsv_expected})", flush=True) + return a + + +def _gunzip(gz_path: Path, out_path: Path) -> Path: + if not out_path.exists(): + print(f" [tms] gunzip {gz_path.name}", flush=True) + with gzip.open(gz_path, "rb") as f_in, open(out_path, "wb") as f_out: + shutil.copyfileobj(f_in, f_out) + return out_path + + +def _load_tms_official(gz: Path, method_tag: str) -> ad.AnnData | None: + tmp = SCRATCH / gz.name.replace(".gz", "") + _gunzip(gz, tmp) + a_all = ad.read_h5ad(tmp) + if "tissue" not in a_all.obs.columns: + print(f" [tms {method_tag}] no tissue column, skipping", flush=True); return None + mask = a_all.obs["tissue"].astype(str).values == "Marrow" + a = a_all[mask].copy(); del a_all + a.var_names_make_unique() + if not sp.issparse(a.X): + a.X = sp.csr_matrix(a.X) + a.X = a.X.astype("float32") + a.obs["dataset"] = "tms_GSE132042" + a.obs["source_sample"] = a.obs.get("mouse.id", pd.Series(method_tag, index=a.obs.index)) \ + .astype(str) + f"_{method_tag}" + a.obs["paper_label"] = a.obs["cell_ontology_class"].astype(str) + a.obs.index = [f"tms_{method_tag}_{i}" for i in a.obs.index.astype(str)] + print(f" [tms {method_tag}] {a.n_obs} Marrow cells, " + f"{a.obs['paper_label'].astype(str).ne('').sum()} labeled", flush=True) + return a + + +def load_tms() -> ad.AnnData: + facs_gz = TMS_DIR / "GSM4505405_tabula-muris-senis-facs-official-raw-obj.h5ad.gz" + dropl_gz = TMS_DIR / "GSM4505404_tabula-muris-senis-droplet-official-raw-obj.h5ad.gz" + parts = [] + for gz, tag in [(facs_gz, "facs"), (dropl_gz, "droplet")]: + if gz.exists(): + a = _load_tms_official(gz, tag) + if a is not None: parts.append(a) + else: + print(f" [tms] missing {gz.name}", flush=True) + if not parts: + raise RuntimeError("no TMS official h5ad found") + a = ad.concat(parts, join="outer", index_unique=None) + a = qc(a, "tms") + return a + + +def canonicalize(a: ad.AnnData, mapping: dict) -> ad.AnnData: + pl = a.obs["paper_label"].astype(str).values + canon = np.empty(len(pl), dtype=object) + source = np.empty(len(pl), dtype=object) + for i, p in enumerate(pl): + p_norm = p.strip() + if p_norm in {"", "nan", "NaN", "UNK", "unassigned", "Unknown", "unknown"}: + canon[i] = "" + source[i] = "abstain" + elif p_norm in mapping: + m = mapping[p_norm] + if m == "": + canon[i] = "" + source[i] = "abstain" + else: + canon[i] = m + source[i] = "paper" + else: + # unmapped labels must abstain — never leak into vocab + canon[i] = "" + source[i] = "abstain" + a.obs["canonical_label"] = canon + a.obs["label_source"] = source + return a + + +def rank_hvgs(a: ad.AnnData) -> pd.Series: + x = a.copy() + try: + sc.pp.highly_variable_genes(x, n_top_genes=K_TOP_PER_DS, flavor="seurat_v3", + subset=False, check_values=False) + if "variances_norm" in x.var: + return x.var["variances_norm"].fillna(-np.inf) + return x.var["dispersions_norm"].fillna(-np.inf) + except Exception: + sc.pp.normalize_total(x, target_sum=1e4); sc.pp.log1p(x) + X = x.X.toarray() if sp.issparse(x.X) else x.X + return pd.Series(np.asarray(X.var(axis=0)).ravel(), index=x.var_names) + + +def build_hvgs(ds_list: list[ad.AnnData]) -> list[str]: + ranks = {} + for i, a in enumerate(ds_list): + name = str(a.obs["dataset"].iloc[0]) + key = f"{name}_{i}" + print(f"[hvg] rank {key}", flush=True) + ranks[key] = rank_hvgs(a) + R = pd.DataFrame(ranks) + top_hits = pd.DataFrame({d: R[d].rank(ascending=False, method="min") <= K_TOP_PER_DS + for d in R.columns}).fillna(False) + hit_count = top_hits.sum(axis=1) + eligible = hit_count[hit_count >= max(1, len(ds_list) // 2)] + print(f"[hvg] {len(eligible)} genes top-{K_TOP_PER_DS} in >=half datasets", flush=True) + mean_score = R.mean(axis=1, skipna=True) + shared_by_score = mean_score.loc[eligible.index].sort_values(ascending=False).index + all_genes = set(R.index) + forced = [g for g in MUST_INCLUDE if g in all_genes] + print(f"[hvg] forced markers present: {len(forced)}/{len(MUST_INCLUDE)}", flush=True) + picked = list(forced) + for g in shared_by_score: + if g not in picked: + picked.append(g) + if len(picked) >= K_TARGET: + break + return picked[:K_TARGET] + + +def normalize_and_stats(ds_list, shared_genes): + G = len(shared_genes) + gene_idx = {g: i for i, g in enumerate(shared_genes)} + running_sum = np.zeros(G, dtype=np.float64) + running_sq = np.zeros(G, dtype=np.float64) + N = 0 + for a in ds_list: + common = [g for g in shared_genes if g in a.var_names] + a_s = a[:, common].copy() + sc.pp.normalize_total(a_s, target_sum=1e4); sc.pp.log1p(a_s) + X = a_s.X.toarray() if sp.issparse(a_s.X) else a_s.X + cols = np.array([gene_idx[g] for g in common]) + running_sum[cols] += X.sum(axis=0) + running_sq[cols] += (X ** 2).sum(axis=0) + N += X.shape[0] + del a_s, X + mu = running_sum / N + sig = np.sqrt(np.maximum(running_sq / N - mu ** 2, 1e-6)) + return mu, sig, N, gene_idx + + +def transform_dataset(a, shared_genes, gene_idx, mu, sig): + G = len(shared_genes) + common = [g for g in shared_genes if g in a.var_names] + a_s = a[:, common].copy() + raw_counts = np.asarray(a_s.X.sum(axis=1)).ravel() if sp.issparse(a_s.X) else a_s.X.sum(axis=1) + sc.pp.normalize_total(a_s, target_sum=1e4); sc.pp.log1p(a_s) + X = a_s.X.toarray() if sp.issparse(a_s.X) else a_s.X + Xf = np.zeros((X.shape[0], G), dtype=np.float32) + cols = np.array([gene_idx[g] for g in common]) + Xf[:, cols] = X.astype(np.float32) + Xz = np.clip((Xf - mu.astype(np.float32)) / sig.astype(np.float32), -10, 10) + return Xf, Xz, raw_counts, common + + +def main(): + print("=== corpus v3 build ===", flush=True) + print("[1/6] loading Weinreb …", flush=True) + a_wei = load_weinreb() + a_wei = canonicalize(a_wei, WEINREB_MAP) + + print("[2/6] loading Baccin …", flush=True) + a_bac = load_baccin() + a_bac = canonicalize(a_bac, BACCIN_MAP) + + print("[3/6] loading TMS …", flush=True) + a_tms = load_tms() + a_tms = canonicalize(a_tms, TMS_MAP) + + ds_list = [a_wei, a_bac, a_tms] + + cov_rows = [] + for a in ds_list: + name = str(a.obs["dataset"].iloc[0]) + n = a.n_obs + n_paper = int((a.obs["label_source"] == "paper").sum()) + n_abstain = int((a.obs["label_source"] == "abstain").sum()) + cov_rows.append({"dataset": name, "n_cells": n, + "n_paper": n_paper, "n_abstain": n_abstain, + "paper_frac": round(n_paper / n, 3)}) + cov_df_pre = pd.DataFrame(cov_rows) + print("\n[coverage pre-drop]\n" + cov_df_pre.to_string(index=False), flush=True) + + combined_labels = np.concatenate([a.obs["canonical_label"].astype(str).values for a in ds_list]) + from collections import Counter + label_ct = Counter([l for l in combined_labels if l]) + print("\n[class-count] combined (before rare-drop):", flush=True) + for k, v in sorted(label_ct.items(), key=lambda x: -x[1]): + print(f" {k:16s} {v}", flush=True) + + kept = {k for k, v in label_ct.items() if v >= MIN_CELLS_PER_CLASS} + dropped = {k: v for k, v in label_ct.items() if v < MIN_CELLS_PER_CLASS} + print(f"\n[class-drop] < {MIN_CELLS_PER_CLASS} cells -> abstain:", flush=True) + for k, v in sorted(dropped.items(), key=lambda x: -x[1]): + print(f" {k:16s} {v}", flush=True) + print(f"[class-final] {len(kept)} kept: {sorted(kept)}", flush=True) + + for a in ds_list: + canon = a.obs["canonical_label"].astype(str).values.copy() + source = a.obs["label_source"].astype(str).values.copy() + for i, c in enumerate(canon): + if c and c not in kept: + canon[i] = "" + source[i] = "abstain" + a.obs["canonical_label"] = canon + a.obs["label_source"] = source + + print("\n[4/6] ranking HVGs …", flush=True) + shared_genes = build_hvgs(ds_list) + print(f"[hvg] final HVG count: {len(shared_genes)}", flush=True) + + print("[5/6] computing per-gene mean/std …", flush=True) + mu, sig, N, gene_idx = normalize_and_stats(ds_list, shared_genes) + print(f"[stats] mean/std over {N} cells", flush=True) + + print("[5/6] sample-fitting PCA …", flush=True) + rng = np.random.default_rng(0) + fit_chunks = [] + ds_transforms = [] + for a in ds_list: + Xf, Xz, raw_counts, common = transform_dataset(a, shared_genes, gene_idx, mu, sig) + ds_transforms.append((Xf, Xz, raw_counts, common)) + n = Xz.shape[0] + take = min(SAMPLE_PER_DS, n) + idx = rng.choice(n, size=take, replace=False) + fit_chunks.append(Xz[idx]) + Xfit = np.vstack(fit_chunks); del fit_chunks + pca = PCA(n_components=N_PCA, random_state=42).fit(Xfit) + print(f"[pca] explained var={pca.explained_variance_ratio_.sum():.3f}", flush=True) + # sign fix so components are stable across runs + for k in range(N_PCA): + top = int(np.argmax(np.abs(pca.components_[k]))) + if pca.components_[k, top] < 0: pca.components_[k] *= -1 + + print("\n[6/6] emitting corpus …", flush=True) + parts = [] + for a, (Xf, Xz, raw_counts, common) in zip(ds_list, ds_transforms): + Z = pca.transform(Xz).astype(np.float32) + obs = a.obs.copy() + obs["total_counts"] = raw_counts.astype(np.float32) + obs["missing_hvg_frac"] = 1.0 - len(common) / len(shared_genes) + a2 = ad.AnnData(X=sp.csr_matrix(Xf.astype(np.float32)), obs=obs, + var=pd.DataFrame(index=shared_genes)) + a2.obsm["X_pca"] = Z + parts.append(a2) + corpus = ad.concat(parts, join="outer", label="_batch", index_unique=None) + corpus.uns["shared_hvgs"] = shared_genes + + # coerce obs to str; h5py vlen crashes otherwise (mirrors 09_) + for c in list(corpus.obs.columns): + try: + corpus.obs[c] = corpus.obs[c].astype(str) + except Exception: + del corpus.obs[c] + + HARM.mkdir(parents=True, exist_ok=True) + out_h5 = HARM / "corpus_v3.h5ad" + corpus.write_h5ad(out_h5, compression="gzip") + print(f"[emit] wrote {out_h5} ({corpus.n_obs} x {corpus.n_vars})", flush=True) + + np.savez(HARM / "corpus_stats_v3.npz", + shared_hvgs=np.asarray(shared_genes, dtype=object), + mean=mu.astype(np.float32), std=sig.astype(np.float32), n_cells=N) + with open(HARM / "pca_basis_v3.pkl", "wb") as fh: + pickle.dump(pca, fh) + print(f"[emit] wrote corpus_stats_v3.npz + pca_basis_v3.pkl", flush=True) + + cov_rows_post = [] + for a in ds_list: + name = str(a.obs["dataset"].iloc[0]) + source = a.obs["label_source"].astype(str).values + n = a.n_obs + n_paper = int((source == "paper").sum()) + n_abstain = int((source == "abstain").sum()) + cov_rows_post.append({"dataset": name, "n_cells": n, + "n_paper": n_paper, "n_abstain": n_abstain, + "paper_frac": round(n_paper / n, 3)}) + cov_df = pd.DataFrame(cov_rows_post) + cov_df.to_csv(HARM / "per_dataset_coverage_v3.tsv", sep="\t", index=False) + print("\n[coverage post-drop]\n" + cov_df.to_string(index=False), flush=True) + + print("\n=== verification ===", flush=True) + cl = corpus.obs["canonical_label"].astype(str) + src = corpus.obs["label_source"].astype(str) + bad_tokens = {"UNK", "unassigned", "Unknown", "unknown", "nan", "NaN"} + leaked = cl[cl.isin(bad_tokens)] + print(f" leaked bad tokens in canonical_label: {len(leaked)} (must be 0)", flush=True) + assert len(leaked) == 0 + + final_counts = cl[cl != ""].value_counts() + print(f" final canonical classes ({len(final_counts)}):", flush=True) + print(final_counts.to_string(), flush=True) + + assert (final_counts >= MIN_CELLS_PER_CLASS).all(), \ + "some final class fell below threshold; verify class-drop logic" + + class_names = list(final_counts.index) + assert len(class_names) == len(set(class_names)), "duplicate class names found" + + print(f" label_source counts: {src.value_counts().to_dict()}", flush=True) + assert ((cl != "") == (src == "paper")).all(), "label_source / canonical_label mismatch" + + print("\nALL VERIFICATIONS PASSED", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/hematopoiesis/11_filter_paper_only.py b/scripts/hematopoiesis/11_filter_paper_only.py new file mode 100644 index 0000000000000000000000000000000000000000..800c24e34ad428af3ebd6accfa9ab1e86b816510 --- /dev/null +++ b/scripts/hematopoiesis/11_filter_paper_only.py @@ -0,0 +1,56 @@ +"""filter corpus_v3 to paper-only cells; refit hvg stats + pca.""" +from pathlib import Path +import warnings, pickle, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp +warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 +from sklearn.decomposition import PCA + +ROOT = Path("/home/bcheng/PRISM") +CORP = ROOT / "data/corpus/hematopoiesis/harmonized" + +print("[load] corpus_v3.h5ad", flush=True) +a = ad.read_h5ad(CORP / "corpus_v3.h5ad") +print(f"[start] {a.n_obs:,} cells, label_source dist: {a.obs['label_source'].value_counts().to_dict()}", flush=True) + +mask = (a.obs["label_source"].astype(str) == "paper").values +b = a[mask].copy() +print(f"[filter] kept {b.n_obs:,} paper cells (dropped {int((~mask).sum()):,})", flush=True) + +counts = b.obs["canonical_label"].astype(str).value_counts() +print(f"[classes] {len(counts)} classes present, min = {counts.min()}, max = {counts.max()}") +# drop classes below 150 cells +keep_classes = counts[counts >= 150].index.tolist() +if len(keep_classes) < len(counts): + dropped_cls = counts[counts < 150].index.tolist() + print(f"[drop] {len(dropped_cls)} classes < 150 cells: {dropped_cls}", flush=True) + b = b[b.obs["canonical_label"].astype(str).isin(keep_classes)].copy() + print(f"[filter] kept {b.n_obs:,} cells across {len(keep_classes)} classes", flush=True) + +b.write_h5ad(CORP / "corpus_v3.h5ad") +print(f"[save] {CORP / 'corpus_v3.h5ad'} ({b.n_obs:,} × {b.n_vars:,})", flush=True) + +old_stats = np.load(CORP / "corpus_stats_v3.npz", allow_pickle=True) +hvgs = [str(g) for g in old_stats["shared_hvgs"]] +hvg2i = {g: i for i, g in enumerate(hvgs)} +common = [g for g in b.var_names.astype(str) if g in hvg2i] +b_c = b[:, common].copy() +sc.pp.normalize_total(b_c, target_sum=1e4); sc.pp.log1p(b_c) +X = b_c.X.toarray().astype(np.float32) if sp.issparse(b_c.X) else b_c.X.astype(np.float32) +Xf = np.zeros((b.n_obs, len(hvgs)), dtype=np.float32) +Xf[:, np.array([hvg2i[g] for g in common])] = X +mu = Xf.mean(axis=0); sig = Xf.std(axis=0) + 1e-6 +np.savez(CORP / "corpus_stats_v3.npz", shared_hvgs=np.asarray(hvgs, dtype=object), mean=mu, std=sig) +Xz = np.clip((Xf - mu) / sig, -10, 10) +rng = np.random.default_rng(0) +idx = rng.choice(Xz.shape[0], size=min(30000, Xz.shape[0]), replace=False) +pca = PCA(n_components=50, random_state=0).fit(Xz[idx]) +with open(CORP / "pca_basis_v3.pkl", "wb") as f: pickle.dump(pca, f) +print(f"[save] refit stats + pca (explained variance {pca.explained_variance_ratio_.sum():.4f})", flush=True) + +print(f"\n[final] {b.n_obs:,} cells, {len(counts)} classes:") +for cls, n in b.obs["canonical_label"].value_counts().items(): + print(f" {cls:<20} {n:>8,}") + +print(f"\n[per-dataset]:") +for d in sorted(b.obs["dataset"].astype(str).unique()): + sub = b.obs[b.obs["dataset"] == d] + print(f" {d:<30} n={len(sub):>8,} classes={sub['canonical_label'].nunique()}") diff --git a/scripts/hematopoiesis/README.md b/scripts/hematopoiesis/README.md new file mode 100644 index 0000000000000000000000000000000000000000..a396166650894b2029677ec6829c6c579f6b38a0 --- /dev/null +++ b/scripts/hematopoiesis/README.md @@ -0,0 +1,39 @@ +# Pan-Hematopoiesis PANDA + +Multi-dataset HSC + hematopoietic lineage corpus for cell identity + zero-shot transfer + +mechanistic discovery. + +## Corpus datasets + +| Accession | Study | Cells | Species | Modality | Role | +|---|---|---:|---|---|---| +| GSE140802 | Weinreb LARRY 2020 | 75,339 | mouse | 10x + clonal barcodes | Anchor (local) | +| GSE72857 | Paul 2015 | ~10,000 | mouse | MARS-seq | Foundational myeloid branch reference | +| GSE81682 | Nestorowa 2016 | ~1,656 | mouse | Smart-seq2 | FACS-index-sorted HSPC anchor | +| GSE89754 | Tusi 2018 | ~4,763 | mouse | inDrops | Erythroid trajectory + PBA fate probs | +| GSE107727 | Dahlin 2018 | ~44,802 | mouse | 10x | Large LSK/Kit+ landscape + Kit-mutant contrast | + +## Held-out labeled validation + +**Nestorowa 2016 index-sort labels** — held-out from training and evaluated with the +FACS-index-sorted ground truth (LT-HSC / MPP / GMP / MEP / CLP) which is orthogonal to +mRNA. Provides a clean cross-lab / cross-platform accuracy measurement. + +## Held-out unlabeled discovery target + +**Dahlin 2018 W41-mutant subset** — subset of GSE107727 where Kit-mutant hematopoiesis is +stress-tested. Weakly labeled; PANDA predictions probe how mutation shifts the identity +distribution. + +## Ontology (10 classes) + +- LT-HSC +- MPP (multipotent progenitor) +- CMP (common myeloid progenitor) +- MEP (megakaryocyte-erythroid progenitor) +- GMP (granulocyte-monocyte progenitor) +- erythroid +- myeloid (Neu + Mono + DC merged) +- lymphoid (B/T/NK) +- megakaryocyte +- basophil-mast diff --git a/scripts/hematopoiesis/__init__.py b/scripts/hematopoiesis/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/scripts/hematopoiesis/_loaders/__init__.py b/scripts/hematopoiesis/_loaders/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/scripts/hematopoiesis/known_hsc_markers.yaml b/scripts/hematopoiesis/known_hsc_markers.yaml new file mode 100644 index 0000000000000000000000000000000000000000..72f8e7e5d7373fa10432f8ad243bc237453b1dcd --- /dev/null +++ b/scripts/hematopoiesis/known_hsc_markers.yaml @@ -0,0 +1,18 @@ +# canonical mouse HSC lineage markers for corpus labeling. + +classes: + LT-HSC: [Procr, Fgd5, Hlf, Mllt3, Mecom, Meis1, Mpl, Slamf1, Lmo2] + MPP: [Cd48, Flt3, Slamf1, Meis1] + CMP: [Cebpa, Csf1r, Gata1, Cd44] + MEP: [Gata1, Klf1, Gata2, Zfpm1, Epor] + GMP: [Mpo, Elane, Cebpe, Cd177, Ms4a3] + erythroid: [Hbb-b1, Hbb-b2, Hba-a1, Hba-a2, Klf1, Gypa] + myeloid: [Elane, Mpo, S100a8, S100a9, Csf3r, Ltf] + lymphoid: [Cd19, Ms4a1, Ebf1, Rag1, Il7r] + megakaryocyte: [Itga2b, Pf4, Gp1ba, Zfpm1, Nfe2] + basophil-mast: [Cpa3, Prss34, Mrgprb1, Kit, Mcpt8] + +assignment: + min_score: 0.10 + min_margin: 0.05 + cluster_resolution: 0.8 diff --git a/scripts/pan_skin/01_download_tier_a.sh b/scripts/pan_skin/01_download_tier_a.sh new file mode 100644 index 0000000000000000000000000000000000000000..5895b8eaf8853060953a39356f61a9da0550db92 --- /dev/null +++ b/scripts/pan_skin/01_download_tier_a.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# tier a — anchor skin datasets. writes to data/corpus/pan_skin/tier_a, never /tmp. +set -uo pipefail + +BASE=/home/bcheng/PRISM/data/corpus/pan_skin/tier_a +LOG=/home/bcheng/PRISM/logs/pan_skin +mkdir -p "$BASE" "$LOG" + +download() { + local acc=$1 url=$2 out=$3 + local outfile="$BASE/$out" + if [ -s "$outfile" ]; then + echo "[$acc] cached: $outfile ($(du -h "$outfile" | cut -f1))" + return 0 + fi + echo "[$acc] downloading $url" + wget --quiet -c -O "$outfile.part" "$url" \ + && mv "$outfile.part" "$outfile" \ + && echo "[$acc] done: $outfile ($(du -h "$outfile" | cut -f1))" \ + || { echo "[$acc] FAILED"; rm -f "$outfile.part"; return 1; } +} + +# GSE214695 — Aldrich Dev Cell 2023, En1-cKO volar snRNA-seq, ~45k cells (sibling of Dingwall). +download GSE214695_tar \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE214nnn/GSE214695/suppl/GSE214695_RAW.tar" \ + aldrich_GSE214695_RAW.tar +download GSE214695_annot \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE214nnn/GSE214695/suppl/GSE214695_cell_annotation.csv.gz" \ + aldrich_GSE214695_cell_annotation.csv.gz + +# GSE131498 - Ge/Gupta 2020 Theranostics, E13.5/E16.5/P0 dorsal skin, ~15k cells. +download GSE131498 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE131nnn/GSE131498/suppl/GSE131498_scRNA_seq_skin.txt.gz" \ + ge_gupta_GSE131498_expression.txt.gz + +# GSE67602 - Joost 2016 Cell Systems Smart-seq2 platform anchor, ~1.4k cells. +download GSE67602 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE67nnn/GSE67602/suppl/GSE67602_Joost_et_al_expression.txt.gz" \ + joost_GSE67602_expression.txt.gz + +# GSE142471 — Haensel 2020 Cell Reports, adult mouse skin homeostasis + wound, 5 samples. +if [ ! -s "$BASE/haensel_GSE142471_RAW.tar" ]; then + download GSE142471 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE142nnn/GSE142471/suppl/GSE142471_RAW.tar" \ + haensel_GSE142471_RAW.tar +else + echo "[GSE142471] already cached" +fi + +echo +ls -lh "$BASE" diff --git a/scripts/pan_skin/02_download_tier_b.sh b/scripts/pan_skin/02_download_tier_b.sh new file mode 100644 index 0000000000000000000000000000000000000000..306128ea26850bdb7838dc5e8d5200a28b0a7aa1 --- /dev/null +++ b/scripts/pan_skin/02_download_tier_b.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# tier b - depth/platform stress + wound recap + FACS holdout. +set -uo pipefail + +BASE=/home/bcheng/PRISM/data/corpus/pan_skin/tier_b +LOG=/home/bcheng/PRISM/logs/pan_skin +mkdir -p "$BASE" "$LOG" + +download() { + local acc=$1 url=$2 out=$3 + local outfile="$BASE/$out" + if [ -s "$outfile" ]; then + echo "[$acc] cached: $outfile ($(du -h "$outfile" | cut -f1))" + return 0 + fi + echo "[$acc] downloading $url" + wget --quiet -c -O "$outfile.part" "$url" \ + && mv "$outfile.part" "$outfile" \ + && echo "[$acc] done: $outfile ($(du -h "$outfile" | cut -f1))" \ + || { echo "[$acc] FAILED"; rm -f "$outfile.part"; return 1; } +} + +# GSE108097 - MCA Han 2018 Microwell-seq low-depth anchor. skin partitions only (large tar). +download GSE108097 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE108nnn/GSE108097/suppl/GSE108097_RAW.tar" \ + mca_GSE108097_RAW.tar + +# GSE141814 - WIHN wound-induced hair neogenesis +download GSE141814 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE141nnn/GSE141814/suppl/GSE141814_RAW.tar" \ + wihn_GSE141814_RAW.tar + +# GSE124901 - Ge/Fuchs Cell 2017 stem cell lineage. FACS-orthogonal holdout. +download GSE124901 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE124nnn/GSE124901/suppl/GSE124901_RAW.tar" \ + ge_fuchs_GSE124901_RAW.tar + +# GSE201447 - Merkel/touch dome (volar biology) +download GSE201447 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE201nnn/GSE201447/suppl/GSE201447_RAW.tar" \ + merkel_GSE201447_RAW.tar + +echo +ls -lh "$BASE" diff --git a/scripts/pan_skin/03_download_tier_c.sh b/scripts/pan_skin/03_download_tier_c.sh new file mode 100644 index 0000000000000000000000000000000000000000..228b30bea2e3ed630bf412c943da74ed5c40eed4 --- /dev/null +++ b/scripts/pan_skin/03_download_tier_c.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +# tier c — eccrine + placode/DC substate anchors. cross-species/stage to Dingwall (no leak). +set -euo pipefail +BASE="${1:-data/corpus/pan_skin/tier_c}" +mkdir -p "$BASE" + +download() { + local id="$1"; local url="$2"; local out="$3" + if [ -s "$BASE/$out" ]; then + echo "[$id] cached" + return 0 + fi + echo "[$id] downloading $url -> $out" + curl -fL --retry 3 -o "$BASE/$out" "$url" \ + || { echo "[$id] FAILED"; return 1; } +} + +# GSE70288 — Sennett 2015 Dev Cell, E14.5 bulk-RNA per-population (Pc=placode, DC=dermal condensate). +# used as marker refs only, not per-cell scRNA. +download GSE70288 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE70nnn/GSE70288/suppl/GSE70288_repFpkmMatrix_E14.5_embryonic_skin.txt.gz" \ + sennett_GSE70288_E14.5_bulk_fpkm.txt.gz + +# GSE221927 — Tie 2024 JID, aged-human eccrine sweat gland scRNA-seq. +download GSE221927 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE221nnn/GSE221927/suppl/GSE221927_RAW.tar" \ + tie_GSE221927_RAW.tar + +# GSE202352 — Wiedemann 2023 Cell Reports, palm/sole/hip human skin regional atlas. +download GSE202352 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE202nnn/GSE202352/suppl/GSE202352_RAW.tar" \ + wiedemann_GSE202352_palm_sole_hip_RAW.tar + +echo "[done] tier_c downloads complete" diff --git a/scripts/pan_skin/06_build_per_dataset_h5ads.py b/scripts/pan_skin/06_build_per_dataset_h5ads.py new file mode 100644 index 0000000000000000000000000000000000000000..764ced5fc7673f8e69f95986d16759179bdaa786 --- /dev/null +++ b/scripts/pan_skin/06_build_per_dataset_h5ads.py @@ -0,0 +1,64 @@ +"""build per-dataset h5ads (raw counts + basic qc) under harmonized/.""" +from __future__ import annotations +from pathlib import Path +import sys +import warnings +warnings.filterwarnings("ignore") + +import scanpy as sc + +import importlib.util +spec = importlib.util.spec_from_file_location( + "pan_skin_loaders", "/home/bcheng/PRISM/panda/data/pan_skin_loaders.py") +_mod = importlib.util.module_from_spec(spec); spec.loader.exec_module(_mod) +ALL_LOADERS = _mod.ALL_LOADERS + +OUT = Path("/home/bcheng/PRISM/data/corpus/pan_skin/harmonized") +OUT.mkdir(parents=True, exist_ok=True) + + +def qc(adata, name): + n0 = adata.n_obs + sc.pp.filter_cells(adata, min_genes=200) + sc.pp.filter_genes(adata, min_cells=3) + if adata.n_obs == 0: + print(f" [{name}] EMPTY after QC — skipping") + return None + adata.var["mt"] = adata.var_names.str.startswith(("mt-", "MT-")) + if adata.var["mt"].any(): + sc.pp.calculate_qc_metrics( + adata, qc_vars=["mt"], percent_top=None, log1p=False, inplace=True + ) + adata = adata[adata.obs["pct_counts_mt"] < 20].copy() + print(f" [{name}] n_cells: {n0} -> {adata.n_obs}, n_genes: {adata.n_vars}") + return adata + + +def main(): + for name, loader in ALL_LOADERS.items(): + out_path = OUT / f"{name}.h5ad" + if out_path.exists(): + print(f"[{name}] cached at {out_path}") + continue + print(f"[{name}] loading …") + try: + adata = loader() + except Exception as exc: + print(f"[{name}] LOADER FAILED: {exc!r}") + continue + print(f" raw: {adata.shape}") + adata = qc(adata, name) + if adata is None: + continue + # dataset id in obs + adata.obs["dataset"] = name + import scipy.sparse as sp + if not sp.issparse(adata.X): + adata.X = sp.csr_matrix(adata.X) + adata.X = adata.X.astype("float32") + adata.write_h5ad(out_path, compression="gzip") + print(f" saved to {out_path} ({out_path.stat().st_size/1e6:.1f} MB)") + + +if __name__ == "__main__": + main() diff --git a/scripts/pan_skin/07_build_shared_hvgs_and_pca.py b/scripts/pan_skin/07_build_shared_hvgs_and_pca.py new file mode 100644 index 0000000000000000000000000000000000000000..729d0d427546be81eed8a374b72beef4dd267b46 --- /dev/null +++ b/scripts/pan_skin/07_build_shared_hvgs_and_pca.py @@ -0,0 +1,168 @@ +"""shared-HVG picker + PCA fit. union of top-4000 HVGs across >=3 datasets, plus forced skin markers.""" +from __future__ import annotations +from pathlib import Path +import warnings, pickle +warnings.filterwarnings("ignore") + +import numpy as np +import pandas as pd +import scanpy as sc +import scipy.sparse as sp +from sklearn.decomposition import PCA +import anndata as ad + +HARM = Path("/home/bcheng/PRISM/data/corpus/pan_skin/harmonized") +K_TOP_PER_DS = 4000 +K_TARGET = 4000 +N_PCA = 50 +SAMPLE_PER_DS = 5000 + +MUST_INCLUDE = [ + # placode / eccrine progenitor + "Wnt10b", "Shh", "Foxi3", "Edar", "Sox9", "Bmp4", "Wnt10a", + "Foxa1", "En1", "Tfap2b", "Nkx3-1", "Foxn1", + # HF matrix / IRS / ORS + "Msx2", "Krt31", "Krt71", "Krt17", "Krt5", "Krt14", + "Cutl1", "Lgr5", "Cd34", "Nfatc1", "Runx3", + # basal / epidermal + "Trp63", "Ovol1", "Ivl", "Flg", "Lor", + # DP / dermal + "Sox2", "Corin", "Bmp6", "Wif1", "Igfbp3", "Prrx1", "Tbx15", + # melanocyte / merkel / neural crest + "Mitf", "Dct", "Tyr", "Pmel", "Sox10", "Atoh1", "Piezo2", + # signaling regulators + "Wnt3", "Wnt7b", "Bmpr1a", "Bmpr1b", "Fgf9", "Fgf10", +] + + +def rank_hvgs(a): + x = a.copy() + try: + sc.pp.highly_variable_genes(x, n_top_genes=K_TOP_PER_DS, flavor="seurat", + subset=False, check_values=False) + return x.var["variances_norm"].fillna(-np.inf) if "variances_norm" in x.var else \ + x.var["dispersions_norm"].fillna(-np.inf) + except Exception: + sc.pp.normalize_total(x, target_sum=1e4); sc.pp.log1p(x) + X = x.X.toarray() if sp.issparse(x.X) else x.X + return pd.Series(np.asarray(X.var(axis=0)).ravel(), index=x.var_names) + + +def load_and_norm(f, shared_genes, gene_idx, mu, sig, G): + a = ad.read_h5ad(f) + raw_counts = np.asarray(a.X.sum(axis=1)).ravel() if sp.issparse(a.X) else a.X.sum(axis=1) + common = [g for g in shared_genes if g in a.var_names] + a_s = a[:, common].copy() + sc.pp.normalize_total(a_s, target_sum=1e4) + sc.pp.log1p(a_s) + X = a_s.X.toarray() if sp.issparse(a_s.X) else a_s.X + Xf = np.zeros((X.shape[0], G), dtype=np.float32) + cols = [gene_idx[g] for g in common] + Xf[:, cols] = X.astype(np.float32) + Xz = np.clip((Xf - mu.astype(np.float32)) / sig.astype(np.float32), -10, 10) + return a, Xf, Xz, raw_counts, common + + +def main(): + files = sorted(HARM.glob("*.h5ad")) + files = [f for f in files if not f.name.startswith("corpus")] + + # rank per dataset + ranks = {} + for f in files: + print(f"[hvg] rank {f.name}", flush=True) + a = ad.read_h5ad(f) + ranks[f.stem] = rank_hvgs(a) + del a + + R = pd.DataFrame(ranks) # rows = union of genes, cols = datasets + top_hits = pd.DataFrame({d: R[d].rank(ascending=False, method="min") <= K_TOP_PER_DS + for d in R.columns}).fillna(False) + hit_count = top_hits.sum(axis=1) + + eligible = hit_count[hit_count >= 3] + print(f"[hvg] {len(eligible)} genes hit top-{K_TOP_PER_DS} in >=3 datasets") + + mean_score = R.mean(axis=1, skipna=True) + shared_ranked = eligible.sort_values(ascending=False).index + shared_by_score = mean_score.loc[shared_ranked].sort_values(ascending=False).index + + all_genes = set(R.index) + forced = [g for g in MUST_INCLUDE if g in all_genes] + print(f"[hvg] {len(forced)}/{len(MUST_INCLUDE)} must-include markers present in union") + + picked = list(forced) + for g in shared_by_score: + if g not in picked: + picked.append(g) + if len(picked) >= K_TARGET: + break + shared_genes = picked[:K_TARGET] + print(f"[hvg] final shared HVG count: {len(shared_genes)}") + forced_kept = [g for g in forced if g in shared_genes] + print(f"[hvg] forced markers retained: {len(forced_kept)}/{len(forced)}") + + (HARM / "shared_hvgs.txt").write_text("\n".join(shared_genes) + "\n") + + G = len(shared_genes) + gene_idx = {g: i for i, g in enumerate(shared_genes)} + running_sum = np.zeros(G, dtype=np.float64) + running_sq = np.zeros(G, dtype=np.float64) + N = 0 + for f in files: + a = ad.read_h5ad(f) + common = [g for g in shared_genes if g in a.var_names] + a_s = a[:, common].copy() + sc.pp.normalize_total(a_s, target_sum=1e4) + sc.pp.log1p(a_s) + X = a_s.X.toarray() if sp.issparse(a_s.X) else a_s.X + cols = [gene_idx[g] for g in common] + running_sum[cols] += X.sum(axis=0) + running_sq[cols] += (X ** 2).sum(axis=0) + N += X.shape[0] + del a, a_s, X + mu = running_sum / N + sig = np.sqrt(np.maximum(running_sq / N - mu ** 2, 1e-6)) + np.savez(HARM / "corpus_stats.npz", shared_hvgs=np.array(shared_genes), + mean=mu.astype(np.float32), std=sig.astype(np.float32), n_cells=N) + print(f"[stats] mean/std over {N} cells") + + fit_chunks = [] + rng = np.random.default_rng(0) + for f in files: + _, _, Xz, _, _ = load_and_norm(f, shared_genes, gene_idx, mu, sig, G) + n = Xz.shape[0] + take = min(SAMPLE_PER_DS, n) + idx = rng.choice(n, size=take, replace=False) + fit_chunks.append(Xz[idx]) + Xfit = np.vstack(fit_chunks) + pca = PCA(n_components=N_PCA, random_state=42) + pca.fit(Xfit) + print(f"[pca] explained var={pca.explained_variance_ratio_.sum():.3f}") + # sign-align top loading positive for determinism + for k in range(N_PCA): + top = int(np.argmax(np.abs(pca.components_[k]))) + if pca.components_[k, top] < 0: + pca.components_[k] *= -1 + with open(HARM / "pca_basis.pkl", "wb") as fh: + pickle.dump(pca, fh) + + parts = [] + for f in files: + a, Xf, Xz, raw_counts, common = load_and_norm(f, shared_genes, gene_idx, mu, sig, G) + Z = pca.transform(Xz).astype(np.float32) + obs = a.obs.copy() + obs["total_counts"] = raw_counts.astype(np.float32) + obs["missing_hvg_frac"] = 1.0 - len(common) / G + a2 = ad.AnnData(X=sp.csr_matrix(Xf.astype(np.float32)), obs=obs, + var=pd.DataFrame(index=shared_genes)) + a2.obsm["X_pca"] = Z + parts.append(a2) + corpus = ad.concat(parts, join="outer", label="_batch") + corpus.uns["shared_hvgs"] = shared_genes + corpus.write_h5ad(HARM / "corpus.h5ad", compression="gzip") + print(f"[emit] wrote corpus.h5ad ({corpus.n_obs} x {corpus.n_vars})") + + +if __name__ == "__main__": + main() diff --git a/scripts/pan_skin/08_assign_labels.py b/scripts/pan_skin/08_assign_labels.py new file mode 100644 index 0000000000000000000000000000000000000000..5171629eccdd1f3a6d3eebfb104c2971b67e5738 --- /dev/null +++ b/scripts/pan_skin/08_assign_labels.py @@ -0,0 +1,69 @@ +"""marker-scored 18-class labels per corpus cell. writes canonical_label + leiden to corpus.h5ad.""" +from __future__ import annotations +from pathlib import Path +import warnings, yaml, sys +warnings.filterwarnings("ignore") + +import numpy as np +import pandas as pd +import scanpy as sc +import anndata as ad + +CORPUS = Path("/home/bcheng/PRISM/data/corpus/pan_skin/harmonized/corpus.h5ad") +YAML = Path("/home/bcheng/PRISM/scripts/pan_skin/known_skin_tfs.yaml") + + +def main(): + with open(YAML) as f: + cfg = yaml.safe_load(f) + classes = list(cfg["classes"].keys()) + markers = cfg["classes"] + min_score = cfg["assignment"]["min_score"] + min_margin = cfg["assignment"]["min_margin"] + resolution = cfg["assignment"]["cluster_resolution"] + + a = ad.read_h5ad(CORPUS) + print(f"[label] loaded corpus: {a.shape}") + + a.obs["canonical_label"] = pd.Categorical(["UNK"] * a.n_obs, categories=classes + ["UNK"]) + a.obs["leiden"] = "0" + + for ds_name, sub in a.obs.groupby("dataset"): + print(f"[label] {ds_name}: n={sub.shape[0]}") + idx = sub.index + sa = a[idx].copy() + + sc.pp.neighbors(sa, use_rep="X_pca", n_neighbors=15) + sc.tl.leiden(sa, resolution=resolution, key_added="leiden") + a.obs.loc[idx, "leiden"] = ds_name + "_" + sa.obs["leiden"].astype(str) + + for c, m in markers.items(): + present = [g for g in m if g in sa.var_names] + if not present: + sa.obs[f"score_{c}"] = -np.inf + continue + sc.tl.score_genes(sa, gene_list=present, score_name=f"score_{c}", + random_state=0, use_raw=False) + S = sa.obs[[f"score_{c}" for c in classes]].values + cluster_labels = sa.obs["leiden"].values + assigned = np.array(["UNK"] * sa.n_obs, dtype=object) + for cl in np.unique(cluster_labels): + mask = cluster_labels == cl + mean_scores = S[mask].mean(axis=0) + order = np.argsort(mean_scores)[::-1] + top, second = mean_scores[order[0]], mean_scores[order[1]] + if top >= min_score and (top - second) >= min_margin: + assigned[mask] = classes[order[0]] + a.obs.loc[idx, "canonical_label"] = assigned + print(f" {ds_name} label breakdown:", + pd.Series(assigned).value_counts().to_dict()) + + print("[label] corpus-wide breakdown:") + print(a.obs.groupby(["dataset", "canonical_label"], observed=True).size().unstack(fill_value=0)) + + a.write_h5ad(CORPUS, compression="gzip") + print(f"[label] wrote {CORPUS}") + + +if __name__ == "__main__": + main() diff --git a/scripts/pan_skin/08b_curated_label_override.py b/scripts/pan_skin/08b_curated_label_override.py new file mode 100644 index 0000000000000000000000000000000000000000..7470bdf412c2e70c6d4abbea2a021a17b6a552c1 --- /dev/null +++ b/scripts/pan_skin/08b_curated_label_override.py @@ -0,0 +1,70 @@ +"""override marker labels with paper-curated ones for sulic + merkel.""" +from __future__ import annotations +from pathlib import Path +import warnings +warnings.filterwarnings("ignore") + +import numpy as np +import pandas as pd +import anndata as ad + +CORPUS = Path("/home/bcheng/PRISM/data/corpus/pan_skin/harmonized/corpus.h5ad") + + +SULIC_MAP = { + "Epithelium": "basal-IFE", + "Placode1": "HF-placode", "Placode2": "HF-placode", + "PlacodeI": "HF-placode", "PlacodeII": "HF-placode", + "PlacodeIII": "HF-placode", "PlacodeIV": "HF-placode", +} + +MERKEL_MAP = { + # touch dome paper labels (per GSE201447_Combined_clusters_cell_types.txt.gz) + "Interfollicular epidermis BII (IFE BII)": "basal-IFE", + "Interfollicular epidermis BI (IFE BI)": "basal-IFE", + "Interfollicular epidermis SP": "spinous", + "Interfollicular epidermis G": "granular", + "Merkel cell": "Merkel", + "Merkel cells": "Merkel", + "K17+ Basal": "HF-ORS", + "Touch dome": "Merkel", + "Neuron": "Merkel", # innervating Merkel-associated +} + + +def main(): + a = ad.read_h5ad(CORPUS) + print(f"[override] corpus: {a.shape}", flush=True) + old = a.obs["canonical_label"].astype(str).copy() + + # sulic loader propagates `sample` = "Epithelium"/"Placode1"/"Placode2" + sulic_mask = a.obs["dataset"] == "sulic_GSE212673" + print(f"[override] sulic sample values:", + a.obs.loc[sulic_mask, "sample"].value_counts().to_dict()) + mapped = a.obs.loc[sulic_mask, "sample"].map(SULIC_MAP) + ok = mapped.notna() + idx = a.obs_names[sulic_mask][ok.values] + a.obs.loc[idx, "canonical_label"] = mapped[ok].values + print(f"[override] sulic: {int(ok.sum())} cells relabeled via sample column") + print(" breakdown:", mapped[ok].value_counts().to_dict()) + + # merkel override + if "cell_type" in a.obs.columns: + mask = (a.obs["dataset"] == "merkel_GSE201447") & a.obs["cell_type"].isin(MERKEL_MAP) + new = a.obs.loc[mask, "cell_type"].map(MERKEL_MAP) + n = int(mask.sum()) + print(f"[override] merkel: {n} cells relabeled via cell_type") + print(" breakdown:", new.value_counts().to_dict()) + a.obs.loc[mask, "canonical_label"] = new.astype(a.obs["canonical_label"].dtype) + + changed = (old != a.obs["canonical_label"].astype(str)).sum() + print(f"[override] total changed: {changed}") + print("[override] new corpus-wide breakdown:") + print(a.obs.groupby(["dataset","canonical_label"], observed=True).size().unstack(fill_value=0)) + + a.write_h5ad(CORPUS, compression="gzip") + print(f"[override] wrote {CORPUS}") + + +if __name__ == "__main__": + main() diff --git a/scripts/pan_skin/10_build_corpus.py b/scripts/pan_skin/10_build_corpus.py new file mode 100644 index 0000000000000000000000000000000000000000..7d3a3cdb3ba2cda96e27779a0bfa5a77416e383e --- /dev/null +++ b/scripts/pan_skin/10_build_corpus.py @@ -0,0 +1,538 @@ +"""pan-skin corpus v3 — 100% paper-labeled. drops datasets without external annotations.""" +from __future__ import annotations +from pathlib import Path +import warnings, importlib.util, pickle, tempfile, tarfile, re, sys +warnings.filterwarnings("ignore") + +import numpy as np +import pandas as pd +import scanpy as sc +import anndata as ad +import scipy.sparse as sp +import yaml +from sklearn.decomposition import PCA + +ROOT = Path("/home/bcheng/PRISM") +CORPUS_ROOT = ROOT / "data/corpus/pan_skin" +TIER_A, TIER_B, TIER_C = (CORPUS_ROOT / "tier_a"), (CORPUS_ROOT / "tier_b"), (CORPUS_ROOT / "tier_c") +HARM = CORPUS_ROOT / "harmonized" +HARM.mkdir(parents=True, exist_ok=True) +EXT = ROOT / "data/external_labels" + +# import loaders (uses fixed Merkel join) +spec = importlib.util.spec_from_file_location( + "pan_skin_loaders", str(ROOT / "panda/data/pan_skin_loaders.py")) +LDR = importlib.util.module_from_spec(spec); spec.loader.exec_module(LDR) + + +# 18-class canonical vocabulary +CANONICAL_CLASSES = [ + "basal-IFE", "spinous", "granular", + "HF-placode", "HF-primary-germ", "HF-matrix", "HF-DP", "HF-ORS", "HF-IRS", + "sebaceous", "eccrine-placode", "eccrine-duct", + "melanocyte", "Merkel", "endothelial", + "fibroblast-papillary", "fibroblast-reticular", "immune", +] + +# legacy-label rename map (applied post-canonicalization to any label source). +# note: eccrine-placode kept as canonical class — appears only in Aldrich holdout (~5 cells in v2). +LEGACY_RENAME = { + "hair-placode": "HF-placode", + "eccrine-ductal": "eccrine-duct", + "basal-multipotent": "basal-IFE", + "eden-dermal-niche": "fibroblast-papillary", + "fibroblast": "fibroblast-reticular", + "eccrine-secretory": "eccrine-duct", + "dermal-condensate": "HF-DP", + "nascent-eccrine-gland": "eccrine-duct", +} + +# per-dataset canonicalization maps +SULIC_MAP = { + "Epithelium": "basal-IFE", + "Placode1": "HF-placode", "Placode2": "HF-placode", + "PlacodeI": "HF-placode", "PlacodeII": "HF-placode", + "PlacodeIII": "HF-placode", "PlacodeIV": "HF-placode", +} + +# merkel touch-dome (22 cluster names) +MERKEL_MAP = { + "Interfollicular epidermis Basal (IFE B)": "basal-IFE", + "Interfollicular epidermis BI (IFE BI)": "basal-IFE", + "Interfollicular epidermis BII (IFE BII)": "basal-IFE", + "Interfollicular epidermis DI (IFE DI)": "spinous", + "Interfollicular epidermis DI/KII (IFE DI/KII)": "spinous", + "Interfollicular epidermis KII (IFE KII)": "granular", + "Sebaceous gland (SG)": "sebaceous", + "Outer/Inner Bulge (OB/IB)": "HF-ORS", + "Upper Hair Follicle II/V (uHF II/V)": "HF-ORS", + "Upper Hair Follicle/Infundibular (uHF I/INFU B)": "HF-ORS", + "Fibroblast III/I (FIB III/I)": "fibroblast-reticular", + "Fibroblasts (FIB)": "fibroblast-reticular", + "Blood vascular Endothelial Cells (BEC)": "endothelial", + "Lymphatic Endothelial Cells (LEC)": "endothelial", + "Dendritic Cells (DC)": "immune", + "Langerhans Cells (LC)": "immune", + "abT/Innate Lymphoid Cells (abT/ILC)": "immune", + "gdT Cells (gdTC)": "immune", + "Innate (NK/ILCs) Cells": "immune", + "Granulocytes": "immune", + "Double Negative SLN (DN)": "immune", + "Double Negative Sromal Cells (EC)": "fibroblast-reticular", +} + +# MCA neonatal skin — annotation col from MCA_CellAssignments.csv, tissue==NeonatalSkin_1 +MCA_MAP = { + "Keratinocyte(Neonatal-Skin)": "basal-IFE", + "Epithelial cell(Neonatal-Skin)": "basal-IFE", + "Osteoblast_Ppic high(Neonatal-Skin)": "fibroblast-reticular", + "Osteoblast_Dlk1 high(Neonatal-Skin)": "fibroblast-reticular", + "Stromal cell_Gas6 high(Neonatal-Skin)": "fibroblast-reticular", + "Stromal cell_Akr1c18 high(Neonatal-Skin)": "fibroblast-reticular", + "Adipocyte(Neonatal-Skin)": "fibroblast-reticular", + "Brown adipose tissue_Cidea high(Neonatal-Skin)": "fibroblast-reticular", + "Brown adipose tissue_Cox8b high(Neonatal-Skin)": "fibroblast-reticular", + "Brown adipose tissue_mt-Nd5 high(Neonatal-Skin)": "fibroblast-reticular", + "Muscle cell_Lrrc15 high(Neonatal-Skin)": "fibroblast-reticular", + "Muscle cell_Actc1 high(Neonatal-Skin)": "fibroblast-reticular", + "Smooth muscle cell_Acta2 high(Neonatal-Skin)": "fibroblast-reticular", + "Endothelial cell(Neonatal-Skin)": "endothelial", + "Lymphatic vessel endothelial cell(Neonatal-Skin)": "endothelial", + "Melanocyte(Neonatal-Skin)": "melanocyte", + "Macrophage_Pf4 high(Neonatal-Skin)": "immune", + "Macrophage_Cd74 high(Neonatal-Skin)": "immune", + "Macrophage_Lyz2 high(Neonatal-Skin)": "immune", + "Mast cell(Neonatal-Skin)": "immune", + "Neutrophil(Neonatal-Skin)": "immune", + # Erythroblast / Neuron / Dividing cell -> unmapped (abstain) +} + +# Joost 2016 back skin epithelium +JOOST_MAP = { + "IFE basal cells (IFE B)": "basal-IFE", + "IFE differentiated cells I (IFE DI)": "spinous", + "IFE differentiated cells II (IFE DII)":"spinous", + "IFE keratinized layer I (IFE KI)": "granular", + "IFE keratinized layer II (IFE KII)": "granular", + "Inner bulge (IB)": "HF-ORS", + "Outer bulge (OB)": "HF-ORS", + "Upper hair follicle I (uHF I)": "HF-ORS", + "Upper hair follicle II (uHF II)": "HF-ORS", + "Upper hair follicle III (uHF III)": "HF-ORS", + "Sebaceous gland (SG)": "sebaceous", + "Langerhans cells (LH)": "immune", + "T cells (TC)": "immune", +} + +# haensel (joost/annusver GSE142471) — meta_Total_Cells_integration_UW_WO_CellRep2020_FigS2E +HAENSEL_MAP = { + "Basal I": "basal-IFE", + "Basal II": "basal-IFE", + "Prolif. Basal": "basal-IFE", + "Spinous": "spinous", + "HF I": "HF-ORS", + "HF II": "HF-ORS", + "HFSC": "HF-ORS", + "Fibroblast I": "fibroblast-papillary", # superficial papillary + "Fibroblast II": "fibroblast-reticular", + "Fibroblast III": "fibroblast-reticular", + "Fibroblast IV": "fibroblast-reticular", + "Fibroblast V": "fibroblast-reticular", + "Fibroblast VI": "fibroblast-reticular", + "Fibroblast VII": "fibroblast-reticular", + "Myofibroblast": "fibroblast-reticular", + "Endothelial": "endothelial", + "Dendritic cell": "immune", + "Langerhans": "immune", + "Macrophage I": "immune", + "Macrophage II": "immune", + "Macrophage III": "immune", + "T cell I": "immune", + "T cell II": "immune", + # Skeletal muscle -> unmapped (abstain) +} + +# run → sample mapping for Haensel obs_name reconstruction +HAENSEL_RUN_TO_SAMPLE = { + "bs_1": "Un-Wounded_1_scRNA-Seq", + "bs_2": "Un-Wounded_2_scRNA-Seq", + "sw_1": "Wounded_1_scRNA-Seq", + "sw_2": "Wounded_2_scRNA-Seq", + "sw_3": "Wounded_3_scRNA-Seq", +} + +def qc(a, name): + n0 = a.n_obs + sc.pp.filter_cells(a, min_genes=200) + sc.pp.filter_genes(a, min_cells=3) + a.var["mt"] = a.var_names.str.startswith(("mt-", "MT-")) + if a.var["mt"].any(): + sc.pp.calculate_qc_metrics(a, qc_vars=["mt"], percent_top=None, + log1p=False, inplace=True) + a = a[a.obs["pct_counts_mt"] < 20].copy() + print(f" [{name}] QC {n0} -> {a.n_obs}", flush=True) + return a + + +# custom loaders for tier_c (not in ALL_LOADERS) +def load_tie(): + parts = [] + D = TIER_C / "tie_GSE221927_RAW_extracted" + for h5 in sorted(D.glob("*_filtered_feature_bc_matrix.h5")): + sample = h5.name.split("_")[0] + a = sc.read_10x_h5(str(h5)); a.var_names_make_unique() + a.obs["source_sample"] = sample + a.obs.index = [f"{sample}_{bc}" for bc in a.obs.index.astype(str)] + parts.append(a) + a = ad.concat(parts, join="outer") + a.uns["dataset"] = "tie_GSE221927"; a.uns["organism"] = "human" + return a + + +def load_wiedemann(): + D = TIER_C / "wiedemann_GSE202352_palm_sole_hip_RAW_extracted" + parts = [] + for tarpath in sorted(list(D.glob("GSM*_filtered_feature_bc_matrix.tar.gz")) + + list(D.glob("GSM*.filtered_feature_bc_matrix.tar.gz"))): + name = tarpath.name.lower() + region = ("palm" if "palm" in name else "sole" if "sole" in name + else "hip" if "hip" in name else "unknown") + sample = tarpath.name.split("_")[0] + "_" + region + with tempfile.TemporaryDirectory() as tmp: + with tarfile.open(tarpath) as tar: + tar.extractall(tmp) + # skip macOS ._resource-fork files + mtx = [p for p in Path(tmp).rglob("matrix.mtx*") + if not p.name.startswith("._")] + if not mtx: continue + a = sc.read_10x_mtx(str(mtx[0].parent)) + a.var_names_make_unique() + a.obs["source_sample"] = sample + a.obs["region"] = region + a.obs.index = [f"{sample}_{bc}" for bc in a.obs.index.astype(str)] + parts.append(a) + a = ad.concat(parts, join="outer") + a.uns["dataset"] = "wiedemann_GSE202352"; a.uns["organism"] = "human" + return a + + +# paper-label attachment (writes a.obs['paper_label']) +def attach_sulic(a): + a.obs["paper_label"] = a.obs["sample"].astype(str) + return a + +def attach_merkel(a): + a.obs["paper_label"] = a.obs["cell_type"].astype(str) + return a + +def attach_mca(a): + csv = pd.read_csv(EXT / "mca/MCA_CellAssignments.csv") + csv = csv[csv["Tissue"] == "Neonatal-Skin"] + lbl = csv.set_index("Cell.name")["Annotation"] + a.obs["paper_label"] = a.obs_names.to_series().map(lbl).fillna("").astype(str) + return a + +def attach_joost2016(a): + csv = pd.read_csv(EXT / "joost2016/Joost_EPI_clusterID.csv") + csv["CellID"] = csv["CellID"].astype(str).str.replace("-", "_", regex=False) + lbl = csv.set_index("CellID")["Cluster"] + a.obs["paper_label"] = a.obs_names.to_series().map(lbl).fillna("").astype(str) + return a + +def attach_haensel(a): + meta = pd.read_csv( + EXT / "haensel/meta_Total_Cells_integration_UW_WO_CellRep2020_FigS2E.txt", + sep="\t") + # Row col has per-Run -N suffixes from Seurat merge (bs_1->-1, sw_1->-2, ...); + # cellranger barcodes.tsv is always -1, so strip -N and re-add -1 for the join. + meta["sample"] = meta["Run"].map(HAENSEL_RUN_TO_SAMPLE) + meta = meta.dropna(subset=["sample"]) + bc_stripped = meta["Row"].astype(str).str.replace(r"-\d+$", "", regex=True) + meta["obsname"] = meta["sample"].astype(str) + "_" + bc_stripped + "-1" + lbl = meta.set_index("obsname")["final_labels"] + a.obs["paper_label"] = a.obs_names.to_series().map(lbl).fillna("").astype(str) + return a + + +# marker scoring for un-labeled cells +def marker_score(a, markers, classes, min_score=0.10, min_margin=0.05, + resolution=0.8, seed=0): + """cluster + score + vote → per-cell canonical label or ''.""" + a_local = a.copy() + # auto-uppercase markers for human datasets (uppercase-dominant var_names) + upper_frac = np.mean([g.isupper() for g in a_local.var_names[:200]]) + if upper_frac > 0.5: + markers = {c: [g.upper() for g in m] for c, m in markers.items()} + sc.pp.normalize_total(a_local, target_sum=1e4); sc.pp.log1p(a_local) + sc.pp.highly_variable_genes(a_local, n_top_genes=3000, flavor="seurat", + subset=False) + sc.pp.scale(a_local, max_value=10) + n_pc = min(50, a_local.n_obs - 1, a_local.n_vars - 1) + sc.tl.pca(a_local, n_comps=n_pc, random_state=seed) + try: + sc.pp.neighbors(a_local, use_rep="X_pca", n_neighbors=15, + random_state=seed) + sc.tl.leiden(a_local, resolution=resolution, key_added="leiden", + random_state=seed) + except Exception: + a_local.obs["leiden"] = "0" + for c, m in markers.items(): + present = [g for g in m if g in a_local.var_names] + if not present: + a_local.obs[f"score_{c}"] = -np.inf + continue + sc.tl.score_genes(a_local, gene_list=present, + score_name=f"score_{c}", random_state=seed, + use_raw=False) + S = a_local.obs[[f"score_{c}" for c in classes]].values + cl = a_local.obs["leiden"].astype(str).values + out = np.array([""] * a_local.n_obs, dtype=object) + for c_id in np.unique(cl): + mask = cl == c_id + m = S[mask].mean(axis=0) + order = np.argsort(m)[::-1] + top, second = m[order[0]], m[order[1]] + if top >= min_score and (top - second) >= min_margin: + out[mask] = classes[order[0]] + return out + + +# HVG picker for shared-genes step +def rank_hvgs(a, K=4000): + x = a.copy() + try: + sc.pp.highly_variable_genes(x, n_top_genes=K, flavor="seurat_v3", + subset=False, check_values=False) + return (x.var["variances_norm"].fillna(-np.inf) + if "variances_norm" in x.var + else x.var["dispersions_norm"].fillna(-np.inf)) + except Exception: + sc.pp.normalize_total(x, target_sum=1e4); sc.pp.log1p(x) + X = x.X.toarray() if sp.issparse(x.X) else x.X + return pd.Series(np.asarray(X.var(axis=0)).ravel(), index=x.var_names) + + +MUST_INCLUDE = [ + "Wnt10b","Shh","Foxi3","Edar","Sox9","Bmp4","Wnt10a", + "Foxa1","En1","Tfap2b","Nkx3-1","Foxn1", + "Msx2","Krt31","Krt71","Krt17","Krt5","Krt14","Cutl1","Lgr5","Cd34","Nfatc1","Runx3", + "Trp63","Ovol1","Ivl","Flg","Lor", + "Sox2","Corin","Bmp6","Wif1","Igfbp3","Prrx1","Tbx15", + "Mitf","Dct","Tyr","Pmel","Sox10","Atoh1","Piezo2","Krt20","Krt18","Krt19", + "Aqp5","Chga","Chgb","Pecam1","Cdh5","Kdr","Vwf","Cldn5", + "Col1a1","Col1a2","Dcn","Lum","Pdgfra","Dpp4","Lef1","Postn","Fbn1","Fap", + "Ptprc","Cd3e","Cd8a","Cd4","Adgre1","Itgam", + "Wnt3","Wnt7b","Bmpr1a","Bmpr1b","Fgf9","Fgf10", + "Scd1","Scd3","Elovl3","Fasn", +] + + +def main(): + cfg = yaml.safe_load(open("/home/bcheng/PRISM/scripts/pan_skin/known_skin_tfs.yaml")) + marker_classes = list(cfg["classes"].keys()) + marker_genes = cfg["classes"] + marker_cfg = cfg["assignment"] + + # paper-labeled datasets only + LOADERS = [ + ("sulic_GSE212673", LDR.load_sulic_GSE212673, attach_sulic, True), + ("merkel_GSE201447", LDR.load_merkel_GSE201447, attach_merkel, True), + ("mca_GSE108097_neonatal", LDR.load_mca_neonatal_skin, attach_mca, True), + ("joost_GSE67602", LDR.load_joost_GSE67602, attach_joost2016, True), + ("joost_annusver_GSE142471", LDR.load_joost_annusver_GSE142471, attach_haensel, True), + ] + + per_dataset_maps = { + "sulic_GSE212673": SULIC_MAP, + "merkel_GSE201447": MERKEL_MAP, + "mca_GSE108097_neonatal": MCA_MAP, + "joost_GSE67602": JOOST_MAP, + "joost_annusver_GSE142471": HAENSEL_MAP, + } + + ADATAS = {} # name -> raw+labeled AnnData (raw counts in .X) + stats_rows = [] + + for name, loader, attacher, has_paper in LOADERS: + print(f"\n### {name} ###", flush=True) + a = loader() + print(f" raw {a.shape}", flush=True) + a = qc(a, name) + if a is None or a.n_obs == 0: continue + a.obs["dataset"] = name + if not sp.issparse(a.X): + a.X = sp.csr_matrix(a.X) + a.X = a.X.astype("float32") + + a.obs["canonical_label"] = "" + a.obs["label_source"] = "abstain" + + n_paper = n_marker = 0 + if has_paper and attacher is not None: + a = attacher(a) + raw = a.obs["paper_label"].astype(str) + m = per_dataset_maps[name] + canonical = raw.map(m).fillna("").astype(str) + canonical = canonical.replace(LEGACY_RENAME) + canonical = canonical.where(canonical.isin(CANONICAL_CLASSES), "") + got = canonical != "" + a.obs.loc[got.values, "canonical_label"] = canonical[got].values + a.obs.loc[got.values, "label_source"] = "paper" + n_paper = int(got.sum()) + + # v3: paper-labeled only. drop abstains/unmapped/no-label. + keep = a.obs["label_source"] == "paper" + n_dropped = int((~keep).sum()) + a = a[keep].copy() + n_abstain = 0 # by construction — corpus is 100% paper-labeled + unique_labels = sorted(a.obs["canonical_label"].unique()) + stats_rows.append(dict(dataset=name, n_cells=a.n_obs, + n_paper=n_paper, n_marker=n_marker, + n_abstain=n_abstain, n_dropped=n_dropped, + n_labels=len(unique_labels))) + print(f" paper={n_paper} marker={n_marker} dropped={n_dropped} " + f"labels={unique_labels}", flush=True) + ADATAS[name] = a + + # drop rare classes (<100 cells corpus-wide) + print("\n### Corpus-wide class filter (<100 cells) ###", flush=True) + all_labels = pd.concat([a.obs["canonical_label"] for a in ADATAS.values()]) + corpus_counts = all_labels.value_counts() + rare = set(corpus_counts[corpus_counts < 100].index) + if rare: + print(f" Dropping rare classes (<100 cells corpus-wide): " + f"{ {c: int(corpus_counts[c]) for c in rare} }", flush=True) + new_ADATAS = {} + for name, a in ADATAS.items(): + keep = ~a.obs["canonical_label"].isin(rare) + n_dropped_rare = int((~keep).sum()) + if n_dropped_rare: + print(f" {name}: dropped {n_dropped_rare} rare-class cells", + flush=True) + a = a[keep].copy() + if a.n_obs: + new_ADATAS[name] = a + ADATAS = new_ADATAS + + print("\n\n### Per-dataset coverage ###") + stats_df = pd.DataFrame(stats_rows) + print(stats_df.to_string(index=False)) + + print("\n### Shared HVGs ###", flush=True) + ranks = {} + for name, a in ADATAS.items(): + print(f" hvg rank {name}", flush=True) + ranks[name] = rank_hvgs(a, K=4000) + R = pd.DataFrame(ranks) + K_TOP = 4000; K_TGT = 4000 + top_hits = pd.DataFrame({d: R[d].rank(ascending=False, method="min") <= K_TOP + for d in R.columns}).fillna(False) + hit_count = top_hits.sum(axis=1) + eligible = hit_count[hit_count >= 3].index + mean_score = R.mean(axis=1, skipna=True) + ranked = mean_score.loc[eligible].sort_values(ascending=False).index + + all_genes = set(R.index) + forced = [g for g in MUST_INCLUDE if g in all_genes] + print(f" must-include kept: {len(forced)}/{len(MUST_INCLUDE)}", flush=True) + picked = list(forced) + for g in ranked: + if g not in picked: picked.append(g) + if len(picked) >= K_TGT: break + shared_genes = picked[:K_TGT] + print(f" final shared HVG count: {len(shared_genes)}", flush=True) + G = len(shared_genes) + gene_idx = {g: i for i, g in enumerate(shared_genes)} + + print("\n### Corpus normalization stats ###", flush=True) + running_sum = np.zeros(G, dtype=np.float64) + running_sq = np.zeros(G, dtype=np.float64) + N = 0 + per_ds_Xf = {} + for name, a in ADATAS.items(): + common = [g for g in shared_genes if g in a.var_names] + a_s = a[:, common].copy() + sc.pp.normalize_total(a_s, target_sum=1e4); sc.pp.log1p(a_s) + X = a_s.X.toarray() if sp.issparse(a_s.X) else a_s.X + Xf = np.zeros((X.shape[0], G), dtype=np.float32) + cols = [gene_idx[g] for g in common] + Xf[:, cols] = X.astype(np.float32) + per_ds_Xf[name] = (Xf, common) + running_sum[cols] += X.sum(axis=0) + running_sq[cols] += (X ** 2).sum(axis=0) + N += X.shape[0] + del a_s, X + mu = running_sum / N + sig = np.sqrt(np.maximum(running_sq / N - mu ** 2, 1e-6)) + print(f" mean/std over {N} cells", flush=True) + + print("\n### PCA fit (30k subsample) ###", flush=True) + rng = np.random.default_rng(0) + per_ds_Xz = {} + fit_chunks = [] + target_total = 30_000 + per_ds_take = max(1, target_total // len(ADATAS)) + for name, (Xf, common) in per_ds_Xf.items(): + Xz = np.clip((Xf - mu.astype(np.float32)) / sig.astype(np.float32), + -10, 10) + per_ds_Xz[name] = Xz + n = Xz.shape[0] + take = min(per_ds_take, n) + idx = rng.choice(n, size=take, replace=False) + fit_chunks.append(Xz[idx]) + Xfit = np.vstack(fit_chunks) + print(f" Xfit shape: {Xfit.shape}", flush=True) + pca = PCA(n_components=50, random_state=42) + pca.fit(Xfit) + print(f" explained var: {pca.explained_variance_ratio_.sum():.3f}", flush=True) + for k in range(pca.n_components_): + top = int(np.argmax(np.abs(pca.components_[k]))) + if pca.components_[k, top] < 0: + pca.components_[k] *= -1 + + print("\n### Assembling corpus ###", flush=True) + parts = [] + for name, a in ADATAS.items(): + Xf, common = per_ds_Xf[name] + Xz = per_ds_Xz[name] + Z = pca.transform(Xz).astype(np.float32) + obs = a.obs.copy() + obs["missing_hvg_frac"] = 1.0 - len(common) / G + a2 = ad.AnnData(X=sp.csr_matrix(Xf.astype(np.float32)), + obs=obs, + var=pd.DataFrame(index=shared_genes)) + a2.obsm["X_pca"] = Z + parts.append(a2) + corpus = ad.concat(parts, join="outer", label="_batch") + corpus.uns["shared_hvgs"] = shared_genes + corpus.uns["canonical_classes"] = CANONICAL_CLASSES + corpus.uns["corpus_version"] = "v3" + out_h5ad = HARM / "corpus_v3.h5ad" + corpus.write_h5ad(out_h5ad, compression="gzip") + print(f" wrote {out_h5ad} {corpus.shape}", flush=True) + + np.savez(HARM / "corpus_stats_v3.npz", + shared_hvgs=np.array(shared_genes), + mean=mu.astype(np.float32), std=sig.astype(np.float32), n_cells=N) + with open(HARM / "pca_basis_v3.pkl", "wb") as fh: + pickle.dump(pca, fh) + print(f" wrote corpus_stats_v3.npz + pca_basis_v3.pkl", flush=True) + + print("\n### Verification ###", flush=True) + ct = corpus.obs["canonical_label"].value_counts() + print("Canonical label counts:") + print(ct.to_string()) + small = ct[(ct.index != "") & (ct < 100)] + forbidden = [c for c in ["UNK", "unassigned", "hair-placode", "eccrine-ductal", ""] + if c in ct.index] + print(f"\nClasses with <100 cells: {small.to_dict()}") + print(f"Forbidden classes present: {forbidden}") + print(f"\nLabel source breakdown (must be 100% paper):") + print(corpus.obs["label_source"].value_counts().to_string()) + print("\nPer-dataset x canonical_label:") + print(corpus.obs.groupby(["dataset","canonical_label"], observed=True) + .size().unstack(fill_value=0).to_string()) + + +if __name__ == "__main__": + main() diff --git a/scripts/pan_skin/20_train_panda.py b/scripts/pan_skin/20_train_panda.py new file mode 100644 index 0000000000000000000000000000000000000000..dad7fcd70bd08d2407f540c3b0cd0ce6b428dba2 --- /dev/null +++ b/scripts/pan_skin/20_train_panda.py @@ -0,0 +1,196 @@ +"""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 + +sys.path.insert(0, "/home/bcheng/PRISM") +from panda.pan_skin.model import ( + PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, prototype_infonce +) + +CORPUS = Path("/home/bcheng/PRISM/data/corpus/pan_skin/harmonized/corpus.h5ad") +OUT = Path("/home/bcheng/PRISM/checkpoints/pan_skin") +V2_CK = Path("/home/bcheng/PRISM/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 # 0 = uniform CE weights, 1 = full inverse-sqrt-frequency +GUARANTEED_PER_CLASS = 6 # rare class cells per batch (if class has >=6 cells) +NATURAL_SLOTS = 96 # additional slots at natural frequency + + +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: + # copy prototypes for classes shared with v2 + 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() diff --git a/scripts/pan_skin/30_zero_shot_aldrich.py b/scripts/pan_skin/30_zero_shot_aldrich.py new file mode 100644 index 0000000000000000000000000000000000000000..791ffb49b16a375f7d69f017f7fb3686bcbcfda9 --- /dev/null +++ b/scripts/pan_skin/30_zero_shot_aldrich.py @@ -0,0 +1,152 @@ +"""zero-shot Aldrich (GSE220977, ~25k WT + En1-cKO volar snRNA-seq) through frozen PANDA.""" +from __future__ import annotations +from pathlib import Path +import warnings, json, sys, pickle +warnings.filterwarnings("ignore") + +import numpy as np +import pandas as pd +import anndata as ad +import scanpy as sc +import scipy.sparse as sp +import torch +import torch.nn.functional as F + +from pathlib import Path as _P_root +ROOT = _P_root(__file__).resolve().parents[2] +ROOT_STR = str(ROOT) +sys.path.insert(0, ROOT_STR) +from panda import PANDAEncoder + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +CKPT_DIR = Path(f"{ROOT_STR}/checkpoints/pan_skin") +HARM = Path(f"{ROOT_STR}/data/corpus/pan_skin/harmonized") +OUT_DIR = Path(f"{ROOT_STR}/discovery/pan_skin/marker") +OUT_DIR.mkdir(parents=True, exist_ok=True) +TARGET = Path(f"{ROOT_STR}/data/processed/skin/adata_processed.h5ad") + + +def project_to_corpus(a: ad.AnnData, shared_hvgs, mu, sig): + """z-scored dense matrix aligned to shared_hvgs, zero-impute missing.""" + n = a.n_obs + G = len(shared_hvgs) + hvg_to_ix = {g: i for i, g in enumerate(shared_hvgs)} + var_names = list(a.var_names.astype(str)) + common = [g for g in var_names if g in hvg_to_ix] + present_frac = len(common) / G + print(f"[proj] {len(common)}/{G} shared HVGs present ({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((n, G), dtype=np.float32) + cols = [hvg_to_ix[g] for g in common] + Xf[:, cols] = X + Xz = np.clip((Xf - mu.astype(np.float32)) / sig.astype(np.float32), -10, 10) + return Xz, present_frac + + +def main(): + ck = torch.load(CKPT_DIR / "panda_final.pt", map_location=DEVICE, weights_only=False) + classes = ck["classes"]; datasets = ck["datasets"] + with (CKPT_DIR / "label_encoding.json").open() as f: + enc = json.load(f) + model = PANDAEncoder(n_pca=50, n_classes=len(classes), + n_datasets=len(datasets)).to(DEVICE).eval() + model.load_state_dict(ck["model"]) + prototypes = torch.tensor(ck["prototypes"]).to(DEVICE) + print(f"[model] classes: {classes}", flush=True) + + 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_basis = pickle.load(f) + + a = ad.read_h5ad(TARGET) + if a.raw is not None: + raw = a.raw.to_adata(); raw.obs = a.obs.copy(); a = raw + print(f"[target] Aldrich shape: {a.shape}, " + f"genotype counts: {a.obs['genotype'].value_counts().to_dict()}", flush=True) + + Xz, present_frac = project_to_corpus(a, shared_hvgs, mu, sig) + counts = np.asarray(a.X.sum(axis=1)).ravel() + logc = np.log10(counts + 1); logc = (logc - logc.mean()) / (logc.std() + 1e-6) + Xpca = pca_basis.transform(Xz).astype(np.float32) + mhf = np.full(a.n_obs, 1.0 - present_frac, dtype=np.float32) + + batch = 4096 + N = a.n_obs + all_z, all_logits = [], [] + with torch.no_grad(): + for i in range(0, N, batch): + xb = torch.from_numpy(Xpca[i:i+batch]).to(DEVICE) + aux = torch.from_numpy(np.stack([mhf[i:i+batch], logc[i:i+batch]], 1).astype(np.float32)).to(DEVICE) + out = model(xb, aux, lam_dann=0.0) + all_z.append(out["z"].cpu().numpy()) + all_logits.append(out["logits"].cpu().numpy()) + Z = np.concatenate(all_z, axis=0) + logits = np.concatenate(all_logits, axis=0) + + protos_np = prototypes.cpu().numpy() + protos_np = protos_np / (np.linalg.norm(protos_np, axis=1, keepdims=True) + 1e-8) + cos = Z @ protos_np.T # (N, K) + + tau = 0.07 + proto_probs = np.exp((cos - cos.max(axis=1, keepdims=True)) / tau) + proto_probs = proto_probs / proto_probs.sum(axis=1, keepdims=True) + + pred_ix = cos.argmax(axis=1) + pred_conf = cos.max(axis=1) + entropy = -(proto_probs * np.log(proto_probs + 1e-12)).sum(axis=1) + + abstain = pred_conf < 0.3 + pred_label = np.array([classes[i] for i in pred_ix], dtype=object) + pred_label[abstain] = "UNK/abstain" + + # BBSE label-shift correction + prior_train = np.bincount(np.arange(len(classes)), minlength=len(classes)).astype(float) + 1.0 + prior_train /= prior_train.sum() + prior_test = proto_probs.mean(axis=0) + ratio = np.log(prior_test / prior_train + 1e-8) + bbse_logits = np.log(proto_probs + 1e-12) + ratio[None, :] + bbse_probs = np.exp(bbse_logits - bbse_logits.max(axis=1, keepdims=True)) + bbse_probs = bbse_probs / bbse_probs.sum(axis=1, keepdims=True) + bbse_pred = bbse_probs.argmax(axis=1) + + a.obs["pred_label"] = pred_label + a.obs["pred_conf"] = pred_conf.astype(np.float32) + a.obs["pred_entropy"] = entropy.astype(np.float32) + a.obs["abstain"] = abstain + a.obs["pred_bbse_label"] = [classes[i] for i in bbse_pred] + + tbl = a.obs[[c for c in ["genotype","sample","cell_type","fate_label", + "pred_label","pred_conf","pred_entropy", + "abstain","pred_bbse_label"] if c in a.obs.columns]].copy() + tbl.to_csv(OUT_DIR / "50_aldrich_predictions.csv") + print(f"[out] wrote {OUT_DIR/'50_aldrich_predictions.csv'} ({len(tbl)} rows)", flush=True) + + print("\n[summary] naive pred_label breakdown:") + print(a.obs["pred_label"].value_counts().head(20)) + print(f"[summary] abstain rate: {abstain.mean():.1%}") + if "genotype" in a.obs.columns: + print("\n[summary] pred_label x genotype:") + print(pd.crosstab(a.obs["pred_label"], a.obs["genotype"])) + print("\n[summary] BBSE-corrected pred_label breakdown:") + print(a.obs["pred_bbse_label"].value_counts().head(20)) + + # slim h5ad: obs + obsm only + slim = ad.AnnData(X=sp.csr_matrix((a.n_obs, 1), dtype=np.float32), obs=a.obs) + slim.obsm["Z_projection"] = Z.astype(np.float32) + slim.obsm["proto_cos"] = cos.astype(np.float32) + slim.obsm["proto_probs"] = proto_probs.astype(np.float32) + slim.obsm["bbse_probs"] = bbse_probs.astype(np.float32) + slim.uns["classes"] = classes + slim.uns["prototypes"] = protos_np + slim.write_h5ad(OUT_DIR / "50_aldrich_projections.h5ad", compression="gzip") + print(f"[out] wrote {OUT_DIR/'50_aldrich_projections.h5ad'}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/pan_skin/40_heldout_5fold_cv.py b/scripts/pan_skin/40_heldout_5fold_cv.py new file mode 100644 index 0000000000000000000000000000000000000000..efc7dc264cbd1a2dda2f3648cbf4ed46bbf9131f --- /dev/null +++ b/scripts/pan_skin/40_heldout_5fold_cv.py @@ -0,0 +1,207 @@ +"""5-fold CV on pan-skin corpus. retrain per fold, eval by prototype-cosine on held-out 20%.""" +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 torch +import torch.nn.functional as F +from torch.utils.data import Dataset, DataLoader +from sklearn.model_selection import StratifiedKFold +from sklearn.metrics import (accuracy_score, f1_score, roc_auc_score, + classification_report) + +sys.path.insert(0, "/home/bcheng/PRISM") +from panda.pan_skin.model import ( + PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, prototype_infonce +) + +CORPUS = Path("/home/bcheng/PRISM/data/corpus/pan_skin/harmonized/corpus.h5ad") +OUT = Path("/home/bcheng/PRISM/discovery/pan_skin/marker") + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +N_FOLDS = 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} + 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_pick = self.rng.choice(self.classes, p=self.p) + batch.append(int(self.rng.choice(self.by_cls[c_pick]))) + yield batch + def __len__(self): return self.n_batches + + +def train_one_fold(X, y, d, mhf, logc, classes, datasets, tr, te, fold_id, log_prefix): + torch.cuda.empty_cache() + Xtr, ytr, dtr, mtr, ltr = X[tr], y[tr], d[tr], mhf[tr], logc[tr] + ds = CorpusDataset(Xtr, ytr, dtr, mtr, ltr) + sampler = HybridSampler(ytr, dtr, n_batches=100, seed=fold_id) + loader = DataLoader(ds, batch_sampler=sampler, num_workers=0) + + counts_per = np.bincount(ytr, minlength=len(classes)) + inv_sqrt = 1.0 / np.sqrt(counts_per + 1) + inv_sqrt = inv_sqrt / inv_sqrt.mean() + class_w_np = 0.5 * inv_sqrt + 0.5 * np.ones_like(inv_sqrt) + class_w = torch.tensor(class_w_np, dtype=torch.float32).to(DEVICE) + + 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): + for e in range(stage_epochs[stage]): + if e % 10 == 0: + print(f"{log_prefix} s{stage} ep {e}/{stage_epochs[stage]}", flush=True) + 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 + 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) + + model.eval() + Xte, yte = X[te], y[te] + with torch.no_grad(): + Xt = torch.from_numpy(Xte.astype(np.float32)).to(DEVICE) + aux = torch.zeros(len(te), 2, device=DEVICE) + out = model(Xt, aux, lam_dann=0.0) + z = out["z"] + cos = z @ model.prototypes.T + pred = cos.argmax(dim=1).cpu().numpy() + probs = torch.softmax(cos / 0.07, dim=1).cpu().numpy() + + acc = accuracy_score(yte, pred) + f1 = f1_score(yte, pred, average="macro", zero_division=0) + try: + auc = roc_auc_score(np.eye(len(classes))[yte], probs, average="macro", multi_class="ovr") + except Exception: + auc = float("nan") + print(f"{log_prefix} acc={acc:.4f} macro_f1={f1:.4f} macro_auc={auc:.4f}", flush=True) + return acc, f1, auc, pred, yte + + +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)} + X = np.asarray(a.obsm["X_pca"]) + 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)]) + 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) + + print(f"[cv] corpus: {a.shape} n_classes={len(classes)} n_datasets={len(datasets)}", + flush=True) + print(f"[cv] class counts: {dict(zip(classes, np.bincount(y, minlength=len(classes)).tolist()))}", + flush=True) + + skf = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=42) + accs, f1s, aucs = [], [], [] + all_preds = [] + partial_path = OUT / "60_heldout_5fold_partial.json" + for fold, (tr, te) in enumerate(skf.split(X, y)): + t0 = time.time() + try: + acc, f1, auc, pred, yte = train_one_fold( + X, y, d, mhf, logc, classes, datasets, tr, te, + fold_id=fold, log_prefix=f"[fold {fold}]") + except Exception as exc: + import traceback; traceback.print_exc() + print(f"[fold {fold}] FAILED: {exc}", flush=True) + continue + print(f"[fold {fold}] wall={time.time()-t0:.0f}s", flush=True) + accs.append(acc); f1s.append(f1); aucs.append(auc) + all_preds.append({"fold": fold, "test_idx": te.tolist(), + "pred": pred.tolist(), "true": yte.tolist()}) + with open(partial_path, "w") as f: + json.dump({"folds_done": fold + 1, "accs": accs, "f1s": f1s, "aucs": aucs}, f) + + print(f"\n[cv] 5-FOLD ACC: {np.mean(accs):.4f} +- {np.std(accs):.4f}") + print(f"[cv] 5-FOLD F1: {np.mean(f1s):.4f} +- {np.std(f1s):.4f}") + print(f"[cv] 5-FOLD AUC: {np.mean(aucs):.4f} +- {np.std(aucs):.4f}") + + all_y_true = np.concatenate([np.array(p["true"]) for p in all_preds]) + all_y_pred = np.concatenate([np.array(p["pred"]) for p in all_preds]) + print("\n[cv] Concatenated held-out classification report:") + rep = classification_report(all_y_true, all_y_pred, target_names=classes, + digits=3, zero_division=0, output_dict=True) + print(classification_report(all_y_true, all_y_pred, target_names=classes, + digits=3, zero_division=0)) + + result = { + "mean_acc": float(np.mean(accs)), "std_acc": float(np.std(accs)), + "mean_f1": float(np.mean(f1s)), "std_f1": float(np.std(f1s)), + "mean_auc": float(np.mean(aucs)), "std_auc": float(np.std(aucs)), + "per_fold_acc": accs, "per_fold_f1": f1s, "per_fold_auc": aucs, + "per_class_report": rep, + "n_folds": N_FOLDS, "n_classes": len(classes), + "protocol": "StratifiedKFold(5) retrain from scratch per fold; " + "eval by prototype-cosine argmax on held-out 20%.", + } + (OUT / "60_heldout_5fold_cv.json").write_text(json.dumps(result, indent=2)) + print(f"\n[cv] wrote {OUT}/60_heldout_5fold_cv.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/pan_skin/91_retrain_without_sulic.py b/scripts/pan_skin/91_retrain_without_sulic.py new file mode 100644 index 0000000000000000000000000000000000000000..1c911853591a7ca327a31a255d8d23d6bbcfe18f --- /dev/null +++ b/scripts/pan_skin/91_retrain_without_sulic.py @@ -0,0 +1,194 @@ +"""retrain pca+marker with sulic GSE212673 excluded; score zero-shot on held-out sulic.""" +from pathlib import Path +import warnings, json, sys, pickle, subprocess, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp, torch, torch.nn.functional as F, yaml +warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 +sys.path.insert(0, "/home/bcheng/PRISM") +from panda import PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, subcenter_angular_infonce, prototype_repulsion +from sklearn.decomposition import PCA +from sklearn.metrics import accuracy_score, f1_score, classification_report + +ROOT = Path("/home/bcheng/PRISM") +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +CORP = ROOT / "data/corpus/pan_skin/harmonized" +OUT_CORP = CORP # write suffixed variants alongside + + +def build_no_sulic_corpus(): + a = ad.read_h5ad(CORP / "corpus.h5ad") + keep = a.obs["dataset"] != "sulic_GSE212673" + a2 = a[keep].copy() + print(f"[corpus] {a.n_obs} → {a2.n_obs} after removing sulic ({(~keep).sum()} cells dropped)", flush=True) + a2.write_h5ad(CORP / "corpus_no_sulic.h5ad") + + # refit mean/std + pca on training-only subset (same hvgs) + stats = np.load(CORP / "corpus_stats.npz", allow_pickle=True) + hvgs = [str(g) for g in stats["shared_hvgs"]] + hvg2i = {g: i for i, g in enumerate(hvgs)} + common = [g for g in a2.var_names.astype(str) if g in hvg2i] + a2c = a2[:, common].copy() + sc.pp.normalize_total(a2c, target_sum=1e4); sc.pp.log1p(a2c) + X = a2c.X.toarray().astype(np.float32) if sp.issparse(a2c.X) else a2c.X.astype(np.float32) + Xf = np.zeros((a2.n_obs, len(hvgs)), dtype=np.float32) + Xf[:, np.array([hvg2i[g] for g in common])] = X + mu = Xf.mean(axis=0); sig = Xf.std(axis=0) + 1e-6 + np.savez(CORP / "corpus_stats_no_sulic.npz", shared_hvgs=np.asarray(hvgs, dtype=object), + mean=mu, std=sig) + Xz = np.clip((Xf - mu) / sig, -10, 10) + n_pca = 50 + rng = np.random.default_rng(0) + idx = rng.choice(Xz.shape[0], size=min(30000, Xz.shape[0]), replace=False) + pca = PCA(n_components=n_pca, random_state=0).fit(Xz[idx]) + with open(CORP / "pca_basis_no_sulic.pkl", "wb") as f: pickle.dump(pca, f) + print(f"[corpus] refit stats + pca on {Xz.shape[0]} cells (subsample {len(idx)} for pca)", flush=True) + return a2, hvgs, mu, sig, pca + + +def prepare_batches(adata, hvgs, mu, sig, pca, marker_genes=None, variant="pca"): + 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() + 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) + + 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(a, hvgs, mu, sig, pca, variant, marker_genes, epochs=8, batch=256, lr=1e-3): + Xpca, Xmark, y, classes, y_dset, datasets, log10cz = prepare_batches(a, hvgs, mu, sig, pca, marker_genes, variant) + print(f"[train] {variant} n={a.n_obs} K={len(classes)} datasets={len(datasets)}", 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) + if stage >= 3: + L = L + 0.5 * prototype_repulsion(model.prototypes.detach().clone()) + 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 {variant}] epoch {epoch}/{epochs} stage={stage} loss={np.mean(losses):.4f}", flush=True) + + ck_dir = ROOT / f"checkpoints/pan_skin_no_sulic/{variant}" + ck_dir.mkdir(parents=True, exist_ok=True) + torch.save({"model": model.state_dict(), "classes": classes, "datasets": datasets, + "marker_genes": marker_genes if variant == "marker" else [], + "prototypes": model.prototypes.detach().cpu().numpy()}, + ck_dir / "panda_final.pt") + print(f"[save] {ck_dir}/panda_final.pt", flush=True) + return model, classes, datasets + + +def infer_on_sulic(variant, classes, marker_genes, hvgs, mu, sig, pca): + """subset to sulic, run inference, evaluate.""" + a = ad.read_h5ad(CORP / "corpus.h5ad") + a_s = a[a.obs["dataset"] == "sulic_GSE212673"].copy() + Xpca, Xmark, y, cls_in_sulic, _, _, _ = prepare_batches(a_s, hvgs, mu, sig, pca, marker_genes, variant) + y_true = a_s.obs["canonical_label"].astype(str).values + + ck = torch.load(ROOT / f"checkpoints/pan_skin_no_sulic/{variant}/panda_final.pt", + map_location=DEVICE, weights_only=False) + n_markers = len(ck.get("marker_genes", [])) if variant == "marker" else 0 + model = PANDAEncoder(variant=variant, n_pca=50, n_markers=n_markers, + n_classes=len(ck["classes"]), n_sub=3, + n_datasets=len(ck["datasets"])).to(DEVICE).eval() + model.load_state_dict(ck["model"]) + + preds, probs = [], [] + with torch.no_grad(): + for i in range(0, a_s.n_obs, 4096): + xb = torch.from_numpy(Xpca[i:i+4096]).to(DEVICE) + xmb = torch.from_numpy(Xmark[i:i+4096]).to(DEVICE) if Xmark is not None else None + aux = torch.zeros(len(xb), 2, device=DEVICE) + out = model(xb, aux, x_markers=xmb, lam_dann=0.0) + mc = model.max_sub_cos(out["z"]) + preds.append(mc.argmax(dim=1).cpu().numpy()) + probs.append(F.softmax(mc / 0.07, dim=1).cpu().numpy()) + pred = np.array([ck["classes"][i] for i in np.concatenate(preds)]) + probs = np.concatenate(probs) + + mask = np.isin(y_true, ck["classes"]) + acc = accuracy_score(y_true[mask], pred[mask]) + f1 = f1_score(y_true[mask], pred[mask], average="macro", zero_division=0) + rep = classification_report(y_true[mask], pred[mask], zero_division=0, output_dict=True) + print(f"[eval-{variant}] sulic zero-shot n={mask.sum()} acc={acc:.4f} macro-f1={f1:.4f}", flush=True) + + out = ROOT / f"discovery/pan_skin/{variant}" + out.mkdir(parents=True, exist_ok=True) + (out / "96_sulic_zero_shot.json").write_text(json.dumps({ + "variant": variant, "n_cells_eval": int(mask.sum()), + "acc": float(acc), "macro_f1": float(f1), + "per_class": rep, + "predicted_dist": pd.Series(pred).value_counts().to_dict(), + "true_dist": pd.Series(y_true).value_counts().to_dict(), + }, indent=2, default=str)) + pd.DataFrame({"cell_id": a_s.obs_names, "true_label": y_true, "pred_label": pred, + "max_cos": probs.max(axis=1)}).to_csv(out / "96_sulic_predictions.csv", index=False) + + +def main(): + print("[1/3] build no-sulic corpus", flush=True) + a, hvgs, mu, sig, pca = build_no_sulic_corpus() + + print("\n[2/3] train pca variant", flush=True) + train(a, hvgs, mu, sig, pca, "pca", [], epochs=8) + + print("\n[3/3] train marker variant", flush=True) + marker_yaml = yaml.safe_load(open(ROOT / "panda/markers.yaml")) + marker_genes = marker_yaml["pan_skin"] + train(a, hvgs, mu, sig, pca, "marker", marker_genes, epochs=8) + + print("\n[eval] pca zero-shot on sulic", flush=True) + infer_on_sulic("pca", None, [], hvgs, mu, sig, pca) + print("\n[eval] marker zero-shot on sulic", flush=True) + infer_on_sulic("marker", None, marker_genes, hvgs, mu, sig, pca) + + +if __name__ == "__main__": + main() diff --git a/scripts/pan_skin/92_retrain_with_sulic_anchor.py b/scripts/pan_skin/92_retrain_with_sulic_anchor.py new file mode 100644 index 0000000000000000000000000000000000000000..f34fb5de55f65039ca308fc5dbef98454b4889c3 --- /dev/null +++ b/scripts/pan_skin/92_retrain_with_sulic_anchor.py @@ -0,0 +1,205 @@ +"""retrain with 500-cell sulic anchor (300 HF-placode + 200 basal-IFE); eval on 4,183 held-out.""" +from pathlib import Path +import warnings, json, sys, pickle, numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp, torch, torch.nn.functional as F, yaml +warnings.filterwarnings("ignore"); sc.settings.verbosity = 0 +sys.path.insert(0, "/home/bcheng/PRISM") +from panda import (PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, + subcenter_angular_infonce, prototype_repulsion) +from sklearn.decomposition import PCA +from sklearn.metrics import accuracy_score, f1_score, classification_report + +ROOT = Path("/home/bcheng/PRISM") +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +CORP_DIR = ROOT / "data/corpus/pan_skin/harmonized" + + +def split_sulic(seed=0): + """300 HF-placode + 200 basal-IFE to anchor, rest held out.""" + a = ad.read_h5ad(CORP_DIR / "corpus.h5ad") + su_mask = (a.obs["dataset"] == "sulic_GSE212673").values + a_sulic = a[su_mask].copy() + a_nonsu = a[~su_mask].copy() + a_sulic.obs["_orig_idx"] = np.arange(a_sulic.n_obs) + rng = np.random.default_rng(seed) + hf = np.where(a_sulic.obs["canonical_label"].astype(str).values == "HF-placode")[0] + bi = np.where(a_sulic.obs["canonical_label"].astype(str).values == "basal-IFE")[0] + print(f"[split] sulic total {a_sulic.n_obs} (HF-placode {len(hf)}, basal-IFE {len(bi)})", flush=True) + hf_anchor = rng.choice(hf, size=min(300, len(hf)), replace=False) + bi_anchor = rng.choice(bi, size=min(200, len(bi)), replace=False) + anchor_ix = np.concatenate([hf_anchor, bi_anchor]) + test_ix = np.setdiff1d(np.arange(a_sulic.n_obs), anchor_ix) + print(f"[split] anchor {len(anchor_ix)} / held-out {len(test_ix)}", flush=True) + a_anchor = a_sulic[anchor_ix].copy() + a_test = a_sulic[test_ix].copy() + return a_nonsu, a_anchor, a_test + + +def build_anchor_corpus(): + a_nonsu, a_anchor, a_test = split_sulic() + a_full = ad.concat([a_nonsu, a_anchor], join="outer") + print(f"[build] anchor-augmented corpus {a_full.shape} = {a_nonsu.n_obs} non-sulic + {a_anchor.n_obs} anchor", flush=True) + a_full.write_h5ad(CORP_DIR / "corpus_with_sulic_anchor.h5ad") + + stats = np.load(CORP_DIR / "corpus_stats.npz", allow_pickle=True) + hvgs = [str(g) for g in stats["shared_hvgs"]] + hvg2i = {g: i for i, g in enumerate(hvgs)} + common = [g for g in a_full.var_names.astype(str) if g in hvg2i] + a_c = a_full[:, 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_full.n_obs, len(hvgs)), dtype=np.float32) + Xf[:, np.array([hvg2i[g] for g in common])] = X + mu = Xf.mean(axis=0); sig = Xf.std(axis=0) + 1e-6 + np.savez(CORP_DIR / "corpus_stats_with_sulic_anchor.npz", + shared_hvgs=np.asarray(hvgs, dtype=object), mean=mu, std=sig) + Xz = np.clip((Xf - mu) / sig, -10, 10) + rng = np.random.default_rng(0) + idx = rng.choice(Xz.shape[0], size=min(30000, Xz.shape[0]), replace=False) + pca = PCA(n_components=50, random_state=0).fit(Xz[idx]) + with open(CORP_DIR / "pca_basis_with_sulic_anchor.pkl", "wb") as f: pickle.dump(pca, f) + return a_full, hvgs, mu, sig, pca, a_test + + +def prepare_batches(adata, hvgs, mu, sig, pca, marker_genes=None, variant="pca", need_labels=True): + 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() + 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) + Xf[:, np.array([hvg2i[g] for g in common])] = 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) + if not need_labels: + return Xpca, Xmark + 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(a, hvgs, mu, sig, pca, variant, marker_genes, epochs=8, batch=256, lr=1e-3): + Xpca, Xmark, y, classes, y_dset, datasets, log10cz = prepare_batches(a, hvgs, mu, sig, pca, marker_genes, variant) + print(f"[train] {variant} n={a.n_obs} K={len(classes)} datasets={len(datasets)}", 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) + if stage >= 3: + L = L + 0.5 * prototype_repulsion(model.prototypes.detach().clone()) + 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 {variant}] epoch {epoch}/{epochs} stage={stage} loss={np.mean(losses):.4f}", flush=True) + ck_dir = ROOT / f"checkpoints/pan_skin_anchor/{variant}" + ck_dir.mkdir(parents=True, exist_ok=True) + torch.save({"model": model.state_dict(), "classes": classes, "datasets": datasets, + "marker_genes": marker_genes if variant == "marker" else [], + "prototypes": model.prototypes.detach().cpu().numpy()}, + ck_dir / "panda_final.pt") + print(f"[save] {ck_dir}/panda_final.pt", flush=True) + + +def infer_on_test(variant, marker_genes, hvgs, mu, sig, pca, a_test): + Xpca, Xmark = prepare_batches(a_test, hvgs, mu, sig, pca, marker_genes, variant, need_labels=False) + ck = torch.load(ROOT / f"checkpoints/pan_skin_anchor/{variant}/panda_final.pt", + map_location=DEVICE, weights_only=False) + classes = ck["classes"] + model = PANDAEncoder(variant=variant, n_pca=50, + n_markers=len(marker_genes) if variant == "marker" else 0, + n_classes=len(classes), n_sub=3, + n_datasets=len(ck["datasets"])).to(DEVICE).eval() + model.load_state_dict(ck["model"]) + preds, probs = [], [] + with torch.no_grad(): + for i in range(0, a_test.n_obs, 4096): + xb = torch.from_numpy(Xpca[i:i+4096]).to(DEVICE) + xmb = torch.from_numpy(Xmark[i:i+4096]).to(DEVICE) if Xmark is not None else None + aux = torch.zeros(len(xb), 2, device=DEVICE) + out = model(xb, aux, x_markers=xmb, lam_dann=0.0) + mc = model.max_sub_cos(out["z"]) + preds.append(mc.argmax(dim=1).cpu().numpy()) + probs.append(F.softmax(mc / 0.07, dim=1).cpu().numpy()) + pred = np.array([classes[i] for i in np.concatenate(preds)]) + probs = np.concatenate(probs) + y_true = a_test.obs["canonical_label"].astype(str).values + mask = np.isin(y_true, classes) + acc = accuracy_score(y_true[mask], pred[mask]) + f1 = f1_score(y_true[mask], pred[mask], average="macro", zero_division=0) + rep = classification_report(y_true[mask], pred[mask], zero_division=0, output_dict=True) + hf_recall = float(((pred == "HF-placode") & (y_true == "HF-placode")).sum() / max(1, (y_true == "HF-placode").sum())) + bi_recall = float(((pred == "basal-IFE") & (y_true == "basal-IFE")).sum() / max(1, (y_true == "basal-IFE").sum())) + print(f"[eval-{variant}] n={mask.sum()} acc={acc:.4f} macro-f1={f1:.4f} HF-placode-recall={hf_recall:.4f} basal-IFE-recall={bi_recall:.4f}", flush=True) + out = ROOT / f"discovery/pan_skin/{variant}" + out.mkdir(parents=True, exist_ok=True) + (out / "97_sulic_anchor_zero_shot.json").write_text(json.dumps({ + "variant": variant, + "n_test": int(mask.sum()), + "acc": float(acc), "macro_f1": float(f1), + "HF_placode_recall": hf_recall, "basal_IFE_recall": bi_recall, + "per_class": rep, + "pred_dist": pd.Series(pred).value_counts().to_dict(), + "true_dist": pd.Series(y_true).value_counts().to_dict(), + }, indent=2, default=str)) + pd.DataFrame({"cell_id": a_test.obs_names, "true_label": y_true, "pred_label": pred, + "max_cos": probs.max(axis=1)}).to_csv(out / "97_sulic_anchor_predictions.csv", index=False) + + +def main(): + print("[1/3] build anchor corpus", flush=True) + a, hvgs, mu, sig, pca, a_test = build_anchor_corpus() + + print("\n[2/3] train pca variant", flush=True) + train(a, hvgs, mu, sig, pca, "pca", [], epochs=8) + + print("\n[3/3] train marker variant", flush=True) + marker_genes = yaml.safe_load(open(ROOT / "panda/markers.yaml"))["pan_skin"] + train(a, hvgs, mu, sig, pca, "marker", marker_genes, epochs=8) + + print("\n[eval] pca on held-out sulic slice", flush=True) + infer_on_test("pca", [], hvgs, mu, sig, pca, a_test) + print("\n[eval] marker on held-out sulic slice", flush=True) + infer_on_test("marker", marker_genes, hvgs, mu, sig, pca, a_test) + + +if __name__ == "__main__": + main() diff --git a/scripts/pan_skin/93_add_belote_anchor.py b/scripts/pan_skin/93_add_belote_anchor.py new file mode 100644 index 0000000000000000000000000000000000000000..e957ef49c0cac004426ffa7dc39bf13ae1333c81 --- /dev/null +++ b/scripts/pan_skin/93_add_belote_anchor.py @@ -0,0 +1,276 @@ +"""add belote 2021 (GSE151091) human melanocyte anchor (~1k cells) to corpus_v3, holdout ~6k.""" +from __future__ import annotations +from pathlib import Path +import warnings, pickle, gzip, sys +warnings.filterwarnings("ignore") + +import numpy as np +import pandas as pd +import scanpy as sc +import anndata as ad +import scipy.sparse as sp +from sklearn.decomposition import PCA + +ROOT = Path("/home/bcheng/PRISM") +HARM = ROOT / "data/corpus/pan_skin/harmonized" +HELD = ROOT / "data/corpus/pan_skin/held_out_labeled" +EXT = ROOT / "data/external_labels/melanocyte_anchor" +HELD.mkdir(parents=True, exist_ok=True) + +# HGNC→mouse special cases; otherwise capitalize (ORF-style keeps case pattern). +HGNC_TO_MOUSE_SPECIAL = { + "TP53": "Trp53", + "TP63": "Trp63", + "TP73": "Trp73", +} + + +def hgnc_to_mouse(sym: str) -> str: + if sym in HGNC_TO_MOUSE_SPECIAL: + return HGNC_TO_MOUSE_SPECIAL[sym] + low = sym.lower() + if "orf" in low: + # Corf style + return sym[:1].upper() + sym[1:].lower() + return sym.capitalize() + + +def read_belote_labels(): + df = pd.read_csv(EXT / "canonical_labels.csv") + print(f"[belote] {len(df)} labeled cells", flush=True) + print(df["canonical"].value_counts().to_string(), flush=True) + return df + + +def read_belote_counts(keep_cells: set[str]) -> ad.AnnData: + """read genes×cells csv, subset to labeled cells, return cells×genes AnnData.""" + csv = EXT / "GSE151091_raw_matrix.csv.gz" + print(f"[belote] reading {csv.name} (streaming)...", flush=True) + with gzip.open(csv, "rt") as fh: + header = fh.readline().rstrip("\n").split(",") + cell_ids = header[1:] # first entry is empty for row index + keep_col_idx = [i for i, c in enumerate(cell_ids) if c in keep_cells] + keep_cell_ids = [cell_ids[i] for i in keep_col_idx] + print(f"[belote] {len(keep_cell_ids)}/{len(cell_ids)} labeled cells found in matrix", + flush=True) + if len(keep_cell_ids) != len(keep_cells): + missing = keep_cells - set(keep_cell_ids) + print(f" WARN {len(missing)} labeled cells not in matrix", flush=True) + + # +1 shifts past the row-index col (unnamed first col in the source csv) + usecols = [0] + [i + 1 for i in keep_col_idx] + print(f"[belote] pd.read_csv usecols={len(usecols)} (genes×cells) ...", flush=True) + dtypes = {c: np.float32 for c in keep_cell_ids} + df = pd.read_csv(csv, index_col=0, usecols=usecols, dtype=dtypes, + low_memory=False) + print(f"[belote] loaded {df.shape} (genes × cells), transposing...", flush=True) + df = df[keep_cell_ids] + X = df.values.T.astype(np.float32) # cells × genes + genes = df.index.astype(str).tolist() + print(f"[belote] X {X.shape} (cells × genes), nnz~{(X>0).mean()*100:.1f}%", + flush=True) + a = ad.AnnData(X=sp.csr_matrix(X), + obs=pd.DataFrame(index=keep_cell_ids), + var=pd.DataFrame(index=genes)) + a.var_names_make_unique() + return a + + +def convert_human_to_mouse(a: ad.AnnData, mouse_hvgs: list[str]) -> ad.AnnData: + """map var_names HGNC→mouse and sum-collapse duplicates; log coverage vs corpus HVGs.""" + n_hgnc = a.n_vars + mouse_names = [hgnc_to_mouse(g) for g in a.var_names.astype(str)] + m_index = pd.Index(mouse_names) + unique_mouse = m_index.unique().tolist() + print(f"[map] {n_hgnc} HGNC → {len(unique_mouse)} unique mouse symbols", + flush=True) + if len(unique_mouse) != n_hgnc: + mouse_to_col = {g: i for i, g in enumerate(unique_mouse)} + proj = sp.lil_matrix((n_hgnc, len(unique_mouse)), dtype=np.float32) + for c, mg in enumerate(mouse_names): + proj[c, mouse_to_col[mg]] = 1.0 + proj = proj.tocsr() + X2 = a.X @ proj + a = ad.AnnData(X=X2, obs=a.obs.copy(), + var=pd.DataFrame(index=unique_mouse)) + else: + a.var_names = unique_mouse + + hvg_set = set(mouse_hvgs) + present = [g for g in a.var_names if g in hvg_set] + print(f"[map] mouse-var coverage on corpus HVGs: " + f"{len(present)}/{len(mouse_hvgs)} " + f"({len(present)/len(mouse_hvgs)*100:.1f}%)", flush=True) + return a + + +def stratified_split(labels: pd.Series, targets: dict, seed=0): + rng = np.random.default_rng(seed) + anchor_ix, test_ix = [], [] + for cls, n in targets.items(): + pool = np.where(labels.values == cls)[0] + rng.shuffle(pool) + take = min(n, len(pool)) + anchor_ix.append(pool[:take]) + test_ix.append(pool[take:]) + return np.concatenate(anchor_ix), np.concatenate(test_ix) + + +def normalize_on_hvgs(a: ad.AnnData, shared_hvgs: list[str]): + """dense log1p normalize_total on shared HVGs (missing filled 0), plus missing_hvg_frac.""" + G = len(shared_hvgs) + hvg2i = {g: i for i, g in enumerate(shared_hvgs)} + common = [g for g in a.var_names if g in hvg2i] + a_c = a[:, common].copy() + sc.pp.normalize_total(a_c, target_sum=1e4) + sc.pp.log1p(a_c) + X = a_c.X.toarray() if sp.issparse(a_c.X) else a_c.X + Xf = np.zeros((a.n_obs, G), dtype=np.float32) + cols = np.array([hvg2i[g] for g in common], dtype=np.int64) + Xf[:, cols] = X.astype(np.float32) + missing_frac = 1.0 - len(common) / G + return Xf, missing_frac + + +def main(): + print("[1/6] load corpus_v3", flush=True) + corpus = ad.read_h5ad(HARM / "corpus_v3.h5ad") + shared_hvgs = list(corpus.var_names.astype(str)) + print(f" corpus {corpus.shape}, shared HVGs {len(shared_hvgs)}", flush=True) + print(f" class counts:\n{corpus.obs['canonical_label'].value_counts().to_string()}", + flush=True) + + print("\n[2/6] read Belote labels + counts", flush=True) + labels_df = read_belote_labels() + keep_cells = set(labels_df["cell_id"].astype(str)) + a_bel = read_belote_counts(keep_cells) + a_bel = a_bel[labels_df["cell_id"].values].copy() + a_bel.obs = labels_df.set_index("cell_id").loc[a_bel.obs_names].copy() + print(f" belote AnnData {a_bel.shape}", flush=True) + + print("\n[3/6] convert human HGNC → mouse", flush=True) + a_bel = convert_human_to_mouse(a_bel, shared_hvgs) + + print("\n[4/6] stratified anchor/held-out split", flush=True) + targets = {"melanocyte": 700, "melanoblast": 250, "melanocyte-precursor": 50} + anchor_ix, test_ix = stratified_split(a_bel.obs["canonical"], targets, seed=0) + print(f" anchor {len(anchor_ix)} | held-out {len(test_ix)}", flush=True) + a_anchor = a_bel[anchor_ix].copy() + a_test = a_bel[test_ix].copy() + + def _attach_meta(a): + a.obs["dataset"] = "belote_GSE151091_anchor" + a.obs["canonical_label"] = a.obs["canonical"].astype(str).values + a.obs["label_source"] = "paper" + a.obs["species"] = "human" + # preserve dev_stage as stage_tag (fet/neo/adt) + a.obs["stage_tag"] = a.obs["dev_stage"].astype(str).values + a.obs["paper_label"] = a.obs["canonical"].astype(str).values + return a + a_anchor = _attach_meta(a_anchor) + a_test = _attach_meta(a_test) + a_test.obs["dataset"] = "belote_GSE151091_holdout" + + print("\n[5/6] normalize anchor on shared HVGs + concat with corpus_v3", flush=True) + Xf_anchor, missing_frac_anchor = normalize_on_hvgs(a_anchor, shared_hvgs) + print(f" anchor missing_hvg_frac = {missing_frac_anchor:.3f}", flush=True) + + anchor_norm = ad.AnnData( + X=sp.csr_matrix(Xf_anchor.astype(np.float32)), + obs=a_anchor.obs.copy(), + var=pd.DataFrame(index=shared_hvgs), + ) + anchor_norm.obs["missing_hvg_frac"] = missing_frac_anchor + + corpus_aug = ad.concat([corpus, anchor_norm], join="outer", label="_batch", + merge="unique") + corpus_aug.uns["canonical_classes"] = list(corpus.uns.get("canonical_classes", [])) + corpus_aug.uns["shared_hvgs"] = shared_hvgs + corpus_aug.uns["corpus_version"] = "belote_anchor" + print(f" augmented corpus {corpus_aug.shape}", flush=True) + + print("\n refitting per-gene mean/std over augmented corpus...", flush=True) + X = corpus_aug.X + n = corpus_aug.n_obs + mu = np.asarray(X.mean(axis=0)).ravel().astype(np.float64) + Xsq = X.multiply(X) + ex2 = np.asarray(Xsq.mean(axis=0)).ravel().astype(np.float64) + sig = np.sqrt(np.maximum(ex2 - mu**2, 1e-6)) + print(f" mean/std over {n} cells", flush=True) + + print("\n refitting PCA (30k stratified subsample)...", flush=True) + rng = np.random.default_rng(0) + per_ds_take = max(1, 30_000 // corpus_aug.obs["dataset"].nunique()) + fit_idx = [] + for ds, grp in corpus_aug.obs.groupby("dataset", observed=True): + pool = np.where(corpus_aug.obs["dataset"].values == ds)[0] + take = min(per_ds_take, len(pool)) + fit_idx.append(rng.choice(pool, size=take, replace=False)) + fit_idx = np.concatenate(fit_idx) + print(f" Xfit shape: ({len(fit_idx)}, {len(shared_hvgs)})", flush=True) + Xfit_raw = corpus_aug.X[fit_idx].toarray().astype(np.float32) + Xfit_z = np.clip((Xfit_raw - mu.astype(np.float32)) / sig.astype(np.float32), + -10, 10) + pca = PCA(n_components=50, random_state=42) + pca.fit(Xfit_z) + for k in range(pca.n_components_): + top = int(np.argmax(np.abs(pca.components_[k]))) + if pca.components_[k, top] < 0: + pca.components_[k] *= -1 + print(f" explained var: {pca.explained_variance_ratio_.sum():.3f}", flush=True) + + print(" computing X_pca for augmented corpus...", flush=True) + X_all = corpus_aug.X.toarray().astype(np.float32) + Xz_all = np.clip((X_all - mu.astype(np.float32)) / sig.astype(np.float32), + -10, 10) + corpus_aug.obsm["X_pca"] = pca.transform(Xz_all).astype(np.float32) + + print("\n[6/6] write outputs", flush=True) + out_h5ad = HARM / "corpus_v3.h5ad" + corpus_aug.write_h5ad(out_h5ad, compression="gzip") + print(f" {out_h5ad} {corpus_aug.shape}", flush=True) + np.savez(HARM / "corpus_stats_v3.npz", + shared_hvgs=np.array(shared_hvgs), + mean=mu.astype(np.float32), + std=sig.astype(np.float32), + n_cells=int(n)) + with open(HARM / "pca_basis_v3.pkl", "wb") as fh: + pickle.dump(pca, fh) + print(f" corpus_stats_v3.npz + pca_basis_v3.pkl", flush=True) + + # held-out slice normalized on shared HVGs + Xf_test, missing_test = normalize_on_hvgs(a_test, shared_hvgs) + a_test_norm = ad.AnnData( + X=sp.csr_matrix(Xf_test.astype(np.float32)), + obs=a_test.obs.copy(), + var=pd.DataFrame(index=shared_hvgs), + ) + a_test_norm.obs["missing_hvg_frac"] = missing_test + a_test_norm.uns["shared_hvgs"] = shared_hvgs + test_path = HELD / "belote_GSE151091_test.h5ad" + a_test_norm.write_h5ad(test_path, compression="gzip") + print(f" {test_path} {a_test_norm.shape}", flush=True) + + print("\n### Verification ###", flush=True) + ct = corpus_aug.obs["canonical_label"].value_counts() + print("Final canonical_label counts:") + print(ct.to_string()) + print(f"\nn classes: {ct.size} total cells: {corpus_aug.n_obs}") + print(f"\nLabel source (must be 100% paper):") + print(corpus_aug.obs["label_source"].value_counts().to_string()) + print(f"\nDatasets:") + print(corpus_aug.obs["dataset"].value_counts().to_string()) + forbidden = [c for c in ["UNK", "unassigned", "", "abstain", "unmapped"] + if c in ct.index] + small = ct[ct < 150] + print(f"\nForbidden classes present: {forbidden}") + print(f"Classes with <150 cells: {small.to_dict()}") + + added = [c for c in ["melanocyte", "melanoblast", "melanocyte-precursor"] + if c in ct.index] + print(f"\nAdded classes ({len(added)}): " + f"{ {c: int(ct[c]) for c in added} }") + + +if __name__ == "__main__": + main() diff --git a/scripts/pan_skin/README.md b/scripts/pan_skin/README.md new file mode 100644 index 0000000000000000000000000000000000000000..592eec6ee5bba45b6f3d7cef27c5a291f2f0b32a --- /dev/null +++ b/scripts/pan_skin/README.md @@ -0,0 +1,19 @@ +# scripts/pan_skin + +canonical end-to-end pipeline for the pan-skin corpus. steps are numbered in the +order they need to run. everything writes into `data/corpus/pan_skin/`, +`checkpoints/pan_skin/`, and `figures/`. + +| step | script | what it does | +|---|---|---| +| 01-03 | `01_download_tier_a.sh` .. `03_download_tier_c.sh` | fetch tier A/B/C raw h5ads + counts | +| 06 | `06_build_per_dataset_h5ads.py` | per-dataset QC + h5ad build | +| 07 | `07_build_shared_hvgs_and_pca.py` | shared HVGs + PCA basis over the union | +| 08 | `08_assign_labels.py` (+ `08b_curated_label_override.py`) | marker-scored labels, Sulic overrides | +| 10 | `10_build_corpus.py` | merged harmonized corpus.h5ad | +| 20 | `20_train_panda.py` | PANDA-MLP with hybrid sampling | +| 30 | `30_zero_shot_aldrich.py` | inference on Aldrich / Dingwall (held-out target) | +| 40 | `40_heldout_5fold_cv.py` | 5-fold CV, PANDA retrained from scratch per fold | + +`run_all.sh` chains these in order. marker panels: `known_skin_tfs.yaml`; ontology +map: `ontology_map.yaml`. diff --git a/scripts/pan_skin/__init__.py b/scripts/pan_skin/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/scripts/pan_skin/_loaders/__init__.py b/scripts/pan_skin/_loaders/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/scripts/pan_skin/known_skin_tfs.yaml b/scripts/pan_skin/known_skin_tfs.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c8178c705aa798067a0a556a0fc8fe26601d39c9 --- /dev/null +++ b/scripts/pan_skin/known_skin_tfs.yaml @@ -0,0 +1,29 @@ +# canonical mouse skin marker genes used ONLY for label assignment during corpus build. +# not used as input features to the classifier (PCA-only ethos preserved). +# each canonical class -> ordered marker list; label = argmax(sc.score_genes) with min threshold. + +classes: + basal-IFE: [Krt14, Krt5, Krt15, Trp63, Itga6, Itgb1] + spinous: [Krt1, Krt10, Ivl, Dsg1a, Dsg1b] + granular: [Lor, Flg, Flg2, Lce3a, Lce3d, Ivl] + HF-placode: [Edar, Wnt10b, Shh, Lef1, Foxi3, Sox9, Bmp4] + HF-primary-germ: [Lgr5, Lef1, Wnt10b, Krt17, Krt71, Nfatc1] + HF-matrix: [Msx2, Ovol1, Shh, Wnt10a, Krt31, Krt71, Cutl1] + HF-DP: [Sox2, Sox18, Corin, Bmp6, Bmp7, Wif1, Igfbp3] + HF-ORS: [Krt17, Krt5, Krt6a, Krt15, Cd34, Sox9, Lgr5, Nfatc1] + HF-IRS: [Trichohyalin, Tchh, Krt71, Krt73, Krt74, Msx2] + sebaceous: [Scd1, Scd3, Mgst1, Elovl3, Awat2, Adipoq, Fasn] + eccrine-placode: [En1, Foxi3, Foxa1, Krt18, Ncam1, Nkx3-1] + eccrine-duct: [Krt18, Krt19, Krt8, Aqp5, Muc5b] + melanocyte: [Mitf, Dct, Tyr, Pmel, Mlana, Sox10] + Merkel: [Atoh1, Krt20, Krt18, Sox2, Chga, Chgb, Piezo2] + endothelial: [Pecam1, Cdh5, Tie1, Tek, Kdr, Vwf, Cldn5] + fibroblast-papillary: [Col1a1, Pdgfra, Dpp4, Lef1, Prrx1, Fap] + fibroblast-reticular: [Col1a2, Dcn, Lum, Vim, Postn, Fbn1] + immune: [Ptprc, Cd3e, Cd8a, Cd4, Cd19, Ms4a1, Adgre1, Itgam] + +# assignment rules +assignment: + min_score: 0.10 # z-scored marker score below this -> "UNK" + min_margin: 0.05 # top-2 gap must exceed this else "UNK" + cluster_resolution: 0.8 # for datasets w/o labels -> Leiden then vote diff --git a/scripts/pan_skin/ontology_map.yaml b/scripts/pan_skin/ontology_map.yaml new file mode 100644 index 0000000000000000000000000000000000000000..53504c56664e2df1626a2a0b53ac3816ba98b788 --- /dev/null +++ b/scripts/pan_skin/ontology_map.yaml @@ -0,0 +1,186 @@ +# pan-skin 18-class ontology map +# canonical → per-dataset raw label(s). exact label strings from published metadata columns. +# every raw string maps to at most one canonical class. cells whose raw label is not listed +# are dropped from cross-dataset supcon (but kept for unsupervised regularizers). + +canonical_classes: + - basal-IFE # interfollicular basal keratinocyte + - spinous + - granular + - HF-placode # early hair follicle induction (E13.5-E14.5 in mouse) + - HF-primary-germ # hair germ (E15-E16 mouse; distinct from placode) + - HF-matrix + - HF-DP # dermal papilla (mesenchymal — kept for negative pairing) + - HF-ORS # outer root sheath + - HF-IRS # inner root sheath + - sebaceous + - eccrine-placode # sweat gland placode (volar-specific; anchor for Dingwall) + - eccrine-duct + - melanocyte # mature melanocyte (Belote anchor: neo/adt, non-cycling) + - melanoblast # fetal melanocyte precursor (Belote anchor: fet, non-cycling) + - melanocyte-precursor # cycling melanocyte (Belote anchor: cyc_mel / cyc_foll_mel) + - Merkel + - endothelial + - fibroblast-papillary + - fibroblast-reticular + - immune + +datasets: + + sulic_GSE212673: + stage: E14.5 + tissue: dorsal skin + label_column: paper_subtype + map: + Epithelium: basal-IFE + Placode1: HF-placode + Placode2: HF-placode + PlacodeI: HF-placode + PlacodeII: HF-placode + PlacodeIII: HF-placode + PlacodeIV: HF-placode + + aldrich_GSE228872: # verify vs GSE214695 before download + stage: P0 + tissue: volar hindpaw + label_column: cell_type + map: + basal_IFE: basal-IFE + eccrine_placode: eccrine-placode + eccrine_duct: eccrine-duct + hair_placode: HF-placode + Merkel: Merkel + + jacob_kasper_EMTAB11920: + stage: E12.5-E14.5 + tissue: dorsal skin + label_column: annotation + map: + basal keratinocyte: basal-IFE + periderm: basal-IFE # provisional; may split if needed + placode: HF-placode + + ge_gupta_GSE131498: + stage: E13.5-P0 + tissue: dorsal skin + label_column: cluster_label + map: + Basal: basal-IFE + Placode: HF-placode + Dermal_condensate: HF-DP + Matrix: HF-matrix + IFE: basal-IFE + IRS: HF-IRS + ORS: HF-ORS + DP: HF-DP + + haensel_GSE142471: + stage: adult telogen/anagen + tissue: back skin + label_column: cell_type + map: + IFE-B: basal-IFE + IFE-S: spinous + IFE-G: granular + Bulge: HF-ORS + HG: HF-primary-germ + Matrix: HF-matrix + IRS: HF-IRS + ORS: HF-ORS + DP: HF-DP + Sebaceous: sebaceous + Melanocyte: melanocyte + + joost_GSE67602: + stage: adult + tissue: back skin + label_column: cell_type + platform: Smart-seq2 + map: + IFE-B: basal-IFE + IFE-DI: spinous + IFE-DII: spinous + IFE-KI: granular + IFE-KII: granular + uHF-I: HF-ORS + uHF-II: HF-ORS + HFSC: HF-ORS + HG: HF-primary-germ + Bulge: HF-ORS + + mca_GSE108097: + stage: E14.5 + neonatal + adult + tissue: skin + label_column: mca_cluster + platform: Microwell-seq + + wihn_GSE141814: + stage: adult wound + tissue: wound bed + label_column: cell_type + + ge_fuchs_GSE124901: + stage: adult + tissue: back skin + label_column: facs_gate + role: holdout + + placode_subtypes_GSE224714: + stage: E13.5-E15.5 + tissue: dorsal skin + label_column: placode_subtype + role: fine_grained_placode + + merkel_touchdome_GSE201447: + stage: adult + tissue: touch dome + label_column: cell_type + map: + Merkel: Merkel + basal: basal-IFE + + human_fetal_GSE247376: + stage: 7-17 PCW + tissue: fetal skin + organism: human + label_column: annotation + subsample_to: 100000 + + human_reynolds_EMTAB8142: + stage: fetal + adult + tissue: skin + organism: human + label_column: full_clustering + subsample_to: 100000 + + human_ji_cellxgene: + stage: adult + tissue: multi-anatomical + organism: human + label_column: cell_type + subsample_to: 75000 + + human_wang_basal_GSE147482: + stage: adult + tissue: interfollicular epidermis + organism: human + label_column: state + + belote_GSE151091_anchor: + stage: fet + neo + adt # dev_stage preserved in obs['stage_tag'] + tissue: skin (FACS-enriched melanocytes; cutaneous / acral / foreskin / follicular) + organism: human # HGNC symbols converted to mouse via str.capitalize() + # (+ minor TP53/TP63/TP73 → Trp53/Trp63/Trp73 fixups). + # Mouse-HVG coverage vs corpus_v3 shared HVGs: 2258/2532 (89.2%) + label_column: canonical # from canonical_map.py (class_1 + class_2 rule) + role: anchor # 1,000 stratified cells fold into corpus_v3; + # 6,088 held out at data/corpus/pan_skin/held_out_labeled/ + # belote_GSE151091_test.h5ad for zero-shot melanocyte eval. + anchor_composition: + melanocyte: 700 # class_1=mel, dev_stage in {neo, adt} + melanoblast: 250 # class_1=mel, dev_stage=fet + melanocyte-precursor: 50 # class_1=cyc_mel or class_2=cyc_foll_mel + map: + melanocyte: melanocyte + melanoblast: melanoblast + melanocyte-precursor: melanocyte-precursor diff --git a/scripts/pan_skin/run_all.sh b/scripts/pan_skin/run_all.sh new file mode 100644 index 0000000000000000000000000000000000000000..a94fc7d4231a3da4c03ee4d819405725dda20873 --- /dev/null +++ b/scripts/pan_skin/run_all.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +# canonical end-to-end pipeline +set -euo pipefail +cd /home/bcheng/PRISM +mkdir -p logs/pan_skin +LDLIBS=$(python -c "import site,os; print(os.path.join(site.getsitepackages()[0],'nvidia','cusparselt','lib'))" 2>/dev/null || echo "") +export LD_LIBRARY_PATH="$LDLIBS:${LD_LIBRARY_PATH:-}" +run() { echo "==== $1 ===="; python -u "$2" 2>&1 | tail -10; } + +bash scripts/pan_skin/01_download_tier_a.sh > logs/pan_skin/01_download.log 2>&1 +bash scripts/pan_skin/02_download_tier_b.sh > logs/pan_skin/02_download.log 2>&1 +run 06_per_dataset scripts/pan_skin/06_build_per_dataset_h5ads.py +run 07_hvg_pca scripts/pan_skin/07_build_shared_hvgs_and_pca.py +run 08a_marker_labels scripts/pan_skin/08_assign_labels.py +run 08b_curated_override scripts/pan_skin/08b_curated_label_override.py +run 20_train scripts/pan_skin/20_train_panda.py +run 30_zero_shot scripts/pan_skin/30_zero_shot_aldrich.py +run 40_heldout_cv scripts/pan_skin/40_heldout_5fold_cv.py +run 44_en1_cko scripts/analysis/44_en1_cko_contrast.py +run 45_markers scripts/analysis/45_marker_refinement.py +run 49_melanocyte scripts/analysis/49_melanocyte_deep_dive.py +run 57_pathways scripts/analysis/57_multiclass_pathway_analysis.py +echo "done" diff --git a/scripts/pancreas/01_download.sh b/scripts/pancreas/01_download.sh new file mode 100644 index 0000000000000000000000000000000000000000..378bc6c18f790d6ba964df610a635bb895502b5e --- /dev/null +++ b/scripts/pancreas/01_download.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash +# pancreas corpus + held-out downloads +set -uo pipefail +BASE=/home/bcheng/PRISM/data/corpus/pancreas/tier_a +LOG=/home/bcheng/PRISM/logs/pancreas +mkdir -p "$BASE" "$LOG" + +download() { + local acc=$1 url=$2 out=$3 + local outfile="$BASE/$out" + if [ -s "$outfile" ]; then echo "[$acc] cached"; return 0; fi + echo "[$acc] downloading" + wget --quiet -c -O "$outfile.part" "$url" && mv "$outfile.part" "$outfile" \ + && echo "[$acc] done: $(du -h "$outfile" | cut -f1)" \ + || { echo "[$acc] FAILED"; rm -f "$outfile.part"; return 1; } +} + +# GSE84133 Baron 2016 - mouse+human islet inDrops +download GSE84133 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE84nnn/GSE84133/suppl/GSE84133_RAW.tar" \ + baron_GSE84133_RAW.tar + +# GSE85241 Muraro 2016 - human CEL-Seq2 +download GSE85241 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE85nnn/GSE85241/suppl/GSE85241_cellsystems_dataset_4donors_updated.csv.gz" \ + muraro_GSE85241_cellsystems.csv.gz + +# GSE81076 Grun 2016 - human CEL-Seq +download GSE81076 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE81nnn/GSE81076/suppl/GSE81076_D2_3_7_10_17.txt.gz" \ + grun_GSE81076.txt.gz + +# GSE101099 Byrnes 2018 - inDrops mouse +download GSE101099 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE101nnn/GSE101099/suppl/GSE101099_RAW.tar" \ + byrnes_GSE101099_RAW.tar + +# GSE114412 Sharon 2019 - hPSC diff, held-out unlabeled +mkdir -p "$BASE/../held_out_unlabeled" +if [ ! -s "$BASE/../held_out_unlabeled/sharon_GSE114412_RAW.tar" ]; then + wget --quiet -c -O "$BASE/../held_out_unlabeled/sharon_GSE114412_RAW.tar.part" \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE114nnn/GSE114412/suppl/GSE114412_RAW.tar" \ + && mv "$BASE/../held_out_unlabeled/sharon_GSE114412_RAW.tar.part" \ + "$BASE/../held_out_unlabeled/sharon_GSE114412_RAW.tar" +fi + +echo +echo "download attempt complete:" +ls -lh "$BASE" "$BASE/../held_out_unlabeled" 2>&1 diff --git a/scripts/pancreas/02_build_per_dataset.py b/scripts/pancreas/02_build_per_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..b0cc3e2110e5f9d59e0377aeff12f57069cfa75b --- /dev/null +++ b/scripts/pancreas/02_build_per_dataset.py @@ -0,0 +1,58 @@ +"""per-dataset h5ads for the pancreas corpus.""" +from __future__ import annotations +from pathlib import Path +import warnings, sys +warnings.filterwarnings("ignore") +import scanpy as sc + +sys.path.insert(0, "/home/bcheng/PRISM") +import importlib.util +spec = importlib.util.spec_from_file_location( + "loaders", "/home/bcheng/PRISM/panda/data/pancreas_loaders.py") +_mod = importlib.util.module_from_spec(spec); spec.loader.exec_module(_mod) +ALL_LOADERS = _mod.ALL_LOADERS + +OUT = Path("/home/bcheng/PRISM/data/corpus/pancreas/harmonized") +OUT.mkdir(parents=True, exist_ok=True) + +# Baron is held-out labeled, Sharon held-out unlabeled — excluded here +TRAINING_CORPUS = ["bastidas_GSE132188", "muraro_GSE85241", + "grun_GSE81076", "byrnes_GSE101099"] + + +def qc(a, name): + n0 = a.n_obs + sc.pp.filter_cells(a, min_genes=200) + sc.pp.filter_genes(a, min_cells=3) + a.var["mt"] = a.var_names.str.startswith(("mt-", "MT-")) + if a.var["mt"].any(): + sc.pp.calculate_qc_metrics(a, qc_vars=["mt"], percent_top=None, + log1p=False, inplace=True) + a = a[a.obs["pct_counts_mt"] < 25].copy() + print(f" [{name}] {n0} -> {a.n_obs} cells, {a.n_vars} genes", flush=True) + return a + + +def main(): + for name in TRAINING_CORPUS: + out_path = OUT / f"{name}.h5ad" + if out_path.exists(): + print(f"[{name}] cached", flush=True); continue + print(f"[{name}] loading …", flush=True) + try: + a = ALL_LOADERS[name]() + except Exception as exc: + print(f"[{name}] LOADER FAILED: {exc!r}", flush=True); continue + print(f" raw: {a.shape}", flush=True) + a = qc(a, name) + a.obs["dataset"] = name + import scipy.sparse as sp + if not sp.issparse(a.X): a.X = sp.csr_matrix(a.X) + a.X = a.X.astype("float32") + a.var_names_make_unique() + a.write_h5ad(out_path, compression="gzip") + print(f" saved to {out_path} ({out_path.stat().st_size/1e6:.1f} MB)", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/pancreas/03_shared_hvgs_and_pca.py b/scripts/pancreas/03_shared_hvgs_and_pca.py new file mode 100644 index 0000000000000000000000000000000000000000..1fef54b4be507e2ffde1faa57fdd0627be093c16 --- /dev/null +++ b/scripts/pancreas/03_shared_hvgs_and_pca.py @@ -0,0 +1,154 @@ +"""shared HVGs + sample-fit PCA for the pancreas corpus.""" +from __future__ import annotations +from pathlib import Path +import warnings, pickle +warnings.filterwarnings("ignore") + +import numpy as np +import pandas as pd +import scanpy as sc +import scipy.sparse as sp +from sklearn.decomposition import PCA +import anndata as ad + +HARM = Path("/home/bcheng/PRISM/data/corpus/pancreas/harmonized") +K_TOP_PER_DS = 4000 +K_TARGET = 3500 +N_PCA = 50 +SAMPLE_PER_DS = 5000 + +# force-include canonical markers (see known_pancreas_markers.yaml) +MUST_INCLUDE = [ + "Ins1", "Ins2", "Gcg", "Sst", "Ppy", "Ghrl", # hormones + "Mafa", "Mafb", "Arx", "Pdx1", "Nkx6-1", "Nkx2-2", "Neurog3", "Neurod1", + "Pax4", "Pax6", "Foxa2", "Hnf1a", "Hnf1b", "Fev", "Isl1", # TFs + "Krt19", "Sox9", "Muc1", "Cftr", "Spp1", # ductal + "Cpa1", "Ptf1a", "Amy1", "Amy2a", "Prss1", "Cel", # acinar + "Pecam1", "Cdh5", "Tie1", "Tek", "Kdr", # endothelial + "Ptprc", "Cd3e", "Cd8a", "Cd4", "Adgre1", "Itgam", # immune + "Chga", "Chgb", "Slc30a8", "Ucn3", "Hhex", "Rbp4", +] + + +def rank_hvgs(a): + x = a.copy() + try: + sc.pp.highly_variable_genes(x, n_top_genes=K_TOP_PER_DS, flavor="seurat", + subset=False, check_values=False) + return x.var["variances_norm"].fillna(-np.inf) if "variances_norm" in x.var else \ + x.var["dispersions_norm"].fillna(-np.inf) + except Exception: + sc.pp.normalize_total(x, target_sum=1e4); sc.pp.log1p(x) + X = x.X.toarray() if sp.issparse(x.X) else x.X + return pd.Series(np.asarray(X.var(axis=0)).ravel(), index=x.var_names) + + +def load_and_norm(f, shared_genes, gene_idx, mu, sig, G): + a = ad.read_h5ad(f) + raw_counts = np.asarray(a.X.sum(axis=1)).ravel() if sp.issparse(a.X) else a.X.sum(axis=1) + common = [g for g in shared_genes if g in a.var_names] + a_s = a[:, common].copy() + sc.pp.normalize_total(a_s, target_sum=1e4) + sc.pp.log1p(a_s) + X = a_s.X.toarray() if sp.issparse(a_s.X) else a_s.X + Xf = np.zeros((X.shape[0], G), dtype=np.float32) + cols = [gene_idx[g] for g in common] + Xf[:, cols] = X.astype(np.float32) + Xz = np.clip((Xf - mu.astype(np.float32)) / sig.astype(np.float32), -10, 10) + return a, Xf, Xz, raw_counts, common + + +def main(): + files = sorted(HARM.glob("*.h5ad")) + files = [f for f in files if not f.name.startswith("corpus")] + print(f"[hvg] found {len(files)} datasets", flush=True) + + ranks = {} + for f in files: + print(f"[hvg] rank {f.name}", flush=True) + a = ad.read_h5ad(f) + ranks[f.stem] = rank_hvgs(a) + del a + + R = pd.DataFrame(ranks) + top_hits = pd.DataFrame({d: R[d].rank(ascending=False, method="min") <= K_TOP_PER_DS + for d in R.columns}).fillna(False) + hit_count = top_hits.sum(axis=1) + eligible = hit_count[hit_count >= max(2, len(files) // 2)] + print(f"[hvg] {len(eligible)} genes top-{K_TOP_PER_DS} in >=half datasets", flush=True) + + mean_score = R.mean(axis=1, skipna=True) + shared_by_score = mean_score.loc[eligible.index].sort_values(ascending=False).index + + all_genes = set(R.index) + forced = [g for g in MUST_INCLUDE if g in all_genes] + print(f"[hvg] forced markers present: {len(forced)}/{len(MUST_INCLUDE)}", flush=True) + + picked = list(forced) + for g in shared_by_score: + if g not in picked: + picked.append(g) + if len(picked) >= K_TARGET: + break + shared_genes = picked[:K_TARGET] + print(f"[hvg] final HVG count: {len(shared_genes)}", flush=True) + (HARM / "shared_hvgs.txt").write_text("\n".join(shared_genes) + "\n") + + G = len(shared_genes) + gene_idx = {g: i for i, g in enumerate(shared_genes)} + running_sum = np.zeros(G, dtype=np.float64) + running_sq = np.zeros(G, dtype=np.float64) + N = 0 + for f in files: + a = ad.read_h5ad(f) + common = [g for g in shared_genes if g in a.var_names] + a_s = a[:, common].copy() + sc.pp.normalize_total(a_s, target_sum=1e4) + sc.pp.log1p(a_s) + X = a_s.X.toarray() if sp.issparse(a_s.X) else a_s.X + cols = [gene_idx[g] for g in common] + running_sum[cols] += X.sum(axis=0) + running_sq[cols] += (X ** 2).sum(axis=0) + N += X.shape[0] + del a, a_s, X + mu = running_sum / N + sig = np.sqrt(np.maximum(running_sq / N - mu ** 2, 1e-6)) + np.savez(HARM / "corpus_stats.npz", shared_hvgs=np.array(shared_genes), + mean=mu.astype(np.float32), std=sig.astype(np.float32), n_cells=N) + print(f"[stats] mean/std over {N} cells", flush=True) + + fit_chunks = [] + rng = np.random.default_rng(0) + for f in files: + _, _, Xz, _, _ = load_and_norm(f, shared_genes, gene_idx, mu, sig, G) + n = Xz.shape[0] + take = min(SAMPLE_PER_DS, n) + idx = rng.choice(n, size=take, replace=False) + fit_chunks.append(Xz[idx]) + Xfit = np.vstack(fit_chunks) + pca = PCA(n_components=N_PCA, random_state=42).fit(Xfit) + print(f"[pca] explained var={pca.explained_variance_ratio_.sum():.3f}", flush=True) + for k in range(N_PCA): + top = int(np.argmax(np.abs(pca.components_[k]))) + if pca.components_[k, top] < 0: pca.components_[k] *= -1 + with open(HARM / "pca_basis.pkl", "wb") as fh: pickle.dump(pca, fh) + + parts = [] + for f in files: + a, Xf, Xz, raw_counts, common = load_and_norm(f, shared_genes, gene_idx, mu, sig, G) + Z = pca.transform(Xz).astype(np.float32) + obs = a.obs.copy() + obs["total_counts"] = raw_counts.astype(np.float32) + obs["missing_hvg_frac"] = 1.0 - len(common) / G + a2 = ad.AnnData(X=sp.csr_matrix(Xf.astype(np.float32)), obs=obs, + var=pd.DataFrame(index=shared_genes)) + a2.obsm["X_pca"] = Z + parts.append(a2) + corpus = ad.concat(parts, join="outer", label="_batch") + corpus.uns["shared_hvgs"] = shared_genes + corpus.write_h5ad(HARM / "corpus.h5ad", compression="gzip") + print(f"[emit] wrote corpus.h5ad ({corpus.n_obs} x {corpus.n_vars})", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/pancreas/04_assign_labels.py b/scripts/pancreas/04_assign_labels.py new file mode 100644 index 0000000000000000000000000000000000000000..e69e3df5785d8d33b3ff14016c2f7e273ae13f51 --- /dev/null +++ b/scripts/pancreas/04_assign_labels.py @@ -0,0 +1,74 @@ +"""assign canonical pancreas labels via marker scoring.""" +from __future__ import annotations +from pathlib import Path +import warnings, yaml +warnings.filterwarnings("ignore") +import numpy as np, pandas as pd, scanpy as sc, anndata as ad + +CORPUS = Path("/home/bcheng/PRISM/data/corpus/pancreas/harmonized/corpus.h5ad") +YAML = Path("/home/bcheng/PRISM/scripts/pancreas/known_pancreas_markers.yaml") + +CURATED_MAP = { + # Bastidas cell_type + "alpha": "alpha", "beta": "beta", "delta": "delta", + "progenitor": "endocrine-progenitor", "other": "other", + # Baron/Muraro/Grun + "gamma": "gamma", "epsilon": "epsilon", "ductal": "ductal", "acinar": "acinar", + "endothelial": "endothelial", "immune": "immune", "activated_stellate": "other", + "quiescent_stellate": "other", "schwann": "other", "mast": "immune", "PP": "gamma", + "T_cell": "immune", "macrophage": "immune", +} + + +def main(): + with open(YAML) as f: cfg = yaml.safe_load(f) + classes = list(cfg["classes"].keys()) + markers = cfg["classes"] + min_score = cfg["assignment"]["min_score"] + min_margin = cfg["assignment"]["min_margin"] + res = cfg["assignment"]["cluster_resolution"] + + a = ad.read_h5ad(CORPUS) + print(f"[label] corpus: {a.shape}", flush=True) + + a.obs["canonical_label"] = pd.Categorical(["UNK"] * a.n_obs, categories=classes + ["UNK"]) + a.obs["leiden"] = "0" + + for ds_name, sub in a.obs.groupby("dataset"): + idx = sub.index + sa = a[idx].copy() + sc.pp.neighbors(sa, use_rep="X_pca", n_neighbors=15) + sc.tl.leiden(sa, resolution=res, key_added="leiden") + a.obs.loc[idx, "leiden"] = ds_name + "_" + sa.obs["leiden"].astype(str) + for c, m in markers.items(): + present = [g for g in m if g in sa.var_names] + if not present: sa.obs[f"score_{c}"] = -np.inf; continue + sc.tl.score_genes(sa, gene_list=present, score_name=f"score_{c}", + random_state=0, use_raw=False) + S = sa.obs[[f"score_{c}" for c in classes]].values + cluster_labels = sa.obs["leiden"].values + assigned = np.array(["UNK"] * sa.n_obs, dtype=object) + for cl in np.unique(cluster_labels): + mask = cluster_labels == cl + mean_scores = S[mask].mean(axis=0) + order = np.argsort(mean_scores)[::-1] + top, second = mean_scores[order[0]], mean_scores[order[1]] + if top >= min_score and (top - second) >= min_margin: + assigned[mask] = classes[order[0]] + a.obs.loc[idx, "canonical_label"] = assigned + print(f"[label] {ds_name}: {pd.Series(assigned).value_counts().to_dict()}", flush=True) + + # prefer paper cell_type when present + if "cell_type" in a.obs.columns: + mapped = a.obs["cell_type"].astype(str).map(CURATED_MAP) + ok = mapped.notna() + a.obs.loc[ok, "canonical_label"] = mapped[ok].astype(a.obs["canonical_label"].dtype) + print(f"[label] curated override: {int(ok.sum())} cells", flush=True) + + print("[label] final corpus breakdown:") + print(a.obs.groupby(["dataset", "canonical_label"], observed=True).size().unstack(fill_value=0)) + a.write_h5ad(CORPUS, compression="gzip") + + +if __name__ == "__main__": + main() diff --git a/scripts/pancreas/05_train_panda.py b/scripts/pancreas/05_train_panda.py new file mode 100644 index 0000000000000000000000000000000000000000..8fd0e05f6c1a2af654347ff4891a074468e3503c --- /dev/null +++ b/scripts/pancreas/05_train_panda.py @@ -0,0 +1,133 @@ +"""train PANDA on the pancreas 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 + +sys.path.insert(0, "/home/bcheng/PRISM") +from panda.pan_skin.model import ( + PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, prototype_infonce +) + +CORPUS = Path("/home/bcheng/PRISM/data/corpus/pancreas/harmonized/corpus.h5ad") +OUT = Path("/home/bcheng/PRISM/checkpoints/pancreas") +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() diff --git a/scripts/pancreas/06_zero_shot_baron.py b/scripts/pancreas/06_zero_shot_baron.py new file mode 100644 index 0000000000000000000000000000000000000000..8c182806a440da6302bda37eeb98c68fd2901671 --- /dev/null +++ b/scripts/pancreas/06_zero_shot_baron.py @@ -0,0 +1,109 @@ +"""zero-shot PANDA on held-out Baron 2016 mouse pancreas.""" +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 pathlib import Path as _P_root +ROOT = _P_root(__file__).resolve().parents[2] +ROOT_STR = str(ROOT) +sys.path.insert(0, ROOT_STR) +from panda import PANDAEncoder +from panda.data.pancreas_loaders import load_baron + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +CKPT = Path(f"{ROOT_STR}/checkpoints/pancreas") +HARM = Path(f"{ROOT_STR}/data/corpus/pancreas/harmonized") +OUT = Path(f"{ROOT_STR}/discovery/pancreas/marker") +OUT.mkdir(parents=True, exist_ok=True) + + +def project(a, shared_hvgs, mu, sig): + G = len(shared_hvgs); n = a.n_obs + hvg2i = {g: i for i, g in enumerate(shared_hvgs)} + common = [g for g in a.var_names.astype(str) if g in hvg2i] + present_frac = len(common) / G + print(f"[proj] {len(common)}/{G} HVGs present ({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((n, 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) + return Xz, present_frac + + +def main(): + 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"]) + print(f"[model] classes: {classes}", flush=True) + + 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) + + a = load_baron() + print(f"[target] Baron shape: {a.shape}", flush=True) + print(f"[target] labels (assigned_cluster head): " + f"{a.obs.get('assigned_cluster', pd.Series(dtype=str)).value_counts().head(12).to_dict()}", + flush=True) + + Xz, present_frac = project(a, shared_hvgs, mu, sig) + counts = np.asarray(a.X.sum(axis=1)).ravel() + logc = np.log10(counts + 1); logc = (logc - logc.mean()) / (logc.std() + 1e-6) + Xpca = pca.transform(Xz).astype(np.float32) + mhf = np.full(a.n_obs, 1.0 - present_frac, dtype=np.float32) + + all_z, batch = [], 4096 + with torch.no_grad(): + for i in range(0, a.n_obs, batch): + xb = torch.from_numpy(Xpca[i:i+batch]).to(DEVICE) + aux = torch.from_numpy(np.stack([mhf[i:i+batch], logc[i:i+batch]], 1).astype(np.float32)).to(DEVICE) + out = model(xb, aux, lam_dann=0.0) + all_z.append(out["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_ix = cos.argmax(axis=1) + pred_label = np.array([classes[i] for i in pred_ix], dtype=object) + conf = cos.max(axis=1) + + a.obs["pred_label"] = pred_label + a.obs["pred_conf"] = conf.astype(np.float32) + a.obs.to_csv(OUT / "baron_predictions.csv") + print(f"\n[summary] Baron pred_label breakdown:") + print(pd.Series(pred_label).value_counts()) + + if "assigned_cluster" in a.obs.columns: + norm_map = { + "alpha": "alpha", "beta": "beta", "delta": "delta", "gamma": "gamma", + "epsilon": "epsilon", "ductal": "ductal", "acinar": "acinar", + "endothelial": "endothelial", "activated_stellate": "other", + "quiescent_stellate": "other", "schwann": "other", + "mast": "immune", "macrophage": "immune", "t_cell": "immune", "T_cell": "immune", + } + true_lbl = a.obs["assigned_cluster"].astype(str).str.lower().map(norm_map).fillna("other") + mask = true_lbl.isin(classes) + if mask.sum() > 0: + from sklearn.metrics import classification_report, accuracy_score + true = true_lbl[mask].values + pred = pred_label[mask.values] + print(f"\n[eval] shared-class accuracy: {accuracy_score(true, pred):.4f} on {mask.sum()} cells") + print(classification_report(true, pred, digits=3, zero_division=0)) + with open(OUT / "baron_accuracy.json", "w") as f: + json.dump({"acc": accuracy_score(true, pred), + "n_evaluated": int(mask.sum())}, f, indent=2) + else: + print("[eval] no shared classes; skipping") + + +if __name__ == "__main__": + main() diff --git a/scripts/pancreas/07_heldout_5fold_cv.py b/scripts/pancreas/07_heldout_5fold_cv.py new file mode 100644 index 0000000000000000000000000000000000000000..90e01c981cb1530adcc91d57e6c11a0179de04a8 --- /dev/null +++ b/scripts/pancreas/07_heldout_5fold_cv.py @@ -0,0 +1,207 @@ +"""5-fold CV: retrain PANDA per fold, eval by prototype-cosine on held-out 20%.""" +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 torch +import torch.nn.functional as F +from torch.utils.data import Dataset, DataLoader +from sklearn.model_selection import StratifiedKFold +from sklearn.metrics import (accuracy_score, f1_score, roc_auc_score, + classification_report) + +sys.path.insert(0, "/home/bcheng/PRISM") +from panda.pan_skin.model import ( + PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, prototype_infonce +) + +CORPUS = Path("/home/bcheng/PRISM/data/corpus/pancreas/harmonized/corpus.h5ad") +OUT = Path("/home/bcheng/PRISM/discovery/pancreas/marker") + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +N_FOLDS = 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} + 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_pick = self.rng.choice(self.classes, p=self.p) + batch.append(int(self.rng.choice(self.by_cls[c_pick]))) + yield batch + def __len__(self): return self.n_batches + + +def train_one_fold(X, y, d, mhf, logc, classes, datasets, tr, te, fold_id, log_prefix): + torch.cuda.empty_cache() + Xtr, ytr, dtr, mtr, ltr = X[tr], y[tr], d[tr], mhf[tr], logc[tr] + ds = CorpusDataset(Xtr, ytr, dtr, mtr, ltr) + sampler = HybridSampler(ytr, dtr, n_batches=100, seed=fold_id) + loader = DataLoader(ds, batch_sampler=sampler, num_workers=0) + + counts_per = np.bincount(ytr, minlength=len(classes)) + inv_sqrt = 1.0 / np.sqrt(counts_per + 1) + inv_sqrt = inv_sqrt / inv_sqrt.mean() + class_w_np = 0.5 * inv_sqrt + 0.5 * np.ones_like(inv_sqrt) + class_w = torch.tensor(class_w_np, dtype=torch.float32).to(DEVICE) + + 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): + for e in range(stage_epochs[stage]): + if e % 10 == 0: + print(f"{log_prefix} s{stage} ep {e}/{stage_epochs[stage]}", flush=True) + 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 + 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) + + model.eval() + Xte, yte = X[te], y[te] + with torch.no_grad(): + Xt = torch.from_numpy(Xte.astype(np.float32)).to(DEVICE) + aux = torch.zeros(len(te), 2, device=DEVICE) + out = model(Xt, aux, lam_dann=0.0) + z = out["z"] + cos = z @ model.prototypes.T + pred = cos.argmax(dim=1).cpu().numpy() + probs = torch.softmax(cos / 0.07, dim=1).cpu().numpy() + + acc = accuracy_score(yte, pred) + f1 = f1_score(yte, pred, average="macro", zero_division=0) + try: + auc = roc_auc_score(np.eye(len(classes))[yte], probs, average="macro", multi_class="ovr") + except Exception: + auc = float("nan") + print(f"{log_prefix} acc={acc:.4f} macro_f1={f1:.4f} macro_auc={auc:.4f}", flush=True) + return acc, f1, auc, pred, yte + + +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)} + X = np.asarray(a.obsm["X_pca"]) + 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)]) + 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) + + print(f"[cv] corpus: {a.shape} n_classes={len(classes)} n_datasets={len(datasets)}", + flush=True) + print(f"[cv] class counts: {dict(zip(classes, np.bincount(y, minlength=len(classes)).tolist()))}", + flush=True) + + skf = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=42) + accs, f1s, aucs = [], [], [] + all_preds = [] + partial_path = OUT / "61_pancreas_heldout_5fold_partial.json" + for fold, (tr, te) in enumerate(skf.split(X, y)): + t0 = time.time() + try: + acc, f1, auc, pred, yte = train_one_fold( + X, y, d, mhf, logc, classes, datasets, tr, te, + fold_id=fold, log_prefix=f"[fold {fold}]") + except Exception as exc: + import traceback; traceback.print_exc() + print(f"[fold {fold}] FAILED: {exc}", flush=True) + continue + print(f"[fold {fold}] wall={time.time()-t0:.0f}s", flush=True) + accs.append(acc); f1s.append(f1); aucs.append(auc) + all_preds.append({"fold": fold, "test_idx": te.tolist(), + "pred": pred.tolist(), "true": yte.tolist()}) + with open(partial_path, "w") as f: + json.dump({"folds_done": fold + 1, "accs": accs, "f1s": f1s, "aucs": aucs}, f) + + print(f"\n[cv] 5-FOLD ACC: {np.mean(accs):.4f} +- {np.std(accs):.4f}") + print(f"[cv] 5-FOLD F1: {np.mean(f1s):.4f} +- {np.std(f1s):.4f}") + print(f"[cv] 5-FOLD AUC: {np.mean(aucs):.4f} +- {np.std(aucs):.4f}") + + all_y_true = np.concatenate([np.array(p["true"]) for p in all_preds]) + all_y_pred = np.concatenate([np.array(p["pred"]) for p in all_preds]) + print("\n[cv] Concatenated held-out classification report:") + rep = classification_report(all_y_true, all_y_pred, target_names=classes, + digits=3, zero_division=0, output_dict=True) + print(classification_report(all_y_true, all_y_pred, target_names=classes, + digits=3, zero_division=0)) + + result = { + "mean_acc": float(np.mean(accs)), "std_acc": float(np.std(accs)), + "mean_f1": float(np.mean(f1s)), "std_f1": float(np.std(f1s)), + "mean_auc": float(np.mean(aucs)), "std_auc": float(np.std(aucs)), + "per_fold_acc": accs, "per_fold_f1": f1s, "per_fold_auc": aucs, + "per_class_report": rep, + "n_folds": N_FOLDS, "n_classes": len(classes), + "protocol": "StratifiedKFold(5) retrain from scratch per fold; " + "eval by prototype-cosine argmax on held-out 20%.", + } + (OUT / "61_pancreas_heldout_5fold_cv.json").write_text(json.dumps(result, indent=2)) + print(f"\n[cv] wrote {OUT}/61_pancreas_heldout_5fold_cv.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/pancreas/08_add_baron_split.py b/scripts/pancreas/08_add_baron_split.py new file mode 100644 index 0000000000000000000000000000000000000000..1e8ad8b04b90c0a5b595fc65d0e40b5d165df3c5 --- /dev/null +++ b/scripts/pancreas/08_add_baron_split.py @@ -0,0 +1,81 @@ +"""baron mouse: stratified 50/50 split — train goes to corpus, test held out.""" +from __future__ import annotations +from pathlib import Path +import warnings, sys +warnings.filterwarnings("ignore") +import numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp +from sklearn.model_selection import train_test_split + +from pathlib import Path as _P_root +ROOT = _P_root(__file__).resolve().parents[2] +ROOT_STR = str(ROOT) +sys.path.insert(0, ROOT_STR) +import importlib.util +spec = importlib.util.spec_from_file_location( + "loaders", f"{ROOT_STR}/panda/data/pancreas_loaders.py") +_mod = importlib.util.module_from_spec(spec); spec.loader.exec_module(_mod) + +HARM = Path(f"{ROOT_STR}/data/corpus/pancreas/harmonized") +HELDOUT = Path(f"{ROOT_STR}/data/corpus/pancreas/held_out_labeled") +HELDOUT.mkdir(parents=True, exist_ok=True) + + +def qc(a, name): + n0 = a.n_obs + sc.pp.filter_cells(a, min_genes=200) + sc.pp.filter_genes(a, min_cells=3) + print(f" [{name}] {n0} -> {a.n_obs} cells, {a.n_vars} genes", flush=True) + return a + + +def main(): + print("[baron] loading Baron 2016 GSE84133 …", flush=True) + a = _mod.load_baron() + print(f" raw: {a.shape}", flush=True) + # mouse subset only (loader sets organism) + if "organism" in a.obs.columns: + m_mask = a.obs["organism"] == "mouse" + else: + m_mask = a.obs.index.str.contains("mouse") + a_m = a[m_mask].copy() + print(f" mouse subset: {a_m.shape}", flush=True) + a_m = qc(a_m, "baron_mouse") + a_m.obs["dataset"] = "baron_GSE84133_mouse" + if "assigned_cluster" in a_m.obs.columns: + norm_map = { + "alpha": "alpha", "beta": "beta", "delta": "delta", "gamma": "gamma", + "epsilon": "epsilon", "ductal": "ductal", "acinar": "acinar", + "endothelial": "endothelial", "activated_stellate": "other", + "quiescent_stellate": "other", "schwann": "other", + "mast": "immune", "macrophage": "immune", "t_cell": "immune", + "T_cell": "immune", "b_cell": "immune", + } + a_m.obs["canonical_label"] = a_m.obs["assigned_cluster"].astype(str).str.lower().map(norm_map).fillna("other") + else: + a_m.obs["canonical_label"] = "unknown" + + print(f"[baron] canonical_label counts: " + f"{a_m.obs['canonical_label'].value_counts().to_dict()}", flush=True) + + y = a_m.obs["canonical_label"].astype(str).values + idx = np.arange(a_m.n_obs) + tr_idx, te_idx = train_test_split(idx, test_size=0.5, stratify=y, random_state=42) + a_train = a_m[tr_idx].copy() + a_test = a_m[te_idx].copy() + print(f"[baron] train: {a_train.shape}, test: {a_test.shape}", flush=True) + + for x in (a_train, a_test): + if not sp.issparse(x.X): x.X = sp.csr_matrix(x.X) + x.X = x.X.astype("float32") + x.var_names_make_unique() + + a_train.write_h5ad(HARM / "baron_GSE84133_mouse_train.h5ad", compression="gzip") + a_test.write_h5ad(HELDOUT / "baron_GSE84133_mouse_test.h5ad", compression="gzip") + print(f"[baron] wrote train -> {HARM/'baron_GSE84133_mouse_train.h5ad'} " + f"({(HARM/'baron_GSE84133_mouse_train.h5ad').stat().st_size/1e6:.1f} MB)", flush=True) + print(f"[baron] wrote test -> {HELDOUT/'baron_GSE84133_mouse_test.h5ad'} " + f"({(HELDOUT/'baron_GSE84133_mouse_test.h5ad').stat().st_size/1e6:.1f} MB)", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/pancreas/09_download.sh b/scripts/pancreas/09_download.sh new file mode 100644 index 0000000000000000000000000000000000000000..731d3dfd49011025d992570cba1da350d32407a0 --- /dev/null +++ b/scripts/pancreas/09_download.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +# pancreas v2: MIA adult islets + Yu sequential EP states +set -euo pipefail +BASE="${1:-data/corpus/pancreas/tier_v2}" +mkdir -p "$BASE" + +download() { + local id="$1"; local url="$2"; local out="$3" + if [ -s "$BASE/$out" ]; then + echo "[$id] cached" + return 0 + fi + echo "[$id] downloading $url -> $out" + curl -fL --retry 3 -o "$BASE/$out" "$url" \ + || { echo "[$id] FAILED"; return 1; } +} + +# GSE139627 - Yu 2021, Ngn3-lineage EP1-4 states +download GSE139627 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE139nnn/GSE139627/suppl/GSE139627_RAW.tar" \ + yu_GSE139627_seqEP_RAW.tar + +# GSE211796 - Hrovatin 2023 Mouse Islet Atlas, 4-month adult (~70k cells) +download GSE211796 \ + "https://ftp.ncbi.nlm.nih.gov/geo/series/GSE211nnn/GSE211796/suppl/GSE211796_RAW.tar" \ + mia_GSE211796_4mo_RAW.tar + +echo "[done] pancreas v2 downloads complete" diff --git a/scripts/pancreas/11_build_corpus.py b/scripts/pancreas/11_build_corpus.py new file mode 100644 index 0000000000000000000000000000000000000000..f94ce76480d1363f1f141eba522d341c6aebe80b --- /dev/null +++ b/scripts/pancreas/11_build_corpus.py @@ -0,0 +1,600 @@ +"""corpus v3: paper-label-first + Veres SC-islet. writes _v3 outputs only.""" +from __future__ import annotations +from pathlib import Path +import warnings, sys, pickle, json, gzip, glob, hashlib +import numpy as np, pandas as pd, anndata as ad, scanpy as sc +import scipy.sparse as sp +warnings.filterwarnings("ignore") + +ROOT = Path("/home/bcheng/PRISM") +V1_CORPUS = ROOT / "data/corpus/pancreas/harmonized/corpus.h5ad" +OUT_DIR = ROOT / "data/corpus/pancreas/harmonized" +HELD_OUT_L = ROOT / "data/corpus/pancreas/held_out_labeled" +LBL_DIR = ROOT / "data/external_labels" +SHARON_DIR = ROOT / "data/corpus/pancreas/held_out_unlabeled/sharon_extract" + +OUT_DIR.mkdir(parents=True, exist_ok=True) +HELD_OUT_L.mkdir(parents=True, exist_ok=True) + +MIN_CLASS_SIZE = 150 +N_HVG = 3500 +N_PCA = 64 +RNG = np.random.default_rng(20260723) + + +# canonical vocabulary +BASTIDAS_MAP = { # cell_type col + "alpha": "alpha", "beta": "beta", "delta": "delta", + "progenitor": "endocrine-progenitor", + "other": None, +} +BARON_MAP = { # assigned_cluster col + "alpha": "alpha", "beta": "beta", "delta": "delta", "gamma": "gamma", + "ductal": "ductal", "endothelial": "endothelial", + "quiescent_stellate": "mesenchyme", "activated_stellate": "mesenchyme", + "macrophage": "immune", "immune_other": "immune", + "B_cell": "immune", "T_cell": "immune", + "schwann": None, # <150 total, drop +} +# byrnes Cluster_Names + "__exocrine__" tag for exo files +BYRNES_MAP = { + "Gcg": "alpha", "Ins1": "beta", + "Ghrl": "epsilon", "Ghrl/Gcg": "epsilon", + "Ngn3": "endocrine-progenitor", "Ngn3/Pax4": "endocrine-progenitor", + "Sox4/Spp1": "endocrine-progenitor", + "Fev": "Fev-EP", "Fev/Chgb": "Fev-EP", "Fev/Cacna2d1": "Fev-EP", + "Pax4/Neurod2": "endocrine-progenitor-primed", + "Pax4/Cacna2d1": "endocrine-progenitor-primed", + "Chgb/Pdx1": "pancreatic-progenitor", + "Gng12/Pdx1": "pancreatic-progenitor", + "Spp1": "ductal", "Sst/Ppy": "delta", + "__exocrine__": "exocrine", +} +# yu "Putative Cell Type" +YU_MAP = { + "Mesenchyme": "mesenchyme", "Acinar cell": "acinar", + "β": "beta", "βearly": "beta", "βlate": "beta", + "Tip": "Tip", "Duct": "ductal", "Trunk": "Trunk", "Trunk/Duct": "ductal", + "α": "alpha", "α-2nd": "alpha", + "EP1": "endocrine-progenitor", "EP2": "Fev-EP", + "EP3": "endocrine-progenitor", + "EP4": "endocrine-progenitor-primed", + "EP4early": "endocrine-progenitor-primed", + "EP4late": "endocrine-progenitor-primed", + "Immune cell": "immune", "δ": "delta", "Endothelium": "endothelial", + "α/PP-Pro-II": "epsilon", "α/PP-Pro-I (ε)": "epsilon", + "α/PP-Pro": "epsilon", + "PP": "gamma", + "Neuron": None, "-": None, +} +# mia cell_type_reannotatedIntegrated (Hrovatin 2023) +MIA_MAP = { + "beta": "adult-beta", "alpha": "adult-alpha", + "delta": "delta", "gamma": "gamma", + "endothelial": "endothelial", "immune": "immune", "ductal": "ductal", + "stellate a.": "mesenchyme", "stellate q.": "mesenchyme", + # mixed / low-quality: drop + "beta+delta": None, "alpha+delta": None, "beta+gamma": None, + "endo. prolif.": None, "delta+gamma": None, "alpha+beta": None, + "lowQ": None, "schwann": None, "E non-endo.": None, +} +# veres Assigned_cluster (21 raw labels) +VERES_MAP = { + "sc_alpha": "alpha_progenitor", + "sc_beta": "beta_progenitor", + "sc_ec": "epsilon", + "pdx1": "pancreatic-progenitor", + "nkx61": "endocrine-progenitor-primed", + "neurog3": "endocrine-progenitor", + "fev_high_isl_neg": "Fev-EP", + "sst_hhex": "delta", + "foxj1": "ductal", + "repl": "proliferating", + "exo": "exocrine", + "alpha": "alpha", "beta": "beta", "delta": "delta", "gamma": "gamma", + "acinar": "acinar", "ductal": "ductal", + "stellate": "mesenchyme", + "endothelial": "endothelial", "immune": "immune", + "other": None, +} + +# special-cased HGNC → mouse (default is str.capitalize()) +HGNC_TO_MOUSE_SPECIAL = { + "INS": "Ins1", # human INS -> mouse Ins1 (Ins2 stays zero for these cells) +} + + +def load_byrnes_labels(): + """byrnes v2 endo/exo seurat CSVs, keyed by (corpus_sample, bare_barcode).""" + D = LBL_DIR / "byrnes" + files = { + "E12_v2_endocrine_seur_ob_meta.csv": ("GSM3140915_E12_v2", "endo"), + "E12_v2_exocrine_seur_ob_meta.csv": ("GSM3140915_E12_v2", "exo"), + "E14_v2_endocrine_seur_ob_meta.csv": ("GSM3140916_E14_v2", "endo"), + "E14_v2_exocrine_seur_ob_meta.csv": ("GSM3140916_E14_v2", "exo"), + "E17_v2_agg_endocrine_seur_ob_meta.csv":("__E17_agg__", "endo"), + "E17_v2_agg_exocrine_seur_ob_meta.csv": ("__E17_agg__", "exo"), + "E14_fev_lineage_endocrine_seur_ob_meta.csv": ("GSM3140916_E14_v2", "endo"), + } + rows = [] + for f, (corpus_sample, kind) in files.items(): + df = pd.read_csv(D / f) + for _, r in df.iterrows(): + bc = str(r["cell"]) + if kind == "exo": + paper = "__exocrine__" + else: + paper = r.get("Cluster_Names", np.nan) + if not isinstance(paper, str) or paper == "" or paper == "nan": + continue # E17 endo aggregate & Fev endo have no cluster names + # E17 agg -1/-2 → split back into two GSMs; corpus barcodes always end in "-1" + if corpus_sample == "__E17_agg__": + if bc.endswith("-1"): + smp = "GSM3140917_E17_1_v2" + elif bc.endswith("-2"): + smp = "GSM3140918_E17_2_v2" + else: + smp = "GSM3140917_E17_1_v2" + bc_bare = bc.rsplit("-", 1)[0] + "-1" + else: + smp = corpus_sample + bc_bare = bc if bc.endswith("-1") else bc + "-1" + rows.append({"sample": smp, "bc": bc_bare, + "byrnes_paper_label": paper}) + lbl = pd.DataFrame(rows) + # cell may appear in both individual + merged files + lbl = lbl.drop_duplicates(subset=["sample", "bc"], keep="first") + return lbl + + +def load_yu_labels(): + df = pd.read_csv(LBL_DIR / "yu/yu_mouse_10x_cells.tsv", sep="\t") + df = df[["Batch", "Barcode", "Putative Cell Type"]].dropna( + subset=["Batch", "Barcode"]) + df = df.rename(columns={"Batch": "batch_short", + "Barcode": "bc_bare_16", + "Putative Cell Type": "yu_paper_label"}) + return df + + +def load_mia_labels(): + df = pd.read_csv(LBL_DIR / "mia/mia_GSE211796_labels.tsv", sep="\t") + def parse(cid): + parts = cid.split("-1-", 1) + if len(parts) != 2: + return None, None + return parts[0] + "-1", parts[1].split("-", 1)[0] # (bc, mouseN) + df["bc"], df["mouse"] = zip(*df["cell_id"].apply(parse)) + df = df[["mouse", "bc", "cell_type_reannotatedIntegrated"]].rename( + columns={"cell_type_reannotatedIntegrated": "mia_paper_label"}) + return df.dropna(subset=["mouse", "bc"]) + + +# veres raw counts loader (GSE114412 / Sharon 2019) + +def build_veres_adata(v1_var_names): + """load veres raw_indrops + Assigned_cluster, HGNC→mouse-capitalize, subset to v1 gene space.""" + files = sorted(glob.glob(str(SHARON_DIR / "*.raw_indrops_counts.tsv.gz"))) + if not files: + raise FileNotFoundError(f"No raw Veres files in {SHARON_DIR}") + + # aggregate label metadata (labels + Stage/CellLine/CellCondition) + meta_files = sorted(glob.glob(str(SHARON_DIR / "*.cell_metadata.tsv.gz"))) + meta_parts = [] + for mf in meta_files: + m = pd.read_csv(mf, sep="\t") + src = Path(mf).name.replace(".cell_metadata.tsv.gz", "") + m["src_file"] = src + # normalise Stage-like column + stage_col = None + for cand in ["Stage", "CellCondition", "CellProtocol"]: + if cand in m.columns: + stage_col = cand + break + if stage_col: + m["Stage"] = m[stage_col].astype(str) + else: + m["Stage"] = "unknown" + m = m[["library.barcode", "Assigned_cluster", "Stage", "src_file"]] + meta_parts.append(m) + meta = pd.concat(meta_parts, ignore_index=True) + keep_bc = set(meta["library.barcode"]) + print(f"[veres] {len(keep_bc)} labeled cells across {len(meta_files)} files", + flush=True) + + # line-by-line + per-file gene mapping (files carry different HGNC subsets, e.g. ES_iPS_comparison) + mouse_set = set(v1_var_names) + mouse_to_idx = {g: i for i, g in enumerate(v1_var_names)} + all_rows = [] + all_bc = [] + n_v = len(v1_var_names) + for fpath in files: + stem = Path(fpath).name.replace(".raw_indrops_counts.tsv.gz", "") + n_seen = n_kept = 0 + rows_this_file = [] + bc_this_file = [] + with gzip.open(fpath, "rt") as fh: + hdr_line = fh.readline().rstrip("\n") + raw_genes = hdr_line.split("\t")[1:] + keep_cols = [] + keep_out = [] + for c_idx, g in enumerate(raw_genes): + m = HGNC_TO_MOUSE_SPECIAL.get(g, g.capitalize()) + if m in mouse_set: + keep_cols.append(c_idx) + keep_out.append(mouse_to_idx[m]) + keep_cols = np.array(keep_cols, dtype=np.int64) + keep_out = np.array(keep_out, dtype=np.int64) + print(f"[veres] {stem}: {len(raw_genes)} raw genes → " + f"{len(keep_cols)} mouse-mapped", flush=True) + n_raw = len(raw_genes) + for line in fh: + n_seen += 1 + tab1 = line.find("\t") + bc = line[:tab1] + if bc not in keep_bc: + continue + vals = np.fromstring(line[tab1+1:], sep="\t", dtype=np.float32) + if len(vals) != n_raw: + continue + # scatter mapped genes into mouse space + sub = vals[keep_cols] + nz = sub > 0 + if nz.any(): + cols = keep_out[nz] + data = sub[nz] + # multiple HGNC → same mouse gene: sum + row_dense = np.zeros(n_v, dtype=np.float32) + np.add.at(row_dense, cols, data) + row_sp = sp.csr_matrix(row_dense) + else: + row_sp = sp.csr_matrix((1, n_v), dtype=np.float32) + rows_this_file.append(row_sp) + bc_this_file.append(bc) + n_kept += 1 + print(f"[veres] {stem}: scanned {n_seen}, kept {n_kept}", flush=True) + if rows_this_file: + all_rows.extend(rows_this_file) + all_bc.extend(bc_this_file) + + if not all_rows: + raise RuntimeError("Veres: no cells kept") + X = sp.vstack(all_rows).tocsr() + + obs = pd.DataFrame({"library.barcode": all_bc}).merge( + meta.drop_duplicates("library.barcode"), + on="library.barcode", how="left") + obs.index = ["veres_" + b for b in obs["library.barcode"]] + obs["dataset"] = "veres_GSE114412" + obs["source_sample"] = obs["src_file"] + obs["organism"] = "human_scislet" # SC-islet, mapped into mouse gene space + obs["sample"] = obs["src_file"] + obs["barcode"] = obs["library.barcode"] + + a = ad.AnnData(X=X, obs=obs, + var=pd.DataFrame(index=list(v1_var_names))) + + # match v1 corpus scale + sc.pp.normalize_total(a, target_sum=1e4) + sc.pp.log1p(a) + print(f"[veres] final adata: {a.shape}", flush=True) + return a + + +def attach_paper_labels(v1: ad.AnnData) -> pd.Series: + """returns (paper_label, label_source); cells without paper label get NaN + 'unlabeled'.""" + paper = pd.Series(np.nan, index=v1.obs.index, dtype=object) + src = pd.Series("unlabeled", index=v1.obs.index, dtype=object) + + # bastidas: cell_type already is a paper label + is_bast = v1.obs["dataset"] == "bastidas_GSE132188" + paper.loc[is_bast] = v1.obs.loc[is_bast, "cell_type"].astype(str).values + src.loc[is_bast] = "paper" + + # baron: assigned_cluster + is_bar = v1.obs["dataset"] == "baron_GSE84133_mouse" + paper.loc[is_bar] = v1.obs.loc[is_bar, "assigned_cluster"].astype(str).values + src.loc[is_bar] = "paper" + + # byrnes: join v2 endo/exo CSVs by (sample, bc) + byrnes_lbl = load_byrnes_labels() + is_byr = v1.obs["dataset"] == "byrnes_GSE101099" + byr = v1[is_byr].obs.copy() + byr["bc"] = byr.index.str.split("_").str[-1] + byr["sample"] = byr["sample"].astype(str) + m = byr.reset_index().merge(byrnes_lbl, on=["sample", "bc"], how="left")\ + .set_index(byr.reset_index().columns[0]) + matched = m["byrnes_paper_label"].notna() + print(f"[byrnes] label coverage: {matched.sum():,}/{len(byr):,} " + f"({matched.mean()*100:.1f}%)", flush=True) + paper.loc[byr.index[matched.values]] = m.loc[matched, "byrnes_paper_label"].values + src.loc[byr.index[matched.values]] = "paper" + + # yu: join TableS1 mouse_10x by (batch_short, bc_bare_16) + yu_lbl = load_yu_labels() + is_yu = v1.obs["dataset"] == "yu_GSE139627" + yu = v1[is_yu].obs.copy() + yu["bc_bare_16"] = yu.index.str.split("_").str[-1].str.replace( + "-1", "", regex=False) + yu["batch_short"] = yu["source_sample"].astype(str).apply( + lambda s: s.replace("Mouse_10xGenomics_", "").split("_", 1)[1] + if "Mouse_10xGenomics_" in s else s) + m = yu.reset_index().merge(yu_lbl, on=["batch_short", "bc_bare_16"], + how="left").set_index( + yu.reset_index().columns[0]) + matched = m["yu_paper_label"].notna() + print(f"[yu] label coverage: {matched.sum():,}/{len(yu):,} " + f"({matched.mean()*100:.1f}%)", flush=True) + paper.loc[yu.index[matched.values]] = m.loc[matched, "yu_paper_label"].values + src.loc[yu.index[matched.values]] = "paper" + + # mia: join by (mouse_id, bc) + mia_lbl = load_mia_labels() + is_mia = v1.obs["dataset"] == "mia_GSE211796" + mia = v1[is_mia].obs.copy() + mia["bc"] = mia.index.str.split("_").str[-1] + mia["mouse"] = mia["source_sample"].astype(str).str.split("_").str[-1] + m = mia.reset_index().merge(mia_lbl, on=["mouse", "bc"], how="left")\ + .set_index(mia.reset_index().columns[0]) + matched = m["mia_paper_label"].notna() + print(f"[mia] label coverage: {matched.sum():,}/{len(mia):,} " + f"({matched.mean()*100:.1f}%)", flush=True) + paper.loc[mia.index[matched.values]] = m.loc[matched, "mia_paper_label"].values + src.loc[mia.index[matched.values]] = "paper" + + return paper, src + + +def canonicalise(paper_label: str, dataset: str) -> str | None: + if paper_label is None or (isinstance(paper_label, float) and np.isnan(paper_label)): + return None + if paper_label in ("nan", "NaN", "UNK", "unassigned", "", "unknown"): + return None + pmap = { + "bastidas_GSE132188": BASTIDAS_MAP, + "baron_GSE84133_mouse": BARON_MAP, + "byrnes_GSE101099": BYRNES_MAP, + "yu_GSE139627": YU_MAP, + "mia_GSE211796": MIA_MAP, + "veres_GSE114412": VERES_MAP, + }.get(dataset, {}) + return pmap.get(paper_label, None) + + +def apply_min_class_filter(labels: pd.Series, min_size: int = MIN_CLASS_SIZE): + """drop labels below threshold.""" + vc = labels.value_counts() + drop = vc[vc < min_size].index.tolist() + report = [] + if drop: + for cls in drop: + report.append(f" drop '{cls}' (n={vc[cls]})") + labels = labels.where(~labels.isin(drop), other=None) + return labels, report + + +def main(): + print(f"Loading v1 corpus from {V1_CORPUS}", flush=True) + v1 = ad.read_h5ad(V1_CORPUS) + print(f" v1: {v1.shape}", flush=True) + + print("\n[1/6] Attaching paper labels ...", flush=True) + paper, src = attach_paper_labels(v1) + v1.obs["paper_label"] = paper.values + v1.obs["label_source"] = src.values + + # drop every v1 cell without a paper label + print("\n Per-dataset paper-label coverage (before filter):", flush=True) + before_counts = {} + for ds in v1.obs["dataset"].astype(str).unique(): + sub = v1.obs[v1.obs["dataset"] == ds] + n_tot = len(sub) + n_pap = (sub["label_source"] == "paper").sum() + before_counts[ds] = (n_tot, n_pap) + print(f" {ds:30s} total={n_tot:>7d} paper-labeled={n_pap:>7d} " + f"({n_pap/max(n_tot,1)*100:.1f}%)", flush=True) + is_paper_v1 = v1.obs["label_source"].values == "paper" + v1 = v1[is_paper_v1].copy() + print(f" v1 after paper-only filter: {v1.shape}", flush=True) + + # veres is 100% paper-labeled by construction (drops 'other') + print("\n[2/6] Loading Veres GSE114412 (raw + labels) ...", flush=True) + veres = build_veres_adata(list(v1.var_names)) + veres.obs["paper_label"] = veres.obs["Assigned_cluster"].astype(str).values + veres.obs["label_source"] = "paper" + + print("\n[3/6] Concatenating corpus ...", flush=True) + keep_obs = ["dataset", "sample", "source_sample", "organism", + "paper_label", "label_source"] + for c in keep_obs: + if c not in v1.obs.columns: + v1.obs[c] = np.nan + if c not in veres.obs.columns: + veres.obs[c] = np.nan + v1_slim = ad.AnnData(X=v1.X, obs=v1.obs[keep_obs].copy(), + var=v1.var.copy()) + v1_slim.obs_names = v1.obs_names + v1_slim.var_names = v1.var_names + ve_slim = ad.AnnData(X=veres.X, obs=veres.obs[keep_obs].copy(), + var=veres.var.copy()) + ve_slim.obs_names = veres.obs_names + ve_slim.var_names = veres.var_names + + corpus = ad.concat([v1_slim, ve_slim], join="outer", axis=0, + merge="same", label=None) + _, keep = np.unique(corpus.obs_names, return_index=True) + if len(keep) != corpus.n_obs: + print(f" removed {corpus.n_obs - len(keep)} duplicate obs_names", + flush=True) + corpus = corpus[np.sort(keep)].copy() + print(f" concat: {corpus.shape}", flush=True) + + print("\n[4/6] Canonicalising labels ...", flush=True) + canon = corpus.obs.apply( + lambda r: canonicalise(r["paper_label"], r["dataset"]), axis=1) + corpus.obs["canonical_label_raw"] = canon.values + print(" raw canonical distribution (top):", flush=True) + print(canon.value_counts().head(25).to_string(), flush=True) + + # drop cells whose paper_label wasn't canonicalisable (e.g. bastidas/veres "other", MIA mixed) + keep_can = canon.notna() + print(f" dropping {(~keep_can).sum():,} paper-labeled cells with " + f"un-mappable class (e.g. 'other', mixed)", flush=True) + corpus = corpus[keep_can.values].copy() + canon = canon[keep_can] + + canon2, drop_report = apply_min_class_filter(canon, MIN_CLASS_SIZE) + for r in drop_report: + print(r, flush=True) + keep_final = canon2.notna() + corpus = corpus[keep_final.values].copy() + corpus.obs["canonical_label"] = canon2[keep_final].values + + cl = corpus.obs["canonical_label"] + bad = cl.isin(["UNK", "unassigned", "nan", "NaN", ""]) + assert not bad.any(), f"bad tokens leaked: {cl.value_counts()}" + assert cl.notna().all(), "NaN labels leaked into canonical_label" + ls_uniq = set(corpus.obs["label_source"].astype(str).unique()) + assert ls_uniq == {"paper"}, f"non-paper label sources present: {ls_uniq}" + print(f"\n ALL {corpus.n_obs:,} cells carry canonical_label + label_source='paper'", + flush=True) + final_vc = cl.value_counts() + print(" FINAL vocabulary:", flush=True) + for k, v in final_vc.items(): + print(f" {k:35s} {v:>7d}", flush=True) + + print("\n[5/6] Per-dataset coverage table (before → after):", flush=True) + cov_rows = [] + veres_total = int((corpus.obs["dataset"] == "veres_GSE114412").sum()) + for ds, (n_before_total, n_before_paper) in before_counts.items(): + n_after = int((corpus.obs["dataset"] == ds).sum()) + cov_rows.append({"dataset": ds, + "n_before": n_before_total, + "n_paper": n_before_paper, + "n_after_v3": n_after, + "pct_kept": round(n_after / max(n_before_total, 1) * 100, 1)}) + # veres comes from Sharon extract, not v1 + cov_rows.append({"dataset": "veres_GSE114412", + "n_before": 57297, + "n_paper": 57297, + "n_after_v3": veres_total, + "pct_kept": round(veres_total / 57297 * 100, 1)}) + cov_df = pd.DataFrame(cov_rows) + print(cov_df.to_string(index=False), flush=True) + cov_df.to_csv(OUT_DIR / "coverage_v3.csv", index=False) + + # veres split: stratified by Stage, 12,297 held out + print("\n[6/6] Splitting Veres → train (45k) + held-out (12,297) ...", + flush=True) + ve_mask = corpus.obs["dataset"] == "veres_GSE114412" + ve_idx = corpus.obs.index[ve_mask].to_numpy() + ve_stage = corpus.obs.loc[ve_idx, "source_sample"].astype(str).values + hold_n = 12297 + hold_ids = [] + from collections import Counter + stage_counts = Counter(ve_stage) + total = len(ve_idx) + for stage, n_stage in stage_counts.items(): + share = int(round(hold_n * n_stage / total)) + pool = ve_idx[ve_stage == stage] + take = min(share, len(pool)) + chosen = RNG.choice(pool, size=take, replace=False) + hold_ids.extend(chosen.tolist()) + hold_ids = set(hold_ids[:hold_n]) + is_hold = corpus.obs.index.isin(hold_ids) + print(f" held-out size: {is_hold.sum():,} (target 12,297)", flush=True) + + held = corpus[is_hold].copy() + held.write_h5ad(HELD_OUT_L / "veres_GSE114412_test.h5ad") + print(f" wrote {HELD_OUT_L / 'veres_GSE114412_test.h5ad'}", flush=True) + + train_mask = ~is_hold + corpus_train = corpus[train_mask].copy() + print(f" training corpus: {corpus_train.shape}", flush=True) + + print("\n Selecting HVGs ...", flush=True) + # drop genes present in <10 cells (seurat-flavor edge case) + if sp.issparse(corpus_train.X): + gene_cells = np.asarray((corpus_train.X > 0).sum(axis=0)).ravel() + else: + gene_cells = (corpus_train.X > 0).sum(axis=0) + nz_mask = gene_cells >= 10 + print(f" {(~nz_mask).sum()} genes in <10 cells; kept {nz_mask.sum()}", + flush=True) + hvg_target = min(N_HVG, int(nz_mask.sum())) + tmp = corpus_train[:, nz_mask].copy() + tmp.X = tmp.X.astype(np.float32) + # variance-based HVG selection — scanpy's dispersion-vs-mean binning blows up + # on log-normalised data with heterogeneous library sizes across studies + if sp.issparse(tmp.X): + Xc = tmp.X.tocsc() + mean = np.asarray(Xc.mean(axis=0)).ravel() + sq = Xc.multiply(Xc).mean(axis=0) + sq = np.asarray(sq).ravel() + var = sq - mean ** 2 + else: + mean = tmp.X.mean(axis=0) + var = tmp.X.var(axis=0) + # rank by normalised dispersion + disp = var / (mean + 1e-8) + top = np.argsort(-disp)[:hvg_target] + top_mask_local = np.zeros(tmp.n_vars, dtype=bool) + top_mask_local[top] = True + hvg = np.zeros(corpus_train.n_vars, dtype=bool) + hvg[np.where(nz_mask)[0]] = top_mask_local + print(f" HVG: {hvg.sum()} genes (variance-ranked)", flush=True) + corpus_train.var["highly_variable"] = hvg + + print(" Fitting PCA basis ...", flush=True) + from sklearn.decomposition import PCA + X_hvg = corpus_train.X[:, hvg] + if sp.issparse(X_hvg): + X_hvg_d = X_hvg.toarray() + else: + X_hvg_d = np.asarray(X_hvg) + pca = PCA(n_components=min(N_PCA, X_hvg_d.shape[1]-1), random_state=0) + Xp = pca.fit_transform(X_hvg_d) + corpus_train.obsm["X_pca"] = Xp.astype(np.float32) + print(f" PCA: {Xp.shape}, ev={pca.explained_variance_ratio_.sum():.3f}", + flush=True) + + print("\nWriting outputs ...", flush=True) + corpus_path = OUT_DIR / "corpus_v3.h5ad" + corpus_train.write_h5ad(corpus_path) + print(f" wrote {corpus_path}", flush=True) + + np.savez(OUT_DIR / "corpus_stats_v3.npz", + n_cells=corpus_train.n_obs, n_genes=corpus_train.n_vars, + hvg_mask=hvg, + n_labeled=int(corpus_train.obs["canonical_label"].notna().sum()), + class_counts=json.dumps( + corpus_train.obs["canonical_label"].dropna() + .value_counts().to_dict()), + var_names=np.array(list(corpus_train.var_names))) + print(f" wrote {OUT_DIR / 'corpus_stats_v3.npz'}", flush=True) + + with open(OUT_DIR / "pca_basis_v3.pkl", "wb") as f: + pickle.dump({"pca": pca, + "hvg_names": list(corpus_train.var_names[hvg]), + "n_components": pca.n_components_}, f) + print(f" wrote {OUT_DIR / 'pca_basis_v3.pkl'}", flush=True) + + print("\n=== VERIFICATION ===", flush=True) + dup = corpus_train.obs_names.duplicated().sum() + print(f" duplicates: {dup}", flush=True) + bad_tokens = corpus_train.obs["canonical_label"].isin( + ["UNK", "unassigned", "nan"]).sum() + print(f" bad tokens (UNK/unassigned/nan) in canonical_label: {bad_tokens}", + flush=True) + final_vc = corpus_train.obs["canonical_label"].dropna().value_counts() + below = (final_vc < MIN_CLASS_SIZE).sum() + print(f" classes below {MIN_CLASS_SIZE}: {below}", flush=True) + print(f" final class list ({len(final_vc)} classes):", flush=True) + for k, v in final_vc.items(): + print(f" {k:35s} {v:>7d}", flush=True) + + print("\nDone.", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/pancreas/README.md b/scripts/pancreas/README.md new file mode 100644 index 0000000000000000000000000000000000000000..c2aad7f8f2547a92c5d7fde829795c87a18e296b --- /dev/null +++ b/scripts/pancreas/README.md @@ -0,0 +1,42 @@ +# Pan-Pancreas PANDA + +Multi-dataset corpus for pancreas cell identity + zero-shot transfer + mechanistic discovery. + +## Corpus datasets + +| Accession | Study | Cells | Species | Modality | Role | +|---|---|---:|---|---|---| +| GSE132188 | Bastidas-Ponce 2019 | 33,896 | mouse | 10x v2 | Anchor E15.5 endocrine progenitor time course (local) | +| GSE84133 | Baron 2016 | ~10,000 | mouse+human | inDrops | Adult islet reference | +| GSE85241 | Muraro 2016 | ~2,126 | human | CEL-Seq2 | Deep-coverage adult islet | +| E-MTAB-5061 | Segerstolpe 2016 | ~3,514 | human | Smart-seq2 | Deep + T2D covariate | +| GSE81076 | Grun 2016 | ~1,595 | human | CEL-Seq | Modality diversity | +| GSE101099 | Byrnes 2018 | ~12,000 | mouse | inDrops | Independent-lab time course | + +## Held-out labeled validation + +**Baron 2016 mouse subset** — has full canonical islet labels (alpha/beta/delta/gamma/epsilon/ +ductal/acinar/endothelial/immune) but is a completely independent scRNA-seq study from a +different lab. Zero-shot transfer of PANDA-pancreas predictions to Baron gives an +apples-to-apples accuracy measurement across labs. + +## Held-out unlabeled discovery target + +**Veres 2019 hPSC-directed differentiation (GSE114412)** — human pluripotent stem cell +differentiation to SC-beta cells with an off-target enterochromaffin population. Weakly +labeled by day-of-differentiation only; PANDA's cell-type predictions become hypothesis- +generating annotations. + +## Ontology (11 classes) + +- alpha +- beta +- delta +- gamma (PP) +- epsilon +- endocrine-progenitor +- ductal +- acinar +- endothelial +- immune +- other/mesenchymal diff --git a/scripts/pancreas/__init__.py b/scripts/pancreas/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/scripts/pancreas/_loaders/__init__.py b/scripts/pancreas/_loaders/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/scripts/pancreas/known_pancreas_markers.yaml b/scripts/pancreas/known_pancreas_markers.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a7d520e2428cc898566bedcaf5f66228b7f245b4 --- /dev/null +++ b/scripts/pancreas/known_pancreas_markers.yaml @@ -0,0 +1,19 @@ +# canonical marker gene sets used only for label assignment during pan-pancreas corpus build. +# not used as input features to the classifier. + +classes: + alpha: [Gcg, Arx, Irx1, Irx2, Mafb, Ttr] + beta: [Ins1, Ins2, Mafa, Nkx6-1, Pdx1, Slc30a8, Ucn3] + delta: [Sst, Hhex, Rbp4, Pdx1] + gamma: [Ppy, Pyy, Meis2] + epsilon: [Ghrl, Arx] + endocrine-progenitor: [Neurog3, Neurod1, Pax4, Fev, Chga, Chgb] + ductal: [Krt19, Sox9, Hnk1st, Muc1, Cftr, Spp1] + acinar: [Cpa1, Ptf1a, Amy1, Amy2a, Prss1, Cel] + endothelial: [Pecam1, Cdh5, Tie1, Tek, Kdr] + immune: [Ptprc, Cd3e, Cd8a, Cd4, Adgre1, Itgam] + other: [] +assignment: + min_score: 0.10 + min_margin: 0.05 + cluster_resolution: 0.8 diff --git a/scripts/push_all.sh b/scripts/push_all.sh new file mode 100644 index 0000000000000000000000000000000000000000..474f96fa68b97f279192bc0f54368e6316424515 --- /dev/null +++ b/scripts/push_all.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +# push current tree to both github and hugging face. +# requires GH_TOKEN and HF_TOKEN in env (never hard-code secrets here). +set -e + +: "${GH_TOKEN:?set GH_TOKEN}" +: "${HF_TOKEN:?set HF_TOKEN}" + +# github (code + paper + figures, no data) +git push "https://${GH_TOKEN}@github.com/bryanc5864/PRISM.git" main + +# hugging face: use the hf CLI for structured uploads. +# omit --include so we push everything except caches / build junk. +hf auth login --token "$HF_TOKEN" >/dev/null 2>&1 + +for d in panda scripts figures; do + [ -d "$d" ] && hf upload bryan7264/PANDA "$d" "$d" --repo-type model \ + --exclude "**/__pycache__/**" --exclude "**/.cache/**" \ + --commit-message "sync $d" || true +done + +for f in PAPER.tex PAPER.pdf README.md; do + [ -f "$f" ] && hf upload bryan7264/PANDA "$f" "$f" --repo-type model \ + --commit-message "sync $f" || true +done + +echo "push complete" diff --git a/scripts/sulic/README.md b/scripts/sulic/README.md new file mode 100644 index 0000000000000000000000000000000000000000..c3dd7ed050ae3784f0781afa647a87d3757ba972 --- /dev/null +++ b/scripts/sulic/README.md @@ -0,0 +1,11 @@ +# scripts/sulic + +Sulic-in-PANDA held-out reproduction, kept in its own dir because it uses a +different CV design from the main pipeline. + +- `sulic_panda_heldout.py` -- runs Test A (binary placode vs epithelium) and + Test C (4-way placode subtype), 5-fold stratified CV, model retrained per fold. +- `sulic_panda_heldout_results.json` -- last-run per-fold + mean results. + +both tests use the pan-skin PANDA architecture (`panda/model.py`) +straight from the corresponding checkpoint's trunk sizing. diff --git a/scripts/sulic/sulic_panda_heldout.py b/scripts/sulic/sulic_panda_heldout.py new file mode 100644 index 0000000000000000000000000000000000000000..28eb7c737220367171b0396f9b2bbe5fc4db9b4d --- /dev/null +++ b/scripts/sulic/sulic_panda_heldout.py @@ -0,0 +1,198 @@ +"""sulic-in-panda held-out 5-fold cv. Test A: binary facs (placode vs epi). Test C: 4-way placode subtype.""" +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 scanpy as sc +import scipy.sparse as sp +import torch +import torch.nn.functional as F +from torch.utils.data import Dataset, DataLoader +from sklearn.decomposition import PCA +from sklearn.model_selection import StratifiedKFold +from sklearn.preprocessing import StandardScaler +from sklearn.metrics import roc_auc_score, accuracy_score + +sys.path.insert(0, "/home/bcheng/PRISM") +from panda.pan_skin.model import ( + PANDAEncoder, supcon_loss, vicreg_loss, prototype_infonce +) + +SULIC_H5AD = Path("/home/bcheng/PRISM/data/processed/sulic/adata_sulic_clustered.h5ad") +OUT = Path("/home/bcheng/PRISM/scripts/sulic") +OUT.mkdir(parents=True, exist_ok=True) + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +N_FOLDS = 5 + + +class SulicDataset(Dataset): + def __init__(self, X, y): + self.X = X.astype(np.float32); self.y = y.astype(np.int64) + 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.zeros(1, dtype=torch.int64), # single-dataset + torch.zeros(2, dtype=torch.float32)) # no aux + + +class PxKSampler: + def __init__(self, y, P=None, K=16, n_batches=80, seed=0): + self.y = np.asarray(y) + self.classes = np.unique(self.y) + self.P = P or len(self.classes) + self.K = K + self.n_batches = n_batches + self.rng = np.random.default_rng(seed) + self.by_cls = {c: np.where(self.y == c)[0] for c in self.classes} + def __iter__(self): + for _ in range(self.n_batches): + classes_p = self.rng.choice(self.classes, + size=min(self.P, len(self.classes)), + replace=False) + batch = [] + for c in classes_p: + idx = self.by_cls[c] + take = self.K + pick = self.rng.choice(idx, size=take, replace=(len(idx) < take)) + batch.extend(pick.tolist()) + yield batch + def __len__(self): return self.n_batches + + +def prepare_pca(a, n_pca=50): + if a.raw is not None: + a = a.raw.to_adata() + sc.pp.normalize_total(a, target_sum=1e4) + sc.pp.log1p(a) + sc.pp.highly_variable_genes(a, n_top_genes=2000, flavor="seurat", subset=False) + a = a[:, a.var["highly_variable"]].copy() + X = a.X.toarray() if sp.issparse(a.X) else a.X + scaler = StandardScaler().fit(X) + Xz = np.clip(scaler.transform(X), -10, 10) + pca = PCA(n_components=n_pca, random_state=42).fit(Xz) + Xp = pca.transform(Xz).astype(np.float32) + return a, Xp + + +def train_fold(Xp, y, classes, tr, te, fold_id, ensemble_seeds=5): + K = len(classes) + all_probs = [] + for seed in range(ensemble_seeds): + torch.manual_seed(fold_id * 100 + seed) + np.random.seed(fold_id * 100 + seed) + torch.cuda.empty_cache() + model = PANDAEncoder(n_pca=Xp.shape[1], n_classes=K, + n_datasets=1).to(DEVICE) + opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4) + ds = SulicDataset(Xp[tr], y[tr]) + sampler = PxKSampler(y[tr], K=16, n_batches=80, seed=seed) + loader = DataLoader(ds, batch_sampler=sampler, num_workers=0) + for stage, ne in enumerate([15, 20, 25]): + for e in range(ne): + for X_b, y_b, _, aux_b in loader: + X_b, y_b = X_b.to(DEVICE), y_b.to(DEVICE) + aux_b = aux_b.to(DEVICE) + out = model(X_b, aux_b, lam_dann=0.0) + L_sup = supcon_loss(out["z"], y_b) + L_vic = vicreg_loss(out["z"]) + L_ce = F.cross_entropy(out["logits"], y_b, 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 + 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) + model.eval() + with torch.no_grad(): + Xt = torch.from_numpy(Xp[te]).to(DEVICE) + aux = torch.zeros(len(te), 2, device=DEVICE) + out = model(Xt, aux, lam_dann=0.0) + cos = out["z"] @ model.prototypes.T + probs = torch.softmax(cos / 0.07, dim=1).cpu().numpy() + all_probs.append(probs) + ensemble_probs = np.mean(all_probs, axis=0) + pred = ensemble_probs.argmax(axis=1) + yte = y[te] + return pred, ensemble_probs, yte + + +def evaluate_test(name, X, y_bin_or_multi, class_list): + print(f"\n=== {name} ===", flush=True) + skf = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=42) + aurocs, accs = [], [] + for fold, (tr, te) in enumerate(skf.split(X, y_bin_or_multi)): + t0 = time.time() + pred, probs, yte = train_fold(X, y_bin_or_multi, class_list, tr, te, fold) + acc = accuracy_score(yte, pred) + if len(class_list) == 2: + auc = roc_auc_score(yte, probs[:, 1]) + else: + try: + auc = roc_auc_score(np.eye(len(class_list))[yte], probs, + average="macro", multi_class="ovr") + except Exception: + auc = float("nan") + print(f"[{name} fold {fold}] acc={acc:.4f} AUROC={auc:.4f} " + f"wall={time.time()-t0:.0f}s", flush=True) + aurocs.append(auc); accs.append(acc) + print(f"[{name}] MEAN AUROC = {np.mean(aurocs):.4f} +- {np.std(aurocs):.4f}", flush=True) + print(f"[{name}] MEAN ACC = {np.mean(accs):.4f} +- {np.std(accs):.4f}", flush=True) + return {"aurocs": aurocs, "accs": accs, + "mean_auroc": float(np.mean(aurocs)), + "std_auroc": float(np.std(aurocs)), + "mean_acc": float(np.mean(accs)), + "std_acc": float(np.std(accs))} + + +def main(): + print(f"[sulic-panda] loading {SULIC_H5AD}", flush=True) + a = ad.read_h5ad(SULIC_H5AD) + print(f"[sulic-panda] shape {a.shape}, samples: {a.obs['sample'].value_counts().to_dict()}", + flush=True) + + a_p, Xp = prepare_pca(a, n_pca=50) + print(f"[sulic-panda] Xp {Xp.shape}", flush=True) + + y_A = (a.obs["sample"].isin(["Placode1", "Placode2"])).astype(int).values + print(f"[sulic-panda] Test A class balance: {np.bincount(y_A).tolist()}", flush=True) + resA = evaluate_test("TestA", Xp, y_A, ["Epithelium", "Placode"]) + + if "placode_enriched" in a.obs.columns: + mask_p = (a.obs["placode_enriched"] == 1).values + sub = a[mask_p].copy() + if "paper_subtype" not in sub.obs.columns: + # fallback: kmeans on placode-cell Xp gives 4 pseudo-subtypes + print("[sulic-panda] paper_subtype missing — deriving 4-way clustering on Xp", flush=True) + from sklearn.cluster import KMeans + Xp_sub = Xp[mask_p] + km = KMeans(n_clusters=4, random_state=42, n_init=10).fit(Xp_sub) + paper_subtype = np.array([f"PlacodeK{i}" for i in km.labels_]) + else: + paper_subtype = sub.obs["paper_subtype"].astype(str).values + cls = sorted(np.unique(paper_subtype)) + y_C = np.array([cls.index(v) for v in paper_subtype], dtype=np.int64) + Xp_C = Xp[mask_p] + print(f"[sulic-panda] Test C n={len(y_C)}, classes={cls}, " + f"counts={np.bincount(y_C).tolist()}", flush=True) + resC = evaluate_test("TestC", Xp_C, y_C, cls) + else: + resC = None + + # save + result = {"testA": resA, "testC": resC} + with open(OUT / "sulic_panda_heldout_results.json", "w") as f: + json.dump(result, f, indent=2) + print(f"\n[sulic-panda] wrote {OUT}/sulic_panda_heldout_results.json", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/sulic/sulic_panda_heldout_results.json b/scripts/sulic/sulic_panda_heldout_results.json new file mode 100644 index 0000000000000000000000000000000000000000..1c76931c6a8d8f38341f3cc716c692ad8743c97d --- /dev/null +++ b/scripts/sulic/sulic_panda_heldout_results.json @@ -0,0 +1,42 @@ +{ + "testA": { + "aurocs": [ + 0.9998735687988025, + 0.9997746150729335, + 0.9999493517017828, + 0.9981990665584416, + 0.9999746347402597 + ], + "accs": [ + 0.991462113127001, + 0.9967982924226254, + 0.9946638207043756, + 0.9957264957264957, + 0.9957264957264957 + ], + "mean_auroc": 0.999554247374444, + "std_auroc": 0.0006811572660512661, + "mean_acc": 0.9948754435413987, + "std_acc": 0.0018352948649496838 + }, + "testC": { + "aurocs": [ + 0.9987616622158157, + 0.9939322298633184, + 0.9976578438776973, + 0.993164500030369, + 0.9992729716996815 + ], + "accs": [ + 0.9692058346839546, + 0.9497568881685575, + 0.9707792207792207, + 0.961038961038961, + 0.9675324675324676 + ], + "mean_auroc": 0.9965578415373763, + "std_auroc": 0.0025237815189043593, + "mean_acc": 0.9636626744406322, + "std_acc": 0.007701972730473335 + } +} \ No newline at end of file