scripts/analysis/66_dahlin_kit_mutant.py

"""zero-shot HSC PANDA on dahlin 2018: WT vs Kit W41/W41 class enrichment."""
from __future__ import annotations
from pathlib import Path
import warnings, json, sys, pickle
warnings.filterwarnings("ignore")
import numpy as np, pandas as pd, anndata as ad, scanpy as sc, scipy.sparse as sp
import torch
from scipy import stats
import os as _os
from pathlib import Path as _Path
PANDA_ROOT = _Path(_os.environ.get("PANDA_ROOT", str(_Path(__file__).resolve().parents[2])))
sys.path.insert(0, str(PANDA_ROOT))
from panda.model import PANDAEncoder

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
CKPT = Path(str(PANDA_ROOT / "checkpoints/hematopoiesis"))
HARM = Path(str(PANDA_ROOT / "data/corpus/hematopoiesis/harmonized"))
DAHLIN = Path(str(PANDA_ROOT / "data/corpus/hematopoiesis/held_out_unlabeled/dahlin_extract"))
OUT = Path(str(PANDA_ROOT / "discovery/hematopoiesis/marker"))

GENOTYPE_MAP = {
    "SIGAB1": "WT", "SIGAC1": "WT", "SIGAD1": "WT",
    "SIGAF1": "WT", "SIGAG1": "WT", "SIGAH1": "WT",
    "SIGAG8": "Kit_W41", "SIGAH8": "Kit_W41",
}

# Dahlin sorted TWO different FACS gates (GEO source_name, GSE107727):
#   SIGAB1/C1/D1 -> "LSK_sample_1..3"                     (WT, LSK gate)
#   SIGAF1/G1/H1 -> "Lin_negative_cKit_positive_sample_4..6" (WT, LK gate)
#   SIGAG8/H8    -> "Lin_negative_cKit_positive_sample_39/40" (W41, LK gate)
# LSK is a subset of LK enriched for immature progenitors, so the two WT
# groups differ enormously (WT-LSK is ~84% MPP / ~1% erythroid; WT-LK is
# ~36% / ~35%). Pooling them into one WT baseline against LK-only mutants
# makes the sorting gate, not the genotype, the dominant contrast: the pure
# WT-LK-vs-WT-LSK effect is log2fc +4.79 erythroid / +4.22 myeloid, i.e.
# LARGER than the Kit-W41 effect it was being attributed to.
# Dahlin themselves compare W41 LK vs WT LK. We now do the same.
GATE_MAP = {
    "SIGAB1": "LSK", "SIGAC1": "LSK", "SIGAD1": "LSK",
    "SIGAF1": "LK", "SIGAG1": "LK", "SIGAH1": "LK",
    "SIGAG8": "LK", "SIGAH8": "LK",
}
GATE_MATCHED = "LK"  # restrict the genotype contrast to this gate


def load_dahlin_all():
    print("[dahlin] loading 8 samples …", flush=True)
    parts = []
    for f in sorted(DAHLIN.glob("*.txt.gz")):
        gsm = f.name.split("_")[0]
        sample = f.name.split("_")[1].split(".")[0]
        genotype = GENOTYPE_MAP.get(sample, "unknown")
        print(f"  {sample} ({genotype})", flush=True)
        df = pd.read_csv(f, sep="\t", compression="gzip", index_col=0)
        X = sp.csr_matrix(df.values.T.astype(np.float32))
        obs = pd.DataFrame(index=[f"{sample}_{bc}" for bc in df.columns.astype(str)])
        obs["sample"] = sample
        obs["genotype"] = genotype
        obs["gate"] = GATE_MAP.get(sample, "unknown")
        obs["dataset"] = "dahlin_GSE107727"
        var = pd.DataFrame(index=df.index.astype(str))
        var["ensmusg"] = var.index.values
        a = ad.AnnData(X=X, obs=obs, var=var)
        parts.append(a)
    return ad.concat(parts, join="outer", label="_batch")


def convert_ensembl_to_symbol(a):
    import mygene
    mg = mygene.MyGeneInfo()
    ids = a.var_names.astype(str).tolist()
    print(f"[dahlin] querying {len(ids)} ENSMUSG IDs …", flush=True)
    res = mg.querymany(ids, scopes="ensembl.gene", fields="symbol",
                       species="mouse", verbose=False)
    id2sym = {r["query"]: r["symbol"] for r in res if "symbol" in r}
    syms = pd.Series(a.var_names.astype(str)).map(id2sym).values
    keep = pd.notna(syms)
    print(f"[dahlin] mapped {int(keep.sum())}/{len(a.var_names)} genes", flush=True)
    a = a[:, keep].copy()
    a.var_names = syms[keep]
    a.var_names_make_unique()
    return a


def main():
    a = load_dahlin_all()
    print(f"[dahlin] concat shape: {a.shape}", flush=True)
    a = convert_ensembl_to_symbol(a)
    print(f"[dahlin] after symbol conversion: {a.shape}", flush=True)

    ck = torch.load(CKPT / "panda_final.pt", map_location=DEVICE, weights_only=False)
    classes = ck["classes"]; datasets = ck["datasets"]
    model = PANDAEncoder(n_pca=50, n_classes=len(classes),
                         n_datasets=len(datasets)).to(DEVICE).eval()
    model.load_state_dict(ck["model"])

    stats_ = np.load(HARM / "corpus_stats.npz", allow_pickle=True)
    shared_hvgs = [str(g) for g in stats_["shared_hvgs"]]
    mu, sig = stats_["mean"], stats_["std"]
    with open(HARM / "pca_basis.pkl", "rb") as f: pca = pickle.load(f)

    G = len(shared_hvgs)
    hvg2i = {g: i for i, g in enumerate(shared_hvgs)}
    common = [g for g in a.var_names.astype(str) if g in hvg2i]
    frac = len(common) / G
    print(f"[proj] {len(common)}/{G} HVGs present ({frac:.1%})", flush=True)
    a_c = a[:, common].copy()
    sc.pp.normalize_total(a_c, target_sum=1e4); sc.pp.log1p(a_c)
    X = a_c.X.toarray().astype(np.float32) if sp.issparse(a_c.X) else a_c.X.astype(np.float32)
    Xf = np.zeros((a.n_obs, G), dtype=np.float32)
    cols = [hvg2i[g] for g in common]; Xf[:, cols] = X
    Xz = np.clip((Xf - mu.astype(np.float32)) / sig.astype(np.float32), -10, 10)
    Xpca = pca.transform(Xz).astype(np.float32)

    all_z = []
    with torch.no_grad():
        for i in range(0, a.n_obs, 4096):
            xb = torch.from_numpy(Xpca[i:i+4096]).to(DEVICE)
            aux = torch.zeros(len(xb), 2, device=DEVICE)
            all_z.append(model(xb, aux, lam_dann=0.0)["z"].cpu().numpy())
    Z = np.concatenate(all_z, axis=0)
    protos = ck["prototypes"]
    protos = protos / (np.linalg.norm(protos, axis=1, keepdims=True) + 1e-8)
    cos = Z @ protos.T
    pred = np.array([classes[i] for i in cos.argmax(axis=1)], dtype=object)
    conf = cos.max(axis=1)
    a.obs["pred_label"] = pred
    a.obs["pred_conf"] = conf.astype(np.float32)

    print(f"\n[dahlin] predicted class distribution overall:")
    print(a.obs["pred_label"].value_counts())
    print(f"\n[dahlin] per genotype:")
    xt = pd.crosstab(a.obs["pred_label"], a.obs["genotype"], normalize="columns")
    print(xt.round(4))
    xt.to_csv(OUT / "66_dahlin_class_per_genotype.csv")

    def enrichment(mask_a, mask_b, label_a, label_b):
        rows = []
        n_a = int(mask_a.sum()); n_b = int(mask_b.sum())
        for c in classes:
            n_c_a = int(((pred == c) & mask_a).sum())
            n_c_b = int(((pred == c) & mask_b).sum())
            odds, p = stats.fisher_exact(np.array(
                [[n_c_a, n_a - n_c_a], [n_c_b, n_b - n_c_b]]))
            f_a = (n_c_a + 1) / (n_a + 2); f_b = (n_c_b + 1) / (n_b + 2)
            rows.append({"class": c, f"n_{label_a}": n_c_a, f"n_{label_b}": n_c_b,
                         f"pct_{label_a}": round(100 * n_c_a / max(n_a, 1), 3),
                         f"pct_{label_b}": round(100 * n_c_b / max(n_b, 1), 3),
                         "log2_fold": round(np.log2(f_a / f_b), 3),
                         "fisher_p": p})
        return pd.DataFrame(rows).sort_values("log2_fold"), n_a, n_b

    gate = a.obs["gate"].values
    geno = a.obs["genotype"].values

    # PRIMARY, gate-matched: W41 LK vs WT LK (what Dahlin themselves compare)
    print(f"\n[dahlin] PRIMARY gate-matched Fisher enrichment "
          f"(Kit_W41 {GATE_MATCHED} vs WT {GATE_MATCHED}):")
    df, n_kit, n_wt = enrichment((geno == "Kit_W41") & (gate == GATE_MATCHED),
                                 (geno == "WT") & (gate == GATE_MATCHED),
                                 "Kit_W41", "WT")
    print(f"  Kit_W41 n={n_kit}, WT n={n_wt} (both {GATE_MATCHED} gate)")
    print(df.to_string(index=False))
    df.to_csv(OUT / "66_dahlin_enrichment.csv", index=False)

    # NEGATIVE CONTROL: WT LK vs WT LSK — same genotype, gate only. Quantifies
    # how much of any "genotype" signal is really the sorting gate.
    print("\n[dahlin] NEGATIVE CONTROL (WT LK vs WT LSK, genotype held constant):")
    dfg, n_lk, n_lsk = enrichment((geno == "WT") & (gate == "LK"),
                                  (geno == "WT") & (gate == "LSK"), "WT_LK", "WT_LSK")
    print(f"  WT_LK n={n_lk}, WT_LSK n={n_lsk}")
    print(dfg.to_string(index=False))
    dfg.to_csv(OUT / "66_dahlin_gate_negative_control.csv", index=False)

    # legacy pooled contrast, kept for provenance only
    dfp, _, _ = enrichment(geno == "Kit_W41", geno == "WT", "Kit_W41", "WT")
    dfp.to_csv(OUT / "66_dahlin_enrichment_pooled_CONFOUNDED.csv", index=False)

    a.obs.to_csv(OUT / "66_dahlin_predictions.csv")
    print(f"\n[dahlin] complete. Outputs in {OUT}/")


if __name__ == "__main__":
    main()