{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyerrors as pe\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use('./base_style.mplstyle')\n",
    "plt.rc('text', usetex=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read data from the pcac example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "p_obs = {}\n",
    "p_obs['f_P'] = pe.load_object('./data/B1k2_f_P.p')\n",
    "\n",
    "# f_A can be accesed via p_obs['f_A']\n",
    "\n",
    "[o.gamma_method() for o in p_obs['f_P']];"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now define a custom fit function, in this case a single exponential. __Here we need to use the autograd wrapped version of numpy__  (imported as anp) to use automatic differentiation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import autograd.numpy as anp\n",
    "def func_exp(a, x):\n",
    "    y = a[1] * anp.exp(-a[0] * x)\n",
    "    return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fit single exponential to f_P. The kwarg `resplot` generates a figure which visualizes the fit with residuals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fit with 2 parameters\n",
      "Method: Levenberg-Marquardt\n",
      "`xtol` termination condition is satisfied.\n",
      "chisquare/d.o.f.: 0.00287692704517733\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 28244 characters]\n",
      "text/plain": [
       "<Figure size 800x494.438 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mass, Matrix element:\n",
      "[Obs[0.2102(63)], Obs[14.24(66)]]\n"
     ]
    }
   ],
   "source": [
    "# Specify fit range for single exponential fit\n",
    "start_se = 8\n",
    "stop_se = 19\n",
    "\n",
    "a = pe.fits.standard_fit(np.arange(start_se, stop_se), p_obs['f_P'][start_se:stop_se], func_exp, resplot=True)\n",
    "[o.gamma_method() for o in a]\n",
    "print('Mass, Matrix element:')\n",
    "print(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The covariance of the two fit parameters can be computed in the following way"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Covariance:  0.003465486601483565\n",
      "Normalized covariance:  0.8360758153764554\n"
     ]
    }
   ],
   "source": [
    "cov_01 = pe.fits.covariance(a[0], a[1])\n",
    "print('Covariance: ', cov_01)\n",
    "print('Normalized covariance: ', cov_01 / a[0].dvalue / a[1].dvalue)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Effective mass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate the effective mass for comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "m_eff_f_P = []\n",
    "for i in range(len(p_obs['f_P']) - 1):\n",
    "    m_eff_f_P.append(np.log(p_obs['f_P'][i] / p_obs['f_P'][i+1]))\n",
    "    m_eff_f_P[i].gamma_method()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate the corresponding plateau and compare the two results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Effective mass:\n",
      "Obs[0.2114(52)]\n",
      "Fitted mass:\n",
      "Obs[0.2102(63)]\n"
     ]
    }
   ],
   "source": [
    "m_eff_plateau = np.mean(m_eff_f_P[start_se: stop_se]) # Plateau from 8 to 16\n",
    "m_eff_plateau.gamma_method()\n",
    "print('Effective mass:')\n",
    "m_eff_plateau.print(0)\n",
    "print('Fitted mass:')\n",
    "a[0].print(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 25292 characters]\n",
      "text/plain": [
       "<Figure size 640x395.55 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "pe.plot_corrs([m_eff_f_P], plateau=[a[0], m_eff_plateau], xrange=[3.5, 19.5], prange=[start_se, stop_se], label=['Effective mass'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fitting two exponentials"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also fit the data with two exponentials where the second term describes the cutoff effects imposed by the boundary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def func_2exp(a, x):\n",
    "    y = a[1] * anp.exp(-a[0] * x) + a[3] * anp.exp(-a[2] * x)\n",
    "    return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can trigger the computation of $\\chi^2/\\chi^2_\\text{exp}$ with the kwarg `expected_chisquare` which takes into account correlations in the data and non-linearities in the fit function and should give a more reliable measure for goodness of fit than $\\chi^2/\\text{d.o.f.}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fit with 4 parameters\n",
      "Method: Levenberg-Marquardt\n",
      "`xtol` termination condition is satisfied.\n",
      "chisquare/d.o.f.: 0.05399877210985092\n",
      "chisquare/expected_chisquare: 0.7915235152330492\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 28248 characters]\n",
      "text/plain": [
       "<Figure size 800x494.438 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fit result:\n",
      "[Obs[0.2146(65)], Obs[15.15(88)], Obs[0.623(60)], Obs[-9.64(74)]]\n"
     ]
    }
   ],
   "source": [
    "# Specify fit range for double exponential fit\n",
    "start_de = 2\n",
    "stop_de = 21\n",
    "\n",
    "a = pe.fits.standard_fit(np.arange(start_de, stop_de), p_obs['f_P'][start_de:stop_de], func_2exp, initial_guess=[0.21, 14.0, 0.6, -10], resplot=True, expected_chisquare=True)\n",
    "[o.gamma_method() for o in a]\n",
    "print('Fit result:')\n",
    "print(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fitting with x-errors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first generate pseudo data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(Obs[-0.08(35)], Obs[0.19(25)])\n",
      "(Obs[1.85(35)], Obs[0.34(25)])\n",
      "(Obs[4.01(35)], Obs[-1.39(25)])\n",
      "(Obs[6.10(35)], Obs[-1.30(25)])\n",
      "(Obs[8.08(35)], Obs[-0.37(25)])\n"
     ]
    }
   ],
   "source": [
    "ox = []\n",
    "oy = []\n",
    "for i in range(0,10,2):\n",
    "    ox.append(pe.pseudo_Obs(i + 0.35 * np.random.normal(), 0.35, str(i)))\n",
    "    oy.append(pe.pseudo_Obs(np.sin(i) + 0.25 * np.random.normal() - 0.2 * i + 0.17, 0.25, str(i)))\n",
    "\n",
    "[o.gamma_method() for o in ox + oy]\n",
    "[print(o) for o in zip(ox, oy)];"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And choose a function to fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def func(a, x):\n",
    "    y = a[0] + a[1] * x + a[2] * anp.sin(x)\n",
    "    return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then fit this function to the data and get the fit parameter as Obs with the function `odr_fit` which uses orthogonal distance regression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fit with 3 parameters\n",
      "Method: ODR\n",
      "Sum of squares convergence\n",
      "Residual variance: 0.5988333933914471\n",
      "Parameter 1 : Obs[-0.01(29)]\n",
      "Parameter 2 : Obs[-0.165(55)]\n",
      "Parameter 3 : Obs[0.89(23)]\n"
     ]
    }
   ],
   "source": [
    "beta = pe.fits.odr_fit(ox, oy, func)\n",
    "\n",
    "pe.Obs.e_tag_global = 1 # Makes sure that the different samples with name length 1 are treated as ensembles and not as replica\n",
    "\n",
    "for i, item in enumerate(beta):\n",
    "    item.gamma_method()\n",
    "    print('Parameter', i + 1, ':', item)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the visulization we determine the value of the fit function in a range of x values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 21224 characters]\n",
      "text/plain": [
       "<Figure size 640x395.55 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "x_t = np.arange(min(ox).value - 1, max(ox).value + 1, 0.01)\n",
    "y_t = func([o.value for o in beta], x_t)\n",
    "\n",
    "plt.errorbar([e.value for e in ox], [e.value for e in oy], xerr=[e.dvalue for e in ox], yerr=[e.dvalue for e in oy], marker='D', lw=1, ls='none', zorder=10)\n",
    "plt.plot(x_t, y_t, '--')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also take a look at how much the inidividual ensembles contribute to the uncetainty of the fit parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameter 0\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 16492 characters]\n",
      "text/plain": [
       "<Figure size 640x395.55 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Parameter 1\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 17308 characters]\n",
      "text/plain": [
       "<Figure size 640x395.55 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Parameter 2\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 16944 characters]\n",
      "text/plain": [
       "<Figure size 640x395.55 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "for i, item in enumerate(beta):\n",
    "    print('Parameter', i)\n",
    "    item.plot_piechart()\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fitting with priors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When extracting energy levels and matrix elements from correlation functions one is interested in using as much data is possible in order to decrease the final error estimate and also have better control over systematic effects from higher states. This can in principle be achieved by fitting a tower of exponentials to the data. However, in practice it can be very difficult to fit a function with 6 or more parameters to noisy data. One way around this is to cnostrain the fit parameters with Bayesian priors. The principle idea is that any parameter which is determined by the data is almost independent of the priors while the additional parameters which would let a standard fit collapse are essentially constrained by the priors."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first generate fake data as a tower of three exponentials with noise which increases with temporal separation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 17500 characters]\n",
      "text/plain": [
       "<Figure size 640x395.55 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "m1 = 0.18\n",
    "m2 = 0.5\n",
    "m3 = 0.8\n",
    "\n",
    "A1 = 180\n",
    "A2 = 300\n",
    "A3 = 500\n",
    "\n",
    "px = []\n",
    "py = []\n",
    "for i in range(40):\n",
    "    px.append(i)\n",
    "    val = (A1 * np.exp(-m1 * i) + A2 * np.exp(-m2 * i) + A3 * np.exp(-m3 * i))\n",
    "    err = 0.03 * np.sqrt(i + 1)\n",
    "    tmp = pe.pseudo_Obs(val * (1 + err * np.random.normal()), val * err, 'e1')\n",
    "    py.append(tmp)\n",
    "    \n",
    "[o.gamma_method() for o in py];\n",
    "\n",
    "pe.plot_corrs([py], logscale=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As fit function we choose the sum of three exponentials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def func_3exp(a, x):\n",
    "    y = a[1] * anp.exp(-a[0] * x) + a[3] * anp.exp(-a[2] * x) + a[5] * anp.exp(-a[4] * x)\n",
    "    return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can specify the priors in a string format or alternatively input `Obs` from a previous analysis. It is important to choose the priors wide enough, otherwise they can influence the final result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "priors = ['0.2(4)', '200(500)', \n",
    "          '0.6(1.2)', '300(550)',\n",
    "          '0.9(1.8)', '400(700)']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is important to chose a sufficiently large value of `Obs.e_tag_global`, as every prior is given an ensemble id."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "pe.Obs.e_tag_global = 5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The fit can then be performed by calling `prior_fit` which in comparison to the standard fit requires the priors as additional input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fit with 6 parameters\n",
      "Method: migrad\n",
      "chisquare/d.o.f.: 1.100354109100944\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 29940 characters]\n",
      "text/plain": [
       "<Figure size 800x494.438 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Obs[0.1795(32)], Obs[186(14)], Obs[0.578(70)], Obs[597(170)], Obs[1.42(83)], Obs[239(173)]]\n"
     ]
    }
   ],
   "source": [
    "beta_p = pe.fits.prior_fit(px, py, func_3exp, priors, resplot=True)\n",
    "[o.gamma_method() for o in beta_p]\n",
    "print(beta_p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now observe how far the individual fit parameters are constrained by the data or the priors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 15196 characters]\n",
      "text/plain": [
       "<Figure size 640x395.55 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 16196 characters]\n",
      "text/plain": [
       "<Figure size 640x395.55 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 17836 characters]\n",
      "text/plain": [
       "<Figure size 640x395.55 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 18944 characters]\n",
      "text/plain": [
       "<Figure size 640x395.55 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 16348 characters]\n",
      "text/plain": [
       "<Figure size 640x395.55 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 18580 characters]\n",
      "text/plain": [
       "<Figure size 640x395.55 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "[o.plot_piechart() for o in beta_p];"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
