07_gaussian_mixture_models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Music Machine Learning - Gaussian Mixture Models\n",
    "\n",
    "### Author: Philippe Esling (esling@ircam.fr)\n",
    "\n",
    "In this course we will cover\n",
    "1. Unsupervised learning and [clustering](#clustering).\n",
    "2. Motivating the need for [latent variable](#latent) models.\n",
    "2. Approaching the problem naïvely with the [k-Means](#kmeans) algorithm.\n",
    "\n",
    "We begin with the standard imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns; sns.set()\n",
    "import numpy as np\n",
    "from helper_plot import hdr_plot_style\n",
    "hdr_plot_style()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining mixture models\n",
    "\n",
    "We have seen previously how to fit a Gaussian, so we could apply the same reasoning and try to fit a single multivariate Gaussian to our dataset\n",
    "$$\n",
    "\\begin{align}\n",
    "p(\\mathbf{x} \\vert \\theta) &= \\mathcal{N}\\left(\\mathbf{x}\\vert \\mu, \\Sigma\\right) \\\\\n",
    "\\theta &= \\{\\mu, \\Sigma\\}\n",
    "\\end{align}\n",
    "$$\n",
    "However, we can quickly see that our data contains several groups, and the single Gaussian seems largely not adapted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from helper_plot import fit_multivariate_gaussian\n",
    "from sklearn.datasets.samples_generator import make_blobs\n",
    "X, y_true = make_blobs(n_samples=400, centers=4, cluster_std=0.60, random_state=0)\n",
    "X = X[:, ::-1]; \n",
    "fit_multivariate_gaussian(X);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instead of using a single Gaussian, we could use several Gaussians, for instance, if we use three Gaussians, we have the following model\n",
    "$$\n",
    "\\begin{align}\n",
    "p(\\mathbf{x} \\vert \\theta) &= \\pi_{1} \\mathcal{N}\\left(\\mathbf{x}\\vert \\mu_{1}, \\Sigma_{1}\\right) \\\\\n",
    "&+ \\pi_{2} \\mathcal{N}\\left(\\mathbf{x}\\vert \\mu_{2}, \\Sigma_{2}\\right)\\\\\n",
    "&+ \\pi_{3} \\mathcal{N}\\left(\\mathbf{x}\\vert \\mu_{3}, \\Sigma_{3}\\right)\\\\\n",
    "&+ \\pi_{4} \\mathcal{N}\\left(\\mathbf{x}\\vert \\mu_{4}, \\Sigma_{4}\\right)\\\\\n",
    "\\end{align}\n",
    "$$\n",
    "So now the parameters are both the parameters of each Gaussian, along with the probabilities of belonging to each Gaussian\n",
    "$$\n",
    "\\theta = \\{\\pi_{1}, \\pi_{2}, \\pi_{3}, \\pi_{4}, \\mu_{1}, \\mu_{2}, \\mu_{3}, \\mu_{4}, \\Sigma_{1}, \\Sigma_{2}, \\Sigma_{3}, \\Sigma_{4}\\}\n",
    "$$\n",
    "Note that if we solve this model, called a _Gaussian Mixture Model_ (GMM), we can find the probabilities of each point belonging to one of the corresponding cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from helper_plot import fit_gaussian_mixture\n",
    "fit_gaussian_mixture(X);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compared to a single Gaussian, a GMM has obviously a higher flexibility, but also a larger number of parameters (which makes it harder to optimize). Generally speaking, the full distribution of a GMM is defined as follows\n",
    "$$\n",
    "p(\\mathbf{x} \\vert \\theta) = \\sum_{c} \\pi_{c} \\mathcal{N}\\left(\\mathbf{x}\\vert \\mu_{c}, \\Sigma_{c}\\right)\n",
    "$$\n",
    "\n",
    "We can see that any *mixture model*, is defined through two main components.\n",
    "\n",
    "**Base distribution**: $p(\\mathbf{x}\\vert z=c, \\theta) = \\mathcal{N}\\left(\\mathbf{x}\\vert \\mu_{c}, \\Sigma_{c}\\right)$\n",
    "\n",
    "**Mixing probabilities** $p(z=c \\vert \\theta) = \\pi_{c}$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimizing this model\n",
    "\n",
    "If we want to optimize our probabilistic model, we have to solve a typical *maximum likelihood* problem (considering that points $x_{i}$ are independent)\n",
    "$$\n",
    "\\underset{\\theta}{\\text{max. }} p(x \\vert \\theta) = \\prod_{i=1}^{N} p(x_{i} \\vert \\theta)\n",
    "$$\n",
    "In the case of a GMM, we could try to replace our full model inside this problem, which leads to a rather complicated formulation.\n",
    "$$\n",
    "\\underset{\\theta}{\\text{max. }} \\prod_{i=1}^{N} \\sum_{c} \\pi_{c} \\mathcal{N}\\left(x_{i}\\vert \\mu_{c}, \\Sigma_{c}\\right)\n",
    "$$\n",
    "\n",
    "Furthermore, we have to include the following constraints\n",
    "$$\n",
    "\\sum_{c} \\pi_{c} = 1; \\pi_{c} \\geq 0, \\forall c \\\\\n",
    "\\Sigma_{c} \\succcurlyeq 0\n",
    "$$\n",
    "\n",
    "Although, we could rely on optimization, we need to incorporate hard to follow constraints, which leads to slow and unstable training. Here, the reasoning is way easier if we introduce a latent variable"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Introducing a latent variable\n",
    "\n",
    "We introduce the latent variable $z$, which _explains_ a given datapoint. In our case, this variable tells us _from which group_ the point was generated.\n",
    "\n",
    "In our case, we can simply say that our latent variable corresponds to the mixing probabilities\n",
    "$$\n",
    "\\begin{align}\n",
    "p(z=c\\vert\\theta) &= \\pi_{c} \\\\\n",
    "p(\\mathbf{x}\\vert z=c,\\theta) &= \\mathcal{N}\\left(\\mathbf{x}\\vert \\mu_{c}, \\Sigma_{c}\\right)\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Therefore, our original model can be retrieved as a marginalization of the latent variable, as\n",
    "$$\n",
    "p(\\mathbf{x} \\vert \\theta) = \\sum_{c} p(\\mathbf{x} \\vert z=c,\\theta) p(z=c\\vert\\theta)\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Estimating parameters\n",
    "\n",
    "Now in order to optimize the parameters, we should need to optimize both the $\\theta$ and $z$ variables. If we knew the cluster belonging (latent) variable $z$, then the optimization would be easy, since we know that\n",
    "$$\n",
    "p(\\mathbf{x}\\vert z=c, \\theta) = \\mathcal{N}\\left(\\mathbf{x}\\vert \\mu_{c}, \\Sigma_{c}\\right)\n",
    "$$\n",
    "And computing the parameters of this Gaussian is pretty easy, as we simply need to update the values depending on the corresponding points (weighted mean)\n",
    "$$\n",
    "\\mu_{c} = \\frac{\\sum_{i} p(z=i\\vert \\mathbf{x}_{i}, \\theta) \\bx_{i}}{\\sum_{i} p(z=i\\vert \\mathbf{x}_{i}, \\theta)}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now in order to optimize the source variable $z$, the optimization would be easy if we knew the parameters as we have\n",
    "$$\n",
    "p(z = c \\vert \\mathbf{x}, \\theta) = \\frac{p(\\mathbf{x} \\vert z = c, \\theta) p(z = c \\vert \\theta)}{Z}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Expectation-Maximization\n",
    "\n",
    "As we have seen in the slides, a natural alternating solution emerges in the form of the **Expectation-Maximization** algorithm\n",
    "1. Start with a random set of Gaussians parameters $\\theta^{(0)}$\n",
    "2. Repeat until convergence\n",
    "    1. *E-step*: for each point, find weights encoding the probability of membership in each cluster by computing $p(z=c \\vert \\mathbf{x}_{i}, \\theta^{(t)})$\n",
    "    2. *M-step*: for each cluster, update its location, normalization, and shape based on *all* data points, making use of the weights, hence updating Gaussian parameters $\\theta^{(t+1)}$\n",
    "\n",
    "The result of this is that each cluster is associated not with a hard-edged sphere, but with a smooth Gaussian model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Deriving the E-Step\n",
    "As discussed in the slides, we want to minimize the gap between our lower bound and the true distribution, which can be written as\n",
    "$$\n",
    "\\mathcal{D} = \\log p(\\mathbf{x} \\vert \\theta) - \\mathcal{L}(\\theta\\vert q)\n",
    "$$\n",
    "Developing this expression, we find out that\n",
    "$$\n",
    "\\mathcal{D} = \\sum_{i=1}^{N} \\mathcal{D}_{KL} \\left[ q(z_{i}) \\Vert p(z_{i} \\vert \\mathbf{x}_{i}, \\theta) \\right]\n",
    "$$\n",
    "Hence, we need to minimize the KL divergence, which amounts to setting the approximate function as the posterior distribution\n",
    "$$\n",
    "q(z_{i}) = p(z_{i} \\vert \\mathbf{x}_{i}, \\theta)\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Deriving the M-step\n",
    "\n",
    "Based on the approximation that we want to maximize\n",
    "$$\n",
    "\\mathcal{L}(\\theta, q) = \\sum_{i=1}^{N} \\sum_{c} q(t_{i}=c) \\log\\frac{p(x_{i}, t_{i}=c \\vert \\theta)}{q(t_{i}=c)}\n",
    "$$\n",
    "We have seen that expending this maximization led to\n",
    "$$\n",
    "\\mathcal{L}(\\theta, q) = \\mathbb{E}_{q} \\log p(x,t\\vert \\theta) + const\n",
    "$$\n",
    "Note that the fonction is usually concave, which makes it easy to optimize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true
   },
   "source": [
    "## Using GMM with `scikit-learn`\n",
    "\n",
    "By using the EM algorithm, a GMM will find a mixture of multi-dimensional Gaussian probability distributions that best model any input dataset, which can elegantly solve our clustering problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "editable": true,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "from sklearn.mixture import GaussianMixture\n",
    "gmm = GaussianMixture(n_components=4).fit(X)\n",
    "labels = gmm.predict(X)\n",
    "plt.figure(figsize=(10,8))\n",
    "plt.scatter(X[:, 0], X[:, 1], c=labels, s=40, cmap='magma', edgecolor='w');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true
   },
   "source": [
    "But because GMM contains a probabilistic model under the hood, it is also possible to find probabilistic cluster assignments—in Scikit-Learn this is done using the ``predict_proba`` method.\n",
    "This returns a matrix of size ``[n_samples, n_clusters]`` which measures the probability that any point belongs to the given cluster:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "editable": true,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "probs = gmm.predict_proba(X)\n",
    "print(probs[:5].round(3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true
   },
   "source": [
    "We can visualize this uncertainty by, for example, making the size of each point proportional to the certainty of its prediction; looking at the following figure, we can see that it is precisely the points at the boundaries between clusters that reflect this uncertainty of cluster assignment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "editable": true,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "size = 100 * (probs.max(1) ** 2)  # square emphasizes differences\n",
    "plt.figure(figsize=(10,8))\n",
    "plt.scatter(X[:, 0], X[:, 1], c=labels, cmap='magma', s=size, edgecolor='w');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true
   },
   "source": [
    "With this in place, we can take a look at what the four-component GMM gives us for our initial data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "editable": true,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "from helper_plot import plot_gmm\n",
    "gmm = GaussianMixture(n_components=4, random_state=42)\n",
    "plot_gmm(gmm, X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true
   },
   "source": [
    "Similarly, we can use the GMM approach to fit our stretched dataset. Here, we already see that allowing for a full covariance the model will fit even very oblong, stretched-out clusters (as opposed to the *hard* version of kMeans)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "editable": true,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "rng = np.random.RandomState(13)\n",
    "X_stretched = np.dot(X, rng.randn(2, 2))\n",
    "gmm = GaussianMixture(n_components=4, covariance_type='full', random_state=42)\n",
    "plot_gmm(gmm, X_stretched)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true
   },
   "source": [
    "This makes clear that GMM addresses the two main practical issues with *k*-means encountered before."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "dL3A2sntBEDj"
   },
   "source": [
    "## Implementing EM for GMM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0xPS_VdpBEDk"
   },
   "source": [
    "Now, we will derive and implement formulas for the GMM ourselves. To do so, we will use samples from a Gaussian mixture model with unknown mean, variance, and priors. We will also have to derive both the E-step and M-step for our selected model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "AmgeHTCfBEDp"
   },
   "source": [
    "### Reminder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RgAwWi1nBEDq"
   },
   "source": [
    "We recall that the EM algorithm is a coordinate descent optimization of a lower bound \n",
    "$$\n",
    "\\underset{\\theta}{\\text{argmax }} \\mathcal{L}(\\theta, q) = \\underset{\\theta}{\\text{argmax }} \\int q(\\mathbf{z}) \\log\\frac{p(\\mathbf{x}, \\mathbf{z} \\vert \\theta)}{q(\\mathbf{z})}d\\mathbf{z}\n",
    "$$\n",
    "\n",
    "**E-step**: In the E-step, we use a lower-bounding approximation to find the latent variables $\\mathbf{z}$, while using the current parameters estimate $\\theta^{(t)}$\n",
    "$$\n",
    "\\underset{q}{\\text{argmax }} \\mathcal{L}(\\theta, q) \\Leftrightarrow \\underset{q}{\\text{argmin }} \\mathcal{D}_{KL} \\left[q(\\mathbf{z}) \\Vert p(\\mathbf{z}\\vert\\mathbf{x}, \\theta) \\right]\n",
    "$$\n",
    "\n",
    "We have seen that the optimal solution is $q(\\mathbf{z}) = p(\\mathbf{z}\\vert\\mathbf{x}, \\theta)$, which means that we can safely use the conditional probability for this step. \n",
    "\n",
    "**M-step**: In the M-step, we use our updated knowledge on the latent variable to update the parameters. Hence, we have to optimize\n",
    "$$\n",
    "\\underset{\\theta}{\\text{argmax }} \\mathcal{L}(\\theta, q) \\Leftrightarrow \\underset{\\theta}{\\text{argmax }} \\mathbb{E}_{q(\\mathbf{z})}\\log p(\\mathbf{x},\\mathbf{z} | \\theta)\n",
    "$$\n",
    "\n",
    "For our GMM model, $\\theta$ is a set of parameters that consists of mean vectors $\\mu_c$, covariance matrices $\\Sigma_c$ and priors $\\pi_c$ for each component. The *latent variable* $z$ are indices of components to which each data point is assigned, i.e. $z_i$  is the cluster index for object $\\mathbf{x}_i$. Therefore, the joint distribution of our GMM can be written as follows\n",
    "\n",
    "$$\\log p(\\mathbf{z}, X \\mid \\theta) =  \\sum\\limits_{i=1}^N \\log p(z_i, x_i \\mid \\theta) = \\sum\\limits_{i=1}^N \\sum\\limits_{c=1}^C q(z_i = c) \\log \\left (\\pi_c \\, \\mathcal{N}(x_i \\mid \\mu_c, \\Sigma_c)\\right)\n",
    "$$\n",
    "where $\\mathcal{N}(x_i \\mid \\mu_c, \\Sigma_c) = \\frac{1}{\\sqrt{(2\\pi)^n|\\boldsymbol\\Sigma_c|}}\n",
    "\\exp\\left(-\\frac{1}{2}({x}-{\\mu_c})^T{\\boldsymbol\\Sigma_c}^{-1}({x}-{\\mu_c})\n",
    "\\right)$ is the probability density function (pdf) of the normal distribution."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ELmi1nAtBEDr"
   },
   "source": [
    "### E-step\n",
    "In this step we need to estimate the posterior distribution over the latent variables with fixed values of parameters: $q_i(z_i) = p(z_i \\mid \\mathbf{x}_i, \\theta)$. We assume that $z_i$ equals to the cluster index of the true component of the $\\mathbf{x}_i$ object. To do so we need to compute $\\pi_{ic} = p(z_i = c \\mid \\mathbf{x}_i, \\theta)$. Note that $\\sum\\limits_{c=1}^C\\pi_{ic}=1$.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KL87Q8-TBEDu"
   },
   "source": [
    "***\n",
    "**Exercise**\n",
    "\n",
    "1. Implement E-step for GMM using template below.\n",
    "\n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wCUCaD28BEDw"
   },
   "outputs": [],
   "source": [
    "from scipy.stats import multivariate_normal\n",
    "\n",
    "def E_step(X, pi, mu, sigma):\n",
    "    \"\"\"\n",
    "    Performs E-step on GMM model\n",
    "    Each input is numpy array:\n",
    "    X: (N x d), data points\n",
    "    pi: (C), mixture component weights \n",
    "    mu: (C x d), mixture component means\n",
    "    sigma: (C x d x d), mixture component covariance matrices\n",
    "    \n",
    "    Returns:\n",
    "    gamma: (N x C), probabilities of clusters for objects\n",
    "    \"\"\"\n",
    "    N = X.shape[0] # number of objects\n",
    "    C = pi.shape[0] # number of clusters\n",
    "    d = mu.shape[1] # dimension of each object\n",
    "    pi_update = np.zeros((N, C)) # distribution q(T)\n",
    "    \n",
    "    ######################\n",
    "    # YOUR CODE GOES HERE\n",
    "    ######################\n",
    "    \n",
    "    return pi_update"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xlhktOlMBED1"
   },
   "outputs": [],
   "source": [
    "pi0 = np.random.randn(4, X.shape[0])\n",
    "mu0 = np.random.randn(4, 2)\n",
    "sigma0 = np.zeros((4, 2, 2))\n",
    "for i in range(4):\n",
    "    sigma0[i] = np.eye(2, 2)\n",
    "pi = E_step(X, pi0, mu0, sigma0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "fERrQWSCBED5"
   },
   "source": [
    "### M-step\n",
    "\n",
    "In M-step we need to maximize $\\mathbb{E}_{q(T)}\\log p(X,T | \\theta)$ with respect to $\\theta$. In our model this means that we need to find optimal values of $\\pi$, $\\mu$, $\\Sigma$. To do so, you need to compute the derivatives and \n",
    "set them to zero. You should start by deriving formulas for $\\mu$ as it is the easiest part. Then move on to $\\Sigma$. Here it is crucial to optimize function w.r.t. to $\\Lambda = \\Sigma^{-1}$ and then inverse obtained result. Finaly, to compute $\\pi$, you will need <a href=\"https://www3.nd.edu/~jstiver/FIN360/Constrained%20Optimization.pdf\">Lagrange Multipliers technique</a> to satisfy constraint $\\sum\\limits_{i=1}^{n}\\pi_i = 1$.\n",
    "\n",
    "<br>\n",
    "<b>Important note:</b> You will need to compute derivatives of scalars with respect to matrices. To refresh this technique from previous courses, see <a href=\"https://en.wikipedia.org/wiki/Matrix_calculus\"> wiki article</a> about it . Main formulas of matrix derivatives can be found in <a href=\"http://www2.imm.dtu.dk/pubdb/views/edoc_download.php/3274/pdf/imm3274.pdf\">Chapter 2 of The Matrix Cookbook</a>. For example, there you may find that $\\frac{\\partial}{\\partial A}\\log |A| = A^{-T}$.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "A31OYSogBED6"
   },
   "source": [
    "<b>Task 2:</b> Implement M-step for GMM using template below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mhOr5I1bBED7"
   },
   "outputs": [],
   "source": [
    "def M_step(X, pi_update):\n",
    "    \"\"\"\n",
    "    Performs M-step on GMM model\n",
    "    Each input is numpy array:\n",
    "    X: (N x d), data points\n",
    "    gamma: (N x C), distribution q(T)  \n",
    "    \n",
    "    Returns:\n",
    "    pi: (C)\n",
    "    mu: (C x d)\n",
    "    sigma: (C x d x d)\n",
    "    \"\"\"\n",
    "    N = X.shape[0] # number of objects\n",
    "    C = pi_update.shape[1] # number of clusters\n",
    "    d = X.shape[1] # dimension of each object\n",
    "    \n",
    "    ######################\n",
    "    # YOUR CODE GOES HERE\n",
    "    ######################\n",
    "    \n",
    "    return pi, mu, sigma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6i5cv65SBED-"
   },
   "outputs": [],
   "source": [
    "pi_update = E_step(X, pi0, mu0, sigma0)\n",
    "pi, mu, sigma = M_step(X, pi_update)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "svyzjt7XBEEC"
   },
   "source": [
    "### Loss function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "qexOXBLUBEED"
   },
   "source": [
    "Finally, we need some function to track convergence. We will use variational lower bound $\\mathcal{L}$ for this purpose. We will stop our EM iterations when $\\mathcal{L}$ will saturate. Usually, you will need only about 10-20 iterations to converge. It is also useful to check that this function never decreases during training. If it does, you have a bug in your code.\n",
    "\n",
    "<b>Task 3:</b> Implement a function that will compute $\\mathcal{L}$ using template below.\n",
    "\n",
    "$$\\mathcal{L} = \\sum_{i=1}^{N} \\sum_{c=1}^{C} q(t_i =c) (\\log \\pi_c + \\log f_{\\!\\mathcal{N}}(x_i \\mid \\mu_c, \\Sigma_c)) - \\sum_{i=1}^{N} \\sum_{c=1}^{K} q(t_i =c) \\log q(t_i =c)$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "B5tKCZe0BEEE"
   },
   "outputs": [],
   "source": [
    "def compute_vlb(X, pi, mu, sigma, pi_update):\n",
    "    \"\"\"\n",
    "    Each input is numpy array:\n",
    "    X: (N x d), data points\n",
    "    gamma: (N x C), distribution q(T)  \n",
    "    pi: (C)\n",
    "    mu: (C x d)\n",
    "    sigma: (C x d x d)\n",
    "    \n",
    "    Returns value of variational lower bound\n",
    "    \"\"\"\n",
    "    N = X.shape[0] # number of objects\n",
    "    C = pi_update.shape[1] # number of clusters\n",
    "    d = X.shape[1] # dimension of each object\n",
    "\n",
    "    \n",
    "    ######################\n",
    "    # YOUR CODE GOES HERE\n",
    "    ######################\n",
    "    \n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xNva3XRTBEEI"
   },
   "outputs": [],
   "source": [
    "pi, mu, sigma = pi0, mu0, sigma0\n",
    "pi_update = E_step(X, pi, mu, sigma)\n",
    "pi, mu, sigma = M_step(X, pi_update)\n",
    "loss = compute_vlb(X, pi, mu, sigma, pi_update)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "I5p8XC-eBEEM"
   },
   "source": [
    "### Bringing it all together"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "cQsdsNVCBEEU"
   },
   "source": [
    "Now that we have E step, M step and VLB, we can implement the training loop. We will initialize values of $\\pi$, $\\mu$ and $\\Sigma$ to some random numbers, train until $\\mathcal{L}$ stops changing, and return the resulting points. We also know that the EM algorithm converges to local optima. To find a better local optima, we will restart the algorithm multiple times from different (random) starting positions. Each training trial should stop either when maximum number of iterations is reached or when relative improvement is smaller than given tolerance ($|\\frac{\\mathcal{L}_i-\\mathcal{L}_{i-1}}{\\mathcal{L}_{i-1}}| \\le \\text{rtol}$).\n",
    "\n",
    "Remember, that initial (random) values of $\\pi$ that you generate must be non-negative and sum up to 1. Also, $\\Sigma$ matrices must be symmetric and positive semi-definite. If you don't know how to generate those matrices, you can use $\\Sigma=I$ as initialization.\n",
    "\n",
    "You will also sometimes get numerical errors because of component collapsing. The easiest way to deal with this problems is to restart the procedure.\n",
    "\n",
    "<b>Task 4:</b> Implement training procedure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Q1AAcyl0BEEW"
   },
   "outputs": [],
   "source": [
    "import math\n",
    "def train_EM(X, C, rtol=1e-3, max_iter=100, restarts=10):\n",
    "    '''\n",
    "    Starts with random initialization *restarts* times\n",
    "    Runs optimization until saturation with *rtol* reached\n",
    "    or *max_iter* iterations were made.\n",
    "    \n",
    "    X: (N, d), data points\n",
    "    C: int, number of clusters\n",
    "    '''\n",
    "    N = X.shape[0] # number of objects\n",
    "    d = X.shape[1] # dimension of each object\n",
    "    best_loss = None\n",
    "    best_pi = None\n",
    "    best_mu = None\n",
    "    best_sigma = None\n",
    "\n",
    "    for _ in range(restarts):\n",
    "        try:\n",
    "            ######################\n",
    "            # YOUR CODE GOES HERE\n",
    "            ######################\n",
    "            \n",
    "        except np.linalg.LinAlgError:\n",
    "            print(\"Singular matrix: components collapsed\")\n",
    "            pass\n",
    "\n",
    "    return best_loss, best_pi, best_mu, best_sigma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0f8A5sbmBEEZ"
   },
   "outputs": [],
   "source": [
    "best_loss, best_pi, best_mu, best_sigma = train_EM(X, 4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "tYd6tPHKBEEd"
   },
   "source": [
    "If you implemented all the steps correctly, your algorithm should converge in about 20 iterations. Let's plot the clusters to see it. We will assign a cluster label as the most probable cluster index. This can be found using a matrix $\\gamma$ computed on last E-step. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mK_M6QLnBEEe"
   },
   "outputs": [],
   "source": [
    "gamma = E_step(X, best_pi, best_mu, best_sigma)\n",
    "labels = gamma.argmax(axis=1)\n",
    "colors = np.array([(31, 119, 180), (255, 127, 14), (44, 160, 44), (50, 50, 50)]) / 255.\n",
    "plt.figure(figsize=(10, 8))\n",
    "plt.scatter(X[:, 0], X[:, 1], c=colors[labels], s=50, edgecolors='w')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true
   },
   "source": [
    "## GMM as *Density Estimation*\n",
    "\n",
    "Though GMM is often categorized as a clustering algorithm, fundamentally it is an algorithm for *density estimation*.\n",
    "That is to say, the result of a GMM fit to some data is technically not a clustering model, but a generative probabilistic model describing the distribution of the data.\n",
    "\n",
    "As an example, consider some data generated from Scikit-Learn's ``make_moons`` function, which we saw in [In Depth: K-Means Clustering](05.11-K-Means.ipynb):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "editable": true,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "from sklearn.datasets import make_moons\n",
    "Xmoon, ymoon = make_moons(200, noise=.05, random_state=0)\n",
    "plt.figure(figsize=(10, 8))\n",
    "plt.scatter(Xmoon[:, 0], Xmoon[:, 1], edgecolor='w', s=50);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true
   },
   "source": [
    "If we try to fit this with a two-component GMM viewed as a clustering model, the results are not particularly useful:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "editable": true,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "gmm2 = GaussianMixture(n_components=2, covariance_type='full', random_state=0)\n",
    "plot_gmm(gmm2, Xmoon)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true
   },
   "source": [
    "But if we instead use many more components and ignore the cluster labels, we find a fit that is much closer to the input data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "editable": true,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "gmm16 = GaussianMixture(n_components=16, covariance_type='full', random_state=0)\n",
    "plot_gmm(gmm16, Xmoon, label=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true
   },
   "source": [
    "Here the mixture of 16 Gaussians serves not to find separated clusters of data, but rather to model the overall *distribution* of the input data.\n",
    "This is a generative model of the distribution, meaning that the GMM gives us the recipe to generate new random data distributed similarly to our input.\n",
    "For example, here are 400 new points drawn from this 16-component GMM fit to our original data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "editable": true,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "Xnew = gmm16.sample(400)[0]\n",
    "plt.figure(figsize=(10, 8))\n",
    "plt.scatter(Xnew[:, 0], Xnew[:, 1], s=40, edgecolor='w');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true
   },
   "source": [
    "GMM is convenient as a flexible means of modeling an arbitrary multi-dimensional distribution of data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true
   },
   "source": [
    "### How many components?\n",
    "\n",
    "The fact that GMM is a generative model gives us a natural means of determining the optimal number of components for a given dataset.\n",
    "A generative model is inherently a probability distribution for the dataset, and so we can simply evaluate the *likelihood* of the data under the model, using cross-validation to avoid over-fitting.\n",
    "Another means of correcting for over-fitting is to adjust the model likelihoods using some analytic criterion such as the [Akaike information criterion (AIC)](https://en.wikipedia.org/wiki/Akaike_information_criterion) or the [Bayesian information criterion (BIC)](https://en.wikipedia.org/wiki/Bayesian_information_criterion).\n",
    "Scikit-Learn's ``GMM`` estimator actually includes built-in methods that compute both of these, and so it is very easy to operate on this approach.\n",
    "\n",
    "Let's look at the AIC and BIC as a function as the number of GMM components for our moon dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "editable": true,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "n_components = np.arange(1, 21)\n",
    "models = [GaussianMixture(n, covariance_type='full', random_state=0).fit(Xmoon)\n",
    "          for n in n_components]\n",
    "\n",
    "plt.figure(figsize=(10, 8))\n",
    "plt.plot(n_components, [m.bic(Xmoon) for m in models], label='BIC')\n",
    "plt.plot(n_components, [m.aic(Xmoon) for m in models], label='AIC')\n",
    "plt.legend(loc='best')\n",
    "plt.xlabel('n_components');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true
   },
   "source": [
    "The optimal number of clusters is the value that minimizes the AIC or BIC, depending on which approximation we wish to use. The AIC tells us that our choice of 16 components above was probably too many: around 8-12 components would have been a better choice.\n",
    "As is typical with this sort of problem, the BIC recommends a simpler model.\n",
    "\n",
    "Notice the important point: this choice of number of components measures how well GMM works *as a density estimator*, not how well it works *as a clustering algorithm*.\n",
    "I'd encourage you to think of GMM primarily as a density estimator, and use it for clustering only when warranted within simple datasets."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "em_assignment.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.6"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}