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