hydro: drop twi dask+cupy host round-trip, vectorize watershed init loops

benchmarks/benchmarks/twi.py

@@ -1,20 +1,29 @@
 from xrspatial import flow_direction, flow_accumulation, slope, twi
+from xrspatial.utils import has_cuda_and_cupy
 
 from .common import get_xr_dataarray
 
 
 class TWI:
-    params = ([100, 300, 1000], ["numpy", "dask"])
+    params = ([100, 300, 1000], ["numpy", "dask", "cupy", "dask+cupy"])
     param_names = ("nx", "type")
 
     def setup(self, nx, type):
+        if type in ("cupy", "dask+cupy") and not has_cuda_and_cupy():
+            raise NotImplementedError()
+
         ny = nx // 2
         elev = get_xr_dataarray((ny, nx), "numpy")
         flow_dir = flow_direction(elev)
         self.flow_accum = flow_accumulation(flow_dir)
         self.slope_agg = slope(elev)
 
-        if type == "dask":
+        if type in ("cupy", "dask+cupy"):
+            import cupy
+            self.flow_accum.data = cupy.asarray(self.flow_accum.data)
+            self.slope_agg.data = cupy.asarray(self.slope_agg.data)
+
+        if type in ("dask", "dask+cupy"):
             import dask.array as da
             chunks = (max(1, ny // 2), max(1, nx // 2))
             self.flow_accum.data = da.from_array(

xrspatial/hydro/tests/test_twi_d8.py

@@ -180,3 +180,39 @@ def test_numpy_equals_dask_cupy(self):
         np.testing.assert_allclose(
             result_np.data, result_dc.data.compute().get(),
             equal_nan=True, rtol=1e-10)
+
+    def test_dask_cupy_stays_on_device(self):
+        """The dask+cupy path must not round-trip chunks through the host.
+
+        Regression test for the removed ``_twi_dask_cupy`` wrapper, which
+        pulled every chunk to numpy and pushed it back with cp.asarray.
+        The lazy result must keep a cupy meta, and twi must add the same
+        graph layers on dask+cupy input as on plain dask input (the
+        round-trip added a host-transfer map_blocks layer per input plus
+        a device-transfer layer on the output).
+        """
+        import cupy
+
+        np.random.seed(42)
+        fa_data = np.random.uniform(1, 1000, (6, 6)).astype(np.float64)
+        sl_data = np.random.uniform(0.1, 60, (6, 6)).astype(np.float64)
+
+        fa_da, sl_da = _make_twi_rasters(fa_data, sl_data, backend='dask',
+                                         chunks=(3, 3))
+        fa_dc, sl_dc = _make_twi_rasters(fa_data, sl_data,
+                                         backend='dask+cupy', chunks=(3, 3))
+
+        result_da = twi(fa_da, sl_da)
+        result_dc = twi(fa_dc, sl_dc)
+
+        assert isinstance(result_dc.data._meta, cupy.ndarray)
+
+        def added_layers(result, *inputs):
+            input_layers = set()
+            for arr in inputs:
+                input_layers |= set(arr.data.__dask_graph__().layers)
+            return len(set(result.data.__dask_graph__().layers)
+                       - input_layers)
+
+        assert (added_layers(result_dc, fa_dc, sl_dc)
+                == added_layers(result_da, fa_da, sl_da))

xrspatial/hydro/tests/test_watershed_d8.py

@@ -95,6 +95,24 @@ def test_nan_handling():
     assert np.isnan(result.data[1, 1])
 
 
+def test_pour_point_on_nan_flow_dir():
+    """A pour point placed on a NaN flow_dir cell is nodata, not a label."""
+    flow_dir = np.array([
+        [1.0, 0.0],
+        [np.nan, 1.0],
+    ], dtype=np.float64)
+    pour_points = np.full((2, 2), np.nan, dtype=np.float64)
+    pour_points[1, 0] = 9.0  # on the NaN flow_dir cell
+    pour_points[0, 1] = 7.0
+
+    fd_da = create_test_raster(flow_dir)
+    pp_da = create_test_raster(pour_points)
+    result = watershed(fd_da, pp_da)
+
+    assert np.isnan(result.data[1, 0])
+    assert result.data[0, 1] == 7.0
+
+
 def test_linear_chain():
     """Row of cells flowing east to a pour point at the end."""
     N = 6

xrspatial/hydro/twi_d8.py

@@ -18,7 +18,7 @@
     da = None
 
 from xrspatial.utils import (_validate_raster, get_dataarray_resolution, has_cuda_and_cupy,
-                             is_cupy_array, is_dask_cupy)
+                             is_cupy_array)
 
 # Minimum tan(slope) clamp: tan(0.001°)
 _TAN_MIN = np.tan(np.radians(0.001))
@@ -64,13 +64,13 @@ def twi_d8(flow_accum: xr.DataArray,
     fa_data = flow_accum.data
     sl_data = slope_agg.data
 
-    if has_cuda_and_cupy() and is_dask_cupy(flow_accum):
-        out = _twi_dask_cupy(fa_data, sl_data, cellsize)
-
-    elif has_cuda_and_cupy() and is_cupy_array(fa_data):
+    if has_cuda_and_cupy() and is_cupy_array(fa_data):
         out = _twi_cupy(fa_data, sl_data, cellsize)
 
     elif da is not None and isinstance(fa_data, da.Array):
+        # Covers numpy- and cupy-backed dask arrays alike: TWI is purely
+        # elementwise, so cupy chunks run the same expression natively
+        # with no host round-trip.
         out = _twi_dask(fa_data, sl_data, cellsize)
 
     elif isinstance(fa_data, np.ndarray):
@@ -106,25 +106,9 @@ def _twi_cupy(fa, sl, cellsize):
 
 
 def _twi_dask(fa, sl, cellsize):
+    """Elementwise TWI on dask arrays (numpy- or cupy-backed chunks)."""
     import dask.array as _da
     sca = fa * cellsize
     tan_slope = _da.tan(_da.radians(sl))
     tan_slope = _da.where(tan_slope < _TAN_MIN, _TAN_MIN, tan_slope)
     return _da.log(sca / tan_slope)
-
-
-def _twi_dask_cupy(fa, sl, cellsize):
-    import cupy as cp
-    fa_np = fa.map_blocks(
-        lambda b: b.get(), dtype=fa.dtype,
-        meta=np.array((), dtype=fa.dtype),
-    )
-    sl_np = sl.map_blocks(
-        lambda b: b.get(), dtype=sl.dtype,
-        meta=np.array((), dtype=sl.dtype),
-    )
-    result = _twi_dask(fa_np, sl_np, cellsize)
-    return result.map_blocks(
-        cp.asarray, dtype=result.dtype,
-        meta=cp.array((), dtype=result.dtype),
-    )

xrspatial/hydro/watershed_d8.py

@@ -54,7 +54,9 @@ def _to_numpy_f64(arr):
 #   state  (int8)        -> 1
 #   path_r (int64)       -> 8
 #   path_c (int64)       -> 8
-# Total ~33 bytes/pixel.  The caller's ``flow_dir`` and ``pour_points``
+# Total ~33 bytes/pixel.  The vectorized init also holds two boolean
+# masks plus their conjunction (~3 B/px) transiently; that fits in the
+# 50% guard headroom.  The caller's ``flow_dir`` and ``pour_points``
 # arrays already live in RAM before dispatch and are not double-counted.
 _BYTES_PER_PIXEL = 33
 
@@ -1043,17 +1045,11 @@ def watershed_d8(flow_dir: xr.DataArray,
         h, w = fd.shape
         # Init labels and state: pour points → resolved (state 3),
         # NaN flow_dir → nodata (state 0), others → unresolved (state 1)
-        labels = np.full((h, w), np.nan, dtype=np.float64)
-        state = np.zeros((h, w), dtype=np.int8)
-        for r in range(h):
-            for c in range(w):
-                if fd[r, c] != fd[r, c]:  # NaN
-                    pass  # state 0, label NaN
-                elif pp[r, c] == pp[r, c]:  # not NaN → pour point
-                    labels[r, c] = pp[r, c]
-                    state[r, c] = 3
-                else:
-                    state[r, c] = 1  # unresolved
+        fd_valid = ~np.isnan(fd)
+        pp_valid = ~np.isnan(pp)
+        labels = np.where(fd_valid & pp_valid, pp, np.nan)
+        state = np.where(fd_valid,
+                         np.where(pp_valid, 3, 1), 0).astype(np.int8)
         out = _watershed_cpu(fd, labels, state, h, w)
 
     elif has_cuda_and_cupy() and is_cupy_array(data):

xrspatial/hydro/watershed_dinf.py

@@ -219,17 +219,11 @@ def _watershed_dinf_cupy(flow_dir_data, pour_points_data):
     fd_np = _to_numpy_f64(flow_dir_data)
     pp_np = _to_numpy_f64(pour_points_data)
     h, w = fd_np.shape
-    labels = np.full((h, w), np.nan, dtype=np.float64)
-    state = np.zeros((h, w), dtype=np.int8)
-    for r in range(h):
-        for c in range(w):
-            if fd_np[r, c] != fd_np[r, c]:
-                pass
-            elif pp_np[r, c] == pp_np[r, c]:
-                labels[r, c] = pp_np[r, c]
-                state[r, c] = 3
-            else:
-                state[r, c] = 1
+    fd_valid = ~np.isnan(fd_np)
+    pp_valid = ~np.isnan(pp_np)
+    labels = np.where(fd_valid & pp_valid, pp_np, np.nan)
+    state = np.where(fd_valid,
+                     np.where(pp_valid, 3, 1), 0).astype(np.int8)
     out = _watershed_dinf_cpu(fd_np, labels, state, h, w)
     return cp.asarray(out)
 
@@ -699,17 +693,11 @@ def watershed_dinf(flow_dir_dinf: xr.DataArray,
         fd = data.astype(np.float64)
         pp = np.asarray(pp_data, dtype=np.float64)
         h, w = fd.shape
-        labels = np.full((h, w), np.nan, dtype=np.float64)
-        state = np.zeros((h, w), dtype=np.int8)
-        for r in range(h):
-            for c in range(w):
-                if fd[r, c] != fd[r, c]:
-                    pass
-                elif pp[r, c] == pp[r, c]:
-                    labels[r, c] = pp[r, c]
-                    state[r, c] = 3
-                else:
-                    state[r, c] = 1
+        fd_valid = ~np.isnan(fd)
+        pp_valid = ~np.isnan(pp)
+        labels = np.where(fd_valid & pp_valid, pp, np.nan)
+        state = np.where(fd_valid,
+                         np.where(pp_valid, 3, 1), 0).astype(np.int8)
         out = _watershed_dinf_cpu(fd, labels, state, h, w)
 
     elif has_cuda_and_cupy() and is_cupy_array(data):

xrspatial/hydro/watershed_mfd.py

@@ -227,17 +227,11 @@ def _watershed_mfd_cupy(fractions_data, pour_points_data):
     fr_np = _to_numpy_f64(fractions_data)
     pp_np = _to_numpy_f64(pour_points_data)
     _, h, w = fr_np.shape
-    labels = np.full((h, w), np.nan, dtype=np.float64)
-    state = np.zeros((h, w), dtype=np.int8)
-    for r in range(h):
-        for c in range(w):
-            if fr_np[0, r, c] != fr_np[0, r, c]:
-                pass
-            elif pp_np[r, c] == pp_np[r, c]:
-                labels[r, c] = pp_np[r, c]
-                state[r, c] = 3
-            else:
-                state[r, c] = 1
+    fr_valid = ~np.isnan(fr_np[0])
+    pp_valid = ~np.isnan(pp_np)
+    labels = np.where(fr_valid & pp_valid, pp_np, np.nan)
+    state = np.where(fr_valid,
+                     np.where(pp_valid, 3, 1), 0).astype(np.int8)
     out = _watershed_mfd_cpu(fr_np, labels, state, h, w)
     return cp.asarray(out)
 
@@ -701,17 +695,11 @@ def watershed_mfd(flow_dir_mfd: xr.DataArray,
         _check_memory(H, W)
         fr = data.astype(np.float64)
         pp = np.asarray(pp_data, dtype=np.float64)
-        labels = np.full((H, W), np.nan, dtype=np.float64)
-        state = np.zeros((H, W), dtype=np.int8)
-        for r in range(H):
-            for c in range(W):
-                if fr[0, r, c] != fr[0, r, c]:
-                    pass
-                elif pp[r, c] == pp[r, c]:
-                    labels[r, c] = pp[r, c]
-                    state[r, c] = 3
-                else:
-                    state[r, c] = 1
+        fr_valid = ~np.isnan(fr[0])
+        pp_valid = ~np.isnan(pp)
+        labels = np.where(fr_valid & pp_valid, pp, np.nan)
+        state = np.where(fr_valid,
+                         np.where(pp_valid, 3, 1), 0).astype(np.int8)
         out = _watershed_mfd_cpu(fr, labels, state, H, W)
 
     elif has_cuda_and_cupy() and is_cupy_array(data):