{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from mpes import fprocessing as fp, visualization as vis\n",
    "import matplotlib.cm as cm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "fdir = r'../data/data_20180605_119.h5'\n",
    "axes = ['X', 'Y', 't']\n",
    "# Pick very odd number of bins to create the artefact\n",
    "bins = [91, 91, 230]\n",
    "ranges = [(300, 1700), (300, 1700), (68500, 72500)]\n",
    "enum_max = 4e7\n",
    "hp = fp.hdf5Processor(fdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Binning without data jittering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "binned_nojitter = hp.localBinning(axes=axes, nbins=bins, ranges=ranges, amax=enum_max)\n",
    "binned_nj = binned_nojitter['binned']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Binning with data jittering\n",
    "Jitter amplitude used = 0.1 * bin size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "binned_jitter = hp.localBinning(axes, bins, ranges, jittered=True, \\\n",
    "               jitter_amplitude=[1e-1]*3, amax=enum_max);\n",
    "binned_jt = binned_jitter['binned']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparison plots\n",
    "#### (1) 2D cuts of the band structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 92396 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "f, axs = plt.subplots(1, 2, figsize=(10,5))\n",
    "vis.colormesh2d((binned_nj[50,::-1,::-1]).T, colormap='terrain_r',\\\n",
    "           vmin=0, cbar=True, plotaxes=axs[0])\n",
    "vis.colormesh2d(binned_jt[50,::-1,::-1].T, colormap='terrain_r',\\\n",
    "           vmin=0, cbar=True, plotaxes=axs[1])\n",
    "axs[0].set_title('Without jittering', fontsize=20)\n",
    "axs[1].set_title('With jittering', fontsize=20);\n",
    "#plt.savefig('Jittering_2D.png', bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 43148 characters]\n",
      "text/plain": [
       "<Figure size 360x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fd, axd = plt.subplots(1, 1, figsize=(5,5))\n",
    "vis.colormesh2d(binned_nj[50,::-1,::-1].T - binned_jt[50,::-1,::-1].T, colormap='RdBu_r', \\\n",
    "                cscale={'vmin':-20, 'vmax':20}, plotaxes=axd);\n",
    "# plt.savefig('Jittering_2D_diff.png', bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### (2) Angle-integrated energy dispersion curve (EDC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 56792 characters]\n",
      "text/plain": [
       "<Figure size 720x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "nj_1d_EDC = binned_nj.sum(axis=(0,1))\n",
    "jt_1d_EDC = binned_jt.sum(axis=(0,1))\n",
    "fedc, axedc = plt.subplots(figsize=(10,6))\n",
    "plt.plot(binned_nojitter['t'], nj_1d_EDC, lw=2, label='without jittering')\n",
    "plt.plot(binned_jitter['t'], jt_1d_EDC, 'r', lw=2, label='with jittering')\n",
    "plt.plot(binned_jitter['t'], nj_1d_EDC - jt_1d_EDC, 'k', lw=2, label='difference')\n",
    "axedc.set_xlabel('Energy (a.u.)', fontsize=15)\n",
    "axedc.set_ylabel('Photoemission intensity (a.u.)', fontsize=15)\n",
    "plt.legend();\n",
    "#plt.savefig('Jittering_1D.png', bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### (3) Energy-integrated momentum dispersion curve (MDC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 48904 characters]\n",
      "text/plain": [
       "<Figure size 720x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "nj_1d_MDC = binned_nj[:,:,100:].sum(axis=(0,2))\n",
    "jt_1d_MDC = binned_jt[:,:,100:].sum(axis=(0,2))\n",
    "fmdc, axmdc = plt.subplots(figsize=(10,6))\n",
    "plt.plot(binned_nojitter['Y'], nj_1d_MDC, lw=2, \\\n",
    "         label='without jittering')\n",
    "plt.plot(binned_jitter['Y'], jt_1d_MDC, 'r', lw=2, \\\n",
    "         label='with jittering')\n",
    "plt.plot(binned_jitter['Y'], nj_1d_MDC - jt_1d_MDC, 'k', lw=2, label='difference')\n",
    "axmdc.set_xlabel('Momentum k$_y$ (a.u.)', fontsize=15)\n",
    "axmdc.set_ylabel('Photoemission intensity (a.u.)', fontsize=15)\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Combine the plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 79416 characters]\n",
      "text/plain": [
       "<Figure size 360x360 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# rect [left, bottom, width, height]\n",
    "f, ax = plt.subplots(figsize=(5, 5))\n",
    "vis.colormesh2d(binned_nj[50,::-1,::-1].T - binned_jt[50,::-1,::-1].T, colormap='RdBu_r', \\\n",
    "                cscale={'vmin':-20, 'vmax':20}, plotaxes=ax);\n",
    "ax.set_xticks([])\n",
    "ax.set_yticks([])\n",
    "ax.xaxis.set_label_position(\"top\")\n",
    "ax.set_xlabel('Momentum ($\\AA^{-1}$)', fontsize=15)\n",
    "ax.set_ylabel('Energy (eV)', fontsize=15)\n",
    "\n",
    "ax_edc = f.add_axes([0.98, 0.08, 0.25, 0.88])\n",
    "ax_edc.plot(nj_1d_EDC[::-1], binned_nojitter['t'], lw=2, label='without jittering')\n",
    "ax_edc.plot(jt_1d_EDC[::-1], binned_jitter['t'], 'r', lw=2, label='with jittering')\n",
    "ax_edc.plot(nj_1d_EDC[::-1] - jt_1d_EDC[::-1], binned_jitter['t'], 'k', lw=2, label='difference')\n",
    "ax_edc.set_xticks([])\n",
    "ax_edc.set_yticks([])\n",
    "ax_edc.set_xlabel('Intensity', fontsize=15)\n",
    "ax_edc.xaxis.set_label_position(\"top\")\n",
    "\n",
    "ax_mdc = f.add_axes([0.1, -0.2, 0.85, 0.25])\n",
    "ax_mdc.plot(binned_nojitter['Y'], nj_1d_MDC, lw=2, \\\n",
    "         label='w/o jittering')\n",
    "ax_mdc.plot(binned_jitter['Y'], jt_1d_MDC, 'r', lw=2, \\\n",
    "         label='w/ jittering')\n",
    "ax_mdc.plot(binned_jitter['Y'], nj_1d_MDC - jt_1d_MDC, 'k', lw=2, label='difference')\n",
    "ax_mdc.set_xticks([])\n",
    "ax_mdc.set_yticks([])\n",
    "ax_mdc.set_ylabel('Intensity', fontsize=15)\n",
    "\n",
    "# ax_lg = f.add_axes([0.98, -0.2, 0.25, 0.25])\n",
    "ax_mdc.legend(loc='lower right', bbox_to_anchor=[1.1, 0, 0.25, 0.25], frameon=False, labelspacing=1.5)\n",
    "# ax_lg.set_xticks([])\n",
    "# ax_lg.set_yticks([])\n",
    "# plt.savefig('./Tutfiles/Jittered_Panel.png', dpi=300, bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parallel version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob as g\n",
    "from time import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "hfiles = g.glob(r'../data/data_114_parts/data*.h5')\n",
    "_, hfsorted = fp.sortNamesBy(hfiles, r'\\d+', slicerange=(-10, None))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "parp = fp.parallelHDF5Processor(hfsorted[:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:09<00:00,  9.74s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Binning took about 9 s.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "tstart = time()\n",
    "parp.parallelBinning(axes=axes, nbins=bins, ranges=ranges, scheduler='threads',\n",
    "                     binning_kwds={'jittered':True, 'jitter_amplitude':[2e-1]*3}, ret=False)\n",
    "tend = time()\n",
    "print('Binning took about {} s.'.format(int(tend - tstart)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = parp.combineResults()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f51b4ba6fd0>"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 51140 characters]\n",
      "text/plain": [
       "<Figure size 720x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(10,6))\n",
    "plt.imshow(result['binned'][50,...].T, cmap='terrain_r', aspect=0.4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f51b4a9b4e0>]"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 16604 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(result['binned'].sum(axis=(0,1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f51b4a71320>]"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 15300 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(result['binned'][:,:,100:].sum(axis=(0,2)))"
   ]
  },
  {
   "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
}
