Issue #1539 fix well cleanup unstructured grids (#1540)

JoerivanEngelen · web-flow · commit 698b43467085 · 2025-05-27T09:28:38.000+02:00
Fixes #1539 # Description Changes the following: - Fix bug where UgridDataArrays where silently converted to regular DataArrays (caused by the call to ``.data_vars``). - Add unstructured grid test case to well tests where easily applicable. # 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
diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst
@@ -40,6 +40,9 @@ Fixed
   relative paths for the extra files block.
 - :class:`imod.msw.IdfMapping` swapped order of y_grid and x_grid in dictionary
   for writing the correct order of coordinates in idf_svat.inp.
+- Fixed bug in :meth:`imod.mf6.Well.cleanup` and
+  :meth:`imod.mf6.LayeredWell.cleanup` which caused an error when called with an
+  unstructured discretization.
 
 Changed
 ~~~~~~~
diff --git a/imod/mf6/wel.py b/imod/mf6/wel.py
@@ -713,8 +713,10 @@ def _cleanup(
         cleanup_func: Callable,
         **cleanup_kwargs,
     ) -> None:
+        # Work around mypy error, .data_vars cannot be used with xu.UgridDataset
+        dict_to_broadcast: dict[str, GridDataArray] = dict(**dis.dataset)  # type: ignore
         # Top and bottom should be forced to grids with a x, y coordinates
-        top, bottom = broadcast_to_full_domain(**dict(dis.dataset.data_vars))
+        top, bottom = broadcast_to_full_domain(**dict_to_broadcast)
         # Collect point variable datanames
         point_varnames = list(self._write_schemata.keys())
         if "concentration" not in self.dataset.keys():
diff --git a/imod/tests/test_mf6/test_mf6_wel.py b/imod/tests/test_mf6/test_mf6_wel.py
@@ -23,6 +23,7 @@
 from imod.mf6.write_context import WriteContext
 from imod.schemata import ValidationError
 from imod.tests.fixtures.flow_basic_fixture import BasicDisSettings
+from imod.typing.grid import ones_like
 
 times = [
     datetime(1981, 11, 30),
@@ -34,6 +35,23 @@
 ]
 
 
+class BasicGridCases:
+    @staticmethod
+    def case_structured(basic_dis):
+        idomain, top, bottom = basic_dis
+        return (
+            imod.mf6.StructuredDiscretization,
+            idomain,
+            top,
+            bottom,
+        )
+
+    @staticmethod
+    def case_unstructured(basic_unstructured_dis):
+        idomain, top, bottom = basic_unstructured_dis
+        return imod.mf6.VerticesDiscretization, idomain, top, bottom
+
+
 class GridAgnosticWellCases:
     def case_well_stationary(self, well_high_lvl_test_data_stationary):
         obj = imod.mf6.Well(*well_high_lvl_test_data_stationary)
@@ -330,21 +348,22 @@ def test_to_mf6_pkg__validate_filter_top(well_high_lvl_test_data_stationary):
     )
 
 
+@parametrize_with_cases("dis", cases=BasicGridCases)
 def test_to_mf6_pkg__logging_with_message(
-    tmp_path, basic_dis, well_high_lvl_test_data_transient
+    tmp_path, dis, well_high_lvl_test_data_transient
 ):
     # Arrange
     logfile_path = tmp_path / "logfile.txt"
-    idomain, top, bottom = basic_dis
+    _, idomain, top, bottom = dis
     modified_well_fixture = list(well_high_lvl_test_data_transient)
 
     # create an idomain where layer 1 is active and layer 2 and 3 are inactive.
     # layer 1 has a bottom at -5, layer 2 at -35 and layer 3 at -135
     # so only wells that have a filter top above -5 will end up in the simulation
     active = idomain == 1
 
-    active.loc[1, :, :] = True
-    active.loc[2:, :, :] = False
+    active.loc[1, ...] = True
+    active.loc[2:, ...] = False
 
     # modify the well filter top and filter bottoms so that
     # well 0 is not placed
@@ -393,7 +412,7 @@ def test_to_mf6_pkg__logging_with_message(
 
         wel = imod.mf6.Well(*modified_well_fixture)
 
-        k = xr.ones_like(idomain)
+        k = ones_like(idomain)
         _ = wel.to_mf6_pkg(active, top, bottom, k)
 
     # the wells that were fully or partially placed should not appear in the log message
@@ -415,15 +434,16 @@ def test_to_mf6_pkg__logging_with_message(
         assert "id = 11 x" in log
 
 
+@parametrize_with_cases("dis", cases=BasicGridCases)
 def test_to_mf6_pkg__logging_without_message(
-    tmp_path, basic_dis, well_high_lvl_test_data_transient
+    tmp_path, dis, well_high_lvl_test_data_transient
 ):
     # This test activates logging, and then converts a high level well package to
     # an MF6 package, in such a way that all the wells can be placed.
     # Logging is active, and the log file should not include the "Some wells were not placed"
     # message
     logfile_path = tmp_path / "logfile.txt"
-    idomain, top, bottom = basic_dis
+    _, idomain, top, bottom = dis
 
     with open(logfile_path, "w") as sys.stdout:
         imod.logging.configure(
@@ -435,10 +455,10 @@ def test_to_mf6_pkg__logging_without_message(
 
         wel = imod.mf6.Well(*well_high_lvl_test_data_transient)
         active = idomain == 1
-        k = xr.ones_like(idomain)
+        k = ones_like(idomain)
 
-        active.loc[1, :, :] = True
-        active.loc[2:, :, :] = True
+        active.loc[1, ...] = True
+        active.loc[2:, ...] = True
 
         # Act
         _ = wel.to_mf6_pkg(active, top, bottom, k)
@@ -448,32 +468,36 @@ def test_to_mf6_pkg__logging_without_message(
         assert "Some wells were not placed" not in log
 
 
-def test_to_mf6_pkg__error_on_all_wells_removed(
-    basic_dis, well_high_lvl_test_data_transient
-):
+@parametrize_with_cases("dis", cases=BasicGridCases)
+def test_to_mf6_pkg__error_on_all_wells_removed(dis, well_high_lvl_test_data_transient):
     """Drop all wells, then run to_mf6_pkg"""
-    idomain, top, bottom = basic_dis
+    _, idomain, top, bottom = dis
 
     wel = imod.mf6.Well(*well_high_lvl_test_data_transient)
     wel.dataset = wel.dataset.drop_sel(index=wel.dataset["index"])
     active = idomain == 1
-    k = xr.ones_like(idomain)
+    k = ones_like(idomain)
 
     with pytest.raises(ValidationError, match="No wells were assigned in package"):
         wel.to_mf6_pkg(active, top, bottom, k)
 
 
-def test_to_mf6_pkg__error_on_well_removal(
-    basic_dis, well_high_lvl_test_data_transient
-):
+@parametrize_with_cases("dis", cases=BasicGridCases)
+def test_to_mf6_pkg__error_on_well_removal(dis, well_high_lvl_test_data_transient):
     """Set k values at well location x=81 to 1e-3, causing it to be dropped.
     Should throw error if error_on_well_removal = True"""
-    idomain, top, bottom = basic_dis
+    dis_pkg_type, idomain, top, bottom = dis
 
     wel = imod.mf6.Well(minimum_k=0.9, *well_high_lvl_test_data_transient)
     active = idomain == 1
-    k = xr.ones_like(idomain)
-    k.loc[{"x": 85.0}] = 1e-3
+    k = ones_like(idomain).astype(float)
+    # Set k values at well location x=81 to 1e-3, causing it to be dropped.
+    if dis_pkg_type is imod.mf6.VerticesDiscretization:
+        dim = k.ugrid.grid.face_dimension
+        indices = k.sel(layer=1).ugrid.sel(x=85.0).coords[dim].data
+        k.loc[{dim: indices}] = 1e-3
+    else:
+        k.loc[{"x": 85.0}] = 1e-3
 
     with pytest.raises(ValidationError, match="x = 81"):
         wel.to_mf6_pkg(active, top, bottom, k, strict_well_validation=True)
@@ -521,17 +545,18 @@ 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):
+@parametrize_with_cases("dis", cases=BasicGridCases)
+def test_cleanup(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)
+
+    dis_pkg_type, idomain, top, bottom = dis
+    if "layer" in top.dims:
+        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
-    )
+    dis_normal = dis_pkg_type(top, bottom, idomain)
+    dis_deep = dis_pkg_type(top - deep_offset, bottom - deep_offset, idomain)
 
     # Nothing to be cleaned up with default discretization
     wel.cleanup(dis_normal)
