{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2021-2024 Lawrence Livermore National Security, LLC and other MuyGPyS\n",
    "Project Developers. See the top-level COPYRIGHT file for details.\n",
    "\n",
    "SPDX-License-Identifier: MIT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fast Posterior Mean Tutorial\n",
    "\n",
    "This notebook walks through the fast posterior mean workflow presented in Fast Gaussian Process Posterior Mean Prediction via Local Cross Validation and Precomputation (Dunton et. al 2022) and explains the relevant components of `MuyGPyS`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "for m in sys.modules.keys():\n",
    "    if m.startswith(\"Muy\"):\n",
    "        sys.modules.pop(m)\n",
    "%env MUYGPYS_BACKEND=numpy\n",
    "%env MUYGPYS_FTYPE=64"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cell below uses the same code as that found in univariate_regression_tutorial.ipynb. This includes generating the synthetic data from a GP and training two MuyGPs models to fit the data using Bayesian optimization. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import timeit\n",
    "\n",
    "from MuyGPyS._test.gp import benchmark_sample, BenchmarkGP\n",
    "from MuyGPyS._test.sampler import UnivariateSampler, print_fast_results\n",
    "from MuyGPyS.neighbors import NN_Wrapper\n",
    "from MuyGPyS.gp import MuyGPS\n",
    "from MuyGPyS.gp.deformation import Isotropy, l2\n",
    "from MuyGPyS.gp.hyperparameter import AnalyticScale, Parameter\n",
    "from MuyGPyS.gp.kernels import Matern\n",
    "from MuyGPyS.gp.noise import HomoscedasticNoise\n",
    "from MuyGPyS.gp.tensors import fast_nn_update, make_fast_predict_tensors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will assume that we have already optimized the a MuyGPs model following the [Univariate Regression Tutorial](./univariate_regression_tutorial.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kernel = Matern(\n",
    "    smoothness=Parameter(2.0),\n",
    "    deformation=Isotropy(\n",
    "        l2,\n",
    "        length_scale=Parameter(0.05),\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use that kernel to simulate a curve and then compare the prediction times for both the conventional regression and the\n",
    "[fast kernel regression method](https://arxiv.org/abs/2205.10879v1)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(0)\n",
    "measurement_noise = 1e-4\n",
    "sampler = UnivariateSampler(\n",
    "    data_count=3000,\n",
    "    train_ratio=0.1,\n",
    "    kernel=kernel,\n",
    "    noise=HomoscedasticNoise(1e-14),\n",
    "    measurement_noise=HomoscedasticNoise(measurement_noise),\n",
    ")\n",
    "train_features, test_features = sampler.features()\n",
    "test_count = test_features.shape[0]\n",
    "train_count = train_features.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_responses, test_responses = sampler.sample()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll visualize the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler.plot_sample()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then prepare a `MuyGPS` object and a nearest neighbors index.\n",
    "We could use a single `MuyGPS` object, but in this case we create a second one for the fast regression because a larger noise prior helps to stabilize the computations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nbrs_lookup = NN_Wrapper(train_features, nn_count=10, nn_method=\"exact\",algorithm=\"ball_tree\")\n",
    "muygps = MuyGPS(\n",
    "    kernel=kernel,\n",
    "    noise=HomoscedasticNoise(1e-4),\n",
    "    scale=AnalyticScale(),\n",
    ")\n",
    "muygps_fast = MuyGPS(\n",
    "    kernel=kernel, \n",
    "    noise=HomoscedasticNoise(1e-1),\n",
    "    scale=AnalyticScale(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Benchmarking Fast Prediction\n",
    "\n",
    "With set (or learned) hyperparameters, we are able to use the `muygps` object for fast prediction capability.\n",
    "\n",
    "See below a fast posterior mean workflow, using the data structures built up in this example.\n",
    "This workflow uses the compact tensor-making function \n",
    "[make_fast_predict_tensors()](../MuyGPyS/gp/tensors.rst)\n",
    "to succinctly create tensors defining the `pairwise_dists` among each nearest neighbor and the `train_nn_targets_fast` or responses of the nearest neighbors in each set.\n",
    "We then create the `Kin` covariance tensor and form the precomputed coefficients matrix.\n",
    "We then pass the precomputed coefficients matrix, the `nn_indices` matrix of neighborhood indices, and the closest neighbor of each test point to [MuyGPS.fast_posterior_mean()](../MuyGPyS/gp/MuyGPS.rst) in order to obtain our predictions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we obtain the indices of the nearest neighbors of all of the training datapoints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_nn_indices, _ = nbrs_lookup.get_nns(train_features)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then update these neighborhoods with the index of the corresponding training point so that each neighborhood contains the query point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_nn_indices_fast = fast_nn_update(train_nn_indices)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then compute the pairwise distance tensor and target matrix and use them to construct the corresponding kernel tensor and the precomputed target matrix to be used in the fast kernel regression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pairwise_dists_fast, nn_targets_fast = make_fast_predict_tensors(\n",
    "    train_nn_indices_fast,\n",
    "    train_features,\n",
    "    train_responses,\n",
    ")\n",
    "Kin_fast = muygps_fast.kernel(\n",
    "    muygps_fast.kernel.deformation.metric(pairwise_dists_fast)\n",
    ") \n",
    "precomputed_coefficients_matrix = muygps_fast.fast_coefficients(Kin_fast, nn_targets_fast)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The steps so far have involved only the training data, and can be precomputed before encountering the test data.\n",
    "We now find the closest training point to each test point and return the corresponding enriched training points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_indices = np.arange(test_count)\n",
    "test_nn_indices, _ = nbrs_lookup.get_nns(test_features)\n",
    "closest_neighbor = test_nn_indices[:, 0]\n",
    "closest_set = train_nn_indices_fast[closest_neighbor, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use these indices to make the crosswise distance tensor, similar to usual prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "crosswise_dists_fast = muygps.kernel.deformation.crosswise_tensor(\n",
    "    test_features, train_features, test_indices, closest_set\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we compute the crosscovariance and perform fast prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Kcross_fast = muygps_fast.kernel(crosswise_dists_fast)\n",
    "predictions_fast = muygps_fast.fast_posterior_mean(\n",
    "    Kcross_fast,\n",
    "    precomputed_coefficients_matrix[closest_neighbor],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comparison with Conventional Prediction\n",
    "\n",
    "With set (or learned) hyperparameters, we are able to use the `muygps` object to predict the response of test data.\n",
    "Several workflows are supported.\n",
    "\n",
    "See below a simple posterior mean workflow, using the data structures built up in this example.\n",
    "This is very similar to the prediction workflow found in the [univariate regression tutorial](./univariate_regression_tutorial.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    crosswise_dists,\n",
    "    pairwise_dists,\n",
    "    nn_targets,\n",
    ") = muygps.make_predict_tensors(\n",
    "    np.arange(test_count),\n",
    "    test_nn_indices,\n",
    "    test_features,\n",
    "    train_features,\n",
    "    train_responses,\n",
    ")\n",
    "Kcross = muygps.kernel(crosswise_dists)\n",
    "Kin = muygps.kernel(pairwise_dists)\n",
    "predictions = muygps.posterior_mean(\n",
    "    Kin, Kcross, nn_targets\n",
    ")\n",
    "variances = muygps.posterior_variance(Kin, Kcross)\n",
    "confidence_intervals = np.sqrt(variances) * 1.96\n",
    "coverage = np.count_nonzero(np.abs(test_responses - predictions) < confidence_intervals) / test_count"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We compare our two methods in terms of time-to-solution and RMSE.\n",
    "In the conventional workflow we compute the sum of the time it takes to:\n",
    "- identify the nearest neighbors of the test features,\n",
    "- form the relevant kernel tensors, and\n",
    "- solve the posterior means.\n",
    "\n",
    "In the fast posterior mean case, we compute the sum of the time it takes to:\n",
    "- identify the nearest neighbor of each test point,\n",
    "- lookup coefficients in the precomputed coefficient matrix, and\n",
    "- perform the dot product to form posterior means. \n",
    "\n",
    "Note that the fast kernel regression method does not compute a variance, and so its posterior variance, confidence intervals, and coverage are nil."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def timing_posterior_mean():\n",
    "    test_nn_indices, _ = nbrs_lookup.get_nns(test_features)\n",
    "    (\n",
    "        crosswise_dists,\n",
    "        pairwise_dists,\n",
    "        nn_targets,\n",
    "    ) = muygps.make_predict_tensors(\n",
    "        test_indices,\n",
    "        test_nn_indices,\n",
    "        test_features,\n",
    "        train_features,\n",
    "        train_responses,\n",
    "    )\n",
    "    Kcross = muygps.kernel(crosswise_dists)\n",
    "    Kin = muygps.kernel(pairwise_dists)\n",
    "    predictions = muygps.posterior_mean(\n",
    "        Kin, Kcross, nn_targets\n",
    "    )\n",
    "\n",
    "def timing_fast_posterior_mean(): \n",
    "    test_nn_indices_fast, _ = nbrs_lookup.get_nns(test_features)\n",
    "    closest_neighbor = test_nn_indices_fast[:, 0]\n",
    "    closest_set = train_nn_indices_fast[closest_neighbor, :].astype(int)\n",
    "    crosswise_dists = muygps.kernel.deformation.crosswise_tensor(\n",
    "        test_features,\n",
    "        train_features,\n",
    "        test_indices,\n",
    "        closest_set,\n",
    "    )\n",
    "    Kcross = muygps_fast.kernel(crosswise_dists)\n",
    "    predictsion_fast = muygps_fast.fast_posterior_mean(\n",
    "        Kcross, \n",
    "        precomputed_coefficients_matrix[closest_neighbor],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "time_conv = %timeit -o timing_posterior_mean()\n",
    "time_fast = %timeit -o timing_fast_posterior_mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nil_vec = np.zeros(test_count)\n",
    "print_fast_results(\n",
    "    test_responses,\n",
    "    (\"conventional\", time_conv, muygps, predictions),\n",
    "    (\"fast\", time_fast, muygps_fast, predictions_fast),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, we can gain an order of magnitude speed improvement by sacrificing some precision and the posterior variance.\n",
    "We also plot our two methods and compare their results graphically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler.plot_results(\n",
    "    (\"conventional\", predictions, confidence_intervals),\n",
    "    (\"fast\", predictions_fast, np.zeros(test_count)),\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "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.10.9"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}