Lab_Lec3_photon_counts.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Frequentist vs Bayesian Approaches to counting photons an the Metropolis-Hastings algorithm\n",
    "Here we will look into the issue of frequentist vs baysian probabilities. We will solve the problem of counting photons from the two perspectivs. We will also introduce the  Metropolis-Hastings algorithm, and will show how to use it to infer the parameters of a model, by sampling from a posterior distribution. This notebook is based on work by Jake Vanderplas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy import stats\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first generate some fake but realistic dataset for the fluxes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generating some simple photon count data\n",
    "import numpy as np\n",
    "from scipy import stats\n",
    "np.random.seed(1)  # for repeatability\n",
    "\n",
    "F_true = 1000  # true flux, say number of photons measured in 1 second\n",
    "N = 50 # number of measurements\n",
    "F = stats.poisson(F_true).rvs(N)  # N measurements of the flux\n",
    "e = np.sqrt(F)  # errors on Poisson counts estimated via square root"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.errorbar(F, np.arange(N), xerr=e, fmt='ok', ecolor='gray', alpha=0.5)\n",
    "ax.vlines([F_true], 0, N, linewidth=5, alpha=0.2)\n",
    "ax.set_xlabel(\"Flux\");ax.set_ylabel(\"measurement number\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Frequentist View\n",
    "\n",
    "If we assume Gaussian errors, then the probability distribution for a single measurement is a normal distribution:\n",
    "\n",
    "$$P(D_i~|~F_{\\rm true}) = \\frac{1}{\\sqrt{2\\pi e_i^2}} \\exp{\\left[\\frac{-(F_i - F_{\\rm true})^2}{2 e_i^2}\\right]}$$\n",
    "\n",
    "And we construct the likelihood as the product of all single measurements ($D$):\n",
    "\n",
    "$$\\mathcal{L}(D~|~F_{\\rm true}) = \\prod_{i=1}^N P(D_i~|~F_{\\rm true})$$\n",
    "\n",
    "It is easier in general to calculate the logarithm of this expression, since the product becomes a sum and we do not loose generality:\n",
    "\n",
    "$$\\log\\mathcal{L} = -\\frac{1}{2} \\sum_{i=1}^N \\left[ \\log(2\\pi  e_i^2) + \\frac{(F_i - F_{\\rm true})^2}{e_i^2} \\right]$$\n",
    "\n",
    "In this case we can compute the value of $F_true$ that maximizes the likelihood analytically, by setting the derivative of the previous expression to zero, and solving for the parameter, which results in:\n",
    "\n",
    "$$F_{\\rm est} = \\frac{\\sum w_i F_i}{\\sum w_i};~~w_i = 1/e_i^2$$\n",
    "\n",
    "If all erros are equal, this reduces to the mean, as expected:\n",
    "\n",
    "$$F_{\\rm est} = \\frac{1}{N}\\sum_{i=1}^N F_i$$\n",
    "\n",
    "We can evaluate this for our little toy model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w = 1. / e ** 2\n",
    "print(\"\"\"\n",
    "      F_true = {0}\n",
    "      F_est  = {1:.0f} +/- {2:.0f} (based on {3} measurements)\n",
    "      \"\"\".format(F_true, (w * F).sum() / w.sum(), w.sum() ** -0.5, N))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bayesian View\n",
    "In this case, we perform the inference directly on the model parameters, which are assumed to be random variables. We infer the physics, not the data. Therefore, we want to estimate:\n",
    "\n",
    "$$P(F_{\\rm true}~|~D)$$\n",
    "\n",
    "In this case, it makes sense to talk about the probability of the parameters like $F_{\\rm{true}}$. We can use the Bayes' rule now to express this posterior distribution as a function of the likelihood and the *prior*:\n",
    "\n",
    "$$P(F_{\\rm true}~|~D) = \\frac{P(D~|~F_{\\rm true})~P(F_{\\rm true})}{P(D)}$$\n",
    "\n",
    "The prior contains our degree of belief on the parameter having certain value *before* we perform the measurements. This previous knowledge can come from previous similar experiments, or independent theories or measurements that inform us about the value that the parameter *should* have. The new data has a very specific task, then: update our prior believe into a posterior that takes into account the new data. Here we assume that the set of parameters is given by $\\theta$, which in this case consists only of the flux, i.e., $\\theta=\\left[F_{\\rm true}\\right]$. We start with a flat prior:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_prior(theta):\n",
    "    return 1  # flat prior\n",
    "\n",
    "# Define the likelihood\n",
    "def log_likelihood(theta, F, e):\n",
    "    return -0.5 * np.sum(np.log(2 * np.pi * e ** 2)\n",
    "                         + (F - theta[0]) ** 2 / e ** 2)\n",
    "\n",
    "# In log space, the posterior is the sum of likelihood and prior.\n",
    "def log_posterior(theta, F, e):\n",
    "    return log_prior(theta) + log_likelihood(theta, F, e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now use emcee to sample this posterior."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ndim = 1  # number of parameters in the model\n",
    "nwalkers = 50  # number of MCMC walkers\n",
    "nburn = 1000  # \"burn-in\" period to let chains stabilize\n",
    "nsteps = 2000  # number of MCMC steps to take\n",
    "\n",
    "# we'll start at random locations between 0 and 2000\n",
    "starting_guesses = 2000 * np.random.rand(nwalkers, ndim)\n",
    "\n",
    "import emcee\n",
    "sampler = emcee.EnsembleSampler(nwalkers, ndim, log_posterior, args=[F, e])\n",
    "sampler.run_mcmc(starting_guesses, nsteps)\n",
    "\n",
    "sample = sampler.chain  # shape = (nwalkers, nsteps, ndim)\n",
    "sample = sampler.chain[:, nburn:, :].ravel()  # discard burn-in points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot a histogram of the sample\n",
    "plt.hist(sample, bins=50, histtype=\"stepfilled\", alpha=0.3, normed=True)\n",
    "\n",
    "# plot a best-fit Gaussian\n",
    "F_fit = np.linspace(975, 1025)\n",
    "pdf = stats.norm(np.mean(sample), np.std(sample)).pdf(F_fit)\n",
    "\n",
    "plt.plot(F_fit, pdf, '-k')\n",
    "plt.xlabel(\"F\"); plt.ylabel(\"P(F)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A bit more complicated\n",
    "\n",
    "### Frequentist approach\n",
    "\n",
    "Now we will consider a case in which the star is not fixed, but has some variability. We will model that variability with a Gaussian model (which is different from the Gaussian model that describes the measurement errors). The new model is the convolution betweem the instrinsic variability and the measurement uncertainty:\n",
    "\n",
    "$$\n",
    "\\mathcal{L}(D~|~\\theta) = \\prod_{i=1}^N \\frac{1}{\\sqrt{2\\pi(\\sigma^2 + e_i^2)}}\\exp\\left[\\frac{-(F_i - \\mu)^2}{2(\\sigma^2 + e_i^2)}\\right]\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now generate synthetic observations from this model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(42)  # for reproducibility\n",
    "N = 100  # we'll use more samples for the more complicated model\n",
    "mu_true, sigma_true = 1000, 15  # stochastic flux model\n",
    "\n",
    "F_true = stats.norm(mu_true, sigma_true).rvs(N)  # (unknown) true flux\n",
    "F = stats.poisson(F_true).rvs()  # observed flux: true flux plus Poisson errors.\n",
    "e = np.sqrt(F)  # root-N error, as above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.errorbar(F, np.arange(N), xerr=e, fmt='ok', ecolor='gray', alpha=0.5)\n",
    "ax.vlines([F_true], 0, N, linewidth=1, alpha=0.2)\n",
    "ax.set_xlabel(\"Flux\");ax.set_ylabel(\"measurement number\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, we can analytically optimize by setting the derivative with respect to the mean $\\mu$ equal to zero and solve for $\\mu$:\n",
    "\n",
    "$$\n",
    "\\mu_{est} =  \\frac{\\sum w_i F_i}{\\sum w_i};~~w_i = \\frac{1}{\\sigma^2 + e_i^2}\n",
    "$$\n",
    "\n",
    "Note that in this case the optima value for $\\mu$ depens on the value of $\\sigma$. They are correlated. In other words, the most likely value of the mean flux depends on how uncertain your measurements are.\n",
    "\n",
    "Because of the correlation between the parameters, straightforward analytical methods are not longer useful to arrive at the optimal values. We have to do this numerically. Scipy offers a way to do it. Rather than maximizing the likelihood, we minimize its negtive:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_likelihood(theta, F, e):\n",
    "    return -0.5 * np.sum(np.log(2 * np.pi * (theta[1] ** 2 + e ** 2))\n",
    "                         + (F - theta[0]) ** 2 / (theta[1] ** 2 + e ** 2))\n",
    "\n",
    "# maximize likelihood <--> minimize negative likelihood\n",
    "def neg_log_likelihood(theta, F, e):\n",
    "    return -log_likelihood(theta, F, e)\n",
    "\n",
    "from scipy import optimize\n",
    "theta_guess = [900, 5]\n",
    "theta_est = optimize.fmin(neg_log_likelihood, theta_guess, args=(F, e))\n",
    "print(\"\"\"\n",
    "      Maximum likelihood estimate for {0} data points:\n",
    "          mu={theta[0]:.0f}, sigma={theta[1]:.0f}\n",
    "      \"\"\".format(N, theta=theta_est))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bayesian approach\n",
    "\n",
    "Now let's go Bayesian. The likelihood has been defined above. We can now define the prior and the posterior, and use emcee to sample from the resulting posterior."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_prior(theta):\n",
    "    # sigma needs to be positive.\n",
    "    if theta[1] <= 0:\n",
    "        return -np.inf\n",
    "    else:\n",
    "        return 1\n",
    "\n",
    "def log_posterior(theta, F, e):\n",
    "    return log_prior(theta) + log_likelihood(theta, F, e)\n",
    "\n",
    "# same setup as above:\n",
    "ndim, nwalkers = 2, 50\n",
    "nsteps, nburn = 2000, 1000\n",
    "\n",
    "starting_guesses = np.random.rand(nwalkers, ndim)\n",
    "starting_guesses[:, 0] *= 2000  # start mu between 0 and 2000\n",
    "starting_guesses[:, 1] *= 20    # start sigma between 0 and 20\n",
    "\n",
    "sampler = emcee.EnsembleSampler(nwalkers, ndim, log_posterior, args=[F, e])\n",
    "sampler.run_mcmc(starting_guesses, nsteps)\n",
    "\n",
    "sample = sampler.chain  # shape = (nwalkers, nsteps, ndim)\n",
    "sample = sampler.chain[:, nburn:, :].reshape(-1, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us plot the samples and compare with the \"true\" values. Note that the true values fall in the area where most samples are. In Bayesian statistics, rather that estimates like the mean and the standard deviation, the final answer is the posterior itself, tha contains all the information about uncertainties."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from astroML.plotting import plot_mcmc\n",
    "fig = plt.figure()\n",
    "ax = plot_mcmc(sample.T, fig=fig, labels=[r'$\\mu$', r'$\\sigma$'], colors='k')\n",
    "ax[0].plot(sample[:, 0], sample[:, 1], ',k', alpha=0.1)\n",
    "ax[0].plot([mu_true], [sigma_true], 'o', color='red', ms=10);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Another example showing what happens behind the scene in Metropolis-Hastings\n",
    "This is taken from Ridlo W. Wibowo's blog. (http://nbviewer.jupyter.org/github/ridlo/stats_notebook/blob/master/mcmc_metropolis-hastings_pdfsampling.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "from scipy.integrate import quad\n",
    "from scipy.stats import norm\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pdftarget(x, norm=1):\n",
    "    return np.exp(0.4*(x-0.4)*(x-0.4) - 0.08*x*x*x*x)/norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sampling_mcmc_mh(xt, stepsize, nsamp):\n",
    "    samples = np.empty(nsamp)\n",
    "    accept = np.empty(nsamp)\n",
    "    for i in range(nsamp):\n",
    "        xprime = xt + stepsize*np.random.normal() # gaussian proposal distribution\n",
    "        a = pdftarget(xprime)/pdftarget(xt) # symmetric -> gaussian\n",
    "        if a >= 1.0:\n",
    "            xt = xprime\n",
    "            accept[i] = 1\n",
    "        else:\n",
    "            u = np.random.random()\n",
    "            if a >= u:\n",
    "                xt = xprime\n",
    "                accept[i] = 1\n",
    "            else:\n",
    "                accept[i] = 0 # reject xprime, xt = xt\n",
    "        \n",
    "        samples[i] = xt\n",
    "        \n",
    "    return samples, accept"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mu = 0\n",
    "sigma = 2\n",
    "nsamp = 10000\n",
    "\n",
    "samples, accept = sampling_mcmc_mh(mu, sigma, nsamp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nbins = 50\n",
    "xmin, xmax = -5, 5\n",
    "\n",
    "I = quad(pdftarget, -100, +100)\n",
    "x = np.linspace(xmin, xmax, 1000)\n",
    "y = pdftarget(x, I[0])\n",
    "\n",
    "plt.hist(samples, bins=nbins, normed=True, histtype=\"stepfilled\", color=\"blue\", alpha=0.5, linewidth=0)\n",
    "plt.plot(x, y, 'k')\n",
    "plt.xlim([xmin, xmax])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's make an animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_samples(samples, accept, stepsize, x, target_dist, xmin=-5, xmax=5, nbins=50, write=False, filename=\"plot_samp_mcmc.png\", trace=False):\n",
    "    nsamples = len(samples)\n",
    "    ofile = '/home/ridlo/project/stats/mcmc_sampling/'+filename \n",
    "    \n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111)\n",
    "    ax.set_xlim(xmin, xmax)\n",
    "    \n",
    "    #x = np.linspace(xmin, xmax, 1000) \n",
    "    #target_dist = pdftarget(x, normed=normalize)\n",
    "    max_propdist = norm.pdf(0, 0, stepsize) # to draw line\n",
    "    \n",
    "    ymax = 1.1*np.amax(target_dist)\n",
    "    ax.set_ylim(0, ymax)\n",
    "    \n",
    "    ax.plot(x, target_dist, 'k') # draw target dist line\n",
    "    if nsamples > 1:\n",
    "        ax.hist(samples, normed=True, bins=nbins, histtype=\"stepfilled\", color=\"blue\", alpha=0.5, linewidth=0)\n",
    "        if trace:\n",
    "            last_state = samples[-1] # last sample \n",
    "            last_acc = accept[-1]\n",
    "            prev_state = samples[-2]\n",
    "            prev_acc = accept[-2]\n",
    "            \n",
    "            ax.plot(x, norm.pdf(x, prev_state, stepsize), 'g') # draw previous propos dist\n",
    "            ax.axvline(x=prev_state, ymin=0, ymax=(max_propdist)/(ymax), c='g')\n",
    "            \n",
    "            color = 'r'\n",
    "            if last_acc > 0: color = 'k'     \n",
    "            ax.axvline(x=last_state, ymin=0, ymax=(norm.pdf(last_state, prev_state, stepsize))/(ymax), c=color)\n",
    "\n",
    "        ratio_accept = float(len(samples[accept>0]))/float(nsamples)\n",
    "        text = r'$n_{sample} = '+'{0:d}$'.format(nsamples)+'\\n'\n",
    "        text += r'$r_{accept} = '+'{0:0.2f}$'.format(ratio_accept)\n",
    "        ax.annotate(text, xy=(0.7, 0.97), xycoords='axes fraction', ha='left', va='top') \n",
    "    \n",
    "    if write:\n",
    "        plt.savefig(ofile, bbox_inches='tight', dpi=400); plt.close()\n",
    "    else:\n",
    "        plt.show(); plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test\n",
    "mu = 0\n",
    "sigma = 2\n",
    "nsamp = 1000\n",
    "samples, accept = sampling_mcmc_mh(mu, sigma, nsamp)\n",
    "\n",
    "I = quad(pdftarget, -100, +100)\n",
    "x = np.linspace(xmin, xmax, 1000)\n",
    "y = pdftarget(x, I[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_samples(samples, accept, sigma, x, y, trace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import YouTubeVideo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "YouTubeVideo(\"zL2lg_Nfi80\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}