{
 "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data has been written using pyerrors 2.0.0.\n",
      "Format version 0.1\n",
      "Written by fjosw on 2022-01-06 11:27:34 +0100 on host XPS139305, Linux-5.11.0-44-generic-x86_64-with-glibc2.29\n",
      "\n",
      "Description:  SF correlation function f_P on a test ensemble\n"
     ]
    }
   ],
   "source": [
    "fP = pe.Corr(pe.input.json.load_json(\"./data/f_P\"), padding=[1, 1])\n",
    "fP.gamma_method()"
   ]
  },
  {
   "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.0023324250917749687\n",
      "\n",
      " Goodness of fit:\n",
      "χ²/d.o.f. = 0.002332\n",
      "p-value   = 1.0000\n",
      "Fit parameters:\n",
      "0\t 0.2036(92)\n",
      "1\t 16.3(1.3)\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 28204 characters]\n",
      "text/plain": [
       "<Figure size 800x494.438 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "start_fit = 9\n",
    "stop_fit = 18\n",
    "\n",
    "fit_result = fP.fit(func_exp, [start_fit, stop_fit], resplot=True)\n",
    "fit_result.gamma_method()\n",
    "print(\"\\n\", fit_result)"
   ]
  },
  {
   "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.009831165592706342\n",
      "Normalized covariance:  0.8384671239654656\n"
     ]
    }
   ],
   "source": [
    "cov_01 = pe.fits.covariance(fit_result[0], fit_result[1])\n",
    "print('Covariance: ', cov_01)\n",
    "print('Normalized covariance: ', cov_01 / fit_result[0].dvalue / fit_result[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_fP = fP.m_eff()\n",
    "m_eff_fP.tag = r\"Effective mass of f_P\""
   ]
  },
  {
   "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": [
      "Fit with 1 parameters\n",
      "Method: Levenberg-Marquardt\n",
      "`ftol` termination condition is satisfied.\n",
      "chisquare/d.o.f.: 0.13241808096937788\n",
      "\n",
      "Effective mass:\t 0.2057(68)\n",
      "Fitted mass:\t 0.2036(92)\n"
     ]
    }
   ],
   "source": [
    "m_eff_fP.gamma_method()\n",
    "m_eff_plateau = m_eff_fP.plateau([start_fit, stop_fit])\n",
    "m_eff_plateau.gamma_method()\n",
    "print()\n",
    "print('Effective mass:\\t', m_eff_plateau)\n",
    "print('Fitted mass:\\t', fit_result[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now visualize the effective mass compared to the result of the fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 21072 characters]\n",
      "text/plain": [
       "<Figure size 640x395.55 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "m_eff_fP.show(plateau=fit_result[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fitting with x-errors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first generate pseudo data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(Obs[0.53(35)], Obs[0.38(25)])\n",
      "(Obs[1.73(35)], Obs[0.59(25)])\n",
      "(Obs[3.92(35)], Obs[-1.23(25)])\n",
      "(Obs[5.73(35)], Obs[-1.18(25)])\n",
      "(Obs[7.74(35)], Obs[-0.40(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": 11,
   "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": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fit with 3 parameters\n",
      "Method: ODR\n",
      "Sum of squares convergence\n",
      "Residual variance: 0.08780824312692749\n",
      "Parameter 1 : 0.06(25)\n",
      "Parameter 2 : -0.160(47)\n",
      "Parameter 3 : 0.80(18)\n"
     ]
    }
   ],
   "source": [
    "beta = pe.fits.total_least_squares(ox, oy, func)\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": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 19492 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": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameter 0\n",
      "\n",
      "Parameter 1\n",
      "\n",
      "Parameter 2\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 17764 characters]\n",
      "text/plain": [
       "<Figure size 640x395.55 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 16084 characters]\n",
      "text/plain": [
       "<Figure size 640x395.55 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 16860 characters]\n",
      "text/plain": [
       "<Figure size 640x395.55 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i, item in enumerate(beta):\n",
    "    print('Parameter', i)\n",
    "    item.plot_piechart()\n",
    "    print()"
   ]
  },
  {
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
