Week8/HarmonicOscillator/HarmonicOscillatorFit_original.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Advanced Harmonic Oscillator\n",
    "\n",
    "## Description\n",
    "\n",
    "This exercise is about more advanced fitting of data. The idea is to begin with the simplest possible function that even remotely does the job, which is then expanded to accommodate all the features present in the data fitted. The data is from three versions of a damped harmonic oscillator, consisting of a weight hanging from a spring, but with different damping:\n",
    "- No additional damping.\n",
    "- More damping from a round piece of cardboard, which increases the drag.\n",
    "- Damping from friction of the weight against metal.\n",
    "\n",
    "\n",
    "## Your task\n",
    "\n",
    "Take a look at the various dataset, and try to fit them with an appropriate function. By considering patterns in the fit residuals, one can get an idea of what to further include. Play around with the initial parameters of the fits and the various possible additions to the function. Answer / Discuss the questions written at the bottom of this notebook. Remember to try to fit all three datasets.\n",
    "\n",
    "_Note that the initial fit values used here may not be the proper ones to begin with!!!_\n",
    "\n",
    "\n",
    "##  Authors: \n",
    "- Troels Petersen ([email](mailto:petersen@nbi.dk))\n",
    "\n",
    "## Date\n",
    "11th of January 2024"
   ]
  },
  {
   "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                                             # Modules to see files and folders in directories\n",
    "from scipy import stats\n",
    "\n",
    "from mpl_toolkits.axes_grid1.inset_locator import inset_axes\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.path.append('../../External_Functions')\n",
    "from ExternalFunctions import Chi2Regression, BinnedLH, UnbinnedLH\n",
    "from ExternalFunctions import nice_string_output, add_text_to_ax    # Useful functions to print fit results on figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to draw parameters from a fit on the figure in one line\n",
    "def draw_fit_parameters(minuit_obj, axis, x, y, title=None, color='k', chi2=None, ndof=None, prob=None) :\n",
    "    names = []\n",
    "    values = []\n",
    "    if title is not None :\n",
    "        names += [title]\n",
    "        values += ['']\n",
    "    if chi2 is not None and ndof is not None and prob is not None :\n",
    "        names += ['Chi2 / ndof','Prob']\n",
    "        values += [\"{:.3f} / {:d}\".format(chi2, ndof), \"{:.3f}\".format(prob)]\n",
    "    # names += minuit_obj.values.keys()\n",
    "    names += minuit_obj.parameters\n",
    "    values += [\"{:.3f} +/- {:.3f}\".format(minuit_obj.values[val], minuit_obj.errors[val]) for val in minuit_obj.parameters]\n",
    "    d ={}\n",
    "    for n,v in zip(names,values):\n",
    "        d[n] = v\n",
    "    axis.text(x, y, nice_string_output(d), family='monospace', transform=axis.transAxes, fontsize=10, verticalalignment='top', color=color)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SavePlots = False\n",
    "verbose = True\n",
    "Nverbose = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## 1st Dataset - no additional damping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set variables:\n",
    "# -------------------------------------- #\n",
    "tmax = 38.0       # Maximum of time range fitted\n",
    "filename = \"data_HarmOsc1.txt\"\n",
    "    \n",
    "# Load time and distance, assign uncertainty:\n",
    "# -------------------------------------- #\n",
    "time, dist = np.loadtxt(filename, unpack=True)\n",
    "time -= time[0]                          # For ensuring that time starts at 0.0s!\n",
    "edist = np.ones_like(dist)*0.0037        # I have estimated the uncertainty from a calibration with no swings.\n",
    "    \n",
    "# Check loaded data:\n",
    "# -------------------------------------- #\n",
    "if verbose :\n",
    "    for i in range(Nverbose) :\n",
    "        print(\"  Time = %6.3f    Dist = %6.3f \"%(time[i], dist[i]))\n",
    "print(\"  Number of entries read: %d    Time of last read: %6.3f\"%(len(time), time[-1]))\n",
    "    \n",
    "# Sanity check (Did you check the data?):\n",
    "for i in np.where((time < -0.001) | (time > 100.0) | (dist < -5.0) | (dist > 5.0))[0] :\n",
    "    print(\"Warning: Strange value for time and/or dist!\", i, time[i], dist[i])\n",
    "    \n",
    "# Plot the data:\n",
    "# -------------------------------------- #\n",
    "fig1 = plt.figure(figsize=(16, 10))   # Make an empty figure\n",
    "ax1 = fig1.add_axes((.1,.3,.8,.6))    # Add the top subfigure\n",
    "ax1.set_ylabel(\"Position\")\n",
    "\n",
    "ax1.errorbar(time, dist, edist, fmt='k_', label='data', ecolor='k', elinewidth=1, capsize=1, capthick=1)\n",
    "ax1.set_xlim(time[0],time[-1])\n",
    "ax1.set_ylim(top=ax1.get_ylim()[1]*2.5)\n",
    "\n",
    "\n",
    "# Fit the data:\n",
    "# -------------------------------------- #\n",
    "mask = time < tmax\n",
    "    \n",
    "# Fit this data with a (initially too) simple harmonic function:\n",
    "def fit1(x, p0, p1, p2) :\n",
    "    return p0 * np.cos(p1*x+p2)\n",
    "    \n",
    "FitObject1 = Chi2Regression(fit1, time[mask], dist[mask], edist[mask])\n",
    "Minuit1 = Minuit(FitObject1, p0=0.0, p1=0.0, p2=0.0)\n",
    "Minuit1.errordef = 1.0\n",
    "Minuit1.migrad()\n",
    "\n",
    "\n",
    "chi2 = Minuit1.fval\n",
    "ndof = len(time[mask])-len(Minuit1.values[:])\n",
    "prob = stats.chi2.sf(chi2, ndof)\n",
    "\n",
    "# Draw the function on top of the data:\n",
    "x_fit = np.linspace(0, tmax, 1000)\n",
    "y_fit = fit1(x_fit, *Minuit1.values[:])\n",
    "ax1.plot(x_fit, y_fit, 'r-', label='Simple harmonic oscillator')\n",
    "draw_fit_parameters(Minuit1, ax1, 0.05, 0.95, 'Simple harmonic oscillator:','r', chi2, ndof, prob)\n",
    "\n",
    "# Calculate residuals:\n",
    "# -------------------------------------- #\n",
    "dd1 = dist-fit1(time, *Minuit1.values[:])\n",
    "    \n",
    "# Draw the residuals:\n",
    "# -----------------------------------------\n",
    "# Draw residuals as function of time in bottom subfigure\n",
    "ax1.get_yaxis().get_ticklabels()[0].set_visible(False)\n",
    "ax1.get_yaxis().get_ticklabels()[1].set_visible(False) # Remove bottom y-tick on top subfigure to prevent overlapping ticks\n",
    "ax2 = fig1.add_axes((0.1,0.1,0.8,0.2),sharex=ax1)      # Add bottom subfigure for residuals, and have its x-axis follow the top figure\n",
    "ax2.set_xlim(time[0], time[-1])\n",
    "ax2.set_xlabel(\"Time elapsed [s]\")\n",
    "ax2.set_ylabel(\"Position residual\")\n",
    "    \n",
    "ax2.errorbar(time, dd1, edist, fmt='r_', ecolor='r', elinewidth=1, capsize=1, capthick=1)\n",
    "ax2.plot((0,tmax),(0,0), 'w--', lw=1, zorder=10)\n",
    "    \n",
    "# Draw histograms of residuals in a new inset figure\n",
    "axins = inset_axes(ax1, 3.0, 2.0, loc=5, bbox_to_anchor=(0.89, 0.78), bbox_transform=ax1.figure.transFigure)\n",
    "plt.yticks([],[])\n",
    "axins.hist(dd1, bins=120, range=(-0.6, 0.6), histtype='step', linewidth=1, color='r')\n",
    "axins.set_ylim(top=axins.get_ylim()[1]*1.5)\n",
    "    \n",
    "axins.text(0.40, 0.95, \"Residuals\", transform=axins.transAxes, fontsize=10, verticalalignment='top')\n",
    "axins.text(0.01, 0.85, nice_string_output({'Mean': \"{:.3f}\".format(dd1.mean()),\n",
    "                                           'RMS' : \"{:.3f}\".format(dd1.std(ddof=1))}),\n",
    "            family='monospace', transform=axins.transAxes, fontsize=10, verticalalignment='top', color='r')\n",
    "\n",
    "# Finalize the figure\n",
    "if (SavePlots) :\n",
    "    fig1.savefig(\"Fit_HarmOsc1.pdf\", dpi=600)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Questions to 1st dataset:\n",
    "------------------------\n",
    "\n",
    "1. Look at the data file and plot and see if you can by eye (or simple fits) estimate\n",
    "   the size of the uncertainty of the points. It is not easy, but you should be able\n",
    "   to get it to within a factor of 2-3. Try for 2-3 minutes and discuss it with your\n",
    "   neighbor before reading on!\n",
    "\n",
    "           - - - - - -     2-3 minutes (success or failure) later - - - - - -\n",
    "\n",
    "   If you didn't know how to estimate this uncertainty, then try to zoom in on a very\n",
    "   small part of the curve (say xy: [0.1,0.3] x [-0.6,0.4]) in a large (e.g. 15x20) plot,\n",
    "   where it should be possible to fit it with a line (even by eye).\n",
    "   Since you know that for a small enough range of the data, this will be a reasonable\n",
    "   PDF to use, you can extract the error from the Std. of the residuals, which gives you\n",
    "   a rough estimate of the errors on the points (remember the pendulum timing!).\n",
    "   Compare the error you got to the one I put in.\n",
    "\n",
    "\n",
    "2. Once you have tried, set the error to 0.0037, and try to fit a damped harmonic oscillator.\n",
    "   You should write the function yourself. Do so and run the fit before reading on.\n",
    "\n",
    "           - - - - - -     5-10 minutes (success or failure) later - - - - - -\n",
    "\n",
    "    Did you manage to write the fitting function? Also, did you remember to put in a\n",
    "    parameter taking care of the offset from zero? Anyway, run the fit...\n",
    "\n",
    "           - - - - - -     5-10 minutes (success or failure) later - - - - - -\n",
    "\n",
    "3. Did the fit converge? I imagine that it didn't (thought it might have), and my first\n",
    "   guess on why (apart from obvious bugs in the function) would be initial parameters. You\n",
    "   need to set the initial parameters quite accurately, for the fit to work. Think about\n",
    "   this, i.e. in a 5D parameter space, where being just 5% wrong or so will give you nothing\n",
    "   of value?\n",
    "   So now try to evaluate what some good initial values for the fit would be. You can\n",
    "   start by simply making educated guesses, but if that fails, you can instead draw the\n",
    "   function choosing some parameters, until it starts looking like the data you have.\n",
    "   \n",
    "          - - - - - -     10-20 minutes (success or failure) later - - - - - -\n",
    "          \n",
    "My guess is, that once you have a functioning harmonic oscillator fit running, then the fit will look reasonable. Never mind the ChiSquare value - it is hard to know the uncertainty anyway, and there may be a myriad of effects. \n",
    "The fit is not that bad, but clearly the residuals still have a lot of structure, which are thus features not included/accounted for in the fit. So you should try to add complexity to your fitting function.\n",
    "\n",
    "The way to judge, if you are doing something right is to see if the ChiSquare (and the residuals) reduce every time you introduce something in addition. The ChiSquare should decrease \"significantly\" (*) every time, otherwise, the addition does not cover any corresponding effect in the data.\n",
    "\n",
    "(*) Remember, that a reduction in ChiSquare of say 9 correponds to a $\\sqrt(9) = 3$ sigma effect, and thus this is the order of magnitude, it should decrease (at least I would say).\n",
    "\n",
    "As inspiration to the more advanced fit, I could ask you:\n",
    " - Is there any visible damping? How to include that in the fit?\n",
    " - Are the period and/or phase constant?\n",
    " - Are there more than one exponential/oscillation in the system?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## 2nd Dataset\n",
    "\n",
    "This dataset was acquired with the help of an overdamped oscillator (a cardboard panel was added to the weight, thereby increasing the atospheric drag of the system). This not only increases the damping, but also increases the complexity a bit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set variables:\n",
    "# -------------------------------------- #\n",
    "tmax = 36.0       # Maximum of time range fitted\n",
    "filename = \"data_HarmOsc2.txt\"\n",
    "    \n",
    "# Load time and distance, assign uncertainty:\n",
    "# -------------------------------------- #\n",
    "time, dist = np.loadtxt(filename, unpack=True)\n",
    "time -= time[0]      # For ensuring that time starts at 0.0s!\n",
    "edist = np.ones_like(dist)*0.0029\n",
    "\n",
    "# Plot the data:\n",
    "# -------------------------------------- #\n",
    "fig1 = plt.figure(figsize=(16, 10))   # Make an empty figure\n",
    "ax1 = fig1.add_axes((.1,.3,.8,.6))    # Add the top subfigure\n",
    "ax1.set_ylabel(\"Position\")\n",
    "    \n",
    "# Make a graph of the data:\n",
    "ax1.errorbar(time, dist, edist, fmt='k_', ecolor='k', elinewidth=1, capsize=1, capthick=1)\n",
    "ax1.set_xlim(time[0],time[-1])\n",
    "ax1.set_ylim(top=ax1.get_ylim()[1]*1.5);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, make a simple harmonic oscillator fit run. Now your job is to expand on the fitting function\n",
    "and introduce terms to include various effects and thus reduce the Chi2. Set your fit to the range [0.005,36.005], and see how low a Chi2 you can get.\n",
    "\n",
    "Also, can you see a change in the oscillation behaviour at some point? There seems to be a point at which something changes (damping goes from turbulent to regular?). Does your result improve, if you divide your fit into two parts around there? Three? Twenty?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "----\n",
    "\n",
    "## 3rd Dataset\n",
    "\n",
    "This set corresponds to a set up in which the oscillator is __slowed__ by sliding constantly along a metal plate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set variables:\n",
    "# -------------------------------------- #\n",
    "tmax = 12.0       # Maximum of time range fitted\n",
    "filename = \"data_HarmOsc3.txt\"\n",
    "    \n",
    "# Load time and distance, assign uncertainty:\n",
    "# -------------------------------------- #\n",
    "time, dist = np.loadtxt(filename, unpack=True)\n",
    "time -= time[0]      # For ensuring that time starts at 0.0s!\n",
    "edist = np.ones_like(dist)*0.0139\n",
    "\n",
    "# Plot the data:\n",
    "# -------------------------------------- #\n",
    "fig1 = plt.figure(figsize=(16, 10))   # Make an empty figure\n",
    "ax1 = fig1.add_axes((.1,.3,.8,.6))    # Add the top subfigure\n",
    "ax1.set_ylabel(\"Position\")\n",
    "    \n",
    "# Make a graph of the data:\n",
    "ax1.errorbar(time, dist, edist, fmt='k_', ecolor='k', elinewidth=1, capsize=1, capthick=1)\n",
    "ax1.set_ylim(top=ax1.get_ylim()[1]*1.5);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The 3rd dataset exhibits a different type of damping in the oscillation. Fit this, and determine at which point in time the oscillation stops."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### You need to write this functional form of a damped harmonic oscillator:\n",
    " \n",
    " $d(t) = A \\cdot\\sin(\\omega t+\\phi) \\cdot\\exp(-\\gamma t)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#def fit1(x, p0, p1, p2, p3, p4) :\n",
    "#     return p0 + p1*np.sin(p2*x+p3)*np.exp(-p4*x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Possible additions to especially the 2nd dataset could be:\n",
    "#    def fit2(x,p0,p1,p2,p3,p4,p5,p6,p7) :           return p0 + p1 *  np.exp(-p2*x)                    * np.cos(p3+p4*x)         * (1.0 + p5*np.cos(p6*x+p7))\n",
    "#    def fit3(x,p0,p1,p2,p3,p4,p5,p6,p7,p8) :        return p0 + p1 *  np.exp(-p2*x)                    * np.cos(p3+p4*x+p8*x**2) * (1.0 + p5*np.cos(p6*x+p7))\n",
    "#    def fit4(x,p0,p1,p2,p3,p4,p5,p6,p7,p8,p9,p10) : return p0 + p1 * (np.exp(-p2*x)+p9*np.exp(-p10*x)) * np.cos(p3+p4*x+p8*x**2) * (1.0 + p5*np.cos(p6*x+p7))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Learning points:\n",
    "\n",
    "This is an exercise in starting with a simple fit and building a more advanced one.\n",
    "\n",
    "Following this exercise, you should have experience in:\n",
    "1. The idea of building a more and more advanced fit by adding fit terms/parameters.\n",
    "2. The usefulness of the fit residuals to\n",
    "  - determine the effect of an added fit term/parameter.\n",
    "  - guide you in realising what the next effect to model in the fit is.\n",
    "3. Knowing how to judge the significance of an added term/parameter.\n",
    "\n",
    "To some extend, the exercise resembles (data) science in general: You look at the data, and try out a simple model. If it fails, then you try something different, and now you can compare the two. And then you add more details to the models that seem to work, and then you might reach the point, where the model actually describes the data in \"enough\" detail, that you are satisfied."
   ]
  }
 ],
 "metadata": {
  "executable": "/usr/bin/env python",
  "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.13"
  },
  "main_language": "python"
 },
 "nbformat": 4,
 "nbformat_minor": 4
}