workshop.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "572d699c",
   "metadata": {},
   "source": [
    "# Monte Carlo Methods\n",
    " This is a 2 hour workshop exploring the fundamentals of Monte Carlo methods as used in particle physics and astronomy.\n",
    " The workshop is designed to be run in a Jupyter notebook\n",
    " \n",
    " By the end of this workshop you will have explored the core concepts of:\n",
    " 1. Probability distributions\n",
    " 2. Samples\n",
    " 3. Sampling\n",
    " 4. Monte Carlo integration\n",
    " \n",
    " and the fundamental concepts of the core tools we use to do this."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d437b060",
   "metadata": {},
   "source": [
    "## 0. Setup\n",
    " I will presume that none of these commands are unfamiliar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8ef5637",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddb74617",
   "metadata": {},
   "source": [
    "## 1. Probability distributions\n",
    " \n",
    " __Recap:__ a probability distribution $P$ on a variable $x$ is defined such that \n",
    "  $$P(a<x<b) = \\int_a^b P(x) dx,$$\n",
    " or equivalently that $P(x)dx$ represents the probability that $x$ lies in the interval $[x,x+dx]$.\n",
    "\n",
    " Probability distributions are the building blocks of a lot of code, for example a cross section in particle physics:\n",
    " $$ \\sigma = \\int |\\mathcal{M}|^2 d\\Omega$$\n",
    " the matrix element $|\\mathcal{M}|^2$ is a distribution over collision events.\n",
    "\n",
    " In Cosmology, these will be posterior distributions $\\mathcal{P}(\\theta|D)$ generated from cosmological likelihoods $\\mathcal{L}(D|\\theta)$.\n",
    "\n",
    " The [scipy.stats](https://docs.scipy.org/doc/scipy/reference/stats.html) module is full of a whole variety of distributions, which forms the basis of a lot of other software.\n",
    " Let's create a [von Mises distribution](https://en.wikipedia.org/wiki/Von_Mises_distribution) (an example you may not have seen before), rather than starting with a boring normal distribution\n",
    " $$P(x) = \\frac{e^{\\kappa \\cos(x-\\mu)}}{2\\pi I_0(\\kappa)}.$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4e67ccb",
   "metadata": {},
   "outputs": [],
   "source": [
    "kappa = 1\n",
    "mu = 0\n",
    "dist = scipy.stats.vonmises(kappa, mu)\n",
    "x = np.linspace(-np.pi, np.pi, 1000)\n",
    "p = dist.pdf(x)\n",
    "plt.plot(x, p);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83e4b95e",
   "metadata": {},
   "source": [
    "__Exercise 1.1:__ produce a plot showing different choices of `kappa` and `mu`?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1637e8db",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Answer\n",
    "# Write your answer here into this code block"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41cce285",
   "metadata": {},
   "source": [
    "uncomment and execute the below to see the solution (only after you've had a go yourself)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c451ca88",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/1.1.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31b3d326",
   "metadata": {},
   "source": [
    "- __Question 1.1.1:__ What happens to the plot if you set `kappa` very large?\n",
    "   - __Answer 1.1.1:__ _insert_\n",
    " - __Question 1.1.2:__ [extension] Try out a few other distribution\n",
    "   - __Answer 1.1.2:__ _insert_"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "311ca9dd",
   "metadata": {},
   "source": [
    "### Two dimensional distributions\n",
    " These concepts are extended relatively straightforwardly to two-dimensional distributions $P(x,y)\n",
    "\n",
    " For a two-dimensional [von mises Fisher distribution](https://en.wikipedia.org/wiki/Von_Mises%E2%80%93Fisher_distribution) we can plot it as a contour plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f03db884",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the parameters of von-mises in terms of concentration kappa and mean direction (phi0, theta0)\n",
    "\n",
    "def vmf_dist(kappa, phi, theta):\n",
    "    mu = np.array([np.cos(phi)*np.sin(theta),\n",
    "                   np.sin(phi)*np.sin(theta),\n",
    "                   np.cos(theta)])\n",
    "    return scipy.stats.vonmises_fisher(mu, kappa)\n",
    "\n",
    "kappa = 1\n",
    "phi0 = np.pi\n",
    "theta0 = np.pi/2\n",
    "dist = vmf_dist(kappa, phi0, theta0)\n",
    "\n",
    "# Compute a meshgrid \n",
    "phi, theta = np.linspace(0, 2*np.pi, 100), np.arccos(np.linspace(1, -1, 100))\n",
    "phi, theta = np.meshgrid(phi, theta)\n",
    "x, y, z  = np.sin(theta) * np.cos(phi), np.sin(theta) * np.sin(phi), np.cos(theta)\n",
    "\n",
    "# Compute the pdf on the meshgrid with numpy broadcasting\n",
    "x = np.stack([x, y, z], axis=2)\n",
    "pdf = dist.pdf(x)\n",
    "plt.contourf(phi, theta, pdf)\n",
    "plt.colorbar();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0e699b4",
   "metadata": {},
   "source": [
    "__Exercise 1.2:__ Copy and paste the above code and adjust it to get an understanding of the effect of  `kappa` and `mu`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f2fd710",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Answer\n",
    "# Write your answer here into this code block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20ef043f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/1.2.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bce1f4f",
   "metadata": {},
   "source": [
    "- __Question 1.2.1:__ What do the colours represent?\n",
    "   - __Answer 1.2.1:__ _insert_\n",
    " - __Question 1.2.2:__ What is the relevance of the `arccos` in the above?\n",
    "   - __Answer 1.2.2:__ _insert_"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d077d5c4",
   "metadata": {},
   "source": [
    "It is more usual to plot 'one sigma' and 'two sigma' contours."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74b54317",
   "metadata": {},
   "outputs": [],
   "source": [
    "pdf = dist.pdf(x)\n",
    "pdf = pdf.reshape(-1)\n",
    "i = np.argsort(pdf)\n",
    "cdf = pdf[i].cumsum()\n",
    "cdf /= cdf[-1]\n",
    "sigma = np.empty_like(pdf)\n",
    "sigma[i] = np.sqrt(2) * scipy.special.erfinv(1-cdf)\n",
    "sigma = sigma.reshape(phi.shape)\n",
    "plt.contourf(phi, theta, sigma, levels=[0, 1, 2], colors=['black', 'gray'])\n",
    "plt.colorbar();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b89270fc",
   "metadata": {},
   "source": [
    "- __Question 1.3.1:__ What exactly do \"one sigma\" and \"two sigma\" contours mean, and how does the above code compute them?\n",
    "   - __Answer 1.3.1:__ _insert_\n",
    " - __Question 1.3.2:__ How should one choose the colours for the contours? (have a look at the [anesthetic](https://github.com/handley-lab/anesthetic#another-posterior-plotting-tool) README for one option for doing this.\n",
    "   - __Answer 1.3.2:__ _insert_\n",
    "\n",
    " __Beware:__ [corner.py](https://corner.readthedocs.io/en/latest/pages/sigmas/) uses a different definition of sigma."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51530323",
   "metadata": {},
   "source": [
    "### Higher dimensional distributions\n",
    " In most cases we have far more than two variables, which makes plotting the full distribution difficult.\n",
    " \n",
    " The best we can do visually is to plot the _marginal distributions_, i.e. the distributions with each other variable integrated out. Since 1D and 2D distributions are plottable, we typically do this by considering the 1D marginals individually:\n",
    " $$ P(x_i) = \\int P(x) \\prod\\limits_{k\\ne i} dx_k$$\n",
    " and pairwise marginals\n",
    " $$ P(x_i, x_j) = \\int P(x) \\prod\\limits_{k\\ne i, k\\ne j} dx_k$$\n",
    " and arrange these into a \"corner\" or \"triangle\" plot."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97c85e50",
   "metadata": {},
   "source": [
    "This example takes a few seconds to generate and plot -- don't worry about the code details, we will cover them later:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a076dcbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from anesthetic.examples.perfect_ns import planck_gaussian\n",
    "params = ['omegabh2', 'omegach2', 'theta', 'tau', 'logA', 'ns']\n",
    "planck_samples = planck_gaussian()[params].compress()\n",
    "planck_samples.plot_2d(kind='kde');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cadd28f6",
   "metadata": {},
   "source": [
    "The next section builds up to how we go about producing these plots in practice."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf1eb9fe",
   "metadata": {},
   "source": [
    "## 2. Samples: why do sampling?\n",
    " \n",
    " The core concept in numerical inference is that of _samples_. The premise is straightforward -- given some density $P(x)$, generate random numbers whose density in the large-number limit is equal to $P(x)$.\n",
    "\n",
    " `scipy.stats` functions have the 'random variables' `.rvs()` method built in, which does exactly this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "622ad6c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "kappa = 1\n",
    "mu = 0\n",
    "dist = scipy.stats.vonmises(kappa, mu)\n",
    "x = np.linspace(-np.pi, np.pi, 1000)\n",
    "p = dist.pdf(x)\n",
    "plt.plot(x, p);\n",
    "samples = dist.rvs(10000)\n",
    "plt.hist(samples, density=True, bins=50);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb1de611",
   "metadata": {},
   "source": [
    "- __Question 2.1.1:__ what is the relevance of the `density=True` and `bins=50` arguments to `plt.hist`?\n",
    "   - __Answer 2.1.1:__ _insert_"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84f52516",
   "metadata": {},
   "source": [
    "We can also sample from the 2D distribution we defined above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ea339a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "kappa = 4\n",
    "phi0 = 0\n",
    "theta0 = np.pi/2\n",
    "dist = vmf_dist(kappa, phi0, theta0)\n",
    "N = 10000\n",
    "samples = dist.rvs(N)\n",
    "theta, phi = np.arccos(samples[:,2]), np.arctan2(samples[:,1], samples[:,0])\n",
    "plt.plot(phi, theta, '.', markersize=1);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb768620",
   "metadata": {},
   "source": [
    "One of the lesser-know functionalities of matplotlib is that the samples above can also be used to plot contours with the triangulation functionality provided by the `tricontourf` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57efc36b",
   "metadata": {},
   "outputs": [],
   "source": [
    "pdf = dist.pdf(samples)\n",
    "i = np.argsort(pdf)\n",
    "cdf = np.arange(1,N+1)/(N+1)\n",
    "sigma = np.empty_like(pdf)\n",
    "sigma[i] = np.sqrt(2) * scipy.special.erfinv(1-cdf)\n",
    "plt.tricontourf(phi, theta, sigma, levels=[0, 1, 2], colors=['black', 'gray'], alpha=0.8)\n",
    "plt.plot(phi, theta, '.', markersize=1)\n",
    "plt.colorbar();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b47185f",
   "metadata": {},
   "source": [
    "- __Question 2.2.1:__  Why is the sigma calculation different in comparison to the previous example? [hard]\n",
    "   - __Answer 2.2.1:__ _insert_"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2c0a754",
   "metadata": {},
   "source": [
    "One way to see how powerful the above approach of using samples to make plots is, is to consider how we might go about making the triangle plot using the meshgrid approach. To do this, we would have to solve two problems:\n",
    " - __Question 2.3.1:__ How would you use a meshgrid to compute marginal distributions?\n",
    "   - __Answer 2.3.1:__ _insert_\n",
    " - __Question 2.3.2:__ How much more expensive would this be in 6 dimensions?\n",
    "   - __Answer 2.3.2:__ _insert_"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6a3fc96",
   "metadata": {},
   "source": [
    "Sampling solves both of these:\n",
    " - marginal samples are found just by ignoring columns\n",
    " - drawing samples from higher-dimensional distributions is not exponentially harder\n",
    " To see this, let's take a look at the planck_samples, which when printed show an anesthetic (pandas extension) array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0730844d",
   "metadata": {},
   "outputs": [],
   "source": [
    "planck_samples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad1d8294",
   "metadata": {},
   "source": [
    "We can plot samples from the marginal distribution by"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56d77b61",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(*planck_samples[['logA', 'tau']].to_numpy().T, '.');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c01ef03f",
   "metadata": {},
   "source": [
    "To plot the marginal contours, we have to use some form of low-dimensional density estimation. Kernel density estimation (KDE) is the standard (basically putting a small gaussian on each sample and adding these together) but in principle one could use neural density estimators or histograms for more/less advanced versions. Density estimation is an acceptable approximation in one and two dimensions but becomes very innacurate in higher dimensions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f33ae19",
   "metadata": {},
   "source": [
    "There are many packages that implement all of this for you, and I would encourage you to resist the temptation to write your own! (irony noted)\n",
    " - [getdist](https://getdist.readthedocs.io/en/latest/)\n",
    "   - state-of-the-art in KDE edge correction\n",
    "   - industry standard since 2010\n",
    "   - difficult to extend/use\n",
    " - [corner.py](https://corner.readthedocs.io/en/latest/)\n",
    "   - histogram-based foreman-mackay software\n",
    " - [chainconsumer](https://samreay.github.io/ChainConsumer/)\n",
    "   - increasingly popular python tool for MCMC samples\n",
    " - [anesthetic](https://anesthetic.readthedocs.io/en/latest/)\n",
    "   - specialised for nested sampling, but can also do MCMC\n",
    "   - explicitly builds on the numpy/scipy/pandas stack.\n",
    " \n",
    " You can see here that the python package anesthetic by default uses the whitespace above the diagonal to plot samples, in addition to estimates of the 1D and 2D marginals:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e2b1d17",
   "metadata": {},
   "outputs": [],
   "source": [
    "planck_samples[['logA', 'tau']].plot_2d();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68265233",
   "metadata": {},
   "source": [
    "Samples are an extremely powerful tool for performing numerical inference. In particular the following property holds\n",
    " $$ x\\sim P(x)dx \\quad\\Rightarrow\\quad f(x) \\sim P(f)df$$\n",
    " namely, if you have a set of samples in a variable $x$, and you want to know how $y=f(x)$ is distributed, you just assume the answer is $x_i$ and compute $y_i = f(x_i)$ for each sample. Sampling turns uncertainty quantification into repeated forward models."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44a33d2b",
   "metadata": {},
   "source": [
    "__Exercise 2.4:__ if $x$ is normally distributed, plot the distribution of $2^xx$. Bonus question: prove mathematically that in this case x is log-normally distributed, find its scale parameter, and plot this on your plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e5e6099",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58ddc680",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/2.4.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ee743db",
   "metadata": {},
   "source": [
    "One can compute averages by summing over samples:\n",
    " $$ \\langle f \\rangle_P \\quad=\\quad \\int f(x) P(x) dx \\quad\\approx\\quad \\frac{1}{N} \\sum\\limits_{i=1}^N f(x_i)$$\n",
    " Note that the (generally unknown) density __does not__ appear in the final expression, which shows that sampling is our primary tool for avoiding high-dimensional density estimation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2978bfb4",
   "metadata": {},
   "source": [
    "To confuse you on this last point, most methods for generating samples will often create _weighted samples_, namely each sample $x_i$ has an associated weight $w_i$, and the average is computed as:\n",
    " $$ \\langle f \\rangle_P \\quad\\approx\\quad \\tfrac{1}{\\sum_i w_i}\\sum_i w_i f(x_i)$$\n",
    " This weight is __not__ generally the probability density.\n",
    "\n",
    " `anesthetic` is a pandas extension which computes weighted statistics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bd61ab4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from anesthetic import Samples\n",
    "\n",
    "x = np.random.rand(100,3)\n",
    "w = np.random.rand(100)\n",
    "samples = Samples(x, weights=w)\n",
    "samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7a189ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "samples.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "505fb258",
   "metadata": {},
   "outputs": [],
   "source": [
    "samples.std()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87641355",
   "metadata": {},
   "source": [
    "Exercise 2.5: confirm that the the weighted mean and standard deviation are not the same as numpy's default."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db99a627",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da2f6432",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/2.5.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50b62c8f",
   "metadata": {},
   "source": [
    "Getting weighted means and standard deviations correct is a faff, and a common source of bugs in numerical inference code, so it's worth being aware of this/using anesthetic's extensive functionality."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a66b63e6",
   "metadata": {},
   "source": [
    "### Summary\n",
    " Samples are the fundamental building block of numerical inference. You should view them as \"souped-up error bars\", summarising uncertainty in your analysis. If you have an analysis that you could do if you knew all the inputs $x$, and you have samples from the distribution $x$ belongs to, you just do the analysis $N$ times for each sample, get $N$ answers, and the distribution of these answers quantifies your uncertainty.\n",
    "\n",
    " The golden rule of Numerical inference is therefore to __stay in samples__ until the end. You should know by now that in general\n",
    " $$f(\\langle x \\rangle) \\ne \\langle f(x) \\rangle $$\n",
    " so taking an average/summary before the end can bias your inference.\n",
    "\n",
    " You can then use samples to compute averages, and extract marginal distributions by ignoring columns, from which you can compute 1D and 2D contours. This is the basis of the triangle plot.\n",
    "\n",
    " __Question 2.6:__ [hard] what statistical operation can you not get easily from samples?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e76b869",
   "metadata": {},
   "source": [
    "We now move on to techniques as to how one can generate samples from an arbitrary distribution."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49b48433",
   "metadata": {},
   "source": [
    "## 3. Sampling\n",
    " We will assume that we can evaluate the probability density function $P(x)$ (or at least the unnormalised density $P^*(x)\\propto P(x)$) in finite time (i.e. seconds). For initiates, it might be surprising that access to the exact $P(x)$ is not sufficient to generate samples, but it is not.\n",
    "\n",
    "(N.B. The frontier of inference at the moment is simulation based inference, which develops methods for inference when you only have a simulator $f(x)$, but this is beyond the scope of this workshop)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68a19a67",
   "metadata": {},
   "source": [
    "### 3.0 Random sampling\n",
    " Let's first see why random sampling is not sufficient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f08d5769",
   "metadata": {},
   "outputs": [],
   "source": [
    "from handleymcmethods.examples import planck"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba2b978c",
   "metadata": {},
   "source": [
    "__Exercise 3.0.1:__ Generate 10000 samples from planck.prior, and find the largest planck.loglikelihood value in the samples. Repeat this process a few times. What do you notice?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4efd3b83",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc442864",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/3.0.1.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000ff1",
   "metadata": {},
   "source": [
    "The issue is that the prior is much wider than the likelihood/posterior, so random samples are very unlikely to be close to the peak of the likelihood. This same problem would occur for meshgrid sampling.\n",
    "\n",
    " One solution for finding a good region is of course a gradient descent\n",
    " __Exercise 3.0.2:__ Use `scipy.optimize.minimize` to find the maximum of the loglikelihood. How does this compare to the maximum of the other results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bbb0ed3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38556968",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/3.0.2.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7cc1549",
   "metadata": {},
   "source": [
    "Whilst this approach does find the maximum probability point, this does not generate samples, which is what we need for our error bars."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "239446f1",
   "metadata": {},
   "source": [
    "### 3.1 Metropolis Hastings\n",
    "\n",
    " The first approach that can successfully generate samples from a distribution is the [Metropolis-Hastings algorithm](https://en.wikipedia.org/wiki/Metropolis%E2%80%93Hastings_algorithm).\n",
    " The simplest version of this algorithm is as follows:\n",
    " - start at some point $x_0$\n",
    " - at iteration $i$: \n",
    "   - propose a new point $x'$ a random step away from $x_i$\n",
    "   - accept the point with probability\n",
    "   $$ \\alpha = \\frac{P(x')}{P(x_i)}$$\n",
    "   - if the point is accepted, $x_{i+1} = x'$, otherwise $x_{i+1} = x_i$\n",
    "   - stop when you have enough samples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13fd2822",
   "metadata": {},
   "source": [
    "__Exercise 3.1.1:__ Implement the Metropolis-Hastings algorithm for the planck likelihood. \n",
    " - hint: if it's not working, try plotting the the set of points you're generating"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8554a82",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0455ba5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/3.1.1.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0537db0d",
   "metadata": {},
   "source": [
    "Things can be made a lot better by choosing a better method for proposing new points.\n",
    " For example, if you use the true posterior coviariance matrix, the solution converges much better\n",
    " - hint: you can get the covariance matrix from `planck.cov` and the mean from `planck.mean`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0fde094",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69ed8274",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/3.1.2.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70daa7b6",
   "metadata": {},
   "source": [
    "Of coures in practice one doesn't know the answer going in, and therefore typically it has to be learned, either by knowing something about the physics, using another algorithm first, or gradually updating the proposal distribution."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b987df42",
   "metadata": {},
   "source": [
    "- __Question 3.1.3:__ Where does the metropolis hastings algorithm fail?\n",
    "   - __Answer 3.1.3:__ _insert_\n",
    " - __Question 3.1.4:__ How well does this algorithm parallelise?\n",
    "   - __Answer 3.1.4:__ _insert_"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b48a36e3",
   "metadata": {},
   "source": [
    "The full algorithm is as follows:\n",
    " - start at some point $x_0$\n",
    " - propose a new point $x'$ from some proposal distribution $Q(x'|x)$\n",
    " - accept the point with probability\n",
    " $$ \\alpha = \\min\\left(1, \\frac{P(x')Q(x|x')}{P(x)Q(x'|x)}\\right)$$\n",
    " - repeat until convergence criterion is reached\n",
    "\n",
    " Note this includes the generalisation to asymmetric proposal distributions, which is necessary for the algorithm to converge, and more carefully acounts for the fact that the probability shouldn't be greater than 1. Convergence criteria usually include a variation of the Gelman-Rubin statistic.\n",
    " \n",
    " Example implementations of metropolis hastings include\n",
    " - [PyMC](https://www.pymc.io/welcome.html)\n",
    " - [Cobaya](https://cobaya.readthedocs.io/en/latest/)\n",
    " - [CosmoSIS](https://cosmosis.readthedocs.io/en/latest/)\n",
    " - [MontePython](https://monte-python.readthedocs.io/en/latest/)\n",
    "\n",
    " Though in practice, since a metropolis sampler is so easy to write, and aspects of the proposal distribution are so system-specific, many people just choose to write their own!\n",
    "\n",
    " More modern work is exploring the use of neural networks to learn the proposal distribution, which can be much more efficient than the above.\n",
    " - [FlowMC](https://arxiv.org/abs/2211.06397)\n",
    " - [MCMC-diffusion](https://arxiv.org/abs/2309.01454)\n",
    "\n",
    " Further extensions to this approach include ensemble sampling ([emcee](https://emcee.readthedocs.io)), slice sampling ([zeus](https://zeus-mcmc.readthedocs.io/en/latest/)) and many others."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e86292b2",
   "metadata": {},
   "source": [
    "### 3.2 Nested sampling\n",
    "\n",
    " As discussed in the [talk](https://github.com/williamjameshandley/talks/raw/unam_2023/will_handley_unam_2023.pdf), the nested sampling algorithm can be summarised as:\n",
    " - generate nlive samples from the prior\n",
    " - at each iteration, replace the lowest likelihood sample with a new sample from the prior at greater likelihood\n",
    " - stop when the live points have sufficiently compressed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "333898da",
   "metadata": {},
   "source": [
    "__Exercise 3.2.1:__ Implement the nested sampling algorithm for the planck likelihood with 50 live points, using a brute-force prior-sample+rejection approach. How many iterations do you get through? (put in a print statement to see the slow down)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1aa5c8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c8c7978",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/3.2.1.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7fe113c",
   "metadata": {},
   "source": [
    "You should find you get about to about 500 iterations before you run out of patience!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4e43de0",
   "metadata": {},
   "source": [
    "__Exercise 3.2.2:__ This time, implement a more efficient approach by using a box around the live points to generate samples from the prior. To be correct, you should make the box slightly larger than this! Run the algorithm for \n",
    " - __Question 3.2.2:__ What are the failure modes of this method?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32be79e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff8f34ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/3.2.2.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7f8aff9",
   "metadata": {},
   "source": [
    "__Exercise 3.2.3:__ Adjust your algorithm so that it records the dead points, as well as the 'birth contour'. Plot the dead points. Pass these into the anesthetic gui"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2ffc717",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc113061",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/3.2.3.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5568e48",
   "metadata": {},
   "source": [
    "__Exercise 3.2.4:__ Write a non-rejection based sampling algorithm (e.g. metropolis hastings using the covariance of the live points to build a proposal distribution) and compare the speed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27e168b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69bad4c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/3.2.4.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cfa379c",
   "metadata": {},
   "source": [
    "## Interlude\n",
    "\n",
    " The goal of the above 'my first metropolis algorithm' and 'my first nested sampler' exercises is to give you a feel for how at these algorithms work at the base level, and how they can be implemented, so that you can assess the efficiacy of new ideas. In practice you should use one of the more established libraries for doing this, which have been battle-tested and optimised for speed. In particular, some of the things which cause ~O(10) lines of code to expand to ~O(1000) lines are:\n",
    " - robust live point generation\n",
    " - parallelisation\n",
    " - mode identification & handling\n",
    "\n",
    " Most packages focus on the first of these, many (but not all) provide the second, with very few implementing the substantial bookkeeping required for the third.\n",
    "\n",
    " As discussed in the talk, samplers \n",
    "\n",
    " #### Rejection samplers (or 'region samplers')\n",
    " - [MultiNest](https://github.com/farhanferoz/MultiNest)\n",
    "   - ellipsoidal decomposition, parallelised, clustered\n",
    "   - the original nested sampler\n",
    "   - Incorporated into a wide variety of packages (GAMBIT, CosmoSIS, MontePython, Cobaya, etc)\n",
    "   - Johannes Buchner implemented a [python wrapper](https://johannesbuchner.github.io/PyMultiNest/)\n",
    " - [UltraNest](https://johannesbuchner.github.io/UltraNest/)\n",
    "   - kernel-based method for generating live points, paralellised\n",
    "   - Johannes Buchner's new python nested sampler\n",
    "   - incorporates new advances in nested sampling not in MultiNest\n",
    "   - python only (so for very fast likelihoods may be worth compiling multinest)\n",
    "   - also implements many other strategies (including path sampling)\n",
    " - [nessai](https://nessai.readthedocs.io/en/latest/)\n",
    "   - normalising flows to generate live points\n",
    "   - popular in gravitational waves community\n",
    "\n",
    " #### Chain-based samplers (or 'path samplers')\n",
    " - [PolyChord](https://github.com/PolyChord/PolyChordLite)\n",
    "   - slice sampling, parallelised, clustered\n",
    "   - the original path sampler\n",
    " - [pymatnext](https://libatoms.github.io/pymatnest/intro.html)\n",
    "   - specific for materials science\n",
    " \n",
    " #### Other samplers\n",
    " - [dynesty](https://dynesty.readthedocs.io/en/latest/)\n",
    "   - python re-implementation of many of the above\n",
    "   - very popular\n",
    "   - easy to install\n",
    "   - defaults not well chosen for path-based samplers\n",
    "   - although 'dy' is in the name, most nested samplers (polychord, ultranest etc) are now dynamic nested samplers as well\n",
    "\n",
    " __Will Handley's recommendations on choosing samplers:__\n",
    " - if your problem is low-dimensional ~O(10) parameters, use UltraNest\n",
    "   - easier for newcomers to install than multinest\n",
    "   - self-tunes its efficiency parameters, so gives an honest scaling with dimensionality\n",
    " - if your problem is high-dimensional, use PolyChord\n",
    "   - whilst other samplers (dynesty, ultranest) implement slice sampling, neither implement clustering, which is necessary for multi-modal problems\n",
    "   - in particular dynesty is very poorly parallelised (since it is optimised for notebook rather than HPC work).\n",
    "   - the dynesty defaults for slice sampling are too generous (and therefore often give wrong, but faster answers)\n",
    "   - it is slightly harder to install due to the legacy fortran component.\n",
    "\n",
    " ### Practical advice for using nested sampling\n",
    " Nested sampling run-time scales as:\n",
    " $$ T \\propto  f_\\mathrm{sampler} \\times n_\\mathrm{live}/n_\\mathrm{cores}$$\n",
    " providing $n_\\mathrm{cores} \\le n_\\mathrm{live}$.\n",
    " - rejection samplers: the $f_\\mathrm{sampler} = e^{d/d_0}$ where $d_0$ is problem dependent$.\n",
    " - chain-based samplers: the $f_\\mathrm{sampler} = \\mathcal{O}(3) \\times n_\\mathrm{repeat}$\n",
    "\n",
    " (The proportionality constant is $\\mathcal{D}_\\mathrm{KL}(\\mathcal{P}||\\pi) T_\\mathrm{like}$)\n",
    "\n",
    " This scaling is very useful for planning big runs, all you need to do is determine the proportionality constant, which in practice you get from a preliminary run.\n",
    "\n",
    " The strategy is therefore:\n",
    " 1. run a preliminary run with a small number of live points, debug code, and determine $T_0$.\n",
    "    - for rejection samplers, you need a minimum number of samples to train the region proxy (ellipsoidal decomposition, normalising flow), usually ~O(500)\n",
    "    - for path samplers this can be set much lower (as low as 5!)\n",
    "    - Since nested sampling is parallelised up to the number of cores, setting nlive ~ ncores is a good starting point (e.g. on modern HPC machines ~O(50), on your laptop ~O(10)).\n",
    " 2. scale up to a production run (nlive ~ 1000)\n",
    "   - If you had $n_\\mathrm{live}^0$ and $n_\\mathrm{cores}^0$ giving you $T_0$, the full runtime will be\n",
    "    $$T = T_0 \\times \\frac{n_\\mathrm{live}}{n_\\mathrm{live}^0} \\times \\frac{n_\\mathrm{cores}^0}{n_\\mathrm{cores}}$$\n",
    "   - With arbitrary HPC, you can scale the cores with the live points and retain the same walltime $T_0$.\n",
    " 3. check that you ran with high enough $f_sampler$ by halving/doubling it and checking the results are consistent."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d9137d6",
   "metadata": {},
   "source": [
    "General advice on usage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d282112",
   "metadata": {},
   "source": [
    "## References:\n",
    " - [Nature review article](https://arxiv.org/abs/2205.15570): Nested sampling for physical sscientists\n",
    "   - pedagogical introduction by the community to nested sampling\n",
    " - [Technical review article](https://arxiv.org/abs/2101.09675): Nested sampling methods\n",
    "   - technical review by Johannes Buchner (author of UltraNest)\n",
    "   - detailed and complete reference list of the entire nested sampling literature up to early 2023.\n",
    " - [John Skillings original paper](https://projecteuclid.org/journals/bayesian-analysis/volume-1/issue-4/Nested-sampling-for-general-Bayesian-computation/10.1214/06-BA127.full)\n",
    "   - original paper on nested sampling\n",
    "   - a goldmine of insight in John's unique style"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a87c58e1",
   "metadata": {},
   "source": [
    "## 4. Integration\n",
    " \n",
    " If your probability distribution is not normalised, the integration constant is of critical importance, either as a cross section (particle physics), a Bayesian evidence (cosmology) or a partition function (statistical mechanics).\n",
    "\n",
    " Fundamentally we want to compute:\n",
    " $$ Z = \\int P^*(x) dx$$\n",
    " The traditional discussion begins by pointing out that over the domain of $P^*(x)$, the region where $P^*(x)$ is of non-zero probability is very small, which we don't know a-priori where it is.  On the face of it, this doesn't look like a deal-breaker -- we have been developing ever more sophisticated ways of generating samples $x\\sim P$, so we have plenty of points with non-zero $P^*(x)$\n",
    "\n",
    " However, there is another portion of the integral, namely $dx$, which posterior samples __do not__ give us. The challenge is therefore not finding the \"typical set\", or generating points within it, it is measuring its volume."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "644c79b8",
   "metadata": {},
   "source": [
    "### 4.1 Importance sampling\n",
    " The go-to method in particle physics for doing this is importance sampling.\n",
    " The premise here is to find a normalised distribution $Q(x)$ which easy to sample from (for example a scipy distribution), and which is similar to $P^*(x)$ in the region where $P^*(x)$ is non-zero.\n",
    " \n",
    " One then uses the (almost trivial) result:\n",
    " $$ \\int P^*(x) dx = \\int \\frac{P^*(x)}{Q(x)} Q(x) dx = \\left\\langle \\frac{P^*(x)}{Q(x)} \\right\\rangle_Q$$\n",
    " I.e. one generates samples from $x\\sim Q(x)$ and computes the average of $P^*(x)/Q(x)$.\n",
    "\n",
    " You can think of this intuitively as 'trimming off' the regions where $P^*(x)$ is zero, and then computing the average of the remaining regions -- a more sophisticated way of picking a narrower prior\n",
    "\n",
    " If you choose a poor $Q$, then this will be very inefficient, with very few samples contributing to the integral. These weighted samples are in a definite sense exact"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84997a41",
   "metadata": {},
   "source": [
    "__Exercise 4.1.1:__ Take $P^*(x)$ to be the `planck.loglikelihood` added to the `planck.prior.logpdf`, and Q(x) to be a scipy.stats_multivariate_normal with `planck.mean`, `planck.cov`. Compute the integral using importance sampling. You should get `logZ=-1431.403883060199`. You may find it useful to use `scipy.special.logsumexp` to compute the average."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e922f88",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f5fdb80",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/4.1.1.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bbfbb24",
   "metadata": {},
   "source": [
    "__Exercise 4.1.2:__ Now compute the integral with a slightly misspecified proposal $Q$ by choosing a mean offset by a random amount. Compute the effective number of samples using the formula $$ n = \\frac{(\\sum_i w_i)^2 }{ \\sum_i w_i^2} $$ where $w_i$ are the importance weights we are averaging over"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7477592e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "393ef3d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/4.1.2.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef3b1738",
   "metadata": {},
   "source": [
    "__Exercise 4.1.2:__ Try other distributions (e.g. multivariate_t, adjusting the covariance matrix using a wishart distribution, etc). How does the efficiency change?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8290c664",
   "metadata": {},
   "source": [
    "What the above exercise shows is that if you have a good proposal $Q$, then this can be a very efficient way of computing an integral. However, the efficiency drops for even slightly misspecified $Q$, and this problem exponentially worsens in high dimensions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85098f11",
   "metadata": {},
   "source": [
    "### 4.2 Nested sampling\n",
    "\n",
    " Nested sampling provides a more sophisticated way to achieve the above. We begin by making the standard [Lebesgue integral](https://en.wikipedia.org/wiki/Lebesgue_integration) manipulation. First define the volume (/prior volume/measure/CDF) as:\n",
    " $$ X(P^*) = \\int_{P^*(x)>P^*} dx$$\n",
    " this is the volume of the space contained inside each contour of $P^*$.  With this definition, we can say\n",
    " $$ Z = \\int P^* dx = \\int P^* dX \\approx \\sum_i P^*_i \\Delta X_i $$\n",
    " which has transformed the multidimensional integral into a one-dimensional one which can be tractably numerically integrated."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1379dc0e",
   "metadata": {},
   "source": [
    "The power of nested sampling is that we can estimate $X_i$ from the nested sampling compression procedure:\n",
    " $$X_i \\approx \\frac{n}{n+1} X_{i-1} $$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "664fc9c5",
   "metadata": {},
   "source": [
    "__Exercise 4.2.1:__ Compute the nested sampling evidence estimate, and compare to the value you found with importance sampling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "679fe091",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7627bc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load solutions/4.2.1.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37dc0512",
   "metadata": {},
   "source": [
    "It's a little more involved to quantify the $\\approx$ using samples from the distribution of $X_i$, but fortunately anesthetic has implemented all of this for you"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f3941f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = planck_gaussian()\n",
    "samples.logZ(1000).hist()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64d41e02",
   "metadata": {},
   "source": [
    "and much more besides"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48eae018",
   "metadata": {},
   "outputs": [],
   "source": [
    "samples.nlive.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5159df09",
   "metadata": {},
   "source": [
    "and much more besides"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a444aadc",
   "metadata": {},
   "outputs": [],
   "source": [
    "samples.stats()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ed71beb",
   "metadata": {},
   "source": [
    "and much more besides"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c77f4196",
   "metadata": {},
   "outputs": [],
   "source": [
    "samples.stats(1000).plot_2d()"
   ]
  }
 ],
 "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}