{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "title",
   "metadata": {},
   "source": [
    "# Injured PT stGP analysis: left kidney\n",
    "\n",
    "Fit stGP on the left biological replicate (`ident` suffix `L`) of injured proximal tubule cells. Run this notebook with the `stGP` conda environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "setup",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-31T09:36:41.550804Z",
     "iopub.status.busy": "2026-05-31T09:36:41.550555Z",
     "iopub.status.idle": "2026-05-31T09:37:18.100547Z",
     "shell.execute_reply": "2026-05-31T09:37:18.099081Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AnnData object with n_obs × n_vars = 69210 × 299\n",
      "    obs: 'x_centroid', 'y_centroid', 'n_genes', 'n_counts', 'ident', 'region', 'celltype_plot', 'time', 'CN', 'injury_time_days', 'side', 'age'\n",
      "    uns: 'CN_colors', 'celltype_plot_colors', 'ident_colors', 'neighbors', 'pca', 'umap'\n",
      "    obsm: 'X_pca', 'X_pca_harmony', 'X_umap', 'spatial'\n",
      "    varm: 'PCs'\n",
      "    obsp: 'connectivities', 'distances'\n",
      "  ident   time  injury_time_days side\n",
      " Hour4L  Hour4          0.166667    L\n",
      "Hour12L Hour12          0.500000    L\n",
      "  Day2L   Day2          2.000000    L\n",
      " Day14L  Day14         14.000000    L\n",
      " Week6L  Week6         42.000000    L\n"
     ]
    }
   ],
   "source": [
    "import os, sys, warnings, pickle\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import scanpy as sc\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "warnings.filterwarnings(\"ignore\", category=FutureWarning)\n",
    "sys.path.insert(0, \"..\")\n",
    "\n",
    "DATA_PROC = Path(\"data/processed\")\n",
    "RESULTS_DIR = Path(\"Results/stgp\")\n",
    "FIGURES_DIR = Path(\"Figures/Inj_PT_L\")\n",
    "FIGURES_DIR.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "CELLTYPE = \"Inj_PT\"\n",
    "SIDE = \"L\"\n",
    "ANALYSIS_NAME = f\"{CELLTYPE}_{SIDE}\"\n",
    "\n",
    "adata_inj_pt = sc.read_h5ad(DATA_PROC / f\"{CELLTYPE}.h5ad\")\n",
    "adata_inj_pt.obs[\"side\"] = adata_inj_pt.obs[\"ident\"].astype(str).str[-1]\n",
    "adata_inj_pt = adata_inj_pt[adata_inj_pt.obs[\"injury_time_days\"] > 0].copy()\n",
    "adata_inj_pt = adata_inj_pt[adata_inj_pt.obs[\"side\"] == SIDE].copy()\n",
    "adata_inj_pt.obs[\"age\"] = adata_inj_pt.obs[\"injury_time_days\"].copy()\n",
    "\n",
    "print(adata_inj_pt)\n",
    "print(adata_inj_pt.obs[[\"ident\", \"time\", \"injury_time_days\", \"side\"]].drop_duplicates().sort_values(\"injury_time_days\").to_string(index=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "global-style",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-31T09:37:18.103819Z",
     "iopub.status.busy": "2026-05-31T09:37:18.103293Z",
     "iopub.status.idle": "2026-05-31T09:37:18.107800Z",
     "shell.execute_reply": "2026-05-31T09:37:18.107053Z"
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "\n",
    "mpl.rcParams.update({\n",
    "    \"font.family\": \"sans-serif\",\n",
    "    \"font.sans-serif\": [\"Arial\", \"Helvetica\", \"DejaVu Sans\"],\n",
    "    \"font.size\": 11,\n",
    "    \"pdf.fonttype\": 42,\n",
    "    \"ps.fonttype\": 42,\n",
    "    \"svg.fonttype\": \"none\",\n",
    "    \"figure.dpi\": 150,\n",
    "    \"savefig.dpi\": 400,\n",
    "    \"figure.facecolor\": \"white\",\n",
    "    \"axes.facecolor\": \"white\",\n",
    "    \"axes.linewidth\": 1.2,\n",
    "    \"axes.spines.top\": False,\n",
    "    \"axes.spines.right\": False,\n",
    "    \"axes.labelsize\": 11,\n",
    "    \"axes.titlesize\": 12,\n",
    "    \"xtick.direction\": \"out\",\n",
    "    \"ytick.direction\": \"out\",\n",
    "    \"legend.frameon\": False,\n",
    "    \"legend.fontsize\": 9,\n",
    "    \"lines.linewidth\": 1.5,\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "prepare-stgp",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-31T09:37:18.110479Z",
     "iopub.status.busy": "2026-05-31T09:37:18.110333Z",
     "iopub.status.idle": "2026-05-31T09:37:24.860420Z",
     "shell.execute_reply": "2026-05-31T09:37:24.859301Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/byual/.conda/envs/stGP/lib/python3.11/functools.py:909: UserWarning: zero-centering a sparse array/matrix densifies it.\n",
      "  return dispatch(args[0].__class__)(*args, **kw)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prepared Inj_PT_L: 69210 cells across 5 slices\n",
      "  slice       age  n_cells\n",
      " Hour4L  0.166667    15139\n",
      "Hour12L  0.500000    33544\n",
      "  Day2L  2.000000    19449\n",
      " Day14L 14.000000      772\n",
      " Week6L 42.000000      306\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "from stgp.estimation import fit_pfactor_auto\n",
    "from stgp.kernels import (\n",
    "    bandwidth_select_spatial,\n",
    "    bandwidth_select_temporal,\n",
    "    build_K_age,\n",
    "    build_K_spa_list_from_stacked,\n",
    ")\n",
    "from stgp.preprocessing import standardize_coords_list\n",
    "\n",
    "OUT_DIR = RESULTS_DIR / ANALYSIS_NAME\n",
    "OUT_DIR.mkdir(parents=True, exist_ok=True)\n",
    "PKL_PATH = OUT_DIR / \"stgp_result.pkl\"\n",
    "\n",
    "age_arr = pd.to_numeric(adata_inj_pt.obs[\"injury_time_days\"], errors=\"coerce\").to_numpy(float)\n",
    "groups = adata_inj_pt.obs[\"ident\"].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",
    "adata_prep = adata_inj_pt.copy()\n",
    "sc.pp.scale(adata_prep)\n",
    "Y_list = [adata_prep.X[ix] for ix in idx_per_group]\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",
    "sort_ord = np.argsort(ages)\n",
    "ages = ages[sort_ord]\n",
    "slices = uniq[sort_ord]\n",
    "idx_sorted = [idx_per_group[i] for i in sort_ord]\n",
    "nlist = nlist[sort_ord]\n",
    "Y_list = [Y_list[i] for i in sort_ord]\n",
    "\n",
    "coords_list = standardize_coords_list([adata_inj_pt.obsm[\"spatial\"][ix] for ix in idx_per_group])\n",
    "coords_list = [coords_list[i] for i in sort_ord]\n",
    "\n",
    "print(f\"Prepared {ANALYSIS_NAME}: {adata_inj_pt.n_obs} cells across {len(slices)} slices\")\n",
    "print(pd.DataFrame({\"slice\": slices, \"age\": ages, \"n_cells\": nlist}).to_string(index=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "build-kernels",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-31T09:37:24.863931Z",
     "iopub.status.busy": "2026-05-31T09:37:24.863754Z",
     "iopub.status.idle": "2026-05-31T09:39:24.262426Z",
     "shell.execute_reply": "2026-05-31T09:39:24.259475Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gamma_spa = 0.3181 | gamma_age = 0.3904\n"
     ]
    }
   ],
   "source": [
    "gamma_spa = bandwidth_select_spatial(coords_list, frac=0.01, rho=0.5)\n",
    "gamma_age = bandwidth_select_temporal(ages, rho=np.exp(-1.5))\n",
    "print(f\"gamma_spa = {gamma_spa:.4f} | gamma_age = {gamma_age:.4f}\")\n",
    "\n",
    "K_age = build_K_age(ages, gamma_age, kernel=\"ar1\", 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",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fit-stgp",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-31T09:39:24.267360Z",
     "iopub.status.busy": "2026-05-31T09:39:24.266839Z",
     "iopub.status.idle": "2026-05-31T11:29:38.187254Z",
     "shell.execute_reply": "2026-05-31T11:29:38.184428Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=001] dW_rel=2.171e-01 dTheta_rel=8.388e-02 time=2.631e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=002] dW_rel=1.062e-01 dTheta_rel=1.977e-02 time=1.853e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=003] dW_rel=1.269e-01 dTheta_rel=1.709e-02 time=1.433e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=004] dW_rel=3.793e-02 dTheta_rel=1.203e-02 time=1.390e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=005] dW_rel=2.382e-02 dTheta_rel=9.420e-03 time=1.097e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=006] dW_rel=9.642e-02 dTheta_rel=3.352e-02 time=1.372e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=007] dW_rel=1.482e-02 dTheta_rel=1.190e-02 time=1.205e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=008] dW_rel=1.086e-02 dTheta_rel=7.258e-03 time=1.082e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=009] dW_rel=6.806e-02 dTheta_rel=6.524e-03 time=1.229e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=010] dW_rel=2.376e-02 dTheta_rel=3.832e-03 time=1.020e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=011] dW_rel=1.126e-02 dTheta_rel=2.049e-03 time=8.972e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=012] dW_rel=7.652e-03 dTheta_rel=1.144e-03 time=7.338e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=013] dW_rel=5.741e-03 dTheta_rel=7.454e-04 time=6.404e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=014] dW_rel=8.669e-02 dTheta_rel=2.594e-03 time=1.062e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=015] dW_rel=1.170e-02 dTheta_rel=3.616e-03 time=1.264e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=016] dW_rel=7.876e-03 dTheta_rel=1.845e-03 time=9.233e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=017] dW_rel=5.649e-03 dTheta_rel=8.460e-04 time=5.447e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=018] dW_rel=4.251e-03 dTheta_rel=4.966e-04 time=6.370e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=019] dW_rel=3.291e-03 dTheta_rel=4.087e-04 time=5.610e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=020] dW_rel=7.504e-02 dTheta_rel=9.172e-04 time=8.994e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=021] dW_rel=4.454e-02 dTheta_rel=1.640e-03 time=8.880e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=022] dW_rel=4.498e-02 dTheta_rel=1.260e-02 time=1.613e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=023] dW_rel=4.910e-02 dTheta_rel=1.025e-02 time=1.177e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=024] dW_rel=1.543e-02 dTheta_rel=5.719e-03 time=9.460e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=025] dW_rel=8.107e-02 dTheta_rel=8.919e-03 time=1.229e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=026] dW_rel=4.479e-03 dTheta_rel=2.150e-03 time=1.119e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=027] dW_rel=2.500e-03 dTheta_rel=9.286e-04 time=8.642e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=028] dW_rel=1.569e-03 dTheta_rel=5.092e-04 time=8.423e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=029] dW_rel=1.085e-03 dTheta_rel=3.084e-04 time=5.609e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=030] dW_rel=7.869e-04 dTheta_rel=2.051e-04 time=4.658e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=031] dW_rel=5.875e-04 dTheta_rel=1.431e-04 time=5.048e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=032] dW_rel=4.474e-04 dTheta_rel=1.042e-04 time=3.142e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=033] dW_rel=3.450e-04 dTheta_rel=7.981e-05 time=3.013e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=034] dW_rel=2.682e-04 dTheta_rel=6.406e-05 time=3.516e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=035] dW_rel=2.092e-04 dTheta_rel=5.099e-05 time=2.963e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=036] dW_rel=1.639e-04 dTheta_rel=4.025e-05 time=2.040e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=037] dW_rel=1.281e-04 dTheta_rel=3.236e-05 time=2.414e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=038] dW_rel=1.011e-04 dTheta_rel=2.774e-05 time=3.239e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[auto_rank prune] p=8 -> 7 (dropped factors [7])\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=001] dW_rel=2.877e-01 dTheta_rel=9.598e-02 time=3.005e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=002] dW_rel=1.178e-01 dTheta_rel=2.802e-02 time=1.665e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=003] dW_rel=2.279e-02 dTheta_rel=6.407e-03 time=1.130e+02\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=004] dW_rel=1.265e-02 dTheta_rel=2.698e-03 time=7.828e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=005] dW_rel=7.551e-03 dTheta_rel=1.545e-03 time=6.819e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=006] dW_rel=4.718e-03 dTheta_rel=9.438e-04 time=5.467e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=007] dW_rel=3.048e-03 dTheta_rel=5.661e-04 time=4.435e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=008] dW_rel=2.039e-03 dTheta_rel=3.307e-04 time=4.407e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=009] dW_rel=1.408e-03 dTheta_rel=1.945e-04 time=3.429e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=010] dW_rel=1.006e-03 dTheta_rel=1.306e-04 time=2.840e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=011] dW_rel=7.344e-04 dTheta_rel=1.044e-04 time=2.690e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=012] dW_rel=5.496e-04 dTheta_rel=9.075e-05 time=3.719e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=013] dW_rel=4.155e-04 dTheta_rel=8.234e-05 time=2.347e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=014] dW_rel=3.161e-04 dTheta_rel=7.395e-05 time=1.910e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=015] dW_rel=2.413e-04 dTheta_rel=6.326e-05 time=1.871e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=016] dW_rel=1.839e-04 dTheta_rel=5.513e-05 time=1.953e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=017] dW_rel=1.415e-04 dTheta_rel=4.816e-05 time=1.874e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sweep=018] dW_rel=1.065e-04 dTheta_rel=3.972e-05 time=1.615e+01\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Runtime: 6613.8s | programs selected: 7\n",
      "Saved: Results/stgp/Inj_PT_L/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.002,\n",
    "    backfit_tol=1e-4,\n",
    "    prune_energy_frac=0.005,\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",
    "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}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "attach-scores",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-31T11:29:38.203369Z",
     "iopub.status.busy": "2026-05-31T11:29:38.203100Z",
     "iopub.status.idle": "2026-05-31T11:30:21.584193Z",
     "shell.execute_reply": "2026-05-31T11:30:21.582901Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved: Results/stgp/Inj_PT_L/adata_with_scores.h5ad\n",
      "Saved: Results/stgp/Inj_PT_L/W.csv\n"
     ]
    }
   ],
   "source": [
    "ADATA_PATH = OUT_DIR / \"adata_with_scores.h5ad\"\n",
    "\n",
    "adata = adata_inj_pt.copy()\n",
    "all_idx = np.concatenate(idx_sorted)\n",
    "H_arr = np.empty_like(res[\"H\"])\n",
    "b_arr = np.empty_like(res[\"b\"])\n",
    "H_arr[all_idx] = res[\"H\"]\n",
    "b_arr[all_idx] = res[\"b\"]\n",
    "\n",
    "adata.obsm[\"X_stgp\"] = H_arr.astype(np.float32)\n",
    "adata.obsm[\"X_stgp_spatial\"] = b_arr.astype(np.float32)\n",
    "adata.uns[\"stgp\"] = dict(\n",
    "    groups=slices.tolist(),\n",
    "    ages=ages.tolist(),\n",
    "    gamma_age=float(res[\"gamma_age\"]),\n",
    "    gamma_spa=float(res[\"gamma_spa\"]),\n",
    "    p_selected=res[\"W\"].shape[0],\n",
    "    alpha=np.asarray(res[\"alpha\"]).tolist(),\n",
    "    alpha_lower=np.asarray(res[\"alpha_lower\"]).tolist(),\n",
    "    alpha_upper=np.asarray(res[\"alpha_upper\"]).tolist(),\n",
    "    theta=np.asarray(res[\"theta\"]).tolist(),\n",
    "    sigma2e=float(res.get(\"sigma2e\", np.nan)),\n",
    ")\n",
    "adata.write_h5ad(str(ADATA_PATH), compression=\"gzip\")\n",
    "print(f\"Saved: {ADATA_PATH}\")\n",
    "\n",
    "p_sel = res[\"W\"].shape[0]\n",
    "W_df = pd.DataFrame(\n",
    "    res[\"W\"],\n",
    "    index=[f\"stGP{j + 1}\" for j in range(p_sel)],\n",
    "    columns=adata.var_names.astype(str),\n",
    ")\n",
    "W_df.to_csv(OUT_DIR / \"W.csv\")\n",
    "print(f\"Saved: {OUT_DIR / 'W.csv'}\")"
   ]
  }
 ],
 "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
}