Oligodendrocyte Gene Programs in the Aging Human Brain
This tutorial walks through the full stGP pipeline on the human brain MERFISH dataset (Jeffries et al., Nature 2025), using Oligodendrocytes (oli) as the target cell type.
1. Setup
[1]:
import os, sys, warnings, pickle
import numpy as np
import pandas as pd
import scanpy as sc
import matplotlib.pyplot as plt
import scipy.sparse as sp
from pathlib import Path
warnings.filterwarnings("ignore", category=FutureWarning)
sys.path.insert(0, "..")
DATA_QC = Path("data/qc/human_merfish_qc.h5ad")
DATA_PROC = Path("data/processed")
RESULTS_DIR = Path("Results/stgp")
FIGURES_DIR = Path("Figure/oli")
FIGURES_DIR.mkdir(parents=True, exist_ok=True)
CELLTYPE = "oli" # target cell type for this tutorial
adata_oli = sc.read_h5ad(DATA_PROC / f"{CELLTYPE}.h5ad")
2. Fitting stGP
stGP models gene expression in each tissue slice as a sum of p latent programs. Each program is characterised by:
W (gene loadings,
p × G): non-negative gene weights defining the programH (cell scores,
N × p): overall activity of each program in each cellb (spatial field,
N × p): spatially smooth residual componentα (age effect,
p × S): how program amplitude varies across slices/agesθ (GP hyperparameters,
p × 2): spatial amplitude and noise fraction per program
The model rank p is selected automatically by greedy forward selection.
[2]:
import time
from stgp.estimation import 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 standardize_coords_list
OUT_DIR = RESULTS_DIR / CELLTYPE
OUT_DIR.mkdir(parents=True, exist_ok=True)
PKL_PATH = OUT_DIR / "stgp_result.pkl"
def balanced_normalize(Y_list, target_sum=None, eps=1e-12, max_scale=10.0):
Y_list = [np.asarray(Y, dtype=float) for Y in Y_list]
all_lib = np.concatenate([Y.sum(axis=1) for Y in Y_list])
if target_sum is None:
target_sum = np.median(all_lib[all_lib > 0])
X_list = []
for Y in Y_list:
lib = Y.sum(axis=1, keepdims=True)
scale = target_sum / np.maximum(lib, eps)
if max_scale is not None:
scale = np.minimum(scale, max_scale)
X_list.append(np.log1p(Y * scale))
gene_mean = np.mean([X.mean(axis=0) for X in X_list], axis=0)
X_list = [X - gene_mean for X in X_list]
return X_list, gene_mean, target_sum
age_arr = pd.to_numeric(adata_oli.obs["age"], errors="coerce").to_numpy(float)
groups = adata_oli.obs["id_region"].astype(str).to_numpy()
uniq, inv = np.unique(groups, return_inverse=True)
idx_per_group = [np.sort(np.where(inv == t)[0]) for t in range(len(uniq))]
X_raw = adata_oli.X.toarray() if sp.issparse(adata_oli.X) else np.asarray(adata_oli.X)
Y_raw_list = [X_raw[ix, :] for ix in idx_per_group]
Y_list, _, _ = balanced_normalize(Y_raw_list)
nlist = np.array([len(ix) for ix in idx_per_group])
ages = np.array([age_arr[ix[0]] for ix in idx_per_group])
slices = uniq.copy() # Match 02_run_stgp.py: keep np.unique(id_region) order.
[3]:
# ── Build GP kernels ─────────────────────────────────────────────────────
coords_list = standardize_coords_list([adata_oli.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(-2))
print(f" gamma_spa = {gamma_spa:.4f} | gamma_age = {gamma_age:.4f}")
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
)
gamma_spa = 0.1402 | gamma_age = 0.8247
[4]:
from plots import plot_spatial_kernel_corr_combined
fig = plot_spatial_kernel_corr_combined(
adata=adata_oli, bandwidth=float(gamma_spa),
slice_idx=10,
age_unit="years",
)
[5]:
t0 = time.perf_counter()
res = fit_pfactor_auto(
Y_list=Y_list,
Nlist=nlist,
K_age=K_age,
Kspa_list=K_spa_list,
p_max=10,
k=15,
inner_rank1_tol=1e-4,
rel_improve_total_tol=0.01,
backfit_tol=1e-4,
random_state=0,
verbose=1,
)
print(f"Runtime: {time.perf_counter() - t0:.1f}s | programs selected: {res['W'].shape[0]}")
# ── Save results ─────────────────────────────────────────────────────────
res["gamma_age"] = gamma_age
res["gamma_spa"] = gamma_spa
with open(PKL_PATH, "wb") as f:
pickle.dump(res, f)
print(f"Saved: {PKL_PATH}")
[sweep=001] dW_rel=2.075e-01 dTheta_rel=2.775e-02 time=7.170e+00
[sweep=002] dW_rel=1.278e-01 dTheta_rel=3.279e-02 time=7.083e+00
[sweep=003] dW_rel=4.877e-02 dTheta_rel=1.024e-02 time=5.501e+00
[sweep=004] dW_rel=2.982e-02 dTheta_rel=3.660e-03 time=4.010e+00
[sweep=005] dW_rel=2.183e-02 dTheta_rel=2.617e-03 time=3.047e+00
[sweep=006] dW_rel=1.685e-02 dTheta_rel=2.073e-03 time=3.054e+00
[sweep=007] dW_rel=1.336e-02 dTheta_rel=1.652e-03 time=2.526e+00
[sweep=008] dW_rel=1.091e-02 dTheta_rel=1.417e-03 time=2.461e+00
[sweep=009] dW_rel=9.049e-03 dTheta_rel=1.170e-03 time=2.243e+00
[sweep=010] dW_rel=7.649e-03 dTheta_rel=9.906e-04 time=2.376e+00
[sweep=011] dW_rel=6.501e-03 dTheta_rel=8.354e-04 time=1.978e+00
[sweep=012] dW_rel=5.586e-03 dTheta_rel=7.340e-04 time=2.207e+00
[sweep=013] dW_rel=4.824e-03 dTheta_rel=5.922e-04 time=1.981e+00
[sweep=014] dW_rel=4.212e-03 dTheta_rel=5.193e-04 time=2.072e+00
[sweep=015] dW_rel=3.694e-03 dTheta_rel=4.595e-04 time=1.773e+00
[sweep=016] dW_rel=3.234e-03 dTheta_rel=4.281e-04 time=1.757e+00
[sweep=017] dW_rel=2.866e-03 dTheta_rel=3.660e-04 time=1.856e+00
[sweep=018] dW_rel=2.538e-03 dTheta_rel=2.885e-04 time=1.655e+00
[sweep=019] dW_rel=2.225e-03 dTheta_rel=2.816e-04 time=1.349e+00
[sweep=020] dW_rel=2.052e-03 dTheta_rel=2.927e-04 time=1.759e+00
[sweep=021] dW_rel=1.801e-03 dTheta_rel=2.388e-04 time=1.343e+00
[sweep=022] dW_rel=1.615e-03 dTheta_rel=2.287e-04 time=1.396e+00
[sweep=023] dW_rel=1.441e-03 dTheta_rel=1.986e-04 time=1.500e+00
[sweep=024] dW_rel=1.312e-03 dTheta_rel=1.464e-04 time=1.305e+00
[sweep=025] dW_rel=1.160e-03 dTheta_rel=1.348e-04 time=1.372e+00
[sweep=026] dW_rel=1.075e-03 dTheta_rel=1.650e-04 time=1.253e+00
[sweep=027] dW_rel=9.498e-04 dTheta_rel=1.312e-04 time=1.247e+00
[sweep=028] dW_rel=8.510e-04 dTheta_rel=9.249e-05 time=1.187e+00
[sweep=029] dW_rel=7.772e-04 dTheta_rel=1.009e-04 time=1.290e+00
[sweep=030] dW_rel=6.918e-04 dTheta_rel=9.760e-05 time=1.180e+00
[sweep=031] dW_rel=6.705e-04 dTheta_rel=1.316e-04 time=1.155e+00
[sweep=032] dW_rel=5.642e-04 dTheta_rel=6.339e-05 time=9.716e-01
[sweep=033] dW_rel=5.199e-04 dTheta_rel=7.150e-05 time=1.383e+00
[sweep=034] dW_rel=4.661e-04 dTheta_rel=4.235e-05 time=9.610e-01
[sweep=035] dW_rel=4.277e-04 dTheta_rel=8.085e-05 time=1.054e+00
[sweep=036] dW_rel=3.734e-04 dTheta_rel=5.327e-05 time=9.899e-01
[sweep=037] dW_rel=3.609e-04 dTheta_rel=5.843e-05 time=1.199e+00
[sweep=038] dW_rel=3.086e-04 dTheta_rel=4.598e-05 time=1.051e+00
[sweep=039] dW_rel=2.979e-04 dTheta_rel=6.517e-05 time=9.668e-01
[sweep=040] dW_rel=2.442e-04 dTheta_rel=1.585e-05 time=8.364e-01
[sweep=041] dW_rel=2.505e-04 dTheta_rel=6.457e-05 time=9.434e-01
[sweep=042] dW_rel=2.153e-04 dTheta_rel=6.893e-05 time=1.021e+00
[sweep=043] dW_rel=1.813e-04 dTheta_rel=2.887e-05 time=9.115e-01
[sweep=044] dW_rel=2.189e-04 dTheta_rel=8.689e-05 time=1.194e+00
[sweep=045] dW_rel=1.763e-04 dTheta_rel=5.094e-05 time=9.351e-01
[sweep=046] dW_rel=1.620e-04 dTheta_rel=5.464e-05 time=1.053e+00
[sweep=047] dW_rel=1.409e-04 dTheta_rel=4.263e-05 time=8.794e-01
[sweep=048] dW_rel=1.453e-04 dTheta_rel=4.771e-05 time=8.704e-01
[sweep=049] dW_rel=1.222e-04 dTheta_rel=3.036e-05 time=9.026e-01
[sweep=050] dW_rel=1.279e-04 dTheta_rel=4.208e-05 time=9.603e-01
Runtime: 117.7s | programs selected: 4
Saved: Results/stgp/oli/stgp_result.pkl
[6]:
# ── Attach scores to AnnData & save ─────────────────────────────────────────
ADATA_PATH = OUT_DIR / "adata_with_scores.h5ad"
adata = adata_oli.copy()
all_idx = np.concatenate(idx_per_group) # res["H"] rows follow np.unique(id_region) group order.
H_arr = np.empty_like(res["H"])
H_arr[all_idx] = res["H"]
b_arr = np.empty_like(res["b"])
b_arr[all_idx] = res["b"]
adata.obsm["X_stgp"] = H_arr.astype(np.float32)
adata.obsm["X_stgp_spatial"] = b_arr.astype(np.float32)
alpha_arr = np.asarray(res.get("alpha", []))
alpha_lower_arr = np.asarray(res.get("alpha_lower", []))
alpha_upper_arr = np.asarray(res.get("alpha_upper", []))
theta_arr = np.asarray(res.get("theta", []))
p_sel = res["W"].shape[0]
adata.uns["stgp"] = dict(
groups=uniq.tolist(),
ages=ages.tolist(),
gamma_age=float(res["gamma_age"]),
gamma_spa=float(res["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 [],
theta=theta_arr.tolist() if theta_arr.ndim == 2 else [],
sigma2e=float(res.get("sigma2e", np.nan)),
)
adata.write_h5ad(str(ADATA_PATH), compression="gzip")
# Also write W.csv and a long-form active-gene table for enrichment/downstream plots.
prog_labels = [f"stGP{j + 1}" for j in range(p_sel)]
W_df = pd.DataFrame(res["W"], index=prog_labels, columns=adata.var_names.astype(str))
W_df.to_csv(OUT_DIR / "W.csv")
W_long = (
W_df.stack()
.rename_axis(["program", "gene"])
.rename("weight")
.reset_index()
)
W_long = W_long[W_long["weight"] > 0].copy()
W_long = W_long.sort_values(["program", "weight"], ascending=[True, False])
W_long["rank"] = W_long.groupby("program").cumcount() + 1
W_long.to_csv(OUT_DIR / "W_active_genes.csv", index=False)
4. Downstream Analysis
4.1 Gene Loadings (W matrix)
Each row of W is a non-negative weight vector over genes. Genes with positive weights are the defining markers of that program.
[7]:
W_df = pd.read_csv(OUT_DIR / "W.csv", index_col=0)
W_df.index = [f"stGP{i+1}" for i in range(len(W_df))]
print("Top 10 genes per program:")
for prog, row in W_df.iterrows():
top = row[row > 0].sort_values(ascending=False).head(10)
print(f" {prog}: {', '.join(top.index.tolist())}")
Top 10 genes per program:
stGP1: SLC17A7, SLC1A2, MEIS3, C1QL3, APBA2, SATB2, ID2, SYT5, ACTL6B, RNF208
stGP2: GAD1, OLIG1, AQP4, RBPJ, SLC1A2, TOP2B, MED13, SDHA, MAT2A, EXOSC6
stGP3: GAD1, RPL3, RPL8, RPL7A, SOX2, FKBP5, FSD1, RPL10A, PHB2, MRPS26
stGP4: SORCS2, RGS11, MOG, FKBP5, MAT2A, PDIA2, OLIG1, IWS1, SOX2, OLIG2
[8]:
adata_full = sc.read_h5ad("data/qc/human_merfish_qc.h5ad")
temp = adata_full[adata_full.obs["id_region"] == "5887_rep1"].copy()
sc.pl.embedding(
temp,
basis='spatial',
color = 'SLC17A7', vmax = 20
)
[9]:
temp = adata[adata.obs["id_region"] == "5887_rep1"].copy()
sc.pl.embedding(
temp,
basis='spatial',
color = 'SLC17A7', vmax = 20
)
[10]:
# ── Heatmap: top genes across all programs ───────────────────────────────────
n_top = 15
top_genes_per_prog = []
for _, row in W_df.iterrows():
top_genes_per_prog.extend(row[row > 0].sort_values(ascending=False).head(n_top).index.tolist())
top_genes = list(dict.fromkeys(top_genes_per_prog)) # preserve order, deduplicate
W_sub = W_df[top_genes]
# Normalise each row to [0, 1] for visualisation
W_norm = W_sub.div(W_sub.max(axis=1) + 1e-12, axis=0)
fig, ax = plt.subplots(figsize=(min(0.4 * len(top_genes) + 2, 18), 3))
im = ax.imshow(W_norm.values, aspect="auto", cmap="YlOrRd", vmin=0, vmax=1)
ax.set_yticks(range(len(W_df))); ax.set_yticklabels(W_df.index)
ax.set_xticks(range(len(top_genes)))
ax.set_xticklabels(top_genes, rotation=90, fontsize=7.5)
ax.set_title("Gene loadings (W) – top genes per program", fontsize=11)
plt.colorbar(im, ax=ax, shrink=0.8, label="Normalised weight")
plt.tight_layout()
#plt.savefig(FIGURES_DIR / "W_heatmap.png", dpi=200, bbox_inches="tight")
plt.show()
4.2 Spatial Gene-Program Maps
The spatial field b captures the within-slice smooth variation of each program. We tile all tissue sections ordered by donor age so any age-related spatial patterns become visible.
[11]:
from IPython.display import display
from plots import plot_stgp_spatial_programs
scores_df = pd.DataFrame(
adata.obsm["X_stgp"],
index=adata.obs_names,
columns=[f"stGP{j+1}" for j in range(p_sel)],
)
figs = plot_stgp_spatial_programs(
stgp_adata=adata, scores=scores_df,
celltype="Oligodendrocyte", age_unit="years",
ncols=4, fg_dot_size=5.0, dpi=150,
)
for j, fig in enumerate(figs):
out_path = FIGURES_DIR / f"spatial_stGP{j+1}.png"
#fig.savefig(out_path, dpi=150, bbox_inches="tight")
display(fig)
<Figure size 1040x770 with 13 Axes>
<Figure size 1040x770 with 13 Axes>
<Figure size 1040x770 with 13 Axes>
<Figure size 1040x770 with 13 Axes>
4.3 Age Trajectories (α)
α(t) is the posterior mean age effect of each program — it quantifies how the program amplitude changes across the human lifespan (15–87 yr). The shaded band shows the 95% posterior credible interval.
[12]:
from plots import plot_alpha_over_age
stgp_info = adata.uns["stgp"]
ages_slices = np.array(stgp_info["ages"])
alpha = np.array(stgp_info["alpha"]) # (p, n_slices)
alpha_lower = np.array(stgp_info["alpha_lower"])
alpha_upper = np.array(stgp_info["alpha_upper"])
fig, axes = plt.subplots(1, p_sel, figsize=(4.5 * p_sel, 4), constrained_layout=True)
for j, ax in enumerate(np.atleast_1d(axes)):
order = np.argsort(ages_slices)
a, lo, hi = alpha[j][order], alpha_lower[j][order], alpha_upper[j][order]
t = ages_slices[order]
ax.fill_between(t, lo, hi, alpha=0.2, color="#2C7FB8", label="95% CI")
ax.plot(t, lo, lw=0.8, ls="--", color="#2C7FB8", alpha=0.5)
ax.plot(t, hi, lw=0.8, ls="--", color="#2C7FB8", alpha=0.5)
ax.plot(t, a, lw=1.8, color="#2C7FB8")
ax.scatter(t, a, s=32, color="#2C7FB8", zorder=3)
ax.axhline(0, color="0.6", lw=0.7, ls=":")
ax.set_title(f"stGP{j+1}", fontsize=12)
ax.set_xlabel("Age (yr)"); ax.set_ylabel("Age effect α" if j == 0 else "")
ax.spines[["top", "right"]].set_visible(False)
if j == 0:
ax.legend(fontsize=8, frameon=False)
fig.suptitle("Oligodendrocyte – age trajectories (α)", fontsize=13)
#plt.savefig(FIGURES_DIR / "alpha_trajectories.png", dpi=200, bbox_inches="tight")
plt.show()
4.4 GP Parameters (θ)
For each program, θ = [amplitude, noise_fraction] describes the relative strength of the spatial GP component versus cell-intrinsic noise. A high amplitude with low noise fraction indicates a strongly spatially structured program.
[13]:
theta = np.array(stgp_info["theta"]) # (p, 2): [amplitude, noise_frac]
prog_names = [f"stGP{j+1}" for j in range(p_sel)]
fig, axes = plt.subplots(1, 2, figsize=(8, 3), constrained_layout=True)
colors = plt.cm.tab10(np.linspace(0, 0.5, p_sel))
axes[0].bar(prog_names, theta[:, 0], color=colors, edgecolor="white", linewidth=0.6)
axes[0].set_ylabel("Temporal variance (θ₁)"); axes[0].set_title("Temporal variance")
axes[1].bar(prog_names, theta[:, 1], color=colors, edgecolor="white", linewidth=0.6)
axes[1].set_ylabel("Spatial Variance (θ₂)"); axes[1].set_title("Spatial Variance")
for ax in axes:
ax.spines[["top", "right"]].set_visible(False)
fig.suptitle("GP parameters per program", fontsize=12)
#plt.savefig(FIGURES_DIR / "theta_barplot.png", dpi=200, bbox_inches="tight")
plt.show()
4.5 Spatial Visualisation of a Single Slice
We inspect one tissue slice in detail, showing the spatial b field of each program (the smooth within-slice component).
[14]:
# Pick a mid-age slice
slice_ages = [(adata.obs.loc[adata.obs["id_region"] == sid, "age"].iloc[0], sid)
for sid in adata.obs["id_region"].unique()]
slice_ages.sort()
_, example_slice = slice_ages[len(slice_ages) // 2] # middle-aged slice
sub = adata[adata.obs["id_region"].astype(str) == example_slice].copy()
age_val = sub.obs["age"].iloc[0]
print(f"Visualising slice: {example_slice} (age {age_val:.0f} yr, n={sub.n_obs} cells)")
fig, axes = plt.subplots(1, p_sel, figsize=(4.5 * p_sel, 4), constrained_layout=True)
b = sub.obsm["X_stgp_spatial"]
xy = np.asarray(sub.obsm["spatial"])
for j, ax in enumerate(np.atleast_1d(axes)):
v99 = np.nanpercentile(np.abs(b[:, j]), 99)
sc_ref = ax.scatter(xy[:, 0], xy[:, 1], c=b[:, j],
cmap="RdBu_r", vmin=-v99, vmax=v99,
s=6, linewidths=0, rasterized=True)
ax.set_aspect("equal"); ax.axis("off")
ax.set_title(f"stGP{j+1}")
plt.colorbar(sc_ref, ax=ax, shrink=0.7, pad=0.01)
fig.suptitle(f"Spatial field b – slice {example_slice} ({age_val:.0f} yr)", fontsize=12)
#plt.savefig(FIGURES_DIR / f"spatial_b_{example_slice}.png", dpi=150, bbox_inches="tight")
plt.show()
Visualising slice: 5887_rep1 (age 49 yr, n=4293 cells)