{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "71e9082c",
   "metadata": {},
   "source": [
    "# Unsaturated Zone Flow\n",
    "This notebook demonstrates the use of the Unsaturated Zone Flow (UZF) package in nlmod. The UZF-package can be used to simulate important processes in the unsaturated zone. These processes create a delay before precipitation reaches the saturated groundwater. In dry periods the water may even have evaporated before that. This notebook shows the difference in calculated head between a model with regular recharge (RCH) and evapotranspiration (EVT) packages, compared to a model with the UZF-package.\n",
    "\n",
    "We create a 1d model, consisting of 1 row and 1 column with multiple layers, of a real location somewhere in the Netherlands. We use weather data from the KNMI as input for a transient simulation of 3 years with daily timetseps. This notebook can be used to vary the uzf-parameters, change the location (do not forget to alter the drain-elevation as well), or to play with the model timestep."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0610edfd",
   "metadata": {},
   "source": [
    "## Import packages and setup logging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbf2d8ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import packages\n",
    "import os\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import nlmod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "febb2e33",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up pretty logging and show package versions\n",
    "nlmod.util.get_color_logger()\n",
    "nlmod.show_versions()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9813a7da",
   "metadata": {},
   "source": [
    "## Generate a model Dataset\n",
    "We first set the model_name and model workspace, which we will use later to write the model files, and so we can determine the cache-directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cc92261",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"model17\"\n",
    "model_ws = \"17_unsaturated_zone_flow\"\n",
    "figdir, cachedir = nlmod.util.get_model_dirs(model_ws)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8de98855",
   "metadata": {},
   "source": [
    "We define a location with a corresponding drainage elevation. From the location we calculate an extent of 100 by 100 meter, and download REGIS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "341f7f9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = 37_400\n",
    "y = 412_600\n",
    "drn_elev = 4.0\n",
    "\n",
    "# round x and y to 100, so we will download only one cell in regis\n",
    "x = np.floor(x / 100) * 100\n",
    "y = np.floor(y / 100) * 100\n",
    "dx = dy = 100\n",
    "extent = [x, x + dx, y, y + dy]\n",
    "regis = nlmod.read.regis.download_regis(\n",
    "    extent, drop_layer_dim_from_top=False, cachename=\"regis.nc\", cachedir=cachedir\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76deaeb7",
   "metadata": {},
   "source": [
    "As the REGIS-data only contains one cell, we can visualize the properties of the layers in a pandas DataFrame. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbd30a5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "regis.sel(x=regis.x[0], y=regis.y[0]).to_pandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdb20bc9",
   "metadata": {},
   "source": [
    "As you can see, there are some NaN-values in the hydaulic conductivities (`kh` and `kv`). These will be filled when making a model Dataset, using fill-values and anisotropy values. See the info-messages after the commands below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71466503",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = nlmod.to_model_ds(\n",
    "    regis,\n",
    "    model_name=model_name,\n",
    "    model_ws=model_ws,\n",
    "    fill_value_kh=10.0,\n",
    "    fill_value_kv=10.0,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e3b93b3",
   "metadata": {},
   "source": [
    "We then add a time-dimension to our model Dataset and download knmi-data. We will have a calculation period of 3 year with daily timesteps, with a steady-state stress period with the weather of 2019 as a warmup period."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa5fb5c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "time = pd.date_range(\"2020\", \"2023\", freq=\"D\")\n",
    "ds = nlmod.time.set_ds_time(ds, start=\"2019\", time=time, steady_start=True)\n",
    "\n",
    "ds.update(\n",
    "    nlmod.read.knmi.get_recharge(\n",
    "        ds, method=\"separate\", cachename=\"recharge.nc\", cachedir=ds.cachedir\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80994109",
   "metadata": {},
   "source": [
    "## Generate a simulation (sim) and groundwater flow (gwf) object\n",
    "We generate a model using with all basic packages. We add a drainage level at 4.0 m NAP. As the top of our model is at 6.5 m NAP this will create an unsaturated zone of about 2.5 m."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b6b18d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create simulation\n",
    "sim = nlmod.sim.sim(ds)\n",
    "\n",
    "# create time discretisation\n",
    "_ = nlmod.sim.tdis(ds, sim)\n",
    "\n",
    "# create groundwater flow model\n",
    "gwf = nlmod.gwf.gwf(ds, sim, under_relaxation=True)\n",
    "\n",
    "# create ims\n",
    "_ = nlmod.sim.ims(sim)\n",
    "\n",
    "# Create discretization\n",
    "_ = nlmod.gwf.dis(ds, gwf)\n",
    "\n",
    "# create node property flow\n",
    "_ = nlmod.gwf.npf(ds, gwf, icelltype=1)\n",
    "\n",
    "# creat storage\n",
    "_ = nlmod.gwf.sto(ds, gwf, iconvert=1, sy=0.2, ss=1e-5)\n",
    "\n",
    "# Create the initial conditions package\n",
    "_ = nlmod.gwf.ic(ds, gwf, starting_head=1.0)\n",
    "\n",
    "# Create the output control package\n",
    "_ = nlmod.gwf.oc(ds, gwf)\n",
    "\n",
    "# set a drainage level with a resistance of 100 days in the layer that contains the drainage level\n",
    "cond = ds[\"area\"] / 100.0\n",
    "layer = nlmod.layers.get_layer_of_z(ds, drn_elev)\n",
    "_ = nlmod.gwf.drn(ds, gwf, elev=drn_elev, cond=cond, layer=layer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62cbe11e",
   "metadata": {},
   "source": [
    "## Run model with RCH and EVT packages\n",
    "We first run the model with the Recharge (RCH) and Evapotranspiration (EVT) packages, as a reference for the model with the UZF-package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9546d7fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create recharge package\n",
    "nlmod.gwf.rch(ds, gwf)\n",
    "\n",
    "# create evapotranspiration package\n",
    "nlmod.gwf.evt(ds, gwf);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3f051d0",
   "metadata": {},
   "source": [
    "We run this model, read the heads and get the groundwater level (the head in the highest active cell). We save the groundwater level to a variable called `gwl_rch_evt`, which we will plot later on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a0939fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run model\n",
    "nlmod.sim.write_and_run(sim, ds, silent=True)\n",
    "\n",
    "# get heads\n",
    "head_rch = nlmod.gwf.get_heads_da(ds)\n",
    "\n",
    "# calculate groundwater level\n",
    "gwl_rch_evt = nlmod.gwf.output.get_gwl_from_wet_cells(head_rch, botm=ds[\"botm\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34f1a8b6",
   "metadata": {},
   "source": [
    "## Run model with UZF package\n",
    "We then run the model again with the uzf-package. Before creating the uzf-package we remove the RCH- and EVT-packages.\n",
    "\n",
    "We choose some values for the residual water content, the saturated water content and the exponent used in the Brooks-Corey function. Other parameters are left to their defaults. The method `nlmod.gwf.uzf` will generate the UZF-package, using the variable `kv` from `ds` for the saturated vertical hydraulic conductivity, the variable `recharge` from `ds` for the infiltration rate and the variable `evaporation` from `ds` as the potential evapotranspiration rate.\n",
    "\n",
    "There can be multiple layers in the unsaturated zone, just like in the saturated zone. The method `nlmod.gwf.uzf` connects the unsaturated zone cells above each other.\n",
    "\n",
    "Because we want to plot the water content in the subsurface we will add some observations of the water content to the uzf-package. We do this by adding the optional parameter `obs_z` to `nlmod.gwf.uzf`. This will create the observations in the corresponding uzf-cells, at the requested z-values (in m NAP)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52c711d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "gwf.remove_package(\"RCH\")\n",
    "gwf.remove_package(\"EVT\")\n",
    "\n",
    "# create uzf package\n",
    "thtr = 0.1  # residual (irreducible) water content\n",
    "thts = 0.3  # saturated water content\n",
    "eps = 3.5  # exponent used in the Brooks-Corey function\n",
    "# add observations of the water concent every 0.2 m, from 1 meter below the drainage-elevation to the model top\n",
    "obs_z = np.arange(drn_elev - 1, ds[\"top\"].max(), 0.2)[::-1]\n",
    "_ = nlmod.gwf.uzf(\n",
    "    ds,\n",
    "    gwf,\n",
    "    thtr=thtr,\n",
    "    thts=thts,\n",
    "    thti=thtr,\n",
    "    eps=eps,\n",
    "    print_input=True,\n",
    "    print_flows=True,\n",
    "    nwavesets=100,  # Modflow 6 failed and advised to increase this value from the default value of 40\n",
    "    obs_z=obs_z,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7835138",
   "metadata": {},
   "source": [
    "We run the model again, and save the groundwater level to a variable called `gwl_uzf`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39f706a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run model\n",
    "nlmod.sim.write_and_run(sim, ds, silent=True)\n",
    "\n",
    "# get heads\n",
    "head_rch = nlmod.gwf.get_heads_da(ds)\n",
    "\n",
    "# calculate groundwater level\n",
    "gwl_uzf = nlmod.gwf.output.get_gwl_from_wet_cells(head_rch, botm=ds[\"botm\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6938848c",
   "metadata": {},
   "source": [
    "We read the water content from the observations, and use the name of the observations to determine the layer, row, columns and z-value of each observation, which we can use in our plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af7128fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%\n",
    "fname = os.path.join(ds.model_ws, f\"{ds.model_name}.uzf.obs.csv\")\n",
    "obs = pd.read_csv(fname, index_col=0)\n",
    "obs.index = pd.to_datetime(ds.time.start) + pd.to_timedelta(obs.index, \"D\")\n",
    "kind, lay, row, col, z = zip(*[x.split(\"_\") for x in obs.columns])\n",
    "lays = np.array([int(x) for x in lay])\n",
    "rows = np.array([int(x) for x in row])\n",
    "cols = np.array([int(x) for x in col])\n",
    "z = np.array([float(x) for x in z])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d09726b",
   "metadata": {},
   "source": [
    "## Compare models\n",
    "We then make a plot to compare the heads in the two simulations we performed, and plot the water content we calculated in the UZF-calculation, and added observations for. We plot the water content in one vertical cell of the model. Figure layout thanks to Martin Vonk!\n",
    "\n",
    "The figure shows that the effect of precipitation on the groundwater level is less in summer if we also take the effect of the unsaturated zone into account (using UZF). In dry periods precipitation never reaches the groundwater level, as evaporation takes place before that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f25f53b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "row = 0\n",
    "col = 0\n",
    "\n",
    "f, ax = plt.subplot_mosaic(\n",
    "    [[\"P\"], [\"WC\"], [\"H\"]], figsize=(10, 8), height_ratios=[2, 5, 3], sharex=True, layout=\"constrained\"\n",
    ")\n",
    "\n",
    "# top ax\n",
    "p = ds[\"recharge\"][:, row, col].to_pandas() * 1e3\n",
    "p.plot(ax=ax[\"P\"], label=\"Precipitation\", color=\"C9\")\n",
    "e = ds[\"evaporation\"][:, row, col].to_pandas() * 1e3\n",
    "e.plot(ax=ax[\"P\"], label=\"Evaporation\", color=\"C8\")\n",
    "ax[\"P\"].set_ylabel(\"[mm/d]\")\n",
    "ax[\"P\"].set_ylim(bottom=0)\n",
    "ax[\"P\"].legend(loc=(0, 1), ncol=2, frameon=False)\n",
    "ax[\"P\"].grid()\n",
    "\n",
    "# middle ax\n",
    "mask = (rows == row) & (cols == col)\n",
    "XY = np.meshgrid(obs.index, z[mask])\n",
    "theta = obs.loc[:, mask].transpose().values\n",
    "pcm = ax[\"WC\"].pcolormesh(XY[0], XY[1], theta, cmap=\"viridis_r\", vmin=thtr, vmax=thts)\n",
    "ax[\"WC\"].set_ylabel(\"z [m NAP]\")\n",
    "ax[\"WC\"].grid()\n",
    "# set xlim, as pcolormesh increases xlim a bit\n",
    "ax[\"WC\"].set_xlim(ds.time[0], ds.time[-1])\n",
    "nlmod.plot.colorbar_inside(pcm, ax=ax[\"WC\"], label=r\"Moisture Content $\\theta$ [-]\")\n",
    "\n",
    "# bottom ax\n",
    "s = gwl_rch_evt[:, row, col].to_pandas()\n",
    "ax[\"H\"].plot(s.index, s.values, color=\"C1\", linestyle=\"--\", label=\"GWL RCH+EVT\")\n",
    "s = gwl_uzf[:, row, col].to_pandas()\n",
    "ax[\"H\"].plot(s.index, s.values, color=\"C0\", linestyle=\"-\", label=\"GWL UZF\")\n",
    "ax[\"H\"].set_ylabel(\"z [m NAP]\")\n",
    "ax[\"H\"].legend(loc=(0, 0.98), ncol=2, frameon=False)\n",
    "ax[\"H\"].grid()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}