{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Specifying priors\n",
    "\n",
    "This tutorial demonstrates how to specify parameter priors, including:\n",
    "\n",
    "* uniform and log-uniform distributions\n",
    "* gaussian and more complicated distributions\n",
    "* multi-dimensional priors (not factorized)\n",
    "* non-analytic priors\n",
    "* priors on fractions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.stats\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cumulative prior distributions\n",
    "\n",
    "Any 1-dimensional probability distribution is normalised so that its integral is 1. That is, the cumulative distribution goes from 0 to 1. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gaussdistribution = scipy.stats.norm(2, 0.3)\n",
    "uniformdistribution = scipy.stats.uniform(3.5, 1.2)\n",
    "x = np.linspace(0, 5, 400)\n",
    "plt.figure()\n",
    "plt.plot(x, gaussdistribution.pdf(x), '--', label='density (Gauss)')\n",
    "plt.plot(x, gaussdistribution.cdf(x), label='cumulative (Gauss)')\n",
    "plt.plot(x, uniformdistribution.pdf(x), '--', label='density (uniform)')\n",
    "plt.plot(x, uniformdistribution.cdf(x), label='cumulative (uniform)')\n",
    "plt.ylabel('Probability')\n",
    "plt.xlabel('Model parameter x')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transforming from the unit interval\n",
    "\n",
    "We invert the cumulative probability distribution mapping quantiles (0...1) to the corresponding model parameter value.\n",
    "\n",
    "Lets start with the uniform distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform_1d_uniform(quantile):\n",
    "    lower_bound = 3.5\n",
    "    width = 1.2\n",
    "    return quantile * width + lower_bound\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Scipy provides the inverse cumulative probability distributions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform_1d(quantile):\n",
    "    return gaussdistribution.ppf(quantile)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "UltraNest samples from the unit interval to obtain prior samples. Lets try drawing a few examples from our function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uniform_samples = transform_1d_uniform(np.random.uniform(0, 1, size=100000))\n",
    "gauss_samples = transform_1d(np.random.uniform(0, 1, size=100000))\n",
    "plt.figure()\n",
    "plt.hist(uniform_samples, bins=20, histtype='step', density=True, label='Uniform')\n",
    "plt.hist(gauss_samples, bins=100, histtype='step', density=True, label='Gauss')\n",
    "plt.xlabel('Model parameter x')\n",
    "plt.ylabel('Density')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Beautiful! We obtained nice samples that follow the prior distribution."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The unit hypercube\n",
    "\n",
    "Lets specify a prior for UltraNest with multiple parameters:\n",
    "\n",
    "* a uniform distribution from 3.5 to 4.7\n",
    "* a log-uniform distribution from 0.01 to 100\n",
    "* a gaussian distribution around 2.0 +- 0.3\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# out transform function will receive one quantile corresponding to each of the three parameter\n",
    "def transform(quantile_cube):\n",
    "    # prepare the output array, which has the same shape\n",
    "    transformed_parameters = np.empty_like(quantile_cube)\n",
    "    # first parameter: a uniform distribution\n",
    "    transformed_parameters[0] = 3.5 + 1.2 * quantile_cube[0]\n",
    "    # second parameter: a log-uniform distribution\n",
    "    transformed_parameters[1] = 10**(-2 + 4 * quantile_cube[1])\n",
    "    # third parameter: Gaussian\n",
    "    transformed_parameters[2] = mydistribution.ppf(quantile_cube[2])\n",
    "    \n",
    "    return transformed_parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some recommendations:\n",
    "\n",
    "* [scipy.stats](https://docs.scipy.org/doc/scipy/reference/stats.html#continuous-distributions) provides many 1-d distributions that can be used like this.\n",
    "* avoid building scipy.stats objects in the transform, because this is slow -- build them outside first, then only invoke the .ppf method in the transform.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dependent priors\n",
    "\n",
    "In some cases, a previous experiment gives informative priors which we want to incorporate, and they may be inter-dependent. For example, consider a two-dimensional gaussian prior distribution:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "means = np.array([2, 3])\n",
    "cov = np.array([[1, 0.6], [0.6, 0.4]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rv = scipy.stats.multivariate_normal(means, cov)\n",
    "x, y = np.linspace(-1, 5, 400), np.linspace(1.5, 5, 400)\n",
    "X, Y = np.meshgrid(x, y)\n",
    "Z = rv.pdf(np.transpose([X.flatten(), Y.flatten()])).reshape(X.shape)\n",
    "plt.figure()\n",
    "plt.title('Correlated prior')\n",
    "plt.contourf(X, Y, Z, cmap='magma_r')\n",
    "plt.xlabel('Parameter 1')\n",
    "plt.ylabel('Parameter 2');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We recall:\n",
    "\n",
    "* Parameter 1 has a cumulative distribution\n",
    "* At each value of Parameter 1, Parameter 2 has a cumulative distribution.\n",
    "* We can thus specify a dependent distribution using the unit hypercube"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = np.linalg.inv(cov)\n",
    "l, v = np.linalg.eigh(a)\n",
    "rotation_matrix = np.dot(v, np.diag(1. / np.sqrt(l)))\n",
    "\n",
    "def transform_correlated(quantiles):\n",
    "    # sample a independent multivariate gaussian\n",
    "    independent_gaussian = scipy.stats.norm.ppf(quantiles)\n",
    "    # rotate and shift\n",
    "    return means + np.einsum('ij,kj->ki', rotation_matrix, independent_gaussian)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets try sampling!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = transform_correlated(np.random.uniform(0, 1, size=(100, 2)))\n",
    "\n",
    "plt.figure()\n",
    "plt.title('Correlated prior')\n",
    "plt.contourf(X, Y, Z, cmap='magma_r')\n",
    "plt.plot(samples[:,0], samples[:,1], 'o', mew=1, mfc='w', mec='k')\n",
    "plt.xlabel('Parameter 1')\n",
    "plt.ylabel('Parameter 2');\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A similar effect can be achieved by defining transforms in sequence (this is a different prior though):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gauss1 = scipy.stats.norm(2, 1)\n",
    "gauss2 = scipy.stats.norm(0, 0.1)\n",
    "\n",
    "\n",
    "def transform_correlated(quantiles):\n",
    "    parameters = np.empty_like(quantiles)\n",
    "    # first parameter is independent\n",
    "    parameters[0] = gauss1.ppf(quantiles[0])\n",
    "    # second parameter depends on first parameter, here with a shift\n",
    "    parameters[1] = parameters[0] + gauss2.ppf(quantiles[0])\n",
    "    return parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Non-analytic priors\n",
    "\n",
    "Sometimes, the prior may not be easily invertable. For example, when it is given as posterior samples from a previous analysis. I\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "posterior_samples = np.hstack((np.random.uniform(0, 3, 2000), np.random.normal(3, 0.2, 2000)))\n",
    "\n",
    "plt.figure(figsize=(4,2))\n",
    "plt.hist(posterior_samples, histtype='step', bins=100);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case, you can compute the cumulative distribution numerically and invert it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hist, bin_edges = np.histogram(posterior_samples, bins=100)\n",
    "hist_cumulative = np.cumsum(hist / hist.sum())\n",
    "bin_middle = (bin_edges[:-1] + bin_edges[1:]) / 2\n",
    "\n",
    "def transform_histogram(quantile):\n",
    "    return np.interp(quantile, hist_cumulative, bin_middle)\n",
    "\n",
    "samples = transform_histogram(np.random.uniform(size=1000))\n",
    "plt.figure(figsize=(4,2))\n",
    "plt.hist(posterior_samples, histtype='step', bins=100, density=True);\n",
    "plt.hist(samples, histtype='step', bins=100, density=True);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fraction parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some parameters such as fractions (or abundances) may be required to sum to 1. How can we specify such parameters?\n",
    "\n",
    "One option is to use absolute numbers. For example, instead of specifying the total mass and mass fractions for each chemical element, parameterise the mass of each element. This avoids the <=1 constraint, and may be easier to infer. A drawback is that the prior ranges for each element mass may be wide.\n",
    "\n",
    "The other option is to use the right distribution exactly made for this, which samples unbiased under the constraint (sum<=1): The [Dirichlet distribution](https://en.wikipedia.org/wiki/Dirichlet_distribution). Here we assume that the prior on the individual fraction is flat (flat Dirichlet distribution, $\\alpha=1$).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform_dirichlet(quantiles):\n",
    "    # https://en.wikipedia.org/wiki/Dirichlet_distribution#Random_number_generation\n",
    "    # first inverse transform sample from Gamma(alpha=1,beta=1), which is Exponential(1)\n",
    "    gamma_quantiles = -np.log(quantiles)\n",
    "    # dirichlet variables\n",
    "    return gamma_quantiles / gamma_quantiles.sum(axis=1).reshape((-1, 1))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets have a look at the samples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = transform_dirichlet(np.random.uniform(0, 1, size=(400, 3)))\n",
    "\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.view_init(elev=10., azim=-30)\n",
    "ax.plot(samples[:,0], samples[:,1], samples[:,2], 'x ');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The samples nicely lie on the plane where the sum is 1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  Further topics\n",
    "\n",
    "Check out the rest of the documentation and the tutorials.\n",
    "\n",
    "They illustrate how to:\n",
    "\n",
    "* [speeding up transforms with vectorisation](https://johannesbuchner.github.io/UltraNest/performance.html)\n",
    "* [How to use priors in Ultranest](https://johannesbuchner.github.io/UltraNest/usage-spectral-line.html)\n",
    "* and many example tutorials"
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}