scripts/analysis/103_replicate_dingwall_seurat_pipeline.py

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

import os as _os
from pathlib import Path as _Path
PANDA_ROOT = _Path(_os.environ.get("PANDA_ROOT", str(_Path(__file__).resolve().parents[2])))
ROOT = Path(str(PANDA_ROOT))
RAW_H5 = ROOT / "data/raw/GSE220977_combined.h5ad"
DERM_MARKERS = ROOT / "data/external_labels/dingwall_supp/biorxiv_media-3.xlsx"
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()