Week1/ChiSquareTest/AnalyticalLinearFit_original.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analytical linear fit:\n",
    "This Python macro illustrates how to perform a linear fit **analytically**, i.e. without any minimisation (e.g. by Minuit), and thus much faster. This can be done, since one can differentiate the Chi2, set this to zero and solve for both intercept and slope.\n",
    "\n",
    "## References:\n",
    "- Note on course webpage\n",
    "- Bevington: Chapter 6\n",
    "\n",
    "## Author(s), contact(s), and dates:\n",
    "- Author: Troels C. Petersen (NBI)\n",
    "- Email:  petersen@nbi.dk\n",
    "- Date:   16th of November 2023"
   ]
  },
  {
   "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 Chi2Regression\n",
    "from ExternalFunctions import nice_string_output, add_text_to_ax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Program settings:"
   ]
  },
  {
   "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 - more on that later.)\n",
    "save_plots = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Nexp = 10       # Number of datasets fitted\n",
    "Npoints = 9\n",
    "alpha0 = 3.6\n",
    "alpha1 = 0.3\n",
    "sigmay = 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generating and fitting data:\n",
    "In the following, we generate data points following a simple line and fit these both **analytically** and **numerically**. The linear fit (and only this) can be done analytically, as discussed in Barlow chapter 6.2 and outlined here: http://www.nbi.dk/~petersen/Teaching/Stat2019/StraightLineFit.pdf.\n",
    "The numerical fit of the line (and any other function) is done iteratively by Minuit. The code is slightly shorter, but a lot slower, as the code will typically test many (50+) possible combination of fit parameters!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Arrays for storing fit results:\n",
    "array_alpha0 = np.zeros(Nexp)\n",
    "array_alpha1 = np.zeros(Nexp)\n",
    "array_Chi2 = np.zeros(Nexp)\n",
    "array_Prob = np.zeros(Nexp)\n",
    "\n",
    "# Loop, repeating the data generation and fit:\n",
    "for iexp in range(Nexp) : \n",
    "\n",
    "    # Generating data by filling values into (x,y) and associated uncertainties (here chosen to be 0 for x):\n",
    "    x = np.arange(Npoints)+1\n",
    "    ex = np.zeros_like(x)\n",
    "    y = alpha0 + alpha1*x + r.normal(0, sigmay, Npoints) #+ alpha2*x**2\n",
    "    ey = sigmay*np.ones_like(x)\n",
    "\n",
    "    # ------------------------------------------------------------------ #\n",
    "    # Analytical fit:\n",
    "    # ------------------------------------------------------------------ #\n",
    "    # Calculate the relevant sums (see note):\n",
    "    s = len(x)\n",
    "    sx = np.sum(x)\n",
    "    sxx = np.sum(x**2)\n",
    "    sy = np.sum(y)\n",
    "    sxy = np.sum(x*y)\n",
    "    \n",
    "    # Use sums in calculations:\n",
    "    Delta = sxx * s - sx**2\n",
    "    alpha0_calc = (sy  * sxx - sxy * sx) / Delta\n",
    "    alpha1_calc = (sxy * s   - sy  * sx) / Delta\n",
    "    sigma_alpha0_calc = sigmay * np.sqrt(sxx / Delta)\n",
    "    sigma_alpha1_calc = sigmay * np.sqrt(s   / Delta)\n",
    "\n",
    "    # So now you have data points and a fit/theory. How to get the fit quality?\n",
    "    # The answer is to calculate the Chi2 and Ndof, and from these two values get their\n",
    "    # probability using the ChiSquare function (e.g. from scipy):\n",
    "    Chi2_calc = 0.0\n",
    "    for i in range( Npoints ) : \n",
    "        fit_value = alpha0_calc + alpha1_calc * x[i]\n",
    "        Chi2_calc += ((y[i] - fit_value) / ey[i])**2\n",
    "    \n",
    "    Nvar = 2                     # Number of variables (alpha0 and alpha1)\n",
    "    Ndof_calc = Npoints - Nvar   # Number of degrees of freedom = Number of data points - Number of variables\n",
    "    \n",
    "    # From Chi2 and Ndof, one can calculate the probability of obtaining this\n",
    "    # or something worse (i.e. higher Chi2). This is a good function to have!\n",
    "    # We will discuss in class, why/how this function works, and how it plays a central role in statistics.\n",
    "    from scipy import stats\n",
    "    Prob_calc =  stats.chi2.sf(Chi2_calc, Ndof_calc) # The chi2 probability given N degrees of freedom (Ndof)\n",
    "    \n",
    "    # Fill the arrays with fit results (to produce plots of these at the end):\n",
    "    array_alpha0[iexp] = alpha0_calc\n",
    "    array_alpha1[iexp] = alpha1_calc\n",
    "    array_Chi2[iexp] = Chi2_calc\n",
    "    array_Prob[iexp] = Prob_calc\n",
    "    \n",
    "    \n",
    "    # ------------------------------------------------------------------ #\n",
    "    # Numerical fit:\n",
    "    # ------------------------------------------------------------------ #\n",
    "    # Define a fit function:\n",
    "    def fit_function(x, alpha0, alpha1):\n",
    "        return alpha0 + alpha1*x\n",
    "\n",
    "    # Now we define a ChiSquare to be minimised (using probfit), where we set various settings and starting parameters:\n",
    "    chi2_object = Chi2Regression(fit_function, x, y, ey) \n",
    "    minuit = Minuit(chi2_object, alpha0=1, alpha1=1)\n",
    "    minuit.errordef = 1\n",
    "    minuit.print_level = 0\n",
    "    minuit.migrad();  # perform the actual fit\n",
    "    minuit_output = [minuit.fmin, minuit.params] # save the output parameters in case needed\n",
    "    \n",
    "    alpha0_fit = minuit.values['alpha0']\n",
    "    alpha1_fit = minuit.values['alpha1']\n",
    "    sigma_alpha0_fit = minuit.errors['alpha0']\n",
    "    sigma_alpha1_fit = minuit.errors['alpha1']\n",
    "    \n",
    "    # In Minuit, you can just ask the fit function for it:\n",
    "    Chi2_fit = minuit.fval # the chi2 value\n",
    "    Prob_fit = stats.chi2.sf(Chi2_fit, Ndof_calc) # The chi2 probability given N degrees of freedom (Ndof, taken from above!)\n",
    "    \n",
    "    # Let us see what the fit gives for the first couple of datasets:\n",
    "    if (iexp < 10) :\n",
    "        print(f\"  Ana. Fit: a0={alpha0_calc:6.3f}+-{sigma_alpha0_calc:5.3f}  a1={alpha1_calc:5.3f}+-{sigma_alpha1_calc:5.3f}  p={Prob_calc:6.4f}\"+\n",
    "              f\"  Num. Fit: a0={alpha0_fit:6.3f}+-{sigma_alpha0_fit:5.3f}  a1={alpha1_fit:5.3f}+-{sigma_alpha1_fit:5.3f}  p={Prob_fit:6.4f}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(10,6))\n",
    "ax.errorbar(x, y, ey, fmt='ro', ecolor='k', elinewidth=1, capsize=2, capthick=1)\n",
    "ax.plot(x, fit_function(x, *minuit.values), '-r')\n",
    "plt.close()\n",
    "\n",
    "d = {'alpha0':   [alpha0_calc, sigma_alpha0_calc],\n",
    "     'alpha1':   [alpha1_calc, sigma_alpha1_calc],\n",
    "     'Chi2':     Chi2_calc,\n",
    "     'ndf':      Ndof_calc,\n",
    "     'Prob':     Prob_calc,\n",
    "    }\n",
    "\n",
    "text = nice_string_output(d, extra_spacing=2, decimals=3)\n",
    "add_text_to_ax(0.02, 0.95, text, ax, fontsize=14)\n",
    "fig.tight_layout()\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if (save_plots) :\n",
    "    fig.savefig('AnalyticalLinearFit.pdf', dpi=600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "---------------------------------------------------------------------------------- \n",
    "\n",
    "\n",
    "# Questions:\n",
    "\n",
    "1) Do the analytical result(s) (Calc) agree with the numerical one(s) (Fit)?\n",
    "\n",
    "### Advanced questions:\n",
    "\n",
    "2) Can you measure/determine the difference in speed between the two methods?"
   ]
  }
 ],
 "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.7.6"
  },
  "main_language": "python"
 },
 "nbformat": 4,
 "nbformat_minor": 4
}