{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building a model with local grid refinement\n",
    "This notebook shows how `nlmod` can be used to create a model with local grid refinement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import geopandas as gpd\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": [
    "## Model settings\n",
    "Modflow 6 makes it possible to use locally refined grids. In `nlmod` you can use a shapefile and a number of levels to specify where and how much you want to use local grid refinement. Below we use a shapefile of the Planetenweg in IJmuiden and set the refinement levels at 2. This well create a grid with cells of 100x100m except at the Planetenweg where the cells will be refined to 25x25m. See figures below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model settings vertex\n",
    "model_ws = \"03_local_grid_refinement\"\n",
    "model_name = \"IJm_planeten\"\n",
    "figdir, cachedir = nlmod.util.get_model_dirs(model_ws)\n",
    "refine_shp_fname = os.path.abspath(os.path.join(\"data\", \"planetenweg_ijmuiden\"))\n",
    "levels = 2\n",
    "extent = [95000.0, 105000.0, 494000.0, 500000.0]\n",
    "delr = 100.0\n",
    "delc = 100.0\n",
    "steady_state = False\n",
    "steady_start = True\n",
    "transient_timesteps = 5\n",
    "perlen = 1.0\n",
    "start_time = \"2015-1-1\"\n",
    "add_northsea = True\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\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read surface water data\n",
    "gdf_surface_water = nlmod.read.rws.get_gdf_surface_water(extent=extent,\n",
    "                                                         cachedir=cachedir,\n",
    "                                                         cachename='rws_surface_water.pklz')\n",
    "\n",
    "# bathymetry\n",
    "bathymetry_da = nlmod.read.jarkus.download_bathymetry(extent=extent,\n",
    "                                                      cachedir=cachedir,\n",
    "                                                      cachename=\"bathymetry.nc\")\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",
    "Modflow 6 makes it possible to use locally refined grids. In `nlmod` you can use a shapefile and a number of levels to specify where and how much you want to use local grid refinement. Below we use a shapefile of the Planetenweg in IJmuiden and set the refinement levels at 2. This well create a grid with cells of 100x100m except at the Planetenweg where the cells will be refined to 25x25m. See figures below."
   ]
  },
  {
   "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 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",
    "# add time discretisation\n",
    "ds = nlmod.time.set_ds_time(\n",
    "    ds,\n",
    "    start=start_time,\n",
    "    steady=steady_state,\n",
    "    steady_start=steady_start,\n",
    "    perlen=[perlen] * (transient_timesteps + 1),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Local grid refinement\n",
    "The code below applies a local grid refinement to the layer model. The local grid refinement is based on the shapefile 'planetenweg_ijmuiden.shp', which contains a line shape of the Planetenweg, and the levels, which is 2. This means that the model cells at the Planetenweg will get a size of 25 x 25m because we halving the cell size twice (100 / (2^2) = 25). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use gridgen to create vertex grid\n",
    "ds = nlmod.grid.refine(ds, refinement_features=[(refine_shp_fname, \"line\", levels)])\n",
    "\n",
    "# read surface water data\n",
    "gdf_surface_water = nlmod.read.rws.get_gdf_surface_water(ds=ds,\n",
    "                                                         cachedir=cachedir,\n",
    "                                                         cachename='rws_surface_water.pklz')\n",
    "\n",
    "# add northsea to modelgrid\n",
    "if add_northsea:\n",
    "    ds.update(nlmod.read.rws.discretize_northsea(ds, gdf=gdf_surface_water))\n",
    "    ds.update(nlmod.read.jarkus.discretize_bathymetry(ds, bathymetry_da, 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",
    "disv = nlmod.gwf.disv(ds, gwf)\n",
    "\n",
    "# create node property flow\n",
    "npf = nlmod.gwf.npf(ds, gwf, save_flows=True)\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",
    "da_name = \"rws_oppwater\"\n",
    "rws_ds = nlmod.read.rws.discretize_surface_water(\n",
    "    ds, gdf=gdf_surface_water, da_basename=da_name\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=f\"{da_name}_stage\", cond=f\"{da_name}_cond\")\n",
    "\n",
    "# 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)\n",
    "\n",
    "# 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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Write and Run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlmod.sim.write_and_run(sim, ds, write_ds=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualise\n",
    "Using the `ds` and `gwf` variables it is quite easy to visualise model data. Below the modelgrid together with the surface water is shown."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plan_weg_gdf = gpd.read_file(refine_shp_fname + \".shp\")\n",
    "\n",
    "# plot modelgrid\n",
    "ax = nlmod.plot.modelgrid(ds)\n",
    "nlmod.plot.surface_water(ds, ax=ax)\n",
    "plan_weg_gdf.plot(ax=ax, color=\"r\", label=\"Planetenweg\")\n",
    "ax.legend()\n",
    "\n",
    "# plot zoomed modelgrid\n",
    "ax = nlmod.plot.modelgrid(ds)\n",
    "nlmod.plot.surface_water(ds, ax=ax)\n",
    "ax.set_title(\"Planetenweg\")\n",
    "plan_weg_gdf.plot(ax=ax, color=\"r\", label=\"Planetenweg\")\n",
    "ax.set_xlim(100000, 103000)\n",
    "ax.set_ylim(495000, 497500)\n",
    "ax.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model dataset of a vertex model differs from a structured model dataset. The data is stored relative to the cell-id instead of the row and column number. Therefore the model dataset has the dimension icell2d instead of the dimensions x and y. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To plot the same rasters as for the previous model we can use the `nlmod.plot.data_array()` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = nlmod.plot.get_map(extent, nrows=2, ncols=2, figsize=14)\n",
    "\n",
    "nlmod.plot.data_array(ds[\"ahn\"], ds, ax=axes[0][0])\n",
    "nlmod.plot.data_array(ds[\"botm\"][0], ds, ax=axes[0][1])\n",
    "nlmod.plot.data_array(nlmod.layers.get_idomain(ds)[0], ds, ax=axes[1][0])\n",
    "nlmod.plot.data_array(ds[\"edge_mask\"][0], ds, ax=axes[1][1])\n",
    "\n",
    "fig, axes = nlmod.plot.get_map(extent, nrows=2, ncols=2, figsize=(14, 11))\n",
    "nlmod.plot.data_array(ds[\"bathymetry\"], ds, ax=axes[0][0])\n",
    "nlmod.plot.data_array(ds[\"northsea\"], ds, ax=axes[0][1])\n",
    "nlmod.plot.data_array(ds[\"kh\"][1], ds, ax=axes[1][0])\n",
    "nlmod.plot.data_array(ds[\"recharge\"][0], ds, ax=axes[1][1]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can save the entire model as a UGRID NetCDF-file. This can be opened in qgis, as a 'Mesh Layer'. For more information see https://docs.qgis.org/3.16/en/docs/user_manual/working_with_mesh/mesh_properties.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = os.path.join(ds.figdir, \"results.nc\")\n",
    "nlmod.gis.ds_to_ugrid_nc_file(ds, fname)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}