{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modifying model layers\n",
    "This notebook shows methods for modyfying model layers for MODFLOW models, for example by combining layers and splitting layers. Multiple layers can be combined into one layer or one layer can be split into sub-layers based on a fraction of the original thickness."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from shapely.geometry import LineString\n",
    "\n",
    "import nlmod\n",
    "from nlmod.plot import DatasetCrossSection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlmod.util.get_color_logger(\"INFO\")\n",
    "nlmod.show_versions()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compare_layer_models(\n",
    "    ds1,\n",
    "    line,\n",
    "    colors,\n",
    "    ds2=None,\n",
    "    zmin=-200,\n",
    "    zmax=10,\n",
    "    min_label_area=1000,\n",
    "    title1=\"REGIS original\",\n",
    "    title2=\"Modified layers\",\n",
    "    xlabel=\"Distance along x-sec (m)\",\n",
    "    ylabel=\"m NAP\",\n",
    "):\n",
    "    if ds2 is None:\n",
    "        fig, ax1 = plt.subplots(1, 1, figsize=(14, 6))\n",
    "    else:\n",
    "        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 12), sharex=True)\n",
    "    dcs1 = DatasetCrossSection(ds1, line=line, ax=ax1, zmin=zmin, zmax=zmax)\n",
    "    polys2 = dcs1.plot_layers(colors=colors, min_label_area=min_label_area)\n",
    "    dcs1.plot_grid(linewidth=0.5, vertical=False)\n",
    "    ax1.set_ylabel(ylabel)\n",
    "\n",
    "    if ds2 is not None:\n",
    "        ax1.set_title(title1)\n",
    "        dcs2 = DatasetCrossSection(ds2, line=line, ax=ax2, zmin=zmin, zmax=zmax)\n",
    "        polys1 = dcs2.plot_layers(colors=colors, min_label_area=min_label_area)\n",
    "        dcs2.plot_grid(linewidth=0.5, vertical=False)\n",
    "        ax2.set_ylabel(ylabel)\n",
    "        ax2.set_xlabel(xlabel)\n",
    "        ax2.set_title(title2)\n",
    "    else:\n",
    "        ax1.set_xlabel(xlabel)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get data\n",
    "\n",
    "Define an extent to obtain REGIS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "extent = [131000, 136800, 471500, 475700]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download and cache REGIS netCDF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "regis = nlmod.read.regis.download_regis(extent)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take a look at the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "regis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define an line to draw a cross-section"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# diagonal line through extent\n",
    "line = LineString([(extent[0], extent[2]), (extent[1], extent[3])])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get colors for our cross-section plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "colors = nlmod.read.regis.get_legend()[\"color\"].to_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Draw the cross-section for REGIS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compare_layer_models(regis, line, colors)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Split layers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we determine how to split the layers. This is done by creating a list of factors,\n",
    "that is used to determine fractions that add up to 1. The layer will be split into\n",
    "sub-layers from the top down, with each sub-layer getting a thickness equal to the\n",
    "fraction times the original thickness.\n",
    "\n",
    "For example, `(1, 1)` will split the layer into two sub-layers, each getting a\n",
    "thickness equal to 50% of the original layer. In this example the fractions already add\n",
    "up to 1 for each layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# split dictionary\n",
    "split_dict = {\n",
    "    \"PZWAz2\": (0.3, 0.3, 0.4),\n",
    "    \"PZWAz3\": (0.2, 0.2, 0.2, 0.2, 0.2),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate the new layer elevations based on the information above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "regis_split, split_reindexer = nlmod.layers.split_layers_ds(\n",
    "    regis, split_dict, return_reindexer=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "View the resulting Dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "regis_split"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The reindexer dictionary links the new layers to the old layers. This can be convenient\n",
    "for copying data from the original layers to the new sub-layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# key = new layer index\n",
    "# value = original layer index: repeats where layer was split\n",
    "split_reindexer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the cross-section of the original and the new layer model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compare_layer_models(regis, line, colors, ds2=regis_split, title2=\"Split layers\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Combine layers\n",
    "\n",
    "Example how to combine model layers. First find the indices of the layers to combine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "combine_layers = [\n",
    "    tuple(np.argwhere(regis.layer.str.startswith(\"URz\").data).squeeze().tolist()),\n",
    "    tuple(np.argwhere(regis.layer.isin([\"PZWAz2\", \"PZWAz3\"]).data).squeeze().tolist()),\n",
    "]\n",
    "combine_layers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Combine layers using the `combine_layers_ds()` function and passing the layer dataset and the list of layer numbers to combine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "regis_combined = nlmod.layers.combine_layers_ds(regis, combine_layers, kD=None, c=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take a look a the resulting dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "regis_combined"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the new and the old cross-section. Use the layer code and color from the first layer name for the combined layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compare_layer_models(regis, line, colors, ds2=regis_combined, title2=\"Combined layers\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set new model top\n",
    "\n",
    "The `nlmod.layers.set_model_top` changes the top of the model. When the new top is\n",
    "lower than the old top, the new top is burned in the layer model, lowering the top of\n",
    "the top layer(s). Top layers can become incactive, when the thickness is reduced to 0.\n",
    "When the new top is higher than the old top, the thickness of the most upper active\n",
    "layer (not necessarily the first) is increased. This method can be used to change the\n",
    "model top to a digital terrain model with a higher accuracy.\n",
    "\n",
    "First transform the regis-date to a model Dataset, as the next methods need a model\n",
    "Dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = nlmod.to_model_ds(regis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_new = nlmod.layers.set_model_top(ds.copy(deep=True), 5.0)\n",
    "compare_layer_models(ds, line, colors, ds2=ds_new, title2=\"New top\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set layer top\n",
    "`nlmod.layers.set_layer_top` sets the layer top to a specified value or array.\n",
    "\n",
    "This method only changes the shape of the layer, and does not check if all hydrological properties are defined for cells that had a thickness of 0 before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_new = nlmod.layers.set_layer_top(ds.copy(deep=True), \"WAk1\", -40.0)\n",
    "compare_layer_models(ds, line, colors, ds2=ds_new, title2=\"Modified\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set layer bottom\n",
    "`nlmod.layers.set_layer_botm` sets the layer botm to a specified value or array.\n",
    "\n",
    "This method only changes the shape of the layer, and does not check if all hydrological properties are defined for cells that had a thickness of 0 before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set the botm of 'WAk1' to -70 m NAP\n",
    "ds_new = nlmod.layers.set_layer_botm(ds.copy(deep=True), \"WAk1\", -70.0)\n",
    "compare_layer_models(ds, line, colors, ds2=ds_new, title2=\"Modified\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set layer thickness\n",
    "`nlmod.layers.set_layer_thickness` sets the thickness of a layer to a specified value or array.  With a parameter called 'change' you can specify in which direction the layer is changed. The only supported option for now is 'botm', which changes the layer botm. \n",
    "\n",
    "This method only changes the shape of the layer, and does not check if all hydrological properties are defined for cells that had a thickness of 0 before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set the thickness of 'WAk1' to 20 m NAP\n",
    "ds_new = nlmod.layers.set_layer_thickness(ds.copy(deep=True), \"WAk1\", 20)\n",
    "compare_layer_models(ds, line, colors, ds2=ds_new, title2=\"Modified\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set minimum layer thickness\n",
    "`nlmod.layers.set_minimum layer_thickness` increases the thickness of a layer if the thickness is less than a specified value.  With a parameter called 'change' you can specify in which direction the layer is changed. The only supported option for now is 'botm', which changes the layer botm. \n",
    "\n",
    "This method only changes the shape of the layer, and does not check if all hydrological properties are defined for cells that had a thickness of 0 before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set the minimum thickness of 'PZWAz2' to 20 m\n",
    "ds_new = nlmod.layers.set_minimum_layer_thickness(ds.copy(deep=True), \"PZWAz2\", 20.0)\n",
    "compare_layer_models(ds, line, colors, ds2=ds_new, title2=\"Modified\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}