{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Curve fitting and the $\\chi^2$ error surface\n",
    "#### Material to accompany Hughes and Hase Section 6.5\n",
    "\n",
    "Marty Ligare August 2020; modified March 2022 -- NFL, March 2023 -- Tom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy import optimize\n",
    "from scipy import stats\n",
    "\n",
    "import urllib    # for importing data from URL\n",
    "\n",
    "import matplotlib as mpl\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm\n",
    "from matplotlib.ticker import LinearLocator, FormatStrFormatter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Following is an Ipython magic command that puts figures in notebook.\n",
    "%matplotlib notebook\n",
    "\n",
    "# M.L. modification of matplotlib defaults\n",
    "# Changes can also be put in matplotlibrc file, \n",
    "# or effected using mpl.rcParams[]\n",
    "mpl.style.use('classic') \n",
    "plt.rc('figure', figsize = (6, 4.5)) # Reduces overall size of figures\n",
    "plt.rc('axes', labelsize=16, titlesize=14)\n",
    "plt.rc('figure', autolayout = True) # Adjusts supblot params for new size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define  functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(x,m,b):\n",
    "    '''Simple linear function with slope m and intercept b'''\n",
    "    return x*m + b\n",
    "\n",
    "def chi2(x, y, u, m, b):\n",
    "    '''Chisquare as a function of data (x, y, and yerr=u), and model \n",
    "    parameters slope and intercept (m and b)'''\n",
    "    return np.sum((y - f(x, m, b))**2/u**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Linear fit to data for $m$ and $b$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Data to be fit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To load from a file, use\n",
    "# g = urllib.request.urlopen('https://www.eg.bucknell.edu/~phys310/skills/data_analysis/data3.dat')\n",
    "# data = np.loadtxt(g)\n",
    "# x, y, u = data.T\n",
    "# \n",
    "# Or: data = np.loadtxt(\"file.dat\")  \n",
    "# Format:  [[x1,y1,u1], [x2,y2,u2], ... ]   where u1 is uncertainty in y1\n",
    "data = np.array([[1, 2.947032612427293, 0.5],\n",
    " [2, 6.168779380682309, 0.5],\n",
    " [3, 7.1618838821688, 0.5],\n",
    " [4, 9.590549514954866, 0.5],\n",
    " [5, 11.20657, 0.5],\n",
    " [6, 12.3123, 0.5],\n",
    " [7, 15.0221,0.5],\n",
    " [8, 16.9281, 0.5],\n",
    " [9, 18.8221, 0.5],\n",
    " [10., 21.7221, 0.5]])\n",
    "\n",
    "x, y, u = data.T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xc = np.linspace(0,6,201) # quasi-continuous set of x's for function plot\n",
    "plt.figure()\n",
    "plt.title(\"data\",fontsize=14)\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$y$')\n",
    "plt.axhline(0,color='magenta')\n",
    "plt.xlim(0,10) \n",
    "plt.errorbar(x,y,yerr=u,fmt='o');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Perform fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "popt, pcov = optimize.curve_fit(f, x, y, sigma=u, absolute_sigma=True)\n",
    "\n",
    "slope = popt[0]\n",
    "alpha_m = np.sqrt(pcov[0,0])  # Std error in slope\n",
    "intercept = popt[1]\n",
    "alpha_b = np.sqrt(pcov[1,1])  # Std error in intercept\n",
    "\n",
    "print(\"slope =\", slope,\"+/-\", alpha_m,\"\\n\")\n",
    "print(\"intercept =\", intercept,\"+/-\", alpha_b,\"\\n\")\n",
    "\n",
    "print(\"covariance matrix =\",\"\\n\",pcov,\"\\n\")\n",
    "pcov_data = pcov\n",
    "\n",
    "print(\"chi2 =\", chi2(x, y, u,*popt))\n",
    "# print(\"reduced chi2 = chi2/(len(x)-len(popt)) =\", chi2(x, y, u, *popt)/(len(x)-len(popt)))\n",
    "\n",
    "a = chi2(x,y,u,*popt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "xc = np.linspace(0,10,201) # quasi-continuous set of x's function plot\n",
    "plt.figure()\n",
    "plt.title(\"data with best fit line\",fontsize=14)\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$y$')\n",
    "plt.axhline(0, color='magenta')\n",
    "plt.xlim(0,10)  # Pad x-range on plot\n",
    "plt.errorbar(x, y, yerr=u, fmt='o');\n",
    "plt.plot(xc ,f(xc, slope, intercept));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Residuals:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.axhline(0,color='magenta')\n",
    "plt.title('normalized residuals')\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('normalized residuals')\n",
    "plt.grid(True)\n",
    "plt.errorbar(x,(f(x,slope,intercept)-y)/u,1,fmt='o')\n",
    "plt.xlim(0,10);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### What confidence can we have in the model?\n",
    "If normalized residuals look ok:\n",
    "+ <i>Assume</i> the model, along with the best fit parameters are ok (\"null hypothesis\") \n",
    "+ Generate simulated data based on the model, best-fit parameters, and uncertainties.\n",
    "(Use same $x$ values).\n",
    "+ Fit each of the simulated data sets\n",
    "+ What's the probability of getting values of $\\chi^2$ from the simulated data sets \n",
    "that are greater than the value given by the data that was actually measured.\n",
    "In other words, is our value of $\\chi^2$ a reasonable result given the fluctuations we\n",
    "expect in our data?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_sim = 10000               # Number of data sets to simulate\n",
    "count = 0\n",
    "for i in range(n_sim):\n",
    "    mSim = np.array([])      # Array for simulated slopes\n",
    "    bSim = np.array([])      # Array for simulated intercepts\n",
    "    chi_sim = np.array([])   # Array for chisq from simulated sets\n",
    "\n",
    "for i in range(n_sim):\n",
    "    ySim = stats.norm.rvs(slope*x+intercept,u) # Generate simulated data set\n",
    "    popt, pcov = optimize.curve_fit(f, x, ySim, sigma=u, absolute_sigma=True)  # Fit simulated data\n",
    "    bSim = np.append(bSim, popt[1])      # Record intercept\n",
    "    mSim = np.append(mSim, popt[0])      # Record slope\n",
    "    aa = chi2(x,ySim,u,*popt) \n",
    "    chi_sim = np.append(chi_sim, aa)\n",
    "    if aa > a:\n",
    "        count += 1\n",
    "    \n",
    "print(count/n_sim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "1 - stats.chi2.cdf(3.700621, df=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Make \"data\" for contour plot\n",
    "+ Choose ranges of $m$ and $b$ for contour plot\n",
    "+ Make meshgrid of slope and intercept values\n",
    "+ Calculate values of $\\chi^2$ at grid points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = np.linspace(1.5, 2.5, 201) # NOTE:  You need to change the range of m to encompass your best fit value for slope (see above)\n",
    "b = np.linspace(0.5, 2.5, 201) # NOTE:  You need to change the range of b to encompass your best fit value for the intercept\n",
    "M, B = np.meshgrid(m, b, indexing='ij')\n",
    "\n",
    "Z = np.zeros((len(m),len(b)))\n",
    "\n",
    "for i in range(len(m)):\n",
    "    for j in range(len(b)):\n",
    "        Z[i,j] = chi2(x, y, u, m[i],b[j]) - chi2(x, y, u, *popt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "CS = plt.contour(M, B, Z, levels=[1,2.7,4,9])   # contour plots of chi2 = 1, 2.7, 4, 9 above chi2_min\n",
    "#\n",
    "#SC = plt.scatter(mSim, bSim, s=0.01)   # overlay scatterplot of simulted data\n",
    "#\n",
    "plt.xlabel('$m$')\n",
    "plt.ylabel('$b$')\n",
    "plt.grid()\n",
    "plt.axhline(intercept)\n",
    "#plt.axhline(intercept + alpha_b, linestyle='--')\n",
    "#plt.axhline(intercept - alpha_b, linestyle='--')\n",
    "plt.axvline(slope)\n",
    "#plt.axvline(slope + alpha_m,linestyle='--')\n",
    "#plt.axvline(slope - alpha_m,linestyle='--')\n",
    "plt.clabel(CS, inline=1, fontsize=10);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####  Correlation between fit parameters\n",
    "\n",
    "Note that the contours are tilted and squashed; this indicates that there is correlation between the two fit parameters, mainly caused because the y-intercept is not within the data range. These elongated contours are telling you that a small change in $m$ can, to some extent, be compensated for by a corresponding change in $b$.\n",
    "\n",
    "Fits with correlated data aren't great, because they tend to lead to an overestimation of the uncertainties in the fit parameters. In this case, note tha thte uncertainty in $b$ is $\\sim 0.5$ due to this correlation effect.\n",
    "\n",
    "We can do better if we perform our linear fit on a dataset whose x-values are centered on $x=0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_mean = np.mean(x)       # determine the center of the distribution of x values\n",
    "x_shift = x - x_mean      # shift the array so that the x data are centered on x = 0\n",
    "plt.figure()\n",
    "plt.title(\"shifted data\",fontsize=14)\n",
    "plt.xlabel('$x_{shift}$')\n",
    "plt.ylabel('$y$')\n",
    "plt.axhline(0,color='magenta')\n",
    "plt.xlim(0-x_mean,10-x_mean) \n",
    "plt.errorbar(x_shift,y,yerr=u,fmt='o');\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "popt, pcov = optimize.curve_fit(f, x_shift, y, sigma=u, absolute_sigma=True)\n",
    "\n",
    "slope = popt[0]\n",
    "alpha_m = np.sqrt(pcov[0,0])  # Std error in slope\n",
    "intercept = popt[1]\n",
    "alpha_b = np.sqrt(pcov[1,1])  # Std error in intercept\n",
    "\n",
    "print(\"slope =\", slope,\"+/-\", alpha_m,\"\\n\")\n",
    "print(\"intercept =\", intercept,\"+/-\", alpha_b,\"\\n\")\n",
    "\n",
    "print(\"covariance matrix =\",\"\\n\",pcov,\"\\n\")\n",
    "pcov_data = pcov\n",
    "\n",
    "print(\"chi2 =\", chi2(x_shift, y, u,*popt))\n",
    "# print(\"reduced chi2 = chi2/(len(x)-len(popt)) =\", chi2(x, y, u, *popt)/(len(x)-len(popt)))\n",
    "\n",
    "a = chi2(x_shift,y,u,*popt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = np.linspace(1.7, 2.3, 201)\n",
    "b = np.linspace(11.7, 12.7, 201)\n",
    "M, B = np.meshgrid(m, b, indexing='ij')\n",
    "\n",
    "Z = np.zeros((len(m),len(b)))\n",
    "\n",
    "for i in range(len(m)):\n",
    "    for j in range(len(b)):\n",
    "        Z[i,j] = chi2(x_shift, y, u, m[i],b[j]) - chi2(x_shift, y, u, *popt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "CS = plt.contour(M, B, Z, levels=[1,2.7,4,9])   # contour plots of chi2 = 1, 2.7, 4, 9 above chi2_min\n",
    "#\n",
    "#SC = plt.scatter(mSim, bSim, s=0.01)   # overlay scatterplot of simulted data\n",
    "#\n",
    "plt.xlabel('$m$')\n",
    "plt.ylabel('$b$')\n",
    "plt.grid()\n",
    "plt.axhline(intercept)\n",
    "#plt.axhline(intercept + alpha_b, linestyle='--')\n",
    "#plt.axhline(intercept - alpha_b, linestyle='--')\n",
    "plt.axvline(slope)\n",
    "#plt.axvline(slope + alpha_m,linestyle='--')\n",
    "#plt.axvline(slope - alpha_m,linestyle='--')\n",
    "plt.clabel(CS, inline=1, fontsize=10);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Version information\n",
    "\n",
    "+ %version_information is an IPython magic extension for showing version information for dependency modules in a notebook; \n",
    "\n",
    "+ See `https://github.com/jrjohansson/version_information`\n",
    "\n",
    "+ `%version_information` is available on Bucknell computers on the linux network.  You can easily install it on\n",
    "any computer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%load_ext version_information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%version_information numpy, scipy, matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}