{ "cells": [ { "cell_type": "markdown", "id": "00f83aba", "metadata": {}, "source": [ "# MouseBrain Microglia stGP Tutorial\n", "\n", "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`." ] }, { "cell_type": "code", "execution_count": 1, "id": "ee3a6bbd", "metadata": { "execution": { "iopub.execute_input": "2026-07-04T22:38:40.494861Z", "iopub.status.busy": "2026-07-04T22:38:40.494750Z", "iopub.status.idle": "2026-07-04T22:38:43.955471Z", "shell.execute_reply": "2026-07-04T22:38:43.953963Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Working directory: /import/home4/byual/stGP-0529/RealData_MouseBrainMERFISH\n", "Results root: /import/home4/byual/stGP-0529/RealData_MouseBrainMERFISH/Results/stgp\n", "Figures root: /import/home4/byual/stGP-0529/RealData_MouseBrainMERFISH/Figures\n" ] } ], "source": [ "import json\n", "import os\n", "import pickle\n", "import sys\n", "import time\n", "from pathlib import Path\n", "\n", "import matplotlib\n", "matplotlib.use(\"Agg\")\n", "import numpy as np\n", "import pandas as pd\n", "import scanpy as sc\n", "import scipy.sparse as sp\n", "from IPython.display import display\n", "\n", "PROJECT_DIR = Path.cwd()\n", "if PROJECT_DIR.name != \"RealData_MouseBrainMERFISH\":\n", " PROJECT_DIR = Path(\"/home/byual/stGP-0529/RealData_MouseBrainMERFISH\")\n", "os.chdir(PROJECT_DIR)\n", "sys.path.insert(0, str(PROJECT_DIR))\n", "sys.path.insert(0, str(PROJECT_DIR.parent))\n", "\n", "from stgp.estimation import fit_pfactor, fit_pfactor_auto\n", "from stgp.kernels import (\n", " bandwidth_select_spatial,\n", " bandwidth_select_temporal,\n", " build_K_age,\n", " build_K_spa_list_from_stacked,\n", ")\n", "from stgp.preprocessing import log1p_norm_centered_list, standardize_coords_list\n", "\n", "import utils as u\n", "from plots import set_nature_style\n", "\n", "set_nature_style()\n", "print(f\"Working directory: {PROJECT_DIR}\")\n", "print(f\"Results root: {(PROJECT_DIR / u.RESULTS_STGP).resolve()}\")\n", "print(f\"Figures root: {(PROJECT_DIR / u.FIGURES_ROOT).resolve()}\")\n" ] }, { "cell_type": "code", "execution_count": 2, "id": "b3fd0893", "metadata": { "execution": { "iopub.execute_input": "2026-07-04T22:38:43.958729Z", "iopub.status.busy": "2026-07-04T22:38:43.957884Z", "iopub.status.idle": "2026-07-04T22:38:43.979519Z", "shell.execute_reply": "2026-07-04T22:38:43.978683Z" } }, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
celltypeinput_h5adstgp_dirfigure_dirfit_modeprun_stgp_now
0Microgliadata/processed/Microglia.h5adResults/stgp/MicrogliaFigures/Microgliafixed p4False
\n", "
" ], "text/plain": [ " celltype input_h5ad stgp_dir \\\n", "0 Microglia data/processed/Microglia.h5ad Results/stgp/Microglia \n", "\n", " figure_dir fit_mode p run_stgp_now \n", "0 Figures/Microglia fixed p 4 False " ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "# Tutorial parameters.\n", "CELLTYPE = \"Microglia\"\n", "N_PROGRAMS = 4 # Fig. 5 fixes p=4 for matched method benchmarking.\n", "P_MAX = 10 # Used only when N_PROGRAMS is None.\n", "K_NEIGHBORS = 15\n", "SEED = 0\n", "SKIP_EXISTING = True\n", "\n", "safe_ct = u.safe_name(CELLTYPE)\n", "input_h5ad = u.DATA_PROCESSED / f\"{safe_ct}.h5ad\"\n", "stgp_dir = u.RESULTS_STGP / safe_ct\n", "fig_dir = u.FIGURES_ROOT / safe_ct\n", "stgp_pkl = stgp_dir / \"stgp_result.pkl\"\n", "scored_h5ad = stgp_dir / \"adata_with_scores.h5ad\"\n", "W_csv = stgp_dir / \"W.csv\"\n", "\n", "run_stgp = not (SKIP_EXISTING and stgp_pkl.exists() and scored_h5ad.exists() and W_csv.exists())\n", "\n", "display(pd.DataFrame([{\n", " \"celltype\": CELLTYPE,\n", " \"input_h5ad\": str(input_h5ad),\n", " \"stgp_dir\": str(stgp_dir),\n", " \"figure_dir\": str(fig_dir),\n", " \"fit_mode\": \"fixed p\" if N_PROGRAMS is not None else \"auto rank\",\n", " \"p\": N_PROGRAMS,\n", " \"run_stgp_now\": run_stgp,\n", "}]))\n" ] }, { "cell_type": "markdown", "id": "d04ab710", "metadata": {}, "source": [ "## 1. Load cells and organize slices\n", "\n", "stGP is fit to one target cell type. Cells are grouped by `mouse_id`; each mouse contributes one age and one spatial kernel block.\n" ] }, { "cell_type": "code", "execution_count": 3, "id": "2f3b7ee0", "metadata": { "execution": { "iopub.execute_input": "2026-07-04T22:38:43.981497Z", "iopub.status.busy": "2026-07-04T22:38:43.981291Z", "iopub.status.idle": "2026-07-04T22:38:44.367083Z", "shell.execute_reply": "2026-07-04T22:38:44.365810Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Using cached stGP outputs in Results/stgp/Microglia\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Cached scored AnnData: 52,417 cells x 220 genes\n" ] } ], "source": [ "if run_stgp:\n", " if not input_h5ad.exists():\n", " raise FileNotFoundError(f\"Missing processed input: {input_h5ad}\")\n", "\n", " adata = sc.read_h5ad(str(input_h5ad))\n", " age_arr = adata.obs[\"age\"].to_numpy(dtype=float)\n", " groups = adata.obs[\"mouse_id\"].astype(str).to_numpy()\n", " uniq_mice, inv = np.unique(groups, return_inverse=True)\n", " idx_per_group = [np.sort(np.where(inv == t)[0]) for t in range(len(uniq_mice))]\n", " nlist = np.array([len(ix) for ix in idx_per_group])\n", " ages = np.array([age_arr[ix[0]] for ix in idx_per_group])\n", "\n", " summary = pd.DataFrame({\"mouse_id\": uniq_mice, \"age\": ages, \"n_cells\": nlist})\n", " print(f\"Loaded {adata.n_obs:,} {CELLTYPE} cells and {adata.n_vars:,} genes.\")\n", " display(summary.head())\n", " display(summary.describe(include=\"all\"))\n", "else:\n", " print(f\"Using cached stGP outputs in {stgp_dir}\")\n", " adata_scores = sc.read_h5ad(str(scored_h5ad))\n", " print(f\"Cached scored AnnData: {adata_scores.n_obs:,} cells x {adata_scores.n_vars:,} genes\")\n" ] }, { "cell_type": "markdown", "id": "fd3f74cf", "metadata": {}, "source": [ "## 2. Build expression, temporal, and spatial inputs\n", "\n", "The expression matrix is split by mouse, log-normalized/centered, and paired with an age kernel plus one spatial kernel per mouse slice.\n" ] }, { "cell_type": "code", "execution_count": 4, "id": "3328bed5", "metadata": { "execution": { "iopub.execute_input": "2026-07-04T22:38:44.368878Z", "iopub.status.busy": "2026-07-04T22:38:44.368705Z", "iopub.status.idle": "2026-07-04T22:38:44.373963Z", "shell.execute_reply": "2026-07-04T22:38:44.373178Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Skipped kernel construction because cached stGP outputs are present.\n" ] } ], "source": [ "if run_stgp:\n", " X_raw = adata.X.toarray() if sp.issparse(adata.X) else np.asarray(adata.X)\n", " Y_list, _ = log1p_norm_centered_list([X_raw[ix] for ix in idx_per_group])\n", " coords_list = standardize_coords_list([adata.obsm[\"spatial\"][ix] for ix in idx_per_group])\n", "\n", " gamma_spa = bandwidth_select_spatial(coords_list, frac=0.01, rho=0.75)\n", " gamma_age = bandwidth_select_temporal(ages, rho=np.exp(-4))\n", " K_age = build_K_age(ages, gamma_age, kernel=\"rbf\", standardize=True)\n", " K_spa_list = build_K_spa_list_from_stacked(\n", " np.vstack(coords_list), nlist, gamma_spa, standardize=False, jitter=1e-6,\n", " )\n", "\n", " print(f\"gamma_spa = {gamma_spa:.4f}\")\n", " print(f\"gamma_age = {gamma_age:.4f}\")\n", " print(f\"K_age shape: {K_age.shape}\")\n", " print(f\"Spatial kernel blocks: {len(K_spa_list)}\")\n", "else:\n", " print(\"Skipped kernel construction because cached stGP outputs are present.\")\n" ] }, { "cell_type": "markdown", "id": "76275e8a", "metadata": {}, "source": [ "## 3. Fit stGP\n", "\n", "`N_PROGRAMS=4` uses the fixed-rank `fit_pfactor` mode used for Fig. 5 benchmarking. Set `N_PROGRAMS=None` to use `fit_pfactor_auto`.\n" ] }, { "cell_type": "code", "execution_count": 5, "id": "1de83c9e", "metadata": { "execution": { "iopub.execute_input": "2026-07-04T22:38:44.374999Z", "iopub.status.busy": "2026-07-04T22:38:44.374887Z", "iopub.status.idle": "2026-07-04T22:38:44.383398Z", "shell.execute_reply": "2026-07-04T22:38:44.382424Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Loaded cached stGP pickle: Results/stgp/Microglia/stgp_result.pkl\n", "Cached programs: 4\n" ] } ], "source": [ "if run_stgp:\n", " t0 = time.perf_counter()\n", " if N_PROGRAMS is not None:\n", " stgp_res = fit_pfactor(\n", " Y_list, nlist, K_age, K_spa_list,\n", " p=N_PROGRAMS, k=K_NEIGHBORS, verbose=1,\n", " )\n", " else:\n", " stgp_res = fit_pfactor_auto(\n", " Y_list=Y_list, Nlist=nlist, K_age=K_age, Kspa_list=K_spa_list,\n", " p_max=P_MAX, k=K_NEIGHBORS, random_state=SEED, verbose=1,\n", " )\n", " elapsed = time.perf_counter() - t0\n", " print(f\"stGP runtime: {elapsed:.1f} sec\")\n", " print(f\"Selected programs: {stgp_res['W'].shape[0]}\")\n", "else:\n", " with open(stgp_pkl, \"rb\") as f:\n", " stgp_res = pickle.load(f)\n", " elapsed = None\n", " print(f\"Loaded cached stGP pickle: {stgp_pkl}\")\n", " print(f\"Cached programs: {np.asarray(stgp_res['W']).shape[0]}\")\n" ] }, { "cell_type": "markdown", "id": "366af001", "metadata": {}, "source": [ "## 4. Write stGP outputs\n", "\n", "Scores `H` and spatial residuals `b` are mapped back to original cell order and saved in `obsm['X_stgp']` and `obsm['X_stgp_spatial']`.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "c87e116d", "metadata": { "execution": { "iopub.execute_input": "2026-07-04T22:38:44.385534Z", "iopub.status.busy": "2026-07-04T22:38:44.385380Z", "iopub.status.idle": "2026-07-04T22:38:44.404516Z", "shell.execute_reply": "2026-07-04T22:38:44.403787Z" } }, "outputs": [], "source": [ "if run_stgp:\n", " stgp_dir.mkdir(parents=True, exist_ok=True)\n", " stgp_res[\"gamma_age\"] = gamma_age\n", " stgp_res[\"gamma_spa\"] = gamma_spa\n", " with open(stgp_pkl, \"wb\") as f:\n", " pickle.dump(stgp_res, f)\n", "\n", " p_sel = stgp_res[\"W\"].shape[0]\n", " prog_labels = [f\"stGP{j + 1}\" for j in range(p_sel)]\n", " W = pd.DataFrame(stgp_res[\"W\"], index=prog_labels, columns=adata.var_names)\n", " W.to_csv(W_csv)\n", "\n", " all_idx = np.concatenate(idx_per_group)\n", " H_adata = np.empty_like(stgp_res[\"H\"])\n", " H_adata[all_idx] = stgp_res[\"H\"]\n", " b_adata = np.empty_like(stgp_res[\"b\"])\n", " b_adata[all_idx] = stgp_res[\"b\"]\n", " adata.obsm[\"X_stgp\"] = H_adata.astype(np.float32)\n", " adata.obsm[\"X_stgp_spatial\"] = b_adata.astype(np.float32)\n", "\n", " alpha_arr = np.asarray(stgp_res.get(\"alpha\", []))\n", " alpha_lower_arr = np.asarray(stgp_res.get(\"alpha_lower\", []))\n", " alpha_upper_arr = np.asarray(stgp_res.get(\"alpha_upper\", []))\n", " adata.uns[\"stgp\"] = dict(\n", " groups=uniq_mice.tolist(), ages=ages.tolist(),\n", " gamma_age=gamma_age, gamma_spa=gamma_spa, p_selected=p_sel,\n", " alpha=alpha_arr.tolist() if alpha_arr.ndim == 2 else [],\n", " alpha_lower=alpha_lower_arr.tolist() if alpha_lower_arr.ndim == 2 else [],\n", " alpha_upper=alpha_upper_arr.tolist() if alpha_upper_arr.ndim == 2 else [],\n", " )\n", " adata.write_h5ad(scored_h5ad, compression=\"gzip\")\n", "\n", " timing = {\n", " \"method\": \"stgp\", \"celltype\": CELLTYPE,\n", " \"runtime_sec\": round(elapsed, 2), \"status\": \"completed\",\n", " \"p\": N_PROGRAMS, \"p_max\": P_MAX, \"k\": K_NEIGHBORS, \"seed\": SEED,\n", " }\n", " (stgp_dir / \"timing.json\").write_text(json.dumps(timing, indent=2))\n", " u.append_jsonl(u.TIMING_LOG, {**timing, \"out_dir\": str(stgp_dir.resolve())})\n", "else:\n", " W = pd.read_csv(W_csv, index_col=0)\n" ] }, { "cell_type": "markdown", "id": "cfba2ff5", "metadata": {}, "source": [ "## 5. Load stGP and baseline results for figure generation\n", "\n", "Missing baseline directories are skipped. Available methods are summarized before plotting.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "5df77b04", "metadata": { "execution": { "iopub.execute_input": "2026-07-04T22:38:44.406436Z", "iopub.status.busy": "2026-07-04T22:38:44.406292Z", "iopub.status.idle": "2026-07-04T22:38:46.723799Z", "shell.execute_reply": "2026-07-04T22:38:46.723052Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ " [loaded] stGP: (52417, 4)\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ " [loaded] SpatialPCA: (52417, 4)\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ " [loaded] MEFISTO: (52417, 4)\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ " [loaded] STAMP: (51003, 4)\n", " [loaded] Popari: (52417, 4)\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
methodcellsprogramshas_gene_weightsresult_dir
0stGP524174TrueResults/stgp/Microglia
1SpatialPCA524174TrueResults/baselines/spatialpca/Microglia
2MEFISTO524174TrueResults/baselines/mefisto/Microglia
3STAMP510034TrueResults/baselines/stamp/Microglia
4Popari524174TrueResults/baselines/popari/Microglia
\n", "
" ], "text/plain": [ " method cells programs has_gene_weights \\\n", "0 stGP 52417 4 True \n", "1 SpatialPCA 52417 4 True \n", "2 MEFISTO 52417 4 True \n", "3 STAMP 51003 4 True \n", "4 Popari 52417 4 True \n", "\n", " result_dir \n", "0 Results/stgp/Microglia \n", "1 Results/baselines/spatialpca/Microglia \n", "2 Results/baselines/mefisto/Microglia \n", "3 Results/baselines/stamp/Microglia \n", "4 Results/baselines/popari/Microglia " ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "fig_dir.mkdir(parents=True, exist_ok=True)\n", "methods = u._load_methods(safe_ct, CELLTYPE, stgp_dir=stgp_dir)\n", "if not methods:\n", " raise RuntimeError(f\"No method results found for {CELLTYPE}\")\n", "\n", "stgp_method = next((m for m in methods if m.method == \"stGP\"), None)\n", "stgp_res = u._load_stgp_pickle(safe_ct, stgp_dir=stgp_dir) if stgp_method is not None else None\n", "\n", "stamp_m = next((m for m in methods if m.method == \"STAMP\"), None)\n", "if stamp_m is not None and stamp_m.gene_weights is None:\n", " stamp_m.gene_weights = u._infer_gene_weights(stamp_m.adata, stamp_m.scores)\n", "\n", "method_table = pd.DataFrame([\n", " {\n", " \"method\": m.method,\n", " \"cells\": m.adata.n_obs,\n", " \"programs\": m.scores.shape[1],\n", " \"has_gene_weights\": m.gene_weights is not None,\n", " \"result_dir\": str(m.result_dir),\n", " }\n", " for m in methods\n", "])" ] }, { "cell_type": "markdown", "id": "f2c53434", "metadata": {}, "source": [ "## 6. Temporal variance and age-related genes\n", "\n", "These panels ask which stGP programs are age-dominated, how weighted expression changes with age, and which genes drive the age-associated programs.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "708138f8", "metadata": { "execution": { "iopub.execute_input": "2026-07-04T22:38:46.726658Z", "iopub.status.busy": "2026-07-04T22:38:46.726509Z", "iopub.status.idle": "2026-07-04T22:38:54.052899Z", "shell.execute_reply": "2026-07-04T22:38:54.051989Z" } }, "outputs": [], "source": [ "stgp_vdf = None\n", "if stgp_method is not None:\n", " stgp_vdf = u._section1_variance_decomposition(\n", " stgp_method, stgp_res, celltype=CELLTYPE, safe_ct=safe_ct, fig_dir=fig_dir,\n", " )\n", " u._section2_program_scores_by_age(\n", " stgp_method, stgp_vdf, celltype=CELLTYPE, safe_ct=safe_ct, fig_dir=fig_dir,\n", " )" ] }, { "cell_type": "markdown", "id": "015245f2", "metadata": {}, "source": [ "## 7. Spatial program maps, domains, and method similarity\n", "\n", "This section writes spatial maps, age-ordered stack plots, benchmark cluster comparisons, and stGP-vs-baseline gene-program similarity summaries.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "8815cf6d", "metadata": { "execution": { "iopub.execute_input": "2026-07-04T22:38:54.054563Z", "iopub.status.busy": "2026-07-04T22:38:54.054415Z", "iopub.status.idle": "2026-07-04T22:38:54.069494Z", "shell.execute_reply": "2026-07-04T22:38:54.068822Z" } }, "outputs": [], "source": [ "spatial_outputs = sorted((fig_dir / \"spatial_selected_2x2\").rglob(\"*.png\"))\n", "sim_outputs = sorted((fig_dir / \"program_similarity\").rglob(\"*.csv\"))\n", "\n", "if spatial_outputs and sim_outputs:\n", " print(\"Using cached spatial maps and program-similarity tables.\")\n", "else:\n", " u._section4_spatial_maps_and_clustering(\n", " methods, stgp_method, celltype=CELLTYPE, safe_ct=safe_ct, fig_dir=fig_dir,\n", " )\n", " u._section5_program_similarity(\n", " methods, stgp_method, celltype=CELLTYPE, safe_ct=safe_ct, fig_dir=fig_dir,\n", " )\n", " spatial_outputs = sorted((fig_dir / \"spatial_selected_2x2\").rglob(\"*.png\"))\n", " sim_outputs = sorted((fig_dir / \"program_similarity\").rglob(\"*.csv\"))" ] }, { "cell_type": "markdown", "id": "74ea2f6f", "metadata": {}, "source": [ "## 8. Kernel diagnostic, W heatmap, and runtime comparison\n", "\n", "These final upstream figures document the spatial kernel scale, active gene-weight matrix, and method runtimes.\n" ] }, { "cell_type": "code", "execution_count": null, "id": "aac4594b", "metadata": { "execution": { "iopub.execute_input": "2026-07-04T22:38:54.071302Z", "iopub.status.busy": "2026-07-04T22:38:54.071191Z", "iopub.status.idle": "2026-07-04T22:38:55.647278Z", "shell.execute_reply": "2026-07-04T22:38:55.646098Z" } }, "outputs": [], "source": [ "if stgp_method is not None:\n", " u._section6_kernel_diagnostic_and_W_heatmap(\n", " stgp_method, stgp_res, celltype=CELLTYPE, safe_ct=safe_ct, fig_dir=fig_dir,\n", " )\n", "u._section7_timing(celltype=CELLTYPE, safe_ct=safe_ct, fig_dir=fig_dir, stgp_dir=stgp_dir)\n", "\n", "figure_files = sorted(fig_dir.rglob(\"*.png\"))\n", "table_files = sorted(fig_dir.rglob(\"*.csv\"))" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.15" } }, "nbformat": 4, "nbformat_minor": 5 }