{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Horizontal conservative interpolation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nes import *\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "import numpy as np\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_data(nessy, var_name):\n",
    "    nessy.create_spatial_bounds()\n",
    "    lon_bnds, lat_bnds = nessy.get_spatial_bounds_mesh_format()\n",
    "    ax = plt.axes()\n",
    "    plt.pcolormesh(lon_bnds, lat_bnds, nessy.variables[var_name]['data'][0,0], cmap='jet', norm=matplotlib.colors.LogNorm())\n",
    "    "
   ]
  },
  {
   "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 (NAMEE)\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": [
    "# Global CAMS domain\n",
    "source_path = \"/gpfs/projects/bsc32/models/NES_tutorial_data/nox_no_201505.nc\"\n",
    "var_name = \"nox_no\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>geometry</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>FID</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>POLYGON ((-180.00000 -90.00000, -179.89999 -90...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>POLYGON ((-179.89999 -90.00000, -179.80002 -90...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>POLYGON ((-179.79999 -90.00000, -179.70001 -90...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>POLYGON ((-179.69998 -90.00000, -179.60001 -90...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>POLYGON ((-179.60001 -90.00000, -179.50000 -90...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6479995</th>\n",
       "      <td>POLYGON ((179.50000 89.89999, 179.60001 89.899...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6479996</th>\n",
       "      <td>POLYGON ((179.60001 89.89999, 179.69998 89.899...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6479997</th>\n",
       "      <td>POLYGON ((179.70001 89.89999, 179.79999 89.899...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6479998</th>\n",
       "      <td>POLYGON ((179.80002 89.89999, 179.89999 89.899...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6479999</th>\n",
       "      <td>POLYGON ((179.89999 89.89999, 180.00000 89.899...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>6480000 rows × 1 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                  geometry\n",
       "FID                                                       \n",
       "0        POLYGON ((-180.00000 -90.00000, -179.89999 -90...\n",
       "1        POLYGON ((-179.89999 -90.00000, -179.80002 -90...\n",
       "2        POLYGON ((-179.79999 -90.00000, -179.70001 -90...\n",
       "3        POLYGON ((-179.69998 -90.00000, -179.60001 -90...\n",
       "4        POLYGON ((-179.60001 -90.00000, -179.50000 -90...\n",
       "...                                                    ...\n",
       "6479995  POLYGON ((179.50000 89.89999, 179.60001 89.899...\n",
       "6479996  POLYGON ((179.60001 89.89999, 179.69998 89.899...\n",
       "6479997  POLYGON ((179.70001 89.89999, 179.79999 89.899...\n",
       "6479998  POLYGON ((179.80002 89.89999, 179.89999 89.899...\n",
       "6479999  POLYGON ((179.89999 89.89999, 180.00000 89.899...\n",
       "\n",
       "[6480000 rows x 1 columns]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "source_grid = open_netcdf(path=source_path)\n",
    "source_grid.create_shapefile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Flux units: m-2.kg.s-1\n"
     ]
    }
   ],
   "source": [
    "source_grid.keep_vars(var_name)\n",
    "print('Flux units: {0}'.format(source_grid.variables[var_name]['units']))\n",
    "source_grid.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 63272 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_data(source_grid, var_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nox_no total flux: 1.0672498547137366e-06\n",
      "nox_no total mass: 113.08470916748047\n"
     ]
    }
   ],
   "source": [
    "# Flux to mass\n",
    "cell_area = source_grid.calculate_grid_area()\n",
    "\n",
    "for var_aux in source_grid.variables.keys():\n",
    "    print(\"{0} total flux: {1}\".format(var_aux, source_grid.variables[var_aux]['data'].sum()))\n",
    "    source_grid.variables[var_aux]['data'] *= cell_area\n",
    "    source_grid.variables[var_aux]['units'] = 'kg.s-1'\n",
    "    print(\"{0} total mass: {1}\".format(var_aux, source_grid.variables[var_aux]['data'].sum()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 62116 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_data(source_grid, var_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Create destination grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>geometry</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>FID</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>POLYGON ((-11.58393 32.47507, -11.54169 32.478...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>POLYGON ((-11.54169 32.47851, -11.49944 32.481...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>POLYGON ((-11.49944 32.48192, -11.45719 32.485...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>POLYGON ((-11.45719 32.48533, -11.41494 32.488...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>POLYGON ((-11.41494 32.48871, -11.37268 32.492...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>157604</th>\n",
       "      <td>POLYGON ((6.95490 46.70274, 7.00684 46.69873, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>157605</th>\n",
       "      <td>POLYGON ((7.00684 46.69873, 7.05878 46.69470, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>157606</th>\n",
       "      <td>POLYGON ((7.05878 46.69470, 7.11071 46.69066, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>157607</th>\n",
       "      <td>POLYGON ((7.11071 46.69066, 7.16264 46.68659, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>157608</th>\n",
       "      <td>POLYGON ((7.16264 46.68659, 7.21456 46.68250, ...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>157609 rows × 1 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                 geometry\n",
       "FID                                                      \n",
       "0       POLYGON ((-11.58393 32.47507, -11.54169 32.478...\n",
       "1       POLYGON ((-11.54169 32.47851, -11.49944 32.481...\n",
       "2       POLYGON ((-11.49944 32.48192, -11.45719 32.485...\n",
       "3       POLYGON ((-11.45719 32.48533, -11.41494 32.488...\n",
       "4       POLYGON ((-11.41494 32.48871, -11.37268 32.492...\n",
       "...                                                   ...\n",
       "157604  POLYGON ((6.95490 46.70274, 7.00684 46.69873, ...\n",
       "157605  POLYGON ((7.00684 46.69873, 7.05878 46.69470, ...\n",
       "157606  POLYGON ((7.05878 46.69470, 7.11071 46.69066, ...\n",
       "157607  POLYGON ((7.11071 46.69066, 7.16264 46.68659, ...\n",
       "157608  POLYGON ((7.16264 46.68659, 7.21456 46.68250, ...\n",
       "\n",
       "[157609 rows x 1 columns]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lat_1 = 37\n",
    "lat_2 = 43\n",
    "lon_0 = -3\n",
    "lat_0 = 40\n",
    "nx = 397\n",
    "ny = 397\n",
    "inc_x = 4000\n",
    "inc_y = 4000\n",
    "x_0 = -807847.688\n",
    "y_0 = -797137.125\n",
    "dst_nes = create_nes(comm=None, info=False, projection='lcc', lat_1=lat_1, lat_2=lat_2, lon_0=lon_0, lat_0=lat_0,\n",
    "                     nx=nx, ny=ny, inc_x=inc_x, inc_y=inc_y, x_0=x_0, y_0=y_0, times=source_grid.get_full_times())\n",
    "dst_nes.create_shapefile()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Interpolation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1. Without creating weight matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 48s, sys: 1.42 s, total: 1min 49s\n",
      "Wall time: 1min 49s\n"
     ]
    }
   ],
   "source": [
    "%time interp_nes = source_grid.interpolate_horizontal(dst_grid=dst_nes, kind='Conservative')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 63316 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_data(interp_nes, var_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nox_no total mass: 113.08470916748047\n"
     ]
    }
   ],
   "source": [
    "for var_aux in source_grid.variables.keys():\n",
    "    print(\"{0} total mass: {1}\".format(var_aux, source_grid.variables[var_aux]['data'].sum()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2. Creating weight matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating Weight Matrix\n",
      "Weight Matrix done!\n",
      "CPU times: user 1min 49s, sys: 1.9 s, total: 1min 51s\n",
      "Wall time: 1min 52s\n"
     ]
    }
   ],
   "source": [
    "%time wm_nes = source_grid.interpolate_horizontal(dst_grid=dst_nes, kind='Conservative', info=True, weight_matrix_path=\"CONS_WM_NAMEE_to_IP.nc\", only_create_wm=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 202 ms, sys: 26.9 ms, total: 229 ms\n",
      "Wall time: 228 ms\n"
     ]
    }
   ],
   "source": [
    "%time interp_nes = source_grid.interpolate_horizontal(dst_grid=dst_nes, kind='Conservative', wm=wm_nes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 63316 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_data(interp_nes, var_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nox_no total mass: 1.4103307834025396\n"
     ]
    }
   ],
   "source": [
    "for var_aux in source_grid.variables.keys():\n",
    "    print(\"{0} total mass: {1}\".format(var_aux, interp_nes.variables[var_aux]['data'].sum()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Flux conservative interpolation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.0 Using previous example to put the data as flux"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nox_no total mass: 1.4103307834025396\n",
      "nox_no total flux: 8.809495172462768e-08\n"
     ]
    }
   ],
   "source": [
    "# Mass to Flux\n",
    "cell_area = interp_nes.calculate_grid_area()\n",
    "\n",
    "for var_aux in interp_nes.variables.keys():\n",
    "    print(\"{0} total mass: {1}\".format(var_aux, interp_nes.variables[var_aux]['data'].sum()))\n",
    "    interp_nes.variables[var_aux]['data'] /= cell_area\n",
    "    interp_nes.variables[var_aux]['units'] = 'm-2.kg.s-1'\n",
    "    print(\"{0} total flux: {1}\".format(var_aux, interp_nes.variables[var_aux]['data'].sum()))\n",
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 Interpolation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Flux units: m-2.kg.s-1\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 63272 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "source_grid = open_netcdf(path=source_path)\n",
    "source_grid.create_shapefile()\n",
    "source_grid.keep_vars(var_name)\n",
    "print('Flux units: {0}'.format(source_grid.variables[var_name]['units']))\n",
    "source_grid.load()\n",
    "plot_data(source_grid, var_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 6min 12s, sys: 2.3 s, total: 6min 15s\n",
      "Wall time: 6min 16s\n"
     ]
    }
   ],
   "source": [
    "%time interp_nes = source_grid.interpolate_horizontal(dst_grid=dst_nes, kind='Conservative', flux=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nox_no total flux: 8.805639456704551e-08\n"
     ]
    }
   ],
   "source": [
    "for var_aux in interp_nes.variables.keys():\n",
    "    print(\"{0} total flux: {1}\".format(var_aux, interp_nes.variables[var_aux]['data'].sum()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 63348 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_data(interp_nes, var_name)"
   ]
  }
 ],
 "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
