Lab_7_Intro_Bayesian_Inference.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNeebl+f+yiTuKHJB9kw4QD",
      "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 \\min_{\\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": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VyrlQR66XIU_",
        "outputId": "cd4cda12-a8b8-4174-e32e-29f4d95ae67c"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "MLE[p_H] = 0.801 | Posterior estimate E[p_h | N_H, N_T] = 0.800, 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 = 0\n",
        "key = rdm.PRNGKey(seed)\n",
        "\n",
        "Nobs = 5000 # 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",
        "\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",
        "ph = 0.8\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\"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 that 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": "code",
      "source": [
        "# code for above using our ExpFam classes\n",
        "!pip install equinox"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "q00LkeRdnqs_",
        "outputId": "edce06a1-1742-45d4-b584-18aaab2341d1"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting equinox\n",
            "  Downloading equinox-0.11.3-py3-none-any.whl (167 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m167.9/167.9 kB\u001b[0m \u001b[31m4.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: jax>=0.4.13 in /usr/local/lib/python3.10/dist-packages (from equinox) (0.4.23)\n",
            "Collecting jaxtyping>=0.2.20 (from equinox)\n",
            "  Downloading jaxtyping-0.2.25-py3-none-any.whl (39 kB)\n",
            "Requirement already satisfied: typing-extensions>=4.5.0 in /usr/local/lib/python3.10/dist-packages (from equinox) (4.9.0)\n",
            "Requirement already satisfied: ml-dtypes>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from jax>=0.4.13->equinox) (0.2.0)\n",
            "Requirement already satisfied: numpy>=1.22 in /usr/local/lib/python3.10/dist-packages (from jax>=0.4.13->equinox) (1.25.2)\n",
            "Requirement already satisfied: opt-einsum in /usr/local/lib/python3.10/dist-packages (from jax>=0.4.13->equinox) (3.3.0)\n",
            "Requirement already satisfied: scipy>=1.9 in /usr/local/lib/python3.10/dist-packages (from jax>=0.4.13->equinox) (1.11.4)\n",
            "Collecting typeguard<3,>=2.13.3 (from jaxtyping>=0.2.20->equinox)\n",
            "  Downloading typeguard-2.13.3-py3-none-any.whl (17 kB)\n",
            "Installing collected packages: typeguard, jaxtyping, equinox\n",
            "Successfully installed equinox-0.11.3 jaxtyping-0.2.25 typeguard-2.13.3\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from abc import abstractmethod\n",
        "\n",
        "import equinox as eqx\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "from jaxtyping import Array, ArrayLike\n",
        "\n",
        "\n",
        "class ExpFam(eqx.Module):\n",
        "  \"\"\"\n",
        "  Simple base class for Exponential Families. Will keep track of\n",
        "  observations as internal state and provide means to evaluate\n",
        "  the loglikelihood of downstream implementations.\n",
        "\n",
        "  x: Array, our observations\n",
        "  \"\"\"\n",
        "\n",
        "  x: Array\n",
        "\n",
        "  @abstractmethod\n",
        "  def base_measure(self) -> Array:\n",
        "    \"\"\"\n",
        "    Computes the base measure for an implementation of ExpFam.\n",
        "\n",
        "    Returns:\n",
        "      The base measure for an implementation of ExpFam.\n",
        "    \"\"\"\n",
        "    ...\n",
        "\n",
        "  @abstractmethod\n",
        "  def sufficient_statistics(self) -> Array:\n",
        "    \"\"\"\n",
        "    Computes the sufficient statistics (i.e. $T(x)$) for an implementation\n",
        "    of ExpFam.\n",
        "\n",
        "    Returns:\n",
        "      The sufficient statistics for each observation under an implementation\n",
        "      of ExpFam.\n",
        "    \"\"\"\n",
        "    ...\n",
        "\n",
        "  @abstractmethod\n",
        "  def log_partition(self, eta: ArrayLike) -> Array:\n",
        "    \"\"\"\n",
        "    Computes the log partition function (i.e. $A(\\eta)$) for an implementation\n",
        "    of ExpFam with natural parameters $\\eta$.\n",
        "\n",
        "    eta: ArrayLike, the natural parameters to evaluate under the log partition.\n",
        "\n",
        "    Returns:\n",
        "      The value of the log partition function for each observation under an\n",
        "      implementation of ExpFam.\n",
        "    \"\"\"\n",
        "    ...\n",
        "\n",
        "  def loglikelihood(self, eta: ArrayLike) -> Array:\n",
        "    \"\"\"\n",
        "    Computes the log likelihood for each observation $x$ under an implementation\n",
        "    of ExpFam with natural parameters $\\eta$.\n",
        "\n",
        "    eta: ArrayLike, the natural parameters to evaluate under the log likelihood.\n",
        "\n",
        "    Returns:\n",
        "      The log likelihood for each observation under an implementation of ExpFam.\n",
        "    \"\"\"\n",
        "\n",
        "    t_x = self.sufficient_statistics()\n",
        "    log_h_x = jnp.log(self.base_measure())\n",
        "    log_eta = self.log_partition(eta)\n",
        "    return t_x @ eta - log_eta + log_h_x\n",
        "\n",
        "  def logposterior(self, eta: ArrayLike, chi: ArrayLike, nu: ArrayLike) -> Array:\n",
        "    \"\"\"\n",
        "    Computes the log posterior for each $\\eta$ assuming conjugacy under an implementation\n",
        "    of ExpFam with natural parameters $\\eta$.\n",
        "\n",
        "    eta: ArrayLike, the natural parameters to evaluate under the log likelihood.\n",
        "    chi: ArrayLike, the prior natural parameters\n",
        "    nu: ArrayLike, the prior pseudo-count\n",
        "\n",
        "    Returns:\n",
        "      The log posterior for eta, given observations x under an implementation of ExpFam.\n",
        "    \"\"\"\n",
        "    ...\n",
        "\n",
        "  def __call__(self, eta: ArrayLike) -> Array:\n",
        "    return self.loglikelihood(eta)\n",
        "\n",
        "\n",
        "class Normal(ExpFam):\n",
        "\n",
        "  def base_measure(self) -> Array:\n",
        "    return 1. / jnp.sqrt(2 * jnp.pi)\n",
        "\n",
        "  def sufficient_statistics(self) -> Array:\n",
        "    x = self.x\n",
        "    x_sq = self.x ** 2\n",
        "    return jnp.concatenate((x[:, jnp.newaxis], x_sq[:, jnp.newaxis]), axis=1)\n",
        "\n",
        "  def log_partition(self, eta: ArrayLike) -> Array:\n",
        "    eta_1, eta_2 = eta\n",
        "    term1 = -(eta_1**2 / (4 * eta_2))\n",
        "    term2 = -0.5 * jnp.log(- 2 * eta_2)\n",
        "    return term1 + term2"
      ],
      "metadata": {
        "id": "jJ0MuowmdEmD"
      },
      "execution_count": 10,
      "outputs": []
    }
  ]
}