Add native Dask+CuPy backends for hydrology core functions (#952) (#966)

Replace CPU fallback with native GPU tile kernels for flow_accumulation,
watershed, basin, stream_order, stream_link, and snap_pour_point when
running on Dask+CuPy arrays. Each function now runs its existing CUDA
kernels per-tile with seed injection at tile boundaries, keeping data
GPU-resident throughout the iterative tile sweep. Also adds a native
CUDA kernel for snap_pour_point's single-GPU CuPy path.

Author: brendancol

README.md

@@ -286,13 +286,13 @@ In the GIS world, rasters are used for representing continuous phenomena (e.g. e
 | [Flow Direction (D8)](xrspatial/flow_direction.py) | Computes D8 flow direction from each cell toward the steepest downhill neighbor | ✅️ | ✅️ | ✅️ | ✅️ |
 | [Flow Direction (Dinf)](xrspatial/flow_direction_dinf.py) | Computes D-infinity flow direction as a continuous angle toward the steepest downslope facet | ✅️ | ✅️ | ✅️ | ✅️ |
 | [Flow Direction (MFD)](xrspatial/flow_direction_mfd.py) | Partitions flow to all downslope neighbors with an adaptive exponent (Qin et al. 2007) | ✅️ | ✅️ | ✅️ | ✅️ |
-| [Flow Accumulation (D8)](xrspatial/flow_accumulation.py) | Counts upstream cells draining through each cell in a D8 flow direction grid | ✅️ | ✅️ | ✅️ | 🔄 |
+| [Flow Accumulation (D8)](xrspatial/flow_accumulation.py) | Counts upstream cells draining through each cell in a D8 flow direction grid | ✅️ | ✅️ | ✅️ | ✅️ |
 | [Flow Accumulation (MFD)](xrspatial/flow_accumulation_mfd.py) | Accumulates upstream area through all MFD flow paths weighted by directional fractions | ✅️ | ✅️ | ✅️ | 🔄 |
-| [Watershed](xrspatial/watershed.py) | Labels each cell with the pour point it drains to via D8 flow direction | ✅️ | ✅️ | ✅️ | 🔄 |
-| [Basins](xrspatial/watershed.py) | Delineates drainage basins by labeling each cell with its outlet ID | ✅️ | ✅️ | ✅️ | 🔄 |
-| [Stream Order](xrspatial/stream_order.py) | Assigns Strahler or Shreve stream order to cells in a drainage network | ✅️ | ✅️ | ✅️ | 🔄 |
-| [Stream Link](xrspatial/stream_link.py) | Assigns unique IDs to each stream segment between junctions | ✅️ | ✅️ | ✅️ | 🔄 |
-| [Snap Pour Point](xrspatial/snap_pour_point.py) | Snaps pour points to the highest-accumulation cell within a search radius | ✅️ | ✅️ | 🔄 | 🔄 |
+| [Watershed](xrspatial/watershed.py) | Labels each cell with the pour point it drains to via D8 flow direction | ✅️ | ✅️ | ✅️ | ✅️ |
+| [Basins](xrspatial/watershed.py) | Delineates drainage basins by labeling each cell with its outlet ID | ✅️ | ✅️ | ✅️ | ✅️ |
+| [Stream Order](xrspatial/stream_order.py) | Assigns Strahler or Shreve stream order to cells in a drainage network | ✅️ | ✅️ | ✅️ | ✅️ |
+| [Stream Link](xrspatial/stream_link.py) | Assigns unique IDs to each stream segment between junctions | ✅️ | ✅️ | ✅️ | ✅️ |
+| [Snap Pour Point](xrspatial/snap_pour_point.py) | Snaps pour points to the highest-accumulation cell within a search radius | ✅️ | ✅️ | ✅️ | ✅️ |
 | [Flow Path](xrspatial/flow_path.py) | Traces downstream flow paths from start points through a D8 direction grid | ✅️ | ✅️ | 🔄 | 🔄 |
 | [HAND](xrspatial/hand.py) | Computes Height Above Nearest Drainage by tracing D8 flow to the nearest stream cell | ✅️ | ✅️ | 🔄 | 🔄 |
 | [TWI](xrspatial/twi.py) | Topographic Wetness Index: ln(specific catchment area / tan(slope)) | ✅️ | ✅️ | ✅️ | 🔄 |

xrspatial/basin.py

@@ -298,18 +298,33 @@ def _basins_make_pp_block(flow_dir_block, block_info=None):
 
 
 def _basins_dask_cupy(flow_dir_da):
-    """Dask+CuPy basins: convert to numpy, run CPU iterative, convert back."""
+    """Dask+CuPy basins: native GPU via watershed infrastructure."""
     import cupy as cp
+    from xrspatial.watershed import _watershed_dask_cupy
 
-    flow_dir_np = flow_dir_da.map_blocks(
-        lambda b: b.get(), dtype=flow_dir_da.dtype,
-        meta=np.array((), dtype=flow_dir_da.dtype),
-    )
-    result = _basins_dask_iterative(flow_dir_np)
-    return result.map_blocks(
-        cp.asarray, dtype=result.dtype,
-        meta=cp.array((), dtype=result.dtype),
-    )
+    chunks_y = flow_dir_da.chunks[0]
+    chunks_x = flow_dir_da.chunks[1]
+    total_h = sum(chunks_y)
+    total_w = sum(chunks_x)
+
+    def _basins_make_pp_block(flow_dir_block, block_info=None):
+        if block_info is None or 0 not in block_info:
+            return cp.full(flow_dir_block.shape, cp.nan, dtype=cp.float64)
+        row_off = block_info[0]['array-location'][0][0]
+        col_off = block_info[0]['array-location'][1][0]
+        h, w = flow_dir_block.shape
+        chunk_np = flow_dir_block.get() if hasattr(flow_dir_block, 'get') \
+            else np.asarray(flow_dir_block)
+        chunk_np = np.asarray(chunk_np, dtype=np.float64)
+        pp = _basins_init_labels(chunk_np, h, w, total_h, total_w,
+                                  row_off, col_off)
+        return cp.asarray(np.where(pp >= 0, pp, np.nan))
+
+    pour_points_da = da.map_blocks(
+        _basins_make_pp_block, flow_dir_da,
+        dtype=np.float64, meta=cp.array((), dtype=cp.float64))
+
+    return _watershed_dask_cupy(flow_dir_da, pour_points_da)
 
 
 # =====================================================================