mini-projects/mini-project_II_how_many_lines.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mini-project II: How many lines are there"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The overall project goal is to determine how many signal peaks there are in a noisy set of data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Learning goals:\n",
    "* Use advanced Monte Carlo sampling to generate posterior probability distributions and analyze the results.\n",
    "* Compute an evidence ratio and explain what it means.\n",
    "* Employ Bayesian evidence for model selection in the context of a prototype for an experimental nuclear-physics problem."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A spectral line problem"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See e.g. section 4.2 in Sivia for a similar problem formulation. In short, we have data from a spectroscopy experiment that supposedly shows a number of spectral lines. The ideal spectrum, as a function of a scaled coordinate $x$, can be expressed as\n",
    "\n",
    "$$ G(x) = \\sum_{j=1}^M A_j f(x,x_j),$$\n",
    "\n",
    "where $A_j$ is the amplitude of the $j$th line, and $x_j$ represents its position. Here we assume that all the spectral lines are Gaussians of width $W$, then\n",
    "\n",
    "$$ f(x,x_j) = \\frac{1}{\\sqrt{2\\pi W^2}} \\exp \\left[ - \\frac{(x-x_j)^2}{2 W^2} \\right]$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The measurement also includes a background signal which we take to be constant $B(x) = B$. \n",
    " \n",
    "We use the label '$k$' to enumerate the bins $\\{x_k\\}$ (do not confuse $x_k$ and $x_j$). The spectrum according to our model is therefore\n",
    " \n",
    " $$ F_k \\equiv F(x_k) = G(x_k) + B.$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The experimental data is denoted $\\{ D_k\\}$. This data also includes measurement errors $\\varepsilon$ that are assumed to be independent and identically distributed (IID) normal with the same variance $\\sigma_\\mathrm{exp}$. The measured data is then related to the ideal spectrum by\n",
    "\n",
    "$$ D_k \\equiv D(x_k) = G(x_k) + B + \\varepsilon.$$\n",
    "\n",
    "The task is to infer how many spectral lines ($M$) are in the experimental data, and their positions ($x_j$) and amplitudes ($A_j$). The magnitude of the background is not known before the analysis except that it must be $ \\leq 1$.\n",
    "\n",
    "Using scaled parameters the range of peak positions is $[0,1]$, and the range of amplitudes as well as background is also $[0,1]$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution strategy:\n",
    "Our model with $M$ lines has $2M+1$ model parameters that we denote by the vector $\\vec{\\alpha}$. These are the amplitudes, the positions, and the constant background. We order them as follows:\n",
    "\n",
    "$$ \\vec{\\alpha} = (A_0, x_0, A_1, x_1, \\ldots, B).$$\n",
    "\n",
    "The background strength is a *nuisance parameter* in the sense that we're not really interested in its value, we just need to marginalize over it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Subtasks:\n",
    "1. Formulate the problem of how many lines and what are the model parameters in Bayesian language.\n",
    "1. Using the expressions from step 1, try to derive the approximate posterior probability \n",
    "$$\n",
    "p(M | \\{ D_k \\}, I ) \\propto \n",
    "\\frac{M! (4 \\pi)^M }{\\left[ x_\\mathrm{max} - x_\\mathrm{min} \\right]^M \\left[ A_\\mathrm{max} - A_\\mathrm{min} \\right]^M \\sqrt{\\det(\\boldsymbol\\nabla \\boldsymbol\\nabla \\chi^2)} } \n",
    "\\exp \\left( - \\frac{\\chi^2_\\mathrm{min}}{2} \\right).\n",
    "$$\n",
    "1. Argue for how you would compute this probability. For the imagined numerical implementation you should consider that some model parameters enter non-linearly in your likelihood. Discuss with one of the teachers or facilitators.\n",
    "1. Generate data using the code snippet below, but plan to change values to explore the dependence of the analysis on the number and relative position of peaks, noise, background, etc. Note that the random seed is initialized to specific values in this notebook. You should start with this example, but should then run the data generator several times without the seeds to observe the degree of fluctuation.\n",
    "1. Use an MCMC sampler to perform parameter estimation for a known number of signal peaks. Make appropriate plots and analyze the posterior pdf (note that this is the full posterior now, not an approximation).\n",
    "1. Use PyMultinest to both infer parameters and compute the evidence for a model with $M$ peaks. Relevant code is given below (try to understand the interface). You might want to modify some arguments such as  `sampling_efficiency`,  `evidence_tolerance`, and `max_iter=1000` (where the latter is set rather low for starters). Do the results make sense?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import statements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.stats, scipy\n",
    "\n",
    "import seaborn as sns\n",
    "sns.set() # set default plot styles\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters known before the analysis (explore different values for these)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "W = 0.12 # The width of the spectral lines\n",
    "sigma_exp = 0.4 # Variance of IID experimental errors "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fixed prior ranges (Don't change these)\n",
    "xmin=0; xmax=1;\n",
    "Amin=0; Amax=1;\n",
    "Bmax=1.;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters that should be learned from the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(3) # For reproducibility\n",
    "\n",
    "M = 2 # Number of lines\n",
    "\n",
    "A0 = 10**np.random.uniform(low=-.5, high=0, size=M) # Amplitudes of the spectral lines\n",
    "\n",
    "X0 = np.ones(M)*np.random.uniform(low=.1, high=.4, size=1) # Position of the spectral lines\n",
    "\n",
    "# special treatment to place the lines close together\n",
    "for i in np.arange(1,M):\n",
    "    X0[i] += np.random.normal(loc=.3,scale=.1)\n",
    "    X0[i] = min(0.95,X0[i])\n",
    "    X0[i] = max(0.05,X0[i])\n",
    "\n",
    "B = np.random.uniform(0.05, 0.8) # A constant background signal\n",
    "\n",
    "print(\"True values:   Peak amplitude     Peak position\")\n",
    "print(\"-----------    --------------     -------------\")\n",
    "for iM in range(M):\n",
    "    print(f\"     peak #{iM}             {A0[iM]:.2f}              {X0[iM]:.2f}  \")\n",
    "print(f\" Background:             {B:.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(x,x0,w):\n",
    "    '''Simple Gaussian line shape. Position x0, width w.'''\n",
    "    return np.exp(- (x-x0)**2 / (2 * w**2)) / np.sqrt(2*np.pi*w**2)\n",
    "\n",
    "def G(x,amplitudes,positions,width=W):\n",
    "    '''Ideal spectrum composed of N spectral lines with input amplitudes \n",
    "    and positions (arrays of length N), and a single, fixed width.'''\n",
    "    ideal = np.zeros_like(x)\n",
    "    for a,x0 in zip(amplitudes,positions):\n",
    "        ideal += a * f(x,x0,width)\n",
    "    return ideal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x=np.linspace(xmin,xmax,100)\n",
    "plt.plot(x,G(x,A0,X0,W))\n",
    "plt.xlabel('$x$ [x-unit]')\n",
    "plt.ylabel('$y$ [y-unit]');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assert: \n",
    "1. unknown constant background ($B\\le 1$), unknown number of spectral lines with unknown amplitudes and positions.\n",
    "2. known, natural width (W) of the spectral lines\n",
    "3. known variance ($\\sigma_\\mathrm{exp}$) for the IID normal experimental errors\n",
    "4. known and relevant interval:\n",
    "   peak positions within [xmin=0, xmax=1] in position space, \n",
    "   peak amplitudes within [Amin=0, Amax=1], "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'Natural width of spectral lines:                W = {W}')\n",
    "print(f'Variance for IID normal exp errors:             s = {sigma_exp}')\n",
    "print(f'Relevant range in position space:    [xmin, xmax] = [{xmin:.1f}, {xmax:.1f}]')\n",
    "print(f'Relevant range for peak amplitudes:  [Amin, Amax] = [{Amin:.1f}, {Amax:.1f}]')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate experimental data\n",
    "np.random.seed(42) # For reproducibility\n",
    "Ndata = 100 # Number of data points\n",
    "xk = np.linspace(xmin,xmax,Ndata)\n",
    "sk=np.ones_like(xk)*sigma_exp # We assume that the variance is the same for all k\n",
    "ek = np.random.normal(loc=0.0, scale=sigma_exp, size=Ndata)\n",
    "Dk = G(xk,A0,X0,W) + B + ek\n",
    "# Data saved in a file (but also printed below)\n",
    "data = [xk,Dk,sk]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.errorbar(xk,Dk,yerr=sk,fmt='ok',ecolor='gray')\n",
    "plt.xlabel('$x$ [x-unit]')\n",
    "plt.ylabel('$y$ [y-unit]');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameter estimation with your favorite MCMC sampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fill this in!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PyMultinest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pymultinest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# number of lines\n",
    "M=2\n",
    "# number of dimensions our problem has\n",
    "n_params = 2*M + 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prior(cube, ndim, nparams):\n",
    "    '''Prior should transform the unit cube into the parameter cube. \n",
    "    We have a uniform prior [0,1] in all dimensions.'''\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loglike(cube, ndim, nparams):\n",
    "    '''Log likelihood in PyMultinest format. Uses the (hyper)cube object from the prior.'''\n",
    "    background  = cube[-1]\n",
    "    amplitudes=[cube[i] for i in np.arange(0, ndim-1,2)]\n",
    "    positions=[cube[i] for i in np.arange(1, ndim-1,2)]\n",
    "    ymodel = G(data[0],amplitudes,positions) + background \n",
    "    loglikelihood = (-0.5 * ((ymodel - data[1]) / sigma_exp)**2).sum()\n",
    "    return loglikelihood"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run MultiNest\n",
    "base='pm_out_'\n",
    "prefix = f\"{base}case_M{M}\"\n",
    "pymultinest.run(loglike, prior, n_params, outputfiles_basename=prefix, resume = False, \\\n",
    "                verbose = True, evidence_tolerance=1.0, sampling_efficiency=0.8, max_iter=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the PyMultinest analyzer.\n",
    "a = pymultinest.Analyzer(outputfiles_basename=prefix, n_params = n_params)\n",
    "# get the best fit (highest likelihood) point\n",
    "bestfit_params = a.get_best_fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_lnZ = a.get_stats()['global evidence']\n",
    "print('************************')\n",
    "print('MAIN RESULT: Evidence Z ')\n",
    "print('************************')\n",
    "print(f'  log Z for model with {M} lines = {a_lnZ / np.log(10):.1f}')"
   ]
  },
  {
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}