Issue #1187 cleanup wells (#1206)

Fixes #1187

# Description
Adds cleanup utility for wells, as well as a method on the ``Well``
object. ``LayeredWell`` doesn't have ``cleanup`` method yet, as cleaning
up almost purely focuses on adjusting filter screen levels for now. The
function as well as the method are a bit more involved than the other
cleanup utilities, because we are dealing here with a comparison of
point data and grids.

# Checklist
- [x] Links to correct issue
- [x] Update changelog, if changes affect users
- [x] PR title starts with ``Issue #nr``, e.g. ``Issue #737``
- [x] Unit tests were added
- [ ] **If feature added**: Added/extended example

Author: JoerivanEngelen

docs/api/changelog.rst

@@ -56,12 +56,13 @@ Added
 - :class:`imod.mf6.LayeredWell` to specify wells directly to layers instead
   assigning them with filter depths.
 - :func:`imod.prepare.cleanup_drn`, :func:`imod.prepare.cleanup_ghb`,
-  :func:`imod.prepare.cleanup_riv`. These are utility functions to clean up
-  drainage, general head boundaries, and rivers, respectively.
+  :func:`imod.prepare.cleanup_riv`, :func:`imod.prepare.cleanup_wel`. These are
+  utility functions to clean up drainage, general head boundaries, and rivers,
+  respectively.
 - :meth:`imod.mf6.Drainage.cleanup`,
-  :meth:`imod.mf6.GeneralHeadboundary.cleanup`, :meth:`imod.mf6.River.cleanup`
-  convenience methods to call the corresponding cleanup utility functions with
-  the appropriate arguments.
+  :meth:`imod.mf6.GeneralHeadboundary.cleanup`, :meth:`imod.mf6.River.cleanup`,
+  :meth:`imod.mf6.Well.cleanup` convenience methods to call the corresponding
+  cleanup utility functions with the appropriate arguments.
 
 
 Removed

docs/api/mf6.rst

@@ -97,6 +97,7 @@ Flow Packages
     StorageCoefficient
     UnsaturatedZoneFlow
     Well
+    Well.cleanup
     Well.from_imod5_data
     Well.mask
     Well.regrid_like

docs/api/prepare.rst

@@ -51,3 +51,4 @@ Prepare model input
     cleanup_drn
     cleanup_ghb
     cleanup_riv
+    cleanup_wel

imod/mf6/wel.py

@@ -15,6 +15,7 @@
 import xugrid as xu
 
 import imod
+import imod.mf6.utilities
 from imod.logging import init_log_decorator, logger
 from imod.logging.logging_decorators import standard_log_decorator
 from imod.logging.loglevel import LogLevel
@@ -28,6 +29,7 @@
 from imod.mf6.mf6_wel_adapter import Mf6Wel
 from imod.mf6.package import Package
 from imod.mf6.utilities.dataset import remove_inactive
+from imod.mf6.utilities.grid import broadcast_to_full_domain
 from imod.mf6.validation import validation_pkg_error_message
 from imod.mf6.write_context import WriteContext
 from imod.prepare import assign_wells
@@ -973,8 +975,43 @@ def _validate_imod5_depth_information(
             logger.log(loglevel=LogLevel.ERROR, message=log_msg, additional_depth=2)
             raise ValueError(log_msg)
 
-    def cleanup(self, _: StructuredDiscretization | VerticesDiscretization):
-        self.dataset = cleanup_wel(self.dataset)
+    @standard_log_decorator()
+    def cleanup(self, dis: StructuredDiscretization | VerticesDiscretization):
+        """
+        Clean up package inplace. This method calls
+        :func:`imod.prepare.cleanup.cleanup_wel`, see documentation of that
+        function for details on cleanup.
+
+        dis: imod.mf6.StructuredDiscretization | imod.mf6.VerticesDiscretization
+            Model discretization package.
+        """
+        # Top and bottom should be forced to grids with a x, y coordinates
+        top, bottom = broadcast_to_full_domain(**dict(dis.dataset.data_vars))
+        # Collect point variable datanames
+        point_varnames = list(self._write_schemata.keys())
+        if "concentration" not in self.dataset.keys():
+            point_varnames.remove("concentration")
+        point_varnames.append("id")
+        # Create dataset with purely point locations
+        point_ds = self.dataset[point_varnames]
+        # Take first item of irrelevant dimensions
+        point_ds = point_ds.isel(time=0, species=0, drop=True, missing_dims="ignore")
+        # Cleanup well dataframe
+        wells = point_ds.to_dataframe()
+        minimum_thickness = float(self.dataset["minimum_thickness"])
+        cleaned_wells = cleanup_wel(wells, top.isel(layer=0), bottom, minimum_thickness)
+        # Select with ids in cleaned dataframe to drop points outside grid.
+        well_ids = cleaned_wells.index
+        dataset_cleaned = self.dataset.swap_dims({"index": "id"}).sel(id=well_ids)
+        # Assign adjusted screen top and bottom
+        dataset_cleaned["screen_top"] = cleaned_wells["screen_top"]
+        dataset_cleaned["screen_bottom"] = cleaned_wells["screen_bottom"]
+        # Ensure dtype of id is preserved
+        id_type = self.dataset["id"].dtype
+        dataset_cleaned = dataset_cleaned.swap_dims({"id": "index"}).reset_coords("id")
+        dataset_cleaned["id"] = dataset_cleaned["id"].astype(id_type)
+        # Override dataset
+        self.dataset = dataset_cleaned
 
 
 class LayeredWell(GridAgnosticWell):

imod/prepare/__init__.py

@@ -14,7 +14,7 @@
 """
 
 from imod.prepare import hfb, spatial, subsoil, surface_water
-from imod.prepare.cleanup import cleanup_drn, cleanup_ghb, cleanup_riv
+from imod.prepare.cleanup import cleanup_drn, cleanup_ghb, cleanup_riv, cleanup_wel
 from imod.prepare.hfb import (
     linestring_to_square_zpolygons,
     linestring_to_trapezoid_zpolygons,

imod/prepare/cleanup.py

@@ -3,11 +3,13 @@
 from enum import Enum
 from typing import Optional
 
+import pandas as pd
 import xarray as xr
 
 from imod.mf6.utilities.mask import mask_arrays
+from imod.prepare.wells import locate_wells, validate_well_columnnames
 from imod.schemata import scalar_None
-from imod.typing import GridDataArray, GridDataset
+from imod.typing import GridDataArray
 
 
 class AlignLevelsMode(Enum):
@@ -83,7 +85,7 @@ def cleanup_riv(
     idomain: xarray.DataArray | xugrid.UgridDataArray
         MODFLOW 6 model domain. idomain==1 is considered active domain.
     bottom: xarray.DataArray | xugrid.UgridDataArray
-        Grid with`model bottoms
+        Grid with model bottoms
     stage: xarray.DataArray | xugrid.UgridDataArray
         Grid with river stages
     conductance: xarray.DataArray | xugrid.UgridDataArray
@@ -209,11 +211,128 @@ def cleanup_ghb(
     return _cleanup_robin_boundary(idomain, output_dict)
 
 
-def cleanup_wel(wel_ds: GridDataset):
+def _locate_wells_in_bounds(
+    wells: pd.DataFrame, top: GridDataArray, bottom: GridDataArray
+) -> tuple[pd.DataFrame, pd.Series, pd.Series]:
     """
-    Clean up wells
+    Locate wells in model bounds, wells outside bounds are dropped. Returned
+    dataframes and series have well "id" as index.
 
-    - Removes wells where the screen bottom elevation exceeds screen top.
+    Returns
+    -------
+    wells_in_bounds: pd.DataFrame
+        wells in model boundaries. Has "id" as index.
+    xy_top_series: pd.Series
+        model top at well xy location. Has "id" as index.
+    xy_base_series: pd.Series
+        model base at well xy location. Has "id" as index.
     """
-    deactivate = wel_ds["screen_top"] < wel_ds["screen_bottom"]
-    return wel_ds.where(~deactivate, drop=True)
+    id_in_bounds, xy_top, xy_bottom, _ = locate_wells(
+        wells, top, bottom, validate=False
+    )
+    xy_base_model = xy_bottom.isel(layer=-1, drop=True)
+
+    # Assign id as coordinates
+    xy_top = xy_top.assign_coords(id=("index", id_in_bounds))
+    xy_base_model = xy_base_model.assign_coords(id=("index", id_in_bounds))
+    # Create pandas dataframes/series with "id" as index.
+    xy_top_series = xy_top.to_dataframe(name="top").set_index("id")["top"]
+    xy_base_series = xy_base_model.to_dataframe(name="bottom").set_index("id")["bottom"]
+    wells_in_bounds = wells.set_index("id").loc[id_in_bounds]
+    return wells_in_bounds, xy_top_series, xy_base_series
+
+
+def _clip_filter_screen_to_surface_level(
+    cleaned_wells: pd.DataFrame, xy_top_series: pd.Series
+) -> pd.DataFrame:
+    cleaned_wells["screen_top"] = cleaned_wells["screen_top"].clip(upper=xy_top_series)
+    return cleaned_wells
+
+
+def _drop_wells_below_model_base(
+    cleaned_wells: pd.DataFrame, xy_base_series: pd.Series
+) -> pd.DataFrame:
+    is_below_base = cleaned_wells["screen_top"] >= xy_base_series
+    return cleaned_wells.loc[is_below_base]
+
+
+def _clip_filter_bottom_to_model_base(
+    cleaned_wells: pd.DataFrame, xy_base_series: pd.Series
+) -> pd.DataFrame:
+    cleaned_wells["screen_bottom"] = cleaned_wells["screen_bottom"].clip(
+        lower=xy_base_series
+    )
+    return cleaned_wells
+
+
+def _set_inverted_filters_to_point_filters(cleaned_wells: pd.DataFrame) -> pd.DataFrame:
+    # Convert all filters where screen bottom exceeds screen top to
+    # point filters
+    cleaned_wells["screen_bottom"] = cleaned_wells["screen_bottom"].clip(
+        upper=cleaned_wells["screen_top"]
+    )
+    return cleaned_wells
+
+
+def _set_ultrathin_filters_to_point_filters(
+    cleaned_wells: pd.DataFrame, minimum_thickness: float
+) -> pd.DataFrame:
+    not_ultrathin_layer = (
+        cleaned_wells["screen_top"] - cleaned_wells["screen_bottom"]
+    ) > minimum_thickness
+    cleaned_wells["screen_bottom"] = cleaned_wells["screen_bottom"].where(
+        not_ultrathin_layer, cleaned_wells["screen_top"]
+    )
+    return cleaned_wells
+
+
+def cleanup_wel(
+    wells: pd.DataFrame,
+    top: GridDataArray,
+    bottom: GridDataArray,
+    minimum_thickness: float = 0.05,
+) -> pd.DataFrame:
+    """
+    Clean up dataframe with wells, fixes some common mistakes in the following
+    order:
+
+    1. Wells outside grid bounds are dropped
+    2. Filters above surface level are set to surface level
+    3. Drop wells with filters entirely below base
+    4. Clip filter screen_bottom to model base
+    5. Clip filter screen_bottom to screen_top
+    6. Well filters thinner than minimum thickness are made point filters
+
+    Parameters
+    ----------
+    wells: pandas.Dataframe
+        Dataframe with wells to be cleaned up. Requires columns ``"x", "y",
+        "id", "screen_top", "screen_bottom"``
+    top: xarray.DataArray | xugrid.UgridDataArray
+        Grid with model top
+    bottom: xarray.DataArray | xugrid.UgridDataArray
+        Grid with model bottoms
+    minimum_thickness: float
+        Minimum thickness, filter thinner than this thickness are set to point
+        filters
+
+    Returns
+    -------
+    pandas.DataFrame
+        Cleaned well dataframe.
+    """
+    validate_well_columnnames(
+        wells, names={"x", "y", "id", "screen_top", "screen_bottom"}
+    )
+
+    cleaned_wells, xy_top_series, xy_base_series = _locate_wells_in_bounds(
+        wells, top, bottom
+    )
+    cleaned_wells = _clip_filter_screen_to_surface_level(cleaned_wells, xy_top_series)
+    cleaned_wells = _drop_wells_below_model_base(cleaned_wells, xy_base_series)
+    cleaned_wells = _clip_filter_bottom_to_model_base(cleaned_wells, xy_base_series)
+    cleaned_wells = _set_inverted_filters_to_point_filters(cleaned_wells)
+    cleaned_wells = _set_ultrathin_filters_to_point_filters(
+        cleaned_wells, minimum_thickness
+    )
+    return cleaned_wells

imod/prepare/wells.py

@@ -84,7 +84,7 @@ def locate_wells(
     wells: pd.DataFrame,
     top: GridDataArray,
     bottom: GridDataArray,
-    k: Optional[GridDataArray],
+    k: Optional[GridDataArray] = None,
     validate: bool = True,
 ) -> tuple[
     npt.NDArray[np.object_], GridDataArray, GridDataArray, float | GridDataArray
@@ -126,6 +126,22 @@ def locate_wells(
     return id_in_bounds, xy_top, xy_bottom, xy_k
 
 
+def validate_well_columnnames(
+    wells: pd.DataFrame, names: set = {"x", "y", "id"}
+) -> None:
+    missing = names.difference(wells.columns)
+    if missing:
+        raise ValueError(f"Columns are missing in wells dataframe: {missing}")
+
+
+def validate_arg_types_equal(**kwargs):
+    types = [type(arg) for arg in (kwargs.values()) if arg is not None]
+    if len(set(types)) != 1:
+        members = ", ".join([t.__name__ for t in types])
+        names = ", ".join(kwargs.keys())
+        raise TypeError(f"{names} should be of the same type, received: {members}")
+
+
 def assign_wells(
     wells: pd.DataFrame,
     top: GridDataArray,
@@ -166,19 +182,9 @@ def assign_wells(
         Wells with rate subdivided per layer. Contains the original columns of
         ``wells``, as well as layer, overlap, transmissivity.
     """
-
-    names = {"x", "y", "id", "top", "bottom", "rate"}
-    missing = names.difference(wells.columns)
-    if missing:
-        raise ValueError(f"Columns are missing in wells dataframe: {missing}")
-
-    types = [type(arg) for arg in (top, bottom, k) if arg is not None]
-    if len(set(types)) != 1:
-        members = ",".join([t.__name__ for t in types])
-        raise TypeError(
-            "top, bottom, and optionally k should be of the same type, "
-            f"received: {members}"
-        )
+    columnnames = {"x", "y", "id", "top", "bottom", "rate"}
+    validate_well_columnnames(wells, columnnames)
+    validate_arg_types_equal(top=top, bottom=bottom, k=k)
 
     id_in_bounds, xy_top, xy_bottom, xy_k = locate_wells(
         wells, top, bottom, k, validate

imod/tests/test_mf6/test_mf6_riv.py

@@ -71,6 +71,14 @@ def case_unstructured(self):
         return make_dict_unstructured(riv_dict())
 
 
+class DisCases:
+    def case_structured(self):
+        return dis_dict()
+
+    def case_unstructured(self):
+        return make_dict_unstructured(dis_dict())
+
+
 class RivDisCases:
     def case_structured(self):
         return riv_dict(), dis_dict()

imod/tests/test_mf6/test_mf6_wel.py

@@ -394,6 +394,32 @@ def test_is_empty(well_high_lvl_test_data_transient):
     assert wel_empty.is_empty()
 
 
+def test_cleanup(basic_dis, well_high_lvl_test_data_transient):
+    # Arrange
+    wel = imod.mf6.Well(*well_high_lvl_test_data_transient)
+    ds_original = wel.dataset.copy()
+    idomain, top, bottom = basic_dis
+    top = top.isel(layer=0, drop=True)
+    deep_offset = 100.0
+    dis_normal = imod.mf6.StructuredDiscretization(top, bottom, idomain)
+    dis_deep = imod.mf6.StructuredDiscretization(
+        top - deep_offset, bottom - deep_offset, idomain
+    )
+
+    # Nothing to be cleaned up with default discretization
+    wel.cleanup(dis_normal)
+    xr.testing.assert_identical(ds_original, wel.dataset)
+
+    # Cleanup
+    wel.cleanup(dis_deep)
+    assert not ds_original.identical(wel.dataset)
+    # Wells filters should be placed downwards at surface level as point filters
+    np.testing.assert_array_almost_equal(
+        wel.dataset["screen_top"], wel.dataset["screen_bottom"]
+    )
+    np.testing.assert_array_almost_equal(wel.dataset["screen_top"], top - deep_offset)
+
+
 class ClipBoxCases:
     @staticmethod
     def case_clip_xy(parameterizable_basic_dis):
@@ -959,6 +985,26 @@ def test_import_and_convert_to_mf6(imod5_dataset, tmp_path, wel_class):
     mf6_well.write("wel", [], write_context)
 
 
+@parametrize("wel_class", [Well])
+@pytest.mark.usefixtures("imod5_dataset")
+def test_import_and_cleanup(imod5_dataset, wel_class: Well):
+    data = imod5_dataset[0]
+    target_dis = StructuredDiscretization.from_imod5_data(data)
+
+    ntimes = 8399
+    times = list(pd.date_range(datetime(1989, 1, 1), datetime(2013, 1, 1), ntimes + 1))
+
+    # Import grid-agnostic well from imod5 data (it contains 1 well)
+    wel = wel_class.from_imod5_data("wel-WELLS_L3", data, times)
+    assert len(wel.dataset.coords["time"]) == ntimes
+    # Cleanup
+    wel.cleanup(target_dis)
+    # Nothing to be cleaned, single well point is located properly, test that
+    # time coordinate has not been dropped.
+    assert "time" in wel.dataset.coords
+    assert len(wel.dataset.coords["time"]) == ntimes
+
+
 @parametrize("wel_class", [Well, LayeredWell])
 @pytest.mark.usefixtures("well_regular_import_prj")
 def test_import_multiple_wells(well_regular_import_prj, wel_class):