python/ch13/13.7-statsville-revisited.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0b10f741",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.distributions import Normal\n",
    "import math"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79bdab58",
   "metadata": {},
   "source": [
    "Let us revisit the problem of predicting if a resident of Statsville is female based on the height. For this purpose, we have collected a set of height samples from adult female residents in Statsville. Unfortunately, due to unforseen circumstances we have collected a very small sample from the residents. Armed with our knowledge of Bayesian inference, we do not want to let this deter us from trying to build a model.\n",
    "\n",
    "From physical considerations, we can assume that the distribution of heights is Gaussian. Our goal is to estimate the parameters ($\\mu$, $\\sigma$) of this Gaussian.\n",
    "\n",
    "\n",
    "Let us first create the dataset by sampling 5 points from a Gaussian distribution with $\\mu$=152 and $\\sigma$=8. In real life scenarios, we do not know the mean and standard deviation of the true distribution. But for the sake of this example, let's assume that the mean height is 152cm and standard deviation is 8cm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "eeaff6b6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset shape: torch.Size([5, 1])\n"
     ]
    }
   ],
   "source": [
    "torch.random.manual_seed(0)\n",
    "num_samples = 5\n",
    "true_dist = Normal(152, 8)\n",
    "X = true_dist.sample((num_samples, 1))\n",
    "print('Dataset shape: {}'.format(X.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "787cea38",
   "metadata": {},
   "source": [
    "### Maximum Likelihood Estimate\n",
    "\n",
    "If we relied on Maximum Likelihood estimation, our approach would be simply to compute the mean and standard deviation of the dataset, and use this normal distribution as our model.\n",
    "\n",
    "$$\\mu_{MLE} = \\frac{1}{N}\\sum_{i=1}^nx_i$$\n",
    "$$\\sigma_{MLE} = \\frac{1}{N}\\sum_{i=1}^n(x_i - \\mu)^2$$\n",
    "\n",
    "Once we estimate the parameters, we can find out the probability that a sample lies in the range using the following formula\n",
    "$$ p(a < X <= b) = \\int_{a}^b p(X) dX $$\n",
    "\n",
    "However, when the amount of data is low, the MLE estimates are not as reliable. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6bf715a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MLE: mu 149.68 std 11.52\n"
     ]
    }
   ],
   "source": [
    "mle_mu, mle_std = X.mean(), X.std()\n",
    "mle_dist = Normal(mle_mu, mle_std)\n",
    "\n",
    "print(f\"MLE: mu {mle_mu:0.2f} std {mle_std:0.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a291af2",
   "metadata": {},
   "source": [
    "## Bayesian Inference\n",
    "\n",
    "Can we do better than MLE? \n",
    "\n",
    "One potential method to do this is to use Bayesian inference with a good prior. How does one go about selecting a good prior? Well, lets say from another survey, we know that the average and the standard deviation of height of adult female residents in Neighborville, the neighboring town. Additionally, we have no reason to believe that the distribution of heights at Statsville is significantly different.  So we can use this information to \"initialize\" our prior. \n",
    "\n",
    "Lets say the the mean height of adult female resident in Neighborville is 150 cm with a standard deviation of 9 cm.\n",
    "\n",
    "We can use this information as our prior. The prior distribution encodes our beliefs on the parameter values.\n",
    "\n",
    "Given that we are dealing with an unknown mean, and unknown variance, we will model the prior as a Normal Gamma distribution. \n",
    "\n",
    "$$p\\left( \\theta \\middle\\vert X \\right) = p \\left( X \\middle\\vert \\theta \\right) p \\left( \\theta \\right)\\\\\n",
    "p\\left( \\theta \\middle\\vert X \\right) = Normal-Gamma\\left( \\mu_{n}, \\lambda_{n}, \\alpha_{n}, \\beta_{n} \\right) \\\\\n",
    "p \\left( X \\middle\\vert \\theta \\right)  = \\mathbb{N}\\left( \\mu, \\lambda^{ -\\frac{1}{2} } \\right) \\\\\n",
    "p \\left( \\theta \\right) = Normal-Gamma\\left( \\mu_{0}, \\lambda_{0}, \\alpha_{0}, \\beta_{0} \\right)$$\n",
    "\n",
    "We will choose a prior, $p \\left(\\theta \\right)$, such that \n",
    "$$ \\mu_{0} = 150 \\\\\n",
    "   \\lambda_{0} = 100 \\\\\n",
    "   \\alpha_{0} = 100.5 \\\\\n",
    "   \\beta_{0} = 8100 $$\n",
    "   \n",
    "$$p \\left( \\theta \\right) = Normal-Gamma\\left( 150, 100, 100.5 , 8100 \\right)$$\n",
    "\n",
    "\n",
    "We will compute the posterior, $p\\left( \\theta \\middle\\vert X \\right)$,  using Bayesian inference.\n",
    "\n",
    "$$\\mu_{n} = \\frac{ \\left( n \\bar{x} + \\mu_{0} \\lambda_{0} \\right) }{ n + \\lambda_{0} } \\\\\n",
    "\\lambda_{n} = n + \\lambda_{0} \\\\\n",
    "\\alpha_{n} = \\frac{n}{2} + \\alpha_{0} \\\\\n",
    "\\beta_{n} = \\frac{ ns }{ 2 } + \\beta_{ 0 } + \\frac{ n \\lambda_{0} } { 2 \\left( n + \\lambda_{0} \\right) } \\left( \\bar{x} - \\mu_{0} \\right)^{ 2 }$$\n",
    "\n",
    "$$p\\left( \\theta \\middle\\vert X \\right) = Normal-Gamma\\left( \\mu_{n}, \\lambda_{n}, \\alpha_{n}, \\beta_{n} \\right)$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "40eb48f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NormalGamma():\n",
    "    def __init__(self, mu_, lambda_, alpha_, beta_):\n",
    "        self.mu_ = mu_\n",
    "        self.lambda_ = lambda_\n",
    "        self.alpha_ = alpha_\n",
    "        self.beta_ = beta_\n",
    "        \n",
    "    @property\n",
    "    def mean(self):\n",
    "        return self.mu_, self.alpha_/ self.beta_\n",
    "\n",
    "    \n",
    "    @property\n",
    "    def mode(self):\n",
    "        return self.mu_, (self.alpha_-0.5)/ self.beta_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "518babf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def inference_unknown_mean_variance(X, prior_dist):\n",
    "    mu_mle = X.mean()\n",
    "    sigma_mle = X.std()\n",
    "    n = X.shape[0]\n",
    "    # Parameters of the prior\n",
    "    mu_0 = prior_dist.mu_\n",
    "    lambda_0 = prior_dist.lambda_\n",
    "    alpha_0 = prior_dist.alpha_\n",
    "    beta_0 = prior_dist.beta_\n",
    "    \n",
    "    # Parameters of posterior\n",
    "    mu_n = (n * mu_mle + mu_0 * lambda_0) / (lambda_0 + n) \n",
    "    lambda_n = n + lambda_0\n",
    "    alpha_n = n / 2 + alpha_0\n",
    "    beta_n = (n / 2 * sigma_mle ** 2) + beta_0 + (0.5* n * lambda_0  * (mu_mle - mu_0) **2 /(n + lambda_0))  \n",
    "    posterior_dist = NormalGamma(mu_n, lambda_n, alpha_n, beta_n)\n",
    "    \n",
    "    return posterior_dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e4789160",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let us initialize the prior based on our beliefs\n",
    "prior_dist = NormalGamma(150, 100, 10.5, 810)\n",
    "\n",
    "# We compute the posterior distribution\n",
    "posterior_dist = inference_unknown_mean_variance(X, prior_dist)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb2f9930",
   "metadata": {},
   "source": [
    "How do we use the posterior distribution?\n",
    "\n",
    "Note that the posterior distribution is a distribution on the parameters $\\mu$ and $\\lambda$. It is important to note that the posterior and prior are distributions in the parameter space. The likelihood is a distribution on the data space.\n",
    "\n",
    "\n",
    "Once we learn the posterior distribution, one way to use the distribution is to look at the mode of the distribution i.e the parameter values which have the highest probability density. Using these point estimates leads us to Maximum A Posteriori / MAP estimation.\n",
    "\n",
    "As usual, we will obtain the maxima of the posterior probability density function $p\\left( \\mu, \\sigma \\middle\\vert X \\right) = Normal-Gamma\\left(  \\mu, \\sigma ; \\;\\; \\mu_{n}, \\lambda_{n}, \\alpha_{n}, \\beta_{n} \\right) $.\n",
    "\n",
    "This function attains its maxima when\n",
    "\n",
    "$$\\mu = \\mu_{n} \\\\\n",
    "\\lambda = \\frac{ \\alpha_{n} - \\frac{1}{2} } { \\beta_{n} }$$\n",
    "\n",
    "We notice that the MAP estimates for $\\mu$ and $\\sigma$ are better than the MLE estimates. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "138485e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MAP: mu 149.98 std 9.56\n"
     ]
    }
   ],
   "source": [
    "# With the Normal Gamma formulation, the unknown parameters are mu and precision\n",
    "map_mu, map_precision =  posterior_dist.mode\n",
    "\n",
    "# We can compute the standard deviation using precision.\n",
    "map_std = math.sqrt(1 / map_precision)\n",
    "\n",
    "map_dist = Normal(map_mu, map_std)\n",
    "print(f\"MAP: mu {map_mu:0.2f} std {map_std:0.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a59d4d1",
   "metadata": {},
   "source": [
    "How did we arrive at the values of the parameters for the prior distribution? \n",
    "\n",
    "Let us consider the case when we have 0 data points. In this case, posterior will become equal to the prior. If we use the mode of this posterior for our MAP estimate, we see that the mu and std parameters are the same as the $\\mu$ and $\\sigma$ of adult female residents in Neighborville."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "eb3a20da",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prior: mu 150.00 std 9.00\n"
     ]
    }
   ],
   "source": [
    "prior_mu, prior_precision =  prior_dist.mode\n",
    "prior_std = math.sqrt(1 / prior_precision)\n",
    "print(f\"Prior: mu {prior_mu:0.2f} std {prior_std:0.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6d7314d",
   "metadata": {},
   "source": [
    "## Inference\n",
    "\n",
    "Let us say we want to find out the probability that a height between 150 and 155 belongs to an adult female resident. We can now use the  the MAP estimates for $\\mu$ and $\\sigma$ to compute this value. \n",
    "\n",
    "Since our prior was good, we notice that the MAP serves as a better estimator than MLE at low values of n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "50bd53c2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True probability: tensor([0.2449])\n",
      "MAP probability: tensor([0.1995])\n",
      "MLE probability: tensor([0.1669])\n"
     ]
    }
   ],
   "source": [
    "a, b = torch.Tensor([150]), torch.Tensor([155])\n",
    "\n",
    "true_prob = true_dist.cdf(b) - true_dist.cdf(a)\n",
    "print(f'True probability: {true_prob}')\n",
    "\n",
    "map_prob = map_dist.cdf(b) - map_dist.cdf(a)\n",
    "print(f'MAP probability: {map_prob}')\n",
    "\n",
    "mle_prob = mle_dist.cdf(b) - mle_dist.cdf(a)\n",
    "print('MLE probability: {}'.format(mle_prob))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d79bdb99",
   "metadata": {},
   "source": [
    "Let us say we receive more samples, how do we incorporate this information into our model? We can now set the prior to our current posterior and run inference again to obtain the new posterior. This process can be done interatively.\n",
    "\n",
    "$$ p \\left( \\theta \\right)_{n}  = p\\left( \\theta \\middle\\vert X \\right)_{n-1}$$\n",
    "$$ p\\left( \\theta \\middle\\vert X \\right)_{n}=inference\\_unknown\\_mean\\_variance(X_{n}, p \\left( \\theta \\right)_{n})$$\n",
    "\n",
    "We also notice that as the number of data points increases, the MAP starts to converge towards the true values of $\\mu$ and $\\sigma$ respectively"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "787a49fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MAP at batch 0: mu 149.98 std 8.84\n",
      "MAP at batch 5: mu 150.65 std 8.98\n",
      "MAP at batch 10: mu 150.70 std 8.77\n",
      "MAP at batch 15: mu 151.15 std 8.79\n",
      "MAP at batch 19: mu 151.04 std 8.70\n"
     ]
    }
   ],
   "source": [
    "num_batches, batch_size = 20, 10\n",
    "for i in range(num_batches):\n",
    "    X_i = true_dist.sample((batch_size, 1))\n",
    "    prior_i = posterior_dist\n",
    "    posterior_dist = inference_unknown_mean_variance(X_i, prior_i)\n",
    "    map_mu, map_precision =  posterior_dist.mode\n",
    "\n",
    "    # We can compute the standard deviation using precision.\n",
    "    map_std = math.sqrt(1 / map_precision)\n",
    "    map_dist = Normal(map_mu, map_std)\n",
    "    if i % 5 == 0:\n",
    "        print(f\"MAP at batch {i}: mu {map_mu:0.2f} std {map_std:0.2f}\")\n",
    "print(f\"MAP at batch {i}: mu {map_mu:0.2f} std {map_std:0.2f}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}