{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vertical interpolation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nes\n",
    "import matplotlib.pyplot as plt\n",
    "from datetime import datetime\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Interpolation with vertical coordinate and variable to interpolate in one file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1. Read dataset to interpolate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_grid = nes.open_netcdf(path=source_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Selecting only one time step and one variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_grid.sel(time_min=datetime(year=2022, month=11, day=16, hour=0), \n",
    "                time_max=datetime(year=2022, month=11, day=16, hour=0))\n",
    "source_grid.keep_vars([var_name, 'mid_layer_height_agl'])\n",
    "source_grid.load()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the vertical component variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_grid.vertical_var_name = 'mid_layer_height_agl'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 35324 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "\n",
    "ax.plot(source_grid.variables[var_name]['data'][0, :, 146, 90], \n",
    "        source_grid.variables['mid_layer_height_agl']['data'][0, :, 146, 90], \n",
    "        'go--')\n",
    "\n",
    "ax.xaxis.tick_top()\n",
    "ax.set_ylabel('mid_layer_height_agl' + \" ({0})\".format(source_grid.variables['mid_layer_height_agl']['units']))\n",
    "ax.set_xlabel(var_name + \" ({0})\".format(source_grid.variables[var_name]['units']))\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2. Interpolate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "level_list = [0.,50.,100.,250.,500.,750.,1000.,2000.,3000.,5000.]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Linear (recommended & faster)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 731 ms, sys: 21.1 ms, total: 753 ms\n",
      "Wall time: 737 ms\n"
     ]
    }
   ],
   "source": [
    "%time interp_linear = source_grid.interpolate_vertical(level_list, kind='linear')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Nearest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 3.29 s, sys: 680 µs, total: 3.29 s\n",
      "Wall time: 3.3 s\n"
     ]
    }
   ],
   "source": [
    "%time interp_nearest = source_grid.interpolate_vertical(level_list, kind='nearest')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cubic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 9.82 s, sys: 17.8 ms, total: 9.83 s\n",
      "Wall time: 9.87 s\n"
     ]
    }
   ],
   "source": [
    "%time interp_cubic = source_grid.interpolate_vertical(level_list, kind='cubic')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Quadratic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 9.77 s, sys: 35.3 ms, total: 9.81 s\n",
      "Wall time: 9.84 s\n"
     ]
    }
   ],
   "source": [
    "%time interp_quad = source_grid.interpolate_vertical(level_list, kind='quadratic')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 50112 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "\n",
    "ax.plot(interp_linear.variables[var_name]['data'][0, :, 146, 90], \n",
    "        level_list, \n",
    "        color='green', marker='o', label = 'linear')\n",
    "\n",
    "ax.plot(interp_nearest.variables[var_name]['data'][0, :, 146, 90], \n",
    "        level_list, \n",
    "        color='blue', marker='o', label = 'nearest')\n",
    "\n",
    "ax.plot(interp_cubic.variables[var_name]['data'][0, :, 146, 90], \n",
    "        level_list, \n",
    "        color='red', marker='o', label = 'cubic')\n",
    "\n",
    "ax.plot(interp_quad.variables[var_name]['data'][0, :, 146, 90], \n",
    "        level_list, \n",
    "        color='yellow', marker='o', label = 'quadratic')\n",
    "\n",
    "ax.xaxis.tick_top()\n",
    "ax.set_ylabel('mid_layer_height_agl' + \" ({0})\".format(source_grid.variables['mid_layer_height_agl']['units']))\n",
    "ax.set_xlabel(var_name + \" ({0})\".format(source_grid.variables[var_name]['units']))\n",
    "\n",
    "ax.legend()\n",
    "plt.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.9.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
