{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Posterior sampling (MCMC)\n",
    "Instead—or on top—of running the ``Profiler`` to obtain credible intervals for the parameter estimates, we can also run an MCMC sampler to generate a proper posterior sample. This allows estimating the actual marginal distributions for each parameter.\n",
    "\n",
    "This tutorial provides an exemplatory implementation of an MCMC sampler for posterior sampling. We use the [sampler implementation provided by noctiluca](https://noctiluca.readthedocs.io/en/latest/noctiluca.util.html#noctiluca.util.mcmc.Sampler); of course you can adapt this to use your favorite MCMC implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from scipy import stats\n",
    "\n",
    "from noctiluca.util import mcmc\n",
    "import bayesmsd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up\n",
    "We generate an example dataset with 100 trajectories à 20 frames, with $\\text{MSD}(\\Delta t) = \\sqrt{\\Delta t}$ and fit it with a powerlaw MSD, hoping to recover the parameters $\\Gamma = 1$ and $\\alpha = 0.5$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(186875736)\n",
    "data = bayesmsd.gp.generate((bayesmsd.deco.MSDfun(lambda dt: np.sqrt(dt)), 1, 1), T=100, n=20)\n",
    "\n",
    "fit = bayesmsd.lib.NPXFit(data, ss_order=1)\n",
    "fit.parameters['log(σ²) (dim 0)'].fix_to = -np.inf\n",
    "fitres = fit.run(show_progress=True)\n",
    "\n",
    "# Print results\n",
    "for name, val in fitres['params'].items():\n",
    "    print(f\"{name:>15s} = {val:.5f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For reference, let us also run the ``Profiler``. This is just so that we can include the found credible intervals in the plots at the end; otherwise this will be ignored for the rest of the tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiler = bayesmsd.Profiler(fit)\n",
    "profiler.point_estimate = fitres\n",
    "mci = profiler.find_MCI(show_progress=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Writing an MCMC sampler\n",
    "As mentioned above, we rely on the base MCMC implementation provided by [noctiluca](https://noctiluca.readthedocs.io/en/latest). For basic usage, refer to its [reference page](https://noctiluca.readthedocs.io/en/latest/noctiluca.util.html#noctiluca.util.mcmc.Sampler).\n",
    "\n",
    "### Class definition and constructor\n",
    "```py\n",
    "class PosteriorSampler(mcmc.Sampler):\n",
    "    def __init__(self, fit, fitresult, stepsizes):\n",
    "        \"\"\"\n",
    "        stepsizes: {name : stepsize for name in <(independent) parameters>}\n",
    "        \"\"\"\n",
    "        self.fit = fit\n",
    "        if fitresult is None:\n",
    "            self.initial_params = fit.initial_params()\n",
    "        else:\n",
    "            self.initial_params = fitresult['params']\n",
    "        self.stepsizes = stepsizes\n",
    "        \n",
    "        self.min_target_from_fit = self.fit.MinTarget(self.fit)\n",
    "        self.independent_parameters = self.min_target_from_fit.param_names\n",
    "...\n",
    "```\n",
    "We initialize the sampler from a ``Fit`` instance, which holds the data and defines the posterior (likelihood * prior) that we wish to sample. For the most part it makes sense to actually run the fit and then use the MCMC to sample the posterior around the found point estimate. However, you can of course also skip the fit altogether and just start the MCMC from some initial parameter guess; incidentally, ``Fit`` provides such an initial estimate as well.\n",
    "\n",
    "### Implementing the MCMC scheme\n",
    "The ``mcmc.Sampler`` base class takes care of running the MCMC, once we define the basic scheme. This is done by overriding the methods ``logL()`` (which defines the likelihood function) and ``propose_update()`` (which defines the proposal steps).\n",
    "```py\n",
    "...\n",
    "    def logL(self, params):\n",
    "        params_array = self.min_target_from_fit.params_dict2array(params)\n",
    "        return -self.min_target_from_fit(params_array)\n",
    "...\n",
    "```\n",
    "This one is pretty straight-forward: the ``min_target_from_fit`` is essentially the likelihood function (up to a sign, because the fit *min*imizes). So let's move on to the proposal step:\n",
    "```py\n",
    "...\n",
    "    def propose_update(self, current_params):\n",
    "        name = np.random.choice(self.independent_parameters)\n",
    "        new_params = deepcopy(current_params)\n",
    "        step_dist = stats.norm(scale=self.stepsizes[name])\n",
    "        step = step_dist.rvs()\n",
    "        new_params[name] += step\n",
    "\n",
    "        return new_params, step_dist.logpdf(step), step_dist.logpdf(-step)\n",
    "...\n",
    "```\n",
    "This is also not very complicated: we randomly pick a parameter to update and then take a normally distributed step. One could also perform a multivariate update step here (i.e. pick all parameters anew).\n",
    "\n",
    "The ``propose_update()`` method should always return the new parameters and the forward/backward probabilities; the latter are necessary to ensure that the MCMC converges to the correct limiting distribution. Note, however, how in this example the two values will always be identical, since the Gaussian proposal distribution ``step_dist`` is symmetric. In that case we could also omit the evaluation of the distribution altogether and just ``return new_params, 1, 1``.\n",
    "\n",
    "### Additional functionality\n",
    "To make the code from this example directly useable for real applications, we add some convenience.\n",
    "\n",
    "One great way to use the MCMC sampler is to first run the profiler to get an idea for the width of the posterior in parameter space, and only then move on to the MCMC to get an estimate of the actual marginals. This has the advantage that we can skip the burn-in period of the MCMC sampler, making the procedure one bit less dependent on heuristics. To facilitate this workflow, use the below for initalization from a ``Profiler`` and its results (``mci``):\n",
    "```py\n",
    "...\n",
    "    @classmethod\n",
    "    def fromProfiler(cls, profiler, mci):\n",
    "        stepsizes = {name : np.min(np.abs(ci - m)) for name, (m, ci) in mci.items()}\n",
    "        return cls(profiler.fit,\n",
    "                   profiler.point_estimate,\n",
    "                   stepsizes,\n",
    "                  )\n",
    "...\n",
    "```\n",
    "\n",
    "Finally, the ``run()`` method of ``mcmc.Sampler`` usually expects initial values from the user. But we know good starting points, so for convenience we override as follows:\n",
    "```py\n",
    "...\n",
    "    def run(self):\n",
    "        return super().run(self.initial_params)\n",
    "```\n",
    "\n",
    "### Full implementation\n",
    "In summary, we obtain the following implementation of the ``PosteriorSampler``:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PosteriorSampler(mcmc.Sampler):\n",
    "    def __init__(self, fit, fitresult, stepsizes):\n",
    "        \"\"\"\n",
    "        stepsizes: {name : stepsize for name in <(independent) parameters>}\n",
    "        \"\"\"\n",
    "        self.fit = fit\n",
    "        if fitresult is None:\n",
    "            self.initial_params = fit.initial_params()\n",
    "        else:\n",
    "            self.initial_params = fitresult['params']\n",
    "        self.stepsizes = stepsizes\n",
    "        \n",
    "        self.min_target_from_fit = self.fit.MinTarget(self.fit)\n",
    "        self.independent_parameters = self.min_target_from_fit.param_names\n",
    "        \n",
    "    @classmethod\n",
    "    def fromProfiler(cls, profiler, mci):\n",
    "        stepsizes = {name : np.min(np.abs(ci - m)) for name, (m, ci) in mci.items()}\n",
    "        return cls(profiler.fit,\n",
    "                   profiler.point_estimate,\n",
    "                   stepsizes,\n",
    "                  )\n",
    "    \n",
    "    def logL(self, params):\n",
    "        params_array = self.min_target_from_fit.params_dict2array(params)\n",
    "        return -self.min_target_from_fit(params_array)\n",
    "    \n",
    "    def propose_update(self, current_params):\n",
    "        name = np.random.choice(self.independent_parameters)\n",
    "        new_params = deepcopy(current_params)\n",
    "        step_dist = stats.norm(scale=self.stepsizes[name])\n",
    "        step = step_dist.rvs()\n",
    "        new_params[name] += step\n",
    "\n",
    "        return new_params, step_dist.logpdf(step), step_dist.logpdf(-step)\n",
    "        \n",
    "    def run(self):\n",
    "        return super().run(self.initial_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the sampler\n",
    "The posterior sampler can now be used through the interface provided by ``mcmc.Sampler``. The unfamiliar reader can consult the relevant [documentation](https://noctiluca.readthedocs.io/en/latest/noctiluca.util.html#noctiluca.util.mcmc.Sampler); but for the most part the usage is straight-forward: we initialize the ``PosteriorSampler``, configure it to our liking and finally let it run.\n",
    "\n",
    "Note that here we initialize the MCMC just from the fit we ran [above](#Set-up) and guess\\* good step sizes for the two free parameters. Alternatively we could initialize the stepsizes from the measured credible intervals; for this example, however, let's keep MCMC and Profiler independent, such that we can show a fair comparison below.\n",
    "\n",
    "(\\*) by \"guess\" we mean that these are chosen such that the acceptance rate of the sampler is about 50%."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = PosteriorSampler(fit, fitres, stepsizes={'log(Γ) (dim 0)' : 0.05, 'α (dim 0)' : 0.05})\n",
    "sampler.configure(iterations = 1000,\n",
    "                  burn_in = 100,\n",
    "                  log_every = 100,\n",
    "                  show_progress=True,\n",
    "                 )\n",
    "mcmcrun = sampler.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot results\n",
    "fig, axs = plt.subplots(1, 2, figsize=[10, 4])\n",
    "\n",
    "for (ax,\n",
    "     name,\n",
    "     color,\n",
    "     true_value,\n",
    "    ) in zip(axs,\n",
    "             ['log(Γ)', 'α'],\n",
    "             ['slateblue', 'tab:blue'],\n",
    "             [0, 0.5],\n",
    "            ):\n",
    "    \n",
    "    samples = np.array([params[name+' (dim 0)'] for params in mcmcrun.samples])\n",
    "    ax.hist(samples, bins='auto', density=True, color=color, label='posterior')\n",
    "    \n",
    "    ax.axvline(true_value, color='k', label='true value')\n",
    "    \n",
    "    m, ci = mci[name+' (dim 0)']\n",
    "    ax.axvline(m, color='tab:red', label='point estimate')\n",
    "    \n",
    "    ax.set_ylim([0, None])\n",
    "    _, ytop = ax.get_ylim()\n",
    "    yci = 0.33*ytop\n",
    "    ax.errorbar(m, yci, xerr=np.abs(ci-m)[:, None],\n",
    "                color='tab:orange', label='95% CI\\nfrom profiler',\n",
    "                capsize=10, capthick=2,\n",
    "                zorder=3,\n",
    "               )\n",
    "\n",
    "    ax.legend()\n",
    "    ax.set_xlabel(name)\n",
    "    ax.set_ylabel('density')\n",
    "    ax.set_title('Marginal posterior for '+name)\n",
    "\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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}