Last Homework - Part 3 - Expectation–Maximization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "90f1e5fd",
   "metadata": {},
   "source": [
    "# Notebook 16: Expectation Maximization in practice\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cceedf66",
   "metadata": {},
   "source": [
    "## Exercise 1: What if we know everything?\n",
    "<ul>\n",
    "<li>Consider first the case where we have complete knowledge of the experiment, namely, both  $\\textbf{x}$ and  $\\textbf{z}$ are known. How would you intuitively estimate the biases of the two coins $\\boldsymbol{\\theta}=(\\theta_A,\\theta_B)$.?</li>\n",
    "<span style=\"color:blue\">    \n",
    "    $$\n",
    "\\begin{aligned}    \n",
    "\\theta_A &= \\dfrac{\\sum_{i=1}^{n=5} (1-z_i)x_i}{10 \\sum_{i=1}^{n=5} (1-z_i)} \\\\\n",
    "\\theta_B &= \\dfrac{\\sum_{i=1}^{n=5} z_i x_i}{10 \\sum_{i=1}^{n=5} z_i} \\\\\n",
    "\\end{aligned}\n",
    "$$ \n",
    "where $z_i$ is 0 if experiment $i$ used coin $A$ and $1$ if it used coin $B$.\n",
    " </span>\n",
    "<li>What's the likelihood of observing the complete outcomes of these experiments? In other words, what is  $P(\\textbf{x},\\textbf{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.</li>\n",
    "<span style=\"color:blue\">    \n",
    "Let $p_i = (1-z_i)\\theta_A +z_i\\theta_B$. Then\n",
    "$$\n",
    "   P(\\textbf{x},\\textbf{z}|n,\\boldsymbol{\\theta}) = \\dfrac{1}{2^n}\\prod_{i=1}^{n} {10 \\choose x_i} p_i^{x_i}(1-p_i)^{10-x^i} \n",
    "$$\n",
    "</span>\n",
    "\n",
    "<li>What's the Maximum Likelihood Estimator (MLE)? Is this consistent with your intuition?</li>\n",
    "<span style=\"color:blue\"> \n",
    "$$\n",
    "   \\dfrac{\\partial \\log P(\\textbf{x},\\textbf{z}|n,\\boldsymbol{\\theta})}{\\partial \\theta_A} = 0 \n",
    "   \\implies \\sum_{i=1}^n (1-z_i)\\left[ \\dfrac{x_i}{\\theta_A} - \\dfrac{10-x_i}{1-\\theta_A} \\right] = 0\n",
    "   \\implies \\theta_A = \\dfrac{\\sum_{i=1}^{n=5} (1-z_i)x_i}{10 \\sum_{i=1}^{n=5} (1-z_i)} \n",
    "$$    \n",
    "Similarly for $\\theta_B$. We got the same results in the first question.\n",
    "</span>\n",
    "\n",
    "</ul>\t\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "57b2332c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.special import comb\n",
    "import math\n",
    "import time\n",
    "import matplotlib.pyplot as plt\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "id": "4c812823",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True values: (0.8, 0.4)\n",
      "\n",
      "MLE: (0.802, 0.404)\n",
      "Time: 0.0\n",
      "Error: 0.004472135954999584\n",
      "\n",
      "MAP: (0.8020017783902068, 0.40399672566151273)\n",
      "Time: 0.0010271072387695312\n",
      "Error: 0.004470003661843356\n",
      "\n",
      "E-M (without prior): (0.802000000018688, 0.4040000001135043)\n",
      "...converges at iteration 3\n",
      "Time: 0.004998683929443359\n",
      "Error: 0.004472136064878382\n",
      "\n",
      "E-M (with prior): (0.8020017784084726, 0.40399672577363166)\n",
      "...converges at iteration 4\n",
      "Time: 0.00693821907043457\n",
      "Error: 0.0044700037702711725\n"
     ]
    }
   ],
   "source": [
    "def compute_likelihood(obs, n, pheads):\n",
    "    likelihood = comb(n, obs, exact=True)*(pheads**obs)*(1.0-pheads)**(n-obs) \n",
    "    #likelihood = (pheads**obs)*(1.0-pheads)**(n-obs) \n",
    "    return likelihood\n",
    "\n",
    "# Generate experiments\n",
    "num_coin_toss = 100 # each experiment contains num_coin_toss tosses\n",
    "num_exp = 50  # we perform 5 such experiments\n",
    "theta_A_true = 0.8 \n",
    "theta_B_true = 0.4\n",
    "\n",
    "# MLE \n",
    "MLE_A = 0.0\n",
    "MLE_B = 0.0\n",
    "\n",
    "print(\"True values:\",(theta_A_true,theta_B_true))\n",
    "\n",
    "# Generate the outcomes of experiment\n",
    "coin_choice = np.zeros(num_exp) # initialize: 0 for A and 1 for B\n",
    "head_counts = np.zeros(num_exp)\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",
    "tail_counts = num_coin_toss - head_counts\n",
    "\n",
    "########\n",
    "def MLE(zv=None):\n",
    "    if zv is None:\n",
    "        Z = np.sum(coin_choice)\n",
    "        X = np.sum(head_counts)\n",
    "        XZ = np.sum(coin_choice*head_counts)\n",
    "        return (X-XZ)/(num_coin_toss*(num_exp-Z)), XZ/(num_coin_toss*Z)\n",
    "    else:\n",
    "        Z = np.sum(coin_choice)\n",
    "        X = np.sum(head_counts)\n",
    "        XZ = np.sum(coin_choice*head_counts)\n",
    "        return np.sum(zv[0]*head_counts)/(num_coin_toss*np.sum(zv[0])), np.sum(zv[1]*head_counts)/(num_coin_toss*np.sum(zv[1]))\n",
    "def solve_cubic(a,b,c,d):\n",
    "    D0 = b**2 - 3*a*c\n",
    "    D1 = 2*b**3 - 9*a*b*c + 27*a**2*d\n",
    "    C = ( ( D1 + np.sqrt( D1**2 - 4*D0**3 +0j ) )/2 )**(1/3)\n",
    "    xi = (-1+1j*np.sqrt(3))/2\n",
    "    def root(i):\n",
    "        return -( b + xi**i*C + D0/(xi**i*C) )/(3*a)\n",
    "    roots = [root(0), root(1), root(2)]\n",
    "    roots_ = [np.real(i) for i in roots if np.abs(np.imag(i))<1e-6 and np.real(i)<=1 and np.real(i)>=0]\n",
    "    return max(roots_)\n",
    "def MAP(zv=None, uA=0.83, sA=1, uB=0.37, sB=1):\n",
    "    if zv is None:\n",
    "        nA = np.sum(1-coin_choice)\n",
    "        nB = np.sum(coin_choice)\n",
    "\n",
    "        bA = -(1+uA)\n",
    "        cA = uA - num_coin_toss*nA*sA**2\n",
    "        dA = sA**2 * np.sum((1-coin_choice)*head_counts)\n",
    "\n",
    "        bB = -(1+uB)\n",
    "        cB = uB - num_coin_toss*nB*sB**2\n",
    "        dB = sB**2 * np.sum(coin_choice*head_counts)\n",
    "        return solve_cubic(1,bA,cA,dA), solve_cubic(1,bB,cB,dB)\n",
    "    else:\n",
    "        nA = np.sum(zv[0])\n",
    "        nB = np.sum(zv[1])\n",
    "\n",
    "        bA = -(1+uA)\n",
    "        cA = uA - num_coin_toss*nA*sA**2\n",
    "        dA = sA**2 * np.sum(zv[0]*head_counts)\n",
    "\n",
    "        bB = -(1+uB)\n",
    "        cB = uB - num_coin_toss*nB*sB**2\n",
    "        dB = sB**2 * np.sum(zv[1]*head_counts)\n",
    "        return solve_cubic(1,bA,cA,dA), solve_cubic(1,bB,cB,dB)    \n",
    "########\n",
    "\n",
    "start = time.time()\n",
    "res = MLE()\n",
    "print(\"\\nMLE:\", res)\n",
    "end = time.time()\n",
    "print(\"Time:\", end - start)\n",
    "print(\"Error:\", np.sqrt( (res[0]-theta_A_true)**2 + (res[1]-theta_B_true)**2 ))\n",
    "\n",
    "start = time.time()\n",
    "res = MAP()\n",
    "print(\"\\nMAP:\", res)\n",
    "end = time.time()\n",
    "print(\"Time:\", end - start)\n",
    "print(\"Error:\", np.sqrt( (res[0]-theta_A_true)**2 + (res[1]-theta_B_true)**2 ))\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",
    "def EM(pA_initial, pB_initial, method=MLE):\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] = pA_initial # 0.60 # initial guess\n",
    "    pB_heads[0] = pB_initial # 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",
    "        weightsA = np.zeros(num_exp, dtype=float)\n",
    "        weightsB = np.zeros(num_exp, 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",
    "            weightsA[i] = weightA\n",
    "            weightsB[i] = weightB\n",
    "\n",
    "            expectation_A[i] = weightA*np.array([eH, eT])\n",
    "            expectation_B[i] = weightB*np.array([eH, eT])\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",
    "        theta_A, theta_B = method(zv=[ weightsA, weightsB ]) # MLE\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",
    "    if theta_A>theta_B: return theta_A, theta_B, j\n",
    "    else: return theta_B, theta_A, j\n",
    "\n",
    "pA0, pB0 = np.random.random(2)\n",
    "start = time.time()\n",
    "res = EM(pA0, pB0, MLE)\n",
    "end = time.time()\n",
    "print('\\nE-M (without prior):', res[0:2])\n",
    "print('...converges at iteration %d' %res[-1])\n",
    "print(\"Time:\", end - start)\n",
    "print(\"Error:\", np.sqrt( (res[0]-theta_A_true)**2 + (res[1]-theta_B_true)**2 ))\n",
    "\n",
    "pA0, pB0 = random.gauss(0.83, 1), random.gauss(0.37, 1)\n",
    "start = time.time()\n",
    "res = EM(pA0, pB0, MAP)\n",
    "end = time.time()\n",
    "print('\\nE-M (with prior):', res[0:2])\n",
    "print('...converges at iteration %d' %res[-1])\n",
    "print(\"Time:\", end - start)\n",
    "print(\"Error:\", np.sqrt( (res[0]-theta_A_true)**2 + (res[1]-theta_B_true)**2 ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "1423860e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best time: 0.0\n",
      "Optimal initial guesses for EM:\n",
      "[0.82131187 0.38235789]  --->  0.7927018961429952 0.3361422556755914\n",
      "Min error: 0.06427343020338316\n"
     ]
    }
   ],
   "source": [
    "TA = []; TB = []; ERROR=[]; TIME = []; IG = []\n",
    "start = time.time()    \n",
    "for i in range(20):\n",
    "    IG.append( np.random.random(2) )\n",
    "    start = time.time()   \n",
    "    _theta_A, _theta_B, _  = EM(*IG[-1])\n",
    "    end = time.time()\n",
    "    TA.append(_theta_A); TB.append(_theta_B); TIME.append(end - start)\n",
    "    ERROR.append( np.sqrt( (_theta_A-theta_A_true)**2 + (_theta_B-theta_B_true)**2 ) )\n",
    "TA = np.array(TA); TB = np.array(TB)\n",
    "\n",
    "print(\"Best time:\", min(TIME))\n",
    "print(\"Optimal initial guesses for EM:\")\n",
    "for j in [i for i in range(len(ERROR)) if ERROR[i] == min(ERROR)]: print(IG[j], \" ---> \", TA[j], TB[j])\n",
    "print(\"Min error:\", min(ERROR))    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e22b755",
   "metadata": {},
   "source": [
    "## Exercise 2\n",
    "<ul>\n",
    "<li>How fast does EM converge? Is the converged result close to what you'd get from MLE?</li>\n",
    "<span style=\"color:blue\">    \n",
    "The performance of the EM algorithm is sensitive to the initial conditions. For instance, when the initial conditions for $\\theta_A$ and $\\theta_B$ are approximately equal, the algorithm seems to require fewer iterations, but it also produces the worst results. In general, MLE is faster and produces almost the same results as the EM algorithm. In this particular case, the execution time of the EM algorithm can be reduced by removing the binomial coefficients, which are simplified when calculating the weights and, therefore, do not affect the results.\n",
    "</span>\n",
    "<li>Following Exercise 1, what's the objective function we're optimizing in the E-step? Does this function have a unique global maximum?</li>\n",
    "<span style=\"color:blue\">       \n",
    "We want to maximize the log-likelihood. In the first section, we found the values of $\\theta_A$ and $\\theta_B$ that maximize this likelihood. The results of the EM algorithm will be a bit different, since the indicators $1-z_i$ and $z_i$ are replaced by  weights $w^{(A)}_i$ and $w^{(B)}_i$. Furthermore, depending on the initial conditions, the EM algorithm can converge to values close to $(\\theta_A,\\theta_B)$ or to values close to $(\\theta_B,\\theta_A)$.\n",
    "</span>\n",
    "    <li>Compare both the results of MLE and EM to the actual bias (i.e. <code>theta_A_true</code> and <code>theta_B_true</code> in the snippet above), comment on their performance.</li>\n",
    "<span style=\"color:blue\">    \n",
    "The results are sensitive to the initial conditions. In general, MLE and the EM algorithm produce similar results.\n",
    "</span>\n",
    "    \n",
    "</ul>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99f04a64",
   "metadata": {},
   "source": [
    "## Exercise 3\n",
    "<ul>\n",
    "\n",
    "<li> Now instead of having a fixed initial guess of coin biases (i.e. <code>pA_heads[0]</code> and <code>pB_heads[0]</code> 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? </li>\n",
    "<span style=\"color:blue\">    \n",
    "The worst results are obtained when the initial conditions for $\\theta_A$ and $\\theta_B$ are approximately equal.\n",
    " </span>\n",
    "\n",
    "<li> 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. </li>\n",
    "<ol>\n",
    "<li> First derive the MAP estimate for the one-coin-flipping example, namely,\n",
    "$$  \\hat{\\theta}_\\mathrm{MAP}(\\textbf{x})=\\arg\\max_{\\theta \\in[0,1]}\\log P(\\textbf{x}|n,\\theta)+\\log P(\\theta), $$\n",
    " where\n",
    "$$ P(\\textbf{x}|n,\\theta)= \\prod_i^{10} \\mathrm{Binomial}(x_i|n,\\theta) \\\\\n",
    "P(\\theta) = \\mathcal{N}(\\theta|\\mu,\\sigma)\n",
    "$$  \n",
    "<span style=\"color:blue\"> \n",
    "$$\n",
    "  \\dfrac{\\partial \\left[ \\log P(\\textbf{x},\\textbf{z}|n,\\boldsymbol{\\theta}) + \\log P(\\theta) \\right]}{\\partial \\theta_A} = 0 \n",
    "   \\implies \\sum_{i=1}^n (1-z_i)\\left[ \\dfrac{x_i}{\\theta_A} - \\dfrac{10-x_i}{1-\\theta_A} \\right] -\\dfrac{\\theta_A-\\mu_A}{\\sigma_A^2} = 0 \\\\ \n",
    "  \\implies \\theta_A^3 - (1+\\mu_A) \\theta_A^2 + \\left[ \\mu_A  -\\sigma_A^2 10 \\sum_{i=1}^n (1-z_i)\\right]  \\theta_A+ \\sigma_A^2  \\sum_{i=1}^n (1-z_i)x_i  = 0\n",
    "$$ \n",
    "This is a cubic equation with exact solutions that are provided in the function <code>solve_cubic</code> defined above.    \n",
    "</span>  \n",
    "</li> \n",
    "<li> 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(\\theta)=\\mathcal{N}(\\theta_A|0.83,1)\\mathcal{N}(\\theta_B|0.37,1)$. </li>\n",
    "<span style=\"color:blue\"> \n",
    "The performance depends on the initial conditions.In general, it seems that using the prior distribution slightly improves the results.\n",
    "</span>\n",
    "</ol>\t \n",
    "</ul>    "
   ]
  }
 ],
 "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}