{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "    \n",
    "<b>Disclaimer:</b> \n",
    "    \n",
    "The main objective of the <i>Jupyter</i> notebooks is to show how to use the models of the <i>QENS library</i> by\n",
    "    \n",
    "- building a fitting model: composition of models, convolution with a resolution function  \n",
    "- setting and running the fit  \n",
    "- extracting and displaying information about the results  \n",
    "\n",
    "These steps have a minimizer-dependent syntax. That's one of the reasons why different minimizers have been used in the notebooks provided as examples.  \n",
    "But, note thatthe initial guessed parameters might not be optimal, resulting in a poor fit of the reference data.\n",
    "\n",
    "</div> \n",
    "\n",
    "# Two Lorentzian &lowast; resolution with lmfit\n",
    "\n",
    "## Introduction\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "    \n",
    "The objective of this notebook is to show how to use one of the models of \n",
    "the <a href=\"https://github.com/QENSlibrary/QENSmodels\">QENSlibrary</a>, <b>Lorentzian</b>, to perform some fits.\n",
    "<a href=\"https://lmfit.github.io/lmfit-py/\">lmfit</a> is used for fitting.\n",
    "</div>\n",
    "\n",
    "The following example uses the data from IRIS:\n",
    "- workspace_index=0, file: `irs26176_graphite002_red.nxs`  \n",
    "- related instrument resolution data `irs26173_graphite002_res.nxs`   \n",
    "\n",
    "The ISIS sample datasets can be downloaded from [Mantid's website](http://download.mantidproject.org/).\n",
    "The data used for this example are in the sample datafile: `data_2lorentzians.dat` and the instrument resolution datafile `irf_iris.dat`, respectively.\n",
    "\n",
    "This example is based on a [Mantid \"Fitting QENS Peaks\" tutorial](https://www.mantidproject.org/Fitting_QENS_Peaks).\n",
    "\n",
    "The implementation with `lmfit` is based on https://lmfit.github.io/lmfit-py/model.html\n",
    "\n",
    "This example requires an additional Python module `scipy.interpolate` to interpolate the tabulated data of the instrument resolution.\n",
    "\n",
    "### Physical units\n",
    "For information about unit conversion, please refer to the jupyter notebook called `Convert_units.ipynb` in the `tools` folder.\n",
    "\n",
    "The dictionary of units defined in the cell below specify the units of the refined parameters adapted to the convention used in the experimental datafile."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Units of parameters for selected QENS model and experimental data\n",
    "dict_physical_units = {'omega': \"meV\", \n",
    "                       'q': \"1/Angstrom\", \n",
    "                       'hwhm': \"meV\", \n",
    "                       'scale': \"unit_of_signal.meV\",\n",
    "                       'center': \"meV\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import ipywidgets\n",
    "import lmfit\n",
    "from scipy.interpolate import interp1d\n",
    "import QENSmodels\n",
    "\n",
    "path_to_data = './data/'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: load instrument resolution data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "irf_iris = np.loadtxt(path_to_data + 'irf_iris.dat')\n",
    "x_irf = irf_iris[:, 0]\n",
    "y_irf = irf_iris[:, 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: create function for instrument resolution data (cubic interpolation between tabulated data points)\n",
    "Create model: 2 lorentzians convoluted with instrument resolution: 6 parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def irf_gate(x):\n",
    "    \"\"\" \n",
    "    Function defined from the interpolation of instrument resolution data \n",
    "    Used to define fitting model and plot\n",
    "    \"\"\" \n",
    "    f = interp1d(x_irf, y_irf, kind='cubic', bounds_error=False, fill_value='extrapolate')\n",
    "\n",
    "    return f(x)\n",
    "\n",
    "# plot tabulated data and interpolated data\n",
    "xx = np.linspace(-.25, .25, 500)\n",
    "\n",
    "fig0 = plt.figure()\n",
    "plt.plot(x_irf, y_irf, 'b.', label='tabulated data')\n",
    "plt.plot(xx, irf_gate(xx), 'g--', label='extrapolated data')\n",
    "plt.legend()\n",
    "plt.xlabel('Energy transfer (meV)')\n",
    "plt.title('Instrument resolution: plot tabulated data and interpolated data')\n",
    "plt.grid();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: create \"double lorentzian\" profile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_2lorentzians(x, scale1, center1, hwhm1, scale2, center2, hwhm2):\n",
    "    return QENSmodels.lorentzian(\n",
    "        x, \n",
    "        scale1, \n",
    "        center1, \n",
    "        hwhm1\n",
    "    ) + QENSmodels.lorentzian(\n",
    "        x, \n",
    "        scale2, \n",
    "        center2, \n",
    "        hwhm2\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: create convolution function\n",
    "\n",
    "Code from https://github.com/jmborr/qef"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convolve(model, resolution):\n",
    "    c = np.convolve(model, resolution, mode='valid')\n",
    "    if len(model) % len(resolution) == 0:\n",
    "        c = c[:-1]\n",
    "    return c\n",
    "\n",
    "class Convolve(lmfit.CompositeModel):\n",
    "    def __init__(self, resolution, model, **kws):\n",
    "        super(Convolve, self).__init__(resolution, model, convolve, **kws)\n",
    "        self.resolution = resolution\n",
    "        self.model = model\n",
    "\n",
    "    def eval(self, params=None, **kwargs):\n",
    "        res_data = self.resolution.eval(params=params, **kwargs)\n",
    "        \n",
    "        # evaluate model on an extended energy range to avoid boundary effects\n",
    "        independent_var = self.resolution.independent_vars[0]\n",
    "        e = kwargs[independent_var]  # energy values\n",
    "        neg_e = min(e) - np.flip(e[np.where(e > 0)], axis=0)\n",
    "        pos_e = max(e) - np.flip(e[np.where(e < 0)], axis=0)\n",
    "        e = np.concatenate((neg_e, e, pos_e))\n",
    "        kwargs.update({independent_var: e})\n",
    "        model_data = self.model.eval(params=params, **kwargs)\n",
    "        \n",
    "        # Multiply by the X-spacing to preserve normalization\n",
    "        de = (e[-1] - e[0])/(len(e) - 1)  # energy spacing\n",
    "        return de * convolve(model_data, res_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Load reference data - extract x and y values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "two_lorentzians_iris = np.loadtxt(path_to_data + 'data_2lorentzians.dat')\n",
    "xx = two_lorentzians_iris[:, 0]\n",
    "yy = two_lorentzians_iris[:, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Convolve(lmfit.Model(irf_gate), lmfit.Model(model_2lorentzians))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6: Fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = model.fit(yy, x=xx, scale1=1., center1=0., hwhm1=0.25, scale2=5., center2=0., hwhm2=0.02)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot initial configuration\n",
    "fig1 = plt.figure()\n",
    "plt.plot(xx, yy, '+', label='experimental data')\n",
    "plt.plot(xx, result.init_fit, 'k--', label='initial model')\n",
    "plt.legend(loc='upper right')\n",
    "plt.xlabel('Energy transfer (meV)') \n",
    "plt.title('Plot before fitting: experimental data and mode with initial guesses')\n",
    "plt.grid();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 7: Plotting results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Result of fit:\\n', result.fit_report())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot experimental data and best fit using `matplotlib.pyplot`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig2 = plt.figure()\n",
    "plt.plot(xx, yy, '+', label='experimental data')\n",
    "plt.plot(xx, result.best_fit, 'r-', label='best fit')\n",
    "plt.grid()\n",
    "plt.xlabel('Energy transfer (meV)')\n",
    "plt.title('Plot selected fitting results: experimental data and best fit')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Other option: use lmfit features to plot result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "result.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Print values and errors of refined parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for item in result.params.keys():\n",
    "    print(f'{item[:-1]}: {result.params[item].value} +/- {result.params[item].stderr} {dict_physical_units[item[:-1]]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "hide_input": false,
  "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.9"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}