{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Update a wflow model: water demand"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With HydroMT, you can easily read your model and update one or several components of your model using the **update** function of the command line interface (CLI). Here are the steps and some examples on how to **add water demand information** to your wflow_sbm model.\n",
    "\n",
    "All lines in this notebook which starts with ! are executed from the command line. Within the notebook environment the logging messages are shown after completion. You can also copy these lines and paste them in your shell to get more feedback.\n",
    "\n",
    "You can update a wflow_sbm model to include the simulation of water demand and allocation. This can be important in basins where there is substantial human influence, and/or you might want to simulate what will happen if (more) water is consumed. We identify four different sources of water demand:\n",
    "- domestic\n",
    "- industry\n",
    "- livestock\n",
    "- irrigation\n",
    "\n",
    "These four different items are taken into account when adding the required water demand data to your wflow model. Besides computing the water demand, we also need information on the regions where water is allocated. \n",
    "\n",
    "Let's open the example configuration file (**wflow_update_water_demand.yml**) from the model repository [examples folder] and have a look at the main setup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fn_config = \"wflow_update_water_demand.yml\"\n",
    "with open(fn_config, \"r\") as f:\n",
    "    txt = f.read()\n",
    "print(txt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, you see we will be adding the required information to simulate water demand and allocation in wflow, by running three different setup functions:\n",
    "- **setup_allocation_areas**: Adds a map that defines the different regions that will be used to allocate water. By default this is a mix between administrative boundaries and catchment boundaries\n",
    "- **setup_allocation_surfacewaterfrac**: prepare the fraction of surface water used for allocation (can be reduced if groundwater or non conventional water sources are also present in the basin).\n",
    "- **setup_lulcmaps_with_paddy** (or **setup_lulcmaps**): update the landuse to include new parameters (crop coefficient kc and soil water pressure heads h) and add the paddies (rice fields). To allow for water to pool on the surface of the rice fields, soil parameters will also be updated to include an additional thin layer with limited vertical conductivity.\n",
    "- **setup_domestic_demand**: Add domestic water demands (gross and net) from gridded data and downscaled using high resolution population map.\n",
    "- **setup_other_demand**: Adds maps to the wflow schematization that describe how much water is demanded (gross and net amounts) by different sources: domestic (dom), industry (ind), and livestock (lsk). In our case, as we downscale domestic with population, we will here add industry and livestock.\n",
    "- **setup_irrigation**: Adds information about where and when irrigation takes place for paddy and nonpaddy crops.\n",
    "\n",
    "More information on these steps will be given below (after we have updated the model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Update the wflow model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we build a wflow_sbm model for the full Piave basin, as most water demand activities occur in the most downstream locations of the catchment. Please note that this is a basin model (until the full outlet at the ocean), as opposed to the sub-basin models (until a certain location) that are being created in the other examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: copy this line (without !) to your shell for more direct feedback\n",
    "! hydromt build wflow \"./wflow_piave_water_demand\" -r \"{'basin': [12.2051, 45.8331], 'bounds': [11.70, 45.35, 12.95, 46.70]}\" -i wflow_build.yml -d artifact_data -vv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we will update the model using the same configuration file that was shown earlier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! hydromt update wflow wflow_piave_water_demand -i wflow_update_water_demand.yml -d artifact_data -d \"https://github.com/Deltares/hydromt_wflow/releases/download/v0.5.0/wflow_artifacts.yml\" -d ./data/demand/data_catalog.yml -v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NOTE: The second catalog wflow_artifacts.yml points to additional data that are not available in the hydromt artifact_data catalog such as extracts of the glcnmo landuse or the pcrglobwb gridded demands."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Looking at the added layers\n",
    "\n",
    "To understand what layers have been added, we'll plot the new layers below. To do the plotting, we first have to import the required python libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from hydromt_wflow import WflowModel\n",
    "\n",
    "mod = WflowModel(\n",
    "    root=\"wflow_piave_water_demand\",\n",
    "    mode=\"r+\",\n",
    "    data_libs=[\n",
    "        \"artifact_data\",\n",
    "        \"https://github.com/Deltares/hydromt_wflow/releases/download/v0.5.0/wflow_artifacts.yml\",\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Landuse and rice fields\n",
    "\n",
    "Additional new parameters have been added to the model for landuse such as:\n",
    "- **crop_factor**: crop factor map that is used to convert the reference PET to the crop specific evaporation\n",
    "- **h** values: soil water pressure heads at which root water uptake is reduced (Feddes) [cm]. They are different for paddy and non paddy landuse types.\n",
    "\n",
    "For paddies, new maps have been added and soil parameters updated to allow water to pool at the surface. These are the **c** maps and a new **kvfrac** map which is a factor used to multiply the vertical conductivity (`ksatver`) in order to allow for a layer with very low conductivity (for the paddy fields)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 3, sharex=True, sharey=True, figsize=(10, 5))\n",
    "\n",
    "axes[0].set_title(\"Paddy fields [-]\")\n",
    "axes[1].set_title(\"Crop factor [-]\")\n",
    "axes[2].set_title(\"Soil water pressure head h1 [cm]\")\n",
    "\n",
    "rice = mod.grid[\"wflow_landuse\"].where(\n",
    "    mod.grid[\"wflow_landuse\"] == 12, mod.grid[\"wflow_landuse\"].raster.nodata\n",
    ")\n",
    "rice.raster.mask_nodata().plot(ax=axes[0], add_labels=False)\n",
    "mod.grid[\"kc\"].raster.mask_nodata().plot(ax=axes[1], add_labels=False)\n",
    "mod.grid[\"h1\"].raster.mask_nodata().plot(ax=axes[2], add_labels=False)\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A `kvfrac` map is added that adds a layer with very low vertical hydraulic conductivity in the cells with paddy irrigation. The target conductivity value can be set for each layer using the `target_conductivity` parameter, and sets a value for each layer in the wflow model (linked to the `wflow_thicknesslayer` parameter). See the figures below for an indication of these maps and their effect on the vertical conductivity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 3, figsize=(11, 5))\n",
    "\n",
    "ax1, ax2, ax3 = axes.flatten()\n",
    "\n",
    "ax1.set_title(\"kvfrac for layer=0\")\n",
    "ax2.set_title(\"kvfrac for layer=2\")\n",
    "ax3.set_title(\"Effect on kv\")\n",
    "\n",
    "mod.grid[\"kvfrac\"].sel(layer=0).raster.mask_nodata().plot(\n",
    "    ax=ax1, add_labels=False, vmin=0, vmax=1\n",
    ")\n",
    "mod.grid[\"kvfrac\"].sel(layer=2).raster.mask_nodata().plot(\n",
    "    ax=ax2, add_labels=False, vmin=0, vmax=1\n",
    ")\n",
    "\n",
    "# Take the cumulative sum to get to the cumulative layers\n",
    "layers = np.cumsum(mod.get_config(\"model.thicknesslayers\"))\n",
    "layers = np.append(layers, 2000)\n",
    "# Add layers to figure\n",
    "for layer in layers:\n",
    "    ax3.axhline(y=layer, c=\"0.8\")\n",
    "\n",
    "# Position of a paddy pixel\n",
    "lat = 45.63\n",
    "lon = 12.66\n",
    "\n",
    "# Read required layers\n",
    "kv_0 = mod.grid[\"KsatVer\"].sel(latitude=lat, longitude=lon, method=\"nearest\")\n",
    "kvfrac = mod.grid[\"kvfrac\"].sel(latitude=lat, longitude=lon, method=\"nearest\")\n",
    "f = mod.grid[\"f\"].sel(latitude=lat, longitude=lon, method=\"nearest\")\n",
    "\n",
    "# Compute original kv values (without kvfrac)\n",
    "depths = np.arange(0, 2000)\n",
    "original = kv_0.values * np.exp(-f.values * depths)\n",
    "\n",
    "# Compute new kv values\n",
    "corrected = original.copy()\n",
    "idxs = np.where(kvfrac.values != 1)[0]\n",
    "for idx in idxs:\n",
    "    start_depth = layers[idx - 1]\n",
    "    end_depth = layers[idx]\n",
    "    corrected[start_depth:end_depth] *= kvfrac.values[idx]\n",
    "\n",
    "ax3.plot(original, depths, label=\"Without kvfrac\")\n",
    "ax3.plot(corrected, depths, label=\"With kvfrac\")\n",
    "\n",
    "# Flip y-axis for easier understanding of depth profile\n",
    "ax3.set_ylim(2000, 0)\n",
    "ax3.set_xlim(0, ax3.get_xlim()[1])\n",
    "ax3.set_xlabel(\"Vertical conductivity [mm day$^{-1}$]\")\n",
    "ax3.set_ylabel(\"Depth below surface [mm]\")\n",
    "ax3.legend(loc=\"lower right\")\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Industry, livestock and domestic demand\n",
    "\n",
    "For the non_irrigation related demand, we assume those to be already prepared datasets with gross and net (consumption) demands. In this example, we rely on data from the PCR-GLOBWB model. See the images below for an explanation of the data and added layers to the wflow model configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "month = 6\n",
    "\n",
    "fig, axes = plt.subplots(2, 3, sharex=True, sharey=True, figsize=(10, 8))\n",
    "\n",
    "fig.suptitle(f\"Demand values for month {month}\")\n",
    "\n",
    "axes[0][0].set_title(\"Industry [mm/day]\")\n",
    "axes[0][1].set_title(\"Livestock [mm/day]\")\n",
    "axes[0][2].set_title(\"Domestic [mm/day]\")\n",
    "axes[0][0].set_ylabel(\"Gross\")\n",
    "axes[1][0].set_ylabel(\"Net\")\n",
    "\n",
    "# Extracting the min-max ranges for consistent colorbars\n",
    "ind_min = mod.grid[\"industry_gross\"].min()\n",
    "ind_max = mod.grid[\"industry_gross\"].max()\n",
    "lsk_min = mod.grid[\"livestock_gross\"].min()\n",
    "lsk_max = mod.grid[\"livestock_gross\"].max()\n",
    "dom_min = mod.grid[\"domestic_gross\"].min()\n",
    "dom_max = mod.grid[\"domestic_gross\"].max()\n",
    "\n",
    "# Plot industry\n",
    "mod.grid[\"industry_gross\"].sel(time=month).plot(\n",
    "    ax=axes[0][0], add_labels=False, vmin=ind_min, vmax=ind_max\n",
    ")\n",
    "mod.grid[\"industry_net\"].sel(time=month).plot(\n",
    "    ax=axes[1][0], add_labels=False, vmin=ind_min, vmax=ind_max\n",
    ")\n",
    "# Plot livestock\n",
    "mod.grid[\"livestock_gross\"].sel(time=month).plot(\n",
    "    ax=axes[0][1], add_labels=False, vmin=lsk_min, vmax=lsk_max\n",
    ")\n",
    "mod.grid[\"livestock_net\"].sel(time=month).plot(\n",
    "    ax=axes[1][1], add_labels=False, vmin=lsk_min, vmax=lsk_max\n",
    ")\n",
    "# Plot domestic (adjusted the max range to improve plotting)\n",
    "mod.grid[\"domestic_net\"].sel(time=month).plot(\n",
    "    ax=axes[0][2], add_labels=False, vmin=dom_min, vmax=dom_max * 0.3\n",
    ")\n",
    "mod.grid[\"domestic_net\"].sel(time=month).plot(\n",
    "    ax=axes[1][2], add_labels=False, vmin=dom_min, vmax=dom_max * 0.3\n",
    ")\n",
    "\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the figure above, we see a slight difference between the gross and net demand for the industry sector. This means that this sector consumes part of the demanded water, but returns a portion of the water back. On the other hand, the gross and net demands for both livestock and domestic are roughly the same, meaning that the majority of the water will be consumed.\n",
    "\n",
    "For the industry and livestock sectors, we see the relatively low-resolution of the original data. This data is retrieved from the PCR-GLOBWB model, which was provided at a resolution of 0.5x0.5 degrees (roughly 40x40km). The same holds for the domestic demands data, but it was downscaled using population density in the `setup_non_irrigation` workflow (using the dataset provided in `population_fn`, which is an optional step). The population data is used to identify densely populated regions, and downscales this to the cities/villages where the people actually live. To give an impression, see the figure below for an insight into the population data. Please note that this is the data at the original resolution of ~100x100m. The data at the wflow model resolution is provided in the staticmaps (see `mod.grid[\"Population_scaled\"]`).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read original data and slice to model domain\n",
    "pop_ds = mod.data_catalog.get_rasterdataset(\n",
    "    \"worldpop_2020_constrained\", geom=mod.geoms[\"basins\"]\n",
    ")\n",
    "# Get mask of catchment\n",
    "pop_ds_mask = pop_ds.raster.geometry_mask(mod.geoms[\"basins\"])\n",
    "pop_ds = pop_ds.raster.mask_nodata().where(pop_ds_mask)\n",
    "\n",
    "# Plot data\n",
    "fig, ax = plt.subplots(1, 1, figsize=(6, 6))\n",
    "# Add original population data to figure\n",
    "pop_ds.plot(\n",
    "    ax=ax,\n",
    "    vmax=20,\n",
    "    add_labels=False,\n",
    "    cbar_kwargs={\"label\": \"Population per pixel (~100x100m)\"},\n",
    ")\n",
    "# Add basin geometry to plot\n",
    "mod.geoms[\"basins\"].plot(ax=ax, facecolor=\"none\", edgecolor=\"red\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Irrigation maps\n",
    "\n",
    "For the `setup_irrigation` workflow, a number of maps have been added to the model:\n",
    "- **paddy_irrigation_areas**: mask (0/1) of the cells that are considered to be paddy/rice fields\n",
    "- **nonpaddy_irrigation_areas**: mask (0/1) of the cells that are irrigated, but not paddy/rice fields\n",
    "- **nonpaddy_irrigation_trigger**: trigger (0/1) that indicates whether we expect irrigation to occur (based on the LAI, to identify the growing season)\n",
    "- **paddy_irrigation_trigger**: trigger (0/1) that indicates whether we expect irrigation to occur on paddies (based on the LAI, to identify the growing season)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(10, 5))\n",
    "\n",
    "axes[0].set_title(\"Paddy irrigation [-]\")\n",
    "axes[1].set_title(\"Non-paddy irrigation [-]\")\n",
    "\n",
    "mod.grid[\"paddy_irrigation_areas\"].raster.mask_nodata().plot(\n",
    "    ax=axes[0], add_labels=False\n",
    ")\n",
    "mod.grid[\"nonpaddy_irrigation_areas\"].raster.mask_nodata().plot(\n",
    "    ax=axes[1], add_labels=False\n",
    ")\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this figure, we see the cells with paddy (rice fields) and non-paddy irrigated areas. Cells are identified as such when they exceed a fractional threshold of the cell (set by `area_threshold`). A cell can only be either paddy-irrigated, non-paddy-irrigated, or rain-fed. In the third panel, we see the crop factor associated for all of the model pixels.\n",
    "\n",
    "In the figures below, we show the irrigation trigger maps for three months: before the growing season, during the growing season, and in the closing stages of the growing season. A map with a mask (0/1) contains information whether irrigation is allowed to occur during this month."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 3, sharex=True, sharey=True, figsize=(10, 5))\n",
    "\n",
    "fig.suptitle(\"Irrigation trigger [-]\")\n",
    "\n",
    "axes[0].set_title(\"February\")\n",
    "axes[1].set_title(\"July\")\n",
    "axes[2].set_title(\"October\")\n",
    "\n",
    "mod.grid[\"nonpaddy_irrigation_trigger\"].sel(time=2).raster.mask_nodata().plot(\n",
    "    ax=axes[0], add_labels=False, vmin=0, vmax=1\n",
    ")\n",
    "mod.grid[\"nonpaddy_irrigation_trigger\"].sel(time=7).raster.mask_nodata().plot(\n",
    "    ax=axes[1], add_labels=False, vmin=0, vmax=1\n",
    ")\n",
    "mod.grid[\"nonpaddy_irrigation_trigger\"].sel(time=10).raster.mask_nodata().plot(\n",
    "    ax=axes[2], add_labels=False, vmin=0, vmax=1\n",
    ")\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Water allocation regions\n",
    "\n",
    "To define regions where water can be shared and allocated, a merge between catchment and water areas or administrative boundaries is computed. These result in regions where wflow can allocate available water. No water allocation is supported between these regions. To give an impression on how these regions look like, see the following figures for an example.\n",
    "\n",
    "Note: Water areas or regions are generally defined by sub-river-basins within a Country. In order to mimick reality, it is advisable to avoid cross-Country-border abstractions. Whenever information is available, it is strongly recommended to align the water regions with the actual areas managed by water management authorities, such as regional water boards."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read original data and slice to model domain\n",
    "admin = mod.data_catalog.get_geodataframe(\"gadm_level2\", geom=mod.geoms[\"basins\"])\n",
    "\n",
    "fig, ax = plt.subplots(1)\n",
    "\n",
    "ax.set_title(\"Allocation areas\")\n",
    "\n",
    "mod.grid[\"allocation_areas\"].raster.mask_nodata().plot(ax=ax, add_labels=False)\n",
    "admin.plot(ax=ax, facecolor=\"none\", edgecolor=\"red\")\n",
    "mod.geoms[\"rivers\"].plot(ax=ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When merging the wflow basins with the water regions, some small subbasins can be created that do not contain river cells. These small basins will be merged to larger basins. When merging, you can decide if you prefer to merge with the nearest downstream basin, or with any basin in the same water region that does contain river using the ``priotity_basins`` argument. In the previous map, we gave priority to the basins, here is the results if priority is given to water regions instead:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create allocations areas\n",
    "mod.setup_allocation_areas(\n",
    "    waterareas_fn=\"gadm_level2\",\n",
    "    priority_basins=False,\n",
    ")\n",
    "\n",
    "fig, ax = plt.subplots(1)\n",
    "\n",
    "ax.set_title(\"Allocation areas\")\n",
    "\n",
    "mod.grid[\"allocation_areas\"].raster.mask_nodata().plot(\n",
    "    ax=ax, add_labels=False, cmap=\"viridis_r\"\n",
    ")\n",
    "admin.plot(ax=ax, facecolor=\"none\", edgecolor=\"red\")\n",
    "mod.geoms[\"rivers\"].plot(ax=ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we see a couple of distinct administrative boundaries (black lines), some of which already follow the catchment boundaries. When the catchment crosses an administrative boundary, that region receives a different but unique identifier. Note that we are using level 2 boundaries here, which might not be the most realistic for a region of this size."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Surface water frac used for water allocation\n",
    "\n",
    "By default, Wflow will allocate all water demands with water from the surface water (frac_sw_used=1). However, if in a certain areas, groundwater or other non conventional sources can be used, the frac_sw_used of each allocation area can be reduced and prepared using the **setup_allocation_surfacewaterfrac** method.\n",
    "\n",
    "Here we used global data from GLOFAS (Lisflood) for the fraction of grounwater used, presence of groundwater bodies and non conventional sources (0 is Piave). The water allocations are the ones we prepared in the previous step in order to match better the wflow model basin delineation.\n",
    "\n",
    "Let's have a look at the resulting map:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1)\n",
    "\n",
    "ax.set_title(\"Allocation surface water frac used\")\n",
    "\n",
    "mod.grid[\"frac_sw_used\"].raster.mask_nodata().plot(ax=ax, add_labels=False)\n",
    "mod.geoms[\"basins\"].plot(ax=ax, facecolor=\"none\", edgecolor=\"black\")\n",
    "mod.geoms[\"rivers\"].plot(ax=ax)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hydromt-wflow",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}