Issue #1164 cleanup utils (#1172)

Fixes #1164


# Description
This implements the following cleanup utilities for the robin boundary
conditions, namely the DRN, GHB, and RIV.

- Cleanup helper functions, which are part of the public API (so that
people could also use them outside MODFLOW 6)
- Cleanup methods to the River, GeneralHeadBoundary, Drain package for
convenience. These call the corresponding helper function. Save the user
the hassle of manually finding the right arguments to clean up grids.
- Add a utility method: Call function on grids, which separates grid
variables from settings, calls a function on them, and then merges these
dicts again.
- Converted the fixtures in ``test_mf6_riv`` to regular functions, as
they were used as such anyway. Without this, RivDisCases could not be
used outside the ``test_mf6_riv`` module, now it can be imported.

# Checklist
<!---
Before requesting review, please go through this 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

docs/api/changelog.rst

@@ -51,6 +51,14 @@ Added
   :class:`imod.mf6.HorizontalFlowBarrierHydraulicCharacteristic`.
 - :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.
+- :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.
+
 
 Removed
 ~~~~~~~

docs/api/mf6.rst

@@ -69,8 +69,14 @@ Flow Packages
     Buoyancy
     ConstantHead
     Drainage
+    Drainage.mask
+    Drainage.regrid_like
+    Drainage.cleanup
     Evapotranspiration
     GeneralHeadBoundary
+    GeneralHeadBoundary.mask
+    GeneralHeadBoundary.regrid_like
+    GeneralHeadBoundary.cleanup
     HorizontalFlowBarrierHydraulicCharacteristic
     HorizontalFlowBarrierMultiplier
     HorizontalFlowBarrierResistance
@@ -83,6 +89,9 @@ Flow Packages
     NodePropertyFlow
     Recharge
     River
+    River.mask
+    River.regrid_like
+    River.cleanup
     SpecificStorage
     StorageCoefficient
     UnsaturatedZoneFlow

docs/api/prepare.rst

@@ -47,3 +47,7 @@ Prepare model input
     distribute_drn_conductance
     distribute_ghb_conductance
     distribute_riv_conductance
+
+    cleanup_drn
+    cleanup_ghb
+    cleanup_riv

imod/mf6/drn.py

@@ -4,9 +4,10 @@
 
 import numpy as np
 
-from imod.logging import init_log_decorator
+from imod.logging import init_log_decorator, standard_log_decorator
 from imod.mf6.boundary_condition import BoundaryCondition
 from imod.mf6.dis import StructuredDiscretization
+from imod.mf6.disv import VerticesDiscretization
 from imod.mf6.interfaces.iregridpackage import IRegridPackage
 from imod.mf6.npf import NodePropertyFlow
 from imod.mf6.regrid.regrid_schemes import DrainageRegridMethod
@@ -15,6 +16,7 @@
     _regrid_package_data,
 )
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA
+from imod.prepare.cleanup import cleanup_drn
 from imod.prepare.topsystem.allocation import ALLOCATION_OPTION, allocate_drn_cells
 from imod.prepare.topsystem.conductance import (
     DISTRIBUTING_OPTION,
@@ -163,6 +165,20 @@ def _validate(self, schemata, **kwargs):
 
         return errors
 
+    @standard_log_decorator()
+    def cleanup(self, dis: StructuredDiscretization | VerticesDiscretization) -> None:
+        """
+        Clean up package inplace. This method calls
+        :func:`imod.prepare.cleanup.cleanup_drn`, see documentation of that
+        function for details on cleanup.
+
+        dis: imod.mf6.StructuredDiscretization | imod.mf6.VerticesDiscretization
+            Model discretization package.
+        """
+        dis_dict = {"idomain": dis.dataset["idomain"]}
+        cleaned_dict = self._call_func_on_grids(cleanup_drn, dis_dict)
+        super().__init__(cleaned_dict)
+
     @classmethod
     def from_imod5_data(
         cls,

imod/mf6/ghb.py

@@ -4,8 +4,10 @@
 
 import numpy as np
 
-from imod.logging import init_log_decorator
+from imod.logging import init_log_decorator, standard_log_decorator
 from imod.mf6.boundary_condition import BoundaryCondition
+from imod.mf6.dis import StructuredDiscretization
+from imod.mf6.disv import VerticesDiscretization
 from imod.mf6.interfaces.iregridpackage import IRegridPackage
 from imod.mf6.npf import NodePropertyFlow
 from imod.mf6.regrid.regrid_schemes import (
@@ -14,6 +16,7 @@
 )
 from imod.mf6.utilities.regrid import RegridderWeightsCache, _regrid_package_data
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA
+from imod.prepare.cleanup import cleanup_ghb
 from imod.prepare.topsystem.allocation import ALLOCATION_OPTION, allocate_ghb_cells
 from imod.prepare.topsystem.conductance import (
     DISTRIBUTING_OPTION,
@@ -165,6 +168,20 @@ def _validate(self, schemata, **kwargs):
 
         return errors
 
+    @standard_log_decorator()
+    def cleanup(self, dis: StructuredDiscretization | VerticesDiscretization) -> None:
+        """
+        Clean up package inplace. This method calls
+        :func:`imod.prepare.cleanup.cleanup_ghb`, see documentation of that
+        function for details on cleanup.
+
+        dis: imod.mf6.StructuredDiscretization | imod.mf6.VerticesDiscretization
+            Model discretization package.
+        """
+        dis_dict = {"idomain": dis.dataset["idomain"]}
+        cleaned_dict = self._call_func_on_grids(cleanup_ghb, dis_dict)
+        super().__init__(cleaned_dict)
+
     @classmethod
     def from_imod5_data(
         cls,

imod/mf6/package.py

@@ -4,7 +4,7 @@
 import pathlib
 from collections import defaultdict
 from copy import deepcopy
-from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, Union
 
 import cftime
 import jinja2
@@ -81,6 +81,9 @@ def sel(self):
             f"{type(self).__name__}(**{self._pkg_id}.dataset.sel(**selection))"
         )
 
+    def cleanup(self, dis: Any):
+        raise NotImplementedError("Method not implemented for this package.")
+
     @staticmethod
     def _valid(value: Any) -> bool:
         return _is_valid(value)
@@ -630,6 +633,28 @@ def get_non_grid_data(self, grid_names: list[str]) -> dict[str, Any]:
                 result[name] = self.dataset[name].values[()]
         return result
 
+    def _call_func_on_grids(
+        self, func: Callable, dis: dict
+    ) -> dict[str, GridDataArray]:
+        """
+        Call function on dictionary of grids and merge settings back into
+        dictionary.
+
+        Parameters
+        ----------
+        func: Callable
+            Function to call on all grids
+        """
+        grid_varnames = list(self._write_schemata.keys())
+        grids = {
+            varname: self.dataset[varname]
+            for varname in grid_varnames
+            if varname in self.dataset.keys()
+        }
+        cleaned_grids = func(**dis, **grids)
+        settings = self.get_non_grid_data(grid_varnames)
+        return cleaned_grids | settings
+
     def is_splitting_supported(self) -> bool:
         return True
 

imod/mf6/riv.py

@@ -6,10 +6,11 @@
 import xarray as xr
 
 from imod import logging
-from imod.logging import init_log_decorator
+from imod.logging import init_log_decorator, standard_log_decorator
 from imod.logging.loglevel import LogLevel
 from imod.mf6.boundary_condition import BoundaryCondition
 from imod.mf6.dis import StructuredDiscretization
+from imod.mf6.disv import VerticesDiscretization
 from imod.mf6.drn import Drainage
 from imod.mf6.interfaces.iregridpackage import IRegridPackage
 from imod.mf6.regrid.regrid_schemes import RiverRegridMethod
@@ -18,6 +19,7 @@
     _regrid_package_data,
 )
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA
+from imod.prepare.cleanup import cleanup_riv
 from imod.prepare.topsystem.allocation import ALLOCATION_OPTION, allocate_riv_cells
 from imod.prepare.topsystem.conductance import (
     DISTRIBUTING_OPTION,
@@ -184,6 +186,20 @@ def _validate(self, schemata, **kwargs):
 
         return errors
 
+    @standard_log_decorator()
+    def cleanup(self, dis: StructuredDiscretization | VerticesDiscretization) -> None:
+        """
+        Clean up package inplace. This method calls
+        :func:`imod.prepare.cleanup.cleanup_riv`, see documentation of that
+        function for details on cleanup.
+
+        dis: imod.mf6.StructuredDiscretization | imod.mf6.VerticesDiscretization
+            Model discretization package.
+        """
+        dis_dict = {"idomain": dis.dataset["idomain"], "bottom": dis.dataset["bottom"]}
+        cleaned_dict = self._call_func_on_grids(cleanup_riv, dis_dict)
+        super().__init__(cleaned_dict)
+
     @classmethod
     def from_imod5_data(
         cls,

imod/mf6/wel.py

@@ -30,6 +30,7 @@
 from imod.mf6.validation import validation_pkg_error_message
 from imod.mf6.write_context import WriteContext
 from imod.prepare import assign_wells
+from imod.prepare.cleanup import cleanup_wel
 from imod.prepare.layer import create_layered_top
 from imod.schemata import (
     AllValueSchema,
@@ -971,6 +972,9 @@ def _validate_imod5_depth_information(
             logger.log(loglevel=LogLevel.ERROR, message=log_msg, additional_depth=2)
             raise ValueError(log_msg)
 
+    def cleanup(self):
+        self.dataset = cleanup_wel(self.dataset)
+
 
 class LayeredWell(GridAgnosticWell):
     """

imod/prepare/__init__.py

@@ -14,6 +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.hfb import (
     linestring_to_square_zpolygons,
     linestring_to_trapezoid_zpolygons,