{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `qp` practical example\n",
    "\n",
    "_Alex Malz, Phil Marshall, Eric Charles_\n",
    "\n",
    "In this notebook we use the `qp` module to study some photo-Z PDFs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup, reading ensemble"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Imports\n",
    "\n",
    "First let's import the packages we will need for this notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "\n",
    "import matplotlib\n",
    "matplotlib.use('Agg')\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "import qp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Data files\n",
    "\n",
    "Now lets download the data files we will need, if we haven't already"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_url = 'https://slac.stanford.edu/~echarles/qp_example'\n",
    "if not os.path.exists('qp_test_ensemble.hf5'):\n",
    "    os.system('curl -o %s -OL %s/%s' % ('qp_test_ensemble.hf5', base_url, 'qp_test_ensemble.hf5'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Reading ensemble\n",
    "\n",
    "Now we read the ensemble, note that we only need the name of the data file, the name of the metadata file is assumed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ens = qp.read('qp_test_ensemble.hf5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exploration\n",
    "\n",
    "This will show use that the "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Confirm that we have read the ensembles\n",
    "print(\"Ensemble = \", ens)\n",
    "# Print some simple information about the ensemble\n",
    "print(\"Rep      = \", ens.gen_class.name)\n",
    "print(\"NPDF     = \", ens.npdf)\n",
    "print(\"Metadata = \", ens.metadata())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting Example\n",
    "\n",
    "Now we are going to plot some PDFs from the ensemble\n",
    "Note that the first call to the `plot` specifies the x-axis limits, but does not specify the key (i.e., which PDF in the ensemble), so that defaults to 0 (i.e., the first PDF)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "axes = ens.plot(xlim=(0., 0.5), label=\"PDF 0\")\n",
    "_ = ens.plot(key=1, axes=axes, label=\"PDF 1\")\n",
    "_ = ens.plot(key=1300, axes=axes, label=\"PDF 1300\")\n",
    "legend = axes.figure.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Timing benchmarks\n",
    "\n",
    "Now we are going to extract some information from the ensemble and time how long it takes to do so"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# These are the grid points and quantiles at which we will extract values\n",
    "test_xvals = ens.gen_obj.xvals\n",
    "test_quantiles = np.linspace(0, 1, 51)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Time the pdf evaluation for 20000 PDFs\n",
    "pdfs = ens.pdf(test_xvals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Time the cdf (Cumulative distribution function) evaluation for 20000 PDFs\n",
    "cdfs = ens.cdf(test_xvals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "ppfs = ens.ppf(test_quantiles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Time the sf (survival fraction, 1-cdf) evaluation for 20000 PDFs\n",
    "sfs = ens.sf(test_xvals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Time the isf (inverse survival fraction) evaluation for 20000 PDFs\n",
    "isfs = ens.isf(test_quantiles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Time the generation of 100 samples for each of the 20000 PDFs\n",
    "samples = ens.rvs(size=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Changing the grid for the representation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Convert to a grid using 51 grid points\n",
    "ens_g51 = qp.convert(ens, 'interp', xvals=np.linspace(0, 3, 51))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "# Convert to a grid using 21 grid points\n",
    "ens_g21 = qp.convert(ens, 'interp', xvals=np.linspace(0, 1, 21))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = 0\n",
    "axes_g = ens.plot(key, xlim=(0, 0.5), label=\"orig\")\n",
    "_ = ens_g51.plot(key, axes=axes_g, label=\"g51\")\n",
    "_ = ens_g21.plot(key, axes=axes_g, label=\"g21\")\n",
    "leg_g = axes_g.figure.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conversion to other represenations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####  quantile representaion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "# Convert using 51 quantiles\n",
    "ens_q51 = qp.convert(ens, 'quant', quants=np.linspace(0.01, 0.99, 51))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Convert using 21 quantiles\n",
    "ens_q21 = qp.convert(ens, 'quant', quants=np.linspace(0.01, 0.99, 21))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = 0\n",
    "axes_q = ens.plot(key, xlim=(0, 0.5), label=\"orig\")\n",
    "_ = ens_q51.plot(key, axes=axes_q, label=\"q51\")\n",
    "_ = ens_q21.plot(key, axes=axes_q, label=\"q21\")\n",
    "leg_q = axes_q.figure.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Histogram representation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Convert to a histogram using 51 bins\n",
    "ens_h51 = qp.convert(ens, 'hist', bins=np.linspace(0, 3.0, 51))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Convert to a histogram using 21 bins\n",
    "ens_h21 = qp.convert(ens, 'hist', bins=np.linspace(0, 3.0, 21))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = 0\n",
    "axes_h = ens.plot(key, xlim=(0, 0.5), label=\"orig\")\n",
    "_ = ens_h51.plot(key, axes=axes_h, label=\"h51\")\n",
    "_ = ens_h21.plot(key, axes=axes_h, label=\"h21\")\n",
    "leg_h = axes_h.figure.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Other representations\n",
    "\n",
    "`qp` also includes spline-based and Gaussian mixture represenations.   The conversion to these forms much slower, so we will first reduce the base ensemble from 20000 PDFs to 100 PDFs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ens_red = ens[np.arange(100)]\n",
    "print(\"Reduced ensemble has %i PDFs\" % (ens_red.npdf))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Spline representation\n",
    "\n",
    "We can convert to the spline representation a few different ways.  This particular method specifies that we should evaluate each PDF at a grid of points and then use those to construct the spline represenation.  We do this for 2 different grids.  Note how much slower this conversion function is that the ones above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Convert to a histogram using 51 grid points\n",
    "ens_s51 = qp.convert(ens_red, 'spline', xvals=np.linspace(0, 3.0, 51), method=\"xy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ens_s21 = qp.convert(ens_red, 'spline', xvals=np.linspace(0, 3.0, 21), method=\"xy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = 0\n",
    "axes_s = ens_red.plot(key, xlim=(0, 0.5), label=\"orig\")\n",
    "_ = ens_s51.plot(key, axes=axes_s, label=\"s51\")\n",
    "_ = ens_s21.plot(key, axes=axes_s, label=\"s21\")\n",
    "leg_s = axes_s.figure.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Gaussian mixture model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Convert to a gaussian mixture using 301 sample points and 3 components\n",
    "ens_m3 = qp.convert(ens_red, 'mixmod', xvals=np.linspace(0, 3.0, 301), ncomps=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Convert to a gaussian mixture using 301 sample points and 3 components\n",
    "ens_m5 = qp.convert(ens_red, 'mixmod', xvals=np.linspace(0, 3.0, 301), ncomps=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = 0\n",
    "axes_m = ens_red.plot(key, xlim=(0, 0.5), label=\"orig\")\n",
    "_ = ens_m3.plot(key, axes=axes_m, label=\"m3\")\n",
    "_ = ens_m5.plot(key, axes=axes_m, label=\"m5\")\n",
    "leg_m = axes_m.figure.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = 1\n",
    "axes_m1 = ens_red.plot(key, xlim=(0, 0.5), label=\"orig\")\n",
    "_ = ens_m3.plot(key, axes=axes_m1, label=\"m3\")\n",
    "_ = ens_m5.plot(key, axes=axes_m1, label=\"m5\")\n",
    "leg_m1 = axes_m1.figure.legend()"
   ]
  },
  {
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
