Mouse brain (Visium): per-cell-type K selection (reproduce)

Part 1 — compute: synchronized per-cell-type K selection

Uses select_K_synchronized: all replicates advance round-by-round in lockstep; after each round a single keep/stop decision per cell type is made from the replicates and applied to all of them. A cell type keeps its K-th SGP only if the marginal held-out ΔLL is significantly > 0 (one-sided t-test, p < ALPHA) and the K-th loading is reproducible across replicates (mean pairwise |cosine| ≥ STAB_THRESH). 59 cell types.

import os
GPUS = [0, 1, 2, 3, 4]              # one worker per replicate; set None for in-process single-GPU
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(g) for g in GPUS)

import sys, pickle, warnings
warnings.filterwarnings("ignore")
sys.path.append(r"/import/home2/share/yqzeng/MacSGP/codes")

import numpy as np
import pandas as pd
import scanpy as sc
import matplotlib.pyplot as plt
import matplotlib as mpl

from MacSGP_unified.k_selection import select_K_synchronized

mpl.rcParams.update({"figure.dpi": 110, "savefig.dpi": 200, "font.family": "DejaVu Sans",
    "axes.titlesize": 9, "axes.labelsize": 8, "axes.spines.top": False, "axes.spines.right": False})

DATA_PATH = "/import/home2/share/yqzeng/MacSGP/data/Mouth_brain_visium"
SAVE_PATH = "/import/home2/share/yqzeng/MacSGP/Results/Mouth_brain_visium"
os.makedirs(SAVE_PATH, exist_ok=True)

adata_st = sc.read_h5ad(os.path.join(DATA_PATH, "visium_1.h5ad"))
adata_basis = sc.read_h5ad(os.path.join(DATA_PATH, "visium_1_basis.h5ad"))
for k in [("obsm", "count"), ("obsm", "proportion"), ("obs", "library_size"),
          ("obs", "alpha"), ("var", "gamma")]:
    assert k[1] in getattr(adata_st, k[0]), f"missing {k}"
print(adata_st.shape, "| n_celltypes:", adata_basis.shape[0])

SGP_STEPS = 3000
LR = 2e-3
N_LAYERS = 4
K_MAX = 5
HOLDOUT_RATIO = 0.10
N_REPEATS = 5
SEEDS = [20260703, 20260704, 20260705, 20260706, 20260707]
ALPHA = 0.05          # t-test significance for keeping an SGP
MIN_EFFECT = 0.0      # optional minimum marginal held-out ΔLL to keep (absolute)
METHOD = "ttest"      # or "onese"
STAB_THRESH = 0.8     # cross-replicate loading |cos| gate; 0 to disable

gpu_ids = list(range(len(GPUS))) if GPUS is not None else None

res = select_K_synchronized(
    adata_st,
    adata_basis,
    model_kwargs={"training_steps": SGP_STEPS, "lr": LR, "n_layers": N_LAYERS},
    K_max=K_MAX,
    holdout_ratio=HOLDOUT_RATIO,
    n_repeats=N_REPEATS,
    seeds=SEEDS,
    train_kwargs={"fix_gamma_k": True},
    alpha=ALPHA,
    min_effect=MIN_EFFECT,
    method=METHOD,
    stab_thresh=STAB_THRESH,
    gpus=gpu_ids,
    verbose=True,
)

with open(os.path.join(SAVE_PATH, "sync_selection_result.pkl"), "wb") as f:
    pickle.dump(res, f)
cts = res["celltypes"]; Kc = res["recommended_K_per_celltype"]
print("\nselected K per cell type:")
print(pd.Series(Kc).sort_values(ascending=False).to_string())

# ---- fit the FULL data at the selected per-cell-type K ----
from MacSGP_unified.model import Model

celltypes = list(adata_basis.obs.index)
Kc_list = [int(Kc[c]) for c in celltypes]
maxK = max(max(Kc_list), 1)
print("maxK =", maxK, "| n cell types with K>=1:", sum(k >= 1 for k in Kc_list))

full = Model(adata_st, adata_basis, n_SGPs=maxK, seed=1234,
             training_steps=SGP_STEPS, lr=LR, n_layers=N_LAYERS)
full.train_greedy(fixed_Kc=Kc_list, fix_gamma_k=True)   # fit exactly K_c SGPs per cell type
adata_res = full.eval(sort_sgps=False)                  # factors -> adata_res.obsm[celltype], loadings -> varm
adata_res.write_h5ad(os.path.join(SAVE_PATH, "adata_sync_fit.h5ad"))
print("saved fit ->", os.path.join(SAVE_PATH, "adata_sync_fit.h5ad"))

Part 2 — compute: pathway enrichment of the selected SGPs (poles separated)

Restricted to the neuronal subtypes analysed in revision/vis/1.1/mouse_brain_k=10.ipynb. Because these are neuronal subtypes, enrichment uses GO Biological Process (not the cancer-oriented Hallmark set). Mouse gene symbols are mapped to human orthologs by upper-casing (a quick heuristic that recovers most 1:1 orthologs), so the standard human Enrichr libraries apply. Both pole-separated over-representation and directional GSEA (prerank) are produced. Requires internet (Enrichr) + gseapy.

import time
import gseapy as gp

ORGANISM = "mouse"                       # -> map symbols to human orthologs by upper-casing
LIB = "GO_Biological_Process_2021"       # GO-BP for neuronal subtypes
KEEP = ["Ext_Amy_2", "Ext_Hpc_CA1", "Ext_Hpc_CA3", "Ext_Hpc_DG1", "Ext_L23",
        "Ext_L5_1", "Ext_L5_2", "Ext_Med", "Ext_Pir", "Ext_Thal_1", "Ext_Thal_2",
        "Inh_1", "Inh_4"]                # analysed in revision/vis/1.1 (incl. K=1)

FIT_H5AD = os.path.join(SAVE_PATH, "adata_sync_fit.h5ad")
SEL_PKL = os.path.join(SAVE_PATH, "sync_selection_result.pkl")
ENR_DIR = os.path.join(SAVE_PATH, "enrichment"); os.makedirs(ENR_DIR, exist_ok=True)
print("gseapy", gp.__version__)

adata = sc.read_h5ad(FIT_H5AD)
Kc = pickle.load(open(SEL_PKL, "rb"))["recommended_K_per_celltype"]
genes = np.asarray(adata.var_names)
if ORGANISM == "mouse":
    genes = np.array([str(g).upper() for g in genes])   # approx human ortholog

# (cell type, SGP index) -> loading vector, only for kept SGPs (k < K_c)
sgp_loadings = {}
for c in KEEP:
    K = int(Kc.get(c, 0))
    if K < 1 or c not in adata.varm:
        continue
    L = np.asarray(adata.varm[c])
    for k in range(K):                     # include K=1 cell types
        sgp_loadings[(c, k)] = L[:, k]
print(f"{len({c for c, _ in sgp_loadings})} cell types, {len(sgp_loadings)} SGPs")
print({c: int(Kc[c]) for c in KEEP if Kc.get(c, 0) >= 1})

TOP_N = 100; PADJ_CUTOFF = 0.05; N_PERM = 1000; FDR_CUTOFF = 0.25

def pole_genes(loading, pole, n=TOP_N):
    """Top-n genes on the '+' (largest positive) or '-' (largest negative) pole."""
    order = np.argsort(-loading) if pole == "+" else np.argsort(loading)
    return [str(g) for g in genes[order[:n]]]

# Enrichr over-representation for every (cell type, SGP, pole)
rows = []
for (c, k), loading in sgp_loadings.items():
    for pole in ("+", "-"):
        try:
            df = gp.enrichr(gene_list=pole_genes(loading, pole), gene_sets=[LIB],
                            organism="human", outdir=None).results.copy()
        except Exception as ex:
            print(f"  [skip] {c} SGP{k}{pole}: {ex}"); continue
        df.insert(0, "celltype", c); df.insert(1, "SGP", k); df.insert(2, "pole", pole)
        rows.append(df)
        t = df.sort_values("Adjusted P-value").iloc[0]
        print(f"{c} SGP{k}{pole}: {t['Term'][:55]} (padj={t['Adjusted P-value']:.1e})")
        time.sleep(0.4)
enrich = pd.concat(rows, ignore_index=True)
enrich["neglog10_padj"] = -np.log10(enrich["Adjusted P-value"].clip(lower=1e-300))
enrich.to_csv(os.path.join(ENR_DIR, "sgp_enrichment_bipolar.csv"), index=False)
print("saved:", os.path.join(ENR_DIR, "sgp_enrichment_bipolar.csv"), "| rows:", len(enrich))

# ---- GSEA prerank on each SGP's SIGNED loading (positive at top of the ranking) ----
pre_rows = []
for (c, k), loading in sgp_loadings.items():
    rnk = pd.DataFrame({"gene": genes, "score": loading})
    rnk = rnk.reindex(rnk["score"].abs().sort_values(ascending=False).index).drop_duplicates("gene").sort_values("score", ascending=False)
    try:
        pre = gp.prerank(rnk=rnk, gene_sets=LIB, permutation_num=N_PERM, min_size=10, max_size=500,
                         seed=1, threads=8, no_plot=True, outdir=None, verbose=False)
    except Exception as ex:
        print(f"  [skip] {c} SGP{k}: {ex}"); continue
    df = pre.res2d.copy(); df["celltype"] = c; df["SGP"] = k
    for col in ["NES", "FDR q-val", "NOM p-val"]:
        df[col] = pd.to_numeric(df[col], errors="coerce")
    pre_rows.append(df)
    s = df[df["FDR q-val"] < FDR_CUTOFF]
    pos = s[s.NES > 0].sort_values("NES").tail(1); neg = s[s.NES < 0].sort_values("NES").head(1)
    msg = ([f"[-]{neg.iloc[0]['Term'][:40]}"] if len(neg) else []) + ([f"[+]{pos.iloc[0]['Term'][:40]}"] if len(pos) else [])
    print(f"{c} SGP{k}: " + ("  ".join(msg) if msg else "(none<FDR)"))
prerank = pd.concat(pre_rows, ignore_index=True)
prerank.to_csv(os.path.join(ENR_DIR, "sgp_prerank.csv"), index=False)
print("saved:", os.path.join(ENR_DIR, "sgp_prerank.csv"), "| rows:", len(prerank))

Part 3 — figures

Plotting code taken verbatim from mouse_brain_comment1.ipynb, restricted to the cell types analysed in revision/vis/1.1/mouse_brain_k=10.ipynb. Inputs are the precomputed results from Parts 1–2 (sync_selection_result.pkl, adata_sync_fit.h5ad, enrichment/*.csv). Figures are saved to the OUT folder.

import os, pickle
import numpy as np
import pandas as pd
import scanpy as sc
import matplotlib.pyplot as plt
import matplotlib as mpl

RES = "/import/home2/share/yqzeng/MacSGP/Results/Mouth_brain_visium"
OUT = "/import/home2/share/yqzeng/MacSGP/revision_2/reproduce"; os.makedirs(OUT, exist_ok=True)
KEEP_BLUE, STOP_RED, BASE, GRIDC = "#2f6bb5", "#c93b3b", "#8a8f96", "#e5e7eb"

# cell types analysed in revision/vis/1.1 (mouse_brain_k=10) — all have K>0
ANALYSED = ["Ext_Amy_2", "Ext_Hpc_CA1", "Ext_Hpc_CA3", "Ext_Hpc_DG1", "Ext_L23", "Ext_L5_1",
            "Ext_L5_2", "Ext_Med", "Ext_Pir", "Ext_Thal_1", "Ext_Thal_2", "Inh_1", "Inh_4"]

res = pickle.load(open(os.path.join(RES, "sync_selection_result.pkl"), "rb"))
Kc = res["recommended_K_per_celltype"]; pr = res["per_round_dLL"]
adata = sc.read_h5ad(os.path.join(RES, "adata_sync_fit.h5ad"))
enrich = pd.read_csv(os.path.join(RES, "enrichment/sgp_enrichment_bipolar.csv"))
prerank = pd.read_csv(os.path.join(RES, "enrichment/sgp_prerank.csv"))

cts = [c for c in ANALYSED if Kc.get(c, 0) > 0]      # analysed cell types (all K>0)
print("cell types:", {c: int(Kc[c]) for c in cts}, "| total SGPs", sum(Kc[c] for c in cts))

# publication style (Nature Communications) — legible, larger fonts
mpl.rcParams.update({
    'font.family': 'sans-serif',
    'font.sans-serif': ['Arial', 'Helvetica', 'DejaVu Sans'],
    'font.size': 13, 'axes.titlesize': 15, 'axes.labelsize': 14,
    'xtick.labelsize': 12, 'ytick.labelsize': 12, 'legend.fontsize': 12,
    'figure.titlesize': 17, 'axes.linewidth': 0.9,
    'xtick.major.width': 0.9, 'ytick.major.width': 0.9,
    'xtick.major.size': 3.5, 'ytick.major.size': 3.5,
    'axes.spines.top': False, 'axes.spines.right': False,
    'pdf.fonttype': 42, 'ps.fonttype': 42, 'savefig.dpi': 350, 'figure.dpi': 130,
})

1. Cumulative held-out ΔLL vs K

Each panel: cumulative held-out log-likelihood gain vs number of SGPs (mean ± SE over the 5 held-out replicates); the red dashed line marks the selected K*.

ncols = 4; nrows = int(np.ceil(len(cts) / ncols))
fig, axes = plt.subplots(nrows, ncols, figsize=(3.7 * ncols, 3.0 * nrows), squeeze=False)
for i, c in enumerate(cts):
    ax = axes[i // ncols][i % ncols]
    M = pr[c]
    cum = np.cumsum(M, axis=0); cum = np.vstack([np.zeros((1, cum.shape[1])), cum])
    xs = np.arange(cum.shape[0]); mu = cum.mean(1); se = cum.std(1, ddof=1) / np.sqrt(cum.shape[1])
    ax.fill_between(xs, mu - se, mu + se, color="#c9dbee", alpha=0.7, zorder=1)
    ax.plot(xs, mu, color=KEEP_BLUE, lw=2.4, marker="o", ms=7, mec="white", mew=1.2, zorder=2)
    ax.axvline(Kc[c], color=STOP_RED, ls="--", lw=1.6); ax.axhline(0, color=BASE, ls=":", lw=1.0)
    ax.set_xticks(xs.astype(int)); ax.set_title(f"{c}   K*={Kc[c]}", loc="left", fontweight="600")
    if i % ncols == 0: ax.set_ylabel("Cumulative ΔLL")
    ax.set_xlabel("K"); ax.grid(True, axis="y", color=GRIDC, lw=0.6, zorder=0)
for j in range(len(cts), nrows * ncols):
    axes[j // ncols][j % ncols].axis("off")
fig.tight_layout(); fig.savefig(os.path.join(OUT, "mouse_1_heldout_cumulative.pdf"), bbox_inches="tight")
fig.savefig(os.path.join(OUT, "mouse_1_heldout_cumulative.png"), bbox_inches="tight")
plt.show()

2. Spatial distribution of all selected SGP factors

One figure per cell type (name on top). First panel = cell-type proportion (Reds); following panels = each selected SGP factor (diverging coolwarm, centred at 0), restricted to spots where the cell type is present (proportion ≥ threshold; other spots grey).

import re as _re
coords = np.asarray(adata.obsm["spatial"]).astype(float)
sx, sy = coords[:, 0], -coords[:, 1]                 # flip y (image convention)
prop_df = adata.obsm["proportion"]
PROP_THRESH = 0.10
S_VIS, S_MASK = 10, 6

for c in cts:                                        # one separate figure per cell type
    K = Kc[c]
    prop = np.asarray(prop_df[c].values, float)
    mask = prop >= PROP_THRESH
    fac = np.asarray(adata.obsm[c].values)
    ncols = 1 + K                                    # proportion + one column per SGP

    fig, axes = plt.subplots(1, ncols, figsize=(3.4 * ncols, 3.9), squeeze=False)
    axes = axes[0]

    # ---- column 0: cell-type proportion (Reds) ----
    ax0 = axes[0]
    sca0 = ax0.scatter(sx, sy, c=prop, s=S_VIS, cmap="Reds",
                       vmin=0, vmax=float(prop.max()), rasterized=True, linewidths=0)
    cb0 = fig.colorbar(sca0, ax=ax0, orientation="horizontal", fraction=0.045, pad=0.03, shrink=0.7)
    cb0.ax.tick_params(labelsize=13)
    ax0.set_title("Proportion", fontsize=19)
    ax0.set_xticks([]); ax0.set_yticks([]); ax0.set_aspect("equal", adjustable="box")
    for sp in ax0.spines.values(): sp.set_visible(True); sp.set_linewidth(0.9)

    # ---- columns 1..K: SGP factors (coolwarm, diverging about 0) ----
    for k in range(K):
        ax = axes[k + 1]
        vv = fac[:, k][mask]
        fin = vv[np.isfinite(vv)]
        lo = float(fin.min()) if fin.size else -1.0
        hi = float(fin.max()) if fin.size else 1.0
        norm = mpl.colors.TwoSlopeNorm(vmin=min(lo, -1e-9), vcenter=0.0, vmax=max(hi, 1e-9))
        ax.scatter(sx[~mask], sy[~mask], c="lightgray", s=S_MASK, alpha=0.4, rasterized=True, linewidths=0)
        sca = ax.scatter(sx[mask], sy[mask], c=vv, s=S_VIS, cmap="coolwarm",
                         norm=norm, rasterized=True, linewidths=0)
        cb = fig.colorbar(sca, ax=ax, orientation="horizontal", fraction=0.045, pad=0.03, shrink=0.7)
        cb.set_ticks([])
        cb.ax.text(0.0, -0.9, "Low", ha="left", va="top", fontsize=13, transform=cb.ax.transAxes)
        cb.ax.text(1.0, -0.9, "High", ha="right", va="top", fontsize=13, transform=cb.ax.transAxes)
        ax.set_title(f"SGP{k}", fontsize=19)
        ax.set_xticks([]); ax.set_yticks([]); ax.set_aspect("equal", adjustable="box")
        for sp in ax.spines.values(): sp.set_visible(True); sp.set_linewidth(0.9)

    fig.suptitle(c, fontsize=24, fontweight="600", y=1.04)
    plt.subplots_adjust(wspace=0.05)
    fig.tight_layout()
    tag = _re.sub(r"[^0-9A-Za-z]+", "_", c)
    fig.savefig(os.path.join(OUT, f"mouse_2_factor_spatial_{tag}.pdf"), bbox_inches="tight")
    fig.savefig(os.path.join(OUT, f"mouse_2_factor_spatial_{tag}.png"), bbox_inches="tight")
    plt.show()

3. Pathway enrichment of the SGPs

Directional GSEA (prerank on the signed loading, GO-BP); colour = NES (+ pole / − pole), size = \(-\log_{10}\) FDR. Mouse gene symbols were upper-cased to approximate human orthologs.

# directional GSEA prerank dot plot (GO Biological Process)
FDR = 0.25; CLIP_NLQ = 6; TOP_PER_COL_G = 4; TRUNC = 55
pr2 = prerank.copy(); pr2["col"] = pr2.celltype + "·SGP" + pr2.SGP.astype(str)
pr2["FDR q-val"] = pd.to_numeric(pr2["FDR q-val"], errors="coerce"); pr2["NES"] = pd.to_numeric(pr2["NES"], errors="coerce")
pr2["nlq"] = -np.log10(pr2["FDR q-val"].clip(lower=1e-6))
sig = pr2[pr2["FDR q-val"] < FDR]
keep_t = sig.reindex(sig["NES"].abs().sort_values(ascending=False).index).groupby("col").head(TOP_PER_COL_G)["Term"].unique()
d = pr2[pr2.Term.isin(keep_t) & (pr2["FDR q-val"] < FDR)].copy()
if len(d):
    cols = sorted(d.col.unique()); to = d.groupby("Term").NES.apply(lambda s: s.abs().max()).sort_values().index.tolist()
    cix = {c: i for i, c in enumerate(cols)}; tix = {t: i for i, t in enumerate(to)}
    vmax = float(np.nanmax(np.abs(d.NES)))
    fig, ax = plt.subplots(figsize=(0.55 * len(cols) + 5, 0.42 * len(to) + 2.5))
    sc_ = ax.scatter(d.col.map(cix), d.Term.map(tix), s=40 + np.clip(d.nlq, 0, CLIP_NLQ) * 34,
                     c=d.NES, cmap="RdBu_r", vmin=-vmax, vmax=vmax, edgecolor="#444", lw=0.4)
    ax.set_xticks(range(len(cols))); ax.set_xticklabels(cols, rotation=90, fontsize=10)
    ax.set_yticks(range(len(to))); ax.set_yticklabels([t[:TRUNC] for t in to], fontsize=11)
    ax.set_xlim(-.5, len(cols) - .5); ax.set_ylim(-.5, len(to) - .5); ax.grid(True, color="#eee", lw=.5); ax.set_axisbelow(True)
    ax.set_title("Mouse brain — SGP GSEA prerank (GO-BP)\ncolor = NES (+/- pole),  size = -log10 FDR", loc="left", fontsize=14, fontweight="600")
    cb = fig.colorbar(sc_, ax=ax, shrink=0.5); cb.set_label("NES (+ = positive pole)", fontsize=12); cb.ax.tick_params(labelsize=11)
    fig.tight_layout(); fig.savefig(os.path.join(OUT, "mouse_3b_enrichment_prerank.pdf"), bbox_inches="tight")
    fig.savefig(os.path.join(OUT, "mouse_3b_enrichment_prerank.png"), bbox_inches="tight"); plt.show()
else:
    print("no prerank terms below FDR", FDR)