{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Univariate Density Estimation via Dirichlet Process Mixture"
   ]
  },
  {
   "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": [
    "## Data Generation\n",
    "\n",
    "We generate data from a two-component mixture model\n",
    "$$\n",
    "y_i \\sim \\frac{1}{2} \\mathcal N(-3, 1) + \\frac{1}{2} \\mathcal N(3, 1), \\quad i=1, \\ldots, 200\n",
    "$$"
   ]
  },
  {
   "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])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = sample_from_mixture(\n",
    "    np.array([0.5, 0.5]), np.array([-3, 3]), np.array([1, 1]), 200)\n",
    "plt.hist(y)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The statistical model\n",
    "\n",
    "We assume the following model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\\begin{equation}\n",
    "\\begin{aligned}\n",
    "y_i | \\tilde{p} &\\sim f(\\cdot) = \\int_{R \\times R^+} \\mathcal{N}(\\cdot | \\mu, \\sigma^2) \\tilde{p}(d\\mu, d\\sigma^2) \\\\\n",
    "\\tilde{p} &\\sim DP(\\alpha, G_0)\n",
    "\\end{aligned}\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "where $DP(\\alpha, G_0)$ is the Dirichlet Process with base measure $\\alpha G_0$. \n",
    "\n",
    "Given the stick-breaking represetation of the Dirichlet Process, the model is equivalently written as"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\\begin{equation}\n",
    "\\begin{aligned}\n",
    "y_i | \\{w_h\\}_h \\{(\\mu_h, \\sigma^2_h)\\}_h & \\sim f(\\cdot) = \\sum_{h=1}^\\infty w_h \\mathcal{N}(\\cdot | \\mu_h, \\sigma_h^2) \\\\\n",
    "\\{w_h\\}_h &\\sim GEM(\\alpha) \\\\\n",
    " \\{(\\mu_h, \\sigma^2_h)\\}_h &\\sim G_0 \\\\\n",
    "\\end{aligned}\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In pybmix we take advantage of the second representation, and specify a MixtureModel in terms of a Mixing and a Hierarchy. The Mixing is the prior for the weights, while the Hierarchy combines the base measure $G_0$ with the kernel of the mixture (in this case, the univariate Gaussian distribution)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we assume that $\\alpha = 5$ and $G_0(d\\mu, d\\sigma^2) = \\mathcal N(d\\mu | \\mu_0, \\lambda \\sigma^2) \\times IG(d\\sigma^2 | a, b)$, i.e., $G_0$ is a normal-inverse gamma distribution. \n",
    "\n",
    "The parameters $(\\mu_0, \\lambda, a , b)$ of $G_0$ can be set automatically by the method 'make_default_fixed_params' which takes as input the observations and a \"guess\" on the number of clusters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mixing = DirichletProcessMixing(total_mass=5)\n",
    "hierarchy = UnivariateNormal()\n",
    "hierarchy.make_default_fixed_params(y, 2)\n",
    "mixture = MixtureModel(mixing, hierarchy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run MCMC simulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mixture.run_mcmc(y, algorithm=\"Neal2\", niter=2000, nburn=1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get the density estimates\n",
    "\n",
    "1) fix a grid where to estimate the densities\n",
    "\n",
    "2) the method 'estimate_density' returns a matrix of shape [niter - nburn, len(grid)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pybmix.estimators.density_estimator import DensityEstimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid = np.linspace(-6, 6, 500)\n",
    "dens_est = DensityEstimator(mixture)\n",
    "densities = dens_est.estimate_density(grid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot some of the densities and their mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(y, density=True)\n",
    "plt.plot(grid, np.mean(densities, axis=0), lw=3, label=\"predictive density\")\n",
    "idxs = [5, 100, 300]\n",
    "for idx in idxs:\n",
    "    plt.plot(grid, densities[idx, :], \"--\", label=\"iteration: {0}\".format(idx))\n",
    "    \n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "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
}
