{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Groundwater transport modeling\n",
    "This notebook shows how `nlmod` can be used to set up a groundwater transport\n",
    "model. In this example we create a model of a coastal area in the Netherlands\n",
    "where density driven flow caused by the higher salinity of sea water affects\n",
    "the heads."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import packages\n",
    "import flopy as fp\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import xarray as xr\n",
    "\n",
    "import nlmod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up pretty logging and show package versions\n",
    "nlmod.util.get_color_logger(\"INFO\")\n",
    "nlmod.show_versions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set model settings. \n",
    "\n",
    "Note that we set `transport` to True. This variable is passed\n",
    "to the model dataset constructor and indicates that we're building a transport\n",
    "model. This attribute is used by `nlmod` when writing modflow packages so that\n",
    "it is aware that we're working on a transport model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model settings\n",
    "model_ws = \"16_groundwater_transport\"\n",
    "model_name = \"hondsbossche\"\n",
    "\n",
    "figdir, cachedir = nlmod.util.get_model_dirs(model_ws)\n",
    "\n",
    "extent_hbossche = [103700, 106700, 527500, 528500]\n",
    "\n",
    "delr = 100.0\n",
    "delc = 100.0\n",
    "\n",
    "add_northsea = True\n",
    "transport = True\n",
    "\n",
    "start_time = \"2010-1-1\"\n",
    "starting_head = 1.0\n",
    "\n",
    "municipalities = nlmod.read.administrative.download_municipalities(extent=extent_hbossche)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the model area."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = nlmod.plot.get_map(extent_hbossche, background=\"OpenStreetMap.Mapnik\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download REGIS in our model area and create a layer model. Next convert this\n",
    "layer model into a model dataset using grid information using\n",
    "`nlmod.to_model_ds`.\n",
    "\n",
    "Then we add time discretization, add the north sea to our layer model, and set\n",
    "default transport parameters for our transport model. \n",
    "\n",
    "The last step is done with\n",
    "`nlmod.gwt.prepare.set_default_transport_parameters`. In this case we're using\n",
    "chloride concentrations to model salinity effects, so we've set default\n",
    "parameters values for that case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "layer_model = nlmod.read.regis.get_combined_layer_models(\n",
    "    extent_hbossche,\n",
    "    use_regis=True,\n",
    "    regis_botm_layer=\"MSz1\",\n",
    "    use_geotop=False,\n",
    "    cachedir=cachedir,\n",
    "    cachename=\"combined_layer_ds.nc\",\n",
    ")\n",
    "\n",
    "# create a model ds\n",
    "ds = nlmod.to_model_ds(\n",
    "    layer_model,\n",
    "    model_name,\n",
    "    model_ws,\n",
    "    delr=delr,\n",
    "    delc=delc,\n",
    "    transport=transport,\n",
    ")\n",
    "\n",
    "# add time discretisation\n",
    "ds = nlmod.time.set_ds_time(\n",
    "    ds,\n",
    "    start=start_time,\n",
    "    steady=False,\n",
    "    steady_start=True,\n",
    "    perlen=[365.0] * 10,\n",
    ")\n",
    "\n",
    "if ds.transport == 1:\n",
    "    ds = nlmod.gwt.prepare.set_default_transport_parameters(\n",
    "        ds, transport_type=\"chloride\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We download the digital terrain model (AHN4)\n",
    "ahn = nlmod.read.ahn.download_ahn4(ds.extent)\n",
    "# calculate the average surface level in each cell\n",
    "ds[\"ahn\"] = nlmod.resample.structured_da_to_ds(ahn, ds, method=\"average\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We download the surface level below the sea by downloading the vaklodingen\n",
    "vaklodingen = nlmod.read.jarkus.get_dataset_jarkus(\n",
    "    extent_hbossche,\n",
    "    kind=\"vaklodingen\",\n",
    "    time=\"2020\",\n",
    "    cachedir=cachedir,\n",
    "    cachename=\"vaklodingen.nc\",\n",
    ")\n",
    "# calculate the average surface level in each cell\n",
    "ds[\"vaklodingen\"] = nlmod.resample.structured_da_to_ds(\n",
    "    vaklodingen[\"z\"], ds, method=\"average\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate a new top from ahn and vaklodingen\n",
    "new_top = ds[\"ahn\"].where(~ds[\"ahn\"].isnull(), ds[\"vaklodingen\"])\n",
    "ds = nlmod.layers.set_model_top(ds, new_top)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we then determine the part of each cell that is covered by sea from the original ahn\n",
    "ds[\"sea\"] = nlmod.read.rws.calculate_sea_coverage(ahn, ds=ds, method=\"average\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we load chloride concentrations for our model. These are obtained from the\n",
    "NHI salinity dataset, where chloride concentrations for the Netherlands were\n",
    "determined based on observations and modeling. The full dataset `3dchlorde.nc`\n",
    "can be downloaded from here: https://zenodo.org/record/7419219. Here we load a\n",
    "small dataset that was extracted from the full dataset.\n",
    "\n",
    "This dataset does not match our model grid, so we use nearest interpolation get\n",
    "the chloride concentration for each of our model cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cl = xr.open_dataset(\"../../../pwn_diep/data/3dchloride_result.nc\")\n",
    "cl = xr.open_dataset(\"./data/chloride_hbossche.nc\")\n",
    "\n",
    "\n",
    "# interpolate to modelgrid using nearest\n",
    "cli = cl.sel(percentile=\"p50\").interp(x=ds.x, y=ds.y, method=\"nearest\")\n",
    "cli"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The chloride concentration dataset also does not have the same vertical\n",
    "discretization as our model. In order to calculate the mean concentration in\n",
    "each cell in every layer of our model we use\n",
    "`nlmod.layers.aggregate_by_weighted_mean_to_ds` to calculate the weighted mean\n",
    "of the chloride concentration observations in each layer. We also fill the NaNs\n",
    "in the resulting dataset using nearest interpolation.\n",
    "\n",
    "Finally, we add this chloride data array to our model dataset, which now has a\n",
    "chloride concentration for each cell in our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# aggregate chloride to our layer model using weighted mean\n",
    "cli_da = nlmod.layers.aggregate_by_weighted_mean_to_ds(ds, cli, \"3d-chloride\")\n",
    "\n",
    "# interpolate NaNs nearest\n",
    "for ilay in range(cli_da.shape[0]):\n",
    "    cli_da.values[ilay] = nlmod.resample.fillnan_da(\n",
    "        da=cli_da.isel(layer=ilay), method=\"nearest\"\n",
    "    )\n",
    "\n",
    "# set chloride data in model dataset, keep only layer, y and x coordinates\n",
    "ds[\"chloride\"] = (\"layer\", \"y\", \"x\"), cli_da.values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can start building our groundwater model. We start with the Simulation object,\n",
    "time discretization and IMS solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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, complexity=\"MODERATE\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we add the groundwater flow model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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 storage\n",
    "sto = nlmod.gwf.sto(ds, gwf)\n",
    "\n",
    "# Create the initial conditions package\n",
    "ic = nlmod.gwf.ic(ds, gwf, starting_head=starting_head)\n",
    "\n",
    "# Create the output control package\n",
    "oc = nlmod.gwf.oc(ds, gwf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We add general head boundaries to model the North Sea. We want to provide the\n",
    "North Sea with a chloride concentration of 18,000 mgCl-/L. This can be done by\n",
    "passing this value to the auxiliary keyword argument.\n",
    "\n",
    "Note that it is also possible to reference one (or more) data arrays from the\n",
    "model dataset as the auxiliary variable.\n",
    "\n",
    "If an auxiliary variable is provided and the transport attribute of the model\n",
    "dataset is 1 (True), `nlmod` automatically registers the GHB package in the\n",
    "`ssm_sources` attribute, which indicates that we need to add this package as a\n",
    "source (or sink) for our transport model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build ghb package\n",
    "ghb = nlmod.gwf.ghb(\n",
    "    ds,\n",
    "    gwf,\n",
    "    bhead=0.0,\n",
    "    cond=ds[\"sea\"] * ds[\"area\"] / 1.0,\n",
    "    auxiliary=18_000.0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# note that building the GHB added the package to the ssm_sources attribute\n",
    "ds.ssm_sources"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add surface level drains to the model based on the digital elevetion model AHN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build surface level drain package\n",
    "elev = ds[\"ahn\"].where(ds[\"sea\"] == 0)\n",
    "drn = nlmod.gwf.surface_drain_from_ds(ds, gwf, elev=elev, resistance=10.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add recharge based on timeseries measured at meteorolgical stations by KNMI. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download knmi recharge data\n",
    "knmi_ds = nlmod.read.knmi.get_recharge(ds, cachedir=ds.cachedir, cachename=\"recharge\")\n",
    "# update model dataset\n",
    "ds.update(knmi_ds)\n",
    "\n",
    "# create recharge package\n",
    "rch = nlmod.gwf.rch(ds, gwf, mask=ds[\"sea\"] == 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, the transport model is created. Note the following steps:\n",
    "\n",
    "- The buoyancy (BUY) package is added to the groundwater flow model to take\n",
    "into account density effects.\n",
    "- The transport model requires its own IMS solver, which also needs to be\n",
    "registered in the simulation.\n",
    "- The advection (ADV), dispersion (DSP), mass-storage transfer (MST) and\n",
    "source-sink mixing (SSM) packages each obtain information from the model\n",
    "dataset. These variables were defined by\n",
    "`nlmod.gwt.prepare.set_default_transport_parameters`. They can be also be\n",
    "modified or added to the dataset by the user. Another option is to directly\n",
    "pass the variables to the package constructors, in which case the stored values\n",
    "are ignored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if ds.transport:\n",
    "    # BUY: buoyancy package for GWF model\n",
    "    buy = nlmod.gwf.buy(ds, gwf)\n",
    "\n",
    "    # GWT: groundwater transport model\n",
    "    gwt = nlmod.gwt.gwt(ds, sim)\n",
    "\n",
    "    # add IMS for GWT model and register it\n",
    "    ims = nlmod.sim.ims(sim, pname=\"ims_gwt\", filename=f\"{gwt.name}.ims\")\n",
    "    nlmod.sim.register_ims_package(sim, gwt, ims)\n",
    "\n",
    "    # DIS: discretization package\n",
    "    dis_gwt = nlmod.gwt.dis(ds, gwt)\n",
    "\n",
    "    # IC: initial conditions package\n",
    "    ic_gwt = nlmod.gwt.ic(ds, gwt, \"chloride\")\n",
    "\n",
    "    # ADV: advection package\n",
    "    adv = nlmod.gwt.adv(ds, gwt)\n",
    "\n",
    "    # DSP: dispersion package\n",
    "    dsp = nlmod.gwt.dsp(ds, gwt)\n",
    "\n",
    "    # MST: mass transfer package\n",
    "    mst = nlmod.gwt.mst(ds, gwt)\n",
    "\n",
    "    # SSM: source-sink mixing package\n",
    "    ssm = nlmod.gwt.ssm(ds, gwt)\n",
    "\n",
    "    # OC: output control\n",
    "    oc_gwt = nlmod.gwt.oc(ds, gwt)\n",
    "\n",
    "    # GWF-GWT Exchange\n",
    "    gwfgwt = nlmod.gwt.gwfgwt(ds, sim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now write the model files and run the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlmod.sim.write_and_run(sim, ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the model input, specifically the boundary conditions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot using flopy\n",
    "fig, ax = nlmod.plot.get_map(extent_hbossche, background=\"OpenStreetMap.Mapnik\")\n",
    "pmv = fp.plot.PlotMapView(model=gwf, layer=0, ax=ax)\n",
    "# pc = pmv.plot_array(c.isel(time=0), cmap=\"Spectral_r\")\n",
    "pmv.plot_bc(\"GHB\", plotAll=True, alpha=0.1, label=\"GHB\")\n",
    "pmv.plot_bc(\"DRN\", plotAll=True, alpha=0.1, label=\"DRN\")\n",
    "# pmv.plot_bc(\"RCH\", plotAll=True, alpha=0.1, label=\"RCH\")\n",
    "municipalities.plot(edgecolor=\"k\", facecolor=\"none\", ax=ax)\n",
    "pmv.plot_grid(linewidth=0.25)\n",
    "ax.set_xlabel(\"x [km RD]\")\n",
    "ax.set_ylabel(\"y [km RD]\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the calculated heads and concentrations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h = nlmod.gwf.output.get_heads_da(ds)\n",
    "c = nlmod.gwt.output.get_concentration_da(ds)\n",
    "\n",
    "# calculate concentration at groundwater surface\n",
    "ctop = nlmod.gwt.get_concentration_at_gw_surface(c)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the concentration at groundwater surface level."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = nlmod.plot.map_array(\n",
    "    ctop.isel(time=-1),\n",
    "    ds,\n",
    "    ilay=0,\n",
    "    cmap=\"Spectral_r\",\n",
    ")\n",
    "municipalities.plot(edgecolor=\"k\", facecolor=\"none\", ax=ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot a cross-section along (x) showing the calculated concentration in the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = (ds.extent[2] + ds.extent[3]) / 2 + 0.1\n",
    "line = [(ds.extent[0], y), (ds.extent[1], y)]\n",
    "zmin = -150.0\n",
    "zmax = 10.0\n",
    "\n",
    "for time_idx in [0, -1]:\n",
    "    # plot using flopy\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(16, 5))\n",
    "    pmv = fp.plot.PlotCrossSection(model=gwf, line={\"line\": line})\n",
    "    pc = pmv.plot_array(\n",
    "        c.isel(time=time_idx), cmap=\"Spectral_r\", vmin=0.0, vmax=18_000.0\n",
    "    )\n",
    "    pmv.plot_bc(\"GHB\", color=\"k\", zorder=10)\n",
    "    pmv.plot_grid(linewidth=0.25)\n",
    "    cbar = fig.colorbar(pc, ax=ax)\n",
    "    cbar.set_label(\"chloride (mg/L)\")\n",
    "    ax.set_ylim(bottom=-100)\n",
    "    ax.set_xlabel(\"x [m]\")\n",
    "    ax.set_ylabel(\"elevation [m NAP]\")\n",
    "    # convert to pandas timestamp for prettier printing\n",
    "    ax.set_title(f\"time = {pd.Timestamp(c.time.isel(time=time_idx).values)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Converting calculated heads (which represent point water heads) to equivalent\n",
    "freshwater heads, and vice versa, can be done with the following functions in `nlmod`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf = nlmod.gwt.output.freshwater_head(ds, h, c)\n",
    "hp = nlmod.gwt.output.pointwater_head(ds, hf, c)\n",
    "\n",
    "xr.testing.assert_allclose(h, hp)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}