Working on log posterior function
This commit is contained in:
300
borrowing_strenght_hierarchical_model.ipynb
Normal file
300
borrowing_strenght_hierarchical_model.ipynb
Normal file
@@ -0,0 +1,300 @@
|
|||||||
|
{
|
||||||
|
"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 Fig. 1. 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).\n",
|
||||||
|
"\n",
|
||||||
|
"____\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",
|
||||||
|
"___\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": null,
|
||||||
|
"id": "49307215",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"import emcee\n",
|
||||||
|
"import numpy as np\n",
|
||||||
|
"import pandas as pd\n",
|
||||||
|
"\n",
|
||||||
|
"np.random.seed(42)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "af15e048",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"J = 10\n",
|
||||||
|
"reporting_frequency = np.random.choice(\n",
|
||||||
|
" [54, 24, 12],\n",
|
||||||
|
" number_of_hospitals\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": null,
|
||||||
|
"id": "4a09b1d3",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"def get_expected_admissions(mu, sigma):\n",
|
||||||
|
" \"\"\"Get expected number of admissions to a hospital at a time point.\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",
|
||||||
|
" frequency admission rate.\n",
|
||||||
|
" \"\"\"\n",
|
||||||
|
" log_admission_rate = np.random.normal(\n",
|
||||||
|
" mu, sigma\n",
|
||||||
|
" )\n",
|
||||||
|
" admission_rate = np.pow(np.e, log_admission_rate)\n",
|
||||||
|
" #print(admission_rate)\n",
|
||||||
|
" admission_counts = np.random.poisson(admission_rate)\n",
|
||||||
|
" return admission_counts\n"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "2f6b2f3c",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"def get_admission_counts(mu: float, sigma: float) -> pd.DataFrame:\n",
|
||||||
|
" data = {}\n",
|
||||||
|
" max_T = max(reporting_frequency)\n",
|
||||||
|
" for j in range(J): # For each hospital\n",
|
||||||
|
" admission_records = []\n",
|
||||||
|
" for i in range(reporting_frequency[j]): # For each report\n",
|
||||||
|
" admission_records.append(get_expected_admissions(\n",
|
||||||
|
" mu, sigma\n",
|
||||||
|
" ))\n",
|
||||||
|
" admission_records.extend([pd.NA] * (max_T - len(admission_records)))\n",
|
||||||
|
" data[f\"Hospital {j+1}\"] = admission_records\n",
|
||||||
|
" data = pd.DataFrame(data).transpose()\n",
|
||||||
|
" data.index.name = \"Time point\"\n",
|
||||||
|
" return data"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"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": null,
|
||||||
|
"id": "216cdb28",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"m = 85\n",
|
||||||
|
"s = 0.4\n",
|
||||||
|
"tau = 0.1"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "4b5ee692",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"mu_zero = np.random.normal(np.log(m), s)\n",
|
||||||
|
"sigma_zero = np.random.exponential(tau)\n",
|
||||||
|
"\n",
|
||||||
|
"df = get_admission_counts(mu=mu_zero, sigma=sigma_zero)\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": null,
|
||||||
|
"id": "10a1de3d",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"def log_posterior(parameters, m, s, tau):\n",
|
||||||
|
" mu_0, sigma_0 = parameters\n",
|
||||||
|
" if sigma_0 <= 0: # Standard deviation must be strictly positive\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",
|
||||||
|
" # Likelihood for ''\n",
|
||||||
|
" # TODO\n",
|
||||||
|
" return log_prior + log_likelihood"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "1e393c9f",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"# Data: x_obs and y_obs (from previous synthetic data generation)\n",
|
||||||
|
"# Assuming x_obs and y_obs are numpy arrays of length N\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"# Number of dimensions and walkers\n",
|
||||||
|
"ndim = 4 # [theta_0, theta_1, x0, log_Rx2]\n",
|
||||||
|
"nwalkers = 32\n",
|
||||||
|
"\n",
|
||||||
|
"# Initial guess for parameters\n",
|
||||||
|
"initial_guess = [2.0, 0.5, 5.0, Rx**2]\n",
|
||||||
|
"\n",
|
||||||
|
"# Initialize walkers in a small ball around the initial guess\n",
|
||||||
|
"pos = initial_guess + 1e-4 * np.random.randn(nwalkers, ndim)\n",
|
||||||
|
"\n",
|
||||||
|
"# Create the sampler\n",
|
||||||
|
"sampler = emcee.EnsembleSampler(nwalkers, ndim, log_posterior,\n",
|
||||||
|
" args=(x_obs, y_obs, sigma_x, sigma_y,\n",
|
||||||
|
" mu_x0, sigma_x0_squared, alpha_R, beta_R))\n",
|
||||||
|
"\n",
|
||||||
|
"# Run MCMC\n",
|
||||||
|
"nsteps = 5000\n",
|
||||||
|
"sampler.run_mcmc(pos, nsteps, progress=True)\n",
|
||||||
|
"\n",
|
||||||
|
"# Get the samples\n",
|
||||||
|
"samples = sampler.get_chain(discard=1000, thin=10, flat=True)\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.12.12"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"nbformat": 4,
|
||||||
|
"nbformat_minor": 5
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user