Kidney cancer: 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).

import os
# expose the GPUs used for the parallel replicates BEFORE importing torch.
GPUS = [0, 1, 2, 3, 4]              # one worker per replicate; set to 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/kidney_cancer"
SAVE_PATH = "/import/home2/share/yqzeng/MacSGP/results/kidney_cancer"
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"))
# sanity: the fields select_K_synchronized needs must already be present
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, "| celltypes:", list(adata_basis.obs.index))

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"
# stability gate: keep the K-th SGP only if its loading is reproducible across
# replicates (mean pairwise |cosine| >= STAB_THRESH). A real program is recovered
# consistently (~0.9); mask-specific noise is not (~0). Set 0 to disable.
STAB_THRESH = 0.8

# GPUS in the child processes are re-indexed 0..len(GPUS)-1 by CUDA_VISIBLE_DEVICES
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)
print("\nrecommended K per cell type:")
print(res["recommended_K_per_celltype"])

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

Kc = res["recommended_K_per_celltype"]
celltypes = list(adata_basis.obs.index)
Kc_list = [int(Kc[c]) for c in celltypes]
maxK = max(max(Kc_list), 1)
print("selected K:", dict(zip(celltypes, Kc_list)), "| maxK =", maxK)

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)

Each SGP has two anti-correlated poles (a factor/loading sign flip is arbitrary), so we analyse them separately: the top-TOP_N genes on the positive and on the negative pole of every (cell type, SGP), each tested against MSigDB Hallmark / GO-BP / KEGG via Enrichr. In addition, directional GSEA (prerank on the signed loading) assigns each pathway one signed NES. Requires internet (Enrichr) + gseapy.

import time
import gseapy as gp

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)
res = pickle.load(open(SEL_PKL, "rb"))
Kc = res["recommended_K_per_celltype"]
genes = np.asarray(adata.var_names)          # already HGNC symbols

# (cell type, SGP index) -> loading vector, only for kept SGPs (k < K_c)
sgp_loadings = {}
for c, K in Kc.items():
    if K < 1 or c not in adata.varm:
        continue
    L = np.asarray(adata.varm[c])            # [n_genes, maxK]
    for k in range(int(K)):
        sgp_loadings[(c, k)] = L[:, k]
print(f"{len(sgp_loadings)} SGPs; each analysed at + and - pole -> {2*len(sgp_loadings)} gene sets")

TOP_N = 100                       # top genes per pole taken as that pole's signature
LIBRARIES = ["MSigDB_Hallmark_2020", "GO_Biological_Process_2021", "KEGG_2021_Human"]
PADJ_CUTOFF = 0.05

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 ("+", "-"):
        gl = pole_genes(loading, pole)
        try:
            df = gp.enrichr(gene_list=gl, gene_sets=LIBRARIES, 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)
        top = df.sort_values("Adjusted P-value").iloc[0]
        print(f"{c} SGP{k}{pole}: [{top['Gene_set']}] {top['Term']}  (padj={top['Adjusted P-value']:.1e})")
        time.sleep(0.8)

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("\nsaved:", 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_LIB = "MSigDB_Hallmark_2020"
N_PERM = 1000
FDR_SHOW = 0.25

pre_rows = []
for (c, k), loading in sgp_loadings.items():
    rnk = pd.DataFrame({"gene": genes, "score": loading}).sort_values("score", ascending=False)
    try:
        pre = gp.prerank(rnk=rnk, gene_sets=PRE_LIB, permutation_num=N_PERM,
                         min_size=5, max_size=500, seed=1, threads=4,
                         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_SHOW]
    pos = s[s.NES > 0].sort_values("NES").tail(1)
    neg = s[s.NES < 0].sort_values("NES").head(1)
    msg = []
    if len(neg): msg.append(f"[- pole] {neg.iloc[0]['Term']} (NES {neg.iloc[0]['NES']:.2f})")
    if len(pos): msg.append(f"[+ pole] {pos.iloc[0]['Term']} (NES {pos.iloc[0]['NES']:.2f})")
    print(f"{c} SGP{k}: " + ("  ".join(msg) if msg else "(no term below FDR)"))
    time.sleep(0.2)

prerank = pd.concat(pre_rows, ignore_index=True)
prerank.to_csv(os.path.join(ENR_DIR, "sgp_prerank.csv"), index=False)
print("\nsaved:", os.path.join(ENR_DIR, "sgp_prerank.csv"), "| rows:", len(prerank))

Part 3 — figures

Plotting code taken verbatim from kidney_cancer_comment1.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/kidney_cancer"
OUT = "/import/home2/share/yqzeng/MacSGP/revision_2/reproduce"; os.makedirs(OUT, exist_ok=True)
mpl.rcParams.update({"figure.dpi": 110, "savefig.dpi": 200, "font.family": "DejaVu Sans",
    "axes.spines.top": False, "axes.spines.right": False})
KEEP_BLUE, STOP_RED, BASE, GRIDC = "#2f6bb5", "#c93b3b", "#8a8f96", "#e5e7eb"

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 res["celltypes"] if Kc.get(c, 0) > 0]      # K>0 cell types only
print("K>0 cell types:", {c: int(Kc[c]) for c in cts}, "| total SGPs", sum(Kc[c] for c in cts))

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

# ============================================================
# 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,       # embed editable TrueType text (Illustrator-friendly)
    'ps.fonttype': 42,
    'savefig.dpi': 350,
    'figure.dpi': 130,
})

1. Cumulative held-out ΔLL vs K (K>0 cell types)

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.suptitle("Kidney — cumulative held-out ΔLL vs K (dashed = selected K*)", x=0.02, ha="left", y=1.004, fontweight="600")
fig.tight_layout(); fig.savefig(os.path.join(OUT, "kidney_1_heldout_cumulative.pdf"), bbox_inches="tight")
fig.savefig(os.path.join(OUT, "kidney_1_heldout_cumulative.png"), bbox_inches="tight")
plt.show()

2. Spatial distribution of all selected SGP factors

One figure per cell type (name shown on top). The first panel is the cell-type proportion (Reds); the following panels are 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+1}", 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"kidney_2_factor_spatial_{tag}.pdf"), bbox_inches="tight")
    fig.savefig(os.path.join(OUT, f"kidney_2_factor_spatial_{tag}.png"), bbox_inches="tight")
    plt.show()

3. Pathway enrichment of the SGPs

directional GSEA (prerank on the signed loading); colour = NES (+ pole / − pole), size = \(-\log_{10}\) FDR.

# directional GSEA prerank dot plot
FDR = 0.25; CLIP_NLQ = 6; TOP_PER_COL_G = 4
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()
    cols_label = [s[:-1] + str(int(s[-1]) + 1) for s in cols]
    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.68 * len(cols) + 5, 0.46 * 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_label, rotation=90, fontsize=15)
    ax.set_yticks(range(len(to))); ax.set_yticklabels(to, fontsize=15)
    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("GSEAPreranked (Hallmark)\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, "kidney_3b_enrichment_prerank.pdf"), bbox_inches="tight")
    fig.savefig(os.path.join(OUT, "kidney_3b_enrichment_prerank.png"), bbox_inches="tight"); plt.show()
else:
    print("no prerank terms below FDR", FDR)