PY580_Notebook16_Expectation_Maximization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Notebook 16: Expectation Maximization in practice\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Goal \n",
    "The goal of this notebook is to gain intuition for Expectation Maximization using a simple example involving coin tosses.\n",
    "\n",
    "## Overview\n",
    "\n",
    "In Section XIV, we introduce Expectation-Maximization (EM) as a practical way to perform maximum likelihood estimation (MLE) even when some of the data is hidden (i.e in the presence of latent or hidden variables). To better understand EM, in this short notebook we'll explore a very simple coin-tossing example adapted from [Do and Batzoglou, Nat. Biotechnol. (2008)](https://www.nature.com/articles/nbt1406). \n",
    "\n",
    "Suppose that we are given two coins A and B with unkown bias $\\theta_A$ and $\\theta_B$, respectively. Our goal is to estimate the bias vector $\\boldsymbol{\\theta}= (\\theta_A, \\theta_B)$ from the outcomes of the following experiment: \n",
    "\n",
    "<blockquote> \n",
    "First choose one coin at random. Then toss the selected coin 10 times independently and record the number of heads observed. Repeat this procedure 5 times.\n",
    "</blockquote>\n",
    "\n",
    "Formally, let $z_i\\in\\{A,B\\}$ be the coin selected in experiment $i$ and $x_i\\in\\{0,1,\\cdots 10\\}$ be the number heads recorded by tossing $z_i$ 10 times. Since we conduct $n=5$ such experiments, we can summarize the outcomes of these 50 tosses by two vectors: $\\boldsymbol{x}=(x_1,x_2\\cdots, x_5)$ and $\\boldsymbol{z}=(z_1,z_2,\\cdots, z_5)$.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1: What if we know everything?\n",
    "\n",
    " * Consider first the case where we have complete knowledge of the experiment, namely, both $\\boldsymbol{x}$ and $\\boldsymbol{z}$ are known. How would you intuitively estimate the biases of the two coins  $\\boldsymbol{\\theta}= (\\theta_A, \\theta_B)$ ?\n",
    " \n",
    " * What's the likelihood of observing the complete outcomes of these experiments? In other words, what is $P(\\boldsymbol{x},\\boldsymbol{z}| n,\\boldsymbol{\\theta} )$? You may assume this is a Bernoulli trial. Namely, every time coin A(B) is tossed, we have, with probability $\\theta_A$($\\theta_B$), that the outcome is heads.\n",
    " \n",
    " * What's the Maximum Likelihood Estimator (MLE)? Is this consistent with your intuition? \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comparing MLE and EM\n",
    "\n",
    "To test your answer, let's do some numerics! We will compare the MLE estimates of biases with an Expectation Maximization procedure where we do not know ${\\bf z}$. The following code computes our best guess for the biases using MLE -- assuming we know the identity of the coin used -- and compares it estimates arrived at using an EM procedure where we have no knowledge about which coin was being tossed (though we know the same coin was tossed 10 times)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "At iteration 0, theta_A = 0.681216,  theta_B = 0.503301\n",
      "At iteration 1, theta_A = 0.732818,  theta_B = 0.459266\n",
      "At iteration 2, theta_A = 0.760207,  theta_B = 0.425543\n",
      "At iteration 3, theta_A = 0.766395,  theta_B = 0.409645\n",
      "At iteration 4, theta_A = 0.766694,  theta_B = 0.402593\n",
      "At iteration 5, theta_A = 0.766301,  theta_B = 0.399292\n",
      "At iteration 6, theta_A = 0.766019,  theta_B = 0.397709\n",
      "At iteration 7, theta_A = 0.765867,  theta_B = 0.396945\n",
      "E-M converges at iteration 8\n",
      "RESULT:\n",
      "E-M: theta_A = 0.765867,  theta_B = 0.396945\n",
      "MLE with complete data: theta_A = 0.766667,  theta_B = 0.350000\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from scipy.special import comb\n",
    "import math\n",
    "\n",
    "\n",
    "def compute_likelihood(obs, n, pheads): # No surprise, it's Binomial!!!\n",
    "\n",
    "    likelihood = comb(n, obs, exact=True)*(pheads**obs)*(1.0-pheads)**(n-obs)\n",
    "\n",
    "    return likelihood\n",
    "\n",
    "# generate experiments\n",
    "num_coin_toss = 10 # each experiment contains num_coin_toss tosses\n",
    "num_exp = 5  # we perform 5 such experiments\n",
    "theta_A_true = 0.8 \n",
    "theta_B_true = 0.4\n",
    "coin_choice = np.zeros(num_exp) # initialize: 0 for A and 1 for B\n",
    "head_counts = np.zeros(num_exp)\n",
    "\n",
    "# MLE \n",
    "MLE_A = 0.0\n",
    "MLE_B = 0.0\n",
    "\n",
    "# generate the outcomes of experiment\n",
    "for i in np.arange(num_exp):\n",
    "    \n",
    "    if np.random.randint(2) == 0: # coin A is selected\n",
    "        head_counts[i] = np.random.binomial(num_coin_toss , theta_A_true, 1) # toss coin A num_coin_toss times\n",
    "        MLE_A = MLE_A +  head_counts[i] # add the number of heads observed to total headcounts \n",
    "    \n",
    "    else: # coin B is selected \n",
    "        head_counts[i] = np.random.binomial(num_coin_toss , theta_B_true, 1) # toss coin B num_coin_toss times\n",
    "        coin_choice[i] = 1  # record the selection of coin B during experiment i \n",
    "        MLE_B = MLE_B +  head_counts[i] # add the number of heads observed to total headcounts \n",
    "    \n",
    "tail_counts = num_coin_toss - head_counts\n",
    "\n",
    "\n",
    "# MLE is merely the proportion of heads for each coin toss\n",
    "MLE_A = MLE_A / ((num_exp - np.count_nonzero(coin_choice))*num_coin_toss)\n",
    "MLE_B = MLE_B / (np.count_nonzero(coin_choice)*num_coin_toss)\n",
    "\n",
    "\n",
    "\n",
    "# initialize the pA(heads) and pB(heads), namely, coin biases\n",
    "pA_heads = np.zeros(100); \n",
    "pB_heads = np.zeros(100); \n",
    "\n",
    "pA_heads[0] = 0.60 # initial guess\n",
    "pB_heads[0] = 0.50 # initial guess\n",
    "\n",
    "# E-M begins!\n",
    "epsilon = 0.001   # error threshold\n",
    "j = 0 # iteration counter\n",
    "improvement = float('inf')\n",
    "\n",
    "while (improvement > epsilon):\n",
    "    \n",
    "    expectation_A = np.zeros((num_exp,2), dtype=float) \n",
    "    expectation_B = np.zeros((num_exp,2), dtype=float)\n",
    "    \n",
    "    for i in np.arange(min(len(head_counts),len(tail_counts))):\n",
    "        \n",
    "        eH = head_counts[i]\n",
    "        eT = tail_counts[i]\n",
    "        \n",
    "        # E step:\n",
    "        lA = compute_likelihood(eH, num_coin_toss, pA_heads[j])\n",
    "        lB = compute_likelihood(eH, num_coin_toss, pB_heads[j])\n",
    "        \n",
    "        weightA = lA / (lA + lB)\n",
    "        weightB = lB / (lA + lB)\n",
    "              \n",
    "        expectation_A[i] = weightA*np.array([eH, eT])\n",
    "        expectation_B[i] = weightB*np.array([eH, eT])\n",
    "\n",
    "  \n",
    "    # M step\n",
    "    theta_A = np.sum(expectation_A, axis = 0)[0] / np.sum(expectation_A) \n",
    "    theta_B = np.sum(expectation_B, axis = 0)[0] / np.sum(expectation_B) \n",
    "\n",
    "    print('At iteration %d, theta_A = %2f,  theta_B = %2f' % (j, theta_A, theta_B))\n",
    "    \n",
    "    pA_heads[j+1] = sum(expectation_A)[0] / sum(sum(expectation_A)); \n",
    "    pB_heads[j+1] = sum(expectation_B)[0] / sum(sum(expectation_B)); \n",
    "\n",
    "    improvement = max( abs(np.array([pA_heads[j+1],pB_heads[j+1]]) - np.array([pA_heads[j],pB_heads[j]]) ))\n",
    "    j = j+1\n",
    "\n",
    "# END of E-M, print the outcome\n",
    "\n",
    "print('E-M converges at iteration %d' %j)\n",
    "print('RESULT:')\n",
    "print('E-M: theta_A = %2f,  theta_B = %2f' % (theta_A, theta_B))\n",
    "print('MLE with complete data: theta_A = %2f,  theta_B = %2f' % (MLE_A, MLE_B))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "### Exercise 2\n",
    "\n",
    " * How fast does EM converge? Is the converged result close to what you'd get from MLE? \n",
    " \n",
    " * Following Exercise 1, what's the objective function we're optimizing in the E-step? Does this function have a *unique global maximum*? \n",
    " \n",
    " * Compare both the results of MLE and EM to the actual bias (i.e. *theta_A_true*  and *theta_B_true* in the snippet above), comment on their performance.\n",
    " \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Final remarks: a few practical tricks\n",
    "\n",
    "From Exercise 2 and Section XIV, we know that the E-M algorithm often approximates the MLE even in the presence of latent (hidden variables). Like with most optimization methods for nonconcave functions, E-M only guarantees convergence to a local maximum of the objective function. For this reason, its performance can be boosted by running the EM procedure starting with multiple initial parameters. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3\n",
    "\n",
    "* Now instead of having a fixed initial guess of coin biases (i.e. *pA_heads[0]* and *pB_heads[0]* in the snippet), draw these values uniformly at random from $[0,1]$ and run the E-M algorithm. Repeat this twenty times and report what you observed. What's the best initial guess that gives the closest estimate to the true parameters?\n",
    "\n",
    "* As we discussed in Section X (LinReg), **Maximum a posteriori (MAP)** estimation differs from MLE in that it employs an augmented objective function which incorporates a prior distribution over the quantities we want to estimate, and the prior distribution can be think of as a regularizer for the objective fuction used in MLE. Here we will explore how to extend E-M to MAP estimation. \n",
    "\n",
    "  (1) First derive the MAP estimate for the one-coin-flipping example, namely,\n",
    "  $$\n",
    "  \\hat{{\\theta}}_{MAP}(\\boldsymbol{x}) = \\arg\\max_{\\theta\\in[0,1]} \\log P(\\boldsymbol{x}|n,{\\theta} ) + \\log P({\\theta}),\n",
    "  $$\n",
    "  where \n",
    "  $$P(\\boldsymbol{x}|n,{\\theta}) = \\prod_{i=1}^{10} \\text{Binomial} (x_i|n,\\theta)$$\n",
    "  \n",
    "  $$P({\\theta})=\\mathcal{N}(\\theta|\\mu, \\sigma)$$\n",
    "  \n",
    "  (2) Based on (1), now modify the E-M snippet above to incorporate this prior distribution into the **M-step**. Comment on the performance. For the prior choice, try $P(\\boldsymbol{\\theta})=\\mathcal{N}(\\theta_A|0.83, 1)\\mathcal{N}(\\theta_B|0.37, 1)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.8 (default, Apr 13 2021, 12:59:45) \n[Clang 10.0.0 ]"
  },
  "vscode": {
   "interpreter": {
    "hash": "b51f1cc126782fe4aec98aa18c208a32006c2e8c51559c6f4b17791aaf94cabd"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}