Expose multi_stop_search on the Dataset accessor via supports_dataset (#3563)

Author: brendancol

xrspatial/accessor.py

@@ -1947,6 +1947,12 @@ def surface_direction(self, elevation, **kwargs):
         from .surface_distance import surface_direction
         return surface_direction(self._obj, elevation, **kwargs)
 
+    # ---- Pathfinding ----
+
+    def multi_stop_search(self, waypoints, **kwargs):
+        from .pathfinding import multi_stop_search
+        return multi_stop_search(self._obj, waypoints, **kwargs)
+
     # ---- Preview ----
 
     def preview(self, **kwargs):

xrspatial/pathfinding.py

@@ -15,6 +15,7 @@
     dask = None
 
 from xrspatial.cost_distance import _heap_push, _heap_pop
+from xrspatial.dataset_support import supports_dataset
 from xrspatial.utils import (
     _validate_raster,
     get_dataarray_resolution, ngjit,
@@ -1324,6 +1325,7 @@ def _optimize_waypoint_order(surface, waypoints, barriers, x, y,
     return [waypoints[i] for i in order]
 
 
+@supports_dataset
 def multi_stop_search(surface: xr.DataArray,
                       waypoints: list,
                       barriers: list = [],
@@ -1344,8 +1346,10 @@ def multi_stop_search(surface: xr.DataArray,
 
     Parameters
     ----------
-    surface : xr.DataArray
-        2-D elevation / cost surface.
+    surface : xr.DataArray or xr.Dataset
+        2-D elevation / cost surface.  A Dataset routes each data
+        variable through the same waypoints independently and returns
+        a Dataset of the per-variable results.
     waypoints : list of array-like
         Sequence of ``(y, x)`` coordinate pairs to visit.  Must contain
         at least two points.
@@ -1368,9 +1372,10 @@ def multi_stop_search(surface: xr.DataArray,
 
     Returns
     -------
-    xr.DataArray
+    xr.DataArray or xr.Dataset
         Cumulative path cost surface.  Attributes include
         ``waypoint_order``, ``segment_costs``, and ``total_cost``.
+        A Dataset input returns a Dataset of per-variable results.
 
     Raises
     ------

xrspatial/tests/test_accessor.py

@@ -117,6 +117,7 @@ def test_dataset_accessor_has_expected_methods():
         'morph_erode', 'morph_dilate', 'morph_opening', 'morph_closing',
         'morph_gradient', 'morph_white_tophat', 'morph_black_tophat',
         'proximity', 'allocation', 'direction', 'cost_distance',
+        'multi_stop_search',
         'ndvi', 'evi', 'arvi', 'savi', 'nbr', 'sipi',
         'rasterize',
         'validate',
@@ -299,6 +300,17 @@ def test_da_multi_stop_search_kwargs(elevation):
     xr.testing.assert_identical(result, expected)
 
 
+def test_ds_multi_stop_search(elevation):
+    from xrspatial.pathfinding import multi_stop_search
+    ds = xr.Dataset({'a': elevation, 'b': elevation + 100})
+    waypoints = [(0, 0), (5, 5), (9, 9)]
+    expected = multi_stop_search(ds, waypoints)
+    result = ds.xrs.multi_stop_search(waypoints)
+    xr.testing.assert_identical(result, expected)
+    # supports_dataset routes each variable through its own surface
+    assert set(result.data_vars) == {'a', 'b'}
+
+
 # ---------------------------------------------------------------------------
 # 5. Dataset single-input — accessor matches direct call
 # ---------------------------------------------------------------------------