Multiple improvements

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

@@ -1025,27 +1025,28 @@ def fill_top_and_bottom(ds, drop_layer_dim_from_top=True):
 
     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"]
 
-    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]
+    # 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")
+
     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"])
+
+        if drop_layer_dim_from_top:
+            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
 
     return ds
 
@@ -1379,20 +1380,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 +1439,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 +1546,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 = fill_top_and_bottom(ds, drop_layer_dim_from_top=True)
     return ds
 
 

nlmod/gwf/gwf.py

@@ -928,7 +928,7 @@ def buy(ds, gwf, pname="buy", **kwargs):
         ds, "denseref", attr="denseref", value=kwargs.pop("denseref", None)
     )
 
-    pdata = [(0, drhodc, crhoref, f"{ds.model_name}_gwt", "none")]
+    pdata = [(0, drhodc, crhoref, f"{ds.model_name}_gwt", "CONCENTRATION")]
 
     buy = flopy.mf6.ModflowGwfbuy(
         gwf,

nlmod/plot/__init__.py

@@ -5,6 +5,7 @@
     data_array,
     facet_plot,
     geotop_lithok_in_cross_section,
+    geotop_lithok_on_map,
     map_array,
     modelgrid,
     surface_water,

nlmod/plot/plot.py

@@ -2,6 +2,7 @@
 from functools import partial
 
 import flopy as fp
+import xarray as xr
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
@@ -203,7 +204,7 @@ def data_array(da, ds=None, ax=None, rotated=False, edgecolor=None, **kwargs):
         y = da.y
         if rotated:
             if ds is None:
-                raise (Exception("Supply model dataset (ds) for grid information"))
+                raise (ValueError("Supply model dataset (ds) for grid information"))
             if "angrot" in ds.attrs and ds.attrs["angrot"] != 0.0:
                 affine = get_affine_mod_to_world(ds)
                 x, y = affine * np.meshgrid(x, y)
@@ -223,7 +224,7 @@ def geotop_lithok_in_cross_section(
         The voxel-dataset from GeoTOP. It is downloaded with the method
         `nlmod.read.geotop.get_geotop()` if None. The default is None.
     ax : matplotlib.Axes, optional
-        The axes in whcih the cross-section is plotted. Will default to the current axes
+        The axes in which the cross-section is plotted. Will default to the current axes
         if None. The default is None.
     legend : bool, optional
         When True, add a legend to the plot with the lithology-classes. The default is
@@ -259,29 +260,95 @@ def geotop_lithok_in_cross_section(
         lithok_props = geotop.get_lithok_props()
 
     cs = DatasetCrossSection(gt, line, layer="z", ax=ax, **kwargs)
-    lithoks = gt["lithok"].values
-    lithok_un = np.unique(lithoks[~np.isnan(lithoks)])
-    array = np.full(lithoks.shape, np.NaN)
-
-    colors = []
-    for i, lithok in enumerate(lithok_un):
-        lithok = int(lithok)
-        array[lithoks == lithok] = i
-        colors.append(lithok_props.at[lithok, "color"])
-    cmap = ListedColormap(colors)
-    norm = Normalize(-0.5, np.nanmax(array) + 0.5)
+    array, cmap, norm = _get_geotop_cmap_and_norm(gt["lithok"], lithok_props)
     cs.plot_array(array, norm=norm, cmap=cmap)
     if legend:
         # make a legend with dummy handles
-        handles = []
-        for i, lithok in enumerate(lithok_un):
-            label = lithok_props.at[int(lithok), "name"]
-            handles.append(Patch(facecolor=colors[i], label=label))
-        ax.legend(handles=handles, loc=legend_loc)
+        _add_geotop_lithok_legend(lithok_props, ax, lithok=gt["lithok"], loc=legend_loc)
 
     return cs
 
 
+def geotop_lithok_on_map(
+    gt, z=None, ax=None, legend=True, legend_loc=None, lithok_props=None, **kwargs
+):
+    """PLot the lithoclass-data of GeoTOP on a map.
+
+    Parameters
+    ----------
+    gt : xr.Dataset or xr.DataArray
+        The voxel-dataset from GeoTOP or the lithok-DataArray from GeoTOP.
+    z : float, optional
+        The z-value of the lithok-data to plot on the map.
+    ax : matplotlib.Axes, optional
+        The axes in which the array is plotted. Will default to the current axes if
+        None. The default is None.
+    legend : bool, optional
+        When True, add a legend to the plot with the lithology-classes. The default is
+        True.
+    legend_loc : None or str, optional
+        The location of the legend. See matplotlib documentation. The default is None.
+    lithok_props : pd.DataFrame, optional
+        A DataFrame containing the properties of the lithoclasses.
+        Will call nlmod.read.geotop.get_lithok_props() when None. The default is None.
+
+    **kwargs : dict
+        kwargs are passed onto ax.pcolormesh.
+
+    Returns
+    -------
+    qm : matplotlib.collections.QuadMesh
+
+    """
+    if ax is None:
+        ax = plt.gca()
+    if isinstance(gt, xr.Dataset):
+        lithok = gt["lithok"]
+    else:
+        lithok = gt
+    if "z" in lithok.dims:
+        if z is None:
+            raise (ValueError("Select a depth (z) in the GeoTOP data first"))
+        lithok = lithok.sel(z=z, method="nearest")
+    if lithok_props is None:
+        lithok_props = geotop.get_lithok_props()
+
+    array, cmap, norm = _get_geotop_cmap_and_norm(lithok, lithok_props)
+    qm = ax.pcolormesh(lithok.x, lithok.y, array, norm=norm, cmap=cmap, **kwargs)
+    if legend:
+        # make a legend with dummy handles
+        _add_geotop_lithok_legend(lithok_props, ax, lithok=lithok, loc=legend_loc)
+    return qm
+
+
+def _add_geotop_lithok_legend(lithok_props, ax, lithok=None, **kwargs):
+    handles = []
+    if lithok is None:
+        lithoks = lithok_props.index
+    else:
+        lithoks = np.unique(lithok)
+        lithoks = lithoks[~np.isnan(lithoks)]
+    for index in lithoks:
+        color = lithok_props.at[index, "color"]
+        label = lithok_props.at[int(index), "name"]
+        handles.append(Patch(facecolor=color, label=label))
+    return ax.legend(handles=handles, **kwargs)
+
+
+def _get_geotop_cmap_and_norm(lithok, lithok_props):
+    lithok_un = np.unique(lithok)
+    lithok_un = lithok_un[~np.isnan(lithok_un)]
+    array = np.full(lithok.shape, np.NaN)
+    colors = []
+    for i, ilithok in enumerate(lithok_un):
+        ilithok = int(ilithok)
+        array[lithok == ilithok] = i
+        colors.append(lithok_props.at[ilithok, "color"])
+    cmap = ListedColormap(colors)
+    norm = Normalize(-0.5, np.nanmax(array) + 0.5)
+    return array, cmap, norm
+
+
 def _get_figure(ax=None, da=None, ds=None, figsize=None, rotated=True, extent=None):
     # figure
     if ax is not None:

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
+from ..dims.layers import insert_layer, remove_layer, fill_top_and_bottom
 from ..util import MissingValueError
 
 logger = logging.getLogger(__name__)
@@ -61,7 +61,11 @@ def get_kh_kv_table(kind="Brabant"):
 
 @cache.cache_netcdf
 def to_model_layers(
-    geotop_ds, strat_props=None, method_geulen="add_to_layer_below", **kwargs
+    geotop_ds,
+    strat_props=None,
+    method_geulen="add_to_layer_below",
+    drop_layer_dim_from_top=True,
+    **kwargs,
 ):
     """Convert geotop voxel dataset to layered dataset.
 
@@ -86,6 +90,11 @@ def to_model_layers(
         layers, which can fail if a 'geul' is locally both below the top and above the
         bottom of another layer (splitting the layer in two, which is not supported).
         The default is "add_to_layer_below".
+    drop_layer_dim_from_top : bool, optional
+        When True, fill NaN values in top and botm and drop the layer dimension from
+        top. This will transform top and botm to the data model in MODFLOW. An advantage
+        of this data model is that the layer model is consistent by definition, with no
+        possibilities of gaps between layers. The default is True.
     kwargs : dict
         Kwargs are passed to `aggregate_to_ds()`
 
@@ -205,6 +214,9 @@ 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)
+
     if "kh" in geotop_ds and "kv" in geotop_ds:
         aggregate_to_ds(geotop_ds, ds, **kwargs)
 
@@ -235,7 +247,7 @@ def get_geotop(extent, url=GEOTOP_URL, probabilities=False):
         url of geotop netcdf file. The default is
         http://www.dinodata.nl/opendap/GeoTOP/geotop.nc
     probabilities : bool, optional
-        if True, download probability data
+        if True, keep probability data in returned Dataset. The default is False.
 
     Returns
     -------