density-estimation/2-Gaussian-Mixture-Models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "source": [
    "Some packages we'll need below - make sure you have them installed! \n",
    "\n",
    "If you use `conda`:\n",
    "\n",
    "    conda install numpy matplotlib pygmmis\n",
    "\n",
    "If you use `pip`:\n",
    "\n",
    "    pip install numpy matplotlib pygmmis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "# Density estimation\n",
    "\n",
    "## 2) Gaussian Mixture Models (GMMs)\n",
    "\n",
    "If we want to fit the probability density function (PDF) of set of data samples $\\mathbf{x}$ with a mixture model, that's what we need to do: maximize the observed likelihood\n",
    "\n",
    "$$p(\\mathbf{x}\\,|\\,\\mathbf{\\theta}) = \\prod_{i=1}^N\\sum_{k=1}^K \\alpha_k\\, p(\\mathbf{x}_i\\,|\\,\\mathbf{\\theta}_k)$$\n",
    "\n",
    "with respect to the mixing amplitudes $\\alpha_k$ and per-component parameters $\\mathbf{\\theta}_k$. For GMMs, those parameters are a mean $\\mathbf{\\mu}_k$ and a covariance matrix $\\mathbf{\\Sigma}_k$. A key problem is that the association of data samples to GMM components $k$ is unknown. Otherwise we could simply compute the mean and covariance given all samples associated with any component; the mixing amplitude would then be given by the relative abundance of those samples compared to all data.\n",
    "\n",
    "One common solution was proposed by [Dempster et al. (1977)](http://www.jstor.org/stable/2984875): the *Expectation-Maximization* procedure. Without going into specifics, the methods amounts to iterations of updates to all involved quantities that provably converge a local minimum of the likelihood. It's not the fastest method, but it works:\n",
    "\n",
    "$$\\begin{split}\n",
    "\\text{E-step:}\\ \\ q_{ik} &\\leftarrow \\frac{\\alpha_k p_k(x_i)}{\\sum_j \\alpha_j p_j(x_i)}\\\\\n",
    "\\text{M-step:}\\ \\ \\alpha_k &\\leftarrow \\frac{1}{N} \\sum_i q_{ik} \\equiv \\frac{1}{N} q_k\\\\\n",
    " \\mu_k & \\leftarrow \\frac{1}{q_k} \\sum_i q_{ik}x_i\\\\\n",
    " \\Sigma_k & \\leftarrow \\frac{1}{q_k} \\sum_i q_{ik}\\left[ (\\mu_k - x_i)(\\mu_k - x_i)^\\top\\right].\n",
    "\\end{split}$$\n",
    "\n",
    "As you can see, it first determines the probabilistic assignment of samples to components, $q_{ik}$, and then uses that as weights for calculating the other component parameters.\n",
    "\n",
    "We will now work on a simulated data set, mimicking a redshift survey of a galaxy cluster field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# construct fake redshift data\n",
    "z_cl = 0.35\n",
    "w_cl = 0.005\n",
    "frac_bg = 0.7\n",
    "N = 1000\n",
    "bins = np.linspace(0,1,51)\n",
    "\n",
    "def gauss(z,z0,s):\n",
    "    return np.exp(-(z-z0)**2/(2*s**2)) / np.sqrt(2*np.pi) / s\n",
    "def plotTruth():\n",
    "    z = np.linspace(0,1,201)\n",
    "    uniform = N/len(bins)*(frac_bg)*np.ones_like(z)\n",
    "    cluster = N/len(bins)*gauss(z,z_cl,w_cl)\n",
    "    #plt.plot(z, uniform)\n",
    "    plt.plot(z, cluster)\n",
    "    plt.plot(z, uniform+cluster)\n",
    "\n",
    "N_bg = np.random.poisson(N*frac_bg)\n",
    "N_cl = N - N_bg\n",
    "bg = np.random.rand(N_bg)\n",
    "cl = np.random.normal(loc=z_cl, scale=w_cl, size=N_cl)\n",
    "sample = np.concatenate((bg,cl)).reshape(-1,1)\n",
    "plt.hist(sample, bins)\n",
    "plotTruth()\n",
    "plt.ylim(0,200)\n",
    "plt.xlim(0,1)\n",
    "print (\"True background amplitude: %.3f\" % (N_bg / N))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "We have a cluster at $z=0.35$ and a flat distribution of other galaxies. If we fit these data with one, two, or three Gaussian components, the results are, well ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import pygmmis\n",
    "# see the fitter work: iterations towards maximum likelihood \n",
    "import logging\n",
    "logging.basicConfig(format='%(message)s',level=logging.INFO)\n",
    "\n",
    "gmm = pygmmis.GMM(K=2, D=1)\n",
    "pygmmis.fit(gmm, sample)\n",
    "\n",
    "def plotMixture(gmm, background=None):\n",
    "    z = np.linspace(0,1,201)\n",
    "    if background is not None:\n",
    "        uniform = N/len(bins)*bg.amp*np.ones_like(z)\n",
    "    else:\n",
    "        uniform = 0\n",
    "    ps = [N/len(bins)*np.exp(gmm.logL_k(coords=z.reshape(-1,1),k=k)) for k in range(gmm.K)]\n",
    "    for k in range(gmm.K):\n",
    "        plt.plot(z,ps[k])\n",
    "    plt.plot(z, np.sum(ps, axis=0) + uniform)\n",
    "\n",
    "plt.hist(sample, bins)\n",
    "plotMixture(gmm)\n",
    "plt.ylim(0,200)\n",
    "plt.xlim(0,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "You'll get mostly bad fits, which is caused by rather poor initial parameter values for the GMM. With the default initialization, the means are chosen randomly from the samples, and the covariances are such that the components take $1/K$ of the volume of the space covered by the samples. In the case here, we know something about the components, so we could try a custom initialization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "# super-wide component\n",
    "gmm = pygmmis.GMM(K=2, D=1)\n",
    "gmm.amp[:] = 0.5\n",
    "gmm.mean[0][:] = 0.35\n",
    "gmm.mean[1][:] = 0.5\n",
    "gmm.covar[0][:] = 0.01\n",
    "gmm.covar[1][:] = 20\n",
    "pygmmis.fit(gmm, sample, init_method='none')\n",
    "plt.hist(sample, bins)\n",
    "plotMixture(gmm)\n",
    "plt.ylim(0,200)\n",
    "plt.xlim(0,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "Still not great. Even though we start at the right location with one component, it doesn't latch onto the cluster peak. Why? Because the second component is so wrong that the first one starts to cover for it. We can prevent that from happening by freezing the first components:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "# super-wide component\n",
    "gmm = pygmmis.GMM(K=2, D=1)\n",
    "gmm.amp[:] = 0.5\n",
    "gmm.mean[0][:] = 0.35\n",
    "gmm.mean[1][:] = 0.5\n",
    "gmm.covar[0][:] = w_cl*w_cl\n",
    "gmm.covar[1][:] = 20\n",
    "pygmmis.fit(gmm, sample, init_method='none', frozen={\"mean\": [0], \"covar\": [0]})\n",
    "plt.hist(sample, bins)\n",
    "plotMixture(gmm)\n",
    "plt.ylim(0,200)\n",
    "plt.xlim(0,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "Still, this was quite some fine-tuning, and the flat distribution still doesn't look right. Instead, we can make use of the equation for $q_{ik}$ above and make a more general mixture model, in which we add a constant background component. The only thing required here is that we can evaluate $p_\\mathrm{bg}(x_i)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "bg = pygmmis.Background((np.array([0]), np.array([1])), amp=0.5)\n",
    "gmm = pygmmis.GMM(K=1, D=1)\n",
    "pygmmis.fit(gmm, sample, background=bg)\n",
    "print (\"Background amplitude: %.3f\" % bg.amp)\n",
    "print (\"Gaussian component: %.3f +- %.3f\" % (gmm.mean[0], np.sqrt(gmm.covar[0])))\n",
    "plt.hist(sample, bins)\n",
    "plotMixture(gmm, background=bg)\n",
    "plt.ylim(0,200)\n",
    "plt.xlim(0,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "No fine-tuning and a pretty resonable result.\n",
    "\n",
    "Let's add noise to the samples..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "# add noise:\n",
    "std = 0.05\n",
    "disps = std * np.random.rand(N)\n",
    "noisy = (sample[:,0] + np.random.normal(scale=disps)).reshape(-1,1)\n",
    "plt.hist(sample, bins)\n",
    "plt.hist(noisy, bins, fc='None', ec='r')\n",
    "plt.title(\"Noisy data\")\n",
    "plt.xlim(0,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "gmm = pygmmis.GMM(K=1, D=1)\n",
    "bg = pygmmis.Background((np.array([0]), np.array([1])), amp=0.5)\n",
    "pygmmis.fit(gmm, noisy, background=bg)\n",
    "print (\"Background amplitude: %.3f\" % bg.amp)\n",
    "print (\"Gaussian component: %.3f +- %.3f\" % (gmm.mean[0], np.sqrt(gmm.covar[0])))\n",
    "plt.hist(sample, bins)\n",
    "plt.hist(noisy, bins, fc='None', ec='r')\n",
    "plt.title(\"Noisy data\")\n",
    "plotMixture(gmm, background=bg)\n",
    "plt.ylim(0,100)\n",
    "plt.xlim(0,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "The fit now picked up the broadening from the noise. [Bovy et al. (2011)](http://projecteuclid.org/euclid.aoas/1310562737) realized that the data samples could be noisy, $\\mathbf{x}_i \\rightarrow \\mathbf{x}_i + \\mathbf{e}_i$, with Gaussian errors drawn from independent distributions. This insight is a heteroscedastic extension of the known relation that the PDF of $a+b$ is given by the convolution $p(a)\\star p(b)$, which in the case of Gaussian has a width $\\mathbf{\\Sigma}_{a+b} = \\mathbf{\\Sigma}_a + \\mathbf{\\Sigma}_b$. In essence, if you want to get rid of the error, you need to deconvolve from its PDF, which boils down to subtracting the width of its PDF. That you can do this even when every sample has its own error gave rise to the name *Extreme Deconvolution* and to the popularity of the method. It can be used for fitting a very wide range of noisy data, or to fit a straight line in 2D when both $x​$ and $y​$ values have, potentially independent, errors.\n",
    "\n",
    "**Caveat:** For it to properly work, the noise PDFs need to be Gaussian, and the underlying data PDF should be reasonably describable be a GMM with a finite number of components. We will encounter this requirement again later on.\n",
    "\n",
    "pyGMMis has fully implemented the Extreme Deconvolution equations. All you need to do is to provide the covariance of each sample. That means, you have to know the errors for your data!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "gmm = pygmmis.GMM(K=1, D=1)\n",
    "bg = pygmmis.Background((np.array([0]), np.array([1])), amp=0.5)\n",
    "covar = (disps**2).reshape(-1,1,1)\n",
    "pygmmis.fit(gmm, noisy, covar=covar, background=bg)\n",
    "print (\"Background amplitude: %.3f\" % bg.amp)\n",
    "print (\"Gaussian component: %.3f +- %.3f\" % (gmm.mean[0], np.sqrt(gmm.covar[0])))\n",
    "plt.hist(sample, bins)\n",
    "plt.hist(noisy, bins, fc='None', ec='r')\n",
    "plt.title(\"Noisy data\")\n",
    "plotMixture(gmm, background=bg)\n",
    "plt.ylim(0,100)\n",
    "plt.xlim(0,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "Now let's finish up by adding one more concept. We will allow the sampling process to be incomplete, that means that there might be samples that are entirely missing. Think of cases where a foreground star blocks out all the light from neighboring objects, so that the catalog seems to have a big gap there.\n",
    "\n",
    "Here we'll reject a few samples behind the cluster, with a probability to increases with redshift. One can easily imagine that the cluster interferes with the measurement of galaxies in the background."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "def zrange(coords):\n",
    "    return (coords[:,0] >= 0) & (coords[:,0] <= 1)\n",
    "\n",
    "def ramp(coords):\n",
    "    return coords[:,0] - z_cl < np.random.rand(len(coords)) * (1-z_cl)\n",
    "\n",
    "def ramp_range(coords):\n",
    "    return ramp(coords) & zrange(coords)\n",
    "\n",
    "sel = ramp_range(noisy)\n",
    "data_s = noisy[sel]\n",
    "covar_s = covar[sel]\n",
    "plt.hist(noisy, bins)\n",
    "plt.hist(data_s, bins, fc='none', ec='r')\n",
    "plt.title('noisy and selected')\n",
    "plt.ylim(0,100)\n",
    "plt.xlim(0,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "gmm = pygmmis.GMM(K=1, D=1)\n",
    "bg = pygmmis.Background((np.array([0]), np.array([1])), amp=0.5)\n",
    "pygmmis.fit(gmm, data_s, covar=covar_s, background=bg)\n",
    "print (\"Background amplitude: %.3f\" % bg.amp)\n",
    "print (\"Gaussian component: %.3f +- %.3f\" % (gmm.mean[0], np.sqrt(gmm.covar[0])))\n",
    "plt.hist(sample, bins)\n",
    "plt.hist(data_s, bins, fc='None', ec='r')\n",
    "plotMixture(gmm, background=bg)\n",
    "plt.ylim(0,100)\n",
    "plt.xlim(0,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "If left uncorrected, the background amplitude is biased low, the cluster component is broadened to account for some of the samples that are in fact part of the background. If we run many samples, we could even see that the cluster component moves to lower redshift because even samples from the far side of the cluster occasionally get rejected.\n",
    "\n",
    "Here's where we can exploit again that we assume to know the parametric form of the density distribution. We can therefore create samples that are drawn from that distribution. If we now also know when samples are observed and when not, we can **augment the observed samples with an estimate of the missing ones**. In order to do that, you need to specify a function that returns True or False for every test sample. It's called `ramp_range` here, and it's exactly the function we've used before to get the \"selected\" sample.\n",
    "\n",
    "As a minor but important detail, if there's noise on the observed samples, you also have to define a function that gives an estimate of the noise in the missing samples. Because the mean of the noise dispersion is $0.05/2$, we'll simply use this value for every missing sample.\n",
    "\n",
    "The gory details of this augmentation (also called: \"imputation\") method are in [Melchior & Goulding (2016)](http://arxiv.org/abs/1611.05806).\n",
    "\n",
    "However, we ask more and more of the model: fitting multiple components, correcting for noise and incomplete sampling. In this example here, I didn't find a good solution with completely unconstrained, randomly initialized components. I therefore choose to start with a reasonable guess for the Gaussian component. If more samples were available, we could probably get away with completely random initializations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "def zerror(coords):\n",
    "    return (std/2)**2 * np.eye(coords.shape[1])\n",
    "\n",
    "gmm = pygmmis.GMM(K=1, D=1)\n",
    "gmm.amp[0] = 1\n",
    "gmm.mean[0,:] = z_cl\n",
    "gmm.covar[0,:,:] = 0.1\n",
    "bg = pygmmis.Background((np.array([0]), np.array([1])), amp=0.7)\n",
    "pygmmis.fit(gmm, data_s, init_method='none', covar=covar_s, background=bg, sel_callback=ramp_range, covar_callback=zerror)\n",
    "print (\"Background amplitude: %.3f\" % bg.amp)\n",
    "print (\"Gaussian component: %.3f +- %.3f\" % (gmm.mean[0], np.sqrt(gmm.covar[0])))\n",
    "plt.hist(sample, bins)\n",
    "plt.hist(data_s, bins, fc='None', ec='r')\n",
    "plotMixture(gmm, background=bg)\n",
    "plt.ylim(0,100)\n",
    "plt.xlim(0,1)"
   ]
  }
 ],
 "metadata": {
  "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}