Fixes #224 uparea snapping to river

examples/test_stations.csv

@@ -1,4 +1,4 @@
-ID,Name,x,y
-1001,Gauge1,11.9594,45.8925
-1002,Gauge2,12.3395,46.1492
-1003,Gauge3,12.0785,46.1122
+ID,Name,x,y,uparea
+1001,Gauge1,11.9594,45.8925,3642
+1002,Gauge2,12.3395,46.1492,2
+1003,Gauge3,12.0785,46.1122,837

hydromt_wflow/wflow.py

@@ -1112,6 +1112,7 @@ def setup_gauges(
         max_dist: Optional[float] = 10e3,
         wdw: Optional[int] = 3,
         rel_error: Optional[float] = 0.05,
+        abs_error: Optional[float] = 50.0,
         derive_subcatch: Optional[bool] = False,
         basename: Optional[str] = None,
         toml_output: Optional[str] = "csv",
@@ -1187,7 +1188,11 @@ def setup_gauges(
         rel_error: float, optional
             Maximum relative error (default 0.05)
             between the gauge location upstream area and the upstream area of
-            the best fit grid cell, only used if snap_area is True.
+            the best fit grid cell, only used if snap_uparea is True.
+        abs_error: float, optional
+            Maximum absolute error (default 50.0)
+            between the gauge location upstream area and the upstream area of
+            the best fit grid cell, only used if snap_uparea is True.
         derive_subcatch : bool, optional
             Derive subcatch map for gauges, by default False
         basename : str, optional
@@ -1297,8 +1302,10 @@ def setup_gauges(
                 self.grid,
                 gdf_gauges,
                 uparea_name="wflow_uparea",
+                mask=mask,
                 wdw=wdw,
                 rel_error=rel_error,
+                abs_error=abs_error,
                 logger=self.logger,
             )
         else:

hydromt_wflow/workflows/gauges.py

@@ -18,6 +18,7 @@ def gauge_map_uparea(
     ds: xr.Dataset,
     gdf: gpd.GeoDataFrame,
     uparea_name: Optional[str] = "wflow_uparea",
+    mask: Optional[np.ndarray] = None,
     wdw: Optional[int] = 1,
     rel_error: float = 0.05,
     abs_error: float = 50,
@@ -40,6 +41,8 @@ def gauge_map_uparea(
         GeoDataFrame with gauge points and uparea column.
     uparea_name : str, optional
         Name of the upstream area variable in ``ds``, by default "wflow_uparea".
+    mask : np.ndarray, optional
+        Mask cells to apply the uparea snapping, by default None.
     wdw : int, optional
         Window size around the original location to search for the best matching cell,
         by default 1.
@@ -61,13 +64,23 @@ def gauge_map_uparea(
     if uparea_name not in ds:
         raise ValueError(f"uparea_name {uparea_name} not found in ds.")
 
+    ds = ds.copy()
+    # Add mask to ds
+    if mask is not None:
+        ds["mask"] = xr.DataArray(mask, dims=(ds.raster.y_dim, ds.raster.x_dim))
+
     ds_wdw = ds.raster.sample(gdf, wdw=wdw)
+    # Mask valid cells in ds_wdw
+    if mask is not None:
+        ds_wdw[uparea_name] = ds_wdw[uparea_name].where(
+            ds_wdw["mask"], ds_wdw[uparea_name].raster.nodata
+        )
     logger.debug(
         f"Snapping gauges points to best matching uparea cell within wdw (size={wdw})."
     )
     upa0 = xr.DataArray(gdf["uparea"], dims=("index"))
     upa_dff = np.abs(ds_wdw[uparea_name].where(ds_wdw[uparea_name] > 0).load() - upa0)
-    upa_check = np.logical_or((upa_dff / upa0) <= rel_error, upa_dff <= abs_error)
+    upa_check = np.logical_and((upa_dff / upa0) <= rel_error, upa_dff <= abs_error)
     # find best matching uparea cell in window
     i_wdw = upa_dff.argmin("wdw").load()
 

tests/test_model_methods.py

@@ -380,7 +380,7 @@ def test_setup_gauges(example_wflow_model):
     ds_samp = example_wflow_model.grid[["wflow_river", "wflow_uparea"]].raster.sample(
         gdf, wdw=0
     )
-    assert np.all(ds_samp["wflow_river"].values == 1)
+    # assert np.all(ds_samp["wflow_river"].values == 1)
     assert np.allclose(ds_samp["wflow_uparea"].values, gdf["uparea"].values, rtol=0.05)
 
     # Test with/without snapping
@@ -408,6 +408,40 @@ def test_setup_gauges(example_wflow_model):
     assert len(equal) == 1
     assert equal.index.values[0] == 1003
 
+    # Test uparea with/without river snapping
+    example_wflow_model.setup_gauges(
+        gauges_fn=stations_fn,
+        basename="stations_uparea_no_snapping",
+        snap_to_river=False,
+        mask=None,
+        snap_uparea=True,
+        wdw=5,
+        rel_error=0.05,
+    )
+    gdf_no_snap = example_wflow_model.geoms["gauges_stations_uparea_no_snapping"]
+
+    example_wflow_model.setup_gauges(
+        gauges_fn=stations_fn,
+        basename="stations_uparea_snapping",
+        snap_to_river=True,
+        mask=None,
+        snap_uparea=True,
+        wdw=5,
+        rel_error=0.05,
+        abs_error=25,
+    )
+    gdf_snap = example_wflow_model.geoms["gauges_stations_uparea_snapping"]
+
+    ds_samp = example_wflow_model.grid[["wflow_river", "wflow_uparea"]].raster.sample(
+        gdf_snap, wdw=0
+    )
+    assert np.all(ds_samp["wflow_river"].values == 1)
+    ds_samp = example_wflow_model.grid[["wflow_river", "wflow_uparea"]].raster.sample(
+        gdf_no_snap, wdw=0
+    )
+    assert not np.all(ds_samp["wflow_river"].values == 1)
+    assert not gdf_snap.equals(gdf_no_snap)
+
 
 @pytest.mark.parametrize("elevtn_map", ["wflow_dem", "dem_subgrid"])
 def test_setup_rivers(elevtn_map, floodplain1d_testdata, example_wflow_model):