{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Clustering of univariate data via Dirichlet Process Mixture"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "this is a continuation of 'estimate_univ_density'. Make sure to check it before going through this tutorial!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from pybmix.core.mixing import DirichletProcessMixing, StickBreakMixing\n",
    "from pybmix.core.hierarchy import UnivariateNormal\n",
    "from pybmix.core.mixture_model import MixtureModel\n",
    "\n",
    "np.random.seed(2021)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DP and clustering\n",
    "\n",
    "Recall that $\\tilde p \\sim DP(\\alpha, G_0)$ means that $\\tilde p = \\sum_{h=1}^\\infty w_h \\delta_{\\tau_h}$ with $\\{w_h\\}_h \\sim GEM(\\alpha)$ and $\\{\\tau_h\\}_h \\sim G_0$. Hence, realizations from a DP are almost surely discrete probability measures.\n",
    "\n",
    "Hence, sampling "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\\begin{equation}\n",
    "\\begin{aligned}\n",
    "\\theta_1, \\ldots, \\theta_n | \\tilde{p} & \\sim \\tilde{p} \\\\\n",
    "\\tilde{p} &\\sim DP(\\alpha, G_0)\n",
    "\\end{aligned}\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "entails that with positive probability $\\theta_i = \\theta_j$ (with $i \\neq j$). In a sample of size $n$ there will be $k \\geq n$ unique values $\\theta^*_1, \\ldots, \\theta^*_k$ among the $\\theta_i$'s and clusters are defined as $C_j = \\{i : \\theta_i = \\theta^*_j \\}$.\n",
    "\n",
    "When considering a mixture model, the $\\theta_i$'s are not observations but latent variables. In the case of a univariate normal mizture models, $\\theta_i = (\\mu_i, \\sigma^2_i)$ and the model can be written as"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\\begin{equation}\n",
    "\\begin{aligned}\n",
    "    y_i | \\theta_i = (\\mu_i, \\sigma^2_i) &\\sim \\mathcal N(\\mu_i, \\sigma^2_i) \\\\\n",
    "    \\theta_1, \\ldots, \\theta_n | \\tilde{p} &\\sim \\tilde{p} \\\\\n",
    "    \\tilde{p} &\\sim DP(\\alpha, G_0)\n",
    "\\end{aligned}\n",
    "\\end{equation}\n",
    "\n",
    "and the clustering among the observations $y_i$'s is inherited by the clustering among the $\\theta_i$'s."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's go back to the previous example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_from_mixture(weigths, means, sds, n_data):\n",
    "    n_comp = len(weigths)\n",
    "    clus_alloc = np.random.choice(np.arange(n_comp), p=[0.5, 0.5], size=n_data)\n",
    "    return np.random.normal(loc=means[clus_alloc], scale=sds[clus_alloc])\n",
    "\n",
    "y = sample_from_mixture(\n",
    "    np.array([0.5, 0.5]), np.array([-3, 3]), np.array([1, 1]), 200)\n",
    "\n",
    "mixing = DirichletProcessMixing(total_mass=5)\n",
    "hierarchy = UnivariateNormal()\n",
    "hierarchy.make_default_fixed_params(y, 2)\n",
    "mixture = MixtureModel(mixing, hierarchy)\n",
    "\n",
    "mixture.run_mcmc(y, algorithm=\"Neal2\", niter=2000, nburn=1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can extract the cluster allocation MCMC chain very easily"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mcmc_chain = mixture.get_chain()\n",
    "cluster_alloc_chain = mcmc_chain.extract(\"cluster_allocs\")\n",
    "print(cluster_alloc_chain.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "cluster_alloc_chain is a matrix of shape [niter - nburn, ndata]. \n",
    "\n",
    "To get the posterior distribution of the number of clusters, we count in each row the number of unique values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_clust_chain = np.apply_along_axis(lambda x: len(np.unique(x)), 1, \n",
    "                                    cluster_alloc_chain)\n",
    "\n",
    "fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(8, 4))\n",
    "axes[0].vlines(np.arange(len(n_clust_chain)), n_clust_chain - 0.3, n_clust_chain + 0.3)\n",
    "axes[0].set_title(\"Traceplot\")\n",
    "\n",
    "clusgrid = np.arange(1, 10)\n",
    "probas = np.zeros_like(clusgrid)\n",
    "for i, c in enumerate(clusgrid):\n",
    "    probas[i] = np.sum(n_clust_chain == c)\n",
    "\n",
    "probas = probas / np.sum(probas)\n",
    "axes[1].bar(clusgrid, probas)\n",
    "axes[1].set_xticks(clusgrid)\n",
    "axes[1].set_title(\"Posterior number of clusters\")\n",
    "    \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's inspect two iterations: the first one and the last one, and look at the cluster allocations of the first 5 observations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"First iteration: \", cluster_alloc_chain[0][:5])\n",
    "print(\"Last iteration: \", cluster_alloc_chain[-1][:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Observe that the clustering are identicals: the one is made of observations $\\{1, 2, 5\\}$ and the other cluster of observations $\\{3, 4\\}$. However the labels associated to each cluster are differend depending on the iterations: in the first iteration, $\\{1, 2, 5\\}$ are the first cluster (0th cluster) and  $\\{3, 4\\}$ are the second cluster, while in the last iteration the opposite happens.\n",
    "\n",
    "This is due to the so-called \"label-switching\". Usually to interpret the clustering result, a suitable point-estimate is chosen to minimize a loss function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pybmix.estimators.cluster_estimator import ClusterEstimator\n",
    "\n",
    "clus_est = ClusterEstimator(mixture)\n",
    "best_clust = clus_est.get_point_estimate()\n",
    "\n",
    "plt.hist(y, density=True, alpha=0.3)\n",
    "for cluster_idx in clus_est.group_by_cluster(best_clust):\n",
    "    data = y[cluster_idx]\n",
    "    plt.scatter(data, np.zeros_like(data) + 5e-3)\n",
    "    \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note how the posterior mode of the number of clusters is 3, but the point estimate for the best clustering consists of 2 clusters"
   ]
  }
 ],
 "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
