{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pesfit as pf\n",
    "from silx.io.dictdump import dicttoh5, h5todict\n",
    "import matplotlib.pyplot as plt\n",
    "import lmfit.models as lmm\n",
    "from tqdm import notebook as nb\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.1 Fitting of a patch of photoemission spectra (patch fitting)\n",
    "The following code demonstrate a few levels of lineshape fitting automation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(14, 30, 30)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load energy initialization\n",
    "# krec = h5todict(r'E:\\Momentoscope\\WSe2\\krec_lda_kpoint.h5')['recon']\n",
    "krec = h5todict(r'..\\data\\WSe2\\recon\\kpoint\\kpoint_LDA_recon.h5')['recon']\n",
    "krec.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load photoemission data\n",
    "fdata = h5todict(r'..\\data\\WSe2\\pes\\kpoint\\kpoint_symmetrized.h5')\n",
    "Es = fdata['E'][20:100] # Energy values\n",
    "eofs = 0.015 # Energy offset\n",
    "sfts = np.arange(-0.08, 0.09, 0.01) # Choices of random shifts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1.1 Serial patch fitting\n",
    "#### (a) Fitting using explicit coded steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "bpatch = np.zeros((2, 30, 30))\n",
    "spatch = np.zeros((2, 30, 30))\n",
    "gpatch = np.zeros((2, 30, 30))\n",
    "\n",
    "# Initialize the model\n",
    "mdl = pf.fitter.model_generator(peaks={'Voigt':2}, background='None')\n",
    "pars = mdl.make_params()\n",
    "\n",
    "vardict = [{'lp1_':{'amplitude':dict(value=0.2, min=0, max=2, vary=True),\n",
    "                  'sigma':dict(value=0.1, min=0.05, max=2, vary=False),\n",
    "                  'gamma':dict(value=0.02, min=0, max=2, vary=True)}},\n",
    "           \n",
    "           {'lp2_':{'amplitude':dict(value=0.2, min=0, max=2, vary=True),\n",
    "                   'sigma':dict(value=0.1, min=0.05, max=2, vary=False),\n",
    "                   'gamma':dict(value=0.02, min=0, max=2, vary=True)}}]\n",
    "\n",
    "pf.fitter.varsetter(pars, vardict, ret=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6ae4f8b4cb4048dba0238b29680b58ef",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=30.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "for r in nb.tqdm(range(30)):\n",
    "#     print(r)\n",
    "    for c in range(30):\n",
    "\n",
    "        row_id, col_id = r, c\n",
    "\n",
    "        # Update the band position estimates (corresponding to the centers of the lineshape components)\n",
    "        centerdict = [{'lp1_':{'center':dict(value=krec[0,row_id,col_id]+eofs, vary=True)}},\n",
    "                      {'lp2_':{'center':dict(value=krec[1,row_id,col_id]+eofs, vary=True)}}]\n",
    "\n",
    "        pf.fitter.varsetter(pars, centerdict, ret=False)\n",
    "        y = fdata['V'][row_id, col_id, 20:100]\n",
    "        fit_result = mdl.fit(y/y.max(), pars, x=Es)\n",
    "\n",
    "        # Add random shift to the initialization to produce better fits\n",
    "        out = pf.fitter.random_varshift(fit_result, model=mdl, params=pars, yvals=y, xvals=Es, shifts=sfts)\n",
    "        comps = out.eval_components(x=Es)\n",
    "\n",
    "        bpatch[0,r,c] = out.params['lp1_center'].value\n",
    "        bpatch[1,r,c] = out.params['lp2_center'].value\n",
    "\n",
    "        spatch[0,r,c] = out.params['lp1_sigma'].value\n",
    "        spatch[1,r,c] = out.params['lp2_sigma'].value\n",
    "\n",
    "        gpatch[0,r,c] = out.params['lp1_gamma'].value\n",
    "        gpatch[1,r,c] = out.params['lp2_gamma'].value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0.98, 'Band #2 dispersion surface around K')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 30596 characters]\n",
      "text/plain": [
       "<Figure size 936x432 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "band_ind = 2\n",
    "f, ax = plt.subplots(1, 2, figsize=(13, 6))\n",
    "ax[0].imshow(bpatch[band_ind-1,...], origin='lower')\n",
    "ax[1].contourf(bpatch[band_ind-1,...], 20)\n",
    "ax[0].set_title('Pixelated', fontsize=15)\n",
    "ax[1].set_title('Contoured', fontsize=15)\n",
    "plt.suptitle('Band #{} dispersion surface around K'.format(band_ind), fontsize=15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### (b) Fitting using ``pesfit.fitter.pointwise_fitting()`` method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "eofs = 0.015 # Energy offset\n",
    "mdl = pf.fitter.model_generator(peaks={'Voigt':2}, background='None')\n",
    "pars = mdl.make_params()\n",
    "\n",
    "vardict = [{'lp1_':{'amplitude':dict(value=0.2, min=0, max=2, vary=True),\n",
    "                  'sigma':dict(value=0.1, min=0.05, max=2, vary=False),\n",
    "                  'gamma':dict(value=0.02, min=0, max=2, vary=True)}},\n",
    "           \n",
    "           {'lp2_':{'amplitude':dict(value=0.2, min=0, max=2, vary=True),\n",
    "                   'sigma':dict(value=0.1, min=0.05, max=2, vary=False),\n",
    "                   'gamma':dict(value=0.02, min=0, max=2, vary=True)}}]\n",
    "pf.fitter.varsetter(pars, vardict, ret=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "787db3fa00da452aa076712480eccf1d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=900.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# DataFrame torage of fitting parameters\n",
    "df_fit = pd.DataFrame(columns=pars.keys())\n",
    "\n",
    "erange = slice(20, 100)\n",
    "Es = fdata['E'][erange] # Energy values\n",
    "ys = fdata['V'][..., erange] # Patch of photoemission spectra\n",
    "nspec = np.prod(ys.shape[:-1]) # Number of spectra\n",
    "\n",
    "# Partially flatten data to 2D for ease of iteration\n",
    "ys2d = pf.utils.partial_flatten(ys, axis=(0, 1))\n",
    "krec2d = pf.utils.partial_flatten(krec, axis=(1, 2))\n",
    "\n",
    "for d in nb.tqdm(range(nspec)):\n",
    "    \n",
    "    centerdict = [{'lp1_':{'center':dict(value=krec2d[0,d]+eofs, vary=True)}},\n",
    "                  {'lp2_':{'center':dict(value=krec2d[1,d]+eofs, vary=True)}}]\n",
    "\n",
    "    pf.fitter.varsetter(pars, centerdict, ret=False)\n",
    "        \n",
    "    ydata = ys2d[d, :]\n",
    "    # Carry out fitting for each spectrum\n",
    "    out = pf.fitter.pointwise_fitting(Es, ydata, model=mdl, params=pars, jitter_init=False,\n",
    "                                      shifts=np.arange(-0.08, 0.09, 0.01), verbose=True)\n",
    "#     pf.fitter.plot_fit_result(out, Es)\n",
    "\n",
    "    dfout = pf.utils.df_collect(out.params, currdf=df_fit)\n",
    "    df_fit = dfout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "band_ind = 2\n",
    "band_energy = df_fit['lp'+str(band_ind)+'_center'].values.reshape((30, 30))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0.98, 'Band #2 dispersion surface around K')"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 30984 characters]\n",
      "text/plain": [
       "<Figure size 936x432 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "f, ax = plt.subplots(1, 2, figsize=(13, 6))\n",
    "ax[0].imshow(band_energy)\n",
    "ax[1].contourf(band_energy, 20, origin='upper')\n",
    "ax[0].set_title('Pixelated', fontsize=15)\n",
    "ax[1].set_title('Contoured', fontsize=15)\n",
    "plt.suptitle('Band #{} dispersion surface around K'.format(band_ind), fontsize=15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1.2 Fitting using ``pesfit.fitter.PatchFitter`` class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "fdata = h5todict(r'..\\data\\WSe2\\pes\\kpoint\\kpoint_symmetrized.h5')\n",
    "\n",
    "kfit = pf.fitter.PatchFitter(xdata=fdata['E'], ydata=fdata['V'])\n",
    "# kfit.load_spec_data(fdir=r'..\\data\\WSe2\\pes\\kpoint\\kpoint_symmetrized.h5', parts='V')\n",
    "\n",
    "vardict = [{'lp1_':{'amplitude':dict(value=0.2, min=0, max=2, vary=1.),\n",
    "                  'sigma':dict(value=0.1, min=0.05, max=2, vary=False),\n",
    "                  'gamma':dict(value=0.02, min=0, max=2, vary=True),\n",
    "                  'center':dict(vary=True)}},\n",
    "           \n",
    "           {'lp2_':{'amplitude':dict(value=0.2, min=0, max=2, vary=True),\n",
    "                   'sigma':dict(value=0.1, min=0.05, max=2, vary=False),\n",
    "                   'gamma':dict(value=0.02, min=0, max=2, vary=True),\n",
    "                   'center':dict(vary=True)}}]\n",
    "\n",
    "kfit.set_inits(inits_dict=vardict, band_inits=krec, drange=slice(20,100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b9b34e82d3854c4a81e0ef19c6b143ac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=900.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Sequentially fit each line spectrum\n",
    "kfit.sequential_fit(pbar=True, jitter_init=True, shifts=np.arange(-0.08, 0.09, 0.01), nspec=900)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "band_ind = 2\n",
    "band_energy = kfit.df_fit['lp'+str(band_ind)+'_center'].values.reshape((30, 30))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0.98, 'Band #2 dispersion surface around K')"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 31012 characters]\n",
      "text/plain": [
       "<Figure size 936x432 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "f, ax = plt.subplots(1, 2, figsize=(13, 6))\n",
    "ax[0].imshow(band_energy)\n",
    "ax[1].contourf(band_energy, 20, origin='upper')\n",
    "ax[0].set_title('Pixelated', fontsize=15)\n",
    "ax[1].set_title('Contoured', fontsize=15)\n",
    "plt.suptitle('Band #{} dispersion surface around K'.format(band_ind), fontsize=15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'n' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-5-7b70b1c5f504>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m()\u001b[0m\n\u001b[1;32m----> 1\u001b[1;33m \u001b[0mkfit\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mparallel_fit\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mpbar\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;32mTrue\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mjitter_init\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;32mFalse\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mshifts\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0marange\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m-\u001b[0m\u001b[1;36m0.08\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;36m0.09\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;36m0.01\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mnspec\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;36m900\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;32mC:\\ProgramData\\Anaconda3\\lib\\site-packages\\pesfit\\fitter.py\u001b[0m in \u001b[0;36mparallel_fit\u001b[1;34m(self, varkeys, other_initvals, **kwds)\u001b[0m\n\u001b[0;32m    421\u001b[0m         \u001b[0mnspec\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mkwds\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mpop\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m'nspec'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mnspec\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    422\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 423\u001b[1;33m         \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mout\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mparallel_fitting\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mxvals\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mydata2D\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmodel\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0minits_persist\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mband_inits2D\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mnspec\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mprefixes\u001b[0m\u001b[1;33m,\u001b[0m\u001b[0mvarkeys\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mother_initvals\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    424\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    425\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0msave_data\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mfdir\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;34mr'./'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mfname\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;34m''\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mftype\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;34m'h5'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mkeyname\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;34m'fitres'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0morient\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;34m'dict'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mC:\\ProgramData\\Anaconda3\\lib\\site-packages\\pesfit\\fitter.py\u001b[0m in \u001b[0;36mparallel_fitting\u001b[1;34m(xvals, yvals, model, inits_persist, band_inits2D, nspec, prefixes, varkeys, other_initvals, **kwds)\u001b[0m\n\u001b[0;32m    493\u001b[0m             \u001b[1;32mraise\u001b[0m \u001b[0mException\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m'other_initvals has incorrect shape!'\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    494\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 495\u001b[1;33m     \u001b[0mother_inits\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0minits_vary_vals\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;33m...\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mn\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mT\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    496\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    497\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0m_single_fit\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0myvals\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mvarkeys\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mother_inits\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mNameError\u001b[0m: name 'n' is not defined"
     ]
    }
   ],
   "source": [
    "kfit.parallel_fit(pbar=True, jitter_init=False, shifts=np.arange(-0.08, 0.09, 0.01), nspec=900)"
   ]
  },
  {
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
