Issue #1248 Write svat_groundwater to file and update the tests

docs/api/changelog.rst

@@ -6,6 +6,17 @@ All notable changes to this project will be documented in this file.
 The format is based on `Keep a Changelog`_, and this project adheres to
 `Semantic Versioning`_.
 
+Fixed
+~~~~~
+
+- :class:`imod.msw.Sprinkling` now correctly writes source svats to scap_svat.inp file.
+
+
+Changed
+~~~~~~~
+
+- Variable ``max_abstraction_groundwater`` in :class:`imod.msw.Sprinkling` now needs to 
+  have a subunit coordinate.
 
 [0.18.0]
 --------

imod/msw/sprinkling.py

@@ -37,21 +37,20 @@ class Sprinkling(MetaSwapPackage, IRegridPackage):
         "max_abstraction_surfacewater_mm_d": VariableMetaData(8, None, None, str),
         "max_abstraction_groundwater": VariableMetaData(8, 0.0, 1e9, float),
         "max_abstraction_surfacewater": VariableMetaData(8, 0.0, 1e9, float),
-        "svat_groundwater": VariableMetaData(10, None, None, str),
+        "svat_groundwater": VariableMetaData(10, 1, 99999999, int),
         "layer": VariableMetaData(6, 1, 9999, int),
         "trajectory": VariableMetaData(10, None, None, str),
     }
 
-    _with_subunit = ()
-    _without_subunit = (
+    _with_subunit = (
         "max_abstraction_groundwater",
         "max_abstraction_surfacewater",
     )
+    _without_subunit = ()
 
     _to_fill = (
         "max_abstraction_groundwater_mm_d",
         "max_abstraction_surfacewater_mm_d",
-        "svat_groundwater",
         "trajectory",
     )
 
@@ -76,31 +75,42 @@ def _render(
         mf6_dis: StructuredDiscretization,
         mf6_well: Mf6Wel,
     ):
+        def ravel_per_subunit(array: xr.DataArray) -> np.ndarray:
+            # per defined well element, all subunits
+            array_out = array.to_numpy()[:, well_row, well_column].ravel()
+            # per defined well element, per defined subunits
+            return array_out[np.isfinite(array_out)]
+
         if not isinstance(mf6_well, Mf6Wel):
             raise TypeError(rf"well not of type 'Mf6Wel', got '{type(mf6_well)}'")
 
         well_cellid = mf6_well["cellid"]
-        if len(well_cellid.coords["dim_cellid"]) != 3:
-            raise TypeError("Coupling to unstructured grids is not supported.")
 
         well_layer = well_cellid.sel(dim_cellid="layer").data
         well_row = well_cellid.sel(dim_cellid="row").data - 1
         well_column = well_cellid.sel(dim_cellid="column").data - 1
 
-        n_subunit = svat["subunit"].size
-
-        well_svat = svat.data[:, well_row, well_column]
-        well_active = well_svat != 0
-
-        # Tile well_layers for each subunit
-        layer = np.tile(well_layer, (n_subunit, 1))
-
-        data_dict = {"svat": well_svat[well_active], "layer": layer[well_active]}
-
-        for var in self._without_subunit:
-            well_arr = self.dataset[var].data[well_row, well_column]
-            well_arr = np.tile(well_arr, (n_subunit, 1))
-            data_dict[var] = well_arr[well_active]
+        max_rate_per_svat = self.dataset["max_abstraction_groundwater"].where(svat > 0)
+        layer_per_svat = xr.full_like(max_rate_per_svat, np.nan)
+        layer_per_svat.values[:, well_row, well_column] = well_layer
+
+        layer_source = ravel_per_subunit(
+            layer_per_svat.where(max_rate_per_svat > 0)
+        ).astype(dtype=np.int32)
+        svat_source_target = ravel_per_subunit(
+            svat.where(max_rate_per_svat > 0)
+        ).astype(dtype=np.int32)
+
+        data_dict = {
+            "svat": svat_source_target,
+            "layer": layer_source,
+            "svat_groundwater": svat_source_target,
+        }
+
+        for var in self._with_subunit:
+            array = self.dataset[var].where(max_rate_per_svat > 0).to_numpy()
+            array = array[np.isfinite(array)]
+            data_dict[var] = array
 
         for var in self._to_fill:
             data_dict[var] = ""

imod/tests/fixtures/msw_model_fixture.py

@@ -204,8 +204,8 @@ def make_msw_model():
 
     # %% Sprinkling
     msw_model["sprinkling"] = msw.Sprinkling(
-        max_abstraction_groundwater=xr.full_like(msw_grid, 100.0),
-        max_abstraction_surfacewater=xr.full_like(msw_grid, 100.0),
+        max_abstraction_groundwater=xr.full_like(area, 100.0),
+        max_abstraction_surfacewater=xr.full_like(area, 100.0),
     )
 
     # %% Ponding

imod/tests/test_msw/test_sprinkling.py

@@ -11,7 +11,7 @@
 from imod.mf6.wel import derive_cellid_from_points
 
 
-def test_simple_model(fixed_format_parser):
+def test_simple_model_all_svats(fixed_format_parser):
     x = [1.0, 2.0, 3.0]
     y = [1.0, 2.0, 3.0]
     subunit = [0, 1]
@@ -20,20 +20,32 @@ def test_simple_model(fixed_format_parser):
     # fmt: off
     max_abstraction_groundwater = xr.DataArray(
         np.array(
-            [[nan, 100.0, nan],
-             [nan, 200.0, nan],
-             [nan, 300.0, nan]]),
-        dims=("y", "x"),
-        coords={"y": y, "x": x, "dx": dx, "dy": dy}
+            [
+                [[nan, 100.0, nan],
+                [nan, 200.0, nan],
+                [nan, 300.0, nan]],
+                [[nan, 100.0, nan],
+                [nan, 200.0, nan],
+                [nan, 300.0, nan]]
+            ]
+        ),
+        dims=("subunit", "y", "x"),
+        coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy}
     )
 
     max_abstraction_surfacewater = xr.DataArray(
         np.array(
-            [[nan, 100.0, nan],
-             [nan, 200.0, nan],
-             [nan, 300.0, nan]]),
-        dims=("y", "x"),
-        coords={"y": y, "x": x, "dx": dx, "dy": dy}
+            [
+                [[nan, 100.0, nan],
+                [nan, 200.0, nan],
+                [nan, 300.0, nan]],
+                [[nan, 100.0, nan],
+                [nan, 200.0, nan],
+                [nan, 300.0, nan]]
+            ]
+        ),
+        dims=("subunit", "y", "x"),
+        coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy}
     )
 
     svat = xr.DataArray(
@@ -86,6 +98,97 @@ def test_simple_model(fixed_format_parser):
         np.array([100.0, 300.0, 100.0, 200.0]),
     )
     assert_equal(results["layer"], np.array([3, 1, 3, 2]))
+    assert_equal(results["svat_groundwater"], np.array([1, 2, 3, 4]))
+
+
+def test_simple_model_some_svats(fixed_format_parser):
+    x = [1.0, 2.0, 3.0]
+    y = [1.0, 2.0, 3.0]
+    subunit = [0, 1]
+    dx = 1.0
+    dy = 1.0
+    # fmt: off
+    max_abstraction_groundwater = xr.DataArray(
+        np.array(
+            [
+                [[nan, 100.0, nan],
+                [nan, 200.0, nan],
+                [nan, 300.0, nan]],
+                [[nan, nan, nan],
+                [nan, 200.0, nan],
+                [nan, nan, nan]]
+            ]
+        ),
+        dims=("subunit", "y", "x"),
+        coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy}
+    )
+
+    max_abstraction_surfacewater = xr.DataArray(
+        np.array(
+            [
+                [[nan, 100.0, nan],
+                [nan, 200.0, nan],
+                [nan, 300.0, nan]],
+                [[nan, nan, nan],
+                [nan, 200.0, nan],
+                [nan, nan, nan]]
+            ]
+        ),
+        dims=("subunit", "y", "x"),
+        coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy}
+    )
+
+    svat = xr.DataArray(
+        np.array(
+            [
+                [[0, 1, 0],
+                 [0, 0, 0],
+                 [0, 2, 0]],
+
+                [[0, 3, 0],
+                 [0, 4, 0],
+                 [0, 0, 0]],
+            ]
+        ),
+        dims=("subunit", "y", "x"),
+        coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy}
+    )
+    # fmt: on
+    index = (svat != 0).values.ravel()
+
+    # Well
+    well_layer = [3, 2, 1]
+    well_y = [1.0, 2.0, 3.0]
+    well_x = [2.0, 2.0, 2.0]
+    well_rate = [-5.0] * 3
+    cellids = derive_cellid_from_points(svat, well_x, well_y, well_layer)
+    well = Mf6Wel(cellids, well_rate)
+
+    coupler_mapping = msw.Sprinkling(
+        max_abstraction_groundwater,
+        max_abstraction_surfacewater,
+    )
+
+    with tempfile.TemporaryDirectory() as output_dir:
+        output_dir = Path(output_dir)
+        coupler_mapping.write(output_dir, index, svat, None, well)
+
+        results = fixed_format_parser(
+            output_dir / msw.Sprinkling._file_name,
+            msw.Sprinkling._metadata_dict,
+        )
+
+    assert_equal(results["svat"], np.array([1, 2, 4]))
+    assert_almost_equal(
+        results["max_abstraction_groundwater"],
+        np.array([100.0, 300.0, 200.0]),
+    )
+    assert_almost_equal(
+        results["max_abstraction_surfacewater"],
+        np.array([100.0, 300.0, 200.0]),
+    )
+    assert_equal(results["layer"], np.array([3, 1, 2]))
+    assert_equal(results["svat_groundwater"], np.array([1, 2, 4]))
 
 
 def test_simple_model_1_subunit(fixed_format_parser):
@@ -97,20 +200,26 @@ def test_simple_model_1_subunit(fixed_format_parser):
     # fmt: off
     max_abstraction_groundwater = xr.DataArray(
         np.array(
-            [[nan, 100.0, nan],
-             [nan, 200.0, nan],
-             [nan, 300.0, nan]]),
-        dims=("y", "x"),
-        coords={"y": y, "x": x, "dx": dx, "dy": dy}
+            [
+                [[nan, 100.0, nan],
+                 [nan, 200.0, nan],
+                 [nan, 300.0, nan]]
+            ]
+        ),
+        dims=("subunit", "y", "x"),
+        coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy}
     )
 
     max_abstraction_surfacewater = xr.DataArray(
         np.array(
-            [[nan, 100.0, nan],
-             [nan, 200.0, nan],
-             [nan, 300.0, nan]]),
-        dims=("y", "x"),
-        coords={"y": y, "x": x, "dx": dx, "dy": dy}
+            [
+                [[nan, 100.0, nan],
+                 [nan, 200.0, nan],
+                 [nan, 300.0, nan]]
+            ]
+        ),
+        dims=("subunit", "y", "x"),
+        coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy}
     )
 
     svat = xr.DataArray(
@@ -159,3 +268,4 @@ def test_simple_model_1_subunit(fixed_format_parser):
         np.array([100.0, 300.0]),
     )
     assert_equal(results["layer"], np.array([3, 2]))
+    assert_equal(results["svat_groundwater"], np.array([1, 2]))