Expose multi_stop_search on the .xrs DataArray and Dataset accessors (#3564)

Author: brendancol

docs/source/reference/pathfinding.rst

@@ -21,3 +21,10 @@ A* Pathfinding
     :toctree: _autosummary
 
     xrspatial.pathfinding.a_star_search
+
+Multi-Stop Routing
+==================
+.. autosummary::
+    :toctree: _autosummary
+
+    xrspatial.pathfinding.multi_stop_search

xrspatial/accessor.py

@@ -1283,6 +1283,10 @@ def a_star_search(self, start, goal, **kwargs):
         from .pathfinding import a_star_search
         return a_star_search(self._obj, start, goal, **kwargs)
 
+    def multi_stop_search(self, waypoints, **kwargs):
+        from .pathfinding import multi_stop_search
+        return multi_stop_search(self._obj, waypoints, **kwargs)
+
     # ---- Zonal ----
 
     def zonal_stats(self, zones, **kwargs):
@@ -2059,6 +2063,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

@@ -91,7 +91,7 @@ def test_dataarray_accessor_has_expected_methods(elevation):
         'morph_erode', 'morph_dilate', 'morph_opening', 'morph_closing',
         'morph_gradient', 'morph_white_tophat', 'morph_black_tophat',
         'proximity', 'allocation', 'direction', 'cost_distance',
-        'a_star_search',
+        'a_star_search', 'multi_stop_search',
         'zonal_stats', 'zonal_apply', 'zonal_crosstab', 'crop', 'trim',
         'regions',
         'generate_terrain', 'perlin',
@@ -118,6 +118,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',
@@ -272,6 +273,45 @@ def test_ds_morph_gradient(elevation):
     xr.testing.assert_identical(result, expected)
 
 
+# ---------------------------------------------------------------------------
+# 4c. DataArray pathfinding — accessor matches direct call
+# ---------------------------------------------------------------------------
+
+def test_da_a_star_search(elevation):
+    from xrspatial.pathfinding import a_star_search
+    start, goal = (0, 0), (9, 9)
+    expected = a_star_search(elevation, start, goal)
+    result = elevation.xrs.a_star_search(start, goal)
+    xr.testing.assert_identical(result, expected)
+
+
+def test_da_multi_stop_search(elevation):
+    from xrspatial.pathfinding import multi_stop_search
+    waypoints = [(0, 0), (5, 5), (9, 9)]
+    expected = multi_stop_search(elevation, waypoints)
+    result = elevation.xrs.multi_stop_search(waypoints)
+    xr.testing.assert_identical(result, expected)
+
+
+def test_da_multi_stop_search_kwargs(elevation):
+    from xrspatial.pathfinding import multi_stop_search
+    waypoints = [(0, 0), (9, 0), (9, 9)]
+    expected = multi_stop_search(elevation, waypoints, optimize_order=True)
+    result = elevation.xrs.multi_stop_search(waypoints, optimize_order=True)
+    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
 # ---------------------------------------------------------------------------