Add D-inf and MFD variants of flow_path, watershed, and HAND

examples/flood_simulation.py

@@ -40,10 +40,10 @@
 from matplotlib.animation import FuncAnimation
 
 from xrspatial import generate_terrain, slope
-from xrspatial.flow_direction import flow_direction
-from xrspatial.flow_accumulation import flow_accumulation
-from xrspatial.flow_length import flow_length
-from xrspatial.hand import hand
+from xrspatial.hydro.flow_direction_d8 import flow_direction_d8 as flow_direction
+from xrspatial.hydro.flow_accumulation_d8 import flow_accumulation_d8 as flow_accumulation
+from xrspatial.hydro.flow_length_d8 import flow_length_d8 as flow_length
+from xrspatial.hydro.hand_d8 import hand_d8 as hand
 from xrspatial.flood import (
     curve_number_runoff,
     flood_depth,

examples/landslide_risk.py

@@ -43,9 +43,9 @@
 
 from xrspatial import generate_terrain, slope
 from xrspatial.curvature import curvature
-from xrspatial.flow_direction import flow_direction
-from xrspatial.flow_accumulation import flow_accumulation
-from xrspatial.twi import twi
+from xrspatial.hydro.flow_direction_d8 import flow_direction_d8 as flow_direction
+from xrspatial.hydro.flow_accumulation_d8 import flow_accumulation_d8 as flow_accumulation
+from xrspatial.hydro.twi_d8 import twi_d8 as twi
 from xrspatial.terrain_metrics import tpi
 
 # -- Tunable parameters -----------------------------------------------------

examples/watershed_explorer.py

@@ -34,12 +34,12 @@
 import matplotlib.colors as mcolors
 
 from xrspatial import generate_terrain
-from xrspatial.flow_direction import flow_direction
-from xrspatial.flow_accumulation import flow_accumulation
-from xrspatial.stream_order import stream_order
-from xrspatial.snap_pour_point import snap_pour_point
-from xrspatial.watershed import watershed
-from xrspatial.basin import basin
+from xrspatial.hydro.flow_direction_d8 import flow_direction_d8 as flow_direction
+from xrspatial.hydro.flow_accumulation_d8 import flow_accumulation_d8 as flow_accumulation
+from xrspatial.hydro.stream_order_d8 import stream_order_d8 as stream_order
+from xrspatial.hydro.snap_pour_point_d8 import snap_pour_point_d8 as snap_pour_point
+from xrspatial.hydro.watershed_d8 import watershed_d8 as watershed
+from xrspatial.hydro.basin_d8 import basin_d8 as basin
 
 # -- Tunable parameters -----------------------------------------------------
 CELL_SIZE = 30.0                    # metres per pixel

xrspatial/__init__.py

@@ -23,7 +23,8 @@
 from xrspatial.curvature import curvature  # noqa
 from xrspatial.emerging_hotspots import emerging_hotspots  # noqa
 from xrspatial.erosion import erode  # noqa
-from xrspatial.fill import fill  # noqa
+from xrspatial.hydro import fill  # noqa: unified wrapper
+from xrspatial.hydro import fill_d8  # noqa
 from xrspatial.interpolate import idw  # noqa
 from xrspatial.interpolate import kriging  # noqa
 from xrspatial.interpolate import spline  # noqa
@@ -38,23 +39,22 @@
 from xrspatial.flood import flood_depth  # noqa
 from xrspatial.flood import inundation  # noqa
 from xrspatial.flood import travel_time  # noqa
-from xrspatial.flow_accumulation import flow_accumulation  # noqa
-from xrspatial.flow_accumulation_dinf import flow_accumulation_dinf  # noqa
-from xrspatial.flow_accumulation_mfd import flow_accumulation_mfd  # noqa
-from xrspatial.flow_direction import flow_direction  # noqa
-from xrspatial.flow_direction_dinf import flow_direction_dinf  # noqa
-from xrspatial.flow_direction_mfd import flow_direction_mfd  # noqa
-from xrspatial.flow_length import flow_length  # noqa
-from xrspatial.flow_length_dinf import flow_length_dinf  # noqa
-from xrspatial.flow_length_mfd import flow_length_mfd  # noqa
-from xrspatial.flow_path import flow_path  # noqa
+from xrspatial.hydro import flow_accumulation  # noqa: unified wrapper
+from xrspatial.hydro import flow_accumulation_d8, flow_accumulation_dinf, flow_accumulation_mfd  # noqa
+from xrspatial.hydro import flow_direction  # noqa: unified wrapper
+from xrspatial.hydro import flow_direction_d8, flow_direction_dinf, flow_direction_mfd  # noqa
+from xrspatial.hydro import flow_length  # noqa: unified wrapper
+from xrspatial.hydro import flow_length_d8, flow_length_dinf, flow_length_mfd  # noqa
+from xrspatial.hydro import flow_path  # noqa: unified wrapper
+from xrspatial.hydro import flow_path_d8, flow_path_dinf, flow_path_mfd  # noqa
 from xrspatial.focal import mean  # noqa
 from xrspatial.glcm import glcm_texture  # noqa
 from xrspatial.morphology import morph_closing  # noqa
 from xrspatial.morphology import morph_dilate  # noqa
 from xrspatial.morphology import morph_erode  # noqa
 from xrspatial.morphology import morph_opening  # noqa
-from xrspatial.hand import hand  # noqa
+from xrspatial.hydro import hand  # noqa: unified wrapper
+from xrspatial.hydro import hand_d8, hand_dinf, hand_mfd  # noqa
 from xrspatial.hillshade import hillshade  # noqa
 from xrspatial.mahalanobis import mahalanobis  # noqa
 from xrspatial.multispectral import arvi  # noqa
@@ -76,14 +76,14 @@
 from xrspatial.proximity import great_circle_distance  # noqa
 from xrspatial.proximity import manhattan_distance  # noqa
 from xrspatial.proximity import proximity  # noqa
-from xrspatial.sink import sink  # noqa
-from xrspatial.snap_pour_point import snap_pour_point  # noqa
-from xrspatial.stream_link import stream_link  # noqa
-from xrspatial.stream_link_dinf import stream_link_dinf  # noqa
-from xrspatial.stream_link_mfd import stream_link_mfd  # noqa
-from xrspatial.stream_order import stream_order  # noqa
-from xrspatial.stream_order_dinf import stream_order_dinf  # noqa
-from xrspatial.stream_order_mfd import stream_order_mfd  # noqa
+from xrspatial.hydro import sink  # noqa: unified wrapper
+from xrspatial.hydro import sink_d8  # noqa
+from xrspatial.hydro import snap_pour_point  # noqa: unified wrapper
+from xrspatial.hydro import snap_pour_point_d8  # noqa
+from xrspatial.hydro import stream_link  # noqa: unified wrapper
+from xrspatial.hydro import stream_link_d8, stream_link_dinf, stream_link_mfd  # noqa
+from xrspatial.hydro import stream_order  # noqa: unified wrapper
+from xrspatial.hydro import stream_order_d8, stream_order_dinf, stream_order_mfd  # noqa
 from xrspatial.sky_view_factor import sky_view_factor  # noqa
 from xrspatial.slope import slope  # noqa
 from xrspatial.surface_distance import surface_allocation  # noqa
@@ -95,12 +95,16 @@
 from xrspatial.terrain_metrics import roughness  # noqa
 from xrspatial.terrain_metrics import tpi  # noqa
 from xrspatial.terrain_metrics import tri  # noqa
-from xrspatial.twi import twi  # noqa
+from xrspatial.hydro import twi  # noqa: unified wrapper
+from xrspatial.hydro import twi_d8  # noqa
 from xrspatial.polygonize import polygonize  # noqa
 from xrspatial.viewshed import viewshed  # noqa
-from xrspatial.basin import basin  # noqa
-from xrspatial.watershed import basins  # noqa
-from xrspatial.watershed import watershed  # noqa
+from xrspatial.hydro import basin  # noqa: unified wrapper
+from xrspatial.hydro import basin_d8  # noqa
+from xrspatial.hydro import basins  # noqa: backward-compat alias
+from xrspatial.hydro import basins_d8  # noqa
+from xrspatial.hydro import watershed  # noqa: unified wrapper
+from xrspatial.hydro import watershed_d8, watershed_dinf, watershed_mfd  # noqa
 from xrspatial.zonal import apply as zonal_apply  # noqa
 from xrspatial.zonal import crop  # noqa
 from xrspatial.zonal import trim  # noqa

xrspatial/hydro/__init__.py

@@ -0,0 +1,149 @@
+"""Hydrology analysis modules for xarray-spatial.
+
+Includes flow direction, flow accumulation, flow length, flow path,
+watershed delineation, basin labeling, HAND, stream ordering, and
+related utilities for D8, D-infinity, and MFD routing.
+
+Each function family provides a unified wrapper that accepts a
+``routing`` parameter ('d8', 'dinf', or 'mfd') and dispatches to
+the corresponding implementation.  The suffixed variants
+(e.g. ``flow_direction_d8``) are also importable directly.
+"""
+
+# -- concrete D8 implementations ------------------------------------------
+from xrspatial.hydro.basin_d8 import basin_d8  # noqa
+from xrspatial.hydro.fill_d8 import fill_d8  # noqa
+from xrspatial.hydro.flow_accumulation_d8 import flow_accumulation_d8  # noqa
+from xrspatial.hydro.flow_direction_d8 import flow_direction_d8  # noqa
+from xrspatial.hydro.flow_length_d8 import flow_length_d8  # noqa
+from xrspatial.hydro.flow_path_d8 import flow_path_d8  # noqa
+from xrspatial.hydro.hand_d8 import hand_d8  # noqa
+from xrspatial.hydro.sink_d8 import sink_d8  # noqa
+from xrspatial.hydro.snap_pour_point_d8 import snap_pour_point_d8  # noqa
+from xrspatial.hydro.stream_link_d8 import stream_link_d8  # noqa
+from xrspatial.hydro.stream_order_d8 import stream_order_d8  # noqa
+from xrspatial.hydro.twi_d8 import twi_d8  # noqa
+from xrspatial.hydro.watershed_d8 import basins_d8  # noqa
+from xrspatial.hydro.watershed_d8 import watershed_d8  # noqa
+
+# -- concrete D-infinity implementations -----------------------------------
+from xrspatial.hydro.flow_accumulation_dinf import flow_accumulation_dinf  # noqa
+from xrspatial.hydro.flow_direction_dinf import flow_direction_dinf  # noqa
+from xrspatial.hydro.flow_length_dinf import flow_length_dinf  # noqa
+from xrspatial.hydro.flow_path_dinf import flow_path_dinf  # noqa
+from xrspatial.hydro.hand_dinf import hand_dinf  # noqa
+from xrspatial.hydro.stream_link_dinf import stream_link_dinf  # noqa
+from xrspatial.hydro.stream_order_dinf import stream_order_dinf  # noqa
+from xrspatial.hydro.watershed_dinf import watershed_dinf  # noqa
+
+# -- concrete MFD implementations -----------------------------------------
+from xrspatial.hydro.flow_accumulation_mfd import flow_accumulation_mfd  # noqa
+from xrspatial.hydro.flow_direction_mfd import flow_direction_mfd  # noqa
+from xrspatial.hydro.flow_length_mfd import flow_length_mfd  # noqa
+from xrspatial.hydro.flow_path_mfd import flow_path_mfd  # noqa
+from xrspatial.hydro.hand_mfd import hand_mfd  # noqa
+from xrspatial.hydro.stream_link_mfd import stream_link_mfd  # noqa
+from xrspatial.hydro.stream_order_mfd import stream_order_mfd  # noqa
+from xrspatial.hydro.watershed_mfd import watershed_mfd  # noqa
+
+
+# =========================================================================
+# Routing dispatch
+# =========================================================================
+
+class _RoutingDispatch:
+    """Map routing algorithm names to concrete implementations.
+
+    Inspired by ArrayTypeFunctionMapping but keyed on a string
+    (``'d8'``, ``'dinf'``, ``'mfd'``) rather than array type.
+    """
+
+    __slots__ = ('_name', '_impls')
+
+    def __init__(self, name, **impls):
+        self._name = name
+        self._impls = impls
+
+    def __call__(self, *args, routing='d8', **kwargs):
+        try:
+            fn = self._impls[routing]
+        except KeyError:
+            opts = ', '.join(repr(k) for k in self._impls)
+            raise ValueError(
+                f"Unknown routing {routing!r} for {self._name}; "
+                f"expected one of {opts}"
+            ) from None
+        return fn(*args, **kwargs)
+
+
+# -- 8 families with d8 / dinf / mfd variants ----------------------------
+
+flow_direction = _RoutingDispatch(
+    'flow_direction',
+    d8=flow_direction_d8, dinf=flow_direction_dinf, mfd=flow_direction_mfd,
+)
+
+flow_accumulation = _RoutingDispatch(
+    'flow_accumulation',
+    d8=flow_accumulation_d8, dinf=flow_accumulation_dinf,
+    mfd=flow_accumulation_mfd,
+)
+
+flow_length = _RoutingDispatch(
+    'flow_length',
+    d8=flow_length_d8, dinf=flow_length_dinf, mfd=flow_length_mfd,
+)
+
+flow_path = _RoutingDispatch(
+    'flow_path',
+    d8=flow_path_d8, dinf=flow_path_dinf, mfd=flow_path_mfd,
+)
+
+watershed = _RoutingDispatch(
+    'watershed',
+    d8=watershed_d8, dinf=watershed_dinf, mfd=watershed_mfd,
+)
+
+hand = _RoutingDispatch(
+    'hand',
+    d8=hand_d8, dinf=hand_dinf, mfd=hand_mfd,
+)
+
+stream_link = _RoutingDispatch(
+    'stream_link',
+    d8=stream_link_d8, dinf=stream_link_dinf, mfd=stream_link_mfd,
+)
+
+
+# stream_order needs special handling: the ordering param (strahler/shreve)
+# is called `ordering` in d8 and `method` in dinf/mfd.
+
+class _StreamOrderDispatch(_RoutingDispatch):
+    def __call__(self, *args, routing='d8', ordering='strahler', **kwargs):
+        try:
+            fn = self._impls[routing]
+        except KeyError:
+            opts = ', '.join(repr(k) for k in self._impls)
+            raise ValueError(
+                f"Unknown routing {routing!r} for {self._name}; "
+                f"expected one of {opts}"
+            ) from None
+        if routing == 'd8':
+            return fn(*args, ordering=ordering, **kwargs)
+        return fn(*args, method=ordering, **kwargs)
+
+
+stream_order = _StreamOrderDispatch(
+    'stream_order',
+    d8=stream_order_d8, dinf=stream_order_dinf, mfd=stream_order_mfd,
+)
+
+
+# -- 5 D8-only functions (future-proofed with routing param) --------------
+
+basin = _RoutingDispatch('basin', d8=basin_d8)
+basins = _RoutingDispatch('basins', d8=basins_d8)
+sink = _RoutingDispatch('sink', d8=sink_d8)
+snap_pour_point = _RoutingDispatch('snap_pour_point', d8=snap_pour_point_d8)
+fill = _RoutingDispatch('fill', d8=fill_d8)
+twi = _RoutingDispatch('twi', d8=twi_d8)

xrspatial/basin.py → xrspatial/hydro/basin_d8.py

@@ -272,7 +272,7 @@ def _basins_dask_iterative(flow_dir_da):
     Constructs basin pour_points lazily, then delegates to the
     watershed dask infrastructure.
     """
-    from xrspatial.watershed import _watershed_dask_iterative
+    from xrspatial.hydro.watershed_d8 import _watershed_dask_iterative
 
     chunks_y = flow_dir_da.chunks[0]
     chunks_x = flow_dir_da.chunks[1]
@@ -300,7 +300,7 @@ def _basins_make_pp_block(flow_dir_block, block_info=None):
 def _basins_dask_cupy(flow_dir_da):
     """Dask+CuPy basins: native GPU via watershed infrastructure."""
     import cupy as cp
-    from xrspatial.watershed import _watershed_dask_cupy
+    from xrspatial.hydro.watershed_d8 import _watershed_dask_cupy
 
     chunks_y = flow_dir_da.chunks[0]
     chunks_x = flow_dir_da.chunks[1]
@@ -332,7 +332,7 @@ def _basins_make_pp_block(flow_dir_block, block_info=None):
 # =====================================================================
 
 @supports_dataset
-def basin(flow_dir: xr.DataArray,
+def basin_d8(flow_dir: xr.DataArray,
           name: str = 'basin') -> xr.DataArray:
     """Delineate drainage basins: every cell labeled with its outlet ID.
 
@@ -360,7 +360,7 @@ def basin(flow_dir: xr.DataArray,
     data = flow_dir.data
 
     if isinstance(data, np.ndarray):
-        from xrspatial.watershed import _watershed_cpu
+        from xrspatial.hydro.watershed_d8 import _watershed_cpu
         fd = data.astype(np.float64)
         h, w = fd.shape
         labels = _basins_init_labels(fd, h, w, h, w, 0, 0)

xrspatial/fill.py → xrspatial/hydro/fill_d8.py

@@ -36,7 +36,7 @@ class cupy:  # type: ignore[no-redef]
     is_dask_cupy,
     ngjit,
 )
-from xrspatial._boundary_store import BoundaryStore
+from xrspatial.hydro._boundary_store import BoundaryStore
 from xrspatial.dataset_support import supports_dataset
 
 
@@ -449,7 +449,7 @@ def _fill_dask_cupy(dem_data):
 # =====================================================================
 
 @supports_dataset
-def fill(dem: xr.DataArray,
+def fill_d8(dem: xr.DataArray,
          z_limit=None,
          name: str = 'fill') -> xr.DataArray:
     """Fill depressions in a DEM using Planchon-Darboux iterative flooding.

xrspatial/flow_accumulation.py → xrspatial/hydro/flow_accumulation_d8.py

@@ -38,7 +38,7 @@ class cupy:  # type: ignore[no-redef]
     is_dask_cupy,
     ngjit,
 )
-from xrspatial._boundary_store import BoundaryStore
+from xrspatial.hydro._boundary_store import BoundaryStore
 from xrspatial.dataset_support import supports_dataset
 
 
@@ -881,7 +881,7 @@ def _flow_accum_dask_cupy(flow_dir_da):
 # =====================================================================
 
 @supports_dataset
-def flow_accumulation(flow_dir: xr.DataArray,
+def flow_accumulation_d8(flow_dir: xr.DataArray,
                       name: str = 'flow_accumulation') -> xr.DataArray:
     """Compute flow accumulation from a D8 flow direction grid.
 

xrspatial/flow_accumulation_dinf.py → xrspatial/hydro/flow_accumulation_dinf.py

@@ -37,9 +37,9 @@ class cupy:  # type: ignore[no-redef]
     is_dask_cupy,
     ngjit,
 )
-from xrspatial._boundary_store import BoundaryStore
+from xrspatial.hydro._boundary_store import BoundaryStore
 from xrspatial.dataset_support import supports_dataset
-from xrspatial.flow_accumulation import (
+from xrspatial.hydro.flow_accumulation_d8 import (
     _find_ready_and_finalize,
     _preprocess_tiles,
 )

xrspatial/flow_accumulation_mfd.py → xrspatial/hydro/flow_accumulation_mfd.py

@@ -37,7 +37,7 @@ class cupy:  # type: ignore[no-redef]
     is_dask_cupy,
     ngjit,
 )
-from xrspatial._boundary_store import BoundaryStore
+from xrspatial.hydro._boundary_store import BoundaryStore
 from xrspatial.dataset_support import supports_dataset
 
 # Neighbor offsets: E, SE, S, SW, W, NW, N, NE

xrspatial/flow_direction.py → xrspatial/hydro/flow_direction_d8.py

@@ -266,7 +266,7 @@ def _run_dask_cupy(data: da.Array,
 # =====================================================================
 
 @supports_dataset
-def flow_direction(agg: xr.DataArray,
+def flow_direction_d8(agg: xr.DataArray,
                    name: str = 'flow_direction',
                    boundary: str = 'nan') -> xr.DataArray:
     """Compute D8 flow direction for each cell.