F0E-03-EMalgorithm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FLIP (00): Data Science \n",
    "**(Module 04: Probability Theory)**\n",
    "\n",
    "---\n",
    "- Materials in this module include resources collected from various open-source online repositories.\n",
    "- You are free to use,but NOT allowed to change and distribute this package.\n",
    "\n",
    "Prepared by and for \n",
    "**Student Members** |\n",
    "2006-2018 [TULIP Lab](http://www.tulip.org.au), Australia\n",
    "\n",
    "---\n",
    "## Session 1 Probability theory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import glob\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "%matplotlib inline\n",
    "plt.style.use('ggplot')\n",
    "np.random.seed(1234)\n",
    "\n",
    "np.set_printoptions(formatter={'all':lambda x: '%.3f' % x})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from IPython.display import Image\n",
    "from numpy.core.umath_tests import matrix_multiply as mm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from scipy.optimize import minimize\n",
    "from scipy.stats import bernoulli, binom"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Outline\n",
    "----\n",
    "\n",
    "- Review of Jensen's inequality\n",
    "- Concavity of log function\n",
    "- Example of coin tossing with missing informaiton to provide context\n",
    "- Derivation of EM equations\n",
    "- Illustration of EM convergence\n",
    "- Derivation of update equations of coin tossing example\n",
    "- Code for coin tossing example\n",
    "- Derivation of update equatiosn for mixture of Gaussians\n",
    "- Code for mixture of Gaussians\n",
    "- More examples to show wide applicability of EM and exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Jensen's inequality\n",
    "----\n",
    "\n",
    "For a convex function $f$, $E[f(x) \\geq f(E[x])$. Flip the signe for a concave function. \n",
    "\n",
    "A function $f(x)$ is convex if $f''(x) \\geq 0$ everywhere in its domain. For example, if $f(x) = \\log x$, $f''(x) = -1/x^2$, so the log function is concave for $x \\in (0, \\infty]$. A visual illustration of Jensen's inequality is shown below.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "Image(filename='figs/jensen.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When is Jensen's inequality an equality? From the diagram, we can see that this only happens if the function $f(x)$ is a constant! We will make use of this fact later on in the lecture."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Maximum likelihood with complete information\n",
    "----\n",
    "\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. We draw $m$ samples as follows - for each sample, pick one of the coins, flip it 10 times, and record the number of heads and tails. 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. To be very explicit, suppose we drew 5 samples with the $m$ values represented as a vector $x$, and the sequence of coins chosen was $A, A, B, A, B$. Then the complete log likelihood is \n",
    "\n",
    "$$\n",
    "\\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^\\text{th}$ coin, that is - whether we used coin $A$ or $B$ to gnerate the $i^\\text{th}$ sample."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solving for complete likelihood using minimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def neg_loglik(thetas, n, xs, zs):\n",
    "    return -np.sum([binom(n, thetas[z]).logpmf(x) for (x, z) in zip(xs, zs)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "m = 10\n",
    "theta_A = 0.8\n",
    "theta_B = 0.3\n",
    "theta_0 = [theta_A, theta_B]\n",
    "\n",
    "coin_A = bernoulli(theta_A)\n",
    "coin_B = bernoulli(theta_B)\n",
    "\n",
    "xs = map(sum, [coin_A.rvs(m), coin_A.rvs(m), coin_B.rvs(m), coin_A.rvs(m), coin_B.rvs(m)])\n",
    "zs = [0, 0, 1, 0, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "bnds = [(0,1), (0,1)]\n",
    "minimize(neg_loglik, [0.5, 0.5], args=(m, xs, zs), bounds=bnds, method='tnc', options={'maxiter': 100})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Incomplete information\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, if we did not record the coin we used, we have *missing* data and the problem is harder to solve. One way to approach the problem is to ask - can we assign weights $w_i$ to each sample according to how likely it is to be generated from coin $A$ or coin $B$? Intuitively, it seems that the weights should be the posterior distribution of $z_i$, that is\n",
    "$$\n",
    "w_i = p(z_i | x_i; \\theta)\n",
    "$$\n",
    "If we know $z_i$, then we can estiamte $\\theta$ since we have the complete likelihood as above. So the basic idea behind Expectation Maximization (EM) is simply to start with a guess for $\\theta$, then calculate $z$, then update $\\theta$ using this new value for $z$, and repeat till convergence. The derivation below shows why the EM algorithm using this \"alternating\" updates actually works.\n",
    "\n",
    "A verbal outline of the derivtion - first consider the log likelihood function as a curve (surface) where the base is $\\theta$. Find another function $Q$ of $\\theta$ that is a lower bound of the log-likelihood but touches the log likelihodd function at some $\\theta$ (E-step). Next find the value of $\\theta$ that maximizes this function (M-step). Now find yet antoher function of $\\theta$ that is a lower bound of the log-likelihood but touches the log likelihodd function at this new $\\theta$. Now repeat until convergence - at this point, the maxima of the lower bound and likelihood functions are the same and we have found the maximum log likelihood. See illustratoin below. \n",
    "\n",
    "The only remaining step is how to find the functions that are lower bounds of the log likelihood. This will require a little math using Jensen's inequality, and is shown in the next section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Image from http://www.nature.com/nbt/journal/v26/n8/extref/nbt1406-S1.pdf\n",
    "Image(filename='figs/em.png', width=800)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Derivation\n",
    "\n",
    "In the E-step, we identify a function which is a lower bound for the log-likelikelihood\n",
    "\n",
    "\\begin{align}\n",
    "ll &= \\sum_i{\\log p(x_i; \\theta)} && \\text{definition of log likelihood} \\\\\n",
    "&= \\sum_i \\log \\sum_{z_i}{p(x_i, z_i; \\theta)} && \\text{augment with latent variables $z$} \\\\\n",
    "&= \\sum_i \\log \\sum_{z_i} Q(z_i) \\frac{p(x_i, z_i; \\theta)}{Q_i(z_i)} && \\text{$Q_i$ is a distribution for $z_i$} \\\\\n",
    "&= \\sum_i \\log E_{z_i}[\\frac{p(x_i, z_i; \\theta)}{Q_i(z_i)}] && \\text{taking expectations - hence the E in EM} \\\\\n",
    "&\\geq \\sum E_{z_i}[\\log \\frac{p(x_i, z_i; \\theta)}{Q_i(z_i)}] && \\text{Using Jensen's rule for $\\log$ which is concave} \\\\\n",
    "&\\geq \\sum_i \\sum_{z_i} Q_i(z_i) \\log \\frac{p(x_i, z_i; \\theta)}{Q_i(z_i)} && \\text{Q function}\n",
    "\\end{align}\n",
    "\n",
    "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\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{p(x_i, z_i; \\theta)}{Q_i(z_i)} =& c \\\\\n",
    "\\implies Q_i(z_i) &\\propto p(x_i, z_i; \\theta)\\\\\n",
    "\\implies Q_i(z_i) &= \\frac{p(x_i, z_i; \\theta) }{\\sum_{z_i}{p(x_i, z_i; \\theta)}} &&\\text{Since $Q$ is a distribution and sums to 1} \\\\\n",
    "\\implies Q_i(z_i) &= \\frac{p(x_i, z_i; \\theta) }{{p(x_i, \\theta)}} && \\text{marginalizing $z_i$}\\\\\n",
    "\\implies Q_i(z_i) &= p(z_i | x_i; \\theta) && \\text{by definition}\n",
    "\\end{align}\n",
    "\n",
    "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. \n",
    "\n",
    "So we see that EM is an algorihtm for maximum likelikhood optimization when there is missing inforrmaiton - or when it is useful to add latent augmented variables to simplify maximum likelihood calculatoins."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Coin toss example from [What is the expectation maximization algorithm?](http://www.nature.com/nbt/journal/v26/n8/full/nbt1406.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the E-step, we have\n",
    "\n",
    "\\begin{align}\n",
    "w_j &= Q_i(z_i = j) \\\\\n",
    "&= p(z_i = j | x_i; \\theta) \\\\\n",
    "&= \\frac{p(x_i | z_i = j; \\theta) p(z_i = j; \\phi)}  {\\sum_{l=1}^k{p(x_i | z_i = l; \\theta) p(z_i = l; \\phi)}}  && \\text{Baye's rule} \\\\\n",
    "&= \\frac{\\theta_j^h(1-\\theta_j)^{n-h} \\phi_j}{\\sum_{l=1}^k \\theta_l^h(1-\\theta_l)^{n-h} \\phi_l} \\\\\n",
    "&= \\frac{\\theta_j^h(1-\\theta_j)^{n-h} }{\\sum_{l=1}^k \\theta_l^h(1-\\theta_l)^{n-h} } && \\text{assume $\\phi$ is fixed for simplicity}\n",
    "\\end{align}\n",
    "\n",
    "For the M-step, we need to find the value of $\\theta$ that maximises\n",
    "\n",
    "\\begin{align}\n",
    "& \\sum_i \\sum_{z_i} Q_i(z_i) \\log \\frac{p(x_i, z_i; \\theta)}{Q_i(z_i)} \\\\\n",
    "&= \\sum_{i=1}^m \\sum_{j=1}^k w_j \\log \\frac{p(x_i | z_i=j; \\theta) \\, p(z_i = j; \\phi)}{w_j} \\\\\n",
    "&= \\sum_{i=1}^m \\sum_{j=1}^k w_j \\log \\frac{\\theta_j^h(1-\\theta_j)^{n-h} \\phi_j}{w_j} \\\\\n",
    "&= \\sum_{i=1}^m \\sum_{j=1}^k w_j \\left( h \\log \\theta_j + (n-h) \\log (1-\\theta_j) + \\log \\phi_j - \\log w_j \\right)\n",
    "\\end{align}\n",
    "\n",
    "We can differentiate and solve for $\\theta_s$ where the derivative vanishes\n",
    "\n",
    "\\begin{align}\n",
    "\\sum_{i=1}^m w_s \\left( \\frac{h}{\\theta_s} - \\frac{n-h}{1-\\theta_s} \\right) &= 0  \\\\\n",
    "\\implies \\theta_s &= \\frac {\\sum_{i=1}^m w_s h}{\\sum_{i=1}^m w_s n}\n",
    "\\end{align}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementation in code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def mn_ll(y, theta, axis=None):\n",
    "    \"\"\"Log likelihood for multinomial distribution (ignoring constant).\"\"\"\n",
    "    return np.sum(y * np.log(theta), axis=axis)\n",
    "    # return np.sum(y * np.log(theta))\n",
    "\n",
    "def ll(y, theta, p):\n",
    "    \"\"\"Complete log likelihood for mixture.\"\"\"\n",
    "    return np.sum(p * mm_ll(y, theta))\n",
    "\n",
    "def normalize(xs, axis=1):\n",
    "    \"\"\"Return normalized marirx so that sum of row or column (default) entries = 1.\"\"\"\n",
    "    if axis==0:\n",
    "        return xs/xs.sum(0)\n",
    "    else:\n",
    "        return xs/xs.sum(1)[:, None]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "xs = np.array([(5,5), (9,1), (8,2), (4,6), (7,3)])\n",
    "# theta_A, theta_B = np.array([0.6, 0.4]), np.array([0.5, 0.5])\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",
    "    exp_A = []\n",
    "    exp_B = []\n",
    "    ll_new = 0\n",
    "\n",
    "    # E-step: calculate probability distributions over possible completions\n",
    "    for x in xs:\n",
    "\n",
    "        ll_A = mn_ll(x, thetas[0])\n",
    "        ll_B = mn_ll(x, thetas[1])\n",
    "      \n",
    "        denom = np.exp(ll_A) + np.exp(ll_B)       \n",
    "        w_A = np.exp(ll_A)/denom\n",
    "        w_B = np.exp(ll_B)/denom\n",
    "\n",
    "        exp_A.append(np.dot(w_A, x))\n",
    "        exp_B.append(np.dot(w_B, x))\n",
    "        \n",
    "        # update complete log likelihood\n",
    "        ll_new += w_A * ll_A + w_B * ll_B\n",
    "    \n",
    "    # M-step: update values for parameters given current distribution\n",
    "    thetas[0] = np.sum(exp_A, 0)/np.sum(exp_A)\n",
    "    thetas[1] = np.sum(exp_B, 0)/np.sum(exp_B)\n",
    "    # print distribution of z for each x and current parameter estimate\n",
    "    print \"Iteration: %d\" % (i+1)\n",
    "    print \"theta_A = %.2f, theta_B = %.2f, ll = %.2f\" % (thetas[0,0], thetas[1,0], ll_new)\n",
    "\n",
    "    if np.abs(ll_new - ll_old) < tol:\n",
    "        break\n",
    "    ll_old = ll_new"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### After some cleaning up.\n",
    "\n",
    "xs = observed data (sampled from multinomial distribution)\n",
    "wts = conditional distribution of missing data given observed data \n",
    "thetas = parameters to estimate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from scipy.stats import binom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "binom(n=10, p=0.3).logpmf([4, 3, 7, 9, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "xs = 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "xs[0].sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from scipy.stats import binom\n",
    "from numpy.core.umath_tests import inner1d\n",
    "    \n",
    "def _calc_weights(thetas, xs, n):\n",
    "    s = np.array([binom(n=n, p=theta[0]).logpmf(xs[:, 0]) \n",
    "                        for i, theta in enumerate(thetas)])\n",
    "    es = np.exp(s)\n",
    "    wts = es/es.sum(1)[:, None]\n",
    "    ll = inner1d(wts, s).sum()\n",
    "    return wts, ll    \n",
    "    \n",
    "def calc_weights(thetas, xs, n):\n",
    "    s = np.array([binom(n=n, p=theta).logpmf(xs) \n",
    "                        for i, theta in enumerate(thetas)])\n",
    "    es = np.exp(s)\n",
    "    wts = es/es.sum(1)[:, None]\n",
    "    ll = inner1d(wts, s).sum()\n",
    "    return wts, ll\n",
    "    \n",
    "def _calc_thetas(ws, xs):\n",
    "    thetas = np.dot(ws, xs)\n",
    "    # print \"ws\", ws\n",
    "    # print \"xs\", xs\n",
    "    # print \"thetas\", thetas\n",
    "    return thetas/thetas.sum(1)[:, None]\n",
    "\n",
    "def calc_thetas(ws, xs):\n",
    "    thetas = np.dot(ws, xs)\n",
    "    # print \"ws\", ws\n",
    "    # print \"xs\", xs\n",
    "    # print \"thetas\", thetas\n",
    "    return thetas/thetas.sum()\n",
    "\n",
    "def _em(thetas, xs, n):\n",
    "    for i in range(10):\n",
    "        ws, ll = _calc_weights(thetas, xs, n)\n",
    "        thetas = _calc_thetas(ws, xs)\n",
    "        print ll, thetas[:,0]\n",
    "\n",
    "def em(thetas, xs, n):\n",
    "    for i in range(10):\n",
    "        ws, ll = calc_weights(thetas, xs, n)\n",
    "        thetas = calc_thetas(ws, xs)\n",
    "        print ll, thetas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "np.random.seed(1234)\n",
    "\n",
    "n = 1000\n",
    "p0 = 0.2\n",
    "p1 = 0.8\n",
    "xs = np.concatenate([np.random.binomial(n, p0, n/2), np.random.binomial(n, p1, n/2)])\n",
    "xs_ = np.column_stack((xs, n-xs))\n",
    "# np.random.shuffle(xs)\n",
    "print xs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "p = np.random.random(2)\n",
    "p = np.array([0.4, 0.6])\n",
    "thetas = p\n",
    "thetas_ = np.column_stack([p, 1-p])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "em(thetas, xs, n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "thetas = [(0.4,0.6), (0.6, 0.4)]\n",
    "em(thetas, xs, n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "thetas = np.array([[0.5, 0.5], [0.5, 0.5]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "em(thetas, xs, n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "thetas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "calc_weights(thetas, xs, n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def em(loglik, Qfunc, theta0, max_iter=100, tol=1e-3):\n",
    "    ll_old = 0\n",
    "    for i in range(max_iter):\n",
    "        ll_new = 0\n",
    "\n",
    "        # E-step\n",
    "        lls = np.array([mn_ll(x, thetas, axis=1) for x in xs])\n",
    "        ws = normalize(np.exp(lls))\n",
    "\n",
    "        # M-step\n",
    "        thetas = normalize(np.dot(ws.T, xs))\n",
    "\n",
    "        # update complete log likelihoood \n",
    "        ll_new = np.sum(np.dot(w, ll) for (w, ll) in zip(ws, lls))\n",
    "\n",
    "        # print distribution of z for each x and current parameter estimate\n",
    "        print \"Iteration: %d\" % (i+1)\n",
    "        print \"theta_A = %.2f, theta_B = %.2f, ll = %.2f\" % (thetas[0,0], thetas[1,0], ll_new)\n",
    "        print \"Weights\\n\", \\\n",
    "            pd.DataFrame(normalize(ws), columns=['A', 'B'], index=range(1, 1+len(xs))), \"\\n\"\n",
    "\n",
    "        if np.abs(ll_new - ll_old) < tol:\n",
    "            break\n",
    "        ll_old = ll_new"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "xs = np.array([(5,5), (9,1), (8,2), (4,6), (7,3)])\n",
    "# thetas = normalize(np.random.random((2,2)))\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",
    "    exp_A = []\n",
    "    exp_B = []\n",
    "    ll_new = 0\n",
    "\n",
    "    # E-step\n",
    "    lls = np.array([mn_ll(x, thetas, axis=1) for x in xs])\n",
    "    ws = normalize(np.exp(lls))\n",
    "       \n",
    "    print \">>\", ws.shape\n",
    "    print \">>\", xs.shape\n",
    "        \n",
    "    # M-step\n",
    "    thetas = normalize(np.dot(ws.T, xs))\n",
    "    \n",
    "    # update complete log likelihoood \n",
    "    ll_new = np.sum(np.dot(w, ll) for (w, ll) in zip(ws, lls))\n",
    "    \n",
    "    # print distribution of z for each x and current parameter estimate\n",
    "    print \"Iteration: %d\" % (i+1)\n",
    "    print \"theta_A = %.2f, theta_B = %.2f, ll = %.2f\" % (thetas[0,0], thetas[1,0], ll_new)\n",
    "    print \"Weights\\n\", \\\n",
    "        pd.DataFrame(normalize(ws), columns=['A', 'B'], index=range(1, 1+len(xs))), \"\\n\"\n",
    " \n",
    "    if np.abs(ll_new - ll_old) < tol:\n",
    "        break\n",
    "    ll_old = ll_new"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Gaussian mixture models\n",
    "----\n",
    "\n",
    "A mixture of $k$ Gaussians has the following PDF\n",
    "\n",
    "\\begin{align}\n",
    "p(x) = \\sum_{j=1}^k \\alpha_j \\phi(x; \\mu_j, \\Sigma_j)\n",
    "\\end{align}\n",
    "\n",
    "where $\\alpha_j$ is the weight of the $j^\\text{th}$ Gaussain component and \n",
    "\n",
    "\\begin{align}\n",
    "\\phi(x; \\mu, \\Sigma) = \\frac{1}{(2 \\pi)^{d/2}|\\Sigma|^{1/2}} \\exp \\left( -\\frac{1}{2}(x-\\mu)^T\\Sigma^{-1}(x-\\mu) \\right)\n",
    "\\end{align}\n",
    "\n",
    "Suppose we observe $y_1, y2, \\ldots, y_n$ as a sample from a mixture of Gaussians. The log-likeihood is then\n",
    "\n",
    "\\begin{align}\n",
    "l(\\theta) = \\sum_{i=1}^n \\log \\left( \\sum_{j=1}^k \\alpha_j \\phi(y_i; \\mu_j, \\Sigma_j) \\right)\n",
    "\\end{align}\n",
    "\n",
    "where $\\theta = (\\alpha, \\mu, \\Sigma)$\n",
    "\n",
    "There is no closed form for maximizing the parameters of this log-likelihood, and it is hard to maximize directly.\n",
    "\n",
    "Using EM\n",
    "----\n",
    "\n",
    "Suppose we augment with the latent variable $z$ that indicates which of the $k$ Gaussians our observation $y$ came from. The derivation of the E and M steps are the same as for the toy example, only with more algebra.\n",
    "\n",
    "For the E-step, we have\n",
    "\n",
    "\\begin{align}\n",
    "w_j^i &= Q_i(z^i = j) \\\\\n",
    "&= p(z^i = j | y^i; \\theta) \\\\\n",
    "&= \\frac{p(y^i | z^i = j; \\mu, \\Sigma) p(z^i = j; \\alpha)}  {\\sum_{l=1}^k{p(y^i | z^i = l; \\mu, \\Sigma) p(z^i = l; \\alpha)}}  && \\text{Baye's rule} \\\\\n",
    "&= \\frac{\\phi(y^i; \\mu_j, \\Sigma_j) \\alpha_j}{\\sum_{l=1}^k \\phi(y^i; \\mu_l, \\Sigma_l) \\alpha_l}\n",
    "\\end{align}\n",
    "\n",
    "For the M-step, we have to find $\\theta = (w, \\mu, \\Sigma)$ that maximizes $Q$\n",
    "\n",
    "\\begin{align}\n",
    "\\sum_{i=1}^{m}\\sum{j=1}^{k} Q(z^i=j) \\log \\frac{p(x^i | z^i= j; \\mu, \\Sigma) p(z^i=j; \\alpha)}{Q(z^i=j)}\n",
    "\\end{align}\n",
    "\n",
    "By taking derivatives with respect to $(w, \\mu, \\Sigma)$ respectively and solving (remember to use Lagrange multipliers for the constraint that $\\sum_{j=1}^k w_j = 1$), we get\n",
    "\n",
    "\\begin{align}\n",
    "\\alpha_j &= \\frac{1}{m} \\sum_{i=1}^{m} w_j^i \\\\\n",
    "\\mu_j &= \\frac{\\sum_{i=1}^{m} w_j^i x^i}{\\sum_{i=1}^{m} w_j^i} \\\\\n",
    "\\Sigma_j &= \\frac{\\sum_{i=1}^{m} w_j^i (x^i - \\mu)(x^i - \\mu)^T}{\\sum_{i1}^{m} w_j^i}\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from scipy.stats import multivariate_normal as mvn\n",
    "\n",
    "def normalize(xs, axis=None):\n",
    "    \"\"\"Return normalized marirx so that sum of row or column (default) entries = 1.\"\"\"\n",
    "    if axis is None:\n",
    "        return xs/xs.sum()\n",
    "    elif axis==0:\n",
    "        return xs/xs.sum(0)\n",
    "    else:\n",
    "        return xs/xs.sum(1)[:, None]\n",
    "\n",
    "def mix_mvn_pdf(xs, pis, mus, sigmas):\n",
    "    return np.array([pi*mvn(mu, sigma).pdf(xs) for (pi, mu, sigma) in zip(pis, mus, sigmas)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def em_gmm_orig(xs, pis, mus, sigmas, tol=0.01, max_iter=100):\n",
    "\n",
    "    n, p = xs.shape\n",
    "    k = len(pis)\n",
    "\n",
    "    ll_old = 0\n",
    "    for i in range(max_iter):\n",
    "        exp_A = []\n",
    "        exp_B = []\n",
    "        ll_new = 0\n",
    "\n",
    "        # E-step\n",
    "        ws = np.zeros((k, n))\n",
    "        for j in range(len(mus)):\n",
    "            for i in range(n):\n",
    "                ws[j, i] = pis[j] * mvn(mus[j], sigmas[j]).pdf(xs[i])\n",
    "        ws /= ws.sum(0)\n",
    "\n",
    "    #     # vectorized version\n",
    "    #     for j in range(len(mus)):\n",
    "    #         ws[j, :] = pis[j] * mvn(mus[j], sigmas[j]).pdf(xs)\n",
    "    #     ws /= ws.sum(0)\n",
    "\n",
    "        # M-step\n",
    "        pis = np.zeros(k)\n",
    "        for j in range(len(mus)):\n",
    "            for i in range(n):\n",
    "                pis[j] += ws[j, i]\n",
    "        pis /= n\n",
    "\n",
    "    #     # vectorization 1\n",
    "    #     pis = np.zeros(len(mus))\n",
    "    #     for j in range(len(mus)):\n",
    "    #         pis[j] = ws[j,:].sum()\n",
    "    #     pis /= m\n",
    "\n",
    "    #     # vectorization 2\n",
    "    #     pis = ws.sum(axis=1)\n",
    "    #     pis /= m\n",
    "\n",
    "        mus = np.zeros((k, p))\n",
    "        for j in range(k):\n",
    "            for i in range(n):\n",
    "                mus[j] += ws[j, i] * xs[i]\n",
    "            mus[j] /= ws[j, :].sum()\n",
    "\n",
    "    #     # vectorization 1\n",
    "    #     mus = np.zeros((len(mus), p))\n",
    "    #     for j in range(len(mus)):\n",
    "    #         for k in range(m):\n",
    "    #             mus[j] += ws[j, k] * xs[k]\n",
    "    #     mus /= ws.sum(axis=1)[:, None]\n",
    "\n",
    "    #     # Vectorization 2\n",
    "    #     mus = np.zeros((len(mus), p))\n",
    "    #     for j in range(len(mus)):\n",
    "    #         mus[j, :] = np.dot(ws[j, :], xs)\n",
    "    #         mus[j] /= ws[j, :].sum()\n",
    "\n",
    "    #     # Vectorization 3\n",
    "    #     mus = np.dot(ws, xs)\n",
    "    #     mus /= ws.sum(1)[:, None]\n",
    "\n",
    "        sigmas = np.zeros((k, p, p))\n",
    "        for j in range(k):\n",
    "            for i in range(n):\n",
    "                ys = np.reshape(xs[i]- mus[j], (2,1))\n",
    "                sigmas[j] += ws[j, i] * np.dot(ys, ys.T)\n",
    "            sigmas[j] /= ws[j,:].sum()\n",
    "\n",
    "    #     # Vectorization 1\n",
    "    #     sigmas = np.zeros((len(mus), p, p))\n",
    "    #     for j in range(len(mus)):\n",
    "    #         ys = xs - mus[j, :]\n",
    "    #         sigmas[j] = (ws[j,:,None,None] * mm(ys[:,:,None], ys[:,None,:])).sum(axis=0)\n",
    "    #     sigmas /= ws.sum(axis=1)[:,None,None]\n",
    "\n",
    "        # update complete log likelihoood \n",
    "        ll_new = 0.0\n",
    "        for i in range(n):\n",
    "            s = 0\n",
    "            for j in range(k):\n",
    "                s += pis[j] * mvn(mus[j], sigmas[j]).pdf(xs[i])\n",
    "            ll_new += np.log(s)\n",
    "\n",
    "    #     # Vectorization 1\n",
    "    #     ll_new = 0\n",
    "    #     for pi, mu, sigma in zip(pis, mus, sigmas):\n",
    "    #         ll_new += pi*mvn(mu, sigma).pdf(xs)\n",
    "    #     ll_new = np.log(ll_new).sum()\n",
    "\n",
    "        if np.abs(ll_new - ll_old) < tol:\n",
    "            break\n",
    "        ll_old = ll_new\n",
    "\n",
    "    return ll_new, pis, mus, sigmas"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Vectorized version\n",
    "----"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def em_gmm_vect(xs, pis, mus, sigmas, tol=0.01, max_iter=100):\n",
    "\n",
    "    n, p = xs.shape\n",
    "    k = len(pis)\n",
    "\n",
    "    ll_old = 0\n",
    "    for i in range(max_iter):\n",
    "        exp_A = []\n",
    "        exp_B = []\n",
    "        ll_new = 0\n",
    "\n",
    "        # E-step\n",
    "        ws = np.zeros((k, n))\n",
    "        for j in range(k):\n",
    "            ws[j, :] = pis[j] * mvn(mus[j], sigmas[j]).pdf(xs)\n",
    "        ws /= ws.sum(0)\n",
    "\n",
    "        # M-step\n",
    "        pis = ws.sum(axis=1)\n",
    "        pis /= n\n",
    "\n",
    "        mus = np.dot(ws, xs)\n",
    "        mus /= ws.sum(1)[:, None]\n",
    "\n",
    "        sigmas = np.zeros((k, p, p))\n",
    "        for j in range(k):\n",
    "            ys = xs - mus[j, :]\n",
    "            sigmas[j] = (ws[j,:,None,None] * mm(ys[:,:,None], ys[:,None,:])).sum(axis=0)\n",
    "        sigmas /= ws.sum(axis=1)[:,None,None]\n",
    "\n",
    "        # update complete log likelihoood \n",
    "        ll_new = 0\n",
    "        for pi, mu, sigma in zip(pis, mus, sigmas):\n",
    "            ll_new += pi*mvn(mu, sigma).pdf(xs)\n",
    "        ll_new = np.log(ll_new).sum()\n",
    "\n",
    "        if np.abs(ll_new - ll_old) < tol:\n",
    "            break\n",
    "        ll_old = ll_new\n",
    "\n",
    "    return ll_new, pis, mus, sigmas"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Vectorization with Einstein summation notation\n",
    "----"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def em_gmm_eins(xs, pis, mus, sigmas, tol=0.01, max_iter=100):\n",
    "\n",
    "    n, p = xs.shape\n",
    "    k = len(pis)\n",
    "\n",
    "    ll_old = 0\n",
    "    for i in range(max_iter):\n",
    "        exp_A = []\n",
    "        exp_B = []\n",
    "        ll_new = 0\n",
    "\n",
    "        # E-step\n",
    "        ws = np.zeros((k, n))\n",
    "        for j, (pi, mu, sigma) in enumerate(zip(pis, mus, sigmas)):\n",
    "            ws[j, :] = pi * mvn(mu, sigma).pdf(xs)\n",
    "        ws /= ws.sum(0)\n",
    "\n",
    "        # M-step\n",
    "        pis = np.einsum('kn->k', ws)/n\n",
    "        mus = np.einsum('kn,np -> kp', ws, xs)/ws.sum(1)[:, None]\n",
    "        sigmas = np.einsum('kn,knp,knq -> kpq', ws, \n",
    "            xs-mus[:,None,:], xs-mus[:,None,:])/ws.sum(axis=1)[:,None,None]\n",
    "\n",
    "        # update complete log likelihoood \n",
    "        ll_new = 0\n",
    "        for pi, mu, sigma in zip(pis, mus, sigmas):\n",
    "            ll_new += pi*mvn(mu, sigma).pdf(xs)\n",
    "        ll_new = np.log(ll_new).sum()\n",
    "\n",
    "        if np.abs(ll_new - ll_old) < tol:\n",
    "            break\n",
    "        ll_old = ll_new\n",
    "\n",
    "    return ll_new, pis, mus, sigmas"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Comparison of EM routines\n",
    "----"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.random.seed(123)\n",
    "\n",
    "# create data set\n",
    "n = 1000\n",
    "_mus = np.array([[0,4], [-2,0]])\n",
    "_sigmas = np.array([[[3, 0], [0, 0.5]], [[1,0],[0,2]]])\n",
    "_pis = np.array([0.6, 0.4])\n",
    "xs = np.concatenate([np.random.multivariate_normal(mu, sigma, int(pi*n)) \n",
    "                    for pi, mu, sigma in zip(_pis, _mus, _sigmas)])\n",
    "\n",
    "# initial guesses for parameters\n",
    "pis = normalize(np.random.random(2))\n",
    "mus = np.random.random((2,2))\n",
    "sigmas = np.array([np.eye(2)] * 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "ll1, pis1, mus1, sigmas1 = em_gmm_orig(xs, pis, mus, sigmas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "intervals = 101\n",
    "ys = np.linspace(-8,8,intervals)\n",
    "X, Y = np.meshgrid(ys, ys)\n",
    "_ys = np.vstack([X.ravel(), Y.ravel()]).T\n",
    "\n",
    "z = np.zeros(len(_ys))\n",
    "for pi, mu, sigma in zip(pis1, mus1, sigmas1):\n",
    "    z += pi*mvn(mu, sigma).pdf(_ys)\n",
    "z = z.reshape((intervals, intervals))\n",
    "\n",
    "ax = plt.subplot(111)\n",
    "plt.scatter(xs[:,0], xs[:,1], alpha=0.2)\n",
    "plt.contour(X, Y, z, N=10)\n",
    "plt.axis([-8,6,-6,8])\n",
    "ax.axes.set_aspect('equal')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "ll2, pis2, mus2, sigmas2 = em_gmm_vect(xs, pis, mus, sigmas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "intervals = 101\n",
    "ys = np.linspace(-8,8,intervals)\n",
    "X, Y = np.meshgrid(ys, ys)\n",
    "_ys = np.vstack([X.ravel(), Y.ravel()]).T\n",
    "\n",
    "z = np.zeros(len(_ys))\n",
    "for pi, mu, sigma in zip(pis2, mus2, sigmas2):\n",
    "    z += pi*mvn(mu, sigma).pdf(_ys)\n",
    "z = z.reshape((intervals, intervals))\n",
    "\n",
    "ax = plt.subplot(111)\n",
    "plt.scatter(xs[:,0], xs[:,1], alpha=0.2)\n",
    "plt.contour(X, Y, z, N=10)\n",
    "plt.axis([-8,6,-6,8])\n",
    "ax.axes.set_aspect('equal')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "ll3, pis3, mus3, sigmas3 = em_gmm_eins(xs, pis, mus, sigmas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# %timeit em_gmm_orig(xs, pis, mus, sigmas)\n",
    "%timeit em_gmm_vect(xs, pis, mus, sigmas)\n",
    "%timeit em_gmm_eins(xs, pis, mus, sigmas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "intervals = 101\n",
    "ys = np.linspace(-8,8,intervals)\n",
    "X, Y = np.meshgrid(ys, ys)\n",
    "_ys = np.vstack([X.ravel(), Y.ravel()]).T\n",
    "\n",
    "z = np.zeros(len(_ys))\n",
    "for pi, mu, sigma in zip(pis3, mus3, sigmas3):\n",
    "    z += pi*mvn(mu, sigma).pdf(_ys)\n",
    "z = z.reshape((intervals, intervals))\n",
    "\n",
    "ax = plt.subplot(111)\n",
    "plt.scatter(xs[:,0], xs[:,1], alpha=0.2)\n",
    "plt.contour(X, Y, z, N=10)\n",
    "plt.axis([-8,6,-6,8])\n",
    "ax.axes.set_aspect('equal')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [conda root]",
   "language": "python",
   "name": "conda-root-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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}