{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to read and write regular lat-lon grids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nes\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import geopandas as gpd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Open dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Original path: /esarchive/recon/ecmwf/cams_glob_antv42/monthly_mean/nox_no_road_transport/nox_no_201603.nc\n",
    "# Global regular lat-lon grid\n",
    "path = '/gpfs/projects/bsc32/models/NES_tutorial_data/cams_glob_antv42_nox_road_transport_201603.nc'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<nes.nc_projections.latlon_nes.LatLonNes at 0x7fa30ed18340>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nessy = nes.open_netcdf(path=path, info=True)\n",
    "nessy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read projection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Proj('+proj=longlat +ellps=WGS84 +no_defs', preserve_units=True)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nessy.projection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'grid_mapping_name': 'latitude_longitude',\n",
       " 'semi_major_axis': 6371000.0,\n",
       " 'inverse_flattening': 0}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nessy.projection_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[datetime.datetime(2016, 3, 1, 0, 0)]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nessy.time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'data': array([0]), 'units': ''}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nessy.lev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'data': masked_array(data=[-89.95, -89.85, -89.75, ...,  89.75,  89.85,  89.95],\n",
       "              mask=False,\n",
       "        fill_value=1e+20,\n",
       "             dtype=float32),\n",
       " 'dimensions': ('lat',),\n",
       " 'dtype': dtype('float32'),\n",
       " 'units': 'degrees_north',\n",
       " 'axis': 'Y',\n",
       " 'long_name': 'latitude coordinate',\n",
       " 'standard_name': 'latitude',\n",
       " 'bounds': 'lat_bnds'}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nessy.lat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'data': masked_array(data=[-179.95, -179.85, -179.75, ...,  179.75,  179.85,\n",
       "                     179.95],\n",
       "              mask=False,\n",
       "        fill_value=1e+20,\n",
       "             dtype=float32),\n",
       " 'dimensions': ('lon',),\n",
       " 'dtype': dtype('float32'),\n",
       " 'units': 'degrees_east',\n",
       " 'axis': 'X',\n",
       " 'long_name': 'longitude coordinate',\n",
       " 'standard_name': 'longitude',\n",
       " 'bounds': 'lon_bnds'}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nessy.lon"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read variables\n",
    "\n",
    "There is no data because we haven't loaded any variable (lazy mode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'nox_no': {'data': None,\n",
       "  'dimensions': ('time', 'lat', 'lon'),\n",
       "  'dtype': dtype('float32'),\n",
       "  'units': 'm-2.kg.s-1',\n",
       "  'coordinates': 'lat lon',\n",
       "  'cell_measures': 'area: cell_area',\n",
       "  'grid_mapping': 'crs'}}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nessy.variables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we load the data for the variable NOx (NO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rank 000: Loading nox_no var (1/1)\n",
      "Rank 000: Loaded nox_no var ((1, 1, 1800, 3600))\n"
     ]
    }
   ],
   "source": [
    "nessy.keep_vars('nox_no')\n",
    "nessy.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'data': array([[[[0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          ...,\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.],\n",
       "          [0., 0., 0., ..., 0., 0., 0.]]]], dtype=float32),\n",
       " 'dimensions': ('time', 'lat', 'lon'),\n",
       " 'dtype': dtype('float32'),\n",
       " 'units': 'm-2.kg.s-1',\n",
       " 'coordinates': 'lat lon',\n",
       " 'cell_measures': 'area: cell_area',\n",
       " 'grid_mapping': 'crs'}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nessy.variables['nox_no']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "nessy.create_spatial_bounds()\n",
    "lon_bnds, lat_bnds = nessy.get_spatial_bounds_mesh_format()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 275300 characters]",
      "text/plain": [
       "<Figure size 1368x504 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(1, figsize=(19, 7))\n",
    "plt.pcolormesh(lon_bnds, lat_bnds, nessy.variables['nox_no']['data'].sum(axis=(0, 1)), cmap='jet', \n",
    "               norm=matplotlib.colors.LogNorm())\n",
    "countries = gpd.read_file('/gpfs/projects/bsc32/models/NES_tutorial_data/gadm_country_mask/gadm_country_ISO3166.shp')\n",
    "countries.plot(ax=ax, facecolor='none', edgecolor='black', linewidth=0.3)\n",
    "plt.colorbar(ax=ax, label='NOx expressed as NO (m-2·kg·s-1)')\n",
    "ax.margins(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Write dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Write NetCDF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rank 000: Creating regular_grid.nc\n",
      "Rank 000: NetCDF ready to write\n",
      "Rank 000: Dimensions done\n",
      "Rank 000: Cell measures done\n",
      "Rank 000: Writing nox_no var (1/1)\n",
      "Rank 000: Var nox_no created (1/1)\n",
      "Rank 000: Filling nox_no)\n",
      "Rank 000: Var nox_no data (1/1)\n",
      "Rank 000: Var nox_no completed (1/1)\n"
     ]
    }
   ],
   "source": [
    "nessy.to_netcdf('regular_grid.nc', info=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Write shapefile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/esarchive/scratch/avilanova/software/NES/nes/nc_projections/default_nes.py:3058: UserWarning: No vertical level has been specified. The first one will be selected.\n",
      "  warnings.warn(msg)\n",
      "/esarchive/scratch/avilanova/software/NES/nes/nc_projections/default_nes.py:3069: UserWarning: No time has been specified. The first one will be selected.\n",
      "  warnings.warn(msg)\n"
     ]
    }
   ],
   "source": [
    "nessy.to_shapefile('regular_grid_shp')"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
