Lab_7_Intro_Bayesian_Inference.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyM4Z4A8MN7Pqnj/3G3Z3EhJ",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/USCbiostats/PM520/blob/main/Lab_7_Intro_Bayesian_Inference.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Don't Stop Believin', or: Intro to Bayesian Inference\n",
        "$\\newcommand{\\data}{\\text{Data}}$\n",
        "$\\newcommand{\\E}{\\mathbb{E}}$\n",
        "So far, we've focused primarily on [likelihood](https://en.wikipedia.org/wiki/Likelihood_function)-based inference using [maximum likelihood estimation](https://en.wikipedia.org/wiki/Maximum_likelihood_estimation), or MLE. Recall that MLE roughly seeks $$\\hat{\\theta}_{MLE} := \\arg \\max_{\\theta} \\ell(\\theta | \\data) = \\arg \\max_{\\theta} \\log \\Pr(\\data | \\theta),$$\n",
        "where $\\ell(\\theta | \\data)$ is a log-likelihood function and $\\theta$ are the parameters of interest. This procedure operationally reflects identifying a value of $\\theta$ such that our observed $\\data$ is most likely.\n",
        "\n",
        "In contrast to this regime, [_Bayesian_ inference](https://en.wikipedia.org/wiki/Bayesian_inference) operationally reflects updating our _beliefs_ about $\\theta$ conditioned on having observed $\\data$. The celebrated [Bayes' Theorem](https://en.wikipedia.org/wiki/Bayes%27_theorem) states,\n",
        "$$\\Pr(\\theta | \\data) = \\frac{\\Pr(\\data | \\theta) \\Pr(\\theta)}{\\Pr(\\data)},$$\n",
        "where $\\Pr(\\theta | \\data)$ is the [_posterior_ probability](https://en.wikipedia.org/wiki/Posterior_probability) for $\\theta$ and reflects our uncertainty in the values that $\\theta$ may take on, $\\Pr(\\data | \\theta)$ is our likelihood, $\\Pr(\\theta)$ is a [_prior_ probability](https://en.wikipedia.org/wiki/Prior_probability) (or _prior_) over $\\theta$ and $\\Pr(\\data)$ is a [_marginal_ probability/likelihood](https://en.wikipedia.org/wiki/Marginal_likelihood) of the data. In the case that $\\theta$ is continuous we have,\n",
        "$$ \\Pr(\\data) =\\int \\Pr(\\data | \\theta') \\Pr(\\theta')d\\theta'$$ and in the case\n",
        "that $\\theta$ is discrete we have,\n",
        "$$\\Pr(\\data) = \\sum_{\\theta'} \\Pr(\\data | \\theta') \\Pr(\\theta').$$\n",
        "\n",
        "\n",
        "\n",
        "## Example\n",
        "Say we have some examination which determines whether a _sick_ individual is sick with a TP rate of 0.9. The prevalance of this disease in the population is 0.05. If someone takes this test to determine if they are sick, what is the probability they are indeed sick? Formally, $\\newcommand{\\sick}{\\text{sick}}$\n",
        "$$\\begin{align*}\\Pr(\\sick | +) &=\n",
        "  \\frac{\\Pr(+ | \\sick) \\Pr(\\sick)}{\\Pr(+)} =\n",
        "  \\frac{\\Pr(+ | \\sick) \\Pr(\\sick)}{\\Pr(+ | \\sick)\\Pr(\\sick) + \\Pr(+ | \\neg \\sick)\\Pr(\\neg \\sick)} \\\\\n",
        "  &= \\frac{0.9 \\times 0.05}{(0.9 \\times 0.05) + (0.1 \\times 0.95)} \\approx 0.32\n",
        "  \\end{align*}$$"
      ],
      "metadata": {
        "id": "QYQci7HuXVeU"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Bayesian Estimators\n",
        "When performing Bayesian inference, we often seek to identify the posterior expectation of $\\theta$ which is given by, $$\\E[\\theta | \\data] = \\int \\theta \\Pr(\\theta | \\data) d\\theta$$ (or analogously using summation for discrete $\\theta)$. Importantly, this is a _conditional_ expectation! This is the expected value (or mean) given our observations $\\data$, which is different from the unconditional, or _prior_ expectation $\\E[\\theta] = \\int \\theta \\Pr(\\theta) d \\theta$. This further reflects how our beliefs about the values $\n",
        "\\theta$ may take on change after having observed $\\data$.\n",
        "\n",
        "As currently stated, it is somewhat unclear why we should select $\\hat{\\theta} = \\E[\\theta | \\data]$ as our [estimator](https://en.wikipedia.org/wiki/Bayes_estimator), but we can reformulate this procedure under a _risk_ framework. Let's define the risk of some estimate $\\hat{\\theta}$ as, $$\\E[L(\\theta, \\hat{\\theta}) | \\data]$$ where $L(\\theta, \\hat{\\theta})$ is a [loss function](https://en.wikipedia.org/wiki/Loss_function) that reflects some notion of _distance_ between some putative \"true\" value $\\theta$ and our estimate $\\hat{\\theta}$. If we select a quadratic loss, we have\n",
        "$$\\E[(\\theta - \\hat{\\theta})^2 | \\data],$$ which is the [mean squared error](https://en.wikipedia.org/wiki/Mean_squared_error), or MSE. With a bit of algebra, it's clear that our estimator $\\hat{\\theta}$ should result in $\\hat{\\theta} = \\E[\\theta | \\data]$, which is the posterior expectation!\n",
        "\n",
        "An alternative approach may be to seek the values $\\theta$ which maximize our posterior $\\Pr(\\theta | \\data)$, hence its name [maximum a-posteriori](https://en.wikipedia.org/wiki/Maximum_a_posteriori_estimation), or MAP for short. Its definition is given by, $$\\hat{\\theta}_{MAP} := \\arg \\max_\\theta \\log \\Pr(\\theta | \\data).$$ While this can be a valid approach in some sense, it seeks a value that is most-probable given our observations $\\data$, or the posterior [mode](https://en.wikipedia.org/wiki/Mode_(statistics)), which not be a representitive example due to its ignorance of the total density landscape."
      ],
      "metadata": {
        "id": "1mySqE7zlGsR"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Example\n",
        "Let's say we have some possibly biased coin whose outcomes are heads (H) or tails (H). We can model the likelihood of observing the outcomes from a sequence of $N = N_H + N_T$ tosses using the [Binomial distribution](https://en.wikipedia.org/wiki/Binomial_distribution), given by $$\\Pr(N_H, N_T | p_H) = \\begin{pmatrix}N \\\\ N_H\\end{pmatrix} p_H^{N_H}(1 - p_H)^{(N - N_H)}.$$ Furthermore we can model our _prior_ belief of the probability of observing H by a [Beta distribution](https://en.wikipedia.org/wiki/Beta_distribution) as,\n",
        "$$\\Pr(p_H) = \\frac{1}{B(\\alpha, \\beta)}p_H^{\\alpha - 1}(1 - p_H)^{\\beta - 1}.$$\n",
        "\n",
        "Given these two definitions, we can state the posterior probability of $p_H$ given $N_H, N_T$ as,\n",
        "\n",
        "$$\\begin{align*}\n",
        "\\Pr(p_H | N_H, N_T) &= \\frac{\\Pr(N_H, N_T | p_H) \\Pr(p_H)}{\\Pr(N_H, N_T)} \\\\\n",
        "  &\\propto \\Pr(N_H, N_T | p_H) \\Pr(p_H) \\\\\n",
        "  &= \\begin{pmatrix}N \\\\ N_H\\end{pmatrix} p_H^{N_H}(1 - p_H)^{N_T}\n",
        "  \\times \\frac{1}{B(\\alpha, \\beta)}p_H^{\\alpha - 1}(1 - p_H)^{\\beta - 1} \\\\\n",
        "  &\\propto p_H^{N_H}(1 - p_H)^{N_T} \\times p_H^{\\alpha - 1}(1 - p_H)^{\\beta - 1} \\\\\n",
        "  &= p_H^{(N_H + \\alpha - 1)}(1 - p_H)^{(N_T + \\beta - 1)} \\\\\n",
        "  &= \\text{Beta}(p_H | \\widetilde{\\alpha}, \\widetilde{\\beta}),\n",
        "\\end{align*}$$\n",
        "where $\\widetilde{\\alpha} = N_H + \\alpha$ and $\\widetilde{\\beta} = N_T + \\beta$.\n",
        "\n",
        "Here, our posterior is of the same distributional family as its prior. This property is known as [conjugacy](https://en.wikipedia.org/wiki/Conjugate_prior), which we will explore in depth in a moment. In general we will not be so lucky as to have a vast simplification of the distribution to a known form, and we will require sampling (e.g., [MCMC](https://en.wikipedia.org/wiki/Markov_chain_Monte_Carlo) or related approaches) to identify posterior means (or in general its moments)."
      ],
      "metadata": {
        "id": "gHv3JnpunjIP"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VyrlQR66XIU_",
        "outputId": "a6c2b631-eebb-4cfc-f278-d5fd7e8cffd5"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "p_H = 0.406 | MLE[p_H] = 0.405 | Posterior estimate E[p_h | N_H, N_T] = 0.405, given $\\alpha$ = 5, $\\beta$ = 5$\n"
          ]
        }
      ],
      "source": [
        "# let's code up the above to evalute some simulated data!\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import jax.random as rdm\n",
        "\n",
        "seed = 1\n",
        "key = rdm.PRNGKey(seed)\n",
        "\n",
        "Nobs = 50000 # total number of tosses\n",
        "alpha = 5 # prior number of 'psuedo' observations for heads (ie '1')\n",
        "beta = 5  # prior number of 'psuedo' observations for tails (ie '0')\n",
        "# E[ph] = alpha / (alpha + beta)\n",
        "\n",
        "# first simulate a probability for heads (ie sample from p_H ~ Beta(\\alpha, \\beta))\n",
        "key, p_key = rdm.split(key)\n",
        "ph = rdm.beta(p_key, alpha, beta)\n",
        "\n",
        "# then simulate our observations N_H ~ Binomial(N, p_H)\n",
        "key, x_key = rdm.split(key)\n",
        "Nh = rdm.binomial(x_key, Nobs, ph)\n",
        "\n",
        "# then infer the posterior mean for p_H, namely E[p_H | N_H, N_T];\n",
        "# for Beta(\\alpha, \\beta), $\\E[p_H] = \\frac{\\alpha}{\\alpha + \\beta}$\n",
        "# recall that posterior is Beta(Nh | Nh + alpha, Nt + beta)\n",
        "post_mean_ph = (Nh + alpha - 1) / (Nobs + alpha + beta)\n",
        "\n",
        "# let's compare our bayesian estimate for p_H with the MLE estimate\n",
        "mle_ph = Nh / Nobs\n",
        "print(f\"p_H = {ph:.3f} | MLE[p_H] = {mle_ph:.3f} | Posterior estimate E[p_h | N_H, N_T] = {post_mean_ph:.3f}, given $\\\\alpha$ = {alpha}, $\\\\beta$ = {beta}$\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Conjugate Priors and Exponential Families\n",
        "Recall our definition of exponential families, such that the probability of an observation is given by $$\\Pr(x | \\eta) \\propto \\exp(\\eta \\cdot T(x) - A(\\eta)).$$ If our _prior_ over $\\eta$ is  _also_ in the exponential family and can be written as,\n",
        "$$\\Pr(\\eta) \\propto \\exp(\\eta \\cdot \\chi - \\nu A(\\eta)),$$\n",
        "then our posterior simplifies to,\n",
        "$$\\begin{align*}\n",
        "\\Pr(\\eta | x) &\\propto \\Pr(x | \\eta)\\Pr(\\eta)\\\\\n",
        " &\\propto \\exp(\\eta \\cdot T(x) - A(\\eta)) \\exp(\\eta \\chi - \\nu A(\\eta)) \\\\\n",
        " &= \\exp(\\eta \\cdot T(x) - A(\\eta) + \\eta \\cdot \\chi - \\nu A(\\eta)) \\\\\n",
        " &= \\exp(\\eta \\cdot(T(x) + \\chi) - A(\\eta)(\\nu + 1)),\n",
        "\\end{align*}$$ or more generally, given $n$ observations as\n",
        "\n",
        "$$\\begin{align*}\n",
        "\\Pr(\\eta | x) &\\propto \\exp(\\eta \\cdot(\\sum_{i=1}^n T(x_i) + \\chi) - A(\\eta)(\\nu + n))\\\\\n",
        "   &= \\exp(\\eta \\underbrace{\\cdot(T(X) + \\chi)}_{\\widetilde{\\chi}}\n",
        "   -A(\\eta)\\underbrace{(\\nu + n)}_{\\widetilde{\\nu}}),\n",
        " \\end{align*}$$\n",
        " where $T(X) = \\sum_{i=1}^n T(x_i)$ and $\\widetilde{\\chi}$ are the _posterior_ natural parameters and $\\widetilde{\\nu}$ is the number of _psuedo_ observations (i.e. actual number of observations $n$ and prior pseudo number $\\nu$).\n",
        "\n",
        " # Revisit above example\n",
        " Recall, the ExpFam form for Binomial is given by,\n",
        " $$\\begin{align*}\n",
        " \\Pr(N_H | p_h) &\\propto \\exp(\\eta \\cdot T(N_H) - A(\\eta)) \\\\\n",
        "   &= \\exp(\\eta \\times N_H - N \\log(1 + \\exp(\\eta))\\\\\n",
        "   &= \\exp(\\log(\\frac{p_h}{1 - p_h}) \\times N_H - N \\log(1 + \\exp(\\log(\\frac{p_h}{1 - p_h}))) \\\\\n",
        "   &= \\exp(\\log(\\frac{p_h}{1 - p_h}) \\times N_H - N \\log(1 + \\frac{p_h}{1 - p_h})) \\\\\n",
        "   &= \\exp(\\log(\\frac{p_h}{1 - p_h}) \\times N_H) \\exp(- N \\log(1 + \\frac{p_h}{1 - p_h})) \\\\\n",
        "   &= p_h^{N_H}(1 - p_h)^{-N_H} \\frac{1}{1 + \\frac{p_h}{1 - p_h}}^N \\\\\n",
        "   &= p_h^{N_H}(1 - p_h)^{-N_H} (1 - p_h)^{N} \\\\\n",
        "   &= p_h^{N_H}(1 - p_h)^{(N - N_H)} \\\\\n",
        "   &= p_h^{N_H}(1 - p_h)^{N_T}.\n",
        " \\end{align*}$$\n",
        " Similarly, lets recall that the Beta distribution can be written as,\n",
        " $$\\begin{align*}\n",
        " \\Pr\\left(\\log\\left(\\frac{p_h}{1 - p_h}\\right)\\right) = \\Pr(\\eta) &\\propto \\exp(\\chi \\cdot T(\\eta) - A(\\eta)) \\\\\n",
        " \\end{align*}$$\n",
        " TBD: Nick finish notes."
      ],
      "metadata": {
        "id": "3g6cTU8wnnqi"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Bayesian Linear Regression\n",
        "Let's take a look at a straightforward [Bayesian linear regression](https://en.wikipedia.org/wiki/Bayesian_linear_regression) problem, where we have some continuous outcome $y$ for $n$ observations that is a linear function of $n \\times p$ matrix $X$ and $p$ coefficients $\\beta$. We state our distributions assumptions below as,\n",
        "\n",
        "$$\\begin{align*}\n",
        "y | X, \\beta &\\sim N(X\\beta, I_n \\sigma^2) \\\\\n",
        "\\beta &\\sim N(0, I_p \\sigma^2_b).\n",
        "\\end{align*}\n",
        "$$\n",
        "\n",
        "What is the posterior mean of $\\beta$? That is, what is $\\mathbb{E}[\\beta | y, X]$?\n",
        "\n",
        "$$\\begin{align*}\n",
        "\\Pr(\\beta | y) &\\propto \\exp(-\\frac{1}{2\\sigma^2}(y - X \\beta)^T(y - X \\beta))\\exp(-\\frac{1}{2\\sigma^2_b}\\beta^T\\beta)\\\\\n",
        "&\\propto \\exp(-\\frac{1}{2\\sigma^2}(y - X \\beta)^T(y - X \\beta) -\\frac{1}{2\\sigma^2_b}\\beta^T\\beta) \\\\\n",
        "&= \\exp(-\\frac{1}{2\\sigma^2}(y^Ty - 2y^T X \\beta + \\beta^T X^T X \\beta) -\\frac{1}{2\\sigma^2_b}\\beta^T\\beta) \\\\\n",
        "&\\propto \\exp(\\frac{1}{\\sigma^2} y^T X \\beta -\\frac{1}{2\\sigma^2} \\beta^T X^T X \\beta -\\frac{1}{2\\sigma^2_b}\\beta^T\\beta) \\\\\n",
        "&= \\exp(\\frac{1}{\\sigma^2} y^T X \\beta -\\frac{1}{2} \\beta^T \\underbrace{(X^T X \\sigma^{-2} + I \\sigma^{-2}_b)}_{\\Sigma^{-1}} \\beta ) \\\\\n",
        "&= \\exp(\\frac{1}{\\sigma^2} y^T X \\beta -\\frac{1}{2} \\beta^T \\Sigma^{-1}\\beta ) \\\\\n",
        "&= \\exp(\\underbrace{\\frac{1}{\\sigma^2} y^T X \\Sigma}_{\\mu^T} \\Sigma^{-1}\\beta -\\frac{1}{2} \\beta^T \\Sigma^{-1}\\beta ) \\\\\n",
        "&\\propto \\exp(\\underbrace{-\\frac{1}{2}\\mu^T\\Sigma^{-1}\\mu}_{O(1)} +\\mu^T \\Sigma^{-1}\\beta -\\frac{1}{2} \\beta^T \\Sigma^{-1}\\beta ) \\\\\n",
        "&= N(\\beta | \\mu, \\Sigma),\n",
        "\\end{align*}\n",
        "$$\n",
        "where\n",
        "$$\\begin{align*}\n",
        "\\mathbb{E}[\\beta | y] = \\mu &= \\frac{1}{\\sigma^2} \\Sigma X^T y\\\\\n",
        "&= \\frac{1}{\\sigma^2} (X^T X \\sigma^{-2} + I \\sigma^{-2}_b)^{-1} X^T y\\\\\n",
        "&= (X^T X  + I \\frac{\\sigma^2}{\\sigma^2_b})^{-1} X^T y\\\\\n",
        "&= (X^T X  + I \\lambda)^{-1} X^T y.\n",
        " \\end{align*}\n",
        "$$\n",
        "Notice when $\\lambda = \\frac{\\sigma^2}{\\sigma^2_b}$ this is the same result as [Ridge Regression](https://en.wikipedia.org/wiki/Ridge_regression)!\n",
        "\n",
        "These results align due to the mode/mean being the same under Normality, since Ridge regression is an optimization problem (i.e. mode seeking under minimize squared loss + squared penalty on effect sizes).\n",
        "\n",
        "Let's code up a simulation and compare/contrast the various estimators of $\\beta$ under MLE and Posterior for diff $\\sigma^2_b$."
      ],
      "metadata": {
        "id": "rUjjY0PDz8I2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import jax.numpy as jnp\n",
        "import jax.random as rdm\n",
        "\n",
        "# simulate a bayesian linear regression\n",
        "seed = 0\n",
        "key = rdm.PRNGKey(seed)\n",
        "\n",
        "N, P = 500, 10\n",
        "\n",
        "prior_mean = jnp.zeros(P)\n",
        "prior_var = 1e-2\n",
        "\n",
        "# fake X\n",
        "key, x_key = rdm.split(key)\n",
        "X = rdm.normal(x_key, shape=(N, P))\n",
        "\n",
        "# sample beta\n"
      ],
      "metadata": {
        "id": "jJ0MuowmdEmD"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}