Week2/FittingMethods/FittingMethods_original.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fitting with ChiSquare, Binned Likelihood, and Unbinned Likelihood\n",
    "\n",
    "This is an introduction to likelihood fits, both binned and unbinned, trying to show the following points:\n",
    "1. How likelihood fits work, both binned and unbinned.\n",
    "2. When to use which type of fits, also including ChiSquare.\n",
    "3. See how the amount of statistics influences the above choice.\n",
    "\n",
    "Comparisons are naturally done with the ChiSquare. Also, this program serves as an introduction to fitting with iMinuit.\n",
    "\n",
    "The exercise is made to play around with the (statistics of) the signal and background distributions, so that you can get a feel for the advantages and drawbacks of each method... and make sure that you feel comfortable in applying all three.\n",
    "\n",
    "### References:\n",
    "- Barlow: 5.3 + 5.4 + 5.6\n",
    "- Bevington: Chapter 10\n",
    "\n",
    "### Authors: \n",
    "- Troels C. Petersen (Niels Bohr Institute)\n",
    "\n",
    "### Date:    \n",
    "- 23-11-2023 (latest update)\n",
    "\n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np                                     # Matlab like syntax for linear algebra and functions\n",
    "import matplotlib.pyplot as plt                        # Plots and figures like you know them from Matlab\n",
    "import seaborn as sns                                  # Make the plots nicer to look at\n",
    "from iminuit import Minuit                             # The actual fitting tool, better than scipy's\n",
    "import sys                                             # Module to see files and folders in directories\n",
    "from scipy import stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.path.append('../../External_Functions')\n",
    "from ExternalFunctions import UnbinnedLH, BinnedLH, Chi2Regression\n",
    "from ExternalFunctions import nice_string_output, add_text_to_ax   # Useful functions to print fit results on figure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below are some of the general program settings, which are good to have in one place:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = np.random # Random generator\n",
    "r.seed(42)    # Set a random seed (but a fixed one)\n",
    "\n",
    "save_plots = False # For now, don't save plots (once you trust your code, switch on)\n",
    "verbose = True     # For now, print a lot of output (once you trust your code, switch off)\n",
    "Nverbose = 10      # But only print a lot for the first 10 random numbers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Truth values used:\n",
    "\n",
    "We choose some \"god given\" values (the \"truth\") to simulate the data with. The goal of the fits is to estimate these values most accurately, as if you had data, and wanted to extract these values from it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Signal parameters:\n",
    "Npeak = 2000           # Number of random Gaussian points\n",
    "mu_peak = 3.0        # Peak location\n",
    "sigma_peak = 0.3     # Peak width\n",
    "\n",
    "# Background parameters:\n",
    "Nbkg  = 5000           # Number of random Exponential points\n",
    "tau_bkg = 3.0        # Decay constant (could be anything!)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the data:\n",
    "\n",
    "We create some data with a signal peak on an exponential background:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create signal:\n",
    "x_peak = r.normal(loc=mu_peak, scale=sigma_peak, size=Npeak)\n",
    "\n",
    "# Create background:\n",
    "x_bkg  = r.exponential(tau_bkg, Nbkg)\n",
    "\n",
    "# Combine the two:\n",
    "x_all = np.concatenate((x_peak, x_bkg), axis=0)\n",
    "if (verbose) :\n",
    "    print(x_all)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Histogram of data:\n",
    "\n",
    "Since we are considering both ChiSquare (chi2) and Binned Likelihood (bllh), we start with the histogram, which is the input for both of these. Of course, the Unbinned Likelihood (ullh) doesn't care about binning!\n",
    "\n",
    "### Note on binning:\n",
    "It is very important to specify - in a well thought manner - the range and binning to use. Many problems could have been avoided by a few minutes (or even seconds!) of consideration from the beginning. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# General input:\n",
    "Nbins = 100\n",
    "xmin, xmax = 0, 10\n",
    "binwidth = (xmax-xmin)/Nbins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create just a single figure and axes, and a (classic) histogram:\n",
    "fig, ax = plt.subplots(figsize=(16, 6))\n",
    "hist = ax.hist(x_all, bins=Nbins, range=(xmin, xmax), histtype='step', linewidth=2, color='red', label='Data, normal histogram')\n",
    "\n",
    "# Find the x, y and error on y (sy) given the histogram:\n",
    "counts, bin_edges = np.histogram(x_all, bins=Nbins, range=(xmin, xmax))\n",
    "x = (bin_edges[1:] + bin_edges[:-1])/2\n",
    "y = counts\n",
    "sy = np.sqrt(counts)   # NOTE: We (naturally) assume that the bin count is Poisson distributed.\n",
    "                       # This is an approximation, since there is a low count in the last bins.\n",
    "\n",
    "# Did you make sure, that all bins were non-zero???\n",
    "# x = (bin_edges[1:][counts>0] + bin_edges[:-1][counts>0])/2\n",
    "# y = counts[counts>0]\n",
    "# sy = np.sqrt(counts[counts>0])   # NOTE: We (naturally) assume that the bin count is Poisson distributed.\n",
    "\n",
    "# Now create a histogram with uncertainties (better, I would argue):\n",
    "ax.errorbar(x, y, yerr=sy, xerr=0.0, label='Data, with Poisson errors', fmt='.k',  ecolor='k', elinewidth=1, capsize=1, capthick=1)\n",
    "\n",
    "# Set the figure texts; xlabel, ylabel and title.\n",
    "ax.set(xlabel=\"Random numbers\",           # the label of the y axis\n",
    "       ylabel=\"Frequency / 10\",           # the label of the y axis\n",
    "       title =\"Distribution of Gaussian and exponential numbers\")    # the title of the plot\n",
    "ax.legend(loc='best', fontsize=20);       # could also be # loc = 'upper right' e.g."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fit to the data / histogram (1D):\n",
    "\n",
    "Having created the data, we now want to fit the distributions in three ways:\n",
    "1. ChiSquare fit\n",
    "2. Binned Likelihood fit\n",
    "3. Unbinned Likelihood fit\n",
    "\n",
    "We first define the function to be fitted with, a Gaussian and and Exponential PDF in this case:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define background PDF:\n",
    "def exp_pdf(x, tau):\n",
    "    \"\"\"Exponential with lifetime tau\"\"\"\n",
    "    return 1.0 / tau * np.exp(-x/tau)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define signal PDF:\n",
    "def gauss_pdf(x, mu, sigma) :\n",
    "    \"\"\"Gaussian\"\"\"\n",
    "    return 1.0 / np.sqrt(2*np.pi) / sigma * np.exp( -0.5 * (x-mu)**2 / sigma**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define your actual fitting function, with both background and signal in:\n",
    "# which is NOT normalised but has normalisation constants \"N\" in, and includes the bin width:\n",
    "def fit_pdf(x, Nexp, tau, Ngauss, mu, sigma) :\n",
    "    \"\"\"Exponential + Gaussian\"\"\"\n",
    "    return Nexp * binwidth * exp_pdf(x, tau) + Ngauss * binwidth * gauss_pdf(x, mu, sigma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define your actual fitting function, with both background and signal in,\n",
    "# which is NOT normalised but has normalisation constants \"N\" in, but does NOT includes the bin width:\n",
    "def fit_pdf_llhfit(x, Nexp, tau, Ngauss, mu, sigma) :\n",
    "    \"\"\"Exponential + Gaussian\"\"\"\n",
    "    return Nexp * exp_pdf(x, tau) + Ngauss * gauss_pdf(x, mu, sigma)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having defined our fitting function, we now want to make our three types of fits. Make sure that you understand the difference between the three.\n",
    "\n",
    "It is important to use the same bins and bounds as the original histogram. In this case we are also treating the number of events/random numbers as being a fit variable.\n",
    "\n",
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ChiSquare fit:\n",
    "\n",
    "### Setting up a fit:\n",
    "\n",
    "We initialize our fitting object, `chi2fit`, by using the `Chi2Regression` function from External Functions. You can also define your own Chi2 function (see introduction to plotting and fitting in Week0 and below). This first takes the function to fit as input, then the input data.\n",
    "Following this, we make a Minuit object, i.e. the minimisation itself. This is done using __[iMinuit](https://iminuit.readthedocs.io/en/latest/)__ which is a Python-wrapper for the wonderful minimization tool Minuit developed by CERN. Minuit requires the fitting object (here \"chi2fit\") and then GOOD input values and fitting options. Finally, we ask Minuit to actually find the minimum, and we have added a check to see, if this converged or not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "Minuit.print_level = 1    # Print result of fits (generally - can also be moved down to each fit instance)\n",
    "\n",
    "# Defining Chi2 calculation:\n",
    "def chi2_owncalc(Nexp, tau, Ngauss, mu, sigma) :\n",
    "    y_fit = fit_pdf(x, Nexp, tau, Ngauss, mu, sigma)\n",
    "    chi2 = np.sum(((y - y_fit) / sy)**2)\n",
    "    return chi2\n",
    "\n",
    "# Alternatively, use the Chi2 regression from External Functions:\n",
    "chi2fit = Chi2Regression(fit_pdf, x, y, sy)\n",
    "\n",
    "# NOTE: Did you remember to ensure, that all bins had entries in them, i.e. that y>0?\n",
    "\n",
    "minuit_chi2 = Minuit(chi2_owncalc, Nexp=Nbkg, tau=2.9, Ngauss=Npeak, mu=3.1, sigma=0.3)\n",
    "minuit_chi2.errordef = 1.0     # This is the definition for ChiSqaure fits\n",
    "minuit_chi2.migrad()           # This is where the minimisation is carried out! Put \";\" at the end to void output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Extracting fit results - values and $\\chi^2$:\n",
    "\n",
    "Once the fit has converged, the results and $\\chi^2$ probability should be evaluated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First, check if the result was a valid minimum:\n",
    "if (not minuit_chi2.fmin.is_valid) :\n",
    "    print(\"  WARNING: The ChiSquare fit DID NOT converge!!! \")    \n",
    "    \n",
    "# Short hand way of accessing the fit parameters:\n",
    "fit_Ngauss, fit_mu, fit_sigma, fit_Nexp, fit_tau = minuit_chi2.values[:]   # The fitted values of the parameters\n",
    "\n",
    "# Loop to get both parameter values and uncertainties:\n",
    "for name in minuit_chi2.parameters :\n",
    "    value, error = minuit_chi2.values[name], minuit_chi2.errors[name]\n",
    "    print(f\"Fit value: {name} = {value:.5f} +/- {error:.5f}\")\n",
    "\n",
    "# Get Chi2 value:\n",
    "chi2_value = minuit_chi2.fval            # The value minimised, i.e. Chi2 or -2*LogLikeliHood (LLH) value\n",
    "\n",
    "# Get number of degrees-of-freedom (Ndof):\n",
    "N_NotEmptyBin = np.sum(y > 0)\n",
    "Ndof_value = N_NotEmptyBin - minuit_chi2.nfit\n",
    "\n",
    "Prob_value = stats.chi2.sf(chi2_value, Ndof_value) # The chi2 probability given N_DOF degrees of freedom\n",
    "print(f\"Chi2 value: {chi2_value:.1f}   Ndof = {Ndof_value:.0f}    Prob(Chi2,Ndof) = {Prob_value:5.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Drawing fit result on top of data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create just a single figure and axes, along with a histogram with uncertainties:\n",
    "fig, ax = plt.subplots(figsize=(16, 8))  # figsize is in inches\n",
    "ax.errorbar(x, y, yerr=sy, xerr=0.0, label='Data, with Poisson errors', fmt='.k',  ecolor='k', elinewidth=1, capsize=1, capthick=1)\n",
    "\n",
    "# Set the figure texts; xlabel, ylabel and title.\n",
    "ax.set(xlabel=\"Some feature X (some unit)\", # the label of the y axis\n",
    "       ylabel=\"Frequency / 0.1\",  # the label of the y axis\n",
    "       title=\"Distribution of Gaussian and exponential numbers\", # the title of the plot\n",
    "       ylim=[0.0,None]) # Setting the miminum to zero\n",
    "\n",
    "# Adding fit function to plot:\n",
    "x_axis = np.linspace(xmin, xmax, 1000)\n",
    "ax.plot(x_axis, fit_pdf(x_axis, *minuit_chi2.values[:]), '-r', label='Chi2 fit model result') \n",
    "\n",
    "# Adding fit results to plot:\n",
    "d = {'Ngauss':   [minuit_chi2.values['Ngauss'], minuit_chi2.errors['Ngauss']],\n",
    "     'mu':       [minuit_chi2.values['mu'], minuit_chi2.errors['mu']],\n",
    "     'sigma':       [minuit_chi2.values['sigma'], minuit_chi2.errors['sigma']],\n",
    "     'Nexp':     [minuit_chi2.values['Nexp'], minuit_chi2.errors['Nexp']],\n",
    "     'tau':         [minuit_chi2.values['tau'], minuit_chi2.errors['tau']],\n",
    "     'Chi2':     chi2_value,\n",
    "     'ndf':      Ndof_value,\n",
    "     'Prob':     Prob_value,\n",
    "    }\n",
    "\n",
    "text = nice_string_output(d, extra_spacing=2, decimals=3)\n",
    "add_text_to_ax(0.62, 0.95, text, ax, fontsize=20)\n",
    "ax.legend(loc='lower left', fontsize=18); # could also be # loc = 'upper right' e.g.\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Binned Likelihood fit:\n",
    "\n",
    "Next we try to fit the same data with the same function in a *binned likelihood fit*. Note that the fitting object now takes the x-values themselves (not a histogram), but requires numbers of bins and a range/bound."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bllhfit = BinnedLH(fit_pdf_llhfit, x_all, bins=Nbins, bound=(xmin, xmax), extended=True)\n",
    "minuit_bllh = Minuit(bllhfit, Nexp=Nbkg, tau=2.9, Ngauss=Npeak, mu=3.1, sigma=0.3)\n",
    "minuit_bllh.errordef = 0.5     # Value for likelihood fits\n",
    "minuit_bllh.migrad()           # Perform the actual fit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Unbinned Likelihood fit:\n",
    "\n",
    "For the *unbinned likelihood fit*, the input is again the x-values themselves, and bounds are possible, but no binning is suggested, as this is... well... unbinned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "ullhfit = UnbinnedLH(fit_pdf_llhfit, x_all, bound=(xmin, xmax), extended=True)\n",
    "minuit_ullh = Minuit(ullhfit, Nexp=Nbkg, tau=2.9, Ngauss=Npeak, mu=3.1, sigma=0.3)\n",
    "minuit_ullh.errordef = 0.5     # Value for likelihood fits\n",
    "minuit_ullh.migrad()           # Perform the actual fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if (not minuit_bllh.fmin.is_valid) :\n",
    "    print(\"  WARNING: The Binned Likelihood fit DID NOT converge!!! \")    \n",
    "\n",
    "if (not minuit_ullh.fmin.is_valid) :\n",
    "    print(\"  WARNING: The Unbinned Likelihood fit DID NOT converge!!! \")\n",
    "\n",
    "# Notice how both results are plotted with the \"usual\" PDF, as the fit parameters should be correct/similar. \n",
    "ax.plot(x_axis, fit_pdf(x_axis, *minuit_bllh.values[:]), '-b', label='Binned LLH fit model result') \n",
    "ax.plot(x_axis, fit_pdf(x_axis, *minuit_ullh.values[:]), '-g', label='Unbinned LLH fit model result') \n",
    "fig    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Questions:\n",
    "\n",
    "This exercise is meant for you to see to what extend the three fitting methods give the same results when varying the input parameters, especially when the number of signal and background events is low, and also when the binning is (too) coarse.\n",
    "\n",
    "1) Start by drawing all three fits on top of the data. Do they (by eye) give the same result, and does the fit generally look \"good\"? How do you judge that? Also, compare the fitted values of the three fits, and see if they agree within uncertainties?\n",
    "\n",
    "2) Now repeat question 1) for lower statistics cases, for example Ngauss = 20 and Nexp = 50. See to what extend the different methods yields reasonable results (i.e. in accordance with the input values). You will encounter problems to be mended by yourself, such as empty bins (should not be passed to the ChiSquare dividing by \"observed events\", but is fine for the LLH fits) and poor initial values.\n",
    "Possibly also alter the number of bins to 50 or 25 - is that a good choice? Why/why not?\n",
    "\n",
    "---\n",
    "\n",
    "### Advanced question:\n",
    "\n",
    "3) Knowing the PDFs used for producing the data, try to see, if you can produce a goodness-of-fit measure for the unbinned likelihood fit by repeating this type of fit many times on (re-)simulated data.\n",
    "\n",
    "4) Improve on the printing of results on the plots to better align and have a choice of decimals :-)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Learning points:\n",
    "\n",
    "The exercise is meant as an illustration of the three main fitting paradigms:\n",
    "1. ChiSquare (binned)\n",
    "2. Likelihood (binned)\n",
    "3. Likelihood (unbinned)\n",
    "\n",
    "While the ChiSquare is almost always recommendable to start out with, low statistics makes it non-optimal. If possible, use an unbinned likelihood, and if the data is already binned, use a binned likelihood."
   ]
  }
 ],
 "metadata": {
  "executable": "/usr/bin/env python",
  "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.11"
  },
  "main_language": "python"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}