From c5c338661576ec268cd71da29ba33ff683c86919 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Fri, 3 Jul 2026 09:20:14 -0400 Subject: [PATCH 1/2] 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. --- xrspatial/pathfinding.py | 37 +++++++++++++++++++------- xrspatial/tests/test_pathfinding.py | 41 ++++++++++++++++++++++++++++- 2 files changed, 67 insertions(+), 11 deletions(-) diff --git a/xrspatial/pathfinding.py b/xrspatial/pathfinding.py index 8f7a23bd0..276a14846 100644 --- a/xrspatial/pathfinding.py +++ b/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, diff --git a/xrspatial/tests/test_pathfinding.py b/xrspatial/tests/test_pathfinding.py index 4fccd29fc..e0b2ded9a 100644 --- a/xrspatial/tests/test_pathfinding.py +++ b/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 # ===================================================================== From f98949566ebd1d38e66bb0432048a140d09e4fc0 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Fri, 3 Jul 2026 09:25:03 -0400 Subject: [PATCH 2/2] Document output array-type preservation in multi_stop_search (#3630) --- xrspatial/pathfinding.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/xrspatial/pathfinding.py b/xrspatial/pathfinding.py index 276a14846..596c6b080 100644 --- a/xrspatial/pathfinding.py +++ b/xrspatial/pathfinding.py @@ -1383,7 +1383,8 @@ def multi_stop_search(surface: xr.DataArray, Returns ------- xr.DataArray or xr.Dataset - Cumulative path cost surface. Attributes include + Cumulative path cost surface, backed by the same array type as + *surface* (numpy, cupy, dask, or dask+cupy). Attributes include ``waypoint_order``, ``segment_costs``, and ``total_cost``. A Dataset input returns a Dataset of per-variable results.