{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Horizontal interpolation: Nearest Neighbours"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nes\n",
    "import geopandas as gpd\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "countries = gpd.read_file('/gpfs/projects/bsc32/models/NES_tutorial_data/gadm_country_mask/gadm_country_ISO3166.shp')\n",
    "\n",
    "def plot_data(nessy, var_name, zoom=None):\n",
    "    nessy.create_spatial_bounds()\n",
    "    lon_bnds, lat_bnds = nessy.get_spatial_bounds_mesh_format()\n",
    "    \n",
    "    fig, ax = plt.subplots(1, figsize=(19, 7))\n",
    "    plt.pcolormesh(lon_bnds, \n",
    "                   lat_bnds, \n",
    "                   nessy.variables[var_name]['data'].sum(axis=(0, 1)), \n",
    "                   cmap='jet', \n",
    "                   norm=matplotlib.colors.LogNorm())\n",
    "    \n",
    "    countries.plot(ax=ax, facecolor='none', edgecolor='black', linewidth=0.3)\n",
    "    plt.colorbar(ax=ax)\n",
    "    ax.margins(0)\n",
    "    \n",
    "    if zoom is not None:\n",
    "        plt.axis([zoom['x_min'], zoom['x_max'], zoom['y_min'], zoom['y_max']])\n",
    "\n",
    "    return None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Read data to interpolate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Original path: /esarchive/exp/monarch/a4dd/original_files/000/2022111512/MONARCH_d01_2022111512.nc\n",
    "# Rotated grid from MONARCH\n",
    "source_path = '/gpfs/projects/bsc32/models/NES_tutorial_data/MONARCH_d01_2022111512.nc'\n",
    "var_name = 'O3'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_grid = nes.open_netcdf(path=source_path)\n",
    "source_grid.keep_vars(var_name)\n",
    "source_grid.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_data(source_grid, var_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Interpolation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1. Read or Create the destination grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CAMS2_40 Forecast final grid\n",
    "\n",
    "lat_orig = 30.05\n",
    "lon_orig = -24.95\n",
    "inc_lat = 0.1\n",
    "inc_lon = 0.1\n",
    "n_lat = 420\n",
    "n_lon = 700\n",
    "\n",
    "dst_grid = nes.create_nes(comm=None, info=False, projection='regular',\n",
    "                          lat_orig=lat_orig, lon_orig=lon_orig, \n",
    "                          inc_lat=inc_lat, inc_lon=inc_lon, \n",
    "                          n_lat=n_lat, n_lon=n_lon)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2. Interpolate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Nearest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time interp_nn1 = source_grid.interpolate_horizontal(dst_grid, kind='NearestNeighbour', n_neighbours=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_data(interp_nn1, var_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Zoom in plot\n",
    "plot_data(interp_nn1, var_name, zoom={'x_min': -30, 'x_max': 50, 'y_min': 25, 'y_max':75})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Bilinear"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time interp_nn4 = source_grid.interpolate_horizontal(dst_grid, kind='NearestNeighbour', n_neighbours=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_data(interp_nn4, var_name, zoom={'x_min': -30, 'x_max': 50, 'y_min': 25, 'y_max':75})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Bicubic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time interp_nn16 = source_grid.interpolate_horizontal(dst_grid, kind='NearestNeighbour', n_neighbours=16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_data(interp_nn16, var_name, zoom={'x_min': -30, 'x_max': 50, 'y_min': 25, 'y_max':75})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 32 Nearest Neighbours"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time interp_nn32 = source_grid.interpolate_horizontal(dst_grid, kind='NearestNeighbour', n_neighbours=32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_data(interp_nn32, var_name, zoom={'x_min': -30, 'x_max': 50, 'y_min': 25, 'y_max':75})"
   ]
  }
 ],
 "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.9.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
