econmt-bayesian.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(econmt-bayesian)=\n",
    "# Bayesian Inference\n",
    "\n",
    "In this chapter, we'll look at how to perform analysis and regressions using Bayesian techniques.\n",
    "\n",
    "Let's import a few of the packages we'll need first. Two key packages that we'll be using that you might not have seen before are [**PyMC**](https://docs.pymc.io/), a Bayesian inference package.\n",
    "\n",
    "In the next chapter, we'll look at an easier way to do Bayesian inference in Python, but you'll find life easier overall if you start here. We'll also use [**ArviZ**](https://arviz-devs.github.io/) for visualisation of the results from Bayesian inference, a package that builds on **Matplotlib**, but this will get automatically installed when you intall **PyMC**. You should follow the install instructions for **PyMC** and **Bambi** carefully and, if you're confident with using different Python environments, it's a good idea to spin up a new 'bayes' environment to try them out in. In case you need a refresher, the chapter on {ref}`code-preliminaries` covers basic information on how to install new packages.\n",
    "\n",
    "The expert Bayesian may wonder, \"why these packages?\" There are many others available, even just in Python. One reason is that they're simpler to use for beginners while having enormous power for advanced users too, so they're very much in keeping with Python's low floor, high ceiling style. The other is that they're considerably more mature than some of the more recent alternatives based on popular tensor and machine learning packages (though **PyMC** is also built on an underlying tensor package). Finally, they also have flexible computing back-ends (including Jax and Numba, and including using GPUs), which means that super users can make their models go *extremely fast* should they need to—you can find more information on the speed-ups that are possible [here](https://www.pymc-labs.io/blog-posts/pymc-stan-benchmark/).\n",
    "\n",
    "This chapter has benefitted from the many Bayesian blogs and forum conversations across the web, in particular the ones [here](https://elevanth.org/blog/2017/11/28/build-a-better-markov-chain/), [here](https://people.duke.edu/~ccc14/sta-663/MCMC.html#markov-chain-monte-carlo-mcmc), [here](https://github.com/fonnesbeck/Bayes_Computing_Course/blob/master/notebooks/Section2_1-MCMC.ipynb), [here](https://stats.stackexchange.com/questions/4700/what-is-the-difference-between-fixed-effect-random-effect-and-mixed-effect-mode), and [here](https://rlhick.people.wm.edu/stories/bayesian_3.html).\n",
    "\n",
    "Here are our initial imports and settings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import xarray as xr\n",
    "import warnings\n",
    "import pymc as pm\n",
    "import arviz as az\n",
    "\n",
    "# Plot settings\n",
    "plt.style.use(\n",
    "    \"https://github.com/aeturrell/coding-for-economists/raw/main/plot_style.txt\"\n",
    ")\n",
    "az.style.use(\"arviz-darkgrid\")\n",
    "\n",
    "# Pandas: Set max rows displayed for readability\n",
    "pd.set_option(\"display.max_rows\", 23)\n",
    "\n",
    "# Set seed for random numbers\n",
    "seed_for_prng = 78557\n",
    "prng = np.random.default_rng(seed_for_prng)  # prng=probabilistic random number generator\n",
    "# Turn off warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "# check pymc version\n",
    "print(f\"Running on PyMC v{pm.__version__}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51a55374",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import matplotlib_inline.backend_inline\n",
    "\n",
    "matplotlib_inline.backend_inline.set_matplotlib_formats(\"svg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "The biggest difference between the Bayesian and frequentist approaches (that we've seen in the other chapters) is probably that, in Bayesian models, the parameters are not assumed to be fixed but instead are treated as random variables whose uncertainty is described using probability distributions. The data are considered fixed. You might see the 'inverse probability' formulation of a Bayesian model written as $p(\\theta | y)$ where the $y$ are the data, and the $\\theta$ are the model parameters. An interesting aspect of Bayesianism is that there is just one estimator: Bayes' theorem.\n",
    "\n",
    "This is a contrast with the frequentist view, which holds that the data are random but the model parameters are fixed, and models are often expressed as functions, for example $f(y | \\theta)$. Frequentist inference typically involves deriving estimators for the model parameters, and these are usually created to minimise the bias, minimise the variance, or maximise the efficiency. Not so with Bayesian inference.\n",
    "\n",
    "As a reminder, Bayes' theorem says that ${\\displaystyle p(A\\mid B)={\\frac {p(B\\mid A)p(A)}{p(B)}}}$, where $A$ and $B$ are distinct events, $p(A)$ is the probability of event $A$ happening, and $p(A\\mid B)$ is the probability of $A$ happening given $B$ has happened or is true. $p$ could stand in for a single number, a probability density function, or a probability mass function. When dealing with data and model parameters, and ignoring a rescaling factor, this can be written as:\n",
    "\n",
    "$$\n",
    "p({\\boldsymbol{\\theta }}|{\\boldsymbol{y}})\\propto p({\\boldsymbol{y}}|{\\boldsymbol{\\theta }})p({\\boldsymbol{\\theta }}).\n",
    "$$\n",
    "\n",
    "In these equations:\n",
    "\n",
    "1. $p({\\boldsymbol{\\theta}})$ is the prior put on model parameters—what we think the distribution of model parameters will look like.\n",
    "2. $p({\\boldsymbol{y}}|{\\boldsymbol{\\theta }})$ is the likelihood of this data given a particular set of model parameters.\n",
    "3. $p({\\boldsymbol{\\theta }}|{\\boldsymbol{y}})$ is the posterior probability of model parameters given the observed data.\n",
    "\n",
    "Bayesian modelling proceeds as highlighted in the review article, *Bayesian statistics and modelling* {cite:t}`van2021bayesian`:\n",
    "\n",
    "![The Bayesian research cycle.](https://media.springernature.com/lw685/springer-static/image/art%3A10.1038%2Fs43586-020-00001-2/MediaObjects/43586_2020_1_Fig1_HTML.png?as=webp)\n",
    "\n",
    "One key strength of the Bayesian approach is that it preserves uncertainty—by construction, it includes the degree of belief you have in a parameter. This makes it especially useful in cases where uncertainty is important because the uncertainties are propagated all the through to the final results. One disadvantage of the Bayesian approach is that it's not always as fast and doesn't always scale well to very large datasets.\n",
    "\n",
    "Actually estimating $p({\\boldsymbol{y}}|{\\boldsymbol{\\theta }})p({\\boldsymbol{\\theta }})$ is done computationally, at least for all but the most trivial examples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The methods used to do this are all some variation on the *Monte Carlo* method, a tool that goes much wider than Bayesian inference. Monte Carlo techniques use random numbers to choose samples that allow you to estimate properties of interest. It can be used to perform numerical integration and is used for discrete particle simulations in fields like physics and chemistry. {cite:t}`rubinstein2016simulation` provides a good introduction to the Monte Carlo methods more broadly, including simulation; {cite:t}`bielajew2001fundamentals` is a very good free textbook on Monte Carlo methods for particle transport physics.\n",
    "\n",
    "![An illustation of Monte Carlo integration](https://upload.wikimedia.org/wikipedia/commons/thumb/2/20/MonteCarloIntegrationCircle.svg/440px-MonteCarloIntegrationCircle.svg.png)\n",
    "\n",
    "*Yoderj, Mysid; Wikipedia*\n",
    "\n",
    "Say we wanted to compute the integral of a 2D shape, a circle, to get an area. In the figure above, random pairs of numbers (say, $x, y$) are chosen within the unit square. By taking the ratio of the number of points that land within the 2D shape (40) to all samples (50), we can get an approximation of the area by computing (area of square) x ratio of samples $= 3.2 \\approx \\pi$. This type of algorithm is a form of Monte Carlo rejection sampling.\n",
    "\n",
    "Markov Chain Monte Carlo takes this a step further. A Markov Chain is a stochastic process where the chance of being in a particular state (for example, taking a particular sample) depends only on the previous state. These have equilibrium distributions; essentially a long-run average distribution over states. Markov Chain Monte Carlo finds the distribution you're after (the posterior) by recording states from the chain; the more steps, the more closely the distribution of samples will match the distribution you're trying to find out about. There are nice walkthroughs of MCMC available [here](https://people.duke.edu/~ccc14/sta-663/MCMC.html#markov-chain-monte-carlo-mcmc), [here](https://github.com/fonnesbeck/Bayes_Computing_Course/blob/master/notebooks/Section2_1-MCMC.ipynb), and [here](https://rlhick.people.wm.edu/stories/bayesian_3.html).\n",
    "\n",
    "There are many algorithms for undertaking MCMC, with perhaps the most famous being the Metropolis-Hastings algorithm. For this, we have a probability density function to sample from, a function $Q$ that quantifies where to make a transition, and $\\theta_0$, the first guess of the parameters. Starting from $\\theta = \\theta_0$, the algorithm boils down to:\n",
    "1. starting from this point $\\theta$, determine its match with the target distribution by estimating $p = p({\\boldsymbol{y}}|{\\boldsymbol{\\theta}})p({\\boldsymbol{\\theta}})$\n",
    "2. propose a new sample, $\\theta'$, using some transition function $Q$, and check its match to the target (let this match be $p'$)\n",
    "3. compute the ratio of densities, $\\frac{p'}{p}$, as a metric of favourability of the new point\n",
    "4. generate a random number $r \\thicksim \\mathcal{U}[0, 1)$ and compare it to the ratio; accept the new point only if $r < \\frac{p'}{p}$\n",
    "5. repeat\n",
    "\n",
    "The clever trick here is that, over time and with more samples, the points you draw will be from the target (posterior) distribution because, at each loop, we are forcing the samples to stick to the posterior distribution. The set of samples is often called a trace. You can find a nice walkthrough of the Metropolis-Hastings algorithm, with code, [here](https://towardsdatascience.com/from-scratch-bayesian-inference-markov-chain-monte-carlo-and-metropolis-hastings-in-python-ef21a29e25a).\n",
    "\n",
    "The Metropolis-Hastings algorithm is a good place to start to understand what's going on in computational Bayesian inference, but there are more modern and powerful algorithms available today, most notably Hamiltonian Monte Carlo (HMC), which takes its inspiration from physics (where the Hamiltonian is an operator that returns the total energy of the system). Instead of sampling completely randomly, which can be inefficient, HMC sees the sampler get 'rolled' (launched in a smooth trajectory) along the estimates of the posterior. There's a very good explanation of this approach [here](https://elevanth.org/blog/2017/11/28/build-a-better-markov-chain/), and a 'conceptual introduction to HMC' may be found in {cite:t}`betancourt2018conceptual`.\n",
    "\n",
    "The most commonly used HMC sampler is the No U-Turn Sampler (NUTS), and this is what we'll be using for most of our examples.\n",
    "\n",
    "For the rest of this chapter, you won't need to worry about most of this though! The joy of **PyMC** and **Bambi** is that they take care of some of the hard work for us. Expert users will want to dig further into the details, but these packages make running your first Bayesian inference easier with sensible default settings."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fixed vs Random Effects (An Aside)\n",
    "\n",
    "As this is going to come up in this chapter, it's worth having a quick review of the difference between fixed and random effects in regression-based models. There's a bit of disagreement about what people really mean when they say \"fixed effects\" or \"random effects\", so let's define them clearly here.\n",
    "\n",
    "> Random effects are estimated with partial pooling, while fixed effects are not.\n",
    "\n",
    "What does this mean? Well, technically, if an effect is assumed to be a realised value of a random variable, it is called a random effect. We could say that, in frequentist land (where parameters have true values), a single fixed effect model might look like:\n",
    "\n",
    "$$\n",
    "y = \\gamma\\cdot d\n",
    "$$\n",
    "\n",
    "where $d \\in \\{0, 1\\}$, which could be whether an individual belongs to, say, group 0 or group 1 and $\\gamma$ is a coefficient that is to be estimated. $\\gamma$ has a single \"true\" value, and we can admire the value of $\\hat{\\gamma}$ in a regression table once we have done our estimation.\n",
    "\n",
    "In contrast, even in a frequentist model, random effects are allowed to vary: they are drawn from a distribution (or, more accurately, their deviations from the true value are drawn from a distribution). This might be written as\n",
    "\n",
    "$$\n",
    "y_{i}= w_{i}\n",
    "$$\n",
    "\n",
    "and it could be that, instead of a group-specific intercept as in the fixed effects model above, the intercept can take a *range of values* depending on the individual. Although it's not always true, another way of thinking about this is that fixed effects are constant across individuals while random effects can vary.\n",
    "\n",
    "So what's this business about partial pooling? Imagine we have a situation where one group is under-sampled relative to the other groups. With few data points, we may not be able to precisely estimate fixed effects. However, partial pooling means that, if you do have few data points within a group, the group effect estimate can be partially based on the more abundant data from other groups. Random effects can provide a good compromise between estimating an effect by completely pooling all groups, which masks group-level variation, and estimating an effect for all groups completely separately, which could give imprecise estimates for low-sample groups. For example, when estimating means using a random effects approach, subgroup means are allowed to deviate a bit from the mean of the larger group, but not by an arbitrary amount—those deviations follow a distribution, typically a Gaussian (aka Normal) distribution. That's where the \"random\" in random effects comes from: the deviations of subgroups from a \"parent group\" follow the distribution of a random variable.\n",
    "\n",
    "The rule of thumb when thinking about how many different categories should exist for a variable to get the random, rather than fixed, effects treatment is that there should be more than 5 or 6 levels at a minimum in order to apply random effects—fewer than this and it's hard to estimate the variance around the central estimate. Because you could potentially precisely estimated a fixed effect if you had lots of samples in those 5 or 6 levels, you're more likely to reach for random effects when you have fewer samples per class.\n",
    "\n",
    "Random effects are the extension of this partial pooling technique to a wide range of situations: including multiple predictors, mixed continuous and categorical variables, and more.\n",
    "\n",
    "As a conrete example, taken from a [forum post](https://stats.stackexchange.com/questions/4700/what-is-the-difference-between-fixed-effect-random-effect-and-mixed-effect-mode) by Paul of MassMutual, suppose you want to estimate average US household income by ZIP code. You have a large dataset containing observations of household incomes and ZIP codes. Some ZIP codes are well represented in the dataset, but others have only a couple of households.\n",
    "\n",
    "For your initial model you would most likely take the mean income in each ZIP. This will work well when you have lots of data for all ZIPs, but the estimates for your poorly sampled ZIPs will suffer from high variance. You can mitigate this by using a shrinkage estimator (aka partial pooling), which will push extreme values towards the mean income across all ZIP codes.\n",
    "\n",
    "But how much shrinkage/pooling should you do for a particular ZIP? Intuitively, it should depend on the following:\n",
    "\n",
    "1. How many observations you have in that ZIP\n",
    "2. How many observations you have overall\n",
    "3. The individual-level mean and variance of household income across all ZIP codes\n",
    "4. The group-level variance in mean household income across all ZIP codes\n",
    "\n",
    "Now, if you model ZIP code as a random effect, the mean income estimate in all ZIP codes will be subjected to a statistically well-founded shrinkage, taking into account all of these factors.\n",
    "\n",
    "Random effects models automatically handle 4, the variability estimation, for all random effects in the model. This is useful, as it would be hard to do manually. Having accounted for (1)–(4), a random effects model is able to determine the appropriate shrinkage for low-sample groups. It can also handle much more complicated models with many different predictors.\n",
    "\n",
    "**Mixed effects models** are models that combine random and fixed effects."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bayesian Modelling: A Simple Example\n",
    "\n",
    "We're going to set up a very simple example using synthetic data and a known data generating process;\n",
    "\n",
    "$$\n",
    "\\mu = \\alpha + \\beta x\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "$$\n",
    "Y \\mid \\alpha, \\beta, \\sigma \\stackrel{\\text{ind}}{\\thicksim} \\mathcal{N}(\\mu, \\sigma^2)\n",
    "$$\n",
    "\n",
    "For our prior distributions on the parameters, we will assume:\n",
    "\n",
    "- $\\alpha \\thicksim \\mathcal{N}(0, 10)$\n",
    "- $\\beta \\thicksim \\mathcal{N}(0, 10)$\n",
    "- $\\sigma \\thicksim \\mathcal{N_{+}}(1)$\n",
    "\n",
    "where, to be clear, these three parameters are scalars.\n",
    "\n",
    "These priors are known as *weakly informative priors* because they are quite diffuse—the 10th and 90th percentile of a normal distribution with $\\sigma=10$ spans from, roughly, -13 to 13. With a weakly informative prior, a reasonably large amount of data will ensure that the prior will not be important and the posterior will hone in on a sensible range. We're also implicitly assuming that the data are normalised so that 1 is a meaningful change and zero a sensible reference point. You may want to, for example, scale by the standard deviation of the data or take logs of (positive) predictors, so that coeficients can be interpreted as percent changes (more on rescaling later).\n",
    "\n",
    "We know the 'true' values of the parameters, because we're going to run a synthetic example, so we'll set those first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# True parameter values\n",
    "alpha_true, beta_true, sigma_true = 1, 2.5, 1.5\n",
    "\n",
    "# Size of dataset\n",
    "size = 100\n",
    "\n",
    "num_samples = 1000\n",
    "num_chains = 2\n",
    "\n",
    "# Predictor variable: random sample\n",
    "X = prng.standard_normal(size)\n",
    "\n",
    "# Simulate outcome variable\n",
    "Y = alpha_true + beta_true * X + prng.standard_normal(size) * sigma_true"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using PyMC to do Bayesian inference\n",
    "\n",
    "Next, we'll begin the Bayesian modelling part of the code. We'll be using the powerful [**PyMC**](https://docs.pymc.io/) package {cite:t}`salvatier2016probabilistic`; check out the documentation for tutorials and examples. PyMC is likely the most popular package for probabilistic programming in Python, and its computations are built on what was originally a deep learning library. The aim of the package is to enable you to write down models using an intuitive syntax for the data generating process.\n",
    "\n",
    "Let's see how to specify a model in PyMC. We'll start by putting `with pm.Model() as linear_model:`, to set the 'context'; effectively to say, 'what follows will be a model that we'll save as `linear_model`. Any variables created within a given model's context will be automatically assigned to that model. Next, we set our priors, giving the parameters meaningful names. Then, we build our model by writing down our expected outcome and the expected likelihood of observations according to that model. Up till now, we're just rewriting in code all the bits we put down mathematically in the previous subsection. Finally, we run the 'trace': this is the part that estimates the posterior distribution by taking samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with pm.Model() as linear_model:\n",
    "    # Priors for unknown model parameters\n",
    "    alpha = pm.Normal(\"alpha\", mu=0, sigma=10)\n",
    "    beta = pm.Normal(\"beta\", mu=0, sigma=10)\n",
    "    sigma = pm.HalfNormal(\"sigma\", sigma=1)\n",
    "\n",
    "    # Expected value of outcome\n",
    "    mu = alpha + beta * X\n",
    "\n",
    "    # Likelihood (sampling distribution) of observations\n",
    "    # this says, conditional on confounders, we expect normal variation\n",
    "    Y_obs = pm.Normal(\"Y_obs\", mu=mu, sigma=sigma, observed=Y)\n",
    "\n",
    "    # Sample the posterior\n",
    "    trace = pm.sample(num_samples, return_inferencedata=True, chains=num_chains)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So what just happened? The 'NUTS sampler' was auto-assigned. NUTS stands for the 'No-U-turn sampler'. This is a just a fancy way of picking the points to sample $p({\\boldsymbol{y}}|{\\boldsymbol{\\theta }})p({\\boldsymbol{\\theta }})$ to build up a picture of $p({\\boldsymbol{\\theta }}|{\\boldsymbol{y}})$. It's a good default option for sampling.\n",
    "\n",
    "Let's take a quick look at the model we just estimated using the `model_to_graphviz` method—though note that this is a very simple model, and this method comes into its own in more complex cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pm.model_to_graphviz(linear_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We've made some assumptions in assembling this model, and it's a good idea to make these explicit. The first is that the outcome data are independent once controlling for $X$, the second that the outcomes are a linear function of the inputs, and the third that, at a given $X$, the outcome variable $Y$ will vary normally around the average value with a consistent standard deviation $\\sigma$. These are the assumptions for Bayesian normal regression."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Understanding the results from running a Bayesian model\n",
    "\n",
    "As discussed in the Introduction, the set of samples forms a trace, and this trace got returned by the `pm.sample` method above. We can now look at the trace that we created in the 2 separate chains we ran. Here it is as a standalone object to explore:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But it's much easier to inspect some of what it holds using the pre-configured summaries, which are in table and chart form. First, a plot of the traces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "az.plot_trace(trace);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This looks like fairly meaningless noise, so let's walk through what we're seeing here. On the left-hand side, we can see the results of the two chains we ran for each of the three parameters. The x-axis is the parameter value, the y-axis is the probability density (no values shown). We're actually looking at the posterior distributions here! We'll come on to whether they make sense or not later, but for now we just know we've got some estimates of the posterior.\n",
    "\n",
    "On the right-hand side, we have some traces, with the sample number on the x-axis and the parameter value on the y axis. Just looking at these gives our first sense of whether the sampling has gone well or not. Trace plots can give a sense of whether the model converged (you should see a trace that is a random scatter around a mean value), whether the chains are mixing (they should overlap randomly on the plot), and whether there are specific parameters that the model is struggling to estimate. If the sampling is not going well, neither of these will be true and the chains will not converge to occupy a small area, nor will they mix together.\n",
    "\n",
    "From the above chart, it looks like we successfully created some well-behaved samples and have what appears to be a good estimate of the posterior. \n",
    "\n",
    "Now let's look at our actual estimates of these parameters by running the `summary` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "az.summary(trace, round_to=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In Bayesian inference, the coefficient or parameter estimates are extremely easy to interpret. For $\\alpha$, we have a mean of 1.03 with a standard deviation of $\\approx 0.15$, and a 94% confidence window between approximately 0.75 and 1.3. HDI stands for highest density interval, and you can choose a different percentage using the `hdi_prob=` keyword. We can produce a visualisation of the parameter estimates:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "az.plot_posterior(trace, var_names=['alpha'], figsize=(8, 5));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is telling us exactly what it looks like; the estimate is, with good confidence, around 1.\n",
    "\n",
    "There were a bunch of other numbers in the trace summary; what did they mean? Let's run through them:\n",
    "\n",
    "- The Gelman-Rubin $\\hat{R}$ statistic, `r_hat`: this is a useful metric to determine whether the traces achieved stationarity, but it does require a number of chains (at least four) to make sense as a diagnostic. For a given parameter, the statistic is the ratio of the variance in the samples between the chains to the variance of samples within the chains, ie $\\hat{R} = \\frac{\\text{variance between chains}}{\\text{variance within chains}}$. The value of $\\hat{R}$ approaches unity if the chains are properly sampling the target distribution and a rough guide is that it should be less than 1.01.\n",
    "- Effective sample sizes, `ess_bulk` and `ess_tail`: ideally, samples would be independent but MCMC samples are correlated. These statistics give a sense how many *effectively independent* samples we are drawing given that correlation. The statistics are referred to either as effective samples size (ESS) or number of effective samples, $n_\\text{eff}$. The bulk version, `ess_bulk`, is the main statistic, while `ess_tail` is the effective sample size for more extreme values of the posterior (by default, the lower and upper 5th percentiles). The rule of thumb is that both should be $>400$.\n",
    "- The Monte Carlo standard errors (MCSE), `msce_mean` and `mcse_sd`: They are measurements of the standard error of the mean and the standard error of the standard deviation of the sample chains. They provide an estimate as to how accurate the expectation values given from MCMC samples of the mean and standard deviation are. However, some believe these should not be reported.\n",
    "\n",
    "The most useful numbers (apart from the coefficient estimates), then, are `ess_bulk` and `ess_tail`, and $\\hat{R}$. Remember that these statistics are mostly tell us how good the *sampling* is, rather than how appropriate the *model* is."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Checking the Posterior Distribution\n",
    "\n",
    "It would be good to compare the posteriors to the priors, to see how far we got away from those weakly informative normal (and, in one case, half-normal) priors we used and how nice the posteriors look. Remember, it's not a good sign if our posterior is very much influenced by our prior (unless we have a reason to have a strong prior).\n",
    "\n",
    "To do this we need to assemble all of the information on our model, priors, and posteriors in one object, for which there's a couple of extra steps. We'll stash all of the information about our model in the object called `trace`, but we'll add a lot more info to it.\n",
    "\n",
    "The first task is to add samples of the posterior predictive distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with linear_model:\n",
    "    pm.sample_posterior_predictive(trace, extend_inferencedata=True, random_seed=prng)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now if we look at trace again, it has been extended with 'posterior predictive' values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, adding more information, we'll now sample the predictive model based on just the prior, and extend the trace data to include it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "idata = pm.sample_prior_predictive(model=linear_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trace.extend(idata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay, now we can compare our priors and posteriors. Let's do it for just one variable, $\\sigma$, for which our prior was a half normal distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "az.plot_dist_comparison(trace, var_names=[\"sigma\"], figsize=(8, 5));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see here that the weakly informative prior, which has a relatively wide initial span, has been shrunk to a posterior that is in a tight bunch around the coefficient we set in the true data generating process.\n",
    "\n",
    "Another typical check that is run is a check on whether the model can reproduce the patterns observed in the real data. We can use the plot for posterior/prior predictive checks to do this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "az.plot_ppc(trace, group=\"posterior\", figsize=(8, 5));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While there are lots of diagnostic charts available in [arviz](https://arviz-devs.github.io/arviz/), in this case it's useful to roll our own to check visually how our model predicts the data, taking into account the heterogeneity in our data in a way that the posterior predictive check chart we just saw does not.\n",
    "\n",
    "First, we'll extract the samples that we already made of the posterior. Then we'll use `xarray` (a package for working with labelled multi-dimensional data that is a dependency of **PyMC**) to multiply these samples through by the input data, $X$. This will give us an output variable `mu_pp` that has shape (number of chains, number of samples, number of observations). You can see this below (and we'll have to take an average over some of these dimensions to get a vector that has shape 'number of observations' only).\n",
    "\n",
    "Don't worry too much about the details of what's going on; mostly it's just constructing the equation we're trying to estimate with the results of our Bayesian inference but with the added complication of multiple chains and multiple samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "post = trace.posterior\n",
    "mu_pp = post[\"alpha\"] + post[\"beta\"] * xr.DataArray(X, dims=[\"obs_id\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mu_pp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can plot the mean outcomes and the data, along with confidence intervals for the mean outcomes and for all predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(X, Y, \"o\", ms=4, alpha=0.4, label=\"Data\")\n",
    "ax.plot(X, mu_pp.mean(axis=(0, 1)), label=\"Mean outcomes\", alpha=1, lw=0.2, color=\"k\")\n",
    "az.plot_hdi(\n",
    "    X,\n",
    "    mu_pp,\n",
    "    ax=ax,\n",
    "    fill_kwargs={\"alpha\": 0.5, \"color\": \"orchid\", \"label\": \"Mean Outcome 94% HPD\"},\n",
    ")\n",
    "az.plot_hdi(\n",
    "    X,\n",
    "    trace[\"posterior_predictive\"][\"Y_obs\"],\n",
    "    ax=ax,\n",
    "    fill_kwargs={\"alpha\": 0.2, \"color\": \"coral\", \"label\": \"Outcome 94% HPD\"},\n",
    ")\n",
    "ax.set_xlabel(\"Predictor (X)\")\n",
    "ax.set_ylabel(\"Outcome (Y)\")\n",
    "ax.set_title(\"Posterior predictive checks\")\n",
    "ax.legend(ncol=2, fontsize=10);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fixed Effects and Categorical Variables As Regressors\n",
    "\n",
    "How do we go about fitting a model with a dummy or fixed effect using a categorical variable in Bayesian land? (At least when it's an exogenous variable—we'll come back to the case when the outcome variable is discrete later).\n",
    "\n",
    "Let's take the previous example and modify to include a group-specific intercept for two groups—so a dummy variable that is 0 for group 0, and 1 for group 1. Our model will now be:\n",
    "\n",
    "$$\n",
    "\\mu = \\alpha + \\beta x + \\gamma d\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "$$\n",
    "Y \\mid \\alpha, \\beta, \\sigma, \\gamma \\stackrel{\\text{ind}}{\\thicksim} \\mathcal{N}(\\mu, \\sigma^2)\n",
    "$$\n",
    "\n",
    "For our priors, we will assume:\n",
    "\n",
    "- $\\alpha \\thicksim \\mathcal{N}(0, 10)$\n",
    "- $\\beta \\thicksim \\mathcal{N}(0, 10)$\n",
    "- $\\sigma \\thicksim \\mathcal{N_{+}}(1)$ (the half-normal distribution)\n",
    "- $\\gamma \\thicksim \\mathcal{N}(0.5, 10)$\n",
    "\n",
    "where\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "    f(y;\\mu, \\sigma) = \\left\\{\\begin{array}{cll}\n",
    "\\sqrt{\\frac{2}{\\pi\\sigma^2}}\\,\\mathrm{e}^{-(y-\\mu)^2/2\\sigma^2} &  & y \\ge \\mu \\\\[1em]\n",
    "0 & & \\text{otherwise}.\n",
    "\\end{array}\\right.\n",
    "    \\end{align}\n",
    "$$\n",
    "\n",
    "is the half-normal distribution.\n",
    "\n",
    "Note that we are giving $\\gamma$ a mean prior that is not 0, because we do not expect there to be no influence at all from the group 1 contributions.\n",
    "\n",
    "Let's create the synthetic data for this model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# True parameter values\n",
    "α_true, β_true, σ_true, γ_true = 1, 2.5, 1.5, 6\n",
    "\n",
    "# Size of dataset\n",
    "size = 200\n",
    "\n",
    "num_samples = 1000\n",
    "num_chains = 4\n",
    "\n",
    "# Predictor variables\n",
    "X_cat = prng.standard_normal(size)\n",
    "D = prng.binomial(1, 0.4, size)  # This chooses 1 or 0 with 0.4 prob for 1\n",
    "\n",
    "# Simulate outcome variable\n",
    "Y_cat = α_true + β_true * X_cat + γ_true*D + prng.standard_normal(size) * σ_true"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now a model much like we did before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with pm.Model() as linear_model_with_dummy:\n",
    "    # Priors for unknown model parameters\n",
    "    α = pm.Normal(\"α\", mu=0, sigma=10)\n",
    "    β = pm.Normal(\"β\", mu=0, sigma=10)\n",
    "    γ = pm.Normal(\"γ\", mu=0, sigma=10)\n",
    "    σ = pm.HalfNormal(\"σ\", sigma=1)\n",
    "\n",
    "    # Expected value of outcome\n",
    "    μ = α + β * X_cat + γ * D\n",
    "\n",
    "    # Likelihood (sampling distribution) of observations\n",
    "    Y_cat_obs = pm.Normal(\"Y_obs\", mu=μ, sigma=σ, observed=Y_cat)\n",
    "\n",
    "    # Sample the posterior\n",
    "    trace = pm.sample(num_samples, return_inferencedata=True, chains=num_chains)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we check the trace summary statistics—we won't look at the trace visually this time, because we know what the answer should be!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "az.summary(trace, round_to=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is close to the parameters we used to generate the data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Priors\n",
    "\n",
    "This section has benefitted from [Probabilistic Programming and Bayesian Methods for Hackers](https://nbviewer.org/github/CamDavidsonPilon/Probabilistic-Programming-and-Bayesian-Methods-for-Hackers/tree/master/) and the [prior choice recommendations page](https://github.com/stan-dev/stan/wiki/Prior-Choice-Recommendations) of the Stan wiki.\n",
    "\n",
    "### Choosing Priors\n",
    "\n",
    "While **Bambi** (which we'll come to shortly) will choose priors for you, you can override them—and there are many situations where you may wish to do this. And for **PyMC**, you always have to specify your priors. But how do you choose?\n",
    "\n",
    "There are two important paradigms to be aware of when choosing priors: *objective* priors, where the data influence the posterior the most, and *subjective* priors, where the economist expresses their views in the form of the prior.\n",
    "\n",
    "The quintessential example of an *objective prior* is the flat prior, literally a uniform distribution over the entire possible range of the variable that is to be estimated. By choosing this, you are saying that you put equal weight on all possible values. You would very rarely use this in practice, and priors with soft boundaries are generalled preferred over those with hard boundaries.\n",
    "\n",
    "A *subjective* or *informative* prior, however, would have more probability density or mass on certain values, biasing the final estimate toward that part of the range. You might use this if there were lots of previous evidence for something, or if there were extra information outside of the model that you wanted to incorporate. An informative prior reflects a higher degree of certainty about where the final values might end up, and will take a stronger pull to overcome.\n",
    "\n",
    "Inbetween these two extremes are weakly informative priors, which are the default in **Bambi** and a good place to start in **PyMC** models, especially if your data are appropriately normalised.\n",
    "\n",
    "We've already seen some handy priors such as the normal distribution, half-normal distribution, and Student-t distribution (plus it's half equivalent too). You can find a full list of priors supported by **PyMC** [here](https://docs.pymc.io/en/v3/api/distributions.html), in the section on distributions. Some priors to check out are:\n",
    "\n",
    "- Cauchy and half-Cauchy\n",
    "- Beta\n",
    "- Gamma\n",
    "- Binomial (discrete)\n",
    "- Bernoulli (discrete)\n",
    "- Poisson (discrete)\n",
    "\n",
    "If you need a generic prior as a place to start for anything, some people recommend $\\mathcal{N}(0, 1)$ or $\\text{Student-t}(3,0,1)$, assuming your data is appropriately normalised and could be negative or positive.\n",
    "\n",
    "### Scaling Data\n",
    "\n",
    "Aim to keep all the parameters in your model \"scale-free\". This means adjusting them so that a unit increase is meaningful. This will help in numerous ways, including giving a chance for priors like $\\mathcal{N}(0, 1)$ or $\\mathcal{N}(0, 10)$ to work.\n",
    "\n",
    "Some suggested ways of scaling your data could be:\n",
    "\n",
    "- Scale by the standard deviation of the data via $x \\longrightarrow x/\\sigma_{x}$, for example, if you were looking at test scores you might divide by the standard deviation of all the test scores.\n",
    "- In a regression, take logs of (positive-constrained) predictors and outcomes, then the coefficients can be interpreted as elasticities, $x \\longrightarrow \\ln x$\n",
    "- If a variable has a typical value, say 4, then scale by that under a logarithm, $x \\longrightarrow \\ln(x/4)$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Other Bayesian Packages\n",
    "\n",
    "Although **PyMC** and **Bambi** (which will be featured in {ref}`econmt-bayes-bambi`) are the most accessible and, likely, most widely-used Bayesian inference packages in Python, they're certainly not the only ones. Here are a selection of others you might wish to check out:\n",
    "\n",
    "- [PyStan](https://pystan.readthedocs.io/en/latest/), the Python wrapper for popular C++ Bayesian library Stan\n",
    "- [cmdstanpy](https://github.com/stan-dev/cmdstanpy), another Python wrapper for Stan\n",
    "- [Pyro](https://pyro.ai/), \"deep probabilistic modelling, unifying the best of modern deep learning and Bayesian modelling\", built on Facebook's PyTorch deep learning package\n",
    "- [Bean Machine](https://beanmachine.org/), \"A universal probabilistic programming language to enable fast and accurate Bayesian analysis\" from Facebook, also built on PyTorch.\n",
    "- [Tensorflow Probability](https://www.tensorflow.org/probability), \"makes it easy to combine probabilistic models and deep learning on modern hardware (TPU, GPU)\", built on Google's Tensorflow deep learning package\n",
    "- [pomegranite](https://pomegranate.readthedocs.io/en/latest/), \"implements fast and flexible probabilistic models ranging from individual probability distributions to compositional models such as Bayesian networks and hidden Markov models\"."
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "c4570b151692b3082981c89d172815ada9960dee4eb0bedb37dc10c95601d3bd"
  },
  "kernelspec": {
   "display_name": "Python 3.8.12 64-bit ('codeforecon': conda)",
   "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}