Multiple improvements

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

@@ -239,10 +239,7 @@
     "oc = nlmod.gwf.oc(ds, gwf)\n",
     "\n",
     "# create recharge package\n",
-    "rch = nlmod.gwf.rch(ds, gwf)\n",
-    "\n",
-    "# create storage package\n",
-    "sto = nlmod.gwf.sto(ds, gwf)"
+    "rch = nlmod.gwf.rch(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/base.py

@@ -179,7 +179,7 @@ def to_model_ds(
     return ds
 
 
-def extrapolate_ds(ds, mask=None):
+def extrapolate_ds(ds, mask=None, layer="layer"):
     """Fill missing data in layermodel based on nearest interpolation.
 
     Used for ensuring layer model contains data everywhere. Useful for
@@ -190,18 +190,20 @@ def extrapolate_ds(ds, mask=None):
     ----------
     ds : xarray.DataSet
         Model layer DataSet
-    mask: np.ndarray, optional
+    mask : np.ndarray, optional
         Boolean mask for each cell, with a value of True if its value needs to
         be determined. When mask is None, it is determined from the botm-
         variable. The default is None.
+    layer : str, optional
+        The name of the layer dimension. The default is 'layer'.
 
     Returns
     -------
     ds : xarray.DataSet
         filled layermodel
     """
     if mask is None:
-        mask = np.isnan(ds["botm"]).all("layer").data
+        mask = np.isnan(ds["botm"]).all(layer).data
     if not mask.any():
         # all of the model cells are is inside the known area
         return ds
@@ -225,9 +227,9 @@ def extrapolate_ds(ds, mask=None):
         data = ds[key].data
         if ds[key].dims == dims:
             if np.isnan(data[mask]).sum() > 0:  # do not update if no NaNs
-                data[mask] = data[~mask, i]
-        elif ds[key].dims == ("layer",) + dims:
-            for lay in range(len(ds["layer"])):
+                data[mask] = data[~mask][i]
+        elif ds[key].dims == (layer,) + dims:
+            for lay in range(len(ds[layer])):
                 if np.isnan(data[lay, mask]).sum() > 0:  # do not update if no NaNs
                     data[lay, mask] = data[lay, ~mask][i]
         else:

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, inconsistency_threshold=0.001)
+
     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, inconsistency_threshold=0.001)
+
     return ds_combine
 
 
@@ -682,18 +702,11 @@ def set_model_top(ds, top, min_thickness=0.0):
     ds : xarray.Dataset
         The model dataset, containing the new top.
     """
-    if "gridtype" not in ds.attrs:
-        raise (
-            KeyError(
-                "Dataset does not have attribute 'gridtype'. "
-                "Either add attribute or use nlmod functions to build the dataset."
-            )
-        )
     if "layer" in ds["top"].dims:
         raise (ValueError("set_model_top does not support top with a layer dimension"))
     if isinstance(top, (float, int)):
         top = xr.full_like(ds["top"], top)
-    if not top.shape == ds["top"].shape:
+    elif not top.shape == ds["top"].shape:
         raise (
             ValueError("Please make sure the new top has the same shape as the old top")
         )
@@ -982,7 +995,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,51 +1015,144 @@ 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")
-        if drop_layer_dim_from_top:
-            ds["top"] = top_max
     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)
+
+    # then use ffill to fill nans in deeper layers
+    ds["botm"] = ds["botm"].ffill("layer")
 
-    botm = ds["botm"].data
-    # remove nans from botm
-    for lay in range(botm.shape[0]):
-        mask = np.isnan(botm[lay])
-        if lay == 0:
-            # by setting the botm to top_max
-            botm[lay, mask] = top_max.data[mask]
-        else:
-            # by setting the botm to the botm of the layer above
-            botm[lay, mask] = botm[lay - 1, mask]
     if "layer" in ds["top"].dims:
         # remove nans from top by setting it equal to botm
         # which sets the layer thickness to 0
-        top = ds["top"].data
-        mask = np.isnan(top)
-        top[mask] = botm[mask]
+        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,
+    inconsistency_threshold=0.0,
+    return_inconsistencies=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.
+
+    Parameters
+    ----------
+    ds : xr.DataSet
+        model DataSet
+    check : bool, optional
+        If True, checks for inconsistencies 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
+        default is False.
+    inconsistency_threshold : float, optional
+        The threshold, above which to report inconsistencies in the layer model (where
+        the botm of a layer is not equal to top of a deeper layer). The default is 0.0.
+    return_inconsistencies : bool, optional
+        if True, return the difference between the top of layers and the botm of a
+        deeper layer, at the location where inconsitencies have been reported.
+
+    Returns
+    -------
+    ds : xarray.Dataset or numpy.array
+        dataset without a layer dimension in top, if return_inconsistencies is False. If
+        return_inconsistencies is True, a numpy arrays with non-equal top and bottoms is
+        returned.
+    """
+    if "layer" in ds["top"].dims:
+        ds = fill_nan_top_botm(ds)
+        if check:
+            dz = ds["botm"][:-1].data - ds["top"][1:].data
+            inconsistencies = np.abs(dz) > inconsistency_threshold
+            n = inconsistencies.sum()
+            if n > 0:
+                msg = f"Botm of layer is not equal to top of deeper layer in {n} cells"
+                logger.warning(msg)
+                if return_inconsistencies:
+                    return np.vstack(
+                        (
+                            ds["botm"].data[:-1][inconsistencies],
+                            ds["top"].data[1:][inconsistencies],
+                        )
+                    )
+        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
 
 
@@ -1379,20 +1485,17 @@ def insert_layer(ds, name, top, bot, kh=None, kv=None, copy=True):
     needs to be inserted above the existing layer, and if the top or bottom of the
     existing layer needs to be altered.
 
-    For now, this method needs a layer model with a 3d-top, like you get using the
-    method `nlmod.read.get_regis()`, and does not function for a model Dataset with a 2d
-    (structured) or 1d (vertex) top.
-
     When comparing the height of the new layer with an existing layer, there are 7
     options:
 
     1 The new layer is entirely above the existing layer: layer is added completely
     above existing layer. When the bottom of the new layer is above the top of the
     existing layer (which can happen for the first layer), this creates a gap in the
-    layer model.
+    layer model. In the returned Dataset, this gap is added to the existing layer.
 
-    2 part of the new layer lies within an existing layer, bottom is never below: layer
-    is added above the existing layer, and the top of existing layer is lowered.
+    2 part of the new layer lies within an existing layer, but the bottom of the new
+    layer is never below the bottom of the existing layer: layer is added above the
+    existing layer, and the top of existing layer is lowered.
 
     3 there are locations where the new layer is above the bottom of the existing layer,
     but below the top of the existing layer. The new layer splits the existing layer
@@ -1441,18 +1544,21 @@ def insert_layer(ds, name, top, bot, kh=None, kv=None, copy=True):
     shape = ds["botm"].shape[1:]
     assert top.shape == shape
     assert bot.shape == shape
-    msg = "Inserting layers is only supported with a 3d top for now"
-    assert "layer" in ds["top"].dims, msg
+    if copy:
+        # make a copy, so we are sure we do not alter the original DataSet
+        ds = ds.copy(deep=True)
+    if "layer" in ds["top"].dims:
+        msg = "Top in ds has a layer dimension. insert_layer will remove the layer dimension from top in ds."
+        logger.warning(msg)
+    else:
+        ds["top"] = ds["botm"] + calculate_thickness(ds)
     if kh is not None:
         assert kh.shape == shape
     if kv is not None:
         assert kv.shape == shape
     todo = ~(np.isnan(top.data) | np.isnan(bot.data)) & ((top - bot).data > 0)
     if not todo.any():
         logger.warning(f"Thickness of new layer {name} is never larger than 0")
-    if copy:
-        # make a copy, so we are sure we do not alter the original DataSet
-        ds = ds.copy(deep=True)
     isplit = None
     for layer in ds.layer.data:
         if not todo.any():
@@ -1545,6 +1651,8 @@ def insert_layer(ds, name, top, bot, kh=None, kv=None, copy=True):
         if isplit is not None:
             isplit += 1
         ds = _insert_layer_below(ds, None, name, isplit, mask, top, bot, kh, kv, copy)
+    # remove layer dimension from top again
+    ds = remove_layer_dim_from_top(ds, inconsistency_threshold=0.001)
     return ds