{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "064b0618",
   "metadata": {},
   "source": [
    "# Aggregating surface water\n",
    "In `nlmod`, surface water is mostly added to a groundwater model using a GeoDataFrame with surface water features. `nlmod` can generate input from this GeoDataFrame for the DRN- RIV or LAK-packages. Each individual record can directly be added to these packages. Aggregation of surface water features in cells is also possible, using different aggregation-methods.\n",
    "\n",
    "This section demonstrates how to aggrgate surface water features into model cells and add them to a MODFLOW model. In this notebook we perform a steady-state run, in which the stage of the surface water is the mean of the summer and winter stage. The surface water bodies in each cell are aggregated using an area-weighted method and added to the model with the river-package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75f92bdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import geopandas as gpd\n",
    "import flopy\n",
    "import nlmod"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01d6ac40",
   "metadata": {},
   "source": [
    "## Model settings\n",
    "We define some model settings, like the name, the directory of the model files, the model extent and the time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7a972d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"Schoonhoven\"\n",
    "model_ws = \"03_aggregating_surface_water\"\n",
    "figdir, cachedir = nlmod.util.get_model_dirs(model_ws)\n",
    "extent = [116_500, 120_000, 439_000, 442_000]\n",
    "time = pd.date_range(\"2020\", \"2023\", freq=\"MS\")  # monthly timestep"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41c6f097",
   "metadata": {},
   "source": [
    "### layer 'waterdeel' from bgt\n",
    "The location of the surface water bodies is obtained from the GeoDataFrame that was created in the the surface water notebook. We saved this data as a geosjon file and load it here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69294636",
   "metadata": {},
   "outputs": [],
   "source": [
    "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 in the 'data_sources' directory first\"\n",
    "        )\n",
    "    )\n",
    "bgt = gpd.read_file(fname_bgt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f637f3f9",
   "metadata": {},
   "source": [
    "#### Change some values in the GeoDataFrame for this model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57e0e50c",
   "metadata": {},
   "outputs": [],
   "source": [
    "sfw = bgt\n",
    "sfw[\"stage\"] = sfw[[\"winter_stage\", \"summer_stage\"]].mean(1)\n",
    "# use a water depth of 0.5 meter\n",
    "sfw[\"botm\"] = sfw[\"stage\"] - 0.5\n",
    "# set the stage of the Lek to 0.0 m NAP and the botm to -3 m NAP\n",
    "mask = sfw[\"bronhouder\"] == \"L0002\"\n",
    "sfw.loc[mask, \"stage\"] = 0.0\n",
    "sfw.loc[mask, \"botm\"] = -3.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a665f10b",
   "metadata": {},
   "source": [
    "Take a look at the first few rows. For adding surface water features to a MODFLOW model the following attributes must be present:\n",
    "\n",
    "- **stage**: the water level (in m NAP)\n",
    "- **botm**: the bottom elevation (in m NAP)\n",
    "- **c0**: the bottom resistance (in days)\n",
    "\n",
    "The `stage` and the `botm` columns are present in our dataset. The bottom resistance `c0` is rarely known, and is usually estimated when building the model. We will add our estimate later on.\n",
    "\n",
    "    \n",
    "```{note}\n",
    "The NaN's in the dataset indicate that not all parameters are known for each feature. This is not necessarily a problem but this will mean some features will not be converted to model input.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da656ab3",
   "metadata": {},
   "source": [
    "Now use `stage` as the column to color the data. Note the missing features caused by the fact that the stage is undefined (NaN)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f002f2cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = nlmod.plot.get_map(extent)\n",
    "sfw.plot(ax=ax, column=\"stage\", legend=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c37e0e5b",
   "metadata": {},
   "source": [
    "## Build model\n",
    "\n",
    "The next step is to define a model grid and build a model (i.e. create a discretization and define flow parameters)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1a1d5f1",
   "metadata": {},
   "source": [
    "Build the model. We're keeping the model as simple as possible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53d64caa",
   "metadata": {},
   "outputs": [],
   "source": [
    "delr = delc = 50.0\n",
    "start_time = \"2021-01-01\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68199661",
   "metadata": {},
   "outputs": [],
   "source": [
    "# layer model\n",
    "layer_model = nlmod.read.download_regis(\n",
    "    extent, cachedir=cachedir, cachename=\"layer_model.nc\"\n",
    ")\n",
    "layer_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9486033",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a model ds by changing grid of layer_model\n",
    "ds = nlmod.to_model_ds(layer_model, model_name, model_ws, delr=delr, delc=delc)\n",
    "\n",
    "# create model time dataset\n",
    "ds = nlmod.time.set_ds_time(ds, start=start_time, steady=True, perlen=1)\n",
    "\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02adb937",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create simulation\n",
    "sim = nlmod.sim.sim(ds)\n",
    "\n",
    "# create time discretisation\n",
    "tdis = nlmod.sim.tdis(ds, sim)\n",
    "\n",
    "# create ims\n",
    "ims = nlmod.sim.ims(sim)\n",
    "\n",
    "# create groundwater flow model\n",
    "gwf = nlmod.gwf.gwf(ds, sim)\n",
    "\n",
    "# Create discretization\n",
    "dis = nlmod.gwf.dis(ds, gwf)\n",
    "\n",
    "# create node property flow\n",
    "npf = nlmod.gwf.npf(ds, gwf)\n",
    "\n",
    "# Create the initial conditions package\n",
    "ic = nlmod.gwf.ic(ds, gwf, starting_head=1.0)\n",
    "\n",
    "# Create the output control package\n",
    "oc = nlmod.gwf.oc(ds, gwf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb59b5cb",
   "metadata": {},
   "source": [
    "## Add surface water\n",
    "\n",
    "Now that we have a discretization (a grid, and layer tops and bottoms) we can start processing our surface water shapefile to add surface water features to our model. The method to add surface water starting from a shapefile is divided into the following steps:\n",
    "\n",
    "1. Intersect surface water shape with grid. This steps intersects every feature with the grid so we can determine the surface water features in each cell.\n",
    "2. Aggregate parameters per grid cell. Each feature within a cell has its own parameters. For MODFLOW it is often desirable to have one representative set of parameters per cell. These representative parameters are calculated in this step.\n",
    "3. Build stress period data. The results from the previous step are converted to stress period data (generally a list of cellids and representative parameters: `[(cellid), parameters]`) which is used by MODFLOW and flopy to define boundary conditions.\n",
    "4. Create the Modflow6 package\n",
    "\n",
    "The steps are illustrated below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a42a0323",
   "metadata": {},
   "source": [
    "### Intersect surface water shape with grid\n",
    "\n",
    "The first step is to intersect the surface water shapefile with the grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "054176cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "sfw_grid = nlmod.grid.gdf_to_grid(\n",
    "    sfw, gwf, cachedir=ds.cachedir, cachename=\"sfw_grid.pklz\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b50fc596",
   "metadata": {},
   "source": [
    "Plot the result and the model grid and color using `cellid`. It's perhaps a bit hard to see but each feature is cut by the gridlines. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d8a133a",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = nlmod.plot.get_map(extent)\n",
    "sfw_grid.plot(ax=ax, column=\"cellid\")\n",
    "nlmod.plot.modelgrid(ds, ax=ax, lw=0.2);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92aae854",
   "metadata": {},
   "source": [
    "### Aggregate parameters per model cell\n",
    "The next step is to aggregate the parameters for all the features in one grid cell to obtain one representative set of parameters. First, let's take a look at a grid cell containing multiple features. We take the gridcell that contains the most features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d46ae109",
   "metadata": {},
   "outputs": [],
   "source": [
    "cid = sfw_grid.cellid.value_counts().index[0]\n",
    "mask = sfw_grid.cellid == cid\n",
    "sfw_grid.loc[mask]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95be47e3",
   "metadata": {},
   "source": [
    "We can also plot the features within that grid cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0d51b2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(10, 8))\n",
    "sfw_grid.loc[mask].plot(\n",
    "    column=\"identificatie\",\n",
    "    legend=True,\n",
    "    ax=ax,\n",
    "    legend_kwds={\"loc\": \"lower left\", \"ncol\": 2, \"fontsize\": \"x-small\"},\n",
    ")\n",
    "xlim = ax.get_xlim()\n",
    "ylim = ax.get_ylim()\n",
    "gwf.modelgrid.plot(ax=ax)\n",
    "ax.set_xlim(xlim[0], xlim[0] + nlmod.grid.get_delr(ds)[-1] * 1.1)\n",
    "ax.set_ylim(ylim)\n",
    "ax.set_title(f\"Surface water shapes in cell: {cid}\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dadc16ef",
   "metadata": {},
   "source": [
    "Now we want to aggregate the features in each cell to obtain a representative set of parameters (`stage`, `conductance`, `bottom elevation`) to use in the model. There are several aggregation methods. Note that the names of the methods are not representative of the aggregation applied to each parameter. For a full description see the following list:\n",
    "\n",
    "- `'area_weighted'`\n",
    "  - **stage**: area-weighted average of stage in cell\n",
    "  - **cond**: conductance is equal to area of surface water divided by bottom resistance\n",
    "  - **elev**: the lowest bottom elevation is representative for the cell\n",
    "- `'max_area'`\n",
    "  - **stage**: stage is determined by the largest surface water feature in a cell\n",
    "  - **cond**: conductance is equal to area of all surface water features divided by bottom resistance\n",
    "  - **elev**: the lowest bottom elevation is representative for the cell\n",
    "- `'de_lange'`\n",
    "  - **stage**: area-weighted average of stage in cell\n",
    "  - **cond**: conductance is calculated using the formulas derived by De Lange (1999).\n",
    "  - **elev**: the lowest bottom elevation is representative for the cell\n",
    "  \n",
    "Let's try using `area_weighted`. This means the stage is the area-weighted average of all the surface water features in a cell. The conductance is calculated by dividing the total area of surface water in a cell by the bottom resistance (`c0`). The representative bottom elevation is the lowest elevation present in the cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af97529c",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    nlmod.gwf.surface_water.aggregate(sfw_grid, \"area_weighted\")\n",
    "except ValueError as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c13ef10f",
   "metadata": {},
   "source": [
    "The function checks whether the requisite columns are defined in the DataFrame. We need to add a column containing the bottom resistance `c0`. Often a value of 1 day is used as an initial estimate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "462dad4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "sfw_grid[\"c0\"] = 1.0  # days"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7e3a915",
   "metadata": {},
   "source": [
    "Now aggregate the features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b44c3c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "celldata = nlmod.gwf.surface_water.aggregate(\n",
    "    sfw_grid, \"area_weighted\", cachedir=ds.cachedir, cachename=\"celldata.pklz\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "855ede3b",
   "metadata": {},
   "source": [
    "Let's take a look at the result. We now have a DataFrame with cell-id as the index and the three parameters we need for each cell `stage`, `cond` and `rbot`. The area is also given, but is not needed for the groundwater model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bcf3534",
   "metadata": {},
   "outputs": [],
   "source": [
    "celldata.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a12fe67a",
   "metadata": {},
   "source": [
    "### Build stress period data\n",
    "\n",
    "The next step is to take our cell-data and build convert it to 'stress period data' for MODFLOW. This is a data format that defines the parameters in each cell in the following format:\n",
    "\n",
    "```\n",
    "[[(cellid1), param1a, param1b, param1c],\n",
    " [(cellid2), param2a, param2b, param2c],\n",
    " ...]\n",
    "```\n",
    "\n",
    "The required parameters are defined by the MODFLOW-package used:\n",
    "\n",
    "- **RIV**: for the river package `(stage, cond, rbot)`\n",
    "- **DRN**: for the drain package `(stage, cond)`\n",
    "- **GHB**: for the general-head-boundary package `(stage, cond)`\n",
    "\n",
    "We're selecting the RIV package. We don't have a bottom (rbot) for each reach in celldata. Therefore we remove the reaches where rbot is nan (not a number)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f77f3ea2",
   "metadata": {},
   "outputs": [],
   "source": [
    "new_celldata = celldata.loc[~celldata.rbot.isna()]\n",
    "print(f\"removed {len(celldata)-len(new_celldata)} reaches because rbot is nan\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21b0f28c",
   "metadata": {},
   "outputs": [],
   "source": [
    "riv_spd = nlmod.gwf.surface_water.build_spd(new_celldata, \"RIV\", ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8edccd04",
   "metadata": {},
   "source": [
    "Take a look at the stress period data for the river package:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e6a3fc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "riv_spd[:10]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acd8885c",
   "metadata": {},
   "source": [
    "### Create RIV package\n",
    "The final step is to create the river package using flopy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d84280eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "riv = flopy.mf6.ModflowGwfriv(gwf, stress_period_data=riv_spd)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9077a58c",
   "metadata": {},
   "source": [
    "Plot the river boundary condition to see where rivers were added in the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48287cc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use flopy plotting methods\n",
    "fig, ax = plt.subplots(1, 1, figsize=(10, 8), constrained_layout=True)\n",
    "mv = flopy.plot.PlotMapView(model=gwf, ax=ax, layer=0)\n",
    "mv.plot_bc(\"RIV\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91a76816",
   "metadata": {},
   "source": [
    "## Write + run model\n",
    "\n",
    "Now write the model simulation to disk, and run the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64f9b796",
   "metadata": {},
   "outputs": [],
   "source": [
    "nlmod.sim.write_and_run(sim, ds, write_ds=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e79facf",
   "metadata": {},
   "source": [
    "## Visualize results\n",
    "\n",
    "To see whether our surface water was correctly added to the model, let's visualize the results. We'll load the calculated heads, and plot them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8084a14c",
   "metadata": {},
   "outputs": [],
   "source": [
    "head = nlmod.gwf.get_heads_da(ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c37e4e5",
   "metadata": {},
   "source": [
    "Plot the heads in a specific model layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d82805f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# using nlmod plotting methods\n",
    "ax = nlmod.plot.map_array(\n",
    "    head,\n",
    "    ds,\n",
    "    ilay=0,\n",
    "    iper=0,\n",
    "    plot_grid=True,\n",
    "    title=\"Heads top-view\",\n",
    "    cmap=\"RdBu\",\n",
    "    colorbar_label=\"head [m NAP]\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84031738",
   "metadata": {},
   "source": [
    "In cross-section"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1475b56",
   "metadata": {},
   "outputs": [],
   "source": [
    "# using flopy plotting methods\n",
    "col = gwf.modelgrid.ncol // 2\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(10, 3))\n",
    "xs = flopy.plot.PlotCrossSection(model=gwf, ax=ax, line={\"column\": col})\n",
    "qm = xs.plot_array(head[-1], cmap=\"RdBu\")  # last timestep\n",
    "xs.plot_ibound()  # plot inactive cells in red\n",
    "xs.plot_grid(lw=0.25, color=\"k\")\n",
    "ax.set_ylim(bottom=-150)\n",
    "ax.set_ylabel(\"elevation [m NAP]\")\n",
    "ax.set_xlabel(\"distance along cross-section [m]\")\n",
    "ax.set_title(f\"Cross-section along column {col}\")\n",
    "cbar = fig.colorbar(qm, shrink=1.0)\n",
    "cbar.set_label(\"head [m NAP]\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}