Improve remove_thin_layers and add tests (#296)

* Improve remove_thin_layers and add tests

* Update docstring

* Solve too-many-nested-blocks

nlmod/dims/layers.py

@@ -817,10 +817,35 @@ def set_minimum_layer_thickness(ds, layer, min_thickness, change="botm", copy=Tr
     return ds
 
 
-def remove_thin_layers(ds, min_thickness=0.1, copy=True):
-    """Remove layers from cells with a thickness less than min_thickness
+def remove_thin_layers(
+    ds, min_thickness=0.1, update_thickness_every_layer=False, copy=True
+):
+    """
+    Remove cells with a thickness less than min_thickness (setting the thickness to 0)
 
     The thickness of the removed cells is added to the first active layer below
+
+    Parameters
+    ----------
+    ds : xr,Dataset
+        Dataset containing information about layers.
+    min_thickness : float, optional
+        THe minimum thickness of a layer. The default is 0.1.
+    update_thickness_every_layer : bool, optional
+        If True, loop over the layers, from the top down, and remove thin layers, adding
+        the thickness to the first active layer below. If the thickness of this layer is
+        increased more than min_thickness (even if it originally was thinner than
+        min_thickness), the layer is kept. If update_thickness_every_layer is False, all
+        cells that originally were thinner than min_thickness are removed. The default
+        is False.
+    copy : bool, optional
+        If copy=True, data in the return value is always copied, so the original Dataset
+        is not altered. The default is True.
+    Returns
+    -------
+    ds : xr.Dataset
+        Dataset containing information about layers.
+
     """
     if "layer" in ds["top"].dims:
         msg = "remove_thin_layers does not support top with a layer dimension"
@@ -829,10 +854,12 @@ def remove_thin_layers(ds, min_thickness=0.1, copy=True):
         # make a copy, so we are sure we do not alter the original DataSet
         ds = ds.copy(deep=True)
     thickness = calculate_thickness(ds)
-    for lay_org in range(len(ds.layer)):
-        # determine where the layer is too thin
-        mask = thickness[lay_org] < min_thickness
-        if mask.any():
+    if update_thickness_every_layer:
+        for lay_org in range(len(ds.layer)):
+            # determine where the layer is too thin
+            mask = thickness[lay_org] < min_thickness
+            if not mask.any():
+                continue
             # we will set the botm to the top in these cells, so we first get the top
             if lay_org == 0:
                 top = ds["top"]
@@ -843,13 +870,22 @@ def remove_thin_layers(ds, min_thickness=0.1, copy=True):
                 if lay > lay_org:
                     # only keep cells in mask that had no thickness to begin with
                     # we need to increase the botm in these cells as well
-                    mask = mask & (thickness[lay + 1] <= 0)
+                    mask = mask & (thickness[lay] <= 0)
                     if not mask.any():
                         break
                 # set the botm equal to the top in the cells in mask
-                ds["botm"][lay].data[mask] = top.data[mask]
+                ds["botm"][lay] = ds["botm"][lay].where(~mask, top)
             # calculate the thickness again, using the new botms
             thickness = calculate_thickness(ds)
+    else:
+        mask = thickness >= min_thickness
+        # set botm of layers with a small thickness to NaN
+        ds["botm"] = ds["botm"].where(mask)
+        # fill botm of the first layer by top (setting the cell thickness to 0)
+        ds["botm"][0] = ds["botm"][0].fillna(ds["top"])
+        # fill nans in botm by copying values from higher layers
+        ds["botm"] = ds["botm"].ffill("layer")
+
     return ds
 
 

tests/test_009_layers.py

@@ -197,3 +197,38 @@ def test_insert_layer():
 
     # make sure the top of the new layer is equal to the top of the original layer
     assert (ds2["botm"].loc[new_layer] == ds1["botm"].loc[layer]).all()
+
+
+def test_remove_thin_layers():
+    # %% download regis and define min_thickness
+    regis = get_regis_horstermeer()
+
+    min_thickness = 1.0
+
+    # %% test update_thickness_every_layer = False
+    ds_new = nlmod.layers.remove_thin_layers(regis, min_thickness)
+
+    th = nlmod.layers.calculate_thickness(regis)
+    assert th.where(th > 0).min() < min_thickness
+
+    th_new = nlmod.layers.calculate_thickness(ds_new)
+    assert th_new.where(th_new > 0).min() <= min_thickness
+
+    # the following does not assert to True, as there are negative thicknesses in the original regis data
+    # assert th_new.sum() == th.sum()
+
+    # test if all active cells in the new dataset were active in the original dataset
+    assert (th.data[th_new > 0] > 0).all()
+
+    # %% test update_thickness_every_layer = True
+    ds_new2 = nlmod.layers.remove_thin_layers(
+        regis, min_thickness, update_thickness_every_layer=True
+    )
+    th_new2 = nlmod.layers.calculate_thickness(ds_new2)
+
+    assert th_new2.where(th_new2 > 0).min() <= min_thickness
+
+    assert th_new2.sum() == th_new.sum()
+
+    # test if all active cells in the new dataset were active in the original dataset
+    assert (th.data[th_new2 > 0] > 0).all()