Handle cupy and dask+cupy segments in multi_stop_search (#3630)

Segment extraction assumed .get() meant cupy and everything else was
safe for np.asarray(seg.values), which crashes on dask-of-cupy results
(implicit cupy->numpy conversion). Route all backends through a
_segment_to_numpy helper (compute dask, then .get() cupy) in both the
stitching loop and _optimize_waypoint_order, and convert the stitched
output back to the input's array type so cupy and dask inputs no longer
come back numpy-backed. Fixes the no-op conversion expression in
test_multi_stop_cupy_matches_numpy and adds a dask+cupy test.

Author: brendancol

xrspatial/pathfinding.py

@@ -1284,6 +1284,20 @@ def _tour_cost(t):
     return tour, _tour_cost(tour)
 
 
+def _segment_to_numpy(seg_data):
+    """Materialize an a_star_search segment result as a numpy array.
+
+    Handles all four backends: numpy passes through, cupy uses
+    ``.get()``, dask computes first (yielding cupy blocks for
+    dask+cupy, which then also need ``.get()``).
+    """
+    if da is not None and isinstance(seg_data, da.Array):
+        seg_data = seg_data.compute()
+    if hasattr(seg_data, 'get'):  # cupy -> numpy
+        seg_data = seg_data.get()
+    return np.asarray(seg_data)
+
+
 def _optimize_waypoint_order(surface, waypoints, barriers, x, y,
                              connectivity, snap, friction, search_radius):
     """Build pairwise cost matrix and solve TSP with fixed endpoints.
@@ -1306,11 +1320,7 @@ def _optimize_waypoint_order(surface, waypoints, barriers, x, y,
                 snap_start=snap, snap_goal=snap,
                 friction=friction, search_radius=search_radius,
             )
-            seg_data = seg.data
-            if hasattr(seg_data, 'get'):
-                seg_vals = seg_data.get()
-            else:
-                seg_vals = np.asarray(seg.values)
+            seg_vals = _segment_to_numpy(seg.data)
             goal_py, goal_px = _get_pixel_id(waypoints[j], surface, x, y)
             goal_cost = seg_vals[goal_py, goal_px]
             if np.isfinite(goal_cost):
@@ -1452,11 +1462,7 @@ def multi_stop_search(surface: xr.DataArray,
             snap_start=snap, snap_goal=snap,
             friction=friction, search_radius=search_radius,
         )
-        seg_data = seg.data
-        if hasattr(seg_data, 'get'):
-            seg_vals = seg_data.get()  # cupy -> numpy
-        else:
-            seg_vals = np.asarray(seg.values)
+        seg_vals = _segment_to_numpy(seg.data)
 
         goal_py, goal_px = waypoint_pixels[i + 1]
 
@@ -1486,6 +1492,17 @@ def multi_stop_search(surface: xr.DataArray,
         segment_costs.append(float(seg_goal_cost))
         cumulative_cost += seg_goal_cost
 
+    # Match the input's array type, like a_star_search does
+    if _is_dask:
+        chunks = surface_data.chunks
+        path_data = da.from_array(path_data, chunks=chunks)
+        if has_cuda_and_cupy() and is_dask_cupy(surface):
+            import cupy
+            path_data = path_data.map_blocks(cupy.asarray)
+    elif has_cuda_and_cupy() and is_cupy_array(surface_data):
+        import cupy
+        path_data = cupy.asarray(path_data)
+
     path_agg = xr.DataArray(
         path_data,
         coords=surface.coords,

xrspatial/tests/test_pathfinding.py

@@ -978,6 +978,9 @@ def test_multi_stop_dask_matches_numpy():
     path_np = multi_stop_search(agg_np, [wp0, wp1, wp2])
     path_dask = multi_stop_search(agg_dask, [wp0, wp1, wp2])
 
+    # output should stay dask-backed, matching a_star_search
+    assert isinstance(path_dask.data, da.Array)
+
     np.testing.assert_allclose(
         np.asarray(path_dask.values),
         path_np.values,
@@ -999,13 +1002,49 @@ def test_multi_stop_cupy_matches_numpy():
     path_np = multi_stop_search(agg_np, [wp0, wp1, wp2])
     path_cupy = multi_stop_search(agg_cupy, [wp0, wp1, wp2])
 
+    # output should stay cupy-backed, matching a_star_search
+    import cupy
+    assert isinstance(path_cupy.data, cupy.ndarray)
+
     np.testing.assert_allclose(
-        path_cupy.data if not hasattr(path_cupy.data, 'get') else path_cupy.data,
+        path_cupy.data.get(),
         path_np.values,
         equal_nan=True, atol=1e-10,
     )
 
 
+@cuda_and_cupy_available
+@pytest.mark.skipif(not has_dask_array(), reason="Requires dask.Array")
+def test_multi_stop_dask_cupy_matches_numpy():
+    """Dask+CuPy multi-stop should not crash and should match numpy."""
+    data = np.ones((8, 8))
+    wp0 = (7.0, 0.0)
+    wp1 = (4.0, 3.0)
+    wp2 = (0.0, 7.0)
+
+    agg_np = _make_raster(data, backend='numpy')
+    agg_dc = _make_raster(data, backend='dask+cupy', chunks=(4, 4))
+
+    path_np = multi_stop_search(agg_np, [wp0, wp1, wp2])
+    path_dc = multi_stop_search(agg_dc, [wp0, wp1, wp2])
+
+    # output should stay dask-backed with cupy blocks
+    assert isinstance(path_dc.data, da.Array)
+    computed = path_dc.data.compute()
+    assert hasattr(computed, 'get')  # cupy blocks
+    np.testing.assert_allclose(
+        computed.get(),
+        path_np.values,
+        equal_nan=True, atol=1e-10,
+    )
+
+    # optimize_order goes through the same extraction path
+    path_dc_opt = multi_stop_search(
+        agg_dc, [wp0, wp1, wp2], optimize_order=True)
+    computed_opt = path_dc_opt.data.compute()
+    assert np.isfinite(computed_opt.get()).any()
+
+
 # =====================================================================
 # Issue #1439: input validation
 # =====================================================================