add tests

pre-processing/primod/driver_coupling/ribameta.py

@@ -6,6 +6,7 @@
 import imod
 import numpy as np
 import ribasim
+import xarray as xr
 from imod.msw import GridData, MetaSwapModel, Sprinkling
 
 from primod.driver_coupling.driver_coupling_base import DriverCoupling
@@ -15,7 +16,6 @@
 )
 from primod.mapping.svat_basin_mapping import SvatBasinMapping
 from primod.mapping.svat_user_demand_mapping import SvatUserDemandMapping
-import xarray as xr
 
 
 class RibaMetaDriverCoupling(DriverCoupling):

pre-processing/primod/mapping/node_basin_mapping.py

@@ -1,5 +1,3 @@
-from typing import Union
-
 import geopandas as gpd
 import numpy as np
 import pandas as pd
@@ -118,6 +116,7 @@ def __init__(
         basin_ids: pd.Series,
         subgrid_df: pd.DataFrame,
     ):
+        validate_meta_label_column(name, subgrid_df)
         # Use xarray.where() to force the dimension order of conductance, rather than
         # using gridded_basin.where() (which prioritizes the gridded_basin dims)
         conductance = self._ensure_time_invariant_conductance(conductance)
@@ -141,7 +140,7 @@ def __init__(
         )
 
     @staticmethod
-    def _get_subgrid_xy(subgrid: pd.DataFrame) -> Union[NDArray[Float], NDArray[Float]]:
+    def _get_subgrid_xy(subgrid: pd.DataFrame) -> tuple[NDArray[Float], NDArray[Float]]:
         # Check whether columns (optional to Ribasim) are present.
         if "meta_x" not in subgrid or "meta_y" not in subgrid:
             raise ValueError(
@@ -163,6 +162,7 @@ def _get_subgrid_xy(subgrid: pd.DataFrame) -> Union[NDArray[Float], NDArray[Floa
             )
         x = grouped["meta_x"].first().to_numpy()
         y = grouped["meta_y"].first().to_numpy()
+        # At this level the subgrid-df could be a subset of the Ribasim one
         subgrid_id = grouped["subgrid_id"].first().to_numpy()
         return np.column_stack((x, y)), subgrid_id
 
@@ -192,3 +192,11 @@ def _find_nearest_subgrid_elements(
         # FUTURE: maybe do something with the distances returned by query?
         indices: NDArray[Int] = kdtree.query(conductance_xy)[1]
         return indices
+
+
+def validate_meta_label_column(name: str, subgrid_df: pd.DataFrame) -> None:
+    if "meta_label" in subgrid_df.columns:
+        if (subgrid_df["meta_label"] != name).any():
+            raise ValueError(
+                "if column 'meta_label' is defined in subgrid dataframe, all actively coupled packages should be included"
+            )

tests/test_primod/test_ribamod_exchange_creator.py

@@ -172,14 +172,15 @@ def test_get_subgrid_xy():
             "meta_y": [1.0, 1.0, 2.0, 2.0],
         }
     )
-    xy = ActiveNodeBasinMapping._get_subgrid_xy(df)
+    xy, subgrid_id = ActiveNodeBasinMapping._get_subgrid_xy(df)
     assert np.allclose(
         xy,
         [
             [1.0, 1.0],
             [2.0, 2.0],
         ],
     )
+    assert np.allclose(subgrid_id, [1, 2])
 
 
 def test_get_conductance_xy():
@@ -249,3 +250,125 @@ def test_derive_active_coupling():
         ),
         check_dtype=False,  # int32 versus int64 ...
     )
+
+
+def test_node_basin_mapping_stacked() -> None:
+    #  basin definition:  | 1 | 1 | 2 | 2 | 2 |
+    #                     | 1 | 1 | 2 | 2 | 2 |
+    #                     | 1 | 1 | 2 | 2 | 2 |
+    #                     | 1 | 1 | 2 | 2 | 2 |
+
+    #        subgrid id:  |   |   |   |   |   |
+    #                     |   |   |   |   |   |
+    #                     |0,1|2,3|4,5|6,7|8,9|
+    #                     |   |   |   |   |   |
+
+    #
+    #        cond sys 1:  |   |   |   |   |   |
+    #                     |   |   |   |   |   |
+    #                     | x | x | x | x | x |
+    #                     |   |   |   |   |   |
+
+    #        cond sys 2:  |   |   |   |   |   |
+    #                     |   |   |   |   |   |
+    #                     |   |   | x | x | x |
+    #                     |   |   |   |   |   |
+
+    # sys1: subgrid elements 0, 2, 4, 6, 8 should be coupled
+    # sys2: subgrid elements 5, 7, 9 should be coupled
+
+    cond1 = xr.full_like(conductance().isel(layer=0, drop=True), 1.0)
+    cond1[0:2, :] = np.nan
+    cond1[3, :] = np.nan
+
+    cond2 = cond1.copy()
+    cond2[2, 0:2] = np.nan
+
+    gridded_basin = xr.full_like(cond1, 1)
+    gridded_basin[:, 2:5] = 2
+    basin_ids = pd.Series([1, 2])
+    subgrid_df = pd.DataFrame(
+        data={
+            "node_id": np.array([1, 1, 2, 2, 2, 1, 1, 2, 2, 2]),
+            "subgrid_id": np.arange(10),
+            "subgrid_level": np.array(
+                [1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
+            ),
+            "meta_x": np.array([0.5, 0.5, 1.5, 1.5, 2.5, 2.5, 3.5, 3.5, 4.5, 4.5]),
+            "meta_y": np.array([11.5] * 10),
+            "meta_label": np.array(
+                [
+                    "sys1",
+                    "sys2",
+                    "sys1",
+                    "sys2",
+                    "sys1",
+                    "sys2",
+                    "sys1",
+                    "sys2",
+                    "sys1",
+                    "sys2",
+                ]
+            ),
+        }
+    )
+    # check results for sys1
+    mapping = ActiveNodeBasinMapping(
+        "sys1", cond1, gridded_basin, basin_ids, subgrid_df
+    )
+    assert mapping.name == "sys1"
+    table = mapping.dataframe
+    assert isinstance(table, pd.DataFrame)
+    assert table.shape == (5, 3)
+    pd.testing.assert_frame_equal(
+        table,
+        pd.DataFrame(
+            data={
+                "basin_index": [0, 0, 1, 1, 1],
+                "bound_index": [0, 1, 2, 3, 4],
+                "subgrid_index": [0, 2, 4, 6, 8],
+            }
+        ),
+        check_dtype=False,  # int32 versus int64 ...
+    )
+    # and sys2
+    mapping = ActiveNodeBasinMapping(
+        "sys2", cond2, gridded_basin, basin_ids, subgrid_df
+    )
+    assert mapping.name == "sys2"
+    table = mapping.dataframe
+    assert isinstance(table, pd.DataFrame)
+    assert table.shape == (3, 3)
+    pd.testing.assert_frame_equal(
+        table,
+        pd.DataFrame(
+            data={
+                "basin_index": [1, 1, 1],
+                "bound_index": [0, 1, 2],
+                "subgrid_index": [5, 7, 9],
+            }
+        ),
+        check_dtype=False,  # int32 versus int64 ...
+    )
+
+
+def test_node_basin_mapping_stacked_no_label() -> None:
+    cond = xr.full_like(conductance().isel(layer=0, drop=True), 1.0)
+    gridded_basin = xr.full_like(cond, 1)
+    basin_ids = pd.Series([1])
+    subgrid_df = pd.DataFrame(
+        data={
+            "node_id": np.array([1] * 10),
+            "subgrid_id": np.arange(10),
+            "subgrid_level": np.array([1.0] * 10),
+            "meta_x": np.array([0.5, 0.5, 1.5, 1.5, 2.5, 2.5, 3.5, 3.5, 4.5, 4.5]),
+            "meta_y": np.array([11.5] * 10),
+            "meta_label": np.array(["sys1"] * 10),
+        }
+    )
+    _ = ActiveNodeBasinMapping("sys1", cond, gridded_basin, basin_ids, subgrid_df)
+    with pytest.raises(
+        ValueError,
+        match="if column 'meta_label' is defined in subgrid dataframe, all actively coupled packages should be included",
+    ):
+        ActiveNodeBasinMapping("sys2", cond, gridded_basin, basin_ids, subgrid_df)