MouseBrain Microglia stGP Tutorial

This notebook follows the MouseBrain MERFISH Microglia analysis step by step: load the processed cells, build the temporal/spatial kernels, fit or reuse stGP, write Results/stgp/Microglia, and regenerate source figures under Figures/Microglia.

import json
import os
import pickle
import sys
import time
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import numpy as np
import pandas as pd
import scanpy as sc
import scipy.sparse as sp
from IPython.display import display

PROJECT_DIR = Path.cwd()
if PROJECT_DIR.name != "RealData_MouseBrainMERFISH":
    PROJECT_DIR = Path("/home/byual/stGP-0529/RealData_MouseBrainMERFISH")
os.chdir(PROJECT_DIR)
sys.path.insert(0, str(PROJECT_DIR))
sys.path.insert(0, str(PROJECT_DIR.parent))

from stgp.estimation import fit_pfactor, fit_pfactor_auto
from stgp.kernels import (
    bandwidth_select_spatial,
    bandwidth_select_temporal,
    build_K_age,
    build_K_spa_list_from_stacked,
)
from stgp.preprocessing import log1p_norm_centered_list, standardize_coords_list

import utils as u
from plots import set_nature_style

set_nature_style()
print(f"Working directory: {PROJECT_DIR}")
print(f"Results root: {(PROJECT_DIR / u.RESULTS_STGP).resolve()}")
print(f"Figures root: {(PROJECT_DIR / u.FIGURES_ROOT).resolve()}")

Working directory: /import/home4/byual/stGP-0529/RealData_MouseBrainMERFISH
Results root: /import/home4/byual/stGP-0529/RealData_MouseBrainMERFISH/Results/stgp
Figures root: /import/home4/byual/stGP-0529/RealData_MouseBrainMERFISH/Figures

# Tutorial parameters.
CELLTYPE = "Microglia"
N_PROGRAMS = 4       # Fig. 5 fixes p=4 for matched method benchmarking.
P_MAX = 10           # Used only when N_PROGRAMS is None.
K_NEIGHBORS = 15
SEED = 0
SKIP_EXISTING = True

safe_ct = u.safe_name(CELLTYPE)
input_h5ad = u.DATA_PROCESSED / f"{safe_ct}.h5ad"
stgp_dir = u.RESULTS_STGP / safe_ct
fig_dir = u.FIGURES_ROOT / safe_ct
stgp_pkl = stgp_dir / "stgp_result.pkl"
scored_h5ad = stgp_dir / "adata_with_scores.h5ad"
W_csv = stgp_dir / "W.csv"

run_stgp = not (SKIP_EXISTING and stgp_pkl.exists() and scored_h5ad.exists() and W_csv.exists())

display(pd.DataFrame([{
    "celltype": CELLTYPE,
    "input_h5ad": str(input_h5ad),
    "stgp_dir": str(stgp_dir),
    "figure_dir": str(fig_dir),
    "fit_mode": "fixed p" if N_PROGRAMS is not None else "auto rank",
    "p": N_PROGRAMS,
    "run_stgp_now": run_stgp,
}]))

	celltype	input_h5ad	stgp_dir	figure_dir	fit_mode	p	run_stgp_now
0	Microglia	data/processed/Microglia.h5ad	Results/stgp/Microglia	Figures/Microglia	fixed p	4	False

1. Load cells and organize slices

stGP is fit to one target cell type. Cells are grouped by mouse_id; each mouse contributes one age and one spatial kernel block.

if run_stgp:
    if not input_h5ad.exists():
        raise FileNotFoundError(f"Missing processed input: {input_h5ad}")

    adata = sc.read_h5ad(str(input_h5ad))
    age_arr = adata.obs["age"].to_numpy(dtype=float)
    groups = adata.obs["mouse_id"].astype(str).to_numpy()
    uniq_mice, inv = np.unique(groups, return_inverse=True)
    idx_per_group = [np.sort(np.where(inv == t)[0]) for t in range(len(uniq_mice))]
    nlist = np.array([len(ix) for ix in idx_per_group])
    ages = np.array([age_arr[ix[0]] for ix in idx_per_group])

    summary = pd.DataFrame({"mouse_id": uniq_mice, "age": ages, "n_cells": nlist})
    print(f"Loaded {adata.n_obs:,} {CELLTYPE} cells and {adata.n_vars:,} genes.")
    display(summary.head())
    display(summary.describe(include="all"))
else:
    print(f"Using cached stGP outputs in {stgp_dir}")
    adata_scores = sc.read_h5ad(str(scored_h5ad))
    print(f"Cached scored AnnData: {adata_scores.n_obs:,} cells x {adata_scores.n_vars:,} genes")

Using cached stGP outputs in Results/stgp/Microglia
Cached scored AnnData: 52,417 cells x 220 genes

2. Build expression, temporal, and spatial inputs

The expression matrix is split by mouse, log-normalized/centered, and paired with an age kernel plus one spatial kernel per mouse slice.

if run_stgp:
    X_raw = adata.X.toarray() if sp.issparse(adata.X) else np.asarray(adata.X)
    Y_list, _ = log1p_norm_centered_list([X_raw[ix] for ix in idx_per_group])
    coords_list = standardize_coords_list([adata.obsm["spatial"][ix] for ix in idx_per_group])

    gamma_spa = bandwidth_select_spatial(coords_list, frac=0.01, rho=0.75)
    gamma_age = bandwidth_select_temporal(ages, rho=np.exp(-4))
    K_age = build_K_age(ages, gamma_age, kernel="rbf", standardize=True)
    K_spa_list = build_K_spa_list_from_stacked(
        np.vstack(coords_list), nlist, gamma_spa, standardize=False, jitter=1e-6,
    )

    print(f"gamma_spa = {gamma_spa:.4f}")
    print(f"gamma_age = {gamma_age:.4f}")
    print(f"K_age shape: {K_age.shape}")
    print(f"Spatial kernel blocks: {len(K_spa_list)}")
else:
    print("Skipped kernel construction because cached stGP outputs are present.")

Skipped kernel construction because cached stGP outputs are present.

3. Fit stGP

N_PROGRAMS=4 uses the fixed-rank fit_pfactor mode used for Fig. 5 benchmarking. Set N_PROGRAMS=None to use fit_pfactor_auto.

if run_stgp:
    t0 = time.perf_counter()
    if N_PROGRAMS is not None:
        stgp_res = fit_pfactor(
            Y_list, nlist, K_age, K_spa_list,
            p=N_PROGRAMS, k=K_NEIGHBORS, verbose=1,
        )
    else:
        stgp_res = fit_pfactor_auto(
            Y_list=Y_list, Nlist=nlist, K_age=K_age, Kspa_list=K_spa_list,
            p_max=P_MAX, k=K_NEIGHBORS, random_state=SEED, verbose=1,
        )
    elapsed = time.perf_counter() - t0
    print(f"stGP runtime: {elapsed:.1f} sec")
    print(f"Selected programs: {stgp_res['W'].shape[0]}")
else:
    with open(stgp_pkl, "rb") as f:
        stgp_res = pickle.load(f)
    elapsed = None
    print(f"Loaded cached stGP pickle: {stgp_pkl}")
    print(f"Cached programs: {np.asarray(stgp_res['W']).shape[0]}")

Loaded cached stGP pickle: Results/stgp/Microglia/stgp_result.pkl
Cached programs: 4

4. Write stGP outputs

Scores H and spatial residuals b are mapped back to original cell order and saved in obsm['X_stgp'] and obsm['X_stgp_spatial'].

if run_stgp:
    stgp_dir.mkdir(parents=True, exist_ok=True)
    stgp_res["gamma_age"] = gamma_age
    stgp_res["gamma_spa"] = gamma_spa
    with open(stgp_pkl, "wb") as f:
        pickle.dump(stgp_res, f)

    p_sel = stgp_res["W"].shape[0]
    prog_labels = [f"stGP{j + 1}" for j in range(p_sel)]
    W = pd.DataFrame(stgp_res["W"], index=prog_labels, columns=adata.var_names)
    W.to_csv(W_csv)

    all_idx = np.concatenate(idx_per_group)
    H_adata = np.empty_like(stgp_res["H"])
    H_adata[all_idx] = stgp_res["H"]
    b_adata = np.empty_like(stgp_res["b"])
    b_adata[all_idx] = stgp_res["b"]
    adata.obsm["X_stgp"] = H_adata.astype(np.float32)
    adata.obsm["X_stgp_spatial"] = b_adata.astype(np.float32)

    alpha_arr = np.asarray(stgp_res.get("alpha", []))
    alpha_lower_arr = np.asarray(stgp_res.get("alpha_lower", []))
    alpha_upper_arr = np.asarray(stgp_res.get("alpha_upper", []))
    adata.uns["stgp"] = dict(
        groups=uniq_mice.tolist(), ages=ages.tolist(),
        gamma_age=gamma_age, gamma_spa=gamma_spa, p_selected=p_sel,
        alpha=alpha_arr.tolist() if alpha_arr.ndim == 2 else [],
        alpha_lower=alpha_lower_arr.tolist() if alpha_lower_arr.ndim == 2 else [],
        alpha_upper=alpha_upper_arr.tolist() if alpha_upper_arr.ndim == 2 else [],
    )
    adata.write_h5ad(scored_h5ad, compression="gzip")

    timing = {
        "method": "stgp", "celltype": CELLTYPE,
        "runtime_sec": round(elapsed, 2), "status": "completed",
        "p": N_PROGRAMS, "p_max": P_MAX, "k": K_NEIGHBORS, "seed": SEED,
    }
    (stgp_dir / "timing.json").write_text(json.dumps(timing, indent=2))
    u.append_jsonl(u.TIMING_LOG, {**timing, "out_dir": str(stgp_dir.resolve())})
else:
    W = pd.read_csv(W_csv, index_col=0)

5. Load stGP and baseline results for figure generation

Missing baseline directories are skipped. Available methods are summarized before plotting.

fig_dir.mkdir(parents=True, exist_ok=True)
methods = u._load_methods(safe_ct, CELLTYPE, stgp_dir=stgp_dir)
if not methods:
    raise RuntimeError(f"No method results found for {CELLTYPE}")

stgp_method = next((m for m in methods if m.method == "stGP"), None)
stgp_res = u._load_stgp_pickle(safe_ct, stgp_dir=stgp_dir) if stgp_method is not None else None

stamp_m = next((m for m in methods if m.method == "STAMP"), None)
if stamp_m is not None and stamp_m.gene_weights is None:
    stamp_m.gene_weights = u._infer_gene_weights(stamp_m.adata, stamp_m.scores)

method_table = pd.DataFrame([
    {
        "method": m.method,
        "cells": m.adata.n_obs,
        "programs": m.scores.shape[1],
        "has_gene_weights": m.gene_weights is not None,
        "result_dir": str(m.result_dir),
    }
    for m in methods
])

  [loaded] stGP: (52417, 4)
  [loaded] SpatialPCA: (52417, 4)
  [loaded] MEFISTO: (52417, 4)
  [loaded] STAMP: (51003, 4)
  [loaded] Popari: (52417, 4)

	method	cells	programs	has_gene_weights	result_dir
0	stGP	52417	4	True	Results/stgp/Microglia
1	SpatialPCA	52417	4	True	Results/baselines/spatialpca/Microglia
2	MEFISTO	52417	4	True	Results/baselines/mefisto/Microglia
3	STAMP	51003	4	True	Results/baselines/stamp/Microglia
4	Popari	52417	4	True	Results/baselines/popari/Microglia

6. Temporal variance and age-related genes

These panels ask which stGP programs are age-dominated, how weighted expression changes with age, and which genes drive the age-associated programs.

stgp_vdf = None
if stgp_method is not None:
    stgp_vdf = u._section1_variance_decomposition(
        stgp_method, stgp_res, celltype=CELLTYPE, safe_ct=safe_ct, fig_dir=fig_dir,
    )
    u._section2_program_scores_by_age(
        stgp_method, stgp_vdf, celltype=CELLTYPE, safe_ct=safe_ct, fig_dir=fig_dir,
    )

7. Spatial program maps, domains, and method similarity

This section writes spatial maps, age-ordered stack plots, benchmark cluster comparisons, and stGP-vs-baseline gene-program similarity summaries.

spatial_outputs = sorted((fig_dir / "spatial_selected_2x2").rglob("*.png"))
sim_outputs = sorted((fig_dir / "program_similarity").rglob("*.csv"))

if spatial_outputs and sim_outputs:
    print("Using cached spatial maps and program-similarity tables.")
else:
    u._section4_spatial_maps_and_clustering(
        methods, stgp_method, celltype=CELLTYPE, safe_ct=safe_ct, fig_dir=fig_dir,
    )
    u._section5_program_similarity(
        methods, stgp_method, celltype=CELLTYPE, safe_ct=safe_ct, fig_dir=fig_dir,
    )
    spatial_outputs = sorted((fig_dir / "spatial_selected_2x2").rglob("*.png"))
    sim_outputs = sorted((fig_dir / "program_similarity").rglob("*.csv"))

8. Kernel diagnostic, W heatmap, and runtime comparison

These final upstream figures document the spatial kernel scale, active gene-weight matrix, and method runtimes.

if stgp_method is not None:
    u._section6_kernel_diagnostic_and_W_heatmap(
        stgp_method, stgp_res, celltype=CELLTYPE, safe_ct=safe_ct, fig_dir=fig_dir,
    )
u._section7_timing(celltype=CELLTYPE, safe_ct=safe_ct, fig_dir=fig_dir, stgp_dir=stgp_dir)

figure_files = sorted(fig_dir.rglob("*.png"))
table_files = sorted(fig_dir.rglob("*.csv"))