{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook is meant to demonstrate the use of quantile parameterization. Basic usage is covered in the first section. In the second half two different classes of distribution are used to showcase how the various PDF reconstruction algorithms perform."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports and global defaults\n",
    "import qp\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Sets the x axis limits on all plots, rerun notebook to apply changes.\n",
    "STD_X_LIM = (0,5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an Ensemble with a single normal distribution\n",
    "single_norm = qp.Ensemble(qp.stats.norm, data=dict(loc=3, scale=0.5))\n",
    "\n",
    "# Starting with a simple base case. Reasonable quantiles >0, <1.0\n",
    "single_norm_quantiles = np.linspace(0.001, 0.999, 16)\n",
    "single_norm_locations = single_norm.ppf(single_norm_quantiles)\n",
    "\n",
    "# Create a quantile parameterization ensemble\n",
    "quant_dist_single_norm = qp.quant(quants=single_norm_quantiles, locs=single_norm_locations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print out the constructor name\n",
    "print(quant_dist_single_norm.dist.pdf_constructor_name)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To specify a pdf constructor algorithm when creating an Ensemble, pass the name as the `pdf_constructor_name` argument.\n",
    "The available algorithms are defined in the `quant_pdf.py:PDF_CONSTRUCTORS` dictionary.\n",
    "\n",
    "If the `pdf_constructor_name` argument unknown or not provided, the default `piecewise_linear` will be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "quant_dist_single_norm = qp.Ensemble(\n",
    "    qp.quant,\n",
    "    data=dict(\n",
    "        quants=single_norm_quantiles,\n",
    "        locs=single_norm_locations,\n",
    "        check_input=False,\n",
    "        pdf_constructor_name='dual_spline_average'  # <-- This is the magic line\n",
    "    )\n",
    ")\n",
    "\n",
    "# You can subsequently verify the reconstruction algorithm being used.\n",
    "print(quant_dist_single_norm.dist.pdf_constructor_name)\n",
    "\n",
    "# The pdf reconstruction algorithm can be redefined after the Ensemble has been created like so:\n",
    "quant_dist_single_norm.dist.pdf_constructor_name='cdf_spline_derivative'\n",
    "print(quant_dist_single_norm.dist.pdf_constructor_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an Ensemble with 3 normal distributions\n",
    "means = np.array([[1], [2.5], [3]])\n",
    "scales = np.array([[0.25], [0.5], [0.1]])\n",
    "many_norm = qp.Ensemble(qp.stats.norm, data=dict(loc=means, scale=scales))\n",
    "\n",
    "# Starting with a simple base case. Reasonable quantiles >0, <1.0\n",
    "many_norm_quantiles = np.linspace(0.001, 0.999, 16)\n",
    "many_norm_locations = many_norm.ppf(many_norm_quantiles)\n",
    "\n",
    "# Create a quantile parameterization ensemble\n",
    "quant_dist_many_norm = qp.quant(quants=many_norm_quantiles, locs=many_norm_locations)\n",
    "\n",
    "# The following approach is identical, and is basically syntactical sugar.\n",
    "quant_dist_many_norm = qp.convert(many_norm, 'quant', quants=many_norm_quantiles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a grid and call `pdf` to return the y values to reconstruct the initial PDF\n",
    "user_provided_grid = np.linspace(0,4,100)\n",
    "results = quant_dist_many_norm.pdf(user_provided_grid)\n",
    "print(results.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_provided_grid = np.linspace(0,4,100)\n",
    "\n",
    "# Use the following syntax to retrieve a specific distribution's PDF from the Ensemble\n",
    "results = quant_dist_many_norm[1].pdf(user_provided_grid)\n",
    "print(results.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show that we can use different row values, and get different values out.\n",
    "user_provided_grid = np.linspace(0,4,100)\n",
    "user_provided_row = [0]\n",
    "result_1 = quant_dist_many_norm.dist.pdf_constructor.construct_pdf(grid=user_provided_grid, row=user_provided_row)\n",
    "\n",
    "user_provided_row = [1]\n",
    "result_2 = quant_dist_many_norm.dist.pdf_constructor.construct_pdf(grid=user_provided_grid, row=user_provided_row)\n",
    "\n",
    "# Expect that this will be non-zero\n",
    "print(np.sum(result_1 - result_2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show that you can pass in multiple row index values\n",
    "user_provided_grid = np.linspace(0,4,100)\n",
    "user_provided_rows = [[0,1]]\n",
    "results = quant_dist_many_norm.dist.pdf_constructor.construct_pdf(grid=user_provided_grid, row=user_provided_rows)\n",
    "\n",
    "# Expect this to be (2,100), but the `.ravel()` function in the piecewise_*\n",
    "# constructors convert the results of interpolation into an array with shape = (2*100,)...\n",
    "print(results.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show basic plotting\n",
    "_, ax1 = qp.plotting.plot_native(quant_dist_single_norm, xlim=STD_X_LIM)\n",
    "qp.plotting.plot_native(single_norm, axes=ax1)\n",
    "ax1.set_title('Single distribution quantile parameterization')\n",
    "\n",
    "plot_index = 0\n",
    "_, ax2 = qp.plotting.plot_native(quant_dist_many_norm[plot_index], xlim=STD_X_LIM)\n",
    "qp.plotting.plot_native(many_norm[plot_index], axes=ax2)\n",
    "ax2.set_title('Single selected distribution from Ensemble of quantile parameterization')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Examination of different reconstruction algorithms for various types of distributions"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simple Gaussian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Same as the single_norm Ensemble created at the beginning of the notebook.\n",
    "single_norm = qp.Ensemble(qp.stats.norm, data=dict(loc=3, scale=0.5))\n",
    "fig, ax = qp.plotting.plot_native(single_norm, xlim=STD_X_LIM)\n",
    "ax.set_title('Original simple Gaussian distribution')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the quantiles, extract the locations from the original distribution\n",
    "single_norm_quantiles = np.linspace(0.001, 0.999, 16)\n",
    "single_norm_locations = single_norm.ppf(single_norm_quantiles)\n",
    "\n",
    "# Create a quantile parameterization ensemble\n",
    "quant_dist_single_norm = qp.quant(quants=single_norm_quantiles, locs=single_norm_locations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(2,2)\n",
    "plt.setp(ax, xlim=STD_X_LIM)\n",
    "\n",
    "quant_dist_single_norm.dist.pdf_constructor_name = 'piecewise_linear' # Will emit a logger warning that this constructor is already selected\n",
    "qp.plotting.plot_native(quant_dist_single_norm, axes=ax[0,0])\n",
    "qp.plotting.plot_native(single_norm, axes=ax[0,0])\n",
    "\n",
    "\n",
    "quant_dist_single_norm.dist.pdf_constructor_name = 'piecewise_constant'\n",
    "qp.plotting.plot_native(quant_dist_single_norm, axes=ax[0,1])\n",
    "qp.plotting.plot_native(single_norm, axes=ax[0,1])\n",
    "\n",
    "\n",
    "quant_dist_single_norm.dist.pdf_constructor_name = 'cdf_spline_derivative'\n",
    "qp.plotting.plot_native(quant_dist_single_norm, axes=ax[1,0])\n",
    "qp.plotting.plot_native(single_norm, axes=ax[1,0])\n",
    "\n",
    "\n",
    "quant_dist_single_norm.dist.pdf_constructor_name = 'dual_spline_average'\n",
    "qp.plotting.plot_native(quant_dist_single_norm, axes=ax[1,1])\n",
    "qp.plotting.plot_native(single_norm, axes=ax[1,1])\n",
    "\n",
    "ax[0,0].set_title('piecewise_linear')\n",
    "ax[0,1].set_title('piecewise_constant')\n",
    "ax[1,0].set_title('cdf_spline_derivative')\n",
    "ax[1,1].set_title('dual_spline_average')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simple Rayleigh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Same as the single_norm Ensemble created at the beginning of the notebook.\n",
    "single_rayleigh = qp.Ensemble(qp.stats.rayleigh, data=dict(loc=0, scale=1))\n",
    "fig, ax = qp.plotting.plot_native(single_rayleigh, xlim=STD_X_LIM)\n",
    "ax.set_title('Original simple Rayleigh distribution')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the quantiles, extract the locations from the original distribution\n",
    "single_rayleigh_quantiles = np.linspace(0.001, 0.999, 16)\n",
    "single_rayleigh_locations = single_rayleigh.ppf(single_rayleigh_quantiles)\n",
    "\n",
    "# Create a quantile parameterization ensemble\n",
    "quant_dist_single_rayleigh = qp.quant(quants=single_rayleigh_quantiles, locs=single_rayleigh_locations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(2,2)\n",
    "plt.setp(ax, xlim=STD_X_LIM)\n",
    "\n",
    "quant_dist_single_rayleigh.dist.pdf_constructor_name = 'piecewise_linear' # Will emit a logger warning that this constructor is already selected\n",
    "qp.plotting.plot_native(quant_dist_single_rayleigh, axes=ax[0,0])\n",
    "qp.plotting.plot_native(single_rayleigh, axes=ax[0,0])\n",
    "\n",
    "\n",
    "quant_dist_single_rayleigh.dist.pdf_constructor_name = 'piecewise_constant'\n",
    "qp.plotting.plot_native(quant_dist_single_rayleigh, axes=ax[0,1])\n",
    "qp.plotting.plot_native(single_rayleigh, axes=ax[0,1])\n",
    "\n",
    "\n",
    "quant_dist_single_rayleigh.dist.pdf_constructor_name = 'cdf_spline_derivative'\n",
    "qp.plotting.plot_native(quant_dist_single_rayleigh, axes=ax[1,0])\n",
    "qp.plotting.plot_native(single_rayleigh, axes=ax[1,0])\n",
    "\n",
    "\n",
    "quant_dist_single_rayleigh.dist.pdf_constructor_name = 'dual_spline_average'\n",
    "qp.plotting.plot_native(quant_dist_single_rayleigh, axes=ax[1,1])\n",
    "qp.plotting.plot_native(single_rayleigh, axes=ax[1,1])\n",
    "\n",
    "ax[0,0].set_title('piecewise_linear')\n",
    "ax[0,1].set_title('piecewise_constant')\n",
    "ax[1,0].set_title('cdf_spline_derivative')\n",
    "ax[1,1].set_title('dual_spline_average')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Merged Gaussian distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a multi-modal distribution where two Gaussian distributions are merged\n",
    "mu =  np.array([1.7, 2.9])\n",
    "sig = np.array([0.3, 0.4])\n",
    "wt =  np.array([0.4, 0.6])\n",
    "\n",
    "merged_norms = qp.Ensemble(qp.mixmod, data=dict(means=mu, stds=sig, weights=wt))\n",
    "fig, ax = qp.plotting.plot_native(merged_norms, xlim=STD_X_LIM)\n",
    "ax.set_title('Original merged Gaussian distribution')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the quantiles, extract the locations from the original distribution\n",
    "merged_norm_quantiles = np.linspace(0.001, 0.999, 30)\n",
    "merged_norm_locations = merged_norms.ppf(merged_norm_quantiles)\n",
    "\n",
    "# Create a quantile parameterization Ensemble\n",
    "quant_dist_merged_norm = qp.quant(quants=merged_norm_quantiles, locs=merged_norm_locations)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following plots show how each of the different pdf reconstruction algorithms perform in the case of a merged distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(2,2)\n",
    "plt.setp(ax, xlim=STD_X_LIM)\n",
    "\n",
    "quant_dist_merged_norm.dist.pdf_constructor_name = 'piecewise_linear' # Will emit a logger warning that this constructor is already selected\n",
    "qp.plotting.plot_native(quant_dist_merged_norm, axes=ax[0,0])\n",
    "qp.plotting.plot_native(merged_norms, axes=ax[0,0])\n",
    "\n",
    "\n",
    "quant_dist_merged_norm.dist.pdf_constructor_name = 'piecewise_constant'\n",
    "qp.plotting.plot_native(quant_dist_merged_norm, axes=ax[0,1])\n",
    "qp.plotting.plot_native(merged_norms, axes=ax[0,1])\n",
    "\n",
    "\n",
    "quant_dist_merged_norm.dist.pdf_constructor_name = 'cdf_spline_derivative'\n",
    "qp.plotting.plot_native(quant_dist_merged_norm, axes=ax[1,0])\n",
    "qp.plotting.plot_native(merged_norms, axes=ax[1,0])\n",
    "\n",
    "\n",
    "quant_dist_merged_norm.dist.pdf_constructor_name = 'dual_spline_average'\n",
    "qp.plotting.plot_native(quant_dist_merged_norm, axes=ax[1,1])\n",
    "qp.plotting.plot_native(merged_norms, axes=ax[1,1])\n",
    "\n",
    "ax[0,0].set_title('piecewise_linear')\n",
    "ax[0,1].set_title('piecewise_constant')\n",
    "ax[1,0].set_title('cdf_spline_derivative')\n",
    "ax[1,1].set_title('dual_spline_average')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Distinct Gaussian distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a multi-modal distribution with two distinct Gaussian distributions\n",
    "mu =  np.array([1, 3])\n",
    "sig = np.array([0.1, 0.1])\n",
    "wt =  np.array([0.4, 0.6])\n",
    "\n",
    "distinct_norms = qp.Ensemble(qp.mixmod, data=dict(means=mu, stds=sig, weights=wt))\n",
    "fig, ax = qp.plotting.plot_native(distinct_norms, xlim=STD_X_LIM)\n",
    "ax.set_title('Original distinct Gaussian distribution')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the quantiles, extract the locations from the original distribution\n",
    "distinct_norm_quantiles = np.linspace(0.001, 0.999, 30)\n",
    "distinct_norm_locations = distinct_norms.ppf(distinct_norm_quantiles)\n",
    "\n",
    "# Create a quantile parameterization Ensemble\n",
    "quant_dist_distinct_norm = qp.quant(quants=distinct_norm_quantiles, locs=distinct_norm_locations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(2,2)\n",
    "plt.setp(ax, xlim=STD_X_LIM)\n",
    "\n",
    "quant_dist_distinct_norm.dist.pdf_constructor_name = 'piecewise_linear' # Will emit a logger warning that this constructor is already selected\n",
    "qp.plotting.plot_native(quant_dist_distinct_norm, axes=ax[0,0])\n",
    "qp.plotting.plot_native(distinct_norms, axes=ax[0,0])\n",
    "\n",
    "quant_dist_distinct_norm.dist.pdf_constructor_name = 'piecewise_constant'\n",
    "qp.plotting.plot_native(quant_dist_distinct_norm, axes=ax[0,1])\n",
    "qp.plotting.plot_native(distinct_norms, axes=ax[0,1])\n",
    "\n",
    "quant_dist_distinct_norm.dist.pdf_constructor_name = 'cdf_spline_derivative'\n",
    "qp.plotting.plot_native(quant_dist_distinct_norm, axes=ax[1,0])\n",
    "qp.plotting.plot_native(distinct_norms, axes=ax[1,0])\n",
    "\n",
    "quant_dist_distinct_norm.dist.pdf_constructor_name = 'dual_spline_average'\n",
    "qp.plotting.plot_native(quant_dist_distinct_norm, axes=ax[1,1])\n",
    "qp.plotting.plot_native(distinct_norms, axes=ax[1,1])\n",
    "\n",
    "ax[0,0].set_title('piecewise_linear')\n",
    "ax[0,1].set_title('piecewise_constant')\n",
    "ax[1,0].set_title('cdf_spline_derivative')\n",
    "ax[1,1].set_title('dual_spline_average')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PDF Reconstruction algorithm `debug` method\n",
    "Each PDF reconstruction class has a `debug` method which will generally return a tuple of the primary and intermediate calculations used to reconstruct the original PDF. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a multi-modal distribution where two Gaussian distributions are merged - same as the section above, \"Merged Gaussian distributions\".\n",
    "mu =  np.array([1.7, 2.9])\n",
    "sig = np.array([0.3, 0.4])\n",
    "wt =  np.array([0.4, 0.6])\n",
    "\n",
    "merged_norms = qp.Ensemble(qp.mixmod, data=dict(means=mu, stds=sig, weights=wt))\n",
    "fig, ax = qp.plotting.plot_native(merged_norms, xlim=STD_X_LIM)\n",
    "ax.set_title('Original merged Gaussian distribution')\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following diagram, we've used the `debug` method to retrieve the intermediate calculations from the dual spline average PDF reconstructor. \n",
    "\n",
    "The plot shows the original PDF in blue, and the reconstructed points as orange stars. Shown as a green dotted line, the trapezoid reconstruction of the original PDF uses an iterative approach to calculate the unknown corner of the trapezoid approximation of the curve.\n",
    "\n",
    "Noting that the trapezoid reconstruction oscillates around the true PDF at every point leads to fitting two splines (the purple and red dashed lines) on the even and odd points of the trapezoid reconstruction. The final output is the average of the two splines at a given x location."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the quantiles, extract the locations from the original distribution\n",
    "merged_norm_quantiles = np.linspace(0.001, 0.999, 30)\n",
    "merged_norm_locations = merged_norms.ppf(merged_norm_quantiles)\n",
    "\n",
    "# Create a quantile parameterization Ensemble\n",
    "quant_dist_merged_norm = qp.quant(quants=merged_norm_quantiles, locs=merged_norm_locations)\n",
    "\n",
    "# Set the PDF reconstruction algorithm to use\n",
    "quant_dist_merged_norm.dist.pdf_constructor_name = 'dual_spline_average'\n",
    "\n",
    "# Retrieve the reconstructed PDF manually by passing in a \"grid\".\n",
    "x_values = np.linspace(0,5,100)\n",
    "original_y_values = merged_norms.pdf(x_values)\n",
    "y_values = quant_dist_merged_norm.pdf(x_values)\n",
    "\n",
    "# Retrieve the primary data and intermediate calculations for the algorithm.\n",
    "# Here `quantiles`` and `locations`` are the primary data. `p_of_zs`, `y1`, and `y2` are intermediate results.\n",
    "# `p_of_zs` represents the intermediate stepwise calculation of area under the PDF based on the the difference between adjacent quantiles\n",
    "# `y1` and `y2` are splines fit to the odd and even `p_of_zs` values (respectively)\n",
    "quantiles, locations, p_of_zs, y1, y2 = quant_dist_merged_norm.dist.pdf_constructor.debug()\n",
    "\n",
    "plt.plot(x_values, original_y_values, linewidth=5, label='Original PDF')\n",
    "plt.plot(x_values, y_values, marker='*', linestyle='none', label='Reconstructed PDF')\n",
    "plt.plot(np.squeeze(locations), np.squeeze(p_of_zs), linestyle=':', label='Intermediate trapezoid')\n",
    "plt.plot(x_values, np.squeeze(y1), linestyle='--', label='Lower spline fit')\n",
    "plt.plot(x_values, np.squeeze(y2), linestyle='--', label='Upper spline fit')\n",
    "plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))\n",
    "plt.xlim(0,5)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reference material\n",
    "The following code was extracted from the original implementation of `qp` [github link](https://github.com/aimalz/qp). It was originally written using Python 2 and was updated to run on Python 3, and is kept here as a reference.\n",
    "\n",
    "Note: The PDF constructor `cdf_spline_derivative` is a reproduction of this original code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#The following methods were extracted from https://github.com/aimalz/qp/blob/master/qp/pdf.py and modified only so that they would run within this notebook. The math has not changed.\n",
    "from scipy.interpolate import interp1d, InterpolatedUnivariateSpline\n",
    "import sys\n",
    "\n",
    "epsilon = sys.float_info.epsilon\n",
    "infty = sys.float_info.max * epsilon\n",
    "\n",
    "def sandwich(in_arr, ends):\n",
    "    \"\"\"\n",
    "    Adds given values to the ends of a 1D array\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    in_arr: numpy.ndarray, float\n",
    "        original array\n",
    "    ends: numpy.ndarray or tuple or list, float or numpy.ndarray, float\n",
    "        values to be added to the beginning and end\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    out_arr: numpy.ndarray, float\n",
    "        array with front and back concatenations\n",
    "    \"\"\"\n",
    "    if type(ends[0]) == np.ndarray:\n",
    "        prepend = len(ends[0])\n",
    "    else:\n",
    "        prepend = 1\n",
    "    if type(ends[-1]) == np.ndarray:\n",
    "        append = -1 * len(ends[-1])\n",
    "    else:\n",
    "        append = -1\n",
    "    out_arr = np.zeros(prepend + len(in_arr) - append)\n",
    "    out_arr[:prepend] = ends[0]\n",
    "    out_arr[prepend:append] = in_arr\n",
    "    out_arr[append:] = ends[-1]\n",
    "    return out_arr\n",
    "\n",
    "def evaluate_histogram(in_data, threshold=epsilon):\n",
    "    \"\"\"\n",
    "    Produces PDF values given samples\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    in_data: None or tuple, numpy.ndarray, float\n",
    "        tuple of (n+1) bin endpoints x and (n) CDF y between endpoints\n",
    "    threshold: float, optional\n",
    "\n",
    "    vb: boolean, optional\n",
    "        be careful and print progress to stdout?\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    out_data: tuple, float\n",
    "        sorted samples x and corresponding PDF values y\n",
    "    \"\"\"\n",
    "    # x = locs (or in our case redshift values)\n",
    "    # y = first derivatives of delta_quant/delta_loc or in our case p(z)\n",
    "    (x, y) = in_data\n",
    "    dx = threshold\n",
    "    xs = np.zeros(2 * len(y))\n",
    "    ys = xs\n",
    "    # ! xs defines the bin edges.\n",
    "    # ! This is creating the bin edges. xs[0] and xs[1] are the edges of the first bin. \n",
    "    # ! xs[2] = xs[1]+epsilon becomes the beginning of the second bin, etc.\n",
    "    # ! Then we \"repeat\" the y values so that you end up with histogram steps.\n",
    "    xs[::2] = x[:-1] + dx\n",
    "    xs[1::2] = x[1:] - dx\n",
    "    ys = np.repeat(y, 2)\n",
    "    xs = sandwich(xs, (x[0] - dx, x[-1] + dx))\n",
    "    ys = sandwich(ys, (threshold, threshold))\n",
    "    out_data = (xs, ys)\n",
    "    return out_data\n",
    "\n",
    "def evaluate_quantiles(in_data, threshold=epsilon):\n",
    "    \"\"\"\n",
    "    Estimates PDF values given quantile information\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    in_data: tuple, numpy.ndarray, float\n",
    "        tuple of CDF values iy and values x at which those CDFs are achieved\n",
    "    threshold: float, optional\n",
    "        optional minimum threshold for CDF difference\n",
    "    vb: boolean, optional\n",
    "        be careful and print progress to stdout?\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    out_data: tuple, numpy.ndarray, float\n",
    "        values xs and corresponding PDF values ys\n",
    "    \"\"\"\n",
    "\n",
    "    # iy = quants\n",
    "    # x = locs\n",
    "    (iy, x) = in_data\n",
    "\n",
    "    # This is the same as np.diff(x)\n",
    "    dx = x[1:] - x[:-1]\n",
    "    \n",
    "    # This is the same as np.diff(iy)\n",
    "    diy = iy[1:] - iy[:-1]\n",
    "\n",
    "    # this is the numerical first derivative i.e. p(z)\n",
    "    y = diy / dx\n",
    "\n",
    "    # evaluate_histogram(locs, first_derivs)\n",
    "    (xs, ys) = evaluate_histogram((x, y), threshold=threshold)\n",
    "    out_data = (xs[1:-1], ys[1:-1])\n",
    "    return out_data\n",
    "\n",
    "def normalize_quantiles(in_data, threshold=epsilon):\n",
    "    \"\"\"\n",
    "    Evaluates PDF from quantiles including endpoints from linear extrapolation\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    in_data: tuple, numpy.ndarray, float\n",
    "        tuple of CDF values iy corresponding to quantiles and the points x at\n",
    "        which those CDF values are achieved\n",
    "    threshold: float, optional\n",
    "        optional minimum threshold for PDF\n",
    "    vb: boolean, optional\n",
    "        be careful and print progress to stdout?\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    out_data: tuple, ndarray, float\n",
    "        tuple of values x at which CDF is achieved, including extrema, and\n",
    "        normalized PDF values y at x\n",
    "    \"\"\"\n",
    "\n",
    "    # iy = quants\n",
    "    # x = locs\n",
    "    (iy, x) = in_data\n",
    "    (xs, ys) = evaluate_quantiles((iy, x))\n",
    "    # xs = xs[1:-1]\n",
    "    # ys = ys[1:-1]\n",
    "\n",
    "    # ! I believe that this is just using the slope to add end points to the list of x values.\n",
    "    x_min = xs[0] - 2 * iy[0] / ys[0]\n",
    "    x_max = xs[-1] + 2 * (1. - iy[-1]) / ys[-1]\n",
    "    xs = sandwich(xs, (x_min, x_max))\n",
    "    ys = sandwich(ys, (threshold, threshold))\n",
    "    out_data = (xs, ys)\n",
    "    return out_data\n",
    "\n",
    "def normalize_gridded(in_data, thresholds=(epsilon, infty)):\n",
    "    \"\"\"\n",
    "    Removes extreme values from gridded parametrizations\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    in_data: None or tuple, numpy.ndarray, float\n",
    "        tuple of points x at which function is evaluated and the PDF y at those\n",
    "        points\n",
    "    thresholds: tuple, float, optional\n",
    "        optional min/max thresholds for normalization\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    out_data: tuple, numpy.ndarray, float\n",
    "        tuple of input x and normalized y\n",
    "    \"\"\"\n",
    "    if in_data is None:\n",
    "        return in_data\n",
    "    (x, y) = in_data\n",
    "    y[y < thresholds[0]] = thresholds[0]\n",
    "    y[y > thresholds[-1]] = thresholds[-1]\n",
    "    out_data = (x, y)\n",
    "    return out_data\n",
    "\n",
    "def original_qp_interpolation(quants, locs):\n",
    "\n",
    "    # ! These are the variable name substitutions I've made\n",
    "    # quants_and_locs = self.quantiles # A 2d array of quantiles and locations\n",
    "    # scheme = scheme # using 'linear', used to determine the type of interpolation with interp1d\n",
    "\n",
    "    scheme = 'linear'\n",
    "    quants_and_locs = np.array([quants, locs])\n",
    "\n",
    "    if type(scheme) != int:\n",
    "        order = min(5, len(quants_and_locs[0]))\n",
    "    else:\n",
    "        order = scheme\n",
    "\n",
    "    # ! Looks like only min(x) and max(x) are used after this in this function, but x,y are used in `quantile_interpolator` later.\n",
    "    (x, y) = normalize_quantiles(quants_and_locs)\n",
    "    z = np.insert(quants_and_locs[1], 0, min(x))\n",
    "    z = np.append(z, max(x))\n",
    "    q = np.insert(quants_and_locs[0], 0, 0.)\n",
    "    q = np.append(q, 1.)\n",
    "\n",
    "    [x_crit_lo, x_crit_hi] = [quants_and_locs[1][0], quants_and_locs[1][-1]]\n",
    "    [y_crit_lo, y_crit_hi] = [-1., -1.]\n",
    "\n",
    "    try:\n",
    "        while (order>0) and ((y_crit_lo <= 0.) or (y_crit_hi <= 0.)):\n",
    "            inside = InterpolatedUnivariateSpline(z, q, k=order, ext=1).derivative()\n",
    "            [y_crit_lo, y_crit_hi] = inside([x_crit_lo, x_crit_hi])\n",
    "            order -= 1\n",
    "        assert((y_crit_lo > 0.) and (y_crit_hi > 0.))\n",
    "    except AssertionError:\n",
    "        print('ERROR: spline tangents '+str((y_crit_lo, y_crit_hi))+'<0')\n",
    "        if type(scheme) == str:\n",
    "            this_scheme = scheme\n",
    "        else:\n",
    "            this_scheme = 'linear'\n",
    "        inside_int = interp1d(z, q, kind=this_scheme, bounds_error=False, fill_value=epsilon)\n",
    "        derivative = (q[1:] - q[:-1]) / (z[1:] - z[:-1])\n",
    "        derivative = np.insert(derivative, 0, epsilon)\n",
    "        derivative = np.append(derivative, epsilon)\n",
    "        def inside(xf):\n",
    "            nx = len(xf)\n",
    "            yf = np.ones(nx) * epsilon\n",
    "            for n in range(nx):\n",
    "                i = bisect.bisect_left(z, xf[n])\n",
    "                yf[n] = derivative[i]\n",
    "            return(yf)\n",
    "        [y_crit_lo, y_crit_hi] = inside([x_crit_lo, x_crit_hi])\n",
    "        assert((y_crit_lo > 0.) and (y_crit_hi > 0.))\n",
    "    return quants_and_locs, x_crit_lo, x_crit_hi, y_crit_lo, y_crit_hi, x, y, z, inside\n",
    "\n",
    "def quantile_interpolator(quants_and_locs, x_crit_lo, x_crit_hi, y_crit_lo, y_crit_hi, x, y, z, inside, xf):\n",
    "    yf = np.ones(np.shape(xf)) * epsilon\n",
    "    in_inds = ((xf >= quants_and_locs[1][0]) & (xf <= quants_and_locs[1][-1])).nonzero()[0]\n",
    "    lo_inds = ((xf < quants_and_locs[1][0]) & (xf >= z[0])).nonzero()[0]\n",
    "    hi_inds = ((xf > quants_and_locs[1][-1]) & (xf <= z[-1])).nonzero()[0]\n",
    "\n",
    "    try:\n",
    "        yf[in_inds] = inside(xf[in_inds])\n",
    "        assert(np.all(yf >= epsilon))\n",
    "\n",
    "    except AssertionError:\n",
    "        print('ERROR: spline interpolation failed with '+str((xf[in_inds], yf[in_inds])))\n",
    "        try:\n",
    "            alternate = interp1d(x, y, kind='linear', bounds_error=False, fill_value=epsilon)\n",
    "            yf[in_inds] = alternate(xf[in_inds])\n",
    "            assert(np.all(yf >= epsilon))\n",
    "\n",
    "        except AssertionError:\n",
    "            print('ERROR: linear interpolation failed for the '+using+' parametrization with '+str((xf[in_inds], yf[in_inds])))\n",
    "            backup = qp.utils.make_kludge_interpolator((x, y), threshold=epsilon)\n",
    "            yf[in_inds] = backup(xf[in_inds])\n",
    "\n",
    "            assert(np.all(yf >= epsilon))\n",
    "\n",
    "    try:\n",
    "        tan_lo = y_crit_lo / (x_crit_lo - z[0])\n",
    "        yf[lo_inds] = tan_lo * (xf[lo_inds] - z[0])# yf[in_inds[0]] / (xf[in_inds[0]] - z[0])\n",
    "        assert(np.all(yf >= epsilon))\n",
    "    except AssertionError:\n",
    "        print('ERROR: linear extrapolation below failed with '+str((xf[lo_inds], yf[lo_inds]))+' via '+str((tan_lo, x_crit_lo, z[0])))\n",
    "\n",
    "    try:\n",
    "        tan_hi = y_crit_hi / (z[-1] - x_crit_hi)\n",
    "        yf[hi_inds] = tan_hi * (z[-1] - xf[hi_inds])# yf[in_inds[-1]] * (xf[hi_inds] - z[-1]) / (xf[in_inds[-1]] - z[-1])\n",
    "        assert(np.all(yf >= epsilon))\n",
    "    except AssertionError:\n",
    "        print('ERROR: linear extrapolation above failed with '+str((xf[hi_inds], yf[hi_inds]))+' via '+str((tan_hi, z[-1], x_crit_hi)))\n",
    "\n",
    "    return(yf)\n",
    "\n",
    "def approximate(quants, locs, grid=None):\n",
    "        \"\"\"\n",
    "        Interpolates the parametrization to get an approximation to the density.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        points: ndarray\n",
    "            the value(s) at which to evaluate the interpolated function\n",
    "        \n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        points: ndarray, float\n",
    "            the input grid upon which to interpolate\n",
    "        interpolated: ndarray, float\n",
    "            the interpolated points.\n",
    "\n",
    "        Notes\n",
    "        -----\n",
    "        Extrapolation is via the `scheme` while values are positive;\n",
    "        otherwise, extrapolation returns 0.\n",
    "\n",
    "        Example:\n",
    "            x, y = p.approximate(np.linspace(-1., 1., 100))\n",
    "        \"\"\"\n",
    "        if grid is None:\n",
    "            grid = locs\n",
    "\n",
    "        quants_and_locs, x_crit_lo, x_crit_hi, y_crit_lo, y_crit_hi, x, y, z, inside = original_qp_interpolation(quants, locs)\n",
    "\n",
    "        interpolated = quantile_interpolator(quants_and_locs, x_crit_lo, x_crit_hi, y_crit_lo, y_crit_hi, x, y, z, inside, grid)\n",
    "        \n",
    "        interpolated = normalize_gridded((grid, interpolated))\n",
    "\n",
    "        return interpolated"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Example of a single distribution made up of two merged gaussian distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Experimenting with well separated multi-modal Gaussians.\n",
    "mu = np.array([1.7, 2.9])\n",
    "sig = np.array([0.3,0.4])\n",
    "wt = np.array([0.4, 0.6])\n",
    "\n",
    "merged_norms = qp.Ensemble(qp.mixmod, data=dict(means=mu, stds=sig, weights=wt))\n",
    "\n",
    "new_quants = np.linspace(0.001, 0.999, 11)\n",
    "new_locs = merged_norms.ppf(new_quants)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! When this runs it uses 4th degree spline interpolator as `inside` in quantile_interpolator\n",
    "test = approximate(new_quants, new_locs, np.linspace(0,5,100))\n",
    "fig, axes = qp.plotting.plot_native(merged_norms, xlim=(0, 5))\n",
    "axes.plot(test[0], test[1])\n",
    "axes.vlines(new_locs, 0, 0.01)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Example of a single distribution made up of two well separated gaussian distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Experimenting with well separated multi-modal Gaussians.\n",
    "mu = np.array([0.5,2.9])\n",
    "sig = np.array([0.05,0.04])\n",
    "wt = np.array([.4,.6])\n",
    "\n",
    "distinct_norms = qp.Ensemble(qp.mixmod, data=dict(means=mu, stds=sig, weights=wt))\n",
    "\n",
    "new_quants = np.linspace(0.001, 0.999, 15)\n",
    "new_locs = distinct_norms.ppf(new_quants)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When running `approximate` here, the spline fit fails, and we fall back to a linear interpolation over the numerical derivative."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! When this runs it uses linear interpolation of x,y in the first `except` clause in quantile_interpolator\n",
    "test = approximate(new_quants, new_locs, np.linspace(0,5,100))\n",
    "fig, axes = qp.plotting.plot_native(distinct_norms, xlim=(0, 5))\n",
    "axes.plot(test[0], test[1])\n",
    "axes.vlines(new_locs, 0, 0.01)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "qp",
   "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.9"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "5c3e9e8e8c57812b479c84246ef042cffef7a98b6ccc136408e9cfe0d107d268"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
