EM_Algorithm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# EM Algorithm for Clustering and Missing values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## About this notebook: \n",
    "Notebook prepared by **Jesus Perez Colino** Version 0.1, First Released: 21/03/2018, Alpha.  \n",
    "\n",
    "- This work is licensed under a [Creative Commons Attribution-ShareAlike 3.0 Unported License](http://creativecommons.org/licenses/by-sa/3.0/deed.en_US). This work is offered for free, with the hope that it will be useful.\n",
    "\n",
    "\n",
    "- **Summary**: This notebook is just the simplest implementation of the **EM-algorithm** with full description. The objective is not to write the fastest implementation (will come later) but the best possibly explained.\n",
    "\n",
    "\n",
    "- TODO: will be extended soon to vectorized version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------- EM-Algo: Reproducibility conditions for this notebook ---------------\n",
      "Python version:     3.5.5 |Anaconda custom (64-bit)| (default, Mar 12 2018, 16:25:05) \n",
      "[GCC 4.2.1 Compatible Clang 4.0.1 (tags/RELEASE_401/final)]\n",
      "Numpy version:      1.14.2\n",
      "IPython version:    6.2.1\n",
      "-------------------------------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "from sys import version \n",
    "import IPython\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import scipy.stats as scs\n",
    "from scipy.stats import multivariate_normal as mvn\n",
    "from scipy.optimize import minimize\n",
    "\n",
    "print (' EM-Algo: Reproducibility conditions for this notebook '.center(85,'-'))\n",
    "print ('Python version:     ' + version)\n",
    "print ('Numpy version:      ' + np.__version__)\n",
    "print ('IPython version:    ' + IPython.__version__)\n",
    "print ('-'*85)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "## Introduction\n",
    "\n",
    "In this notebook, we study the maximum likelihood estimation by the **EM algorithm**, a special case of the **MM algorithm **(M stands for minorize and the second M for maximize). At the heart of every **EM algorithm** is some notion of missing data. Data can be missing in the ordinary sense, such as a failure to record certain observations on certain cases, or also  data can also be missing in a theoretical sense. In any case, we can think of the **E-step** (expectation) of the algorithm as filling in the missing data. This action replaces the loglikelihood of the observed data by a minorizing function, which is then maximized in the **M-step**. Because the surrogate function is usually much simpler than the likelihood, we can often solve the **M-step** analytically. The price we pay for this simplification is *iteration*. Reconstruction of the missing data is bound to be slightly wrong if the parameters do not already equal their maximum likelihood estimates.\n",
    "\n",
    "One of the advantages of the **EM algorithm** is its numerical stability. As an MM algorithm, any **EM algorithm** leads to a steady increase in the likelihood of the observed data. Thus, the **EM algorithm** avoids wildly overshooting or undershooting the maximum of the likelihood along its current direction of search. Besides this desirable feature, the **EM algorithm** handles parameter constraints gracefully. Constraint satisfaction is by definition built into the solution of the M step. In contrast, competing methods of maximization must employ special techniques to cope with parameter constraints. The **EM algorithm** shares some of the negative features of the more general MM algorithm. For example, the **EM algorithm** often converges at an excruciatingly slow rate in a neighborhood of the maximum point.\n",
    "\n",
    "Examples of applications of the **EM algorithm** include:\n",
    "\n",
    "- Filling in missing data from a sample set\n",
    "- Discovering values of latent variables\n",
    "- Estimating parameters of HMMs\n",
    "- Estimating parameters of finite mixtures [models]\n",
    "- Unsupervised learning of clusters\n",
    "- etc...\n",
    "\n",
    "Let us begin with the simplest case..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Maximum likelihood with complete information\n",
    "Consider an experiment with coin $A$ that has a probability $\\theta_A$ of heads, and a coin $B$ that has a probability $\\theta_B$ of tails. \n",
    "\n",
    "We draw $m$ samples as follows - for each sample, pick one of the coins at random, flip it $n$ times, and record the number of heads and tails (that sum to $n$). If we recorded which coin we used for each sample, we have complete information and can estimate $\\theta_A$ and $\\theta_B$ in closed form. \n",
    "\n",
    "To be more specific, let us assume we drew $5$ samples with the number of heads and tails represented as a vector $x$, and the sequence of coins chosen was $A,A,B,A,B$. Then the complete log-likelihood function is\n",
    "\n",
    "$$\\log L(\\theta;\\mathbf{x}) = \\log p(x_1;\\theta_A)+\\log p(x_2;\\theta_A)+ \\log p(x_3;\\theta_B)+\\log p(x_4;\\theta_A)+\\log p(x_5;\\theta_B)$$\n",
    "\n",
    "\n",
    "where $p(x_i;\\theta)$ is the binomial distribtion PMF with $n=m$ and $p=\\theta$. We will use $z_i$ to indicate the label of the $i$-th coin, that is - whether we used coin $A$ or $B$ to gnerate the $i$-th sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's simulate the experiment in Python\n",
    "\n",
    "m = 10        # number of trials\n",
    "theta_A = 0.8  # probability of HEADS in coin A\n",
    "theta_B = 0.3  # probability of HEADS in coin B\n",
    "theta_0 = [theta_A, theta_B]\n",
    "\n",
    "coin_A = scs.bernoulli(theta_A)\n",
    "coin_B = scs.bernoulli(theta_B)\n",
    "\n",
    "experiment = np.fromiter(map(sum, [coin_A.rvs(m), \n",
    "                                   coin_A.rvs(m), \n",
    "                                   coin_B.rvs(m), \n",
    "                                   coin_A.rvs(m), \n",
    "                                   coin_B.rvs(m)]),np.int)\n",
    "\n",
    "coins_sequence = [0, 0, 1, 0, 1] # A, A, B, A, B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1, 1, 0, 1, 0, 0, 1, 1, 0, 1])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coin_A.rvs(m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([8, 9, 4, 9, 3])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 5 experiments with 10 trials each where we note as 1:HEAD and 0:TAIL\n",
    "# where we are using the sequence of coins: A, A, B, A, B\n",
    "\n",
    "experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.8666666666666667, 0.35)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "maxlikelihood_A = np.sum(experiment[[0,1,3]])/(3.0*m)\n",
    "maxlikelihood_B = np.sum(experiment[[2,4]])/(2.0*m)\n",
    "maxlikelihood_A, maxlikelihood_B\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def neg_loglikelihood(thetas, n, experiment, coins_sequence):\n",
    "    return -np.sum([scs.binom(n, thetas[j]).logpmf(i) for (i, j) \n",
    "                    in zip(experiment, coins_sequence)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "     fun: 6.182734848389383\n",
       "     jac: array([ 9.20152843e-05, -1.40154555e-04])\n",
       " message: 'Converged (|f_n-f_(n-1)| ~= 0)'\n",
       "    nfev: 55\n",
       "     nit: 10\n",
       "  status: 1\n",
       " success: True\n",
       "       x: array([0.86666702, 0.3499984 ])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bnds = [(0,1), (0,1)]\n",
    "minimize(neg_loglikelihood, \n",
    "         [0.5, 0.5], \n",
    "         args=(m, experiment, coins_sequence),\n",
    "         bounds=bnds, \n",
    "         method='tnc', \n",
    "         options={'maxiter': 100})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Maximum likelihood with Incomplete information\n",
    "\n",
    "Let assume now that we keep the result of the previous experiment, head or tail, and we store it in the variable $\\mathbf X$, but assume that we did not record the coin we used. Now we have **missing data**, defined as $\\mathbf Z$, and the problem of estimating $\\mathbf \\theta$ is harder to solve. One way to approach the problem is to ask - can we assign weights $z_j$ defined as $z_j = P(X \\in j)$ for $ j = {A, B}$ to each sample according to how likely it is to be generated from coin $A$ or coin $B$?\n",
    "\n",
    "With some *'prior'* knowledge of $z_j$, we can maximize the likelihod to find $\\theta$. Similarly, given $z_j$, we can calculate what $\\theta$ should be. So the basic idea behind **Expectation Maximization (EM)** is simply to start with a guess for $\\theta$, then calculate the log-likelihood function of the sample, then update $\\theta$ using the value that maximizes the previous log-likelihood function, and *repeat till convergence*. \n",
    "\n",
    "In general, any estimation problem where the **EM-algorithm** has an useful role, has a similar framework: given the statistical experiment/model which generates a set $\\mathbf{X}$ of observed data, a set of unobserved data or missing values $\\mathbf{Z}$, and a vector of unknown parameters $\\boldsymbol\\theta$, along with a likelihood function:\n",
    "\n",
    "$$L(\\boldsymbol\\theta; \\mathbf{X}, \\mathbf{Z}) = p(\\mathbf{X}, \\mathbf{Z}|\\boldsymbol\\theta)$$\n",
    "\n",
    "the maximum likelihood estimate (MLE) of the unknown parameters is determined by the marginal likelihood of the observed data:\n",
    "\n",
    "$$L(\\boldsymbol\\theta; \\mathbf{X}) = p(\\mathbf{X}|\\boldsymbol\\theta) = \\int  p(\\mathbf{X},\\mathbf{Z}|\\boldsymbol\\theta) d\\mathbf{Z}$$\n",
    "\n",
    "However, this quantity is often intractable (e.g. if $\\mathbf{Z}$ is a sequence of events, so that the number of values grows exponentially with the sequence length, making the exact calculation of the sum extremely difficult).\n",
    "\n",
    "The **EM algorithm** seeks to find a maximum likelihood estimation of the marginal likelihood by iteratively applying the **E** (expectation) and **M** (maximization) steps. The derivation below shows why the **EM algorithm** using these “alternating” updates actually works: \n",
    "\n",
    "- *First*, given the current estimate of the parameters $\\theta^{(t)}$, where $t$ is a counting variable of the iteration, consider the log-likelihood function as a curve (surface) where the base is $\\theta$. The idea start finding another function ${\\displaystyle Q(\\theta | \\theta^{(t)})=\\mathbb{E}[\\log L(X |\\theta) | Y=y, \\theta^{(t)}]}$, where $y$ is the actual observed data (incomplete) and $\\theta^{(t)}$ is the current estimated value of $\\theta$ (**E-step**)\n",
    "\n",
    "$$\\begin{align}Q(\\theta|\\theta_n)\n",
    "&= \\mathbb{E}_{\\mathbf{Z}|\\mathbf{X},\\theta^{(t)}} [\\log L(\\theta;\\mathbf{x},\\mathbf{Z}) ] \\\\\n",
    "&= \\mathbb{E}_{\\mathbf{Z}|\\mathbf{X},\\theta^{(t)}} [\\log \\prod_{i=1}^{n}L(\\theta;\\mathbf{x}_i,\\mathbf{z}_i) ] \\\\\n",
    "&= \\mathbb{E}_{\\mathbf{Z}|\\mathbf{X},\\theta^{(t)}} [\\sum_{i=1}^n \\log L(\\theta;\\mathbf{x}_i,\\mathbf{z}_i) ] \\\\\n",
    "&= \\sum_{i=1}^n\\mathbb{E}_{\\mathbf{Z}|\\mathbf{X};\\theta^{(t)}} [\\log L(\\theta;\\mathbf{x}_i,\\mathbf{z}_i) ] \\\\\n",
    "&= \\sum_{i=1}^n \\sum_{j\\in\\{A,B\\}} P(Z_i =j | X_i = \\mathbf{x}_i; \\theta^{(t)}) \\log L(\\theta_j;\\mathbf{x}_i,\\mathbf{z}_i) \n",
    "\\end{align}$$\n",
    "\n",
    "- *Second*, find the value of $\\theta$ that maximizes ${\\displaystyle Q(\\theta | \\theta^{(t)})}$   (**M-step**)\n",
    "\n",
    "$$\\theta^{(t+1)} = \\underset{\\boldsymbol\\theta}{\\operatorname{arg\\,max}} \\ Q(\\theta|\\theta^{(t)})$$\n",
    "\n",
    "This yields to a new parameter estimate $\\theta^{(t+1)}$ that is a lower bound of the log-likelihood but touches the log-likelihood function at this new $\\theta$. And repeat these two steps process until convergence occurs - at this point, the maxima of the lower bound and likelihood functions are the same and we have found the maximum log-likelihood. \n",
    "\n",
    "Therefore, in the **E-step**, we identify a function which is a lower bound for the log-likelikelihood, but how do we choose the distribution $Q_i$? We want the $Q$ function to touch the log-likelihood, and know that Jensen’s inequality is an equality only if the function is constant. So $Q_i$ is just the posterior distribution of $z_i$, and this completes the **E-step**.\n",
    "\n",
    "In the **M-step**, we find the value of $\\theta$ that maximizes the $Q$ function, and then we iterate over the $E$ and $M$ steps until convergence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " \n",
      "Iteration: 1\n",
      "theta_A = 0.713,           theta_B = 0.581,           ll = -32.687\n",
      " \n",
      "Iteration: 2\n",
      "theta_A = 0.745,           theta_B = 0.569,           ll = -31.259\n",
      " \n",
      "Iteration: 3\n",
      "theta_A = 0.768,           theta_B = 0.55,           ll = -30.76\n",
      " \n",
      "Iteration: 4\n",
      "theta_A = 0.783,           theta_B = 0.535,           ll = -30.331\n",
      " \n",
      "Iteration: 5\n",
      "theta_A = 0.791,           theta_B = 0.526,           ll = -30.071\n",
      " \n",
      "Iteration: 6\n",
      "theta_A = 0.795,           theta_B = 0.522,           ll = -29.95\n",
      " \n",
      "Iteration: 7\n",
      "theta_A = 0.796,           theta_B = 0.521,           ll = -29.901\n",
      " \n",
      "Iteration: 8\n",
      "theta_A = 0.796,           theta_B = 0.52,           ll = -29.881\n",
      " \n",
      "Iteration: 9\n",
      "theta_A = 0.797,           theta_B = 0.52,           ll = -29.874\n"
     ]
    }
   ],
   "source": [
    "experiment = np.array([(5,5), (9,1), (8,2), (4,6), (7,3)])\n",
    "thetas = np.array([[0.6, 0.4], [0.5, 0.5]])\n",
    "\n",
    "tol = 0.01\n",
    "max_iter = 100\n",
    "\n",
    "ll_old = 0\n",
    "for i in range(max_iter):\n",
    "    zs_A = []\n",
    "    zs_B = []\n",
    "\n",
    "    vs_A = []\n",
    "    vs_B = []\n",
    "\n",
    "    ll_new = 0\n",
    "\n",
    "    # E-step: calculate probability distributions over possible completions\n",
    "    for x in experiment:\n",
    "\n",
    "        # multinomial (binomial) log-likelihood\n",
    "        ll_A = np.sum([x*np.log(thetas[0])])\n",
    "        ll_B = np.sum([x*np.log(thetas[1])])\n",
    "\n",
    "        denom = np.exp(ll_A) + np.exp(ll_B)\n",
    "        z_A = np.exp(ll_A)/denom\n",
    "        z_B = np.exp(ll_B)/denom\n",
    "\n",
    "        zs_A.append(z_A)\n",
    "        zs_B.append(z_B)\n",
    "\n",
    "        # used for calculating theta\n",
    "        vs_A.append(np.dot(z_A, x))\n",
    "        vs_B.append(np.dot(z_B, x))\n",
    "\n",
    "        # update complete log likelihood\n",
    "        ll_new += z_A * ll_A + z_B * ll_B\n",
    "\n",
    "    # M-step: update values for parameters given current distribution\n",
    "    thetas[0] = np.sum(vs_A, 0)/np.sum(vs_A)\n",
    "    thetas[1] = np.sum(vs_B, 0)/np.sum(vs_B)\n",
    "    \n",
    "    # print distribution of z for each x and current parameter estimate\n",
    "    print(' ')\n",
    "    print('Iteration: {}'.format(i+1))\n",
    "    print(\"theta_A = {0}, \\\n",
    "          theta_B = {1}, \\\n",
    "          ll = {2}\".format(round(thetas[0,0],3),\n",
    "                           round(thetas[1,0],3),\n",
    "                           round(ll_new,3)))\n",
    "\n",
    "    if np.abs(ll_new - ll_old) < tol:\n",
    "        break\n",
    "    ll_old = ll_new"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:py35]",
   "language": "python",
   "name": "conda-env-py35-py"
  },
  "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.5.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}