{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A sampler study"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we perform a short study of how various samplers implemented in pyPESTO perform."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# install if not done yet\n",
    "# !apt install libatlas-base-dev swig\n",
    "# %pip install pypesto[amici,petab,pymc,emcee] --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we show a typical workflow, fully integrating the samplers with a [PEtab](https://github.com/petab-dev/petab) problem, using a toy example of a conversion reaction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import petab\n",
    "\n",
    "import pypesto\n",
    "import pypesto.optimize as optimize\n",
    "import pypesto.petab\n",
    "import pypesto.sample as sample\n",
    "import pypesto.visualize as visualize\n",
    "\n",
    "np.random.seed(0)\n",
    "\n",
    "# import to petab\n",
    "petab_problem = petab.Problem.from_yaml(\n",
    "    \"conversion_reaction/conversion_reaction.yaml\"\n",
    ")\n",
    "# import to pypesto\n",
    "importer = pypesto.petab.PetabImporter(petab_problem)\n",
    "# create problem\n",
    "problem = importer.create_problem(verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "problem.x_names"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Commonly, as a first step, optimization is performed, in order to find good parameter point estimates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = optimize.minimize(problem, n_starts=10, filename=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize.waterfall(result, size=(4, 4));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we perform sampling. Here, we employ a `pypesto.sample.AdaptiveParallelTemperingSampler` sampler, which runs Markov Chain Monte Carlo (MCMC) chains on different temperatures. For each chain, we employ a `pypesto.sample.AdaptiveMetropolisSampler`. For more on the samplers see below or the API documentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = sample.AdaptiveParallelTemperingSampler(\n",
    "    internal_sampler=sample.AdaptiveMetropolisSampler(),\n",
    "    n_chains=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the actual sampling, we call the `pypesto.sample.sample` function. By passing the result object to the function, the previously found global optimum is used as starting point for the MCMC sampling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "result = sample.sample(\n",
    "    problem, n_samples=1000, sampler=sampler, result=result, filename=None\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When the sampling is finished, we can analyse our results. A first thing to do is to analyze the sampling burn-in:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample.geweke_test(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "pyPESTO provides functions to analyse both the sampling process as well as the obtained sampling result. Visualizing the traces e.g. allows to detect burn-in phases, or fine-tune hyperparameters. First, the parameter trajectories  can be visualized:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = visualize.sampling_parameter_traces(result, use_problem_bounds=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, also the log posterior trace can be visualized:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = visualize.sampling_fval_traces(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To visualize the result, there are various options. The scatter plot shows histograms of 1-dim parameter marginals and scatter plots of 2-dimensional parameter combinations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = visualize.sampling_scatter(result, size=[13, 6])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`sampling_1d_marginals` allows to plot e.g. kernel density estimates or histograms (internally using [seaborn](https://seaborn.pydata.org/)):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i_chain in range(len(result.sample_result.betas)):\n",
    "    visualize.sampling_1d_marginals(\n",
    "        result, i_chain=i_chain, suptitle=f\"Chain: {i_chain}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it for the moment on using the sampling pipeline."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1-dim test problem"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To compare and test the various implemented samplers, we first study a 1-dimensional test problem of a gaussian mixture density, together with a flat prior."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "from scipy.stats import multivariate_normal\n",
    "\n",
    "import pypesto\n",
    "import pypesto.sample as sample\n",
    "import pypesto.visualize as visualize\n",
    "\n",
    "\n",
    "def density(x):\n",
    "    return 0.3 * multivariate_normal.pdf(\n",
    "        x, mean=-1.5, cov=0.1\n",
    "    ) + 0.7 * multivariate_normal.pdf(x, mean=2.5, cov=0.2)\n",
    "\n",
    "\n",
    "def nllh(x):\n",
    "    return -np.log(density(x))\n",
    "\n",
    "\n",
    "objective = pypesto.Objective(fun=nllh)\n",
    "problem = pypesto.Problem(objective=objective, lb=-4, ub=5, x_names=[\"x\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The likelihood has two separate modes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xs = np.linspace(-4, 5, 100)\n",
    "ys = [density(x) for x in xs]\n",
    "\n",
    "ax = sns.lineplot(x=xs, y=ys, color=\"C1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Metropolis sampler"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this problem, let us try out the simplest sampler, the `pypesto.sample.MetropolisSampler`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = sample.MetropolisSampler({\"std\": 0.5})\n",
    "result = sample.sample(\n",
    "    problem, 1e3, sampler, x0=np.array([0.5]), filename=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample.geweke_test(result)\n",
    "ax = visualize.sampling_1d_marginals(result)\n",
    "ax[0][0].plot(xs, ys)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The obtained posterior does not accurately represent the distribution, often only capturing one mode. This is because it is hard for the Markov chain to jump between the distribution's two modes. This can be fixed by choosing a higher proposal variation `std`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = sample.MetropolisSampler({\"std\": 1})\n",
    "result = sample.sample(\n",
    "    problem, 1e3, sampler, x0=np.array([0.5]), filename=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample.geweke_test(result)\n",
    "ax = visualize.sampling_1d_marginals(result)\n",
    "ax[0][0].plot(xs, ys)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In general, MCMC have difficulties exploring multimodel landscapes. One way to overcome this is to used parallel tempering. There, various chains are run, lifting the densities to different temperatures. At high temperatures, proposed steps are more likely to get accepted and thus jumps between modes are more likely.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parallel tempering sampler"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In pyPESTO, the most basic parallel tempering algorithm is the `pypesto.sample.ParallelTemperingSampler`. It takes an `internal_sampler` parameter, to specify what sampler to use for performing sampling the different chains. Further, we can directly specify what inverse temperatures `betas` to use. When not specifying the `betas` explicitly, but just the number of chains `n_chains`, an established near-exponential decay scheme is used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = sample.ParallelTemperingSampler(\n",
    "    internal_sampler=sample.MetropolisSampler(), betas=[1, 1e-1, 1e-2]\n",
    ")\n",
    "result = sample.sample(\n",
    "    problem, 1e3, sampler, x0=np.array([0.5]), filename=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample.geweke_test(result)\n",
    "for i_chain in range(len(result.sample_result.betas)):\n",
    "    visualize.sampling_1d_marginals(\n",
    "        result, i_chain=i_chain, suptitle=f\"Chain: {i_chain}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Of interest is here finally the first chain at index `i_chain=0`, which approximates the posterior well."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adaptive Metropolis sampler"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The problem of having to specify the proposal step variation manually can be overcome by using the `pypesto.sample.AdaptiveMetropolisSampler`, which iteratively adjusts the proposal steps to the function landscape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = sample.AdaptiveMetropolisSampler()\n",
    "result = sample.sample(\n",
    "    problem, 1e3, sampler, x0=np.array([0.5]), filename=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample.geweke_test(result)\n",
    "ax = visualize.sampling_1d_marginals(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adaptive parallel tempering sampler"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `pypesto.sample.AdaptiveParallelTemperingSampler` iteratively adjusts the temperatures to obtain good swapping rates between chains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = sample.AdaptiveParallelTemperingSampler(\n",
    "    internal_sampler=sample.AdaptiveMetropolisSampler(), n_chains=3\n",
    ")\n",
    "result = sample.sample(\n",
    "    problem, 1e3, sampler, x0=np.array([0.5]), filename=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample.geweke_test(result)\n",
    "for i_chain in range(len(result.sample_result.betas)):\n",
    "    visualize.sampling_1d_marginals(\n",
    "        result, i_chain=i_chain, suptitle=f\"Chain: {i_chain}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.sample_result.betas"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pymc sampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pypesto.sample.pymc import PymcSampler\n",
    "\n",
    "sampler = PymcSampler()\n",
    "result = sample.sample(\n",
    "    problem, 1e3, sampler, x0=np.array([0.5]), filename=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample.geweke_test(result)\n",
    "for i_chain in range(len(result.sample_result.betas)):\n",
    "    visualize.sampling_1d_marginals(\n",
    "        result, i_chain=i_chain, suptitle=f\"Chain: {i_chain}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If not specified, pymc chooses an adequate sampler automatically. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Emcee sampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = sample.EmceeSampler(nwalkers=10, run_args={\"progress\": True})\n",
    "result = sample.sample(\n",
    "    problem, int(1e3), sampler, x0=np.array([0.5]), filename=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.array([sampler.sampler.get_log_prob(flat=True)]).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample.geweke_test(result)\n",
    "for i_chain in range(len(result.sample_result.betas)):\n",
    "    visualize.sampling_1d_marginals(\n",
    "        result, i_chain=i_chain, suptitle=f\"Chain: {i_chain}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### dynesty sampler"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The [dynesty](https://dynesty.readthedocs.io/en/stable/index.html) package provides nested and dynamic nested samplers. These differ from some of the other samplers that pyPESTO interfaces with. For example, it doesn't make sense to request a certain number of samples. Instead, the sampler runs until stopping criteria have been met."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = sample.DynestySampler(objective_type=\"negloglike\")\n",
    "result = sample.sample(\n",
    "    problem=problem,\n",
    "    n_samples=None,\n",
    "    sampler=sampler,\n",
    "    filename=None,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another difference is that there is no burn-in so, unlike for some other pyPESTO sampler results, the Geweke test is not applied here, and the chain will not appear to be converged. However, by default, pyPESTO returns samples that have been resampled to appear like converged chains from MCMC sampling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize.sampling_1d_marginals(result)\n",
    "visualize.sampling_fval_traces(result, full_trace=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The internal `dynesty` sampler can be saved and restored, for post-sampling analysis. For example, pyPESTO stores resampled MCMC-like samples from the `dynesty` sampler by default. The following code shows how to save and load the internal dynesty sampler, to facilitate post-sampling analysis of both the resampled and original chains. N.B.: when working across different computers, you might prefer to work with the raw sample results via `pypesto.sample.dynesty.save_raw_results` and `load_raw_results`.\n",
    "First, we save the internal sampler."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler.save_internal_sampler(\"dynesty.dill\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we load the internal sampler into some `DynestySampler` object, then set the `sample_result` of some pyPESTO result to the original samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dummy_sampler = sample.DynestySampler()\n",
    "dummy_sampler.restore_internal_sampler(\"dynesty.dill\")\n",
    "dummy_result = pypesto.Result(problem=problem)\n",
    "\n",
    "dummy_result.sample_result = dummy_sampler.get_original_samples()\n",
    "visualize.sampling_1d_marginals(dummy_result)\n",
    "visualize.sampling_fval_traces(dummy_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then set the `sample_result` of the `dummy_result` back to MCMC-like resampled samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dummy_result.sample_result = dummy_sampler.get_samples()\n",
    "visualize.sampling_1d_marginals(dummy_result)\n",
    "visualize.sampling_fval_traces(dummy_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This sampler supports parallelization, see the dynesty documentation for more details: [https://dynesty.readthedocs.io](https://dynesty.readthedocs.io)\n",
    "\n",
    "An example of doing this via pyPESTO is given below\n",
    "\n",
    "    import multiprocessing\n",
    "    import pypesto.sample as sample\n",
    "\n",
    "    if __name__ == '__main__':\n",
    "        with multiprocessing.Manager() as manager:\n",
    "            with manager.Pool() as pool:\n",
    "                sampler_args = {\n",
    "                    'pool': pool,\n",
    "                    'queue_size': multiprocessing.cpu_count(),\n",
    "                }\n",
    "                sampler = sample.DynestySampler(sampler_args=sampler_args)\n",
    "                result = sample.sample(...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2-dim test problem: Rosenbrock banana"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The adaptive parallel tempering sampler with chains running adaptive Metropolis samplers is also able to sample from more challenging posterior distributions. To illustrates this shortly, we use the Rosenbrock function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.optimize as so\n",
    "\n",
    "import pypesto\n",
    "\n",
    "# first type of objective\n",
    "objective = pypesto.Objective(fun=so.rosen)\n",
    "\n",
    "dim_full = 4\n",
    "lb = -5 * np.ones((dim_full, 1))\n",
    "ub = 5 * np.ones((dim_full, 1))\n",
    "\n",
    "problem = pypesto.Problem(objective=objective, lb=lb, ub=ub)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = sample.AdaptiveParallelTemperingSampler(\n",
    "    internal_sampler=sample.AdaptiveMetropolisSampler(), n_chains=10\n",
    ")\n",
    "result = sample.sample(\n",
    "    problem, 2e3, sampler, x0=np.zeros(dim_full), filename=None\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample.geweke_test(result)\n",
    "ax = visualize.sampling_scatter(result)\n",
    "ax = visualize.sampling_1d_marginals(result)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}