{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1030a562",
   "metadata": {},
   "source": [
    "# MouseBrain Downstream Tutorial\n",
    "\n",
    "This notebook expands the downstream analysis after the Microglia stGP fit: spatial proximity tests, cell-type shell enrichment, regression, and pathway/signature enrichment.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f7904f7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T22:59:59.645898Z",
     "iopub.status.busy": "2026-07-04T22:59:59.645790Z",
     "iopub.status.idle": "2026-07-04T23:00:07.967880Z",
     "shell.execute_reply": "2026-07-04T23:00:07.966986Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Working directory: /import/home4/byual/stGP-0529/RealData_MouseBrainMERFISH\n",
      "QC data: data/qc/aging_coronal_qc.h5ad\n",
      "Proximity results: Results/proximity\n",
      "Enrichment results: Results/enrichment\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import os\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",
    "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/RealData_MouseBrainMERFISH\")\n",
    "os.chdir(PROJECT_DIR)\n",
    "sys.path.insert(0, str(PROJECT_DIR))\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\"QC data: {u.DATA_QC}\")\n",
    "print(f\"Proximity results: {u.RESULTS_PROXIMITY}\")\n",
    "print(f\"Enrichment results: {u.RESULTS_ENRICHMENT}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "955413be",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T23:00:07.970381Z",
     "iopub.status.busy": "2026-07-04T23:00:07.969953Z",
     "iopub.status.idle": "2026-07-04T23:00:07.987091Z",
     "shell.execute_reply": "2026-07-04T23:00:07.986341Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>celltype</th>\n",
       "      <th>programs</th>\n",
       "      <th>n_perm</th>\n",
       "      <th>skip_existing</th>\n",
       "      <th>proximity_csv_dir</th>\n",
       "      <th>proximity_fig_dir</th>\n",
       "      <th>enrichment_csv_dir</th>\n",
       "      <th>enrichment_fig_dir</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Microglia</td>\n",
       "      <td>[2]</td>\n",
       "      <td>100</td>\n",
       "      <td>True</td>\n",
       "      <td>Results/proximity/Microglia</td>\n",
       "      <td>Figures/Microglia/proximity</td>\n",
       "      <td>Results/enrichment/Microglia</td>\n",
       "      <td>Figures/Microglia/enrichment</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    celltype programs  n_perm  skip_existing            proximity_csv_dir  \\\n",
       "0  Microglia      [2]     100           True  Results/proximity/Microglia   \n",
       "\n",
       "             proximity_fig_dir            enrichment_csv_dir  \\\n",
       "0  Figures/Microglia/proximity  Results/enrichment/Microglia   \n",
       "\n",
       "             enrichment_fig_dir  \n",
       "0  Figures/Microglia/enrichment  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "CELLTYPE = \"Microglia\"\n",
    "PROGRAMS = [2]       # Fig. 5 focuses on stGP2. Use None to run every program.\n",
    "N_PERM = 100\n",
    "SEED = 0\n",
    "SKIP_EXISTING = True\n",
    "\n",
    "GENE_SET_COLLECTIONS = None  # None = all canonical GO + cell-type signature sets.\n",
    "PADJ_THRESHOLD = 0.1\n",
    "N_TOP = 6\n",
    "TOP_PER_PROGRAM = 5\n",
    "\n",
    "safe_ct = u.safe_name(CELLTYPE)\n",
    "prox_csv_dir = u.RESULTS_PROXIMITY / safe_ct\n",
    "prox_fig_dir = u.FIGURES_ROOT / safe_ct / \"proximity\"\n",
    "enrich_csv_dir = u.RESULTS_ENRICHMENT / safe_ct\n",
    "enrich_fig_dir = u.FIGURES_ROOT / safe_ct / \"enrichment\"\n",
    "for path in [prox_csv_dir, prox_fig_dir, enrich_csv_dir, enrich_fig_dir]:\n",
    "    path.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "display(pd.DataFrame([{\n",
    "    \"celltype\": CELLTYPE,\n",
    "    \"programs\": \"all\" if PROGRAMS is None else PROGRAMS,\n",
    "    \"n_perm\": N_PERM,\n",
    "    \"skip_existing\": SKIP_EXISTING,\n",
    "    \"proximity_csv_dir\": str(prox_csv_dir),\n",
    "    \"proximity_fig_dir\": str(prox_fig_dir),\n",
    "    \"enrichment_csv_dir\": str(enrich_csv_dir),\n",
    "    \"enrichment_fig_dir\": str(enrich_fig_dir),\n",
    "}]))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13a6e98a",
   "metadata": {},
   "source": [
    "## 1. Load target stGP scores and all-cell QC coordinates\n",
    "\n",
    "Proximity tests use the target cell type's stGP spatial residual `b` and all cell coordinates/cell-type labels from the QC AnnData.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1aec3e03",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T23:00:07.989940Z",
     "iopub.status.busy": "2026-07-04T23:00:07.989823Z",
     "iopub.status.idle": "2026-07-04T23:00:15.338630Z",
     "shell.execute_reply": "2026-07-04T23:00:15.337538Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Target: 52,417 Microglia cells\n",
      "Programs: ['stGP1', 'stGP2', 'stGP3', 'stGP4']\n",
      "Regions: ['CC/ACO', 'CTX', 'STR', 'VEN']\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "All QC cells: 981,750\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>program</th>\n",
       "      <th>b_mean</th>\n",
       "      <th>b_sd</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>stGP1</td>\n",
       "      <td>-0.000025</td>\n",
       "      <td>2.349388</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>stGP2</td>\n",
       "      <td>0.109022</td>\n",
       "      <td>2.114654</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>stGP3</td>\n",
       "      <td>-0.133021</td>\n",
       "      <td>2.165237</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>stGP4</td>\n",
       "      <td>0.887728</td>\n",
       "      <td>2.202204</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  program    b_mean      b_sd\n",
       "0   stGP1 -0.000025  2.349388\n",
       "1   stGP2  0.109022  2.114654\n",
       "2   stGP3 -0.133021  2.165237\n",
       "3   stGP4  0.887728  2.202204"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "adata_target = u.load_target_data(CELLTYPE, u.RESULTS_STGP)\n",
    "target = u.extract_target_arrays(adata_target)\n",
    "regions = sorted(set(target[\"region\"].tolist()))\n",
    "\n",
    "print(f\"Target: {len(target['age']):,} {CELLTYPE} cells\")\n",
    "print(f\"Programs: {target['program_labels']}\")\n",
    "print(f\"Regions: {regions}\")\n",
    "\n",
    "adata_all = sc.read_h5ad(str(u.DATA_QC))\n",
    "glob = u.extract_global_arrays(adata_all)\n",
    "print(f\"All QC cells: {adata_all.n_obs:,}\")\n",
    "\n",
    "display(pd.DataFrame({\n",
    "    \"program\": target[\"program_labels\"],\n",
    "    \"b_mean\": target[\"B\"].mean(axis=0),\n",
    "    \"b_sd\": target[\"B\"].std(axis=0),\n",
    "}))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f44b0d0",
   "metadata": {},
   "source": [
    "## 2. Abundance over age\n",
    "\n",
    "Before proximity testing, check whether target/effectors change strongly in abundance across age.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "206bc867",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T23:00:15.341994Z",
     "iopub.status.busy": "2026-07-04T23:00:15.341827Z",
     "iopub.status.idle": "2026-07-04T23:00:16.917471Z",
     "shell.execute_reply": "2026-07-04T23:00:16.916673Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>celltype</th>\n",
       "      <th>total</th>\n",
       "      <th>n_slices</th>\n",
       "      <th>spearman_rho</th>\n",
       "      <th>spearman_p</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>T cell</td>\n",
       "      <td>1008</td>\n",
       "      <td>20</td>\n",
       "      <td>0.645113</td>\n",
       "      <td>2.130693e-03</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Microglia</td>\n",
       "      <td>52417</td>\n",
       "      <td>20</td>\n",
       "      <td>0.373825</td>\n",
       "      <td>1.044513e-01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Neuroblast</td>\n",
       "      <td>3260</td>\n",
       "      <td>20</td>\n",
       "      <td>-0.939850</td>\n",
       "      <td>7.860443e-10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>NSC</td>\n",
       "      <td>2448</td>\n",
       "      <td>20</td>\n",
       "      <td>-0.964635</td>\n",
       "      <td>7.227150e-12</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     celltype  total  n_slices  spearman_rho    spearman_p\n",
       "0      T cell   1008        20      0.645113  2.130693e-03\n",
       "1   Microglia  52417        20      0.373825  1.044513e-01\n",
       "2  Neuroblast   3260        20     -0.939850  7.860443e-10\n",
       "3         NSC   2448        20     -0.964635  7.227150e-12"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wrote abundance figure: Figures/Microglia/proximity/abundance_check.png\n"
     ]
    }
   ],
   "source": [
    "df_counts, df_abund = u.compute_abundance_check(CELLTYPE, glob)\n",
    "df_counts.to_csv(prox_csv_dir / \"abundance_per_slice.csv\", index=False)\n",
    "df_abund.to_csv(prox_csv_dir / \"abundance_summary.csv\", index=False)\n",
    "u.render_abundance(CELLTYPE, df_counts, df_abund, prox_fig_dir)\n",
    "\n",
    "display(df_abund.sort_values(\"spearman_rho\", ascending=False))\n",
    "print(f\"Wrote abundance figure: {prox_fig_dir / 'abundance_check.png'}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1417e8e6",
   "metadata": {},
   "source": [
    "## 3. Matched near-vs-far proximity and permutation null\n",
    "\n",
    "For each program, compare target cells near an effector cell type against far target cells within the same age slice. A spatial permutation null keeps effector counts fixed per slice.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1d4ebb73",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T23:00:16.919319Z",
     "iopub.status.busy": "2026-07-04T23:00:16.919169Z",
     "iopub.status.idle": "2026-07-04T23:00:16.939890Z",
     "shell.execute_reply": "2026-07-04T23:00:16.939156Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "--- Microglia stGP2: matched proximity ---\n",
      "loaded cached matched/permutation tables from Results/proximity/Microglia/stGP2\n",
      "selected effectors: ['Endothelial', 'NSC', 'Astrocyte', 'Neuroblast', 'Ependymal', 'T cell', 'Neuron-Excitatory', 'Oligodendrocyte', 'Neuron-MSN', 'Pericyte', 'Macrophage', 'OPC']\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>effector</th>\n",
       "      <th>program</th>\n",
       "      <th>k</th>\n",
       "      <th>effect</th>\n",
       "      <th>se</th>\n",
       "      <th>z</th>\n",
       "      <th>p</th>\n",
       "      <th>n_near</th>\n",
       "      <th>n_far</th>\n",
       "      <th>n_blocks</th>\n",
       "      <th>q_bh</th>\n",
       "      <th>focus_program</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Astrocyte</td>\n",
       "      <td>stGP1</td>\n",
       "      <td>0</td>\n",
       "      <td>-0.333182</td>\n",
       "      <td>0.201157</td>\n",
       "      <td>-1.656328</td>\n",
       "      <td>9.765550e-02</td>\n",
       "      <td>1122</td>\n",
       "      <td>130</td>\n",
       "      <td>2</td>\n",
       "      <td>1.103932e-01</td>\n",
       "      <td>stGP2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Astrocyte</td>\n",
       "      <td>stGP2</td>\n",
       "      <td>1</td>\n",
       "      <td>2.640462</td>\n",
       "      <td>0.134043</td>\n",
       "      <td>19.698552</td>\n",
       "      <td>2.218769e-86</td>\n",
       "      <td>1122</td>\n",
       "      <td>130</td>\n",
       "      <td>2</td>\n",
       "      <td>5.244364e-86</td>\n",
       "      <td>stGP2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Astrocyte</td>\n",
       "      <td>stGP3</td>\n",
       "      <td>2</td>\n",
       "      <td>1.079732</td>\n",
       "      <td>0.089434</td>\n",
       "      <td>12.072961</td>\n",
       "      <td>1.467561e-33</td>\n",
       "      <td>1122</td>\n",
       "      <td>130</td>\n",
       "      <td>2</td>\n",
       "      <td>2.312520e-33</td>\n",
       "      <td>stGP2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Astrocyte</td>\n",
       "      <td>stGP4</td>\n",
       "      <td>3</td>\n",
       "      <td>-1.990231</td>\n",
       "      <td>0.182373</td>\n",
       "      <td>-10.913000</td>\n",
       "      <td>9.990450e-28</td>\n",
       "      <td>1122</td>\n",
       "      <td>130</td>\n",
       "      <td>2</td>\n",
       "      <td>1.484295e-27</td>\n",
       "      <td>stGP2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Endothelial</td>\n",
       "      <td>stGP1</td>\n",
       "      <td>0</td>\n",
       "      <td>-3.665859</td>\n",
       "      <td>0.245810</td>\n",
       "      <td>-14.913354</td>\n",
       "      <td>2.698483e-50</td>\n",
       "      <td>5482</td>\n",
       "      <td>86</td>\n",
       "      <td>6</td>\n",
       "      <td>5.197079e-50</td>\n",
       "      <td>stGP2</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      effector program  k    effect        se          z             p  \\\n",
       "0    Astrocyte   stGP1  0 -0.333182  0.201157  -1.656328  9.765550e-02   \n",
       "1    Astrocyte   stGP2  1  2.640462  0.134043  19.698552  2.218769e-86   \n",
       "2    Astrocyte   stGP3  2  1.079732  0.089434  12.072961  1.467561e-33   \n",
       "3    Astrocyte   stGP4  3 -1.990231  0.182373 -10.913000  9.990450e-28   \n",
       "4  Endothelial   stGP1  0 -3.665859  0.245810 -14.913354  2.698483e-50   \n",
       "\n",
       "   n_near  n_far  n_blocks          q_bh focus_program  \n",
       "0    1122    130         2  1.103932e-01         stGP2  \n",
       "1    1122    130         2  5.244364e-86         stGP2  \n",
       "2    1122    130         2  2.312520e-33         stGP2  \n",
       "3    1122    130         2  1.484295e-27         stGP2  \n",
       "4    5482     86         6  5.197079e-50         stGP2  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "if PROGRAMS is None:\n",
    "    program_indices = list(range(target[\"n_programs\"]))\n",
    "else:\n",
    "    program_indices = [p - 1 for p in PROGRAMS if 0 < p <= target[\"n_programs\"]]\n",
    "effectors = [c for c in u.ALL_CELLTYPES if c != CELLTYPE]\n",
    "\n",
    "program_context = {}\n",
    "for k in program_indices:\n",
    "    k_label = target[\"program_labels\"][k]\n",
    "    sub_csv = prox_csv_dir / k_label\n",
    "    sub_fig = prox_fig_dir / k_label\n",
    "    sub_csv.mkdir(parents=True, exist_ok=True)\n",
    "    sub_fig.mkdir(parents=True, exist_ok=True)\n",
    "    match_path = sub_csv / \"matched_effects.csv\"\n",
    "    perm_path = sub_csv / \"permutation_null.csv\"\n",
    "\n",
    "    print(f\"\\n--- {CELLTYPE} {k_label}: matched proximity ---\")\n",
    "    if SKIP_EXISTING and match_path.exists() and perm_path.exists():\n",
    "        df_match = pd.read_csv(match_path)\n",
    "        df_perm = pd.read_csv(perm_path)\n",
    "        print(f\"loaded cached matched/permutation tables from {sub_csv}\")\n",
    "    else:\n",
    "        t0 = time.perf_counter()\n",
    "        df_match = u.compute_matched_effect_table(target, glob, regions, effectors)\n",
    "        df_match.to_csv(match_path, index=False)\n",
    "        u.make_standalone_matched_heatmap(df_match, CELLTYPE, k_label, sub_fig / \"matched_heatmap.png\")\n",
    "        print(f\"matched table done in {time.perf_counter() - t0:.1f}s\")\n",
    "\n",
    "        t0 = time.perf_counter()\n",
    "        df_perm, _, _ = u.compute_permutation_null(\n",
    "            target, glob, regions, effectors, k, n_perm=N_PERM, seed=SEED,\n",
    "        )\n",
    "        df_perm.to_csv(perm_path, index=False)\n",
    "        u.make_standalone_effector_ranking(df_match, df_perm, CELLTYPE, k_label, sub_fig / \"effector_ranking.png\")\n",
    "        print(f\"permutation null done in {time.perf_counter() - t0:.1f}s\")\n",
    "\n",
    "    sig_effectors = u._select_significant_effectors(df_match, k_label, q_threshold=u.SIG_Q_THRESHOLD)\n",
    "    if not sig_effectors:\n",
    "        sig_effectors = [e for e in u.DEFAULT_PERM_EFFECTORS if e != CELLTYPE][:4]\n",
    "    print(f\"selected effectors: {sig_effectors}\")\n",
    "\n",
    "    program_context[k] = dict(\n",
    "        label=k_label, sub_csv=sub_csv, sub_fig=sub_fig,\n",
    "        df_match=df_match, df_perm=df_perm, sig_effectors=sig_effectors,\n",
    "    )\n",
    "\n",
    "display(pd.concat([\n",
    "    ctx[\"df_match\"].assign(focus_program=ctx[\"label\"])\n",
    "    for ctx in program_context.values()\n",
    "], ignore_index=True).head())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caf8d24a",
   "metadata": {},
   "source": [
    "## 4. Distance decay, age stratification, shell enrichment, and regression\n",
    "\n",
    "These analyses characterize whether proximity effects are distance-dependent, age-dependent, enriched/depleted in a local shell, and explainable by age/region/cell-type density predictors.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "50733f10",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T23:00:16.941475Z",
     "iopub.status.busy": "2026-07-04T23:00:16.941300Z",
     "iopub.status.idle": "2026-07-04T23:00:17.184555Z",
     "shell.execute_reply": "2026-07-04T23:00:17.183691Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "--- Microglia stGP2: secondary proximity analyses ---\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loaded cached secondary tables from Results/proximity/Microglia/stGP2\n",
      "secondary outputs are available in Results/proximity/Microglia/stGP2 and Figures/Microglia/proximity/stGP2\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>target_celltype</th>\n",
       "      <th>program</th>\n",
       "      <th>demo_age</th>\n",
       "      <th>n_target_cells</th>\n",
       "      <th>n_predictors</th>\n",
       "      <th>R2_full</th>\n",
       "      <th>sig_effectors</th>\n",
       "      <th>top_pro_aging_effector</th>\n",
       "      <th>top_pro_aging_effect</th>\n",
       "      <th>top_pro_rejuv_effector</th>\n",
       "      <th>top_pro_rejuv_effect</th>\n",
       "      <th>n_perm_pass</th>\n",
       "      <th>perm_pass_effectors</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Microglia</td>\n",
       "      <td>stGP2</td>\n",
       "      <td>34.5</td>\n",
       "      <td>52417</td>\n",
       "      <td>17</td>\n",
       "      <td>0.4008</td>\n",
       "      <td>[Endothelial, NSC, Astrocyte, Neuroblast, Epen...</td>\n",
       "      <td>NSC</td>\n",
       "      <td>3.390692</td>\n",
       "      <td>Endothelial</td>\n",
       "      <td>-4.832034</td>\n",
       "      <td>7</td>\n",
       "      <td>Astrocyte|Endothelial|Ependymal|Neuroblast|NSC...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  target_celltype program  demo_age  n_target_cells  n_predictors  R2_full  \\\n",
       "0       Microglia   stGP2      34.5           52417            17   0.4008   \n",
       "\n",
       "                                       sig_effectors top_pro_aging_effector  \\\n",
       "0  [Endothelial, NSC, Astrocyte, Neuroblast, Epen...                    NSC   \n",
       "\n",
       "   top_pro_aging_effect top_pro_rejuv_effector  top_pro_rejuv_effect  \\\n",
       "0              3.390692            Endothelial             -4.832034   \n",
       "\n",
       "   n_perm_pass                                perm_pass_effectors  \n",
       "0            7  Astrocyte|Endothelial|Ependymal|Neuroblast|NSC...  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "program_summaries = []\n",
    "for k, ctx in program_context.items():\n",
    "    k_label = ctx[\"label\"]\n",
    "    sub_csv = ctx[\"sub_csv\"]\n",
    "    sub_fig = ctx[\"sub_fig\"]\n",
    "    sig_effectors = ctx[\"sig_effectors\"]\n",
    "    df_match = ctx[\"df_match\"]\n",
    "    df_perm = ctx[\"df_perm\"]\n",
    "    print(f\"\\n--- {CELLTYPE} {k_label}: secondary proximity analyses ---\")\n",
    "\n",
    "    secondary_paths = {\n",
    "        \"decay\": sub_csv / \"distance_decay.csv\",\n",
    "        \"age\": sub_csv / \"age_stratification.csv\",\n",
    "        \"enrich\": sub_csv / \"proximity_enrichment.csv\",\n",
    "        \"coefs\": sub_csv / \"variance_decomposition_coefs_full.csv\",\n",
    "        \"meta\": sub_csv / \"variance_decomposition_meta.json\",\n",
    "    }\n",
    "    if SKIP_EXISTING and all(p.exists() for p in secondary_paths.values()):\n",
    "        df_decay = pd.read_csv(secondary_paths[\"decay\"])\n",
    "        df_age = pd.read_csv(secondary_paths[\"age\"])\n",
    "        df_enrich = pd.read_csv(secondary_paths[\"enrich\"])\n",
    "        df_coefs = pd.read_csv(secondary_paths[\"coefs\"])\n",
    "        var_meta = json.loads(secondary_paths[\"meta\"].read_text())\n",
    "        demo_age = u.pick_demo_slice(target, k, glob)\n",
    "        print(f\"loaded cached secondary tables from {sub_csv}\")\n",
    "    else:\n",
    "        df_decay = u.compute_distance_decay(target, glob, regions, sig_effectors, k)\n",
    "        df_decay.to_csv(secondary_paths[\"decay\"], index=False)\n",
    "        u.make_standalone_distance_decay(df_decay, CELLTYPE, k_label, sub_fig / \"distance_decay.png\")\n",
    "\n",
    "        df_age = u.compute_age_stratification(target, glob, regions, sig_effectors, k)\n",
    "        df_age.to_csv(secondary_paths[\"age\"], index=False)\n",
    "        u.make_standalone_age_stratification(df_age, CELLTYPE, k_label, sub_fig / \"age_stratification.png\")\n",
    "\n",
    "        df_enrich, _, _, _ = u.compute_proximity_enrichment(target, glob, regions, k)\n",
    "        df_enrich.to_csv(secondary_paths[\"enrich\"], index=False)\n",
    "        u.make_standalone_enrichment(df_enrich, CELLTYPE, k_label, sub_fig / \"proximity_enrichment.png\")\n",
    "\n",
    "        df_coefs, var_meta = u.compute_variance_decomposition(target, glob, regions, k, CELLTYPE)\n",
    "        df_coefs.to_csv(secondary_paths[\"coefs\"], index=False)\n",
    "        secondary_paths[\"meta\"].write_text(json.dumps(var_meta, indent=2))\n",
    "        u.make_standalone_variance(df_coefs, var_meta, CELLTYPE, k_label, sub_fig / \"regression.png\")\n",
    "\n",
    "        demo_age = u.pick_demo_slice(target, k, glob)\n",
    "        for eff_name, eff_tag in [(\"T cell\", \"Tcell\"), (\"NSC\", \"NSC\")]:\n",
    "            if eff_name == CELLTYPE:\n",
    "                continue\n",
    "            eff_demo_age = u.pick_demo_slice(target, k, glob, prefer_old_with_effector=eff_name)\n",
    "            u.make_standalone_spatial(\n",
    "                target, glob, k, eff_demo_age, CELLTYPE, k_label,\n",
    "                sub_fig / f\"spatial_example_{eff_tag}.png\", effector=eff_name,\n",
    "            )\n",
    "        u.make_standalone_near_far_violins(\n",
    "            target, glob, k, CELLTYPE, k_label, sig_effectors, sub_fig / \"near_far_violins.png\",\n",
    "        )\n",
    "\n",
    "    summary = u._summarise_program(df_match, df_perm, CELLTYPE, k_label, demo_age, target, sig_effectors, var_meta)\n",
    "    program_summaries.append(summary)\n",
    "    print(f\"secondary outputs are available in {sub_csv} and {sub_fig}\")\n",
    "\n",
    "display(pd.DataFrame(program_summaries))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdf42b20",
   "metadata": {},
   "source": [
    "## 5. Downstream high-b vs low-b cell-type enrichment\n",
    "\n",
    "This section asks which surrounding cell types are enriched around high spatial residual `b` cells, both across all slices and within each slice.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b0b6b0b3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T23:00:17.186247Z",
     "iopub.status.busy": "2026-07-04T23:00:17.186124Z",
     "iopub.status.idle": "2026-07-04T23:00:17.210723Z",
     "shell.execute_reply": "2026-07-04T23:00:17.209831Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "--- Microglia stGP2: downstream high-b/low-b enrichment ---\n",
      "loaded cached downstream tables from Results/proximity/Microglia/stGP2/downstream\n",
      "{'target_celltype': 'Microglia', 'program': 'stGP2', 'n_effectors_tested': 13, 'n_slices_tested': 20, 'n_all_slice_enrichment_sig': 12, 'top_all_slice_enriched_effector': 'Oligodendrocyte', 'top_all_slice_depleted_effector': 'Neuron-Excitatory', 'n_per_slice_enrichment_sig': 156, 'n_per_slice_matched_sig': 4, 'n_global_matched_sig': 12, 'global_matched_sig_effectors': 'T cell|NSC|Neuroblast|Macrophage|Ependymal|Pericyte|OPC|Endothelial|Astrocyte|Neuron-MSN|Oligodendrocyte|Neuron-Excitatory'}\n"
     ]
    }
   ],
   "source": [
    "for k, ctx in program_context.items():\n",
    "    k_label = ctx[\"label\"]\n",
    "    out_csv = ctx[\"sub_csv\"] / \"downstream\"\n",
    "    out_fig = ctx[\"sub_fig\"] / \"downstream\"\n",
    "    out_csv.mkdir(parents=True, exist_ok=True)\n",
    "    out_fig.mkdir(parents=True, exist_ok=True)\n",
    "    df_match_prog = ctx[\"df_match\"][ctx[\"df_match\"][\"program\"] == k_label].copy()\n",
    "\n",
    "    print(f\"\\n--- {CELLTYPE} {k_label}: downstream high-b/low-b enrichment ---\")\n",
    "    downstream_paths = {\n",
    "        \"all\": out_csv / \"all_slices_enrichment.csv\",\n",
    "        \"slice_enrich\": out_csv / \"per_slice_enrichment.csv\",\n",
    "        \"slice_match\": out_csv / \"per_slice_matched_effects.csv\",\n",
    "        \"summary\": out_csv / \"downstream_summary.json\",\n",
    "    }\n",
    "    if SKIP_EXISTING and all(p.exists() for p in downstream_paths.values()):\n",
    "        df_all = pd.read_csv(downstream_paths[\"all\"])\n",
    "        df_slice_enrich = pd.read_csv(downstream_paths[\"slice_enrich\"])\n",
    "        df_slice_match = pd.read_csv(downstream_paths[\"slice_match\"])\n",
    "        downstream_summary = json.loads(downstream_paths[\"summary\"].read_text())\n",
    "        print(f\"loaded cached downstream tables from {out_csv}\")\n",
    "    else:\n",
    "        counts_by_eff = u._shell_counts_by_effector(target, glob, effectors)\n",
    "        df_all = u.compute_downstream_all_slices_enrichment(\n",
    "            target, glob, effectors, k, counts_by_eff=counts_by_eff,\n",
    "        )\n",
    "        df_slice_enrich = u.compute_downstream_per_slice_enrichment(\n",
    "            target, glob, effectors, k, counts_by_eff=counts_by_eff,\n",
    "        )\n",
    "        df_slice_match = u.compute_downstream_per_slice_matched_effects(target, glob, effectors, k)\n",
    "\n",
    "        df_all.to_csv(downstream_paths[\"all\"], index=False)\n",
    "        df_slice_enrich.to_csv(downstream_paths[\"slice_enrich\"], index=False)\n",
    "        df_slice_match.to_csv(downstream_paths[\"slice_match\"], index=False)\n",
    "\n",
    "        downstream_summary = u.summarise_downstream(\n",
    "            df_all, df_slice_enrich, df_slice_match, df_match_prog, CELLTYPE, k_label,\n",
    "        )\n",
    "        downstream_paths[\"summary\"].write_text(json.dumps(downstream_summary, indent=2))\n",
    "\n",
    "        u.make_downstream_all_slices_figure(\n",
    "            df_all, df_match_prog, CELLTYPE, k_label, out_fig / \"all_slices_celltype_enrichment.png\",\n",
    "        )\n",
    "        u.make_downstream_heatmap(\n",
    "            df_slice_enrich, value_col=\"log2fc\", q_col=\"q_bh_by_slice\",\n",
    "            title=f\"{CELLTYPE} {k_label}: per-slice high-b shell enrichment\",\n",
    "            cbar_label=\"log2 shell enrichment\",\n",
    "            savepath=out_fig / \"per_slice_enrichment_heatmap.png\",\n",
    "        )\n",
    "        u.make_downstream_heatmap(\n",
    "            df_slice_match, value_col=\"effect\", q_col=\"q_bh_by_slice\",\n",
    "            title=f\"{CELLTYPE} {k_label}: per-slice matched proximity effect\",\n",
    "            cbar_label=\"Delta median b (near - far)\",\n",
    "            savepath=out_fig / \"per_slice_matched_effect_heatmap.png\",\n",
    "        )\n",
    "\n",
    "        by_ct_dir = out_fig / \"by_celltype\"\n",
    "        for eff in effectors:\n",
    "            eff_dir = by_ct_dir / u.safe_name(eff)\n",
    "            eff_dir.mkdir(parents=True, exist_ok=True)\n",
    "            u.make_downstream_effector_trend(\n",
    "                df_slice_enrich, df_slice_match, eff, CELLTYPE, k_label, eff_dir / \"slice_trend.png\",\n",
    "            )\n",
    "            demo_age = u.pick_demo_slice(target, k, glob, prefer_old_with_effector=eff)\n",
    "            u.make_downstream_effector_spatial(\n",
    "                target, glob, k, demo_age, CELLTYPE, k_label, eff, eff_dir / \"spatial_overlay.png\",\n",
    "            )\n",
    "\n",
    "    ctx[\"downstream_summary\"] = downstream_summary\n",
    "    print(downstream_summary)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76c6a122",
   "metadata": {},
   "source": [
    "## 6. Cell-type summary tables\n",
    "\n",
    "Concatenate per-program downstream summaries into cell-type and all-cell-type overview tables.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5d79d216",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T23:00:17.212142Z",
     "iopub.status.busy": "2026-07-04T23:00:17.211937Z",
     "iopub.status.idle": "2026-07-04T23:00:17.776282Z",
     "shell.execute_reply": "2026-07-04T23:00:17.775740Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>program</th>\n",
       "      <th>k</th>\n",
       "      <th>effector</th>\n",
       "      <th>hi_threshold</th>\n",
       "      <th>lo_threshold</th>\n",
       "      <th>n_hi</th>\n",
       "      <th>n_lo</th>\n",
       "      <th>hi_mean</th>\n",
       "      <th>lo_mean</th>\n",
       "      <th>log2fc</th>\n",
       "      <th>p</th>\n",
       "      <th>q_bh</th>\n",
       "      <th>R_in</th>\n",
       "      <th>R_out</th>\n",
       "      <th>high_pct</th>\n",
       "      <th>low_pct</th>\n",
       "      <th>min_group</th>\n",
       "      <th>valid</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>stGP2</td>\n",
       "      <td>1</td>\n",
       "      <td>T cell</td>\n",
       "      <td>1.048398</td>\n",
       "      <td>-1.247125</td>\n",
       "      <td>13105</td>\n",
       "      <td>13105</td>\n",
       "      <td>0.026326</td>\n",
       "      <td>0.006410</td>\n",
       "      <td>0.436226</td>\n",
       "      <td>1.136039e-34</td>\n",
       "      <td>2.461419e-34</td>\n",
       "      <td>20.0</td>\n",
       "      <td>50.0</td>\n",
       "      <td>75.0</td>\n",
       "      <td>25.0</td>\n",
       "      <td>30</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>stGP2</td>\n",
       "      <td>1</td>\n",
       "      <td>NSC</td>\n",
       "      <td>1.048398</td>\n",
       "      <td>-1.247125</td>\n",
       "      <td>13105</td>\n",
       "      <td>13105</td>\n",
       "      <td>0.047005</td>\n",
       "      <td>0.003129</td>\n",
       "      <td>0.868570</td>\n",
       "      <td>6.716685e-29</td>\n",
       "      <td>1.091461e-28</td>\n",
       "      <td>20.0</td>\n",
       "      <td>50.0</td>\n",
       "      <td>75.0</td>\n",
       "      <td>25.0</td>\n",
       "      <td>30</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>stGP2</td>\n",
       "      <td>1</td>\n",
       "      <td>Neuroblast</td>\n",
       "      <td>1.048398</td>\n",
       "      <td>-1.247125</td>\n",
       "      <td>13105</td>\n",
       "      <td>13105</td>\n",
       "      <td>0.077757</td>\n",
       "      <td>0.006791</td>\n",
       "      <td>1.169656</td>\n",
       "      <td>7.208538e-44</td>\n",
       "      <td>1.874220e-43</td>\n",
       "      <td>20.0</td>\n",
       "      <td>50.0</td>\n",
       "      <td>75.0</td>\n",
       "      <td>25.0</td>\n",
       "      <td>30</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>stGP2</td>\n",
       "      <td>1</td>\n",
       "      <td>Macrophage</td>\n",
       "      <td>1.048398</td>\n",
       "      <td>-1.247125</td>\n",
       "      <td>13105</td>\n",
       "      <td>13105</td>\n",
       "      <td>0.038306</td>\n",
       "      <td>0.051431</td>\n",
       "      <td>-0.199912</td>\n",
       "      <td>3.426075e-07</td>\n",
       "      <td>4.453897e-07</td>\n",
       "      <td>20.0</td>\n",
       "      <td>50.0</td>\n",
       "      <td>75.0</td>\n",
       "      <td>25.0</td>\n",
       "      <td>30</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>stGP2</td>\n",
       "      <td>1</td>\n",
       "      <td>Ependymal</td>\n",
       "      <td>1.048398</td>\n",
       "      <td>-1.247125</td>\n",
       "      <td>13105</td>\n",
       "      <td>13105</td>\n",
       "      <td>0.074857</td>\n",
       "      <td>0.012514</td>\n",
       "      <td>0.998018</td>\n",
       "      <td>5.836289e-32</td>\n",
       "      <td>1.083882e-31</td>\n",
       "      <td>20.0</td>\n",
       "      <td>50.0</td>\n",
       "      <td>75.0</td>\n",
       "      <td>25.0</td>\n",
       "      <td>30</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  program  k    effector  hi_threshold  lo_threshold   n_hi   n_lo   hi_mean  \\\n",
       "0   stGP2  1      T cell      1.048398     -1.247125  13105  13105  0.026326   \n",
       "1   stGP2  1         NSC      1.048398     -1.247125  13105  13105  0.047005   \n",
       "2   stGP2  1  Neuroblast      1.048398     -1.247125  13105  13105  0.077757   \n",
       "3   stGP2  1  Macrophage      1.048398     -1.247125  13105  13105  0.038306   \n",
       "4   stGP2  1   Ependymal      1.048398     -1.247125  13105  13105  0.074857   \n",
       "\n",
       "    lo_mean    log2fc             p          q_bh  R_in  R_out  high_pct  \\\n",
       "0  0.006410  0.436226  1.136039e-34  2.461419e-34  20.0   50.0      75.0   \n",
       "1  0.003129  0.868570  6.716685e-29  1.091461e-28  20.0   50.0      75.0   \n",
       "2  0.006791  1.169656  7.208538e-44  1.874220e-43  20.0   50.0      75.0   \n",
       "3  0.051431 -0.199912  3.426075e-07  4.453897e-07  20.0   50.0      75.0   \n",
       "4  0.012514  0.998018  5.836289e-32  1.083882e-31  20.0   50.0      75.0   \n",
       "\n",
       "   low_pct  min_group  valid  \n",
       "0     25.0         30   True  \n",
       "1     25.0         30   True  \n",
       "2     25.0         30   True  \n",
       "3     25.0         30   True  \n",
       "4     25.0         30   True  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wrote proximity summary: Results/proximity/Microglia/summary.json\n"
     ]
    }
   ],
   "source": [
    "downstream_summary_df = u.compile_celltype_downstream_summary(\n",
    "    CELLTYPE,\n",
    "    csv_dir=prox_csv_dir,\n",
    "    fig_dir=prox_fig_dir,\n",
    "    program_labels=[program_context[k][\"label\"] for k in program_context],\n",
    ")\n",
    "\n",
    "summary = dict(\n",
    "    target_celltype=CELLTYPE,\n",
    "    n_target_cells=int(len(target[\"age\"])),\n",
    "    n_programs=int(target[\"n_programs\"]),\n",
    "    regions=regions,\n",
    "    n_downstream_summary_rows=int(len(downstream_summary_df)),\n",
    "    programs=program_summaries,\n",
    ")\n",
    "(prox_csv_dir / \"summary.json\").write_text(json.dumps(summary, indent=2))\n",
    "\n",
    "master_rows = []\n",
    "for ct in u.ALL_CELLTYPES:\n",
    "    sj = u.RESULTS_PROXIMITY / u.safe_name(ct) / \"summary.json\"\n",
    "    if sj.exists():\n",
    "        s = json.loads(sj.read_text())\n",
    "        master_rows.extend([p for p in s.get(\"programs\", []) if \"error\" not in p])\n",
    "if master_rows:\n",
    "    pd.DataFrame(master_rows).to_csv(u.RESULTS_PROXIMITY / \"summary_all_celltypes.csv\", index=False)\n",
    "\n",
    "display(downstream_summary_df.head())\n",
    "print(f\"Wrote proximity summary: {prox_csv_dir / 'summary.json'}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f3b9420",
   "metadata": {},
   "source": [
    "## 7. Pathway and cell-type signature enrichment\n",
    "\n",
    "Positive stGP gene weights are tested against the local GMT collections. Per-program bar charts and a combined program-by-term dotplot are written under `Figures/<celltype>/enrichment`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "31fd6888",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T23:00:17.777913Z",
     "iopub.status.busy": "2026-07-04T23:00:17.777790Z",
     "iopub.status.idle": "2026-07-04T23:00:17.792379Z",
     "shell.execute_reply": "2026-07-04T23:00:17.791886Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded W: 4 programs x 220 genes\n",
      "Gene-set collections: ['GO Biological process', 'GO Molecular Function', 'GO Cellular Component', 'Cell-type signatures']\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Th</th>\n",
       "      <th>Cd52</th>\n",
       "      <th>S100a6</th>\n",
       "      <th>Cd3g</th>\n",
       "      <th>Tgfb1</th>\n",
       "      <th>Cd79a</th>\n",
       "      <th>Tpm4</th>\n",
       "      <th>Ercc1</th>\n",
       "      <th>Trp53</th>\n",
       "      <th>Trh</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>stGP1</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.00000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>stGP2</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.03098</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>stGP3</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.00000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.047638</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>stGP4</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.00000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        Th     Cd52  S100a6  Cd3g     Tgfb1  Cd79a  Tpm4  Ercc1  Trp53  Trh\n",
       "stGP1  0.0  0.00000     0.0   0.0  0.000000    0.0   0.0    0.0    0.0  0.0\n",
       "stGP2  0.0  0.03098     0.0   0.0  0.000000    0.0   0.0    0.0    0.0  0.0\n",
       "stGP3  0.0  0.00000     0.0   0.0  0.047638    0.0   0.0    0.0    0.0  0.0\n",
       "stGP4  0.0  0.00000     0.0   0.0  0.000000    0.0   0.0    0.0    0.0  0.0"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "gene_sets = u._resolve_gene_sets(GENE_SET_COLLECTIONS, u.DEFAULT_GENESETS_DIR)\n",
    "W = u._load_W(u.RESULTS_STGP, CELLTYPE)\n",
    "programs = W.index.astype(str).tolist()\n",
    "background_genes = W.columns.astype(str).tolist()\n",
    "\n",
    "print(f\"Loaded W: {len(programs)} programs x {len(background_genes)} genes\")\n",
    "print(f\"Gene-set collections: {list(gene_sets.keys())}\")\n",
    "display(W.iloc[:, : min(10, W.shape[1])])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "14e03e0e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-04T23:00:17.793461Z",
     "iopub.status.busy": "2026-07-04T23:00:17.793355Z",
     "iopub.status.idle": "2026-07-04T23:00:21.733524Z",
     "shell.execute_reply": "2026-07-04T23:00:21.732702Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  [program] stGP1: 10 positive genes ...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  [program] stGP2: 15 positive genes ...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  [program] stGP3: 15 positive genes ...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  [program] stGP4: 15 positive genes ...\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>celltype</th>\n",
       "      <th>program</th>\n",
       "      <th>n_genes_input</th>\n",
       "      <th>Gene_set</th>\n",
       "      <th>Term</th>\n",
       "      <th>Overlap</th>\n",
       "      <th>P-value</th>\n",
       "      <th>Adjusted P-value</th>\n",
       "      <th>Odds Ratio</th>\n",
       "      <th>Combined Score</th>\n",
       "      <th>Genes</th>\n",
       "      <th>gene_set</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Microglia</td>\n",
       "      <td>stGP1</td>\n",
       "      <td>10</td>\n",
       "      <td>m5.go.bp.v2026.1.Mm.symbols.gmt</td>\n",
       "      <td>Actin filament based process</td>\n",
       "      <td>2/16</td>\n",
       "      <td>0.158439</td>\n",
       "      <td>0.460611</td>\n",
       "      <td>3.985801</td>\n",
       "      <td>7.343384</td>\n",
       "      <td>Pik3r1;Prox1</td>\n",
       "      <td>GO Biological process</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Microglia</td>\n",
       "      <td>stGP1</td>\n",
       "      <td>10</td>\n",
       "      <td>m5.go.bp.v2026.1.Mm.symbols.gmt</td>\n",
       "      <td>Actin filament bundle organization</td>\n",
       "      <td>1/2</td>\n",
       "      <td>0.089041</td>\n",
       "      <td>0.356662</td>\n",
       "      <td>22.052632</td>\n",
       "      <td>53.337758</td>\n",
       "      <td>Pik3r1</td>\n",
       "      <td>GO Biological process</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Microglia</td>\n",
       "      <td>stGP1</td>\n",
       "      <td>10</td>\n",
       "      <td>m5.go.bp.v2026.1.Mm.symbols.gmt</td>\n",
       "      <td>Actin filament organization</td>\n",
       "      <td>2/7</td>\n",
       "      <td>0.034656</td>\n",
       "      <td>0.346712</td>\n",
       "      <td>10.989305</td>\n",
       "      <td>36.949025</td>\n",
       "      <td>Pik3r1;Prox1</td>\n",
       "      <td>GO Biological process</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Microglia</td>\n",
       "      <td>stGP1</td>\n",
       "      <td>10</td>\n",
       "      <td>m5.go.bp.v2026.1.Mm.symbols.gmt</td>\n",
       "      <td>Activation of immune response</td>\n",
       "      <td>2/29</td>\n",
       "      <td>0.388978</td>\n",
       "      <td>0.619581</td>\n",
       "      <td>1.962567</td>\n",
       "      <td>1.853120</td>\n",
       "      <td>Pik3r1;Mog</td>\n",
       "      <td>GO Biological process</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Microglia</td>\n",
       "      <td>stGP1</td>\n",
       "      <td>10</td>\n",
       "      <td>m5.go.bp.v2026.1.Mm.symbols.gmt</td>\n",
       "      <td>Activation of innate immune response</td>\n",
       "      <td>1/8</td>\n",
       "      <td>0.315020</td>\n",
       "      <td>0.575638</td>\n",
       "      <td>4.284211</td>\n",
       "      <td>4.948775</td>\n",
       "      <td>Pik3r1</td>\n",
       "      <td>GO Biological process</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    celltype program  n_genes_input                         Gene_set  \\\n",
       "0  Microglia   stGP1             10  m5.go.bp.v2026.1.Mm.symbols.gmt   \n",
       "1  Microglia   stGP1             10  m5.go.bp.v2026.1.Mm.symbols.gmt   \n",
       "2  Microglia   stGP1             10  m5.go.bp.v2026.1.Mm.symbols.gmt   \n",
       "3  Microglia   stGP1             10  m5.go.bp.v2026.1.Mm.symbols.gmt   \n",
       "4  Microglia   stGP1             10  m5.go.bp.v2026.1.Mm.symbols.gmt   \n",
       "\n",
       "                                   Term Overlap   P-value  Adjusted P-value  \\\n",
       "0          Actin filament based process    2/16  0.158439          0.460611   \n",
       "1    Actin filament bundle organization     1/2  0.089041          0.356662   \n",
       "2           Actin filament organization     2/7  0.034656          0.346712   \n",
       "3         Activation of immune response    2/29  0.388978          0.619581   \n",
       "4  Activation of innate immune response     1/8  0.315020          0.575638   \n",
       "\n",
       "   Odds Ratio  Combined Score         Genes               gene_set  \n",
       "0    3.985801        7.343384  Pik3r1;Prox1  GO Biological process  \n",
       "1   22.052632       53.337758        Pik3r1  GO Biological process  \n",
       "2   10.989305       36.949025  Pik3r1;Prox1  GO Biological process  \n",
       "3    1.962567        1.853120    Pik3r1;Mog  GO Biological process  \n",
       "4    4.284211        4.948775        Pik3r1  GO Biological process  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[master] Wrote 80 rows -> Results/enrichment/summary_all_celltypes.csv\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>program</th>\n",
       "      <th>n_pos_genes</th>\n",
       "      <th>n_terms_reported</th>\n",
       "      <th>runtime_sec</th>\n",
       "      <th>status</th>\n",
       "      <th>figure</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>stGP1</td>\n",
       "      <td>10</td>\n",
       "      <td>882</td>\n",
       "      <td>0.84</td>\n",
       "      <td>done</td>\n",
       "      <td>Figures/Microglia/enrichment/Microglia_stGP1_e...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>stGP2</td>\n",
       "      <td>15</td>\n",
       "      <td>1033</td>\n",
       "      <td>1.41</td>\n",
       "      <td>done</td>\n",
       "      <td>Figures/Microglia/enrichment/Microglia_stGP2_e...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>stGP3</td>\n",
       "      <td>15</td>\n",
       "      <td>1741</td>\n",
       "      <td>0.60</td>\n",
       "      <td>done</td>\n",
       "      <td>Figures/Microglia/enrichment/Microglia_stGP3_e...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>stGP4</td>\n",
       "      <td>15</td>\n",
       "      <td>1196</td>\n",
       "      <td>0.59</td>\n",
       "      <td>done</td>\n",
       "      <td>Figures/Microglia/enrichment/Microglia_stGP4_e...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  program  n_pos_genes  n_terms_reported  runtime_sec status  \\\n",
       "0   stGP1           10               882         0.84   done   \n",
       "1   stGP2           15              1033         1.41   done   \n",
       "2   stGP3           15              1741         0.60   done   \n",
       "3   stGP4           15              1196         0.59   done   \n",
       "\n",
       "                                              figure  \n",
       "0  Figures/Microglia/enrichment/Microglia_stGP1_e...  \n",
       "1  Figures/Microglia/enrichment/Microglia_stGP2_e...  \n",
       "2  Figures/Microglia/enrichment/Microglia_stGP3_e...  \n",
       "3  Figures/Microglia/enrichment/Microglia_stGP4_e...  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wrote enrichment summary: Results/enrichment/Microglia/enrichment_summary.json\n"
     ]
    }
   ],
   "source": [
    "all_enrichment_results = []\n",
    "enrichment_stats = []\n",
    "for prog in programs:\n",
    "    stats, per_prog = u._enrich_one_program(\n",
    "        program=prog,\n",
    "        weights=W.loc[prog],\n",
    "        background_genes=background_genes,\n",
    "        gene_sets=gene_sets,\n",
    "        celltype=CELLTYPE,\n",
    "        safe=safe_ct,\n",
    "        fig_dir=enrich_fig_dir,\n",
    "        res_dir=enrich_csv_dir,\n",
    "        padj_threshold=PADJ_THRESHOLD,\n",
    "        n_top=N_TOP,\n",
    "    )\n",
    "    enrichment_stats.append(stats)\n",
    "    if per_prog is not None:\n",
    "        all_enrichment_results.append(per_prog)\n",
    "\n",
    "combined_path = None\n",
    "if all_enrichment_results:\n",
    "    combined = pd.concat(all_enrichment_results, ignore_index=True)\n",
    "    combined.insert(0, \"celltype\", CELLTYPE)\n",
    "    combined_path = enrich_csv_dir / f\"{safe_ct}_combined_enrichment.csv\"\n",
    "    combined.to_csv(combined_path, index=False)\n",
    "    u._plot_program_enrichment_dotplot(\n",
    "        combined,\n",
    "        out=enrich_fig_dir / f\"{safe_ct}_enrichment_program_dotplot.png\",\n",
    "        padj_threshold=PADJ_THRESHOLD,\n",
    "    )\n",
    "    display(combined.head())\n",
    "else:\n",
    "    combined = pd.DataFrame()\n",
    "\n",
    "enrich_summary = dict(\n",
    "    celltype=CELLTYPE,\n",
    "    safe_name=safe_ct,\n",
    "    n_programs=len(programs),\n",
    "    programs=enrichment_stats,\n",
    "    gene_sets=list(gene_sets.keys()),\n",
    "    padj_threshold=float(PADJ_THRESHOLD),\n",
    "    n_top_per_panel=int(N_TOP),\n",
    "    combined_csv=str(combined_path) if combined_path else None,\n",
    ")\n",
    "(enrich_csv_dir / \"enrichment_summary.json\").write_text(json.dumps(enrich_summary, indent=2))\n",
    "u.compile_enrichment_master_summary([CELLTYPE], results_root=u.RESULTS_ENRICHMENT, top_per_program=TOP_PER_PROGRAM)\n",
    "\n",
    "display(pd.DataFrame(enrichment_stats))\n",
    "print(f\"Wrote enrichment summary: {enrich_csv_dir / 'enrichment_summary.json'}\")\n"
   ]
  }
 ],
 "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
}