Several basin related fixes

docs/changelog.rst

@@ -14,9 +14,11 @@ Added
 
 Changed
 -------
+- **setup_outlets**: the IDs of wflow_subcatch are used to define the outlets IDs rather than [1:n]. PR #247
 
 Fixed
 -----
+- Wrong dtype for wflow_subcatch map. PR #247
 
 v0.5.0 (February 2024)
 ======================

hydromt_wflow/wflow.py

@@ -1047,6 +1047,11 @@ def setup_outlets(
     ):
         """Set the default gauge map based on basin outlets.
 
+        If wflow_subcatch is available, the catchment outlets IDs will be matching the
+        wflow_subcatch IDs. If not, then IDs from 1 to number of outlets are used.
+
+        Can also add csv/netcdf output settings in the TOML.
+
         Adds model layers:
 
         * **wflow_gauges** map: gauge IDs map from catchment outlets [-]
@@ -1080,9 +1085,15 @@ def setup_outlets(
             idxs_out = idxs_out[
                 (self.grid[self._MAPS["rivmsk"]] > 0).values.flat[idxs_out]
             ]
+        # Use the wflow_subcatch ids
+        if self._MAPS["basins"] in self.grid:
+            ids = self.grid[self._MAPS["basins"]].values.flat[idxs_out]
+        else:
+            ids = None
         da_out, idxs_out, ids_out = flw.gauge_map(
             self.grid,
             idxs=idxs_out,
+            ids=ids,
             flwdir=self.flwdir,
             logger=self.logger,
         )

hydromt_wflow/workflows/basemaps.py

@@ -214,6 +214,9 @@ def hydrography(
             if basins is None:
                 basins = flwdir_out.basins(idxs=flwdir_out.idxs_pit).astype(np.int32)
             ds_out[basins_name] = xr.Variable(dims, basins, attrs=dict(_FillValue=0))
+        else:
+            # make sure dtype in ds_out is np.int32
+            ds_out[basins_name] = ds_out[basins_name].astype(np.int32)
         # upstream area
         if uparea_name not in ds_out.data_vars:
             uparea = flwdir_out.upstream_area("km2")  # km2

tests/test_model_methods.py

@@ -17,6 +17,33 @@
 EXAMPLEDIR = join(dirname(abspath(__file__)), "..", "examples")
 
 
+def test_setup_basemaps(tmpdir):
+    # Region
+    region = {
+        "basin": [12.2051, 45.8331],
+        "strord": 4,
+        "bounds": [11.70, 45.35, 12.95, 46.70],
+    }
+    mod = WflowModel(
+        root=str(tmpdir.join("wflow_base")),
+        mode="w",
+        data_libs=["artifact_data"],
+    )
+
+    hydrography = mod.data_catalog.get_rasterdataset("merit_hydro")
+    # Change dtype to uint32
+    hydrography["basins"] = hydrography["basins"].astype("uint32")
+
+    # Run setup_basemaps
+    mod.setup_basemaps(
+        region=region,
+        hydrography_fn=hydrography,
+        res=hydrography.raster.res[0],  # no upscaling
+    )
+
+    assert mod.grid["wflow_subcatch"].dtype == "int32"
+
+
 def test_setup_grid(example_wflow_model):
     # Tests on setup_grid_from_raster
     example_wflow_model.setup_grid_from_raster(
@@ -364,6 +391,22 @@ def test_setup_rootzoneclim(example_wflow_model):
     ] == pytest.approx(104.96931418911882, abs=0.5)
 
 
+def test_setup_outlets(example_wflow_model):
+    # Update wflow_subcatch ID
+    new_subcatch = example_wflow_model.grid["wflow_subcatch"].copy()
+    new_subcatch = new_subcatch.where(new_subcatch == new_subcatch.raster.nodata, 1001)
+    example_wflow_model.set_grid(new_subcatch, "wflow_subcatch")
+
+    # Derive outlets
+    example_wflow_model.setup_outlets()
+
+    # Check if the ID is indeed 1001
+    val, count = np.unique(example_wflow_model.grid["wflow_gauges"], return_counts=True)
+    # 0 is no data
+    assert val[1] == 1001
+    assert count[1] == 1
+
+
 def test_setup_gauges(example_wflow_model):
     # uparea rename not in the latest artifact_data version
     example_wflow_model.data_catalog["grdc"].rename = {"area": "uparea"}