{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d731b9d9",
   "metadata": {},
   "source": [
    "# Borrowing of strength in Bayesian hierarchical model\n",
    "We consider the problem of estimating the average number of patients admitted to\n",
    "ER per day from various hospitals across a city. A Bayesian hierarchical model\n",
    "(BHM) allows us to partially pool information across hospitals, reducing\n",
    "variance in small-sample groups and leveraging strength across hospitals to\n",
    "estimate the population-level parameters. Here, we investigate a Poisson-\n",
    "lognormal model.\n",
    "\n",
    "(a) Generate a dataset where the logarithm of the expected number of admissions\n",
    "    per day for hospital $j$, $ln \\lambda_j$, comes from a Normal distribution:\n",
    "\n",
    "$$ ln \\lambda_j \\sim \\mathcal{N}(\\mu_0,\\,\\sigma_0^{2}), \\text{ for } j = 1, ..., J $$\n",
    "\n",
    "The admission counts data for each hospital for day $i$ follow a Poisson\n",
    "distribution:\n",
    "\n",
    "$$ y_{ij} \\mid \\lambda_j \\sim \\text{Poisson}(\\lambda_j), \\text{ for } i = 1, ..., n_j $$\n",
    "\n",
    "where each hospital $j$ has a different reporting frequency, which leads to\n",
    "different $n_j$ for each. To model this, you may draw $n_j$ randomly with equal\n",
    "probability from the set $\\{52, 24, 12\\}$, representing respectively weekly,\n",
    "twice a month and monthly reporting. Do this once, and keep $n_j$ fixed,\n",
    "treating it as known, throughout the exercise.\n",
    "\n",
    "The DAG for this model is shown in the figure.\n",
    "\n",
    "![Figure 1: DAG for hospital admission BHM](./figure_1.pdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5655aaf",
   "metadata": {},
   "source": [
    "Since $\\mu_0$ is the typical log-rate\n",
    "across hospitals, we may choose for its prior\n",
    "$$Pr(\\mu_0) = \\mathcal{N}(\\log m, s^2)$$\n",
    "\n",
    "where $m$ is a plausible choice for the median number of admissions/day, and s\n",
    "controls its spread. Since $\\sigma_0$ controls the variability across hospitals,\n",
    "a prior choice that tends to reinforce shrinkage is $Pr(\\sigma_0) = \\exp(\\tau)$.\n",
    "\n",
    "Fix $m$, $s^2$, $\\tau$ to sensible values, and use an MCMC algorithm of your\n",
    "choice to sample from the model and produce posterior distribution on\n",
    "$\\{\\lambda_j\\}^J_{j=1}, \\mu _0, \\sigma _0$ (a sensible choice might be $J = 10$\n",
    "hospitals)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "49307215",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Callable, Tuple\n",
    "\n",
    "import corner\n",
    "import emcee\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from matplotlib.lines import Line2D\n",
    "from matplotlib.patches import Patch\n",
    "from scipy.special import gammaln\n",
    "from scipy.stats import pearsonr\n",
    "\n",
    "np.random.seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "af15e048",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "We have 10 hospitals reporting admission data.\n",
      "- Hospital number  1 reports admission data 12 times a year.\n",
      "- Hospital number  2 reports admission data 54 times a year.\n",
      "- Hospital number  3 reports admission data 12 times a year.\n",
      "- Hospital number  4 reports admission data 12 times a year.\n",
      "- Hospital number  5 reports admission data 54 times a year.\n",
      "- Hospital number  6 reports admission data 54 times a year.\n",
      "- Hospital number  7 reports admission data 12 times a year.\n",
      "- Hospital number  8 reports admission data 24 times a year.\n",
      "- Hospital number  9 reports admission data 12 times a year.\n",
      "- Hospital number 10 reports admission data 12 times a year.\n"
     ]
    }
   ],
   "source": [
    "J = 10\n",
    "reporting_frequency = np.random.choice(\n",
    "    [54, 24, 12],\n",
    "    J\n",
    ")\n",
    "\n",
    "print(f\"We have {J} hospitals reporting admission data.\")\n",
    "for j in range(J):\n",
    "    print(\n",
    "        f\"- Hospital number {j + 1:2} reports admission data \"\n",
    "        f\"{reporting_frequency[j]} times a year.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4a09b1d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_expected_admission_counts(mu: float, sigma: float, size: int):\n",
    "    \"\"\"Get expected number of admissions to a hospital at 'size' time points.\n",
    "    \n",
    "    The log of the admission rate comes from a normal distribution with mean\n",
    "        'mu' and standard deviation 'sigma'.\n",
    "    The expected number of admissions comes from a Poisson distribution with\n",
    "        parameter 'admission_rate'.\n",
    "    From the same Poisson distribution, 'size' time points will be sampled.\n",
    "    \"\"\"\n",
    "    log_admission_rate = np.random.normal(\n",
    "        mu, sigma\n",
    "    )\n",
    "    admission_rate = np.exp(log_admission_rate)\n",
    "    admission_counts = np.random.poisson(\n",
    "        admission_rate, size=size\n",
    "    )\n",
    "    return admission_counts, admission_rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2f6b2f3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_admission_count_matrix(\n",
    "    mu: float, sigma: float\n",
    ") -> Tuple[pd.DataFrame, List[float]]:\n",
    "    data = {}\n",
    "    true_lambdas = []\n",
    "    total_array_length = max(reporting_frequency)\n",
    "    for j in range(J):  # For each hospital\n",
    "        data_array_length = reporting_frequency[j]\n",
    "        admission_records, admission_rate = get_expected_admission_counts(\n",
    "            mu, sigma, size=data_array_length\n",
    "        )\n",
    "        true_lambdas.append(admission_rate)\n",
    "        padding = np.full(total_array_length - data_array_length, 0, dtype=int)\n",
    "        admission_records = np.concatenate([\n",
    "            admission_records,\n",
    "            padding\n",
    "        ])\n",
    "        data[f\"Hospital {j+1}\"] = admission_records\n",
    "    data = pd.DataFrame(data)\n",
    "    data.index.name = \"Time point\"\n",
    "    return data.transpose(), true_lambdas"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6015dc49",
   "metadata": {},
   "source": [
    "- Prior for $\\mu_0$: we choose $m$ (typical log-rate across hospitals) and $s$\n",
    "(a measure of the spread of $m$); $Pr(\\mu_0) = \\mathcal{N}(\\log m, s^2)$.\n",
    "- Prior for $\\sigma_0$: $Pr(\\sigma_0) = \\text{Exponential}(\\tau)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "216cdb28",
   "metadata": {},
   "outputs": [],
   "source": [
    "m = 85\n",
    "s = 0.4\n",
    "tau = 0.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4b5ee692",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time point    0    1    2    3    4    5    6    7    8    9   ...   44   45  \\\n",
      "Hospital 1   149  147  161  152  144  155  171  161  146  173  ...    0    0   \n",
      "Hospital 2   180  166  160  178  182  147  153  151  154  161  ...  150  152   \n",
      "Hospital 3   174  147  163  149  146  156  144  142  177  166  ...    0    0   \n",
      "Hospital 4   154  178  170  138  178  165  145  167  148  149  ...    0    0   \n",
      "Hospital 5   163  155  156  165  160  160  169  155  183  175  ...  143  168   \n",
      "Hospital 6   163  162  165  161  156  180  136  167  153  151  ...  142  165   \n",
      "Hospital 7   167  165  179  178  164  147  154  161  173  146  ...    0    0   \n",
      "Hospital 8   145  155  174  159  151  176  162  150  164  166  ...    0    0   \n",
      "Hospital 9   183  150  153  150  143  145  165  154  159  146  ...    0    0   \n",
      "Hospital 10  154  167  171  138  181  151  169  148  159  154  ...    0    0   \n",
      "\n",
      "Time point    46   47   48   49   50   51   52   53  \n",
      "Hospital 1     0    0    0    0    0    0    0    0  \n",
      "Hospital 2   160  151  149  150  170  168  165  161  \n",
      "Hospital 3     0    0    0    0    0    0    0    0  \n",
      "Hospital 4     0    0    0    0    0    0    0    0  \n",
      "Hospital 5   141  159  158  164  130  165  176  146  \n",
      "Hospital 6   163  155  176  194  199  178  162  161  \n",
      "Hospital 7     0    0    0    0    0    0    0    0  \n",
      "Hospital 8     0    0    0    0    0    0    0    0  \n",
      "Hospital 9     0    0    0    0    0    0    0    0  \n",
      "Hospital 10    0    0    0    0    0    0    0    0  \n",
      "\n",
      "[10 rows x 54 columns]\n"
     ]
    }
   ],
   "source": [
    "true_mu_0 = np.random.normal(np.log(m), s)\n",
    "true_sigma_0 = np.random.exponential(tau)\n",
    "\n",
    "df, true_lambdas = get_admission_count_matrix(mu=true_mu_0, sigma=true_sigma_0)\n",
    "\n",
    "print(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c8afdbc",
   "metadata": {},
   "source": [
    "We can use AIES (ensemble samplers with affine invariance), introduced by\n",
    "Goodman and Weare in 2010 and implemented by the `emcee` python package, to\n",
    "sample from the joint posterior without analytically deriving it.\n",
    "\n",
    "We have to write a function that takes in all hyper-parameters (both fixed and\n",
    "unknown) and returns the posterior value for those parameters; we will then\n",
    "construct an `emcee.EnsembleSampler` (giving the number of unknown parameters\n",
    "and the known ones) and we will finally run MCMC for a certain number of steps.\n",
    "The sampler will return tuples of sampled parameters.\n",
    "\n",
    "For analytical convenience, we will use $\\ln(Posterior)$, $\\ln(\\text{Prior})$\n",
    "and $\\ln(\\mathcal{L})$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "10a1de3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_posterior(parameters,\n",
    "                  m: float, s: float, tau: float,\n",
    "                  J: int, df: pd.DataFrame):\n",
    "    # lambda_j is an array of length J with the admission rate for each hospital\n",
    "    mu_0, sigma_0, *lambdas = parameters\n",
    "    # Check parameters that have to be strictly positive\n",
    "    if sigma_0 <= 0:  # Standard deviation of a Normal distribution\n",
    "        return -np.inf\n",
    "    if any(l <= 0 for l in lambdas):  # Frequency of a Poisson distribution\n",
    "        return -np.inf\n",
    "    log_prior = 0.0\n",
    "    # Prior for 'mu_0' (normally distributed, parameters 'm' and 's')\n",
    "    log_prior += (\n",
    "        -0.5 * np.log(2 * np.pi * s**2)\n",
    "        -0.5 * ((mu_0 - np.log(m))**2) / s**2\n",
    "    )\n",
    "    # Prior for 'sigma_0' (exponentially distributed, parameter 'tau')\n",
    "    log_prior += (\n",
    "         np.log(tau) - tau * sigma_0\n",
    "    )\n",
    "    log_likelihood = 0.0\n",
    "    # Prior for 'lambdas' (each log-normally distributed, parameters 'mu_0', 'sigma_0')\n",
    "    # and likelihood of 'observed_admissions' (Poisson distributed, parameter 'lambdas[j]')\n",
    "    # We only iterate once over range(J) for the sake of efficiency\n",
    "    for j in range(J):\n",
    "        log_prior += (\n",
    "            -np.log(lambdas[j] * sigma_0 * np.sqrt(2 * np.pi))\n",
    "            - ((np.log(lambdas[j]) - mu_0)**2) / (2 * sigma_0**2)\n",
    "        )\n",
    "        # Get the observed admission counts for hospital j,\n",
    "        #   and cast it to a numpy array for efficiency\n",
    "        # This is equivalent to (but faster than) a for loop\n",
    "        #   for i in range(reporting_frequency[j]):\n",
    "        #       reporting_frequency = df.iloc[j, i]\n",
    "        #       log_likelihood += (...)\n",
    "        observed_admissions = df.iloc[j,:reporting_frequency[j]].to_numpy()\n",
    "        log_likelihood += np.sum(\n",
    "            observed_admissions * np.log(lambdas[j])\n",
    "            - lambdas[j]\n",
    "            - gammaln(observed_admissions+1)\n",
    "        )\n",
    "    return log_prior + log_likelihood"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1e393c9f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5000/5000 [04:03<00:00, 20.51it/s]\n"
     ]
    }
   ],
   "source": [
    "def get_samples_from_full_bhm(initial_mu_0: float, initial_sigma_0: float,\n",
    "                              initial_lambdas: List[float],\n",
    "                              log_posterior: Callable,\n",
    "                              m: float, s: float, tau: float,\n",
    "                              J: int, df: pd.DataFrame):\n",
    "      # Data: observed_admissions (from previous synthetic data generation)\n",
    "      # Assuming observed_admissions is a pandas df (J * max(reporting_frequency))\n",
    "\n",
    "      # Number of dimensions and walkers\n",
    "      ndim = 2 + J  # [mu_0, sigma_0, lambda_1, ..., lambda_J]\n",
    "      num_walkers = 32\n",
    "\n",
    "      # Initial guess for parameters\n",
    "      initial_guess = [\n",
    "            initial_mu_0,\n",
    "            initial_sigma_0,\n",
    "            *initial_lambdas\n",
    "      ]\n",
    "\n",
    "      # Initialize walkers in a small ball around the initial guess\n",
    "      pos = initial_guess + 1e-4 * np.random.randn(num_walkers, ndim)\n",
    "\n",
    "      # Create the sampler\n",
    "      sampler = emcee.EnsembleSampler(\n",
    "            num_walkers, ndim, log_posterior,\n",
    "            args=(m, s, tau, J, df)\n",
    "      )\n",
    "\n",
    "      # Run MCMC\n",
    "      nsteps = 5000\n",
    "      sampler.run_mcmc(pos, nsteps, progress=True)\n",
    "\n",
    "      # Get one sample every 'thin' steps, discarding the burn-in (first 'discard' steps)\n",
    "      return sampler.get_chain(discard=1000, thin=10, flat=True)\n",
    "\n",
    "\n",
    "samples = get_samples_from_full_bhm(\n",
    "      initial_mu_0=true_mu_0 * 0.9, initial_sigma_0=true_sigma_0 * 1.1,\n",
    "      initial_lambdas=[\n",
    "            lambda_j * (0.9 + 0.2 * np.random.random())\n",
    "            for lambda_j in true_lambdas\n",
    "      ],\n",
    "      log_posterior=log_posterior,\n",
    "      m=m, s=s, tau=tau, J=J, df=df\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bcedf091",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 2152916 characters]",
      "text/plain": [
       "<Figure size 2650x2650 with 144 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Labels and true values\n",
    "labels = [r'$\\mu_0$', r'$\\sigma_0$', *[rf'$\\lambda_{{{j + 1}}}$' for j in range(J)]]\n",
    "truths = [true_mu_0, true_sigma_0, *true_lambdas]\n",
    "\n",
    "# Create a corner plot\n",
    "fig = corner.corner(\n",
    "    samples,\n",
    "    labels=labels,\n",
    "    truths=truths,\n",
    "    quantiles=[0.16, 0.5, 0.84],\n",
    "    show_titles=True,\n",
    "    title_kwargs={\"fontsize\": 12}\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66182918",
   "metadata": {},
   "source": [
    "____\n",
    "\n",
    "(b) Determine the marginal posterior distribution for the admission rate for\n",
    "each hospital, and compare it with the posterior estimate in a model with no\n",
    "pooling (i.e., where each hospital’s rate is inferred exclusively from its\n",
    "observed counts, i.e. $\\ln \\lambda_j \\sim \\mathcal{N} (\\mu_0, \\sigma^2_0)$ and\n",
    "the same priors as above, and the posterior for hospital j comes exclusively\n",
    "from its own data). Check that the posterior means for hospitals with a smaller\n",
    "number of records (i.e., $n_j = 12$) exhibit stronger shrinkage towards the\n",
    "global mean, thus demonstrating borrowing of strength, and typically have 68%\n",
    "HPD credible intervals that are shorter than in the pooling model. You may use a\n",
    "violin plot to make this comparison.\n",
    "\n",
    "#### With pooling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "03ee076c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "          Reporting frequency  log lambda 16th percentile  \\\n",
      "Hospital                                                    \n",
      "1                          12                    5.048979   \n",
      "2                          54                    5.060622   \n",
      "3                          12                    5.053617   \n",
      "4                          12                    5.063582   \n",
      "5                          54                    5.059486   \n",
      "6                          54                    5.076310   \n",
      "7                          12                    5.065790   \n",
      "8                          24                    5.068508   \n",
      "9                          12                    5.042735   \n",
      "10                         12                    5.059849   \n",
      "\n",
      "          log lambda 84th percentile  HPD width (68%)      mu_0   sigma_0  \n",
      "Hospital                                                                   \n",
      "1                           5.076839         0.027860  5.072097  0.015922  \n",
      "2                           5.077944         0.017322  5.072097  0.015922  \n",
      "3                           5.079737         0.026119  5.072097  0.015922  \n",
      "4                           5.087553         0.023971  5.072097  0.015922  \n",
      "5                           5.077057         0.017571  5.072097  0.015922  \n",
      "6                           5.096485         0.020175  5.072097  0.015922  \n",
      "7                           5.091870         0.026080  5.072097  0.015922  \n",
      "8                           5.091746         0.023239  5.072097  0.015922  \n",
      "9                           5.074290         0.031554  5.072097  0.015922  \n",
      "10                          5.084888         0.025039  5.072097  0.015922  \n"
     ]
    }
   ],
   "source": [
    "rows = []\n",
    "\n",
    "for j in range(J):\n",
    "    mu_0_j = samples[:, 0]\n",
    "    sigma_0_j = samples[:, 1]\n",
    "    lambda_samples_j = np.log(samples[:, 2 + j])\n",
    "    lower, upper = np.percentile(lambda_samples_j, [16, 84])\n",
    "    rows.append({\n",
    "        \"Hospital\": j + 1,\n",
    "        \"Reporting frequency\": reporting_frequency[j],\n",
    "        \"log lambda 16th percentile\": lower,\n",
    "        \"log lambda 84th percentile\": upper,\n",
    "        \"HPD width (68%)\": upper - lower,\n",
    "        \"mu_0\": np.mean(mu_0_j),\n",
    "        \"sigma_0\": np.mean(sigma_0_j)\n",
    "    })\n",
    "\n",
    "temp = pd.DataFrame(rows)\n",
    "temp = pd.DataFrame(rows).set_index(\"Hospital\")\n",
    "print(temp)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "70b92966",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pearson r = -0.878\n",
      "R^2 = 0.771\n",
      "p-value = 8.381e-04\n"
     ]
    }
   ],
   "source": [
    "x = temp[\"Reporting frequency\"].to_numpy()\n",
    "y = temp[\"HPD width (68%)\"].to_numpy()\n",
    "\n",
    "r, p_value = pearsonr(x, y)\n",
    "r_squared = r**2\n",
    "\n",
    "print(f\"Pearson r = {r:.3f}\")\n",
    "print(f\"R^2 = {r_squared:.3f}\")\n",
    "print(f\"p-value = {p_value:.3e}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "74dbf9b4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 45500 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Fit least-squares line\n",
    "slope, intercept = np.polyfit(x, y, 1)\n",
    "x_line = np.linspace(x.min(), x.max(), 100)\n",
    "y_line = slope * x_line + intercept\n",
    "\n",
    "plt.figure()\n",
    "plt.scatter(x, y)\n",
    "plt.plot(x_line, y_line)\n",
    "plt.xlabel(\"Reporting frequency (n_j)\")\n",
    "plt.ylabel(\"HPD width (68%)\")\n",
    "plt.title(\n",
    "    f\"HPD width vs reporting frequency\\n\"\n",
    "    f\"R² = {r_squared:.3f}, r = {r:.2f}, p = {p_value:.3f}\"\n",
    ")\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d586de09",
   "metadata": {},
   "source": [
    "#### Without pooling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "eb404c06",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3000/3000 [00:17<00:00, 168.07it/s]\n",
      "100%|██████████| 3000/3000 [00:16<00:00, 181.31it/s]\n",
      "100%|██████████| 3000/3000 [00:17<00:00, 170.11it/s]\n",
      "100%|██████████| 3000/3000 [00:18<00:00, 160.65it/s]\n",
      "100%|██████████| 3000/3000 [00:17<00:00, 170.85it/s]\n",
      "100%|██████████| 3000/3000 [00:16<00:00, 181.46it/s]\n",
      "100%|██████████| 3000/3000 [00:20<00:00, 145.86it/s]\n",
      "100%|██████████| 3000/3000 [00:21<00:00, 142.75it/s]\n",
      "100%|██████████| 3000/3000 [00:21<00:00, 140.67it/s]\n",
      "100%|██████████| 3000/3000 [00:21<00:00, 138.94it/s]\n"
     ]
    }
   ],
   "source": [
    "mu_fixed = np.log(m)\n",
    "sigma_fixed = s\n",
    "\n",
    "\n",
    "def log_posterior_without_pooling(parameters: np.ndarray,\n",
    "                                  mu_0: float, sigma_0: float,\n",
    "                                  j: int, df: pd.DataFrame):\n",
    "    lambda_j, *_ = parameters\n",
    "    if lambda_j <= 0:\n",
    "        return -np.inf\n",
    "    # Prior on 'lambda_j' (log-normally distributed, parameters 'mu_0' and 'sigma_0')\n",
    "    log_prior = (\n",
    "        -np.log(lambda_j * sigma_0 * np.sqrt(2 * np.pi))\n",
    "        - (np.log(lambda_j) - mu_0)**2 / (2 * sigma_0**2)\n",
    "    )\n",
    "    observed_admissions = df.iloc[j,:reporting_frequency[j]].to_numpy()\n",
    "    # Likelihood: Poisson\n",
    "    log_likelihood = np.sum(\n",
    "        observed_admissions * np.log(lambda_j)\n",
    "        - lambda_j\n",
    "        - gammaln(observed_admissions + 1)\n",
    "    )\n",
    "\n",
    "    return log_prior + log_likelihood\n",
    "\n",
    "\n",
    "def get_samples_without_pooling(initial_lambdas: List[float],\n",
    "                                log_posterior_without_pooling: Callable,\n",
    "                                mu_fixed: float, sigma_fixed: float,\n",
    "                                J: int, df: pd.DataFrame):\n",
    "    samples_without_pooling = {}\n",
    "    for j in range(J):\n",
    "        # Number of dimensions and walkers\n",
    "        ndim = 1  # Only lambda_j\n",
    "        num_walkers = 32\n",
    "\n",
    "        # Initial guess for parameters\n",
    "        initial_guess = [\n",
    "            initial_lambdas[j]\n",
    "        ]\n",
    "\n",
    "        # Initialize walkers in a small ball around the initial guess\n",
    "        pos = initial_guess + 1e-4 * np.random.randn(num_walkers, ndim)\n",
    "\n",
    "        # Create the sampler\n",
    "        sampler = emcee.EnsembleSampler(\n",
    "            num_walkers, ndim, log_posterior_without_pooling,\n",
    "            args=(mu_fixed, sigma_fixed, j, df)\n",
    "        )\n",
    "\n",
    "        # Run MCMC\n",
    "        nsteps = 3000\n",
    "        sampler.run_mcmc(pos, nsteps, progress=True)\n",
    "\n",
    "        # Get one sample every 'thin' steps, discarding the burn-in (first 'discard' steps)\n",
    "        samples_without_pooling[j] = sampler.get_chain(discard=1000, thin=10, flat=True)\n",
    "    return samples_without_pooling\n",
    "\n",
    "\n",
    "samples_without_pooling = get_samples_without_pooling(\n",
    "    initial_lambdas=[\n",
    "        lambda_j * (0.9 + 0.2 * np.random.random())\n",
    "        for lambda_j in true_lambdas\n",
    "      ],\n",
    "    log_posterior_without_pooling=log_posterior_without_pooling,\n",
    "    mu_fixed=mu_fixed, sigma_fixed=sigma_fixed,\n",
    "    J=J, df=df\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "aa76a65f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "          Reporting frequency  log lambda 16th percentile  \\\n",
      "Hospital                                                    \n",
      "1                          12                    5.017071   \n",
      "2                          54                    5.057248   \n",
      "3                          12                    5.027429   \n",
      "4                          12                    5.057141   \n",
      "5                          54                    5.054694   \n",
      "6                          54                    5.085284   \n",
      "7                          12                    5.070768   \n",
      "8                          24                    5.072616   \n",
      "9                          12                    4.999103   \n",
      "10                         12                    5.044466   \n",
      "\n",
      "          log lambda 84th percentile  HPD width (68%)      mu_0  sigma_0  \n",
      "Hospital                                                                  \n",
      "1                           5.063500         0.046429  4.442651      0.4  \n",
      "2                           5.079072         0.021824  4.442651      0.4  \n",
      "3                           5.072223         0.044794  4.442651      0.4  \n",
      "4                           5.102917         0.045777  4.442651      0.4  \n",
      "5                           5.076253         0.021559  4.442651      0.4  \n",
      "6                           5.106803         0.021519  4.442651      0.4  \n",
      "7                           5.116844         0.046076  4.442651      0.4  \n",
      "8                           5.104699         0.032083  4.442651      0.4  \n",
      "9                           5.046609         0.047507  4.442651      0.4  \n",
      "10                          5.089842         0.045376  4.442651      0.4  \n"
     ]
    }
   ],
   "source": [
    "rows = []\n",
    "\n",
    "for j in range(J):\n",
    "    swp = samples_without_pooling[j]\n",
    "    lambda_samples_j = np.log(swp[:, 0])\n",
    "    lower, upper = np.percentile(lambda_samples_j, [16, 84])\n",
    "    rows.append({\n",
    "        \"Hospital\": j + 1,\n",
    "        \"Reporting frequency\": reporting_frequency[j],\n",
    "        \"log lambda 16th percentile\": lower,\n",
    "        \"log lambda 84th percentile\": upper,\n",
    "        \"HPD width (68%)\": upper - lower,\n",
    "        \"mu_0\": np.mean(mu_fixed),\n",
    "        \"sigma_0\": np.mean(sigma_fixed)\n",
    "    })\n",
    "\n",
    "temp = pd.DataFrame(rows)\n",
    "temp = pd.DataFrame(rows).set_index(\"Hospital\")\n",
    "print(temp)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "0117285c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 97776 characters]",
      "text/plain": [
       "<Figure size 1400x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data = []\n",
    "positions = []\n",
    "groups = []\n",
    "\n",
    "pos, offset, step = 1.0, 0.35, 1.2\n",
    "\n",
    "for j in range(J):\n",
    "    # With pooling\n",
    "    data.append(np.log(samples[:, 2 + j]))\n",
    "    positions.append(pos)\n",
    "    groups.append(\"bhm\")\n",
    "    # Without pooling\n",
    "    data.append(np.log(samples_without_pooling[j][:, 0]))\n",
    "    positions.append(pos + offset)\n",
    "    groups.append(\"nopool\")\n",
    "    pos += step\n",
    "\n",
    "plt.figure(figsize=(14, 6))\n",
    "\n",
    "vp = plt.violinplot(\n",
    "    data,\n",
    "    positions=positions,\n",
    "    widths=0.3,\n",
    "    showmeans=False,\n",
    "    showmedians=True,\n",
    "    showextrema=False\n",
    ")\n",
    "\n",
    "# Color each violin plot according to group (with or without pooling)\n",
    "for body, group in zip(vp[\"bodies\"], groups):\n",
    "    body.set_facecolor(\n",
    "        {\n",
    "            \"bhm\": \"orange\",\n",
    "            \"nopool\": \"darkblue\"\n",
    "        }[group]\n",
    "    )\n",
    "    body.set_alpha(0.7)\n",
    "\n",
    "vp[\"cmedians\"].set_color(\"white\")\n",
    "vp[\"cmedians\"].set_alpha(1.0)\n",
    "vp[\"cmedians\"].set_linewidth(2.0)\n",
    "\n",
    "# Add ground truth\n",
    "x_true_lambdas = [1 + step * j + offset / 2 for j in range(J)]\n",
    "y_true_lambdas = np.log(true_lambdas)\n",
    "plt.scatter(\n",
    "    x_true_lambdas, y_true_lambdas,\n",
    "    color=\"darkgreen\", marker=\"x\",\n",
    "    s=80, linewidths=2,\n",
    "    zorder=100000,\n",
    "    label=\"Ground truth\"\n",
    ")\n",
    "\n",
    "# Add 86% HPD intervals\n",
    "cap_width = 0.06\n",
    "for i, x in enumerate(positions):\n",
    "    lower = np.quantile(data[i], 0.16)\n",
    "    higher = np.quantile(data[i], 0.84)\n",
    "    plt.hlines(\n",
    "        y=[lower, higher],\n",
    "        xmin=x - cap_width,\n",
    "        xmax=x + cap_width,\n",
    "        color=\"black\",\n",
    "        linewidth=2,\n",
    "        zorder=10\n",
    "    )\n",
    "    plt.vlines(\n",
    "        x=x,\n",
    "        ymin=lower,\n",
    "        ymax=higher,\n",
    "        color=\"black\",\n",
    "        linewidth=2,\n",
    "        zorder=10\n",
    "    )\n",
    "\n",
    "# Set axes labels and plot title\n",
    "plt.xticks(\n",
    "    [1 + step * j + offset/2 for j in range(J)],\n",
    "    [f\"Hospital {j+1}\\n(n={reporting_frequency[j]})\" for j in range(J)]\n",
    ")\n",
    "plt.ylabel(r\"$\\ln \\lambda_j$\")\n",
    "plt.title(r\"Distribution of $\\ln \\lambda_y$\"\"\\n\"r\"Full BHM vs w/out pooling\")\n",
    "\n",
    "plt.legend(\n",
    "    handles=[\n",
    "        Patch(color=\"orange\", label=\"With pooling\"),\n",
    "        Patch(color=\"darkblue\", label=\"No pooling\"),\n",
    "        Line2D(\n",
    "            [0], [0],\n",
    "            marker=\"x\",\n",
    "            color=\"darkgreen\",\n",
    "            linestyle=\"None\",\n",
    "            markersize=10,\n",
    "            markeredgewidth=2,\n",
    "            label=\"Ground truth\"\n",
    "        ),\n",
    "        Line2D(\n",
    "            [0], [0],\n",
    "            color=\"black\",\n",
    "            markersize=10,\n",
    "            markeredgewidth=2,\n",
    "            label=\"68% HPD intervals\"\n",
    "        )\n",
    "    ],\n",
    "    loc=\"upper right\"\n",
    ")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd561fbc",
   "metadata": {},
   "source": [
    "___\n",
    "\n",
    "(c) Consider the prior predictive distribution for possible data within the BHM:\n",
    "\n",
    "$$ Pr( y_{ij} \\mid priors) = \\int Pr ( y_{ij} \\mid \\lambda_j) Pr(\\lambda_j \\mid \\mu_0, \\sigma_0) Pr (\\mu_0,\\sigma_0) d \\mu_0 d \\sigma_0$$\n",
    "\n",
    "and simulate from it predictions for the possible counts $y_{ij}$ from the model\n",
    "(before you see any data). Evaluate the degree of shrinkage by using as metric\n",
    "the average shrinkage in the standard deviation of the posterior with respect to\n",
    "the no-pooling scenario, i.e.\n",
    "\n",
    "$$ S = \\frac {1} {J} \\sum_{j=1}^{J}{1 - \\frac{\\text{std}_\\text{BHM}(\\ln \\lambda_j \\mid \\bold y)}{\\text{std}_\\text{no pool}(\\ln \\lambda_j \\mid \\bold y)}}$$\n",
    "\n",
    "where a smaller S corresponds to higher degree of shrinkage.\n",
    "\n",
    "Determine which hyperparameter among $(m, s, \\tau)$ has the most effect on the\n",
    "degree of shrinkage, and tune each to avoid predictions that are too diffuse\n",
    "(i.e., the predictive spread is unreasonably wide) or too narrow (i.e., the\n",
    "predictive spread is over-constraining)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3366ea7e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The degree of shrinkage S for m = 85, s = 0.4, tau = 0.1 is: 0.291\n"
     ]
    }
   ],
   "source": [
    "S = np.mean([\n",
    "    1 - (\n",
    "        np.std(np.log(samples[:, 2 + j]))\n",
    "        / np.std(np.log(samples_without_pooling[j][:, 0]))\n",
    "    )\n",
    "    for j in range(J)\n",
    "])\n",
    "\n",
    "print(\n",
    "    f\"The degree of shrinkage S for m = {m}, s = {s}, tau = {tau} is: {S:.3f}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "9ddb9e81",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Evaluating degree of shrinkage S for m = 93.50000000000001, s = 0.4, tau = 0.1...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5000/5000 [04:07<00:00, 20.24it/s]\n",
      "100%|██████████| 3000/3000 [00:19<00:00, 157.13it/s]\n",
      "100%|██████████| 3000/3000 [00:21<00:00, 140.87it/s]\n",
      "100%|██████████| 3000/3000 [00:19<00:00, 151.15it/s]\n",
      "100%|██████████| 3000/3000 [00:20<00:00, 147.06it/s]\n",
      "100%|██████████| 3000/3000 [00:17<00:00, 173.79it/s]\n",
      "100%|██████████| 3000/3000 [00:16<00:00, 181.20it/s]\n",
      "100%|██████████| 3000/3000 [00:15<00:00, 189.46it/s]\n",
      "100%|██████████| 3000/3000 [00:15<00:00, 189.62it/s]\n",
      "100%|██████████| 3000/3000 [00:16<00:00, 184.59it/s]\n",
      "100%|██████████| 3000/3000 [00:15<00:00, 193.01it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The degree of shrinkage S for m = 93.50000000000001, s = 0.4, tau = 0.1 is: 0.015\n",
      "Evaluating degree of shrinkage S for m = 85, s = 0.44000000000000006, tau = 0.1...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5000/5000 [04:01<00:00, 20.72it/s]\n",
      "100%|██████████| 3000/3000 [00:19<00:00, 157.29it/s]\n",
      "100%|██████████| 3000/3000 [00:17<00:00, 171.69it/s]\n",
      "100%|██████████| 3000/3000 [00:17<00:00, 174.51it/s]\n",
      "100%|██████████| 3000/3000 [00:18<00:00, 163.67it/s]\n",
      "100%|██████████| 3000/3000 [00:19<00:00, 157.06it/s]\n",
      "100%|██████████| 3000/3000 [00:18<00:00, 164.69it/s]\n",
      "100%|██████████| 3000/3000 [00:18<00:00, 159.45it/s]\n",
      "100%|██████████| 3000/3000 [00:17<00:00, 171.62it/s]\n",
      "100%|██████████| 3000/3000 [00:17<00:00, 170.42it/s]\n",
      "100%|██████████| 3000/3000 [00:17<00:00, 171.85it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The degree of shrinkage S for m = 85, s = 0.44000000000000006, tau = 0.1 is: 0.021\n",
      "Evaluating degree of shrinkage S for m = 85, s = 0.4, tau = 0.11000000000000001...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5000/5000 [03:46<00:00, 22.07it/s]\n",
      "100%|██████████| 3000/3000 [00:15<00:00, 188.12it/s]\n",
      "100%|██████████| 3000/3000 [00:16<00:00, 184.13it/s]\n",
      "100%|██████████| 3000/3000 [00:16<00:00, 185.33it/s]\n",
      "100%|██████████| 3000/3000 [00:15<00:00, 189.97it/s]\n",
      "100%|██████████| 3000/3000 [00:16<00:00, 185.13it/s]\n",
      "100%|██████████| 3000/3000 [00:17<00:00, 171.60it/s]\n",
      "100%|██████████| 3000/3000 [00:16<00:00, 178.34it/s]\n",
      "100%|██████████| 3000/3000 [00:16<00:00, 180.16it/s]\n",
      "100%|██████████| 3000/3000 [00:22<00:00, 135.34it/s]\n",
      "100%|██████████| 3000/3000 [00:16<00:00, 178.59it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The degree of shrinkage S for m = 85, s = 0.4, tau = 0.11000000000000001 is: -0.003\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "m = 85\n",
    "s = 0.4\n",
    "tau = 0.1\n",
    "\n",
    "def evaluate_degree_of_shrinkage(m, s, tau):\n",
    "    print(\n",
    "        f\"Evaluating degree of shrinkage S for m = {m}, s = {s}, tau = {tau}...\"\n",
    "    )\n",
    "    true_mu_0 = np.random.normal(np.log(m), s)\n",
    "    true_sigma_0 = np.random.exponential(tau)\n",
    "    df, true_lambdas = get_admission_count_matrix(\n",
    "        mu=true_mu_0, sigma=true_sigma_0\n",
    "    )\n",
    "    initial_mu_0 = true_mu_0 * 0.9\n",
    "    initial_sigma_0 = true_sigma_0 * 1.1\n",
    "    initial_lambdas = [\n",
    "        lambda_j * (0.9 + 0.2 * np.random.random())\n",
    "        for lambda_j in true_lambdas\n",
    "    ]\n",
    "    # Samples from joint posterior (full BHM)\n",
    "    samples_bhm = get_samples_from_full_bhm(\n",
    "        initial_mu_0=initial_mu_0, initial_sigma_0=initial_sigma_0,\n",
    "        initial_lambdas=initial_lambdas,\n",
    "        log_posterior=log_posterior,\n",
    "        m=m, s=s, tau=tau,\n",
    "        J=J, df=df\n",
    "    )\n",
    "    # Samples from non-pooling model\n",
    "    samples_without_pooling = get_samples_without_pooling(\n",
    "        initial_lambdas=initial_lambdas,\n",
    "        log_posterior_without_pooling=log_posterior_without_pooling,\n",
    "        mu_fixed=initial_mu_0, sigma_fixed=initial_sigma_0,\n",
    "        J=J, df=df\n",
    "    )\n",
    "    S = np.mean([\n",
    "        1 - (\n",
    "            np.std(np.log(samples_bhm[:, 2 + j]))\n",
    "            / np.std(np.log(samples_without_pooling[j][:, 0]))\n",
    "        )\n",
    "        for j in range(J)\n",
    "    ])\n",
    "    print(\n",
    "        f\"The degree of shrinkage S for m = {m}, s = {s}, tau = {tau} is: {S:.3f}\"\n",
    "    )\n",
    "\n",
    "\n",
    "evaluate_degree_of_shrinkage(m=(1.1 * m), s=s, tau=tau)\n",
    "evaluate_degree_of_shrinkage(m=m, s=(1.1 * s), tau=tau)\n",
    "evaluate_degree_of_shrinkage(m=m, s=s, tau=(1.1 * tau))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3f149d1",
   "metadata": {},
   "source": [
    "## Tentative interpretation of these results\n",
    "\n",
    "The parameter that affects $S$ more is $\\tau$, which controls the variability of\n",
    "lambdas across hospitals. Therefore, the BHM is most useful when the sources\n",
    "being pooled have similar values for the latent variables (in this case, smaller\n",
    "$\\tau$), which makes sense.\n",
    "\n",
    "Also $s$ (the uncertainty about the population mean admission rate) affects $S$\n",
    "(lower $s$ means more shrinkage).\n",
    "\n",
    "$m$, the population mean admission rate, affects $S$ indirectly through\n",
    "$\\sigma_0$.\n",
    "\n",
    "To support this conclusion, the following table reports the increase of each\n",
    "parameter and the decrease in shrinkage it produced when comparing the full\n",
    "hierarchical model with the model without pooling across hospitals.\n",
    "\n",
    "The smallest change in parameters ($\\tau + 0.01$) also produced the biggest\n",
    "deterioration in BHM performance, supporting the hypothesis that $\\tau$ is the\n",
    "most influent parameter in determining the efficacy of borrowing of strength."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e52ae62",
   "metadata": {},
   "source": [
    "Parameter | 10% increase in parameter | $\\Delta S$\n",
    "| :-----: | :-----------------------: | :--------: |\n",
    "m         | 8.5                       | 27.6%\n",
    "s         | 0.04                      | 27.0%\n",
    "$\\tau$    | 0.01                      | 29.4%"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "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.13.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
