{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot Wflow model maps and forcing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**hydroMT** provides a simple interface to model schematization and forcing data from which we can make beautiful plots:\n",
    "- Raster layers are saved to the model `staticmaps` attribute as a  `xarray.Dataset`\n",
    "- Vector layers are saved to the model `staticgeoms` attribute as a `geopandas.GeoDataFrame`. Note that in case of Wflow these are not used by the model engine, but only for analysis and visualization purposes.\n",
    "- Forcing model layers are saved to model `forcing` attribute as a `dict` of `xarray.DataArray`\n",
    "\n",
    "We rely on the [cartopy](https://scitools.org.uk/cartopy/docs/latest/) package to plot maps. This packages provides a simple interface to plot geographic data and add background satellite imagery."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import xarray as xr\n",
    "import numpy as  np\n",
    "from os.path import join, dirname\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "from matplotlib import cm, colors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import hydromt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "root = 'wflow_piave_subbasin' \n",
    "mod = hydromt.WflowModel(root, mode='r')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot model schematization base maps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we plot the model basemaps (topography map with rivers, lakes, reservoirs, glaciers and gauges geometries). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot maps dependencies\n",
    "import matplotlib.patches as mpatches\n",
    "import cartopy.crs as ccrs\n",
    "import descartes # required to plot polygons\n",
    "import cartopy.io.img_tiles as cimgt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read and mask the model elevation\n",
    "da =mod.staticmaps['wflow_dem'].raster.mask_nodata()\n",
    "da.attrs.update(long_name='elevation', units='m')\n",
    "# read/derive river geometries\n",
    "gdf_riv = mod.rivers\n",
    "# read/derive model basin boundary\n",
    "gdf_bas = mod.basins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use('seaborn-whitegrid') # set nice style\n",
    "# we assume the model maps are in the geographic CRS EPSG:4326\n",
    "proj = ccrs.PlateCarree()\n",
    "# adjust zoomlevel and figure size to your basis size & aspect\n",
    "zoom_level = 10\n",
    "figsize=(10, 8)\n",
    "shaded= False # shaded elevation (looks nicer with more pixels (e.g.: larger basins))!\n",
    "\n",
    "\n",
    "# initialize image with geoaxes\n",
    "fig = plt.figure(figsize=figsize)\n",
    "ax = fig.add_subplot(projection=proj)\n",
    "extent = np.array(da.raster.box.buffer(0.02).total_bounds)[[0, 2, 1, 3]]\n",
    "ax.set_extent(extent, crs=proj)\n",
    "\n",
    "# add sat background image\n",
    "ax.add_image(cimgt.QuadtreeTiles(), zoom_level, alpha=0.5)\n",
    "\n",
    "## plot elevation\\\n",
    "# create nice colormap\n",
    "vmin, vmax = da.quantile([0.0, 0.98]).compute()\n",
    "c_dem = plt.cm.terrain(np.linspace(0.25, 1, 256))\n",
    "cmap = colors.LinearSegmentedColormap.from_list(\"dem\", c_dem)\n",
    "norm = colors.Normalize(vmin=vmin, vmax=vmax)\n",
    "kwargs = dict(cmap=cmap, norm=norm)\n",
    "# plot 'normal' elevation\n",
    "da.plot(transform=proj, ax=ax, zorder=1, cbar_kwargs=dict(aspect=30, shrink=.8), **kwargs)\n",
    "# plot elevation with shades\n",
    "if shaded:\n",
    "    ls = colors.LightSource(azdeg=315, altdeg=45)\n",
    "    dx, dy = da.raster.res\n",
    "    _rgb = ls.shade(\n",
    "        da.fillna(0).values,\n",
    "        norm=kwargs[\"norm\"],\n",
    "        cmap=kwargs[\"cmap\"],\n",
    "        blend_mode=\"soft\",\n",
    "        dx=dx,\n",
    "        dy=dy,\n",
    "        vert_exag=200,\n",
    "    )\n",
    "    rgb = xr.DataArray(\n",
    "        dims=(\"y\", \"x\", \"rgb\"), data=_rgb, coords=da.raster.coords\n",
    "    )\n",
    "    rgb = xr.where(np.isnan(da), np.nan, rgb)\n",
    "    rgb.plot.imshow(transform=proj, ax=ax, zorder=2)\n",
    "                \n",
    "# plot rivers with increasing width with stream order\n",
    "kwargs = dict()\n",
    "for strord in np.unique(gdf_riv['strord']):\n",
    "    if strord == np.unique(gdf_riv['strord']).max():\n",
    "        kwargs.update(label='river')\n",
    "    gdf_riv[gdf_riv['strord']==strord].plot(ax=ax, linewidth=strord/5, color='blue', zorder=3, **kwargs)\n",
    "# plot the basin boundary\n",
    "gdf_bas.boundary.plot(ax=ax, color='k', linewidth=0.3)\n",
    "# plot various vector layers if present\n",
    "if 'gauges' in mod.staticgeoms:\n",
    "    mod.staticgeoms['gauges'].plot(ax=ax, marker='d', markersize=25, facecolor='k', zorder=5, label='gauges')\n",
    "patches = [] # manual patches for legend, see https://github.com/geopandas/geopandas/issues/660\n",
    "if 'lakes' in mod.staticgeoms:\n",
    "    kwargs = dict(facecolor='lightblue', edgecolor='black', linewidth=1, label='lakes')\n",
    "    mod.staticgeoms['lakes'].plot(ax=ax, zorder=4, **kwargs)\n",
    "    patches.append(mpatches.Patch(**kwargs))\n",
    "if 'reservoirs' in mod.staticgeoms:\n",
    "    kwargs = dict(facecolor='blue', edgecolor='black', linewidth=1, label='reservoirs')\n",
    "    mod.staticgeoms['reservoirs'].plot(ax=ax, zorder=4, **kwargs)\n",
    "    patches.append(mpatches.Patch(**kwargs))\n",
    "if 'glaciers' in mod.staticgeoms:\n",
    "    kwargs = dict(facecolor='grey', edgecolor='grey', linewidth=1, label='glaciers')\n",
    "    mod.staticgeoms['glaciers'].plot(ax=ax, zorder=4, **kwargs)\n",
    "    patches.append(mpatches.Patch(**kwargs))\n",
    "    \n",
    "ax.xaxis.set_visible(True)\n",
    "ax.yaxis.set_visible(True)\n",
    "ax.set_ylabel(f\"latitude [degree north]\")\n",
    "ax.set_xlabel(f\"longitude [degree east]\")\n",
    "_ = ax.set_title(f\"wflow base map\")\n",
    "legend = ax.legend(\n",
    "    handles=[*ax.get_legend_handles_labels()[0], *patches],\n",
    "    title=\"Legend\",\n",
    "    loc='lower right',\n",
    "    frameon=True,\n",
    "    framealpha=0.7,\n",
    "    edgecolor='k',\n",
    "    facecolor='white'\n",
    ")\n",
    "\n",
    "# save figure\n",
    "# NOTE create figs folder in model root if it does not exist\n",
    "# fn_out = join(mod.root, \"figs\", \"basemap.png\")\n",
    "# plt.savefig(fn_out, dpi=225, bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot model forcing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we plot the model *basin average* forcing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read wflow forcing; mask region outside the basin and compute the basin average\n",
    "# NOTE: only very limited forcing data is available from the artifacts\n",
    "ds_forcing = xr.merge(mod.forcing.values()).where(mod.staticmaps['wflow_subcatch']>0) \n",
    "ds_forcing = ds_forcing.mean(dim=[ds_forcing.raster.x_dim, ds_forcing.raster.y_dim])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot axes labels\n",
    "_ATTRS = {\n",
    "    \"precip\": {\"standard_name\": \"precipitation\", \"unit\": \"mm.day-1\", \"color\": 'darkblue'},\n",
    "    \"pet\": {\"standard_name\": \"potential evapotranspiration\", \"unit\": \"mm.day-1\", \"color\": 'purple'},\n",
    "    \"temp\": {\"standard_name\": \"temperature\", \"unit\": \"degree C\", \"color\": 'orange'},\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use('seaborn') # set nice style\n",
    "\n",
    "n = len(ds_forcing.data_vars)\n",
    "kwargs0 = dict(sharex=True, figsize=(6, n * 3))\n",
    "\n",
    "fig, axes = plt.subplots(n, 1, **kwargs0)\n",
    "axes = [axes] if n == 1 else axes\n",
    "for i, name in enumerate(ds_forcing.data_vars):\n",
    "    df = ds_forcing[name].squeeze().to_series()\n",
    "    attrs = _ATTRS[name]\n",
    "    longname = attrs.get(\"standard_name\", \"\")\n",
    "    unit = attrs.get(\"unit\", \"\")\n",
    "    if name == 'precip':\n",
    "        axes[i].bar(df.index, df.values, facecolor=attrs['color'])\n",
    "    else:\n",
    "        df.plot.line(ax=axes[i], x=\"time\", color=attrs['color'])\n",
    "    axes[i].set_title(longname)\n",
    "    axes[i].set_ylabel(f'{longname}\\n[{unit}]')\n",
    "\n",
    "# save figure\n",
    "# fn_out = join(mod.root, \"figs\", \"forcing.png\")\n",
    "# plt.savefig(fn_out, dpi=225, bbox_inches=\"tight\")"
   ]
  }
 ],
 "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.9.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}