{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Interpolation of LLC4320 ocean model\n",
    "===============================\n",
    "\n",
    "The interpolation of this object is based on a\n",
    "[R*Tree](https://pangeo-pyinterp.readthedocs.io/en/latest/generated/pyinterp.RTree.html#pyinterp.RTree)\n",
    "structure. To begin with, we start by building this object. By default, this\n",
    "object considers the WGS-84 geodetic coordinate system. But you can define\n",
    "another one using the class\n",
    "[System](https://pangeo-pyinterp.readthedocs.io/en/latest/generated/pyinterp.geodetic.System.html#pyinterp.geodetic.System).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy\n",
    "import cartopy.crs\n",
    "import cartopy.mpl.ticker\n",
    "import intake\n",
    "import matplotlib.pyplot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyinterp\n",
    "mesh = pyinterp.RTree()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we will insert points into the tree. The class allows you to add points\n",
    "using two algorithms. The first one, called\n",
    "[packing](https://pangeo-pyinterp.readthedocs.io/en/latest/generated/pyinterp.RTree.packing.html#pyinterp.RTree.packing),\n",
    "will enable you to enter the values in the tree at once. This mechanism is the\n",
    "recommended solution to create an optimized in-memory structure, both in terms\n",
    "of construction time and queries. When this is not possible, you can insert new\n",
    "information into the tree as you go along using the\n",
    "[insert](https://pangeo-pyinterp.readthedocs.io/en/latest/generated/pyinterp.RTree.packing.html#pyinterp.RTree.insert)\n",
    "method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cat_url = \"https://raw.githubusercontent.com/pangeo-data/pangeo-datastore\" \\\n",
    "    \"/master/intake-catalogs/ocean/llc4320.yaml\"\n",
    "cat = intake.open_catalog(cat_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Grid subsampling (original volume is too huge for this example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "indices = slice(0, None, 8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reads longitudes and latitudes of the grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array = cat.LLC4320_grid.to_dask()\n",
    "lons = array[\"XC\"].isel(i=indices, j=indices)\n",
    "lats = array[\"YC\"].isel(i=indices, j=indices)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reads SSH values for the first time step of the time series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ssh = cat.LLC4320_SSH.to_dask()\n",
    "ssh = ssh[\"Eta\"].isel(time=0, i=indices, j=indices)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Populates the search tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh.packing(\n",
    "    numpy.vstack((lons.values.ravel(), lats.values.ravel())).T,\n",
    "    ssh.values.ravel())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When the tree is created, you can interpolate data with two algorithms:\n",
    "\n",
    "-   [Inverse Distance Weighting](https://pangeo-pyinterp.readthedocs.io/en/latest/generated/pyinterp.RTree.inverse_distance_weighting.html#pyinterp.RTree.inverse_distance_weighting)\n",
    "    or IDW\n",
    "-   [Radial Basis Function](https://pangeo-pyinterp.readthedocs.io/en/latest/generated/pyinterp.RTree.radial_basis_function.html#pyinterp.RTree.radial_basis_function)\n",
    "    or RBF\n",
    "\n",
    "Yon can also search the\n",
    "[nearest neighbors](https://pangeo-pyinterp.readthedocs.io/en/latest/generated/pyinterp.RTree.query.html#pyinterp.RTree.query)\n",
    "on the tree.\n",
    "\n",
    "---\n",
    "**Note**\n",
    "\n",
    "When comparing an RBF to IDW, IDW will never predict values higher than\n",
    "the maximum measured value or lower than the minimum measured value.\n",
    "However, RBFs can predict values higher than the maximum values and\n",
    "lower than the minimum measured values.\n",
    "\n",
    "---\n",
    "\n",
    "In this example, we will under-sample the source grid at 1/32 degree\n",
    "over an area of the globe.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x0, x1 = 80, 170\n",
    "y0, y1 = -45, 30\n",
    "res = 1 / 32.0\n",
    "mx, my = numpy.meshgrid(numpy.arange(x0, x1, res),\n",
    "                        numpy.arange(y0, y1, res),\n",
    "                        indexing=\"ij\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "IDW interpolation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "idw_eta, neighbors = mesh.inverse_distance_weighting(\n",
    "    numpy.vstack((mx.ravel(), my.ravel())).T,\n",
    "    within=True,  # Extrapolation is forbidden\n",
    "    radius=55000,  # In a radius of 5.5 Km\n",
    "    k=8,  # We are looking for at most 8 neighbours\n",
    "    num_threads=0)\n",
    "idw_eta = idw_eta.reshape(mx.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "RBF interpolation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rbf_eta, neighbors = mesh.radial_basis_function(\n",
    "    numpy.vstack((mx.ravel(), my.ravel())).T,\n",
    "    within=True,  # Extrapolation is forbidden\n",
    "    k=11,  # We are looking for at most 11 neighbours\n",
    "    num_threads=0)\n",
    "rbf_eta = rbf_eta.reshape(mx.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualize our interpolated data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = matplotlib.pyplot.figure(figsize=(18, 9))\n",
    "lon_formatter = cartopy.mpl.ticker.LongitudeFormatter(\n",
    "    zero_direction_label=True)\n",
    "lat_formatter = cartopy.mpl.ticker.LatitudeFormatter()\n",
    "ax = fig.add_subplot(121, projection=cartopy.crs.PlateCarree())\n",
    "ax.pcolormesh(mx,\n",
    "              my,\n",
    "              idw_eta,\n",
    "              cmap='terrain',\n",
    "              shading='auto',\n",
    "              transform=cartopy.crs.PlateCarree())\n",
    "ax.coastlines()\n",
    "ax.xaxis.set_major_formatter(lon_formatter)\n",
    "ax.yaxis.set_major_formatter(lat_formatter)\n",
    "ax.set_xticks(numpy.arange(x0, x1, 10.0))\n",
    "ax.set_yticks(numpy.arange(y0, y1, 10))\n",
    "ax.set_title(\"Eta (IDW)\")\n",
    "\n",
    "ax = fig.add_subplot(122, projection=cartopy.crs.PlateCarree())\n",
    "ax.pcolormesh(mx,\n",
    "              my,\n",
    "              rbf_eta,\n",
    "              cmap='terrain',\n",
    "              shading='auto',\n",
    "              transform=cartopy.crs.PlateCarree())\n",
    "ax.coastlines()\n",
    "ax.xaxis.set_major_formatter(lon_formatter)\n",
    "ax.yaxis.set_major_formatter(lat_formatter)\n",
    "ax.set_xticks(numpy.arange(x0, x1, 10.0))\n",
    "ax.set_yticks(numpy.arange(y0, y1, 10))\n",
    "ax.set_title(\"Eta (RBF)\")\n",
    "fig.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}