{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Stromingen example: keeping scripts simple\n",
    "\n",
    "This example is based on the essay _\"Open source grondwatermodellering met\n",
    "MODFLOW 6\"_ that was published in Stromingen (Calje et al., 2022).\n",
    "\n",
    "This example strives to achieve the simplicity of the example psuedo script\n",
    "that was shown in Figure 5 in the article. Some things require a bit more code\n",
    "than in the example, but not by much! Also some data is not yet publicly\n",
    "accessible, i.e. well data, so that has also not yet been implemented in this\n",
    "example. We also changed the extent to build a smaller model (because of\n",
    "computation time)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import Python packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import flopy as fp\n",
    "import geopandas as gpd\n",
    "from pandas import date_range\n",
    "\n",
    "import nlmod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlmod.util.get_color_logger(\"INFO\")\n",
    "nlmod.show_versions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define spatial and temporal properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "extent = [116_500, 120_000, 439_000, 442_000]\n",
    "tmin = \"2010-01-01\"\n",
    "tmax = \"2020-01-01\"\n",
    "freq = \"14D\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# where in the world are we?\n",
    "nlmod.plot.get_map(extent, background=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get data for the current extent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define a model workspace for caching downloaded data\n",
    "model_ws = \"14_stromingen_example\"\n",
    "figdir, cachedir = nlmod.util.get_model_dirs(model_ws)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "layer_model = nlmod.read.regis.get_combined_layer_models(\n",
    "    extent,\n",
    "    use_regis=True,\n",
    "    use_geotop=False,\n",
    "    cachedir=cachedir,\n",
    "    cachename=\"layer_model.nc\",\n",
    ")\n",
    "\n",
    "# wells = nlmod.read.get_wells(extent)  # no well data is available just yet\n",
    "\n",
    "# surface water features and levels\n",
    "fname_bgt = os.path.join(\"..\", \"data_sources\", \"02_surface_water\", \"cache\", \"bgt.gpkg\")\n",
    "if not os.path.isfile(fname_bgt):\n",
    "    raise (\n",
    "        Exception(\n",
    "            f\"{fname_bgt} not found. Please run notebook 02_surface_water.ipynb first\"\n",
    "        )\n",
    "    )\n",
    "sw = gpd.read_file(fname_bgt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate grid with a refinement zone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a model dataset\n",
    "ds = nlmod.to_model_ds(layer_model, \"stromingen\", model_ws=model_ws)\n",
    "\n",
    "# refine model dataset (supply a list of xy-coordinates)\n",
    "xy = [\n",
    "    [\n",
    "        [\n",
    "            (117_500, 439_500),\n",
    "            (117_500, 440_000),\n",
    "            (118_000, 440_000),\n",
    "            (118_000, 439_500),\n",
    "            (117_500, 439_500),\n",
    "        ]\n",
    "    ]\n",
    "]\n",
    "\n",
    "refinement = [(xy, \"polygon\", 1)]\n",
    "ds = nlmod.grid.refine(ds, refinement_features=refinement)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate a model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set model time settings\n",
    "t = date_range(tmin, tmax, freq=freq)\n",
    "ds = nlmod.time.set_ds_time(ds, start=3652, time=t, steady_start=True)\n",
    "\n",
    "# build the modflow6 gwf model\n",
    "gwf = nlmod.gwf.ds_to_gwf(ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add recharge to the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download knmi recharge data\n",
    "knmi_ds = nlmod.read.knmi.get_recharge(ds)\n",
    "\n",
    "# update model dataset\n",
    "ds.update(knmi_ds)\n",
    "\n",
    "# create recharge package\n",
    "rch = nlmod.gwf.rch(ds, gwf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add surface water to the model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set stage for model to mean of summer and winter levels\n",
    "sw[\"stage\"] = sw[[\"winter_stage\", \"summer_stage\"]].mean(axis=1)\n",
    "\n",
    "# use a water depth of 0.5 meter\n",
    "sw[\"rbot\"] = sw[\"stage\"] - 0.5\n",
    "\n",
    "# set the stage of the Lek river to 0.0 m NAP and the botm to -3 m NAP\n",
    "mask = sw[\"bronhouder\"] == \"L0002\"\n",
    "sw.loc[mask, \"stage\"] = 0.0\n",
    "sw.loc[mask, \"rbot\"] = -3.0\n",
    "\n",
    "# we need to mask out the NaNs\n",
    "sw.drop(sw.index[sw[\"rbot\"].isna()], inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# intersect surface water with grid\n",
    "sfw_grid = nlmod.grid.gdf_to_grid(sw, gwf)\n",
    "\n",
    "# add bed resistance to calculate conductance\n",
    "bed_resistance = 1.0  # days\n",
    "sfw_grid[\"cond\"] = sfw_grid.area / bed_resistance\n",
    "sfw_grid.set_index(\"cellid\", inplace=True)\n",
    "\n",
    "# build stress period data for RIV package\n",
    "riv_spd = nlmod.gwf.surface_water.build_spd(sfw_grid, \"RIV\", ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# flopy is used to construct the RIV package directly\n",
    "riv = fp.mf6.ModflowGwfriv(gwf, stress_period_data=riv_spd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlmod.sim.write_and_run(gwf, ds, silent=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Post-processing\n",
    "\n",
    "Plotting the average head in REGIS layer PZWAz3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the computed heads\n",
    "head = nlmod.gwf.output.get_heads_da(ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot on map\n",
    "ax = nlmod.plot.map_array(\n",
    "    head.sel(layer=\"PZWAz3\").mean(dim=\"time\"),\n",
    "    ds,\n",
    "    alpha=0.25,\n",
    "    background=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the GHG in the upper layer, named 'HLc'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use nlmod's calculate_gxg method to calculate the GVG, GLG and GHG.\n",
    "gxg = nlmod.gwf.calculate_gxg(head.sel(layer=\"HLc\"))\n",
    "\n",
    "# plot on map\n",
    "pc = nlmod.plot.map_array(gxg[\"ghg\"], ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### References\n",
    "\n",
    "- Calje, R., F. Schaars, D. Brakenhoff, O. Ebbens. (2022) \"Open source grondwatermodellering met MODFLOW 6\". Stromingen 2022 Vol. 03."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}