{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use N-dimensional Histogram functionality and Interpolation\n",
    "  \n",
    "* could be used for example to read and interpolate an lookup table or IRF.\n",
    "* In this example, we load a sample energy reconstruction lookup-table from a FITS file\n",
    "* In this case it is only in 2D cube (to keep the file size small): `SIZE` vs `IMPACT-DISTANCE`, however the same method will work for any dimensionality"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import matplotlib.pylab as plt\n",
    "import numpy as np\n",
    "from astropy.io import fits\n",
    "from scipy.interpolate import RegularGridInterpolator\n",
    "\n",
    "from ctapipe.utils.datasets import get_dataset_path\n",
    "from ctapipe.utils import Histogram\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### load an example datacube \n",
    "(an energy table generated for a\n",
    "small subset of HESS simulations) to use as a lookup table. Here\n",
    "we will use the `Histogram` class, which automatically loads both\n",
    "the data cube and creates arrays for the coordinates of each\n",
    "axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "testfile = get_dataset_path(\"hess_ct001_energylut.fits.gz\")\n",
    "energy_hdu = fits.open(testfile)['MEAN']\n",
    "energy_table = Histogram.from_fits(energy_hdu)\n",
    "print(energy_table)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### construct an interpolator that we can use to get values at any point:\n",
    "\n",
    "Here we will use a `RegularGridInterpolator`, since it is the most appropriate for this type of data, but others are available (see the SciPy documentation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "centers = [energy_table.bin_centers(ii) for ii in range(energy_table.ndims)]\n",
    "energy_lookup = RegularGridInterpolator(centers, energy_table.hist,\n",
    "                                        bounds_error=False,fill_value=-100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`energy_lookup` is now just a continuous function of `log(SIZE)`,\n",
    "`DISTANCE` in m. \n",
    "\n",
    "### Now plot some curves from the interpolator.  \n",
    "\n",
    "Note that the LUT we used is does not have very high statistics,\n",
    "so the interpolation starts to be affected by noise at the high\n",
    "end. In a real case, we would want to use a table that has been\n",
    "sanitized (smoothed and extrapolated)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "lsize = np.linspace(1.5, 5.0, 100)\n",
    "dists = np.linspace(50, 100, 5)\n",
    "\n",
    "plt.figure()\n",
    "plt.title(\"Variation of energy with size and impact distance\")\n",
    "plt.xlabel(\"SIZE (P.E.)\")\n",
    "plt.ylabel(\"ENERGY (TeV)\")\n",
    "\n",
    "for dist in dists:\n",
    "    plt.plot(10**lsize, 10**energy_lookup((lsize, dist)), '+-',\n",
    "             label=\"DIST={:.1f} m\".format(dist))\n",
    "\n",
    "plt.legend(loc=\"best\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the interpolator, reinterpolate the lookup table onto an $N \\times N$\n",
    "grid (regardless of its original dimensions):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "N = 300\n",
    "xmin, xmax = energy_table.bin_centers(0)[0], energy_table.bin_centers(0)[-1]\n",
    "ymin, ymax = energy_table.bin_centers(1)[0], energy_table.bin_centers(1)[-1]\n",
    "xx,yy = np.linspace(xmin,xmax,N), np.linspace(ymin, ymax, N)\n",
    "X, Y = np.meshgrid( xx,yy )\n",
    "points = list(zip(X.ravel(),Y.ravel()))\n",
    "E = energy_lookup(points).reshape((N, N))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's plot the original table and the new one (E).  The color bar shows $\\log_{10}(E)$ in TeV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12,5))\n",
    "plt.nipy_spectral()\n",
    "\n",
    "# the uninterpolated table\n",
    "plt.subplot(1,2,1)\n",
    "plt.xlim(1.5,5)\n",
    "plt.ylim(0,500)\n",
    "plt.xlabel(\"log10(SIZE)\")\n",
    "plt.ylabel(\"Impact Dist (m)\")\n",
    "plt.pcolormesh(energy_table.bin_centers(0), energy_table.bin_centers(1), energy_table.hist.T)\n",
    "plt.title(\"Raw table, uninterpolated {0}\".format(energy_table.hist.T.shape))\n",
    "cb = plt.colorbar()\n",
    "cb.set_label(\"$\\log_{10}(E/\\mathrm{TeV})$\")\n",
    "\n",
    "# the interpolated table\n",
    "plt.subplot(1,2,2)\n",
    "plt.pcolormesh(np.linspace(xmin,xmax,N), np.linspace(ymin, ymax, N),E)\n",
    "plt.xlim(1.5,5)\n",
    "plt.ylim(0,500)\n",
    "plt.xlabel(\"log10(SIZE)\")\n",
    "plt.ylabel(\"Impact Dist(m)\")\n",
    "plt.title(\"Interpolated to a ({0}, {0}) grid\".format(N))\n",
    "cb = plt.colorbar()\n",
    "cb.set_label(\"$\\log_{10}(E/\\mathrm{TeV})$\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the high-stats central region, we get a nice smooth interpolation function. Of course we can see that there are a few more steps to take before using this table:\n",
    "* need to deal with cases where the table had low stats near the edges (smooth or extrapolate, or set bounds)\n",
    "* may need to smooth the table even where there are sufficient stats, to avoid wiggles"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
