scripts/analysis/98_eden_posthoc_detection.py

"""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
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 import PANDAEncoder

ROOT = Path(str(PANDA_ROOT))
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()