{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "%matplotlib inline"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "4D interpolation\n",
    "=============\n",
    "\n",
    "Interpolation of a four-dimensional regular grid.\n",
    "\n",
    "Quadrivariate\n",
    "----------------\n",
    "\n",
    "The\n",
    "[quadrivariate](https://pangeo-pyinterp.readthedocs.io/en/latest/generated/pyinterp.quadrivariate.html#pyinterp.quadrivariate)\n",
    "interpolation allows obtaining values at arbitrary points in a 4D space of a\n",
    "function defined on a grid.\n",
    "\n",
    "The distribution contains a 4D field `pres_temp_4D.nc` that will be used\n",
    "in this help. This file is located in the `src/pyinterp/tests/dataset` directory\n",
    "at the root of the project.\n",
    "\n",
    "This method performs a bilinear interpolation in 2D space by considering the\n",
    "axes of longitude and latitude of the grid, then performs a linear interpolation\n",
    "in the third and fourth dimensions. Its interface is similar to the\n",
    "[trivariate](https://pangeo-pyinterp.readthedocs.io/en/latest/generated/pyinterp.trivariate.html#pyinterp.trivariate)\n",
    "class except for a fourth axis, which is handled by this object."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import cartopy.crs\n",
    "import matplotlib\n",
    "import matplotlib.pyplot\n",
    "import numpy\n",
    "import pyinterp\n",
    "import pyinterp.backends.xarray\n",
    "import pyinterp.tests\n",
    "import xarray"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "The first step is to load the data into memory and create the\n",
    "interpolator object:"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "ds = xarray.open_dataset(pyinterp.tests.grid4d_path())\n",
    "interpolator = pyinterp.backends.xarray.Grid4D(ds.pressure)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "We will build a new grid that will be used to build a new interpolated\n",
    "grid."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "mx, my, mz, mu = numpy.meshgrid(numpy.arange(-125, -70, 0.5),\n",
    "                                numpy.arange(25, 50, 0.5),\n",
    "                                numpy.datetime64(\"2000-01-01T12:00\"),\n",
    "                                0.5,\n",
    "                                indexing=\"ij\")"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "We interpolate our grid using a\n",
    "[classical](https://pangeo-pyinterp.readthedocs.io/en/latest/generated/pyinterp.backends.xarray.Grid4D.quadrivariate.html#pyinterp.backends.xarray.Grid4D.quadrivariate):"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "quadrivariate = interpolator.quadrivariate(\n",
    "    dict(longitude=mx.flatten(),\n",
    "         latitude=my.flatten(),\n",
    "         time=mz.flatten(),\n",
    "         level=mu.flatten())).reshape(mx.shape)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Bicubic on 4D grid\n",
    "----------------------\n",
    "\n",
    "The grid used organizes the latitudes in descending order. We ask our\n",
    "constructor to flip this axis in order to correctly evaluate the bicubic\n",
    "interpolation from this 4D cube (only necessary to perform a bicubic\n",
    "interpolation)."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "interpolator = pyinterp.backends.xarray.Grid4D(ds.pressure,\n",
    "                                               increasing_axes=True)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "We interpolate our grid using a\n",
    "[bicubic](https://pangeo-pyinterp.readthedocs.io/en/latest/generated/pyinterp.backends.xarray.Grid4D.quadrivariate.html#pyinterp.backends.xarray.Grid4D.bicubic)\n",
    "interpolation in space followed by a linear interpolation in the temporal axis:"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "bicubic = interpolator.bicubic(dict(longitude=mx.flatten(),\n",
    "                                    latitude=my.flatten(),\n",
    "                                    time=mz.flatten(),\n",
    "                                    level=mu.flatten()),\n",
    "                               nx=2,\n",
    "                               ny=2).reshape(mx.shape)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "We transform our result cubes into a matrix."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "quadrivariate = quadrivariate.squeeze(axis=(2, 3))\n",
    "bicubic = bicubic.squeeze(axis=(2, 3))\n",
    "lons = mx[:, 0].squeeze()\n",
    "lats = my[0, :].squeeze()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Let's visualize our results.\n",
    "\n",
    "---\n",
    "**Note**\n",
    "\n",
    "The resolution of the example grid is very low (one pixel every one degree)\n",
    "therefore the calculation window cannot find the required pixels at the edges to\n",
    "calculate the interpolation correctly. See Chapter [Fill NaN\n",
    "Values](https://pangeo-pyinterp.readthedocs.io/en/latest/auto_examples/ex_fill_undef.html)\n",
    "to see how to address this issue.\n",
    "\n",
    "---"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "fig = matplotlib.pyplot.figure(figsize=(5, 4))\n",
    "ax1 = fig.add_subplot(\n",
    "    211, projection=cartopy.crs.PlateCarree(central_longitude=180))\n",
    "pcm = ax1.pcolormesh(lons,\n",
    "                     lats,\n",
    "                     quadrivariate.T,\n",
    "                     cmap='jet',\n",
    "                     transform=cartopy.crs.PlateCarree())\n",
    "ax1.coastlines()\n",
    "ax1.set_title(\"Trilinear\")\n",
    "\n",
    "ax2 = fig.add_subplot(\n",
    "    212, projection=cartopy.crs.PlateCarree(central_longitude=180))\n",
    "pcm = ax2.pcolormesh(lons,\n",
    "                     lats,\n",
    "                     bicubic.T,\n",
    "                     cmap='jet',\n",
    "                     transform=cartopy.crs.PlateCarree())\n",
    "ax2.coastlines()\n",
    "ax2.set_title(\"Spline & Linear in time\")\n",
    "fig.colorbar(pcm, ax=[ax1, ax2], shrink=0.8)\n",
    "fig.show()"
   ],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}