smd_handson_fitting.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"font-size: 3rem; font-weight: bold;\">SMD Hands-On Estimators / Fitting</span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%javascript\n",
    "$.getScript('https://kmahelona.github.io/ipython_notebook_goodies/ipython_notebook_toc.js')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h1 id=\"tocheading\">Table of Contents</h1>\n",
    "<div id=\"toc\"></div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib widget\n",
    "\n",
    "from IPython.display import display, HTML\n",
    "display(HTML(\"<style>.container { width:95% !important; }</style>\"))\n",
    "\n",
    "plt.rcParams['figure.figsize'] = (7.5, 5)\n",
    "plt.rcParams['figure.constrained_layout.use'] = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cov_to_corr(cov):\n",
    "    '''Convert covariance to correlation matrix\n",
    "    \n",
    "    Taken from: https://math.stackexchange.com/a/300775/892886\n",
    "    '''\n",
    "    D = np.diag(1 / np.sqrt(np.diag(cov)))\n",
    "    return D @ cov @ D"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Least Squares\n",
    "\n",
    "## Analytic solution for linear combination of functions\n",
    "\n",
    "(Analog to the exercise on the last sheet)\n",
    "\n",
    "\n",
    "Here, we will fit a function of the form\n",
    "\n",
    "$$\n",
    "f(x) = p_0 + p_1 \\cdot \\sin(x) + p_2 \\cdot \\cos(x)\n",
    "$$\n",
    "\n",
    "to our data.\n",
    "\n",
    "This function is a linear combination of basis functions:\n",
    "$$\n",
    "f(x) = \\sum_{i=0}^2 p_i f_i(x)\n",
    "$$\n",
    "\n",
    "with\n",
    "\n",
    "$$\n",
    "f_0(x) = 1, \\quad f_1(x) = \\sin(x), \\quad f_2(x) = \\cos(x)\n",
    "$$\n",
    "\n",
    "In this case, we can use the analytic solution to the least squares optimization problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def linear_combination(x, funcs, parameters):\n",
    "    '''Evaluate a linear combination of basis functions\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x: number or np.ndarray\n",
    "        The point or points at which to evaluate\n",
    "    funcs: iterable of callables\n",
    "        The basis functions\n",
    "    parameters: iterable of numbers\n",
    "        The coefficients\n",
    "    '''      \n",
    "    return np.sum([p * f(x) for p, f in zip(parameters, funcs, strict=True)], axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define our linear model\n",
    "funcs = [np.ones_like, np.sin, np.cos]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# create some randomized example data points\n",
    "rng = np.random.default_rng(1337)\n",
    "\n",
    "N = 100\n",
    "true_parameters = np.array([2, 1, 0.5])\n",
    "y_unc = rng.uniform(0.1, 0.4, N)\n",
    "\n",
    "x = np.linspace(0, 4 * np.pi, N)\n",
    "y = linear_combination(x, funcs, true_parameters)\n",
    "\n",
    "# add some noise for to simulate measurement uncertainty\n",
    "y += rng.normal(0, y_unc)\n",
    "\n",
    "plt.figure()\n",
    "plt.errorbar(x, y, yerr=y_unc, ls='', marker='.', color='k')\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('y')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create the design matrix $\\boldsymbol{A}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def design_matrix(funcs, x):\n",
    "    '''Create the design matrix for a linear least squares problem'''\n",
    "    return np.column_stack([f(x) for f in funcs])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = design_matrix(funcs, x)\n",
    "A[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the weight matrix $\\boldsymbol{W} = \\mathrm{Cov}^{-1}(\\boldsymbol{y})$ of the measurements.\n",
    "\n",
    "Here we assume that all measured points do not have an uncertainty in $x$ and that the $y$ values are statistically independent (no off-diagonal entries in $\\mathrm{Cov}(\\boldsymbol{y})$).\n",
    "\n",
    "This is a very strong assumption.\n",
    "\n",
    "The linear least squares method yields biased results if the `Covariance` is estimated from the data points themselves. \n",
    "\n",
    "More on that further down in the muon example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# All measurements have a known uncertainty and no correlations\n",
    "cov_y = np.diag(y_unc**2)\n",
    "               \n",
    "W = np.linalg.inv(cov_y)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "plot = ax.matshow(W)\n",
    "fig.colorbar(plot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "W[:5, :5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solve the linear least squares problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_linear_least_squares(A, W, y):\n",
    "    \"\"\"Solve the linear least sqaures problem\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    A : np.ndarray\n",
    "        Design matrix\n",
    "    W : np.ndarray\n",
    "        Weight matrix\n",
    "    y : np.ndarray\n",
    "        Vector of y values\n",
    "    \"\"\"\n",
    "    cov = np.linalg.inv(A.T @ W @ A)\n",
    "    parameters = cov @ A.T @ W @ y\n",
    "    return parameters, cov\n",
    "\n",
    "\n",
    "parameters, cov = solve_linear_least_squares(A, W, y)\n",
    "parameters, cov"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate the $\\chi^2$ over the number of degrees of freedom:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def chisquare_over_ndf(y, A, parameters):\n",
    "    residuals = (y - A @ parameters)\n",
    "    sum_residuals = (residuals.T @ W @ residuals)\n",
    "\n",
    "    ndf = len(y) - len(parameters)\n",
    "    return sum_residuals / ndf\n",
    "\n",
    "\n",
    "chisquare_over_ndf(y, A, parameters), chisquare_over_ndf(y, A, true_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "putting it all together:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def linear_least_squares(x, y, funcs, cov_y=None):\n",
    "    \"\"\"\n",
    "    Perform a linear least squares fit.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x : np.ndarray[ndim=1]\n",
    "        Vector of x values\n",
    "    y : np.ndarray[ndim=1]\n",
    "        Vector of y values\n",
    "    funcs : Sequence[Callable]\n",
    "        The basis functions\n",
    "    cov_y : Optional[np.ndarray[ndim=2]]\n",
    "        The covariance matrix of the y values\n",
    "        \n",
    "    Returns :\n",
    "    parameters : np.ndarray[ndim=1]\n",
    "        The estimated parameters\n",
    "    cov : np.ndarray[ndim=2]\n",
    "        Covariance matrix of the parameters\n",
    "    chisq_ndf : float\n",
    "        The chisquare over the number of degrees of freedom\n",
    "        of the fit result.\n",
    "    \"\"\"\n",
    "    A = design_matrix(funcs, x)\n",
    "    \n",
    "    if cov_y is not None:\n",
    "        W = np.linalg.inv(cov_y)\n",
    "    else:\n",
    "        W = np.eye(len(y))\n",
    "    \n",
    "    parameters, cov = solve_linear_least_squares(A, W, y)\n",
    "    \n",
    "    chisquare_ndf = chisquare_over_ndf(y, A, parameters)\n",
    "    \n",
    "    return parameters, cov, chisquare_ndf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Result:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters, cov, chisquare_ndf = linear_least_squares(x, y, funcs, cov_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_fit = np.linspace(0, 4 * np.pi, 1000)\n",
    "y_fit = linear_combination(x_fit, funcs, parameters)\n",
    "y_truth = linear_combination(x_fit, funcs, true_parameters)\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.errorbar(x, y, yerr=y_unc, ls='', marker='.', label='Data', color='k')\n",
    "\n",
    "ax.plot(x_fit, y_fit, label='Fit-Result')\n",
    "ax.plot(x_fit, y_truth, label='Truth')\n",
    "\n",
    "ax.set(\n",
    "    title=(\n",
    "        rf'$f(x) = {parameters[0]:.2f} + {parameters[1]:.2f} \\cdot \\sin(x) + {parameters[2]:.2f} \\cdot \\cos(x)'\n",
    "        rf', \\quad \\chi^2_\\mathrm{{ndf}} = {chisquare_ndf:.2f}$'\n",
    "    ),\n",
    "    xlabel='x',\n",
    "    ylabel='y',\n",
    ")\n",
    "\n",
    "\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "corr = cov_to_corr(cov)\n",
    "\n",
    "fig = plt.figure(constrained_layout=True)\n",
    "ax = fig.add_subplot(1, 1, 1)\n",
    "m = ax.matshow(corr, cmap='RdBu_r', vmin=-1, vmax=1)\n",
    "fig.colorbar(m)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numerical solution for non-linear functions\n",
    "\n",
    "\n",
    "If the function is not a linear combination of basis functions, the solution\n",
    "can only be found numerically.\n",
    "\n",
    "From the lab courses, you probably know the `scipy.optimize.curve_fit` function, which does exactly this.\n",
    "\n",
    "The Nelder-Mead-Algorithm use by `curve_fit` has the nice property that it is guaranteed to find the analytical\n",
    "solution, if it exists.\n",
    "\n",
    "\n",
    "- https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.curve_fit.html\n",
    "\n",
    "### Example application to our linear problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import curve_fit\n",
    "\n",
    "\n",
    "def func(x, p1, p2, p3):\n",
    "    return p1 + p2 * np.sin(x) + p3 * np.cos(x)\n",
    "\n",
    "\n",
    "# absolute_sigma prevents scaling of errors to match χ²/ndf=1\n",
    "parameters_numeric, cov_numeric = curve_fit(\n",
    "    func,x, y,\n",
    "    sigma=np.full(N, y_unc),\n",
    "    absolute_sigma=True,\n",
    ")\n",
    "\n",
    "print(parameters, parameters_numeric, cov, cov_numeric, sep='\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_fit = np.linspace(0, 4 * np.pi, 1000)\n",
    "y_fit = func(x_fit, *parameters)\n",
    "y_num = func(x_fit, *parameters_numeric)\n",
    "y_truth = func(x_fit, *true_parameters)\n",
    " \n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.errorbar(x, y, yerr=y_unc, ls='', marker='.', label='Data', color='k')\n",
    "\n",
    "ax.plot(x_fit, y_fit, label='Fit-Result (Analytic)')\n",
    "ax.plot(x_fit, y_num, label='Fit-Result (Numeric)')\n",
    "ax.plot(x_fit, y_truth, label='Truth')\n",
    "\n",
    "ax.set(\n",
    "    title=(\n",
    "        rf'$f(x) = {parameters[0]:.2f} + {parameters[1]:.2f} \\cdot \\sin(x) + {parameters[2]:.2f} \\cdot \\cos(x)'\n",
    "        rf', \\quad \\chi^2_\\mathrm{{ndf}} = {chisquare_ndf:.2f}$'\n",
    "    ),\n",
    "    xlabel='x',\n",
    "    ylabel='y',\n",
    ")\n",
    "\n",
    "\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Why not simply always chose the numerical solution then?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "linear_least_squares(x, y, funcs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "curve_fit(func, x, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example with non-linear function \n",
    "\n",
    "\n",
    "From the lab courses: single slit diffraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "df = pd.read_csv('resources/spalt.csv')\n",
    "\n",
    "\n",
    "LASER_WAVELENGTH_NM = 632.8e-9\n",
    "\n",
    "def theory(phi, A0, b):\n",
    "    return (A0 * b * np.sinc(b * np.sin(phi) / LASER_WAVELENGTH_NM))**2\n",
    "\n",
    "\n",
    "# first try with default initial guess (1 for every parameter)\n",
    "p0 = None\n",
    "\n",
    "# now with an \"educated guess\" based on the data and knowledge of the\n",
    "# order of magnitude of the slit size\n",
    "# p0 = [np.sqrt(df['I'].max()) / 1e-4, 1e-4]\n",
    "\n",
    "params, cov = curve_fit(theory, df['phi'], df['I'], p0=p0)\n",
    "\n",
    "\n",
    "x = np.linspace(-0.03, 0.03, 501)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(x, theory(x, *params), label='Fit')\n",
    "\n",
    "ax.plot(df['phi'], df['I'], '.', label='Daten')\n",
    "\n",
    "ax.set(\n",
    "    xlabel=r'$\\varphi \\,\\, / \\,\\, \\mathrm{rad}$',\n",
    "    ylabel=r'$I \\,\\, / \\,\\, \\mathrm{A}$',\n",
    ")\n",
    "ax.legend(loc='best')\n",
    "\n",
    "None # to not have mpl objects in the output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Maximimum-Likelihood-Method\n",
    "\n",
    "\n",
    "## Unbinned Fit of Probability Densities\n",
    "\n",
    "Strongly simplified example of a CERN-Like analysis.\n",
    "\n",
    "\n",
    "We are looking for the mass-peak of a (normally distributed) particle.\n",
    "\n",
    "We also observed a background, which in this simplified example is assumed to be exponentially distributed.\n",
    "\n",
    "We create a simple \"Toy\"-Dataset using Monte Carlo methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(42)\n",
    "\n",
    "E_MIN = 75\n",
    "E_MAX = 175\n",
    "\n",
    "# normally distributed signal\n",
    "higgs_signal = np.random.normal(126, 5, 500)\n",
    "\n",
    "# exponentially distributed background\n",
    "background = np.random.exponential(50, size=20000)\n",
    "\n",
    "# combine signal and background\n",
    "measured = np.append(higgs_signal, background)\n",
    "\n",
    "# remove events outside of \"detector range\"\n",
    "in_range = (E_MIN <= measured) & (measured <= E_MAX)\n",
    "measured = measured[in_range]\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.hist(measured, bins=100)\n",
    "ax.set_xlabel('$m \\,/\\, \\mathrm{GeV}$')\n",
    "ax.margins(x=0)\n",
    "None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Definition of the negative Log-Likelihood\n",
    "\n",
    "We create a superposition of two probability densities, each with a proportion of $p$ and $1 - p$ respectively.\n",
    "\n",
    "\n",
    "We also need to normalize the densities to the observed interval.\n",
    "\n",
    "In this special case, we ignore this for the normal distribution, assuming that it is fully contained in the measurement interval.\n",
    "\n",
    "So:\n",
    "\n",
    "\\begin{align}\n",
    "P_1 &= N(\\mu, \\sigma) \\\\[1ex]\n",
    "P_2 &= \\frac{1}{\\exp(-E_\\mathrm{min} / \\tau) - \\exp(-E_\\mathrm{max} / \\tau)} \\exp(- E / \\tau) \\\\[1ex]\n",
    "P(E | p, \\mu, \\sigma, \\tau) &= p \\cdot P_1(E, \\mu, \\sigma) + (1 - p) P_2(E | \\tau)) \\\\[1ex]\n",
    "\\mathcal{L}(p, \\mu, \\sigma, \\tau) &= \\prod_i P(E_i | p, \\mu, \\sigma, \\tau) \\\\[1ex]\n",
    "-\\log\\mathcal{L}(p, \\mu, \\sigma, \\tau) &= -\\sum_i \\log(P(E_i | p, \\mu, \\sigma, \\tau))\n",
    "\\end{align}\n",
    "\n",
    "In code, using the distribution classes from `scipy.stats`, it looks like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import norm, expon\n",
    "\n",
    "\n",
    "def pdf(x, mean, std, tau, p):\n",
    "    mass_peak = p * norm.pdf(x, mean, std)\n",
    "    \n",
    "    expon_norm = np.exp(-E_MIN / tau) - np.exp(-E_MAX / tau)\n",
    "    background = (1 - p)  / expon_norm * expon.pdf(x, scale=tau)\n",
    "    \n",
    "    return mass_peak + background"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution using `scipy.optimize.minimize` und `numdifftools.Hessian`\n",
    "\n",
    "\n",
    "Using `scipy.optimize.minimize`, we can minimize arbitrary functions.\n",
    "\n",
    "\n",
    "The functions need to have an array of the fit parameters as first argument.\n",
    "\n",
    "Further arguments can be passed using the `args` argument of `minimize`.\n",
    "\n",
    "\n",
    "Naïvely, our negative Log-Likelihood for use with scipy looks like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def neg_log_likelihood(parameters, data):\n",
    "    # we add an epsilon to avoid inf in case of p=0\n",
    "    return -np.sum(np.log(pdf(data, *parameters) + 1e-30))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Naïve, because often, analytic simplifications of the log-likelihood are possible!\n",
    "\n",
    "Analytic simplifications both benefit numerical precision and evaluation speed.\n",
    "\n",
    "It is almost always a good idea, to write down the likelihood on paper and simplify analytically as far as possible."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We find a minimum (and hopefully the global one) using `scipy.optimize.minimize`. \n",
    "\n",
    "As known from the previous lecture, we can estimate the covariance matrix of the parameters from the inverse of the Hessian matrix of the negative log-likelihood,\n",
    "evaluated at the found minimum.\n",
    "\n",
    "\n",
    "The Hessian matrix needs to be determined numerically.\n",
    "`scipy.optimize.minimize` returns a rough estimate, but this is borderline unusable.\n",
    "\n",
    "\n",
    "We use `numdifftools.Hessian` to get the Hessian numerically.\n",
    "\n",
    "\n",
    "`scipy.optimize.minimize` can use multiple algorithms each with benefits and drawbacks under the hood.\n",
    "For more information, see\n",
    " \n",
    "* https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize.html\n",
    "* https://docs.scipy.org/doc/scipy/reference/optimize.minimize-lbfgsb.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import minimize\n",
    "\n",
    "# a very small number to achive > 0 instead of >= 0\n",
    "eps = np.finfo(np.float64).eps\n",
    "\n",
    "result = minimize(\n",
    "    neg_log_likelihood,\n",
    "    args=(measured, ),\n",
    "    x0=[130, 2, 30, 0.2], # here, the initial guess is required\n",
    "    bounds=[\n",
    "        (0, None),    # mean >= 0\n",
    "        (eps, None),  # std > 0\n",
    "        (eps, None),  # tau > 0\n",
    "        (0, 1),       # 0 <= p <= 1\n",
    "    ],\n",
    ")\n",
    "\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numdifftools import Hessian\n",
    "\n",
    "hesse = Hessian(neg_log_likelihood)\n",
    "cov = np.linalg.inv(hesse(result.x, measured))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "higgs_mass = result.x[0]\n",
    "higgs_mass_unc = np.sqrt(cov[0, 0])\n",
    "\n",
    "\n",
    "print(f'Higgs mass is {higgs_mass:.2f} ± {higgs_mass_unc:.2} GeV')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "e = np.linspace(E_MIN, E_MAX, 1000)\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.hist(measured, bins=100, density=True)\n",
    "ax.plot(e, pdf(e, *result.x))\n",
    "ax.axvline(result.x[0], color='C2', lw=2)\n",
    "\n",
    "ax.set_xlabel('$m \\,/\\, \\mathrm{GeV}$')\n",
    "ax.margins(x=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution using iminuit\n",
    "\n",
    "\n",
    "Iminuit provides python bindings for the `Minuit` minimization package from the `ROOT` framework.\n",
    "\n",
    "It does not require a full `ROOT` installation.\n",
    "\n",
    "\"Minuit\" is held to be – at least among particle physisicists – as non-plus-ultra of mimimizers.\n",
    "\n",
    "\n",
    "`iminuit` provides the minimizers and several helper classes for loss functions.\n",
    "This makes it much simpler to perform fits.\n",
    "\n",
    "\n",
    "It can also use likelihood ratio tests via the `minos` interface to estimate parameter uncertainties.\n",
    "(More on that in the next lectures). \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "For starters, let's solve the same problem as above, this time using `iminuit`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from iminuit import Minuit\n",
    "from iminuit.cost import UnbinnedNLL\n",
    "\n",
    "# minuit's UnbinnedNLL takes directly the pdf and the observed data\n",
    "loss = UnbinnedNLL(measured, pdf)\n",
    "\n",
    "m = Minuit(loss, mean=130, std=2, tau=30, p=0.2)\n",
    "\n",
    "# set bounds\n",
    "m.limits['mean'] = (0, None)   # >= 0\n",
    "m.limits['std'] = (eps, None)  # > 0\n",
    "m.limits['tau'] = (eps, None)  # > 0\n",
    "m.limits['p'] = (0, 1)\n",
    "\n",
    "# perform minimization\n",
    "m.migrad()\n",
    "\n",
    "# perform likelihood scan for confidence intervals\n",
    "m.minos()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "higgs_mass = m.values['mean']\n",
    "higgs_mass_unc = m.errors['mean']\n",
    "\n",
    "print(f'Higgs mass is {higgs_mass:.2f} ± {higgs_mass_unc:.2} GeV')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Poisson-Likelihood-Fit using a binned Event Distribution\n",
    "\n",
    "If the single observations are not accessible or the runtime of the analysis is critical for large amounts of data,\n",
    "a so-called \"binned\" fit is also possible.\n",
    "\n",
    "In a binned fit, we estimate the parameters of the distributions from a histogram.\n",
    "\n",
    "\n",
    "Since a histogram is a counting experiment, the single event numbers in each bin each follow a Poisson distribution.\n",
    "\n",
    "\n",
    "The cumulative probability function yields together with the total number of observed events the expected value in each histogram bin,\n",
    "dependent on the fit parameters $\\boldsymbol{\\theta}$.\n",
    "\n",
    "$$\n",
    "\\mathcal{L} = \\prod_{i=1}^N \\mathcal{P}(k=H_i, \\lambda=\\lambda_i(\\boldsymbol{\\theta}))\n",
    "$$\n",
    "\n",
    "with \n",
    "\n",
    "$$\n",
    "\\lambda_i = N_\\mathrm{total} \\cdot (\\mathrm{CDF}(b_i, \\boldsymbol{\\theta}) - \\mathrm{CDF}(a_i, \\boldsymbol{\\theta}))\n",
    "$$\n",
    "\n",
    "Here, $a_i$ und $b_i$ are the bin edges of the $i$-th bin.\n",
    "We then integrate the PDF in each bin and multiply with the total number of events.\n",
    "\n",
    "\n",
    "### Example from \"Lifetime of cosmic Muons\" (Lab Course Experiment V01)\n",
    "\n",
    "The experimental setup of this lab course experiment directly produces a histogram of observed decay times in hardware.\n",
    "\n",
    "We thus cannot perform an unbinned fit, the single values are simply not available.\n",
    "\n",
    "The lab course instruction ask for a least squares fit to the bin heights.\n",
    "\n",
    "This method yields an unbiased estimator but with non-optimal spread, as long as an unweighted fit is performed.\n",
    "\n",
    "In the past, it was recommended to perform a weighted fit assuming $\\sigma_i = \\sqrt{H_i}$.\n",
    "\n",
    "This is wrong! This method consistently yields biased results, due to  under-fluctuations being weighted more strongly than over-fluctuations of the same amount.\n",
    "\n",
    "The correct method is the binned Poisson-likelihood fit as discussed here or the iterative least squares method discussed in SMD-2.\n",
    "\n",
    "A comparison of the different methods is available here: https://gist.github.com/maxnoe/41730e6ca1fac01fc06f0feab5c3566d\n",
    "\n",
    "In the experiment, there is an additional uniformly distributed background from non-decaying, coincident muons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "N = np.genfromtxt(\"resources/muon_data.txt\")\n",
    "t = np.arange(len(N) + 1) / 21.48\n",
    "\n",
    "t = t[5:]\n",
    "N = N[5:]\n",
    "\n",
    "plt.figure()\n",
    "plt.stairs(values=N, edges=t)\n",
    "plt.xlabel('t / µs')\n",
    "plt.ylabel('Number of Events')\n",
    "None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Iminuit` provides the `BinnedNLL` loss function for this use cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from iminuit.cost import BinnedNLL\n",
    "from scipy.stats import uniform\n",
    "\n",
    "T_MIN, T_MAX = t[0], t[-1]\n",
    "\n",
    "def cdf(x, tau, p):\n",
    "    # normalize to 1 in histogram range\n",
    "    cdf_min, cdf_max = expon.cdf([T_MIN, T_MAX], scale=tau) \n",
    "    norm = 1 / (cdf_max - cdf_min)\n",
    "    \n",
    "    signal = p * expon.cdf(x, scale=tau) * norm\n",
    "    background = (1 - p) * uniform.cdf(x, T_MIN, T_MAX)\n",
    "    # combine exponential signal with uniform background\n",
    "    return signal + background\n",
    "\n",
    "\n",
    "# histogram counds, histogram edges and cumulative distribution function\n",
    "loss = BinnedNLL(N, t, cdf)\n",
    "\n",
    "m = Minuit(loss, tau=2, p=0.99)\n",
    "m.limits['tau'] = (eps, None)\n",
    "m.limits['p'] = (0, 1)\n",
    "m.migrad()\n",
    "m.minos()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "muon_lifetime = m.values['tau']\n",
    "muon_lifetime_unc = m.errors['tau']\n",
    "pdg_reference = 2.1969811\n",
    "pdg_reference_unc = 0.0000022\n",
    "\n",
    "print(f'Fit: τ = {muon_lifetime:.3f} ± {muon_lifetime_unc:.3f} µs')\n",
    "print(f'Lit: τ = {pdg_reference:.7f} ± {pdg_reference_unc:.7f} µs')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Likelihood-Scan for uncertainty intervals:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tau, ts, valid = m.mnprofile('tau', size=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "\n",
    "plt.axvline(m.values['tau'], color='C1')\n",
    "plt.plot(m.values['tau'], m.fval, color='C1', marker='o', zorder=3)\n",
    "\n",
    "plt.axvline(m.values['tau'] + m.merrors['tau'].lower, color='C2')\n",
    "plt.axvline(m.values['tau'] + m.merrors['tau'].upper, color='C2')\n",
    "plt.axhline(m.fval + 1, color='C3')\n",
    "plt.axvline(pdg_reference, color='k', label='PDG')\n",
    "plt.plot(tau, ts)\n",
    "\n",
    "plt.xlabel('τ / µs')\n",
    "plt.ylabel('NLL')\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "m.draw_mncontour('tau', 'p', size=250)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}