{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building a model anywhere in the Netherlands  \n",
    "This example notebook shows a basic example of a model created using online data with the `nlmod` package. `nlmod` contains functions to create modflow models anywhere in the Netherlands."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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": [
    "## Model settings\n",
    "We create a modflow model with the name 'IJmuiden'. This model has the following properties:\n",
    "\n",
    "- an extent that covers part of the Northsea, Noordzeekanaal and the small port city IJmuiden.\n",
    "- a structured grid with cells of 100 x 100 m2\n",
    "- the model is a steady state model with a single time step.\n",
    "- starting heads of 1 m NAP in every cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model settings\n",
    "model_ws = \"01_basic_model\"\n",
    "model_name = \"IJmuiden\"\n",
    "figdir, cachedir = nlmod.util.get_model_dirs(model_ws)\n",
    "extent = [95000.0, 105000.0, 494000.0, 500000.0]\n",
    "delr = 100.0\n",
    "delc = 100.0\n",
    "steady_state = True\n",
    "start_time = \"2015-1-1\"\n",
    "starting_head = 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download data\n",
    "First, download the data needed for the model from online sources. It is recommended to store a copy somewhere in case the online dataset changes. For this model we download data from these sources:\n",
    "\n",
    "- regis (layer model)\n",
    "- geotop (layer model)\n",
    "- rijkswaterstaat (surface water)\n",
    "- jarkus (bathymetry)\n",
    "- ahn (digital elevation model)\n",
    "- knmi (precipitation and evaporation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# layer models\n",
    "regis_ds = nlmod.read.regis.download_regis(extent, botm_layer=\"MSz1\",\n",
    "                                            cachedir=cachedir,\n",
    "                                            cachename=\"regis.nc\")\n",
    "\n",
    "geotop_ds = nlmod.read.geotop.download_geotop(extent,\n",
    "                                              cachedir=cachedir,\n",
    "                                              cachename=\"geotop.nc\",\n",
    "                                              chunks=None)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read surface water data\n",
    "gdf_surface_water = nlmod.read.rws.get_gdf_surface_water(\n",
    "    extent=extent, cachedir=cachedir, cachename=\"rws_surface_water.pklz\"\n",
    ")\n",
    "\n",
    "# bathymetry\n",
    "bathymetry_da = nlmod.read.jarkus.download_bathymetry(extent=extent,\n",
    "                                                      cachedir=cachedir,\n",
    "                                                      cachename=\"bathymetry.nc\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Digital elevation model\n",
    "ahn_da = nlmod.read.ahn.download_ahn(extent, cachedir=cachedir, cachename=\"ahn\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# knmi data\n",
    "oc_knmi = nlmod.read.knmi.download_knmi(extent=extent, delr=delr, delc=delc, start='2000-1-1', end='2020-1-1')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create model\n",
    "Now create the modflow model:\n",
    "\n",
    "- Build a layer model using Regis and Geotop. Use MSz1 as the bottom layer; all deeper Regis layers are excluded. Use `nlmod.read.regis.get_layer_names()` to list available layers.\n",
    "- Replace the Holocene layer in Regis with Geotop due to missing hydraulic properties.\n",
    "- Extend Regis/Geotop layers into the North Sea using shoreline extrapolation and Jarkus bathymetry data.\n",
    "- Add large water bodies (e.g. North Sea, IJsselmeer) as general head boundaries.\n",
    "- Add surface drainage using the AHN with a conductance of 1000 m²/d.\n",
    "- Calculate recharge from precipitation and evaporation time series\n",
    "- Add constant head boundaries on all model edges using the starting head."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create layer model\n",
    "layer_model = nlmod.read.regis.add_geotop_to_regis_layers(regis_ds,\n",
    "                                                          geotop_ds,\n",
    "                                                          layers=\"HLc\")\n",
    "\n",
    "# create a dataset from they layer model\n",
    "ds = nlmod.to_model_ds(layer_model, model_name, model_ws, delr=delr, delc=delc)\n",
    "\n",
    "# add time discretisation\n",
    "ds = nlmod.time.set_ds_time(ds, start=start_time, steady=steady_state, perlen=365 * 5)\n",
    "\n",
    "# add northsea and bathymetry to modelgrid\n",
    "ds.update(nlmod.read.rws.discretize_northsea(ds, gdf=gdf_surface_water))\n",
    "ds.update(nlmod.read.jarkus.discretize_bathymetry(ds, bathymetry_da.drop_vars(\"time\"), cachedir=cachedir, cachename=\"bathymetry.nc\"))\n",
    "ds = nlmod.dims.add_bathymetry_to_layer_model(ds)"
   ]
  },
  {
   "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)\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=starting_head)\n",
    "\n",
    "# Create the output control package\n",
    "oc = nlmod.gwf.oc(ds, gwf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# discretize surface water bodies from geodataframe\n",
    "rws_ds = nlmod.read.rws.discretize_surface_water(\n",
    "    ds, gdf=gdf_surface_water, da_basename=\"rws_oppwater\"\n",
    ")\n",
    "\n",
    "# add data to model dataset\n",
    "ds.update(rws_ds)\n",
    "\n",
    "# build ghb package\n",
    "ghb = nlmod.gwf.ghb(ds, gwf, bhead=\"rws_oppwater_stage\", cond=\"rws_oppwater_cond\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# discretize ahn\n",
    "ahn_ds = nlmod.read.ahn.discretize_ahn(ds, ahn_da)\n",
    "\n",
    "# add data to model dataset\n",
    "ds.update(ahn_ds)\n",
    "\n",
    "# build surface level drain package\n",
    "drn = nlmod.gwf.surface_drain_from_ds(ds, gwf, resistance=10.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add constant head cells at model boundaries\n",
    "ds.update(nlmod.grid.mask_model_edge(ds))\n",
    "chd = nlmod.gwf.chd(ds, gwf, mask=\"edge_mask\", head=\"starting_head\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# discretize knmi recharge data\n",
    "knmi_ds = nlmod.read.knmi.discretize_knmi(ds, oc_knmi, cachedir=cachedir, cachename=\"recharge\")\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": [
    "Model data is stored in the variable `ds` which is an `xarray.Dataset`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Write and Run\n",
    "Write the model files and run the model using the function `nlmod.sim.write_and_run)` as shown below. This function has two additional options:\n",
    "\n",
    "1. Saving the model dataset (`write_ds=True`) for faster future loading.\n",
    "2. Saving a copy of the current notebook (nb_path=\"<notebook_name>.ipynb\") alongside the model files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlmod.sim.write_and_run(sim, ds, write_ds=True, script_path=\"01_basic_model.ipynb\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualise\n",
    "Plot the modelgrid and surface water"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = nlmod.plot.modelgrid(ds)\n",
    "nlmod.plot.surface_water(ds, ax=ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot model input data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = nlmod.plot.get_map(ds.extent, nrows=2, ncols=2, figsize=14)\n",
    "ds[\"ahn\"].plot(ax=axes[0][0])\n",
    "ds[\"botm\"][0].plot(ax=axes[0][1])\n",
    "nlmod.layers.get_idomain(ds)[0].plot(ax=axes[1][0])\n",
    "ds[\"edge_mask\"][0].plot(ax=axes[1][1])\n",
    "\n",
    "fig, axes = nlmod.plot.get_map(ds.extent, nrows=2, ncols=2, figsize=14)\n",
    "ds[\"bathymetry\"].plot(ax=axes[0][0])\n",
    "ds[\"northsea\"].plot(ax=axes[0][1])\n",
    "ds[\"kh\"][1].plot(ax=axes[1][0])\n",
    "ds[\"recharge\"].plot(ax=axes[1][1]);"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}