{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to make spatial joins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nes\n",
    "import geopandas as gpd\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from copy import copy\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "countries_path = '/gpfs/projects/bsc32/models/NES_tutorial_data/gadm_country_mask/gadm_country_ISO3166.shp'\n",
    "countries = gpd.read_file(countries_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Read/Create grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rotated grid from MONARCH\n",
    "projection='rotated'\n",
    "centre_lat = 51\n",
    "centre_lon = 10\n",
    "west_boundary = -35\n",
    "south_boundary = -27\n",
    "inc_rlat = 2.\n",
    "inc_rlon = 2.\n",
    "\n",
    "nessy = nes.create_nes(comm=None, info=False, projection=projection,\n",
    "                       centre_lat=centre_lat, centre_lon=centre_lon,\n",
    "                       west_boundary=west_boundary, south_boundary=south_boundary,\n",
    "                       inc_rlat=inc_rlat, inc_rlon=inc_rlon)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "shp = nessy.create_shapefile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 330692 characters]",
      "text/plain": [
       "<Figure size 1900x700 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(1, figsize=(19, 7))\n",
    "shp.plot(ax=ax, facecolor=\"none\", edgecolor=\"blue\")\n",
    "countries.plot(ax=ax, facecolor=\"none\", edgecolor='black', linewidth=0.3)\n",
    "\n",
    "ax.margins(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Spatial join"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1. Centroid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/scratch/tmp/9400020/ipykernel_2995633/2099596389.py:4: UserWarning: Geometry is in a geographic CRS. Results from 'centroid' are likely incorrect. Use 'GeoSeries.to_crs()' to re-project geometries to a projected CRS before this operation.\n",
      "\n",
      "  centroid_shp['geometry'] = centroid_shp['geometry'].centroid\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 335248 characters]",
      "text/plain": [
       "<Figure size 1900x700 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot\n",
    "\n",
    "centroid_shp = copy(nessy.shapefile)\n",
    "centroid_shp['geometry'] = centroid_shp['geometry'].centroid\n",
    "\n",
    "fig, ax = plt.subplots(1, figsize=(19, 7))\n",
    "shp.plot(ax=ax, facecolor=\"none\", edgecolor=\"blue\")\n",
    "countries.plot(ax=ax, facecolor=\"none\", edgecolor='black', linewidth=0.3)\n",
    "\n",
    "centroid_shp.plot(ax=ax, edgecolor='none', facecolor=\"green\", marker=\".\", markersize=15)\n",
    "\n",
    "ax.margins(0)\n",
    "\n",
    "zoom={'x_min': -15, 'x_max': 55, 'y_min': 25, 'y_max':65}\n",
    "plt.axis([zoom['x_min'], zoom['x_max'], zoom['y_min'], zoom['y_max']])\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.24 s, sys: 128 ms, total: 2.37 s\n",
      "Wall time: 2.57 s\n"
     ]
    }
   ],
   "source": [
    "%time nessy.spatial_join(countries_path, method='centroid', var_list='ISO')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot results\n",
    "\n",
    "fig, ax = plt.subplots(1, figsize=(19, 7))\n",
    "\n",
    "nessy.shapefile.plot(ax=ax, column='ISO', cmap='gist_rainbow', legend=True, \n",
    "                     legend_kwds={'ncol': 3, 'loc': 'center left'})\n",
    "\n",
    "countries.plot(ax=ax, facecolor=\"none\", edgecolor='black', linewidth=0.3)\n",
    "shp.plot(ax=ax, facecolor=\"none\", edgecolor=\"grey\", linewidth=0.3)\n",
    "centroid_shp.plot(ax=ax, edgecolor='none', facecolor=\"green\", marker=\".\", markersize=15)\n",
    "\n",
    "leg = ax.get_legend()\n",
    "leg.set_bbox_to_anchor((1., 0.5))\n",
    "\n",
    "ax.margins(0)\n",
    "\n",
    "zoom={'x_min': -15, 'x_max': 55, 'y_min': 25, 'y_max':65}\n",
    "plt.axis([zoom['x_min'], zoom['x_max'], zoom['y_min'], zoom['y_max']])\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2. Intersection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nessy = nes.create_nes(comm=None, info=False, projection=projection,\n",
    "                       centre_lat=centre_lat, centre_lon=centre_lon,\n",
    "                       west_boundary=west_boundary, south_boundary=south_boundary,\n",
    "                       inc_rlat=inc_rlat, inc_rlon=inc_rlon)\n",
    "\n",
    "%time nessy.spatial_join(countries_path, method='intersection', var_list='ISO')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot results\n",
    "\n",
    "fig, ax = plt.subplots(1, figsize=(19, 7))\n",
    "\n",
    "nessy.shapefile.plot(ax=ax, column='ISO', cmap='gist_rainbow', legend=True, \n",
    "                     legend_kwds={'ncol': 3, 'loc': 'center left'})\n",
    "\n",
    "countries.plot(ax=ax, facecolor=\"none\", edgecolor='black', linewidth=0.3)\n",
    "shp.plot(ax=ax, facecolor=\"none\", edgecolor=\"grey\", linewidth=0.3)\n",
    "\n",
    "leg = ax.get_legend()\n",
    "leg.set_bbox_to_anchor((1., 0.5))\n",
    "\n",
    "ax.margins(0)\n",
    "\n",
    "zoom={'x_min': -15, 'x_max': 55, 'y_min': 25, 'y_max':65}\n",
    "plt.axis([zoom['x_min'], zoom['x_max'], zoom['y_min'], zoom['y_max']])\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3. Nearest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nessy = nes.create_nes(comm=None, info=False, projection=projection,\n",
    "                       centre_lat=centre_lat, centre_lon=centre_lon,\n",
    "                       west_boundary=west_boundary, south_boundary=south_boundary,\n",
    "                       inc_rlat=inc_rlat, inc_rlon=inc_rlon)\n",
    "\n",
    "%time nessy.spatial_join(countries_path, method='nearest', var_list='ISO')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot results\n",
    "\n",
    "fig, ax = plt.subplots(1, figsize=(19, 7))\n",
    "\n",
    "nessy.shapefile.plot(ax=ax, column='ISO', cmap='gist_rainbow', legend=True, \n",
    "                     legend_kwds={'ncol': 3, 'loc': 'center left'})\n",
    "\n",
    "countries.plot(ax=ax, facecolor=\"none\", edgecolor='black', linewidth=0.3)\n",
    "shp.plot(ax=ax, facecolor=\"none\", edgecolor=\"grey\", linewidth=0.3)\n",
    "\n",
    "leg = ax.get_legend()\n",
    "leg.set_bbox_to_anchor((1., 0.5))\n",
    "\n",
    "ax.margins(0)\n",
    "\n",
    "zoom={'x_min': -15, 'x_max': 55, 'y_min': 25, 'y_max':65}\n",
    "plt.axis([zoom['x_min'], zoom['x_max'], zoom['y_min'], zoom['y_max']])\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Write string data as NetCDF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nessy.variables['ISO'] = {'data': nessy.shapefile['ISO'].values.reshape([1, 1, nessy.lat['data'].shape[0], nessy.lon['data'].shape[-1]]),\n",
    "                          'dtype': str,\n",
    "                         }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the string lenght\n",
    "nessy.set_strlen(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nessy.to_netcdf('spatial_join.nc')"
   ]
  }
 ],
 "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
}
