{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `qp` Demo\n",
    "\n",
    "_Alex Malz, Phil Marshall, Eric Charles_\n",
    "\n",
    "In this notebook we use the `qp` module to approximate some simple, standard, 1D PDFs using sets of quantiles, samples, and histograms, and assess their relative accuracy. \n",
    "We also show how such analyses can be extended to use \"composite\" PDFs made up of mixtures of standard distributions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import scipy.stats as sps\n",
    "\n",
    "import matplotlib\n",
    "matplotlib.use('Agg')\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Requirements\n",
    "\n",
    "To run `qp`, you will need to first install the module by following the instructions [here](https://github.com/LSSTDESC/qp/blob/main/docs/install.rst)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import qp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Background: the `scipy.stats` module\n",
    "\n",
    "The `scipy.stats` module is the standard for manipulating distribtions so is a natural place to start for implementing 1D PDF parameterizations.  \n",
    "It allows you do define a wide variety of distibutions and uses `numpy` array broadcasting for efficiency."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gaussian (Normal) example\n",
    "\n",
    "Here are some examples of things you can do with the `scipy.stats` module, using a Gaussian or Normal distribution.\n",
    "`loc` and `scale` are the means and standard deviations of the underlying Gaussians.\n",
    "\n",
    "Note the distinction between passing arguments to `norm` and passing arguments to `pdf` to access multiple distributions and their PDF values at multiple points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluate a single distribution's PDF at one value\n",
    "print(\"PDF at one point for one distribution:\", \n",
    "      sps.norm(loc=0, scale=1).pdf(0.5))\n",
    "\n",
    "# evaluate a single distribution's PDF at multiple value\n",
    "print(\"PDF at three points for one distribution:\", \n",
    "      sps.norm(loc=0, scale=1).pdf([0.5, 1., 1.5]))\n",
    "\n",
    "# evalute three distributions' PDFs at one shared value\n",
    "print(\"PDF at one point for three distributions:\", \n",
    "      sps.norm(loc=[0., 1., 2.], scale=1).pdf(0.5))\n",
    "\n",
    "# evalute three distributions' PDFs each at one different value\n",
    "print(\"PDF at one different point for three distributions:\", \n",
    "      sps.norm(loc=[0., 1., 2.], scale=1).pdf([0.5, 1., 1.5]))\n",
    "\n",
    "# evalute three distributions' PDFs each at four different values\n",
    "# (note the change in shape of the argument)\n",
    "print(\"PDF at four different points for three distributions:\\n\",\n",
    "      sps.norm(loc=[0., 1., 2.], scale=1).pdf([[0.5],[1.],[1.5],[2]]))\n",
    "\n",
    "# evalute three distributions' PDFs at each of four different values\n",
    "# (note the change in shape of the argument)\n",
    "print(\"PDF at four different points for three distributions: broadcast reversed\\n\",\n",
    "      sps.norm(loc=[[0.], [1.], [2.]], scale=1).pdf([0.5,1.,1.5,2]))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The `scipy.stats` classes\n",
    "\n",
    "In the `scipy.stats` module, all of the distributions are sub-classes of `scipy.stats.rv_continuous`.  \n",
    "You make an object of a particular sub-type, and then 'freeze' it by passing it shape parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"This is the generic normal distribution class: \", \n",
    "      sps._continuous_distns.norm_gen)\n",
    "\n",
    "ng = sps._continuous_distns.norm_gen()\n",
    "print(\"This is an instance of the generic normal distribution class\", \n",
    "      ng)\n",
    "\n",
    "norm_sp = ng(loc=0, scale=1)\n",
    "print(\"This is a frozen normal distribution, with specific paramters\", \n",
    "      norm_sp, norm_sp.kwds)\n",
    "print(\"The frozen object know what generic distribution it comes from\", \n",
    "      norm_sp.dist)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Properties of distributions\n",
    "\n",
    "`scipy.stats` lets you evaluate multiple properties of distributions.  These include:\n",
    "1. pdf: Probability Density Function\n",
    "2. cdf: Cumulative Distribution Function\n",
    "3. ppf: Percent Point Function (Inverse of CDF)\n",
    "4. sf: Survival Function (1-CDF)\n",
    "5. isf: Inverse Survival Function (Inverse of SF)\n",
    "6. rvs: Random Variates (i.e., sampled values)\n",
    "7. stats: Return mean, variance, optionally: (Fisher’s) skew, or (Fisher’s) kurtosis\n",
    "8. moment: non-central moments of the distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"PDF = \", norm_sp.pdf(0.5))  \n",
    "print(\"CDF = \", norm_sp.cdf(0.5))\n",
    "print(\"PPF = \", norm_sp.ppf(0.6))\n",
    "print(\"SF  = \", norm_sp.sf(0.6))\n",
    "print(\"ISF = \", norm_sp.isf(0.5))\n",
    "print(\"RVS = \", norm_sp.rvs())\n",
    "print(\"stats = \", norm_sp.stats())\n",
    "print(\"M2  = \", norm_sp.moment(2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `qp` parameterizations and visualization functionality\n",
    "\n",
    "The next part of this notebook shows how we can extend the functionality of `scipy.stats` to implement distributions that are based on parameterizations of 1D PDFs, like histograms, interpolations, splines, or mixture models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameterizations from `scipy.stats`\n",
    "\n",
    "`qp` automatically generates classes for all of the `scipy.stats.rv_continuous` distributions, providing feed-through access to all `scipy.stats.rv_continuous` objects but adds on additional attributes and methods specific to parameterization conversions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "qp.stats.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(qp.stats.lognorm_gen)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(qp.stats.lognorm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Native plotting\n",
    "\n",
    "If you have a single distribution you can plot it, the `qp.plotting.plot_native` function will find a nice way to represent the data used to construct the distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loc1 = np.array([[0]])\n",
    "scale1 = np.array([[1]])\n",
    "norm_dist1 = qp.stats.norm(loc=loc1, scale=scale1)\n",
    "fig, axes = qp.plotting.plot_native(norm_dist1, xlim=(-5., 5.))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fig, axes = qp.stats.norm.plot_native(norm_dist1, xlim=(-5., 5.))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `qp` histogram (piecewise constant) parameterization\n",
    "\n",
    "This represents a set of distributions made by interpolating a set of histograms with shared binning.\n",
    "To construct this you need to give the bin edges (shape=(N)) and the bin values (shape=(npdf, N-1)).\n",
    "\n",
    "Note that the native visual representation is different from the Normal distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert to a histogram by computing the bin values by taking the intergral of the CDF\n",
    "xvals = np.linspace(-5, 5, 11)\n",
    "cdf = norm_dist1.cdf(xvals)\n",
    "bin_vals = cdf[:,1:] - cdf[:,0:-1]\n",
    "# Construct histogram PDF using the bin edges and the bin values\n",
    "hist_dist = qp.hist(bins=xvals, pdfs=bin_vals)\n",
    "yvals = hist_dist.pdf(xvals)\n",
    "# Construct a single PDF for plotting\n",
    "hist_dist1 = qp.hist(bins=xvals, pdfs=np.atleast_2d(bin_vals[0]))\n",
    "fig, axes = qp.plotting.plot_native(hist_dist1, xlim=(-5., 5.))\n",
    "leg = fig.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What if you want to evaluate a vector of input values, where each input value is different for each PDF?  In that case you need the shape of the vector of input value to match the implicit shape of the PDFs, which in this case is (2,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xvals_x = np.array([[-1.], [1.]])\n",
    "yvals_x = hist_dist.pdf(xvals_x)\n",
    "print (\"For an input vector of shape %s the output shape is %s\" % (xvals_x.shape, yvals_x.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `qp` quantile parameterization\n",
    "\n",
    "This represents a set of distributions made by interpolating the locations at which various distributions reach a given set of quantiles.\n",
    "To construct this you need to give the quantiles edges (shape=(N)) and the location values (shape=(npdf, N)).\n",
    "\n",
    "Note that the native visual representation is different."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the quantile values to compute the locations for\n",
    "quants = np.linspace(0.01, 0.99, 7)\n",
    "# Compute the corresponding locations\n",
    "locs = norm_dist1.ppf(quants)\n",
    "# Construct the distribution using the quantile value and locations\n",
    "quant_dist = qp.quant(quants=quants, locs=locs)\n",
    "quant_vals = quant_dist.pdf(xvals)\n",
    "print(\"The input and output shapes are:\", xvals.shape, quant_vals.shape)\n",
    "# Construct a single PDF for plotting\n",
    "quant_dist1 = qp.quant(quants=np.atleast_1d(quants), locs=np.atleast_2d(locs[0]))\n",
    "fig, axes = qp.plotting.plot_native(quant_dist1, xlim=(-5., 5.), label=\"quantiles\")\n",
    "leg = fig.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(quants)\n",
    "print(quant_dist.dist.quants)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `qp` interpolated parameterization\n",
    "\n",
    "This represents a set of distributions made by interpolating a set of x and y values. \n",
    "To construct this you need to give the x and y values (both of shape=(npdf, N))\n",
    "\n",
    "Note that the native visual representation is pretty similar to the original one for the Gaussian."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the x-grid locations\n",
    "xvals = np.linspace(-5, 5, 11)\n",
    "# Compute the corresponding y values\n",
    "yvals = norm_dist1.pdf(xvals)\n",
    "# Construct the PDFs using the x grid and y values\n",
    "interp_dist = qp.interp(xvals=xvals, yvals=yvals)\n",
    "interp_vals = interp_dist.pdf(xvals)\n",
    "print(\"The input and output shapes are:\", xvals.shape, interp_vals.shape)\n",
    "# Construct a single PDF for plotting\n",
    "interp_dist1 = qp.interp(xvals=xvals, yvals=np.atleast_2d(yvals[0]))\n",
    "fig, axes = qp.plotting.plot_native(interp_dist1, xlim=(-5., 5.), label=\"interpolated\")\n",
    "leg = fig.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `qp` spline parameterization constructed from kernel density estimate (samples) parameterization\n",
    "\n",
    "This represents a set of distributions made by producing a kernel density estimate from a set of samples.\n",
    "\n",
    "To construct this you need to give the samples edges (shape=(npdf, Nsamples)).\n",
    "\n",
    "Note again that the the native visual represenation is different."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Take 100 random samples from each of 2 PDFs\n",
    "samples = norm_dist1.rvs(size=(2, 1000))\n",
    "# Define points at which to evaluate the kernal density estimate (KDE)\n",
    "xvals_kde = np.linspace(-5., 5., 51)\n",
    "# Use a utility function to construct the KDE, sample it, and they construct a spline\n",
    "kde_dist = qp.spline_from_samples(xvals=xvals_kde, samples=samples)\n",
    "kde_vals = kde_dist.pdf(xvals_kde)\n",
    "print(\"The input and output shapes are:\", xvals.shape, kde_vals.shape)\n",
    "# Construct a single PDF for plotting\n",
    "kde_dist1 = qp.spline_from_samples(xvals=xvals_kde, samples=np.atleast_2d(samples[0]))\n",
    "fig, axes = qp.plotting.plot_native(kde_dist1, xlim=(-5., 5.), label=\"kde\")\n",
    "leg = fig.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `qp` spline parameterization\n",
    "\n",
    "This represents a set of distributions made building a set of splines. Though the parameterization is defined by the spline knots, you can construct this from x and y values (both of shape=(npdf, N)).\n",
    "\n",
    "Note that the native visual representation is pretty similar to the original one for the Gaussian.\n",
    "\n",
    "Note also that the spline knots are stored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To make a spline you need the spline knots, you can get those from the xval, yval values\n",
    "splx, sply, spln = qp.spline_gen.build_normed_splines(np.expand_dims(xvals,0), yvals)\n",
    "spline_dist_orig = qp.spline(splx=splx, sply=sply, spln=spln)\n",
    "# Or we can do these two steps together using one function\n",
    "spline_dist = qp.spline_from_xy(xvals=np.expand_dims(xvals,0), yvals=yvals)\n",
    "spline_vals = spline_dist.pdf(xvals)\n",
    "print(\"The input and output shapes are:\", xvals.shape, spline_vals.shape)\n",
    "print(\"Spline knots\", spline_dist.dist.splx, spline_dist.dist.sply, spline_dist.dist.spln)\n",
    "# Construct a single PDF for plotting\n",
    "spline_dist1 = qp.spline_from_xy(xvals=np.atleast_2d(xvals), yvals=np.atleast_2d(yvals))\n",
    "print(spline_dist1.dist.splx.shape)\n",
    "fig, axes = qp.plotting.plot_native(spline_dist1, xlim=(-5., 5.), label=\"spline\")\n",
    "leg = fig.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overplotting\n",
    "\n",
    "You can visually compare the represenations by plotting them all on the same figure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = qp.plotting.plot_native(norm_dist1, xlim=(-5., 5.), label=\"norm\")\n",
    "qp.plotting.plot_native(hist_dist1, axes=axes)\n",
    "qp.plotting.plot_native(quant_dist1, axes=axes)\n",
    "qp.plotting.plot_native(interp_dist1, axes=axes, label=\"interp\")\n",
    "# qp.plotting.plot_native(kde_dist1, axes=axes)\n",
    "# qp.plotting.plot_native(spline_dist1, axes=axes, label=\"spline\")\n",
    "leg = fig.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The `qp.Ensemble` Class\n",
    "\n",
    "This is the basic element of `qp` - an object representing a set of probability density functions. This class is stored in the module `ensemble.py`.  \n",
    "\n",
    "To create a `qp.Ensemble` you need to specify the class used to represent the PDFs, and provide that data for the specific set of PDFs.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ensembles of distributions\n",
    "\n",
    "`qp` no longer distinguishes between distributions and ensembles thereof -- a single distribution is just a special case of an ensemble with only one member, which takes advantage of computational efficiencies in `scipy`.\n",
    "The shape of the array returned by a call to the pdf function of a distribution depends on the shape of the parameters and evaluate points.  \n",
    "\n",
    "For distributions that take multiple input arrays, `qp` uses te convention that the rows are the individual distributions and the columns are the values of the parameters defining the distributions under a known parameterization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is a trivial extension, with the number of pdfs as a member of the `scipy.stats.norm_gen` distribution.\n",
    "loc = np.array([[0],[1]])\n",
    "scale = np.array([[1],[1]])\n",
    "norm_dist = qp.stats.norm(loc=loc, scale=scale)\n",
    "xvals = np.linspace(-5, 5, 51)\n",
    "yvals = norm_dist.pdf(xvals)\n",
    "print(\"This object represents %i pdfs\" % norm_dist.npdf)\n",
    "print(\"The input and output shapes are:\", xvals.shape, yvals.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print (\"For an input vector of shape %s the output shape is %s\" % (xvals.shape, yvals.shape))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In this case we return an array were the rows are the evaluation points and the columns the different PDFs\n",
    "vector_pdf = qp.stats.norm(loc=[0., 1., 2], scale=1.)\n",
    "vector_pdf.pdf([[0.], [0.5]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is the same, except we use `numpy.expand_dims` to shape the input array of evaluation points\n",
    "vector_pdf = qp.stats.norm(loc=[0., 1., 2], scale=1.)\n",
    "vector_pdf.pdf(np.expand_dims(np.array([0., 0.5]), -1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In this case we return an array were the rows are pdfs and the columns the evaluation points\n",
    "vector_pdf = qp.stats.norm(loc=[[0.], [1.], [2]], scale=1.)\n",
    "vector_pdf.pdf([0., 0.5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is the same, except we use `numpy.expand_dims` to shape the input array of pdf parameters\n",
    "vector_pdf = qp.stats.norm(loc=np.expand_dims([0., 1., 2], -1), scale=1.)\n",
    "vector_pdf.pdf([0., 0.5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we will create 100 Gaussians with means distributed between -1 and 1, and widths distributed between 0.9 and 1.1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "locs = 2* (np.random.uniform(size=(100,1))-0.5)\n",
    "scales = 1 + 0.2*(np.random.uniform(size=(100,1))-0.5)\n",
    "ens_n = qp.Ensemble(qp.stats.norm, data=dict(loc=locs, scale=scales))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using the ensemble\n",
    "\n",
    "All of the methods of the distributions (`pdf`, `cdf` etc.) work the same way for an ensemble as for underlying classes. \n",
    "\n",
    "To isolate a single distribution in the ensemble, use the square brackets operator `[]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vals_n = ens_n.pdf(xvals)\n",
    "print(\"The shapes are: \", xvals.shape, vals_n.shape)\n",
    "fig, axes = qp.plotting.plot_native(ens_n[15], xlim=(-5.,5.))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Converting the ensemble\n",
    "\n",
    "The `qp.Ensemble.convert_to` function lets you convert ensembles to other representations. To do this you have to provide the original ensemble, the class you want to convert to, and any some keyword arguments to specify details about how to convert to the new class, here are some examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-5, 5, 11)\n",
    "quants = np.linspace(0.01, 0.99, 7)\n",
    "print(\"Making hist\")\n",
    "ens_h = ens_n.convert_to(qp.hist_gen, bins=bins)\n",
    "print(\"Making interp\")\n",
    "ens_i = ens_n.convert_to(qp.interp_gen, xvals=bins)\n",
    "print(\"Making spline\")\n",
    "ens_s = ens_n.convert_to(qp.spline_gen, xvals=bins, method=\"xy\")\n",
    "#print(\"Making spline from samples\")\n",
    "#ens_s = ens_n.convert_to(qp.spline_gen, xvals=bins, samples=1000, method=\"samples\")\n",
    "print(\"Making quants\")\n",
    "ens_q = ens_n.convert_to(qp.quant_gen, quants=quants)\n",
    "print(\"Making mixmod\")\n",
    "ens_m = ens_n.convert_to(qp.mixmod_gen, samples=1000, ncomps=3)\n",
    "#print(\"Making flexcode\")\n",
    "#ens_f = ens_n.convert_to(qp.flex_gen, grid=bins, basis_system='cosine')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `qp.convert` function also works the more or less the same way, but with slightly different syntax, where you can use the name of the class instead of the class object.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Making hist\")\n",
    "ens_h2 = qp.convert(ens_n, \"hist\", bins=bins)\n",
    "print(\"Making interp\")\n",
    "ens_i2 = qp.convert(ens_n, \"interp\", xvals=bins)\n",
    "print(\"Making spline\")\n",
    "ens_s2 = qp.convert(ens_n, \"spline\", xvals=bins, method=\"xy\")\n",
    "print(\"Making quants\")\n",
    "ens_q2 = qp.convert(ens_n, \"quant\", quants=quants)\n",
    "print(\"Making mixmod\")\n",
    "ens_m2 = qp.convert(ens_n, \"mixmod\", samples=1000, ncomps=3)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Comparing Parametrizations\n",
    "\n",
    "`qp` supports quantitative comparisons between different distributions, across parametrizations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Qualitative Comparisons: Plotting\n",
    "\n",
    "Let's visualize the PDF object in order to original and the other representaions.  The solid, black line shows the true PDF evaluated between the bounds.  The green rugplot shows the locations of the 1000 samples we took.  The vertical, dotted, blue lines show the percentiles we asked for, and the hotizontal, dotted, red lines show the 10 equally spaced bins we asked for.  Note that the quantiles refer to the probability distribution *between the bounds*, because we are not able to integrate numerically over an infinite range. Interpolations of each parametrization are given as dashed lines in their corresponding colors.  Note that the interpolations of the quantile and histogram parametrizations are so close to each other that the difference is almost imperceptible!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = qp.plotting.plot_native(ens_n[15], xlim=(-5.,5.))\n",
    "qp.plotting.plot_native(ens_h[15], axes=axes)\n",
    "qp.plotting.plot_native(ens_q[15], axes=axes, label='quantile')\n",
    "qp.plotting.plot_native(ens_i[15], axes=axes, label='interp')\n",
    "# qp.plotting.plot_native(ens_s[15], axes=axes, label='spline')\n",
    "qp.plotting.plot_native(ens_m[15], axes=axes, label='mixmod')\n",
    "#qp.qp_plot_native(ens_f[15], axes=axes, label='flex')\n",
    "leg = fig.legend()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also interpolate the function onto an evenly spaced grid point and cache those values with the `gridded` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "grid = np.linspace(-3., 3., 100)\n",
    "gridded = ens_n.pdf(grid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cached_gridded = ens_n.gridded(grid)[1]\n",
    "check = gridded - cached_gridded\n",
    "print(check.min(), check.max())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quantitative Comparisons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "symm_lims = np.array([-1., 1.])\n",
    "all_lims = [symm_lims, 2.*symm_lims, 3.*symm_lims]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "Next, let's compare the different parametrizations to the truth using the Kullback-Leibler Divergence (KLD).  The KLD is a measure of how close two probability distributions are to one another -- a smaller value indicates closer agreement.  It is measured in units of bits of information, the information lost in going from the second distribution to the first distribution.  The KLD calculator here takes in a shared grid upon which to evaluate the true distribution and the interpolated approximation of that distribution and returns the KLD of the approximation relative to the truth, which is not in general the same as the KLD of the truth relative to the approximation.  Below, we'll calculate the KLD of the approximation relative to the truth over different ranges, showing that it increases as it includes areas where the true distribution and interpolated distributions diverge."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for a single pair of pdfs. (the 15th in each ensemble)\n",
    "klds = qp.metrics.calculate_kld(ens_n, ens_s, limits=symm_lims)\n",
    "print(klds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loop over all the other ensemble types\n",
    "ensembles = [ens_n, ens_h, ens_i, ens_s, ens_q, ens_m]\n",
    "for ensemble in ensembles[1:]:\n",
    "    D = []\n",
    "    for lims in all_lims:\n",
    "        klds = qp.metrics.calculate_kld(ens_n, ensemble, limits=lims)\n",
    "        D.append(\"%.2e +- %.2e\" % (klds.mean(), klds.std()))\n",
    "    print(ensemble.gen_class.name + ' approximation: KLD over 1, 2, 3, sigma ranges = ' + str(D))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The progression of KLD values should follow that of the root mean square error (RMSE), another measure of how close two functions are to one another.  The RMSE also increases as it includes areas where the true distribution and interpolated distribution diverge.  Unlike the KLD, the RMSE is symmetric, meaning the distance measured is not that of one distribution from the other but of the symmetric distance between them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for ensemble in ensembles[1:]:\n",
    "    D = []\n",
    "    for lims in all_lims:\n",
    "        rmses = qp.metrics.calculate_rmse(ens_n, ensemble, limits=lims)\n",
    "        D.append(\"%.2e +- %.2e\" % (rmses.mean(), rmses.std()))\n",
    "    print(ensemble.gen_class.name + ' approximation: RMSE over 1, 2, 3, sigma ranges = ' + str(D))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Both the KLD and RMSE metrics suggest that the quantile approximation is better in the high density region, but samples work better when the tails are included. We might expect the answer to the question of which approximation to use to depend on the application, and whether the tails need to be captured or not."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Storing and retreiving ensembles\n",
    "\n",
    "You can store and retrieve ensembles from disk using the `qp.Ensemble.write_to` and `qp.read` methods.\n",
    "\n",
    "These work in two steps, first they convert the Ensemble data to `astropy.table` objects, and then they write the tables.  This means you can store the data in any format support by `astropy`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tabs = ens_n.build_tables()\n",
    "print(tabs.keys())\n",
    "print()\n",
    "print(\"Meta Data\")\n",
    "print(tabs['meta'])\n",
    "print()\n",
    "print(\"Object Data\")\n",
    "print(tabs['data'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a loopback test showing that we get the same results before and after a write/read cycle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "suffix_list = ['_n', '_h', '_i', '_s', '_q', '_m']\n",
    "filetypes = ['fits', 'hf5']\n",
    "for ens, suffix in zip(ensembles, suffix_list):\n",
    "    for ft in filetypes:\n",
    "\n",
    "        outfile = \"test%s.%s\" % (suffix, ft)\n",
    "        metafile = \"test%s_meta.%s\" % (suffix, ft)\n",
    "        \n",
    "        pdf_1 = ens.pdf(bins)        \n",
    "        ens.write_to(outfile)\n",
    "        ens_r = qp.read(outfile)\n",
    "        pdf_2 = ens_r.pdf(bins)\n",
    "\n",
    "        check = pdf_1 - pdf_2\n",
    "        print(suffix, ft, check.min(), check.max())\n",
    "\n",
    "        os.unlink(outfile)\n",
    "        try:\n",
    "            os.unlink(metafile)\n",
    "        except Exception:\n",
    "            pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can compare the moments of each approximation and compare those to the moments of the true distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "which_moments = range(3)\n",
    "all_moments = []\n",
    "for ens in ensembles:\n",
    "    moments = []\n",
    "    for n in which_moments:\n",
    "        moms = qp.metrics.calculate_moment(ens, n, limits=(-3, 3))\n",
    "        moments.append(\"%.2e +- %.2e\" % (moms.mean(), moms.std()))\n",
    "    all_moments.append(moments)\n",
    "    \n",
    "print('moments: '+str(which_moments))\n",
    "for ens, mom in zip(ensembles, all_moments):\n",
    "    print(ens.gen_class.name+': '+str(mom))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "qp_issue_25",
   "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.10.8"
  },
  "vscode": {
   "interpreter": {
    "hash": "9be1a7334e581107f2753dc5e5ebb12b8975e5f80325d35375311bd0b6b51792"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
