{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prior elicitation\n",
    "\n",
    "What is the meaning of the parameters in the various processes? Let's study the prior distribution of the number of clusters!"
   ]
  },
  {
   "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, PitmanYorMixing, StickBreakMixing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dirichlet Process\n",
    "\n",
    "Compare how the prior distribution varies for a DP for different values of the 'total_mass' parameter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nsamples = 100\n",
    "grid = np.arange(1, 30)\n",
    "\n",
    "for total_mass in [0.25, 1, 2, 3, 5, 7]:\n",
    "    dp = DirichletProcessMixing(total_mass=total_mass)\n",
    "    probs = dp.prior_cluster_distribution(grid, nsamples)\n",
    "    plt.plot(grid, probs, label=\"total_mass={0}\".format(total_mass))\n",
    "    \n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pitman-Yor Process\n",
    "\n",
    "The PY process is controlled by two parameters: 'strength' (plays the same role of 'total_mass' in the Dirichlet Process) and 'discount'.\n",
    "For a fixed value of 'strenght', see how the distribution varies for different values of 'discount'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "strength = 1\n",
    "for discount in [0.001, 0.1, 0.2, 0.5]:\n",
    "    py = PitmanYorMixing(strength=strength, discount=discount)\n",
    "    probs = py.prior_cluster_distribution(grid, nsamples)\n",
    "    plt.plot(grid, probs, label=\"strength={0}, discount={1}\".format(strength, discount))\n",
    "    \n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stick Breaking prior"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Truncation of a Dirichlet Process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = [50, 1000]\n",
    "H = [10, 20, 50]\n",
    "grid = np.arange(1, 50)\n",
    "\n",
    "fig, axes = plt.subplots(nrows=len(n), ncols=len(H), figsize=(15, 10))\n",
    "for i, nval in enumerate(n):\n",
    "    for j, hval in enumerate(H):\n",
    "        for total_mass in [0.25, 1, 2, 3, 5, 7]:\n",
    "            sb = StickBreakMixing(n_comp=hval, strength=total_mass)\n",
    "            probs = sb.prior_cluster_distribution(grid, nval, mc_iter=10000)\n",
    "            axes[i][j].plot(grid, probs, label=\"total_mass={0}\".format(total_mass))\n",
    "            axes[i][j].set_title(\"N: {0}, H: {1}\".format(nval, hval), fontsize=18)\n",
    "\n",
    "axes[1, 1].legend(ncol=3, fontsize=14, bbox_to_anchor=(1.5, -0.1))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Truncation of a Pitman-Yor Process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = [200, 1000]\n",
    "H = [10, 20, 50]\n",
    "strength = 1\n",
    "discount = np.linspace(0.2, 0.9, 4)\n",
    "\n",
    "fig, axes = plt.subplots(nrows=len(n), ncols=len(H), figsize=(15, 10))\n",
    "\n",
    "for i, nval in enumerate(n):\n",
    "    for j, hval in enumerate(H):\n",
    "        for d in discount:\n",
    "            sb = StickBreakMixing(n_comp=hval, strength=strength, discount=d)\n",
    "            probs = sb.prior_cluster_distribution(grid, nval, mc_iter=10000)\n",
    "            axes[i][j].plot(grid, probs, label=\"discount={0:.2f}\".format(d))\n",
    "            axes[i][j].set_title(\"N: {0}, H: {1}\".format(nval, hval), fontsize=18)\n",
    "\n",
    "axes[1, 1].legend(ncol=2, fontsize=14, bbox_to_anchor=(1.2, -0.1))\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
}
