Process comments from @dbrakenhoff and other fixes

fix split_layer_ds and combine_layer_ds, and add tests
small fixes to notebooks

docs/examples/04_modifying_layermodels.ipynb

@@ -268,9 +268,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "layer_names = []\n",
-    "colors_new = {}\n",
-    "\n",
     "for j, i in split_reindexer.items():\n",
     "    if j not in colors:\n",
     "        colors[j] = colors[i]"

docs/examples/11_grid_rotation.ipynb

@@ -240,9 +240,6 @@
     "\n",
     "# create recharge package\n",
     "rch = nlmod.gwf.rch(ds, gwf)\n",
-    "\n",
-    "# create storage package\n",
-    "sto = nlmod.gwf.sto(ds, gwf)"
    ]
   },
   {

docs/examples/12_layer_generation.ipynb

@@ -71,7 +71,7 @@
    "outputs": [],
    "source": [
     "f, ax = nlmod.plot.get_map(extent, figsize=5)\n",
-    "nlmod.plot.data_array(regis[\"top\"].max(\"layer\"), edgecolor=\"k\")"
+    "nlmod.plot.data_array(regis[\"top\"], edgecolor=\"k\")"
    ]
   },
   {

docs/examples/17_unsaturated_zone_flow.ipynb

@@ -95,7 +95,7 @@
     "y = np.floor(y / 100) * 100\n",
     "dx = dy = 100\n",
     "extent = [x, x + dx, y, y + dy]\n",
-    "regis = nlmod.read.regis.get_regis(extent, cachename=\"regis.nc\", cachedir=cachedir)"
+    "regis = nlmod.read.regis.get_regis(extent, drop_layer_dim_from_top=False, cachename=\"regis.nc\", cachedir=cachedir)"
    ]
   },
   {

nlmod/dims/layers.py

@@ -240,6 +240,13 @@ def split_layers_ds(
 
     logger.info(f"Splitting layers {list(split_dict)}")
 
+    if "layer" in ds["top"].dims:
+        msg = "Top in ds has a layer dimension. split_layers_ds will remove the layer dimension from top in ds."
+        logger.warning(msg)
+    else:
+        ds = ds.copy()
+        ds["top"] = ds["botm"] + calculate_thickness(ds)
+
     layers_org = layers.copy()
     # add extra layers (keep the original ones for now, as we will copy data first)
     for lay0 in split_dict:
@@ -267,6 +274,9 @@ def split_layers_ds(
     # drop the original layers
     ds = ds.drop_sel(layer=list(split_dict))
 
+    # remove layer dimension from top again
+    ds = remove_layer_dim_from_top(ds)
+
     if return_reindexer:
         # determine reindexer
         reindexer = OrderedDict(zip(layers, layers_org))
@@ -559,6 +569,13 @@ def combine_layers_ds(
     # calculate new tops/bots
     logger.info("Calculating new layer tops and bottoms...")
 
+    if "layer" in ds["top"].dims:
+        msg = "Top in ds has a layer dimension. combine_layers_ds will remove the layer dimension from top in ds."
+        logger.warning(msg)
+    else:
+        ds = ds.copy()
+        ds["top"] = ds["botm"] + calculate_thickness(ds)
+
     da_dict = {}
 
     new_top, new_bot, reindexer = layer_combine_top_bot(
@@ -604,6 +621,9 @@ def combine_layers_ds(
     logger.info("Done! Created new dataset with combined layers!")
     ds_combine = xr.Dataset(da_dict, attrs=attrs)
 
+    # remove layer dimension from top again
+    ds = remove_layer_dim_from_top(ds)
+
     return ds_combine
 
 
@@ -982,7 +1002,7 @@ def fill_nan_top_botm_kh_kv(
     """
 
     # 1
-    ds = fill_top_and_bottom(ds)
+    ds = remove_layer_dim_from_top(ds)
 
     # 2
     if remove_nan_layers:
@@ -1002,32 +1022,27 @@ def fill_nan_top_botm_kh_kv(
     return ds
 
 
-def fill_top_and_bottom(ds, drop_layer_dim_from_top=True):
+def fill_nan_top_botm(ds):
     """
-    Remove Nans in botm variable, and change top from 3d to 2d if necessary.
+    Remove Nans in non-existent layers in botm and top variables
+
+    The NaNs are removed by setting the value to the top and botm of higher/lower
+    layers that do exist.
 
     Parameters
     ----------
     ds : xr.DataSet
         model DataSet
-    drop_layer_dim_from_top : bool, optional
-        If True and top contains a layer dimension, set top to the top of the upper
-        layer (line the definition in MODFLOW). This removes redundant data, as the top
-        of all layers exept the most upper one is also defined as the bottom of previous
-        layers. The default is True.
 
     Returns
     -------
     ds : xarray.Dataset
-        dataset with filled top and bottom data according to modflow definition,
-        with 2d top and 3d bottom.
+        dataset with filled top and botm data
     """
-
     if "layer" in ds["top"].dims:
         top_max = ds["top"].max("layer")
     else:
         top_max = ds["top"]
-
     # fill nans in botm of the first layer
     ds["botm"][0] = ds["botm"][0].where(~ds["botm"][0].isnull(), top_max)
 
@@ -1038,16 +1053,92 @@ def fill_top_and_bottom(ds, drop_layer_dim_from_top=True):
         # remove nans from top by setting it equal to botm
         # which sets the layer thickness to 0
         ds["top"] = ds["top"].where(~ds["top"].isnull(), ds["botm"])
+    return ds
+
+
+def set_nan_top_and_botm(ds):
+    """
+    Sets Nans for non-existent layers in botm and top variables
+
+    Nans are only added to top when it contains a layer dimension.
+
+    Parameters
+    ----------
+    ds : xr.DataSet
+        model DataSet
+
+    Returns
+    -------
+    ds : xarray.Dataset
+        dataset with nans in top and botm at non-existent layers.
+    """
+    thickness = calculate_thickness(ds)
+    mask = thickness > 0
+    if "layer" in ds["top"].dims:
+        ds["top"] = ds["top"].where(mask)
+    ds["botm"] = ds["botm"].where(mask)
+    return ds
+
+
+def remove_layer_dim_from_top(ds, check=True, set_non_existing_layers_to_nan=False):
+    """
+    Change top from 3d to 2d, removing NaNs in top and botm in the process.
+
+    This method sets variable `top` to the top of the upper layer (like the definition
+    in MODFLOW). This removes redundant data, as the top of all layers exept the most
+    upper one is also defined as the bottom of lower layers.
 
-        if drop_layer_dim_from_top:
+    Parameters
+    ----------
+    ds : xr.DataSet
+        model DataSet
+    check : bool, optional
+        If True, checks for inconsistensies in the layer model and report to logger as
+        warning. The defaults is True.
+    set_non_existing_layers_to_nan bool, optional
+        If True, sets the value of the botm-variable to NaN for non-existent layers.
+        This is not recommended, as this might break some procedures in nlmod. The
+        defaults is False.
+
+    Returns
+    -------
+    ds : xarray.Dataset
+        dataset without a layer dimension in top.
+    """
+    if "layer" in ds["top"].dims:
+        ds = fill_nan_top_botm(ds)
+        if check:
             dz = ds["botm"][:-1].data - ds["top"][1:].data
             voids = np.abs(dz) > 0
             if voids.sum() > 0:
                 n = int(voids.sum())
                 msg = f"Botm of layer is not equal to top of deeper layer in {n} cells"
                 logger.warning(msg)
-            ds["top"] = top_max
+        ds["top"] = ds["top"][0]
+    if set_non_existing_layers_to_nan:
+        ds = set_nan_top_and_botm(ds)
+    return ds
+
 
+def add_layer_dim_to_top(ds, set_non_existing_layers_to_nan=True):
+    """
+    Change top from 2d to 3d, setting top and botm to NaN for non-existent layers.
+
+    Parameters
+    ----------
+    ds : xr.DataSet
+        model DataSet
+    set_non_existing_layers_to_nan : bool, optional
+        If True, set the values of top and botm to . The default is True.
+
+    Returns
+    -------
+    ds : xarray.Dataset
+        dataset with a layer dimension in top.
+    """
+    ds["top"] = ds["botm"] + calculate_thickness(ds)
+    if set_non_existing_layers_to_nan:
+        set_nan_top_and_botm(ds)
     return ds
 
 
@@ -1547,7 +1638,7 @@ def insert_layer(ds, name, top, bot, kh=None, kv=None, copy=True):
             isplit += 1
         ds = _insert_layer_below(ds, None, name, isplit, mask, top, bot, kh, kv, copy)
     # remove layer dimension from top again
-    ds = fill_top_and_bottom(ds, drop_layer_dim_from_top=True)
+    ds = remove_layer_dim_from_top(ds)
     return ds
 
 

nlmod/read/geotop.py

@@ -8,7 +8,7 @@
 import xarray as xr
 
 from .. import NLMOD_DATADIR, cache
-from ..dims.layers import insert_layer, remove_layer, fill_top_and_bottom
+from ..dims.layers import insert_layer, remove_layer, remove_layer_dim_from_top
 from ..util import MissingValueError
 
 logger = logging.getLogger(__name__)
@@ -215,7 +215,7 @@ def to_model_layers(
     ds.attrs["geulen"] = geulen
 
     if drop_layer_dim_from_top:
-        ds = fill_top_and_bottom(ds, drop_layer_dim_from_top=drop_layer_dim_from_top)
+        ds = remove_layer_dim_from_top(ds)
 
     if "kh" in geotop_ds and "kv" in geotop_ds:
         aggregate_to_ds(geotop_ds, ds, **kwargs)

nlmod/read/regis.py

@@ -7,7 +7,7 @@
 import xarray as xr
 
 from .. import cache
-from ..dims.layers import calculate_thickness, fill_top_and_bottom
+from ..dims.layers import calculate_thickness, remove_layer_dim_from_top
 from . import geotop
 
 logger = logging.getLogger(__name__)
@@ -169,7 +169,7 @@ def get_regis(
             raise (Exception(msg))
 
     if drop_layer_dim_from_top:
-        ds = fill_top_and_bottom(ds, drop_layer_dim_from_top=drop_layer_dim_from_top)
+        ds = remove_layer_dim_from_top(ds)
 
     # slice data vars
     if variables is not None:
@@ -285,7 +285,7 @@ def add_geotop_to_regis_layers(
         rg = rg.reindex({"layer": layer_order})
 
     # remove the layer dimension from top again
-    rg = fill_top_and_bottom(rg, drop_layer_dim_from_top=True)
+    rg = remove_layer_dim_from_top(rg)
     return rg
 
 

tests/test_009_layers.py

@@ -1,5 +1,6 @@
 import os
 
+import numpy as np
 import matplotlib.pyplot as plt
 from shapely.geometry import LineString
 
@@ -52,12 +53,54 @@ def compare_layer_models(
         ax1.set_xlabel(xlabel)
 
 
-def plot_test(ds, ds_new):
-    line = LineString([(ds.extent[0], ds.extent[2]), (ds.extent[1], ds.extent[3])])
-    colors = nlmod.read.regis.get_legend()["color"].to_dict()
+def plot_test(ds, ds_new, line=None, colors=None):
+    if line is None:
+        line = LineString([(ds.extent[0], ds.extent[2]), (ds.extent[1], ds.extent[3])])
+    if colors is None:
+        colors = nlmod.read.regis.get_legend()["color"].to_dict()
     compare_layer_models(ds, line, colors, ds_new)
 
 
+def test_split_layers():
+    regis = get_regis_horstermeer()
+    split_dict = {"PZWAz2": (0.3, 0.3, 0.4), "PZWAz3": (0.2, 0.2, 0.2, 0.2, 0.2)}
+    regis_split, split_reindexer = nlmod.layers.split_layers_ds(
+        regis, split_dict, return_reindexer=True
+    )
+    assert (
+        nlmod.layers.calculate_thickness(regis).sum()
+        == nlmod.layers.calculate_thickness(regis_split).sum()
+    )
+
+    # diagonal line through extent
+    colors = nlmod.read.regis.get_legend()["color"].to_dict()
+    for j, i in split_reindexer.items():
+        if j not in colors:
+            colors[j] = colors[i]
+    plot_test(regis, regis_split, colors=colors)
+
+
+def test_add_layer_dim_to_top():
+    regis = get_regis_horstermeer()
+    nlmod.layers.add_layer_dim_to_top(regis)
+
+
+def test_combine_layers():
+    regis = get_regis_horstermeer()
+    combine_layers = [
+        tuple(np.argwhere(regis.layer.str.startswith("URz").data).squeeze().tolist()),
+        tuple(
+            np.argwhere(regis.layer.isin(["PZWAz2", "PZWAz3"]).data).squeeze().tolist()
+        ),
+    ]
+    regis_combined = nlmod.layers.combine_layers_ds(
+        regis, combine_layers, kD=None, c=None
+    )
+
+    # plot comparison
+    plot_test(regis, regis_combined)
+
+
 def test_set_layer_thickness(plot=False):
     regis = get_regis_horstermeer()
     ds = nlmod.to_model_ds(regis)