{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f60dd0c8",
   "metadata": {},
   "source": [
    "# Interpolation Examples\n",
    "\n",
    "This Jupyter Notebook demonstrates the usage and functionality of interpolation functions in the PLAID library. It includes examples of:\n",
    "\n",
    "1. Piece-wise linear interpolation\n",
    "2. Piece-wise linear interpolation with mapping\n",
    "3. Vectorized interpolation with mapping\n",
    "4. Vectorized interpolation\n",
    "5. Binary Search\n",
    "\n",
    "This function provides comprehensive examples and tests for interpolation functions, including piece-wise linear interpolation, interpolation with mapping, vectorized interpolation, and binary search."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "663a8965",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import required libraries\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebfc25b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries and functions\n",
    "from plaid.utils.interpolation import (\n",
    "    binary_search,\n",
    "    binary_search_vectorized,\n",
    "    piece_wise_linear_interpolation,\n",
    "    piece_wise_linear_interpolation_vectorized,\n",
    "    piece_wise_linear_interpolation_vectorized_with_map,\n",
    "    piece_wise_linear_interpolation_with_map,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dfa7ce2",
   "metadata": {},
   "source": [
    "## Section 1: Piece-wise Linear Interpolation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1307bbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Init example variables\n",
    "time_indices = np.array([0.0, 1.0, 2.5])\n",
    "vectors = np.array([np.ones(5), 2.0 * np.ones(5), 3.0 * np.ones(5)])\n",
    "\n",
    "print(f\"{time_indices = }\")\n",
    "print(f\"{vectors = }\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ee5f8ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test piece-wise linear interpolation for various inputs\n",
    "result = piece_wise_linear_interpolation(-1.0, time_indices, vectors)\n",
    "print(f\"{result = }\")\n",
    "\n",
    "np.testing.assert_almost_equal(result, [1.0, 1.0, 1.0, 1.0, 1.0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa51f377",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = piece_wise_linear_interpolation(1.0, time_indices, vectors)\n",
    "print(f\"{result = }\")\n",
    "\n",
    "np.testing.assert_almost_equal(result, [2.0, 2.0, 2.0, 2.0, 2.0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42afc33f",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = piece_wise_linear_interpolation(0.4, time_indices, vectors)\n",
    "print(f\"{result = }\")\n",
    "\n",
    "np.testing.assert_almost_equal(result, [1.4, 1.4, 1.4, 1.4, 1.4])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d5cfc24",
   "metadata": {},
   "source": [
    "## Section 2: Piece-wise Linear Interpolation with Mapping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e819dd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Init vectors variables\n",
    "vectors_map = [\"vec1\", \"vec2\", \"vec1\"]\n",
    "vectors_dict = {\"vec1\": np.ones(5), \"vec2\": 2.0 * np.ones(5)}\n",
    "\n",
    "print(f\"{vectors_map = }\")\n",
    "print(f\"{vectors_dict = }\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fc1b653",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test interpolation with mapping to named vectors\n",
    "result = piece_wise_linear_interpolation_with_map(\n",
    "    3.0, time_indices, vectors_dict, vectors_map\n",
    ")\n",
    "print(f\"{result = }\")\n",
    "\n",
    "np.testing.assert_almost_equal(result, [1.0, 1.0, 1.0, 1.0, 1.0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd81717c",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = piece_wise_linear_interpolation_with_map(\n",
    "    1.0, time_indices, vectors_dict, vectors_map\n",
    ")\n",
    "print(f\"{result = }\")\n",
    "\n",
    "np.testing.assert_almost_equal(result, [2.0, 2.0, 2.0, 2.0, 2.0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b0628e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = piece_wise_linear_interpolation_with_map(\n",
    "    0.6, time_indices, vectors_dict, vectors_map\n",
    ")\n",
    "print(f\"{result = }\")\n",
    "\n",
    "np.testing.assert_almost_equal(result, [1.6, 1.6, 1.6, 1.6, 1.6])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fde8c674",
   "metadata": {},
   "source": [
    "## Section 3: Vectorized Interpolation with Mapping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d63769b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Init input values\n",
    "input_values = np.array([-0.1, 2.0, 3.0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "564bbf92",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = piece_wise_linear_interpolation_vectorized_with_map(\n",
    "    input_values, time_indices, vectors_dict, vectors_map\n",
    ")\n",
    "print(f\"{result = }\")\n",
    "\n",
    "expected_result = [\n",
    "    np.array([1.0, 1.0, 1.0, 1.0, 1.0]),\n",
    "    np.array([1.33333333, 1.33333333, 1.33333333, 1.33333333, 1.33333333]),\n",
    "    np.array([1.0, 1.0, 1.0, 1.0, 1.0]),\n",
    "]\n",
    "\n",
    "np.testing.assert_almost_equal(result, expected_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "641d02dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Checks the accuracy of a piecewise linear interpolation function\n",
    "by comparing its output for a set of input values to a set of precomputed\n",
    "expected values.\n",
    "\"\"\"\n",
    "\n",
    "time_indices = np.array(\n",
    "    [0.0, 100.0, 200.0, 300.0, 400.0, 500.0, 600.0, 700.0, 800.0, 900.0, 1000.0, 2000.0]\n",
    ")\n",
    "\n",
    "coefficients = np.array(\n",
    "    [\n",
    "        2000000.0,\n",
    "        2200000.0,\n",
    "        2400000.0,\n",
    "        2000000.0,\n",
    "        2400000.0,\n",
    "        3000000.0,\n",
    "        2500000.0,\n",
    "        2400000.0,\n",
    "        2100000.0,\n",
    "        2800000.0,\n",
    "        4000000.0,\n",
    "        3000000.0,\n",
    "    ]\n",
    ")\n",
    "\n",
    "vals = np.array(\n",
    "    [\n",
    "        -10.0,\n",
    "        0.0,\n",
    "        100.0,\n",
    "        150.0,\n",
    "        200.0,\n",
    "        300.0,\n",
    "        400.0,\n",
    "        500.0,\n",
    "        600.0,\n",
    "        700.0,\n",
    "        800.0,\n",
    "        900.0,\n",
    "        1000.0,\n",
    "        3000.0,\n",
    "        701.4752695491923,\n",
    "    ]\n",
    ")\n",
    "\n",
    "res = np.array(\n",
    "    [\n",
    "        2000000.0,\n",
    "        2000000.0,\n",
    "        2200000.0,\n",
    "        2300000.0,\n",
    "        2400000.0,\n",
    "        2000000.0,\n",
    "        2400000.0,\n",
    "        3000000.0,\n",
    "        2500000.0,\n",
    "        2400000.0,\n",
    "        2100000.0,\n",
    "        2800000.0,\n",
    "        4000000.0,\n",
    "        3000000.0,\n",
    "        2395574.19135242,\n",
    "    ]\n",
    ")\n",
    "\n",
    "for i in range(vals.shape[0]):\n",
    "    assert (\n",
    "        piece_wise_linear_interpolation(vals[i], time_indices, coefficients) - res[i]\n",
    "    ) / res[i] < 1.0e-10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e4b26f3",
   "metadata": {},
   "source": [
    "## Section 4: Vectorized Interpolation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a169135a",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = piece_wise_linear_interpolation_vectorized(\n",
    "    np.array(vals), time_indices, coefficients\n",
    ")\n",
    "\n",
    "expected_result = [\n",
    "    2000000.0,\n",
    "    2000000.0,\n",
    "    2200000.0,\n",
    "    2300000.0,\n",
    "    2400000.0,\n",
    "    2000000.0,\n",
    "    2400000.0,\n",
    "    3000000.0,\n",
    "    2500000.0,\n",
    "    2400000.0,\n",
    "    2100000.0,\n",
    "    2800000.0,\n",
    "    4000000.0,\n",
    "    3000000.0,\n",
    "    2395574.1913524233,\n",
    "]\n",
    "\n",
    "np.testing.assert_almost_equal(result, expected_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d3a56bd",
   "metadata": {},
   "source": [
    "## Section 5: Binary Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56c1170d",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_list = np.array([0.0, 1.0, 2.5, 10.0])\n",
    "val_list = np.array([-1.0, 11.0, 0.6, 2.0, 2.6, 9.9, 1.0])\n",
    "\n",
    "# Apply binary search to find indices for given values within a reference list\n",
    "ref = np.array([0, 3, 0, 1, 2, 2, 1], dtype=int)\n",
    "result = binary_search_vectorized(test_list, val_list)\n",
    "\n",
    "for i, val in enumerate(val_list):\n",
    "    assert binary_search(test_list, val) == ref[i]\n",
    "    assert result[i] == ref[i]"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,py:percent"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}