CRC (Visium HD 16 µm): per-cell-type K selection (reproduce)

Part 1 — compute: synchronized per-cell-type K selection, then fit at the selected K

Runs select_K_synchronized on the FULL Visium HD data (137k spots × 2780 genes) using gene patching + AMP (so the [C, N, G] rate never materialises at full width) and 5 replicates in parallel (one persistent worker per GPU, per-round barrier). A cell type keeps its K-th SGP only if the marginal held-out ΔLL is significantly positive across replicates (t-test, p < ALPHA) and its loading is reproducible across replicates (cross-replicate |cos| ≥ STAB_THRESH). Cell types are pre-filtered by proportion as in crc_16um_additive.py.

Consolidated from crc_16um_sync_selection.py (selection) and crc_16um_sync_fit.py (fit).

import os
GPUS = [1, 2, 3, 4, 5]                       # one persistent worker per replicate
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(g) for g in GPUS)

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

import numpy as np
import pandas as pd
import anndata as ad

from MacSGP_unified.k_selection import select_K_synchronized

DATA_PATH = "/home/yzengbj/factor/data/CRC"
SAVE_PATH = "/import/home2/share/yqzeng/MacSGP/results/CRC"
os.makedirs(SAVE_PATH, exist_ok=True)

# ---- model / training config (aligned with crc_16um_additive.py) ----
SGP_STEPS = 3000
LR = 2e-3
N_LAYERS = 3
COEF_REG = 0.6
PATCH_SIZE = 800          # gene patch width; smaller = less GPU memory, more time
USE_AMP = True

# ---- K-selection config ----
K_MAX = 8
N_REPEATS = 5
HOLDOUT_RATIO = 0.10
SEEDS = [20260701 + i for i in range(N_REPEATS)]
ALPHA = 0.05              # t-test significance to keep an SGP
STAB_THRESH = 0.8         # cross-replicate loading |cos| gate
MIN_EFFECT = 0.0

# ---- cell-type filtering (identical to crc_16um_additive.py) ----
PROP_THRESH = 0.1
CT_THRESH = 100

adata_st = ad.read_h5ad(os.path.join(DATA_PATH, "adata_hd_016.h5ad"))
adata_basis = ad.read_h5ad(os.path.join(DATA_PATH, "adata_basis_016.h5ad"))
print(f"ST: {adata_st.shape} | basis: {adata_basis.shape}")

# keep only cell types that are actually present across enough spots
pdf = adata_st.obsm["proportion"].copy()
pdf[pdf < PROP_THRESH] = 0
keep_ct = pdf.columns[pdf.sum(axis=0) > CT_THRESH]
adata_st.obsm["proportion"] = adata_st.obsm["proportion"][keep_ct]
adata_basis = adata_basis[list(keep_ct)].copy()
print(f"kept {len(keep_ct)}/{pdf.shape[1]} cell types: {list(keep_ct)}")

gpu_ids = list(range(len(GPUS)))          # re-indexed by CUDA_VISIBLE_DEVICES
t0 = time.time()
res = select_K_synchronized(
    adata_st,
    adata_basis,
    model_kwargs={"training_steps": SGP_STEPS, "lr": LR,
                  "n_layers": N_LAYERS, "coef_reg": COEF_REG},
    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,
    stab_thresh=STAB_THRESH,
    gene_patch=True,
    patch_size=PATCH_SIZE,
    use_amp=USE_AMP,
    gpus=gpu_ids,
    verbose=True,
    progress=True,
    progress_every=500,
)
print("elapsed %.1f min" % ((time.time() - t0) / 60))

with open(os.path.join(SAVE_PATH, "sync_selection_result.pkl"), "wb") as f:
    pickle.dump(res, f)
res["curve"].to_csv(os.path.join(SAVE_PATH, "sync_curve.csv"), index=False)
Kc = res["recommended_K_per_celltype"]
pd.DataFrame({"celltype": res["celltypes"],
              "K_selected": [Kc[c] for c in res["celltypes"]]}).to_csv(
    os.path.join(SAVE_PATH, "sync_recommended_K.csv"), index=False)

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 (crc_16um_sync_fit.py)

Refits the whole dataset with exactly K_c SGPs per cell type via train_greedy(fixed_Kc=..., gene_patch=True), storing factors (adata.obsm[celltype]) and loadings (adata.varm[celltype]). Runs on a single GPU (the first of the visible set). Set MIN_EFFECT_FRAC > 0 to re-select a smaller K per cell type from the saved curve without re-running the (expensive) selection.

import torch
from MacSGP_unified.model import Model

SEL_PKL = os.path.join(SAVE_PATH, "sync_selection_result.pkl")

# ---- K source ----
# MIN_EFFECT_FRAC = 0 -> use the recommended K as-is.
# > 0 -> re-select fewer SGPs from the saved curve (relative effect-size gate).
MIN_EFFECT_FRAC = 0.0


def reselect_from_curve(curve, alpha, stab_thresh, frac):
    """Re-derive K per cell type from a saved selection curve, adding a relative
    effect-size floor. Sequential: stop at the first SGP failing any gate."""
    out = {}
    for c in curve["celltype"].unique():
        sub = curve[curve["celltype"] == c].sort_values("K")
        first = float(sub.iloc[0]["dLL_mean"]) if len(sub) else 0.0
        K = 0
        for _, row in sub.iterrows():
            sig = (row["p"] < alpha) if pd.notna(row.get("p")) else (row["dLL_mean"] - row["dLL_se"] > 0)
            stbl = row.get("stability", 1.0) >= stab_thresh
            rel = (int(row["K"]) == 1) or (frac <= 0) or (first > 0 and row["dLL_mean"] >= frac * first)
            if sig and stbl and rel:
                K = int(row["K"])
            else:
                break
        out[c] = K
    return out

# re-read + re-filter exactly as the selection run (keeps this step self-contained)
adata_st = ad.read_h5ad(os.path.join(DATA_PATH, "adata_hd_016.h5ad"))
adata_basis = ad.read_h5ad(os.path.join(DATA_PATH, "adata_basis_016.h5ad"))

pdf = adata_st.obsm["proportion"].copy()
pdf[pdf < PROP_THRESH] = 0
keep_ct = pdf.columns[pdf.sum(axis=0) > CT_THRESH]
adata_st.obsm["proportion"] = adata_st.obsm["proportion"][keep_ct]
adata_basis = adata_basis[list(keep_ct)].copy()
celltypes = list(adata_basis.obs.index)

res = pickle.load(open(SEL_PKL, "rb"))
if MIN_EFFECT_FRAC > 0:
    Kc = reselect_from_curve(res["curve"], ALPHA, STAB_THRESH, MIN_EFFECT_FRAC)
    print(f"re-selected K with min_effect_frac={MIN_EFFECT_FRAC}")
else:
    Kc = res["recommended_K_per_celltype"]

Kc_list = [int(Kc.get(c, 0)) for c in celltypes]
maxK = max(max(Kc_list), 1)
print("fit K per cell type:")
print(pd.Series(dict(zip(celltypes, Kc_list))).sort_values(ascending=False).to_string())
print("maxK =", maxK)

t0 = time.time()
model = Model(adata_st, adata_basis, n_SGPs=maxK, seed=1234,
              training_steps=SGP_STEPS, lr=LR, n_layers=N_LAYERS, coef_reg=COEF_REG)
model.train_greedy(fixed_Kc=Kc_list, fix_gamma_k=True,
                   gene_patch=True, patch_size=PATCH_SIZE, use_amp=USE_AMP)
adata_res = model.eval(sort_sgps=False)
print("fit done in %.1f min" % ((time.time() - t0) / 60))

adata_res.write_h5ad(os.path.join(SAVE_PATH, "adata_sync_fit.h5ad"))
torch.save(model, os.path.join(SAVE_PATH, "model_sync_fit.pt"))
print("saved -> adata_sync_fit.h5ad, model_sync_fit.pt")

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

Restricted to the cell types analysed in revision/vis/1.1/crc_k=10.ipynb (incl. K=1). Human gene symbols; enrichment uses MSigDB Hallmark. Both pole-separated over-representation (top-TOP_N genes of each pole via Enrichr) and directional GSEA (prerank on the signed loading) are produced. Requires internet (Enrichr) + gseapy.

Consolidated from crc_sgp_enrichment.ipynb.

import scanpy as sc
import gseapy as gp

ORGANISM = "human"
LIB = "MSigDB_Hallmark_2020"
KEEP = ["CAF", "Enterocyte", "Goblet", "Macrophage", "Neutrophil", "Plasma",
        "Proliferating Fibroblast", "Tumor III", "Tumor V", "vSM"]   # analysed in revision/vis/1.1 (incl. K=1)
TOP_N = 100; PADJ_CUTOFF = 0.05; N_PERM = 1000; FDR_CUTOFF = 0.25

FIT_H5AD = os.path.join(SAVE_PATH, "adata_sync_fit.h5ad")
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])

# (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})

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'][:50]} (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=5, 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 crc_comment1.ipynb, restricted to the cell types analysed in revision/vis/1.1/crc_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/CRC"
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 (crc_k=10) — all have K>0
ANALYSED = ["CAF", "Enterocyte", "Goblet", "Macrophage", "Neutrophil", "Plasma",
            "Proliferating Fibroblast", "Tumor III", "Tumor V", "vSM"]

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

# analysed cell types, K>0 and present in the fit (only the first Kc columns are trained)
cts = [c for c in ANALYSED if Kc.get(c, 0) > 0 and c in adata.obsm
       and c in adata.obsm["proportion"].columns]
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.1 * 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", fontsize=13, 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, "crc_1_heldout_cumulative.pdf"), bbox_inches="tight")
fig.savefig(os.path.join(OUT, "crc_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 = 4, 2                                  # small markers (137k HD spots)

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"crc_2_factor_spatial_{tag}.pdf"), bbox_inches="tight")
    fig.savefig(os.path.join(OUT, f"crc_2_factor_spatial_{tag}.png"), bbox_inches="tight")
    plt.show()

3. Pathway enrichment of the SGPs

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

# directional GSEA prerank dot plot (MSigDB Hallmark)
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()
    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.5 * 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=9)
    ax.set_yticks(range(len(to))); ax.set_yticklabels(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("CRC — SGP GSEA prerank (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, "crc_3b_enrichment_prerank.pdf"), bbox_inches="tight")
    fig.savefig(os.path.join(OUT, "crc_3b_enrichment_prerank.png"), bbox_inches="tight"); plt.show()
else:
    print("no prerank terms below FDR", FDR)