From 2cf7acd44479a550600cb1511626bc3febcdeb7b Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Sun, 19 Apr 2026 00:28:47 +0000
Subject: [PATCH 1/6] Add horseshoe MCMC demo notebook with math, diagnostics,
 and exercises

Agent-Logs-Url: https://github.com/MacroMagic/bayesian_feature_selection/sessions/138b11ab-6029-445a-a1e3-7ab226609eee

Co-authored-by: MacroMagic <162652138+MacroMagic@users.noreply.github.com>
---
 ...orseshoe_mcmc_feature_selection_demo.ipynb | 430 ++++++++++++++++++
 1 file changed, 430 insertions(+)
 create mode 100644 demo/horseshoe_mcmc_feature_selection_demo.ipynb

diff --git a/demo/horseshoe_mcmc_feature_selection_demo.ipynb b/demo/horseshoe_mcmc_feature_selection_demo.ipynb
new file mode 100644
index 0000000..bf8f90c
--- /dev/null
+++ b/demo/horseshoe_mcmc_feature_selection_demo.ipynb
@@ -0,0 +1,430 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "1769b0f2",
+   "metadata": {},
+   "source": [
+    "# Bayesian Feature Selection Demo (MCMC + Horseshoe Prior)\n",
+    "\n",
+    "This notebook demonstrates Bayesian feature selection with the **regularized horseshoe prior** using the tools in this repository.\n",
+    "\n",
+    "## Learning goals\n",
+    "1. Understand the math behind horseshoe shrinkage and sparsity.\n",
+    "2. Run MCMC inference with `HorseshoeGLM` on simulated sparse data.\n",
+    "3. Evaluate convergence and posterior behavior with **ArviZ**.\n",
+    "4. Inspect posterior betas and feature selection quality.\n",
+    "5. See a GPU-ready workflow template for real datasets.\n",
+    "6. Practice with hands-on exercises that reveal strengths and limits of the method.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fce41faa",
+   "metadata": {},
+   "source": [
+    "## 1) Mathematical intuition: why the horseshoe induces sparsity\n",
+    "\n",
+    "For feature coefficients \\(\beta_j\\), the (regularized) horseshoe hierarchy is:\n",
+    "\n",
+    "\\[\n",
+    "\tau \\sim \text{Half-Cauchy}(0, \tau_0), \\quad\n",
+    "\\lambda_j \\sim \text{Half-Cauchy}(0, 1), \\quad\n",
+    "c^2 \\sim \text{Inverse-Gamma}(1,1)\n",
+    "\\]\n",
+    "\\[\n",
+    "\tilde{\\lambda}_j = \\sqrt{\f",
+    "rac{c^2\\lambda_j^2}{c^2 + \tau^2\\lambda_j^2}}, \\quad\n",
+    "\beta_j \\mid \tau,\\lambda_j,c^2 \\sim \\mathcal{N}(0,\tau^2\tilde{\\lambda}_j^2)\n",
+    "\\]\n",
+    "\n",
+    "For a Gaussian GLM:\n",
+    "\\[\n",
+    "\\eta_i = \u0007lpha + x_i^\top\beta, \\quad\n",
+    "y_i \\sim \\mathcal{N}(\\eta_i, \\sigma^2)\n",
+    "\\]\n",
+    "\n",
+    "### Key sparsity mechanism\n",
+    "- **Global shrinkage** \\(\tau\\): pushes all coefficients toward zero.\n",
+    "- **Local shrinkage** \\(\\lambda_j\\): lets a few relevant features escape shrinkage.\n",
+    "- **Slab regularization** \\(c^2\\): avoids unrealistically huge coefficients.\n",
+    "\n",
+    "A useful view is the shrinkage factor \\(\\kappa_j = 1/(1 + \tau^2\\lambda_j^2)\\):\n",
+    "- Noise features often have \\(\\kappa_j \u0007pprox 1\\) (strong shrinkage to near zero).\n",
+    "- True signals often have smaller \\(\\kappa_j\\), preserving non-zero effects.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "62474249",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import warnings\n",
+    "warnings.filterwarnings('ignore')\n",
+    "\n",
+    "import jax\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import matplotlib.pyplot as plt\n",
+    "import arviz as az\n",
+    "\n",
+    "from sklearn.metrics import precision_score, recall_score, f1_score\n",
+    "from sklearn.preprocessing import StandardScaler\n",
+    "from sklearn.datasets import load_breast_cancer\n",
+    "\n",
+    "from bayesian_feature_selection import HorseshoeGLM, InferenceConfig\n",
+    "\n",
+    "np.random.seed(42)\n",
+    "print('JAX devices:', jax.devices())\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "74623f58",
+   "metadata": {},
+   "source": [
+    "## 2) Simulated sparse data\n",
+    "\n",
+    "We create a high-dimensional-ish problem with many irrelevant features and a few true signals.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e7674182",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Simulated data with sparse truth\n",
+    "n, p = 220, 60\n",
+    "n_signal = 6\n",
+    "signal_idx = np.array([0, 3, 7, 13, 27, 42])\n",
+    "\n",
+    "X = np.random.normal(size=(n, p))\n",
+    "true_beta = np.zeros(p)\n",
+    "true_beta[signal_idx] = np.array([2.8, -2.2, 1.8, -1.6, 1.2, 0.9])\n",
+    "\n",
+    "noise_sd = 1.0\n",
+    "y = X @ true_beta + np.random.normal(0, noise_sd, size=n)\n",
+    "\n",
+    "# standardize features for stable inference\n",
+    "X = StandardScaler().fit_transform(X)\n",
+    "\n",
+    "print(f'n={n}, p={p}, true non-zero features={len(signal_idx)}')\n",
+    "\n",
+    "plt.figure(figsize=(10, 3.5))\n",
+    "plt.stem(np.arange(p), true_beta, basefmt='k-')\n",
+    "plt.title('Ground-truth coefficients (sparse)')\n",
+    "plt.xlabel('Feature index')\n",
+    "plt.ylabel('True beta')\n",
+    "plt.grid(alpha=0.3)\n",
+    "plt.show()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b31dd125",
+   "metadata": {},
+   "source": [
+    "## 3) Fit Horseshoe GLM with MCMC\n",
+    "\n",
+    "We focus on MCMC inference (`method='mcmc'`) for full posterior uncertainty.\n",
+    "\n",
+    "Rule-of-thumb initialization for global scale:\n",
+    "\\[\n",
+    "\tau_0 \u0007pprox \f",
+    "rac{p_0}{p-p_0}\\cdot\f",
+    "rac{1}{\\sqrt{n}}\n",
+    "\\]\n",
+    "where \\(p_0\\) is expected number of relevant features.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "003b1d3a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "expected_p0 = n_signal\n",
+    "scale_global = (expected_p0 / max(p - expected_p0, 1)) / np.sqrt(n)\n",
+    "scale_global = float(np.clip(scale_global, 0.05, 1.0))\n",
+    "\n",
+    "use_gpu = any(d.platform == 'gpu' for d in jax.devices())\n",
+    "print('Using GPU?', use_gpu)\n",
+    "print('scale_global:', round(scale_global, 4))\n",
+    "\n",
+    "model = HorseshoeGLM(family='gaussian', scale_global=scale_global)\n",
+    "config = InferenceConfig(\n",
+    "    method='mcmc',\n",
+    "    num_warmup=300,\n",
+    "    num_samples=500,\n",
+    "    num_chains=1,\n",
+    "    use_gpu=use_gpu,\n",
+    "    progress_bar=False,\n",
+    ")\n",
+    "\n",
+    "model.fit(X, y, config=config)\n",
+    "print('MCMC complete.')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c4b375be",
+   "metadata": {},
+   "source": [
+    "## 4) Feature selection results\n",
+    "\n",
+    "We use posterior inclusion from beta samples (default behavior in this repo) and compare with ground truth.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4e5d1c79",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "importance = model.get_feature_importance(method='beta', threshold=0.5)\n",
+    "importance = importance.sort_values('beta_inclusion_prob', ascending=False).reset_index(drop=True)\n",
+    "\n",
+    "selected_idx = importance.loc[importance['selected'], 'feature_idx'].to_numpy()\n",
+    "true_mask = np.isin(np.arange(p), signal_idx).astype(int)\n",
+    "pred_mask = np.isin(np.arange(p), selected_idx).astype(int)\n",
+    "\n",
+    "precision = precision_score(true_mask, pred_mask, zero_division=0)\n",
+    "recall = recall_score(true_mask, pred_mask, zero_division=0)\n",
+    "f1 = f1_score(true_mask, pred_mask, zero_division=0)\n",
+    "\n",
+    "print(f'Selected features: {len(selected_idx)}')\n",
+    "print(f'Precision: {precision:.3f} | Recall: {recall:.3f} | F1: {f1:.3f}')\n",
+    "\n",
+    "show_cols = ['feature_idx', 'beta_mean', 'beta_lower_95', 'beta_upper_95', 'beta_inclusion_prob', 'selected']\n",
+    "display(importance[show_cols].head(12))\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2a2f824c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Visual comparison: posterior mean vs truth\n",
+    "posterior_mean = np.zeros(p)\n",
+    "posterior_mean[importance['feature_idx'].to_numpy()] = importance['beta_mean'].to_numpy()\n",
+    "\n",
+    "plt.figure(figsize=(11, 4))\n",
+    "plt.plot(true_beta, 'o-', label='true beta', alpha=0.8)\n",
+    "plt.plot(posterior_mean, 's--', label='posterior mean beta', alpha=0.8)\n",
+    "plt.axhline(0, color='k', lw=1)\n",
+    "plt.title('True vs posterior-mean coefficients')\n",
+    "plt.xlabel('Feature index')\n",
+    "plt.ylabel('Coefficient value')\n",
+    "plt.legend()\n",
+    "plt.grid(alpha=0.25)\n",
+    "plt.show()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cfa75552",
+   "metadata": {},
+   "source": [
+    "## 5) ArviZ diagnostics: convergence + beta visualization\n",
+    "\n",
+    "We convert NumPyro output to `InferenceData`, inspect R-hat/ESS, and visualize chains/posteriors.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d3a58def",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "idata = az.from_numpyro(model.mcmc)\n",
+    "\n",
+    "# Summaries for convergence diagnostics\n",
+    "summary = az.summary(idata, var_names=['tau', 'c2', 'beta'])\n",
+    "summary.head()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "19d74972",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Basic convergence checks\n",
+    "rhat_max = summary['r_hat'].dropna().max()\n",
+    "ess_min = summary['ess_bulk'].dropna().min()\n",
+    "\n",
+    "print(f'max R-hat: {rhat_max:.4f}')\n",
+    "print(f'min ESS (bulk): {ess_min:.1f}')\n",
+    "print('Convergence guideline: R-hat near 1.00 and reasonably large ESS.')\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2b02a78a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Trace plot for global parameters\n",
+    "az.plot_trace(idata, var_names=['tau', 'c2'])\n",
+    "plt.tight_layout()\n",
+    "plt.show()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6d33b5c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Beta visualization with ArviZ for top features by inclusion probability\n",
+    "beta_dim = [c for c in idata.posterior.coords if c.startswith('beta_dim_')][0]\n",
+    "top_k = 10\n",
+    "top_features = importance['feature_idx'].head(top_k).tolist()\n",
+    "\n",
+    "az.plot_forest(\n",
+    "    idata,\n",
+    "    var_names=['beta'],\n",
+    "    coords={beta_dim: top_features},\n",
+    "    combined=True,\n",
+    "    hdi_prob=0.95,\n",
+    "    figsize=(8, 5),\n",
+    ")\n",
+    "plt.title('95% HDI for top posterior beta coefficients')\n",
+    "plt.show()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "382f712c",
+   "metadata": {},
+   "source": [
+    "## 6) Real data + GPU-ready workflow (optional run)\n",
+    "\n",
+    "Below is a compact template using a real dataset (`sklearn` breast cancer data) and automatic GPU detection.\n",
+    "\n",
+    "- Set `RUN_REAL_DATA_DEMO = True` to execute.\n",
+    "- Increase MCMC samples/chains for production analysis.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "33ab5273",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "RUN_REAL_DATA_DEMO = False\n",
+    "\n",
+    "if RUN_REAL_DATA_DEMO:\n",
+    "    ds = load_breast_cancer()\n",
+    "    X_real = StandardScaler().fit_transform(ds.data)\n",
+    "    y_real = ds.target.astype(float)\n",
+    "\n",
+    "    use_gpu_real = any(d.platform == 'gpu' for d in jax.devices())\n",
+    "    model_real = HorseshoeGLM(family='binomial', scale_global=0.25)\n",
+    "\n",
+    "    config_real = InferenceConfig(\n",
+    "        method='mcmc',\n",
+    "        num_warmup=200,\n",
+    "        num_samples=300,\n",
+    "        num_chains=1,\n",
+    "        use_gpu=use_gpu_real,\n",
+    "        progress_bar=False,\n",
+    "    )\n",
+    "\n",
+    "    model_real.fit(X_real, y_real, config=config_real)\n",
+    "    idata_real = az.from_numpyro(model_real.mcmc)\n",
+    "\n",
+    "    print('Real-data model fit complete. Device GPU?', use_gpu_real)\n",
+    "    display(az.summary(idata_real, var_names=['tau', 'c2']).head())\n",
+    "else:\n",
+    "    print('Skipped. Set RUN_REAL_DATA_DEMO = True to run the real-data section.')\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "571f2125",
+   "metadata": {},
+   "source": [
+    "## 7) Hands-on exercises and challenge prompts\n",
+    "\n",
+    "Try these short activities to build intuition and discover practical tradeoffs.\n",
+    "\n",
+    "### Exercise A — Sparsity calibration (global shrinkage)\n",
+    "1. Sweep `scale_global` over `[0.05, 0.1, 0.2, 0.5, 1.0]`.\n",
+    "2. Track precision/recall/F1 of selected features.\n",
+    "3. Explain the bias-variance tradeoff you observe.\n",
+    "\n",
+    "### Exercise B — Signal-to-noise stress test\n",
+    "1. Keep the same true coefficients, vary `noise_sd` from `0.3` to `2.5`.\n",
+    "2. Plot recovery metrics vs noise level.\n",
+    "3. Identify where horseshoe starts missing weak signals.\n",
+    "\n",
+    "### Exercise C — Correlated predictors (method limitation)\n",
+    "1. Create blocks of strongly correlated features.\n",
+    "2. Place a true signal in one feature per block.\n",
+    "3. Compare selection under `method='beta'`, `'lambda'`, and `'both'`.\n",
+    "4. Discuss why posterior mass can spread across correlated variables.\n",
+    "\n",
+    "### Exercise D — p >> n scenario (state-of-the-art use case)\n",
+    "1. Set `n=80`, `p=500`, true signals = 8.\n",
+    "2. Use MCMC with fewer warmup/samples first, then increase.\n",
+    "3. Compare runtime, convergence, and recovery quality.\n",
+    "\n",
+    "### Exercise E — Chain quality diagnostics\n",
+    "1. Intentionally set very short warmup/samples.\n",
+    "2. Use ArviZ to inspect R-hat, ESS, and trace behavior.\n",
+    "3. Propose a principled tuning plan to improve inference.\n",
+    "\n",
+    "### Exercise F — Real-data challenge\n",
+    "1. Run the optional breast-cancer section with `RUN_REAL_DATA_DEMO=True`.\n",
+    "2. Rank features by inclusion probability.\n",
+    "3. Compare selected features to coefficients from a lasso baseline.\n",
+    "4. Reflect: when does Bayesian uncertainty add value over point-estimate methods?\n",
+    "\n",
+    "### Exercise G — Prior robustness challenge\n",
+    "1. Replace true coefficients with one very large signal + several weak ones.\n",
+    "2. Evaluate whether regularized horseshoe retains weak effects while controlling noise.\n",
+    "3. Summarize pros/cons for scientific discovery vs pure prediction tasks.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ecd70afc",
+   "metadata": {},
+   "source": [
+    "## 8) Takeaways\n",
+    "\n",
+    "- The horseshoe prior combines **aggressive noise shrinkage** with **heavy-tail signal retention**.\n",
+    "- MCMC gives richer uncertainty quantification for feature relevance than point-estimate methods.\n",
+    "- ArviZ diagnostics are essential before trusting selected features.\n",
+    "- Performance depends on prior scale, SNR, and predictor correlation structure.\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

From 79b1527a55c6e97f67b6e3623e4425aad5726d36 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Sun, 19 Apr 2026 00:31:37 +0000
Subject: [PATCH 2/6] Fix ArviZ notebook plotting compatibility and validate
 demo execution

Agent-Logs-Url: https://github.com/MacroMagic/bayesian_feature_selection/sessions/138b11ab-6029-445a-a1e3-7ab226609eee

Co-authored-by: MacroMagic <162652138+MacroMagic@users.noreply.github.com>
---
 demo/horseshoe_mcmc_feature_selection_demo.ipynb | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/demo/horseshoe_mcmc_feature_selection_demo.ipynb b/demo/horseshoe_mcmc_feature_selection_demo.ipynb
index bf8f90c..5707e7f 100644
--- a/demo/horseshoe_mcmc_feature_selection_demo.ipynb
+++ b/demo/horseshoe_mcmc_feature_selection_demo.ipynb
@@ -290,18 +290,15 @@
    "source": [
     "# Beta visualization with ArviZ for top features by inclusion probability\n",
     "beta_dim = [c for c in idata.posterior.coords if c.startswith('beta_dim_')][0]\n",
-    "top_k = 10\n",
+    "top_k = 8\n",
     "top_features = importance['feature_idx'].head(top_k).tolist()\n",
     "\n",
-    "az.plot_forest(\n",
+    "az.plot_trace(\n",
     "    idata,\n",
     "    var_names=['beta'],\n",
     "    coords={beta_dim: top_features},\n",
-    "    combined=True,\n",
-    "    hdi_prob=0.95,\n",
-    "    figsize=(8, 5),\n",
     ")\n",
-    "plt.title('95% HDI for top posterior beta coefficients')\n",
+    "plt.tight_layout()\n",
     "plt.show()\n"
    ]
   },

From 208719cbd7484cf0fbaa3828db6f722e96888497 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Sun, 19 Apr 2026 00:33:32 +0000
Subject: [PATCH 3/6] Fix LaTeX rendering in demo notebook equations

Agent-Logs-Url: https://github.com/MacroMagic/bayesian_feature_selection/sessions/138b11ab-6029-445a-a1e3-7ab226609eee

Co-authored-by: MacroMagic <162652138+MacroMagic@users.noreply.github.com>
---
 ...orseshoe_mcmc_feature_selection_demo.ipynb | 63 +++++++++----------
 1 file changed, 30 insertions(+), 33 deletions(-)

diff --git a/demo/horseshoe_mcmc_feature_selection_demo.ipynb b/demo/horseshoe_mcmc_feature_selection_demo.ipynb
index 5707e7f..930d86d 100644
--- a/demo/horseshoe_mcmc_feature_selection_demo.ipynb
+++ b/demo/horseshoe_mcmc_feature_selection_demo.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "1769b0f2",
+   "id": "bbae7ab1",
    "metadata": {},
    "source": [
     "# Bayesian Feature Selection Demo (MCMC + Horseshoe Prior)\n",
@@ -20,44 +20,43 @@
   },
   {
    "cell_type": "markdown",
-   "id": "fce41faa",
+   "id": "934d4419",
    "metadata": {},
    "source": [
     "## 1) Mathematical intuition: why the horseshoe induces sparsity\n",
     "\n",
-    "For feature coefficients \\(\beta_j\\), the (regularized) horseshoe hierarchy is:\n",
+    "For feature coefficients \\(\\beta_j\\), the (regularized) horseshoe hierarchy is:\n",
     "\n",
     "\\[\n",
-    "\tau \\sim \text{Half-Cauchy}(0, \tau_0), \\quad\n",
-    "\\lambda_j \\sim \text{Half-Cauchy}(0, 1), \\quad\n",
-    "c^2 \\sim \text{Inverse-Gamma}(1,1)\n",
+    "\\tau \\sim \\text{Half-Cauchy}(0, \\tau_0), \\quad\n",
+    "\\lambda_j \\sim \\text{Half-Cauchy}(0, 1), \\quad\n",
+    "c^2 \\sim \\text{Inverse-Gamma}(1,1)\n",
     "\\]\n",
     "\\[\n",
-    "\tilde{\\lambda}_j = \\sqrt{\f",
-    "rac{c^2\\lambda_j^2}{c^2 + \tau^2\\lambda_j^2}}, \\quad\n",
-    "\beta_j \\mid \tau,\\lambda_j,c^2 \\sim \\mathcal{N}(0,\tau^2\tilde{\\lambda}_j^2)\n",
+    "\\tilde{\\lambda}_j = \\sqrt{\\frac{c^2\\lambda_j^2}{c^2 + \\tau^2\\lambda_j^2}}, \\quad\n",
+    "\\beta_j \\mid \\tau,\\lambda_j,c^2 \\sim \\mathcal{N}(0,\\tau^2\\tilde{\\lambda}_j^2)\n",
     "\\]\n",
     "\n",
     "For a Gaussian GLM:\n",
     "\\[\n",
-    "\\eta_i = \u0007lpha + x_i^\top\beta, \\quad\n",
+    "\\eta_i = \\alpha + x_i^\\top\\beta, \\quad\n",
     "y_i \\sim \\mathcal{N}(\\eta_i, \\sigma^2)\n",
     "\\]\n",
     "\n",
     "### Key sparsity mechanism\n",
-    "- **Global shrinkage** \\(\tau\\): pushes all coefficients toward zero.\n",
+    "- **Global shrinkage** \\(\\tau\\): pushes all coefficients toward zero.\n",
     "- **Local shrinkage** \\(\\lambda_j\\): lets a few relevant features escape shrinkage.\n",
     "- **Slab regularization** \\(c^2\\): avoids unrealistically huge coefficients.\n",
     "\n",
-    "A useful view is the shrinkage factor \\(\\kappa_j = 1/(1 + \tau^2\\lambda_j^2)\\):\n",
-    "- Noise features often have \\(\\kappa_j \u0007pprox 1\\) (strong shrinkage to near zero).\n",
+    "A useful view is the shrinkage factor \\(\\kappa_j = 1/(1 + \\tau^2\\lambda_j^2)\\):\n",
+    "- Noise features often have \\(\\kappa_j \\approx 1\\) (strong shrinkage to near zero).\n",
     "- True signals often have smaller \\(\\kappa_j\\), preserving non-zero effects.\n"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "62474249",
+   "id": "5353239c",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -82,7 +81,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "74623f58",
+   "id": "93bfecd0",
    "metadata": {},
    "source": [
     "## 2) Simulated sparse data\n",
@@ -93,7 +92,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e7674182",
+   "id": "0052282d",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -125,7 +124,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b31dd125",
+   "id": "593bcdb3",
    "metadata": {},
    "source": [
     "## 3) Fit Horseshoe GLM with MCMC\n",
@@ -134,9 +133,7 @@
     "\n",
     "Rule-of-thumb initialization for global scale:\n",
     "\\[\n",
-    "\tau_0 \u0007pprox \f",
-    "rac{p_0}{p-p_0}\\cdot\f",
-    "rac{1}{\\sqrt{n}}\n",
+    "\\tau_0 \\approx \\frac{p_0}{p-p_0}\\cdot\\frac{1}{\\sqrt{n}}\n",
     "\\]\n",
     "where \\(p_0\\) is expected number of relevant features.\n"
    ]
@@ -144,7 +141,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "003b1d3a",
+   "id": "1419405a",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -172,7 +169,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c4b375be",
+   "id": "08812a5b",
    "metadata": {},
    "source": [
     "## 4) Feature selection results\n",
@@ -183,7 +180,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4e5d1c79",
+   "id": "642f9693",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -208,7 +205,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2a2f824c",
+   "id": "0d767a30",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -230,7 +227,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "cfa75552",
+   "id": "aa46a47a",
    "metadata": {},
    "source": [
     "## 5) ArviZ diagnostics: convergence + beta visualization\n",
@@ -241,7 +238,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d3a58def",
+   "id": "a71b8815",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -255,7 +252,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "19d74972",
+   "id": "65fb3a31",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -271,7 +268,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2b02a78a",
+   "id": "589ca7d4",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -284,7 +281,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "6d33b5c3",
+   "id": "de46ceae",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -304,7 +301,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "382f712c",
+   "id": "85175830",
    "metadata": {},
    "source": [
     "## 6) Real data + GPU-ready workflow (optional run)\n",
@@ -318,7 +315,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "33ab5273",
+   "id": "42689732",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -352,7 +349,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "571f2125",
+   "id": "c8343096",
    "metadata": {},
    "source": [
     "## 7) Hands-on exercises and challenge prompts\n",
@@ -399,7 +396,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ecd70afc",
+   "id": "9d776645",
    "metadata": {},
    "source": [
     "## 8) Takeaways\n",

From c9a5c57ac84c3d57b4a4b48af31a6186146104b9 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Sun, 19 Apr 2026 00:35:44 +0000
Subject: [PATCH 4/6] Improve notebook diagnostics robustness and use
 multi-chain MCMC examples

Agent-Logs-Url: https://github.com/MacroMagic/bayesian_feature_selection/sessions/138b11ab-6029-445a-a1e3-7ab226609eee

Co-authored-by: MacroMagic <162652138+MacroMagic@users.noreply.github.com>
---
 demo/horseshoe_mcmc_feature_selection_demo.ipynb | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/demo/horseshoe_mcmc_feature_selection_demo.ipynb b/demo/horseshoe_mcmc_feature_selection_demo.ipynb
index 930d86d..7c67647 100644
--- a/demo/horseshoe_mcmc_feature_selection_demo.ipynb
+++ b/demo/horseshoe_mcmc_feature_selection_demo.ipynb
@@ -158,7 +158,7 @@
     "    method='mcmc',\n",
     "    num_warmup=300,\n",
     "    num_samples=500,\n",
-    "    num_chains=1,\n",
+    "    num_chains=2,\n",
     "    use_gpu=use_gpu,\n",
     "    progress_bar=False,\n",
     ")\n",
@@ -257,8 +257,11 @@
    "outputs": [],
    "source": [
     "# Basic convergence checks\n",
-    "rhat_max = summary['r_hat'].dropna().max()\n",
-    "ess_min = summary['ess_bulk'].dropna().min()\n",
+    "rhat_vals = pd.to_numeric(summary['r_hat'], errors='coerce').dropna()\n",
+    "ess_vals = pd.to_numeric(summary['ess_bulk'], errors='coerce').dropna()\n",
+    "\n",
+    "rhat_max = float(rhat_vals.max()) if len(rhat_vals) else float('nan')\n",
+    "ess_min = float(ess_vals.min()) if len(ess_vals) else float('nan')\n",
     "\n",
     "print(f'max R-hat: {rhat_max:.4f}')\n",
     "print(f'min ESS (bulk): {ess_min:.1f}')\n",
@@ -333,7 +336,7 @@
     "        method='mcmc',\n",
     "        num_warmup=200,\n",
     "        num_samples=300,\n",
-    "        num_chains=1,\n",
+    "        num_chains=2,\n",
     "        use_gpu=use_gpu_real,\n",
     "        progress_bar=False,\n",
     "    )\n",

From 65482a755b8d951eef9fb786fca986899032acc0 Mon Sep 17 00:00:00 2001
From: MacroMagic <162652138+MacroMagic@users.noreply.github.com>
Date: Sat, 18 Apr 2026 21:12:03 -0400
Subject: [PATCH 5/6] Update horseshoe_mcmc_feature_selection_demo.ipynb

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 demo/horseshoe_mcmc_feature_selection_demo.ipynb | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/demo/horseshoe_mcmc_feature_selection_demo.ipynb b/demo/horseshoe_mcmc_feature_selection_demo.ipynb
index 7c67647..e59cfcd 100644
--- a/demo/horseshoe_mcmc_feature_selection_demo.ipynb
+++ b/demo/horseshoe_mcmc_feature_selection_demo.ipynb
@@ -61,7 +61,8 @@
    "outputs": [],
    "source": [
     "import warnings\n",
-    "warnings.filterwarnings('ignore')\n",
+    "warnings.filterwarnings('ignore', category=FutureWarning)\n",
+    "warnings.filterwarnings('ignore', category=DeprecationWarning)\n",
     "\n",
     "import jax\n",
     "import numpy as np\n",

From c9649abcddae214f4ff61728989e0ddb196203f5 Mon Sep 17 00:00:00 2001
From: MacroMagic <162652138+MacroMagic@users.noreply.github.com>
Date: Sat, 18 Apr 2026 21:12:35 -0400
Subject: [PATCH 6/6] Update horseshoe_mcmc_feature_selection_demo.ipynb

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 demo/horseshoe_mcmc_feature_selection_demo.ipynb | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/demo/horseshoe_mcmc_feature_selection_demo.ipynb b/demo/horseshoe_mcmc_feature_selection_demo.ipynb
index e59cfcd..cb4975b 100644
--- a/demo/horseshoe_mcmc_feature_selection_demo.ipynb
+++ b/demo/horseshoe_mcmc_feature_selection_demo.ipynb
@@ -290,7 +290,15 @@
    "outputs": [],
    "source": [
     "# Beta visualization with ArviZ for top features by inclusion probability\n",
-    "beta_dim = [c for c in idata.posterior.coords if c.startswith('beta_dim_')][0]\n",
+    "beta_dims = [d for d in idata.posterior['beta'].dims if d not in ('chain', 'draw')]\n",
+    "if len(beta_dims) != 1:\n",
+    "    raise ValueError(\n",
+    "        f\"Expected 'beta' to have exactly one non-sampling dimension, found {beta_dims!r}\"\n",
+    "    )\n",
+    "beta_dim = beta_dims[0]\n",
+    "if beta_dim not in idata.posterior.coords:\n",
+    "    raise ValueError(f\"Derived beta dimension {beta_dim!r} is not present in posterior coords\")\n",
+    "\n",
     "top_k = 8\n",
     "top_features = importance['feature_idx'].head(top_k).tolist()\n",
     "\n",