Add GPU spill-to-CPU fallback for cost_distance (#863)

* Add GPU spill-to-CPU fallback for cost_distance (CuPy & Dask+CuPy)

Match the A* pathfinding pattern: accept CuPy-backed DataArrays by
transferring to CPU numpy, running the existing numba kernel, and
wrapping the result back as a CuPy array.  Update the README feature
matrix to reflect 🔄 (CPU fallback) support.

* Add GPU test parametrization for cost_distance CuPy/Dask+CuPy fallback

Tests verify the CPU-spill paths produce results identical to numpy:
- cupy matches numpy on random friction grid
- cupy respects max_cost truncation
- cupy input returns cupy-backed result
- dask+cupy matches numpy on multi-source grid

Author: brendancol

README.md

@@ -194,7 +194,7 @@ In the GIS world, rasters are used for representing continuous phenomena (e.g. e
 | Name | Description | NumPy xr.DataArray | Dask xr.DataArray | CuPy GPU xr.DataArray | Dask GPU xr.DataArray |
 |:----------:|:------------|:----------------------:|:--------------------:|:-------------------:|:------:|
 | [Allocation](xrspatial/proximity.py) | Assigns each cell to the identity of the nearest source feature | ✅️ | ✅ | | |
-| [Cost Distance](xrspatial/cost_distance.py) | Computes minimum accumulated cost to the nearest source through a friction surface | ✅️ | ✅ | | |
+| [Cost Distance](xrspatial/cost_distance.py) | Computes minimum accumulated cost to the nearest source through a friction surface | ✅️ | ✅ | 🔄 | 🔄 |
 | [Direction](xrspatial/proximity.py) | Computes the direction from each cell to the nearest source feature | ✅️ | ✅ | | |
 | [Proximity](xrspatial/proximity.py) | Computes the distance from each cell to the nearest source feature | ✅️ | ✅ | | |
 

xrspatial/cost_distance.py

@@ -38,7 +38,10 @@
 except ImportError:
     da = None
 
-from xrspatial.utils import get_dataarray_resolution, ngjit
+from xrspatial.utils import (
+    get_dataarray_resolution, ngjit,
+    has_cuda_and_cupy, is_cupy_array, is_dask_cupy,
+)
 from xrspatial.dataset_support import supports_dataset
 
 # ---------------------------------------------------------------------------
@@ -391,15 +394,50 @@ def cost_distance(
     source_data = raster.data
     friction_data = friction.data
 
-    if da is not None and isinstance(source_data, da.Array):
+    _is_dask = da is not None and isinstance(source_data, da.Array)
+    _is_cupy_backend = (
+        not _is_dask
+        and has_cuda_and_cupy()
+        and is_cupy_array(source_data)
+    )
+    _is_dask_cupy = _is_dask and has_cuda_and_cupy() and is_dask_cupy(raster)
+
+    if _is_dask:
         # Rechunk friction to match raster
         if isinstance(friction_data, da.Array):
             friction_data = friction_data.rechunk(source_data.chunks)
         else:
             friction_data = da.from_array(friction_data,
                                           chunks=source_data.chunks)
 
-    if isinstance(source_data, np.ndarray):
+    if _is_cupy_backend:
+        import cupy
+        # Transfer to CPU, run numpy kernel, transfer back
+        source_np = source_data.get()
+        friction_np = np.asarray(friction.data.get(), dtype=np.float64)
+        result_data = cupy.asarray(
+            _cost_distance_numpy(
+                source_np, friction_np,
+                cellsize_x, cellsize_y, max_cost_f,
+                target_values, dy, dx, dd,
+            )
+        )
+    elif _is_dask_cupy:
+        # Spill each chunk to CPU via map_overlap, then wrap back as dask
+        source_data = source_data.map_blocks(
+            lambda b: b.get(), dtype=source_data.dtype,
+            meta=np.array((), dtype=source_data.dtype),
+        )
+        friction_data = friction_data.map_blocks(
+            lambda b: b.get(), dtype=friction_data.dtype,
+            meta=np.array((), dtype=friction_data.dtype),
+        )
+        result_data = _cost_distance_dask(
+            source_data, friction_data,
+            cellsize_x, cellsize_y, max_cost_f,
+            target_values, dy, dx, dd,
+        )
+    elif isinstance(source_data, np.ndarray):
         if isinstance(friction_data, np.ndarray):
             result_data = _cost_distance_numpy(
                 source_data, friction_data,
@@ -408,7 +446,7 @@ def cost_distance(
             )
         else:
             raise TypeError("friction must be numpy-backed when raster is")
-    elif da is not None and isinstance(source_data, da.Array):
+    elif _is_dask:
         result_data = _cost_distance_dask(
             source_data, friction_data,
             cellsize_x, cellsize_y, max_cost_f,

xrspatial/tests/test_cost_distance.py

@@ -10,10 +10,12 @@
 import xarray as xr
 
 from xrspatial.cost_distance import cost_distance
+from xrspatial.tests.general_checks import cuda_and_cupy_available
+from xrspatial.utils import has_cuda_and_cupy, has_dask_array
 
 
 def _make_raster(data, backend='numpy', chunks=(3, 3)):
-    """Build a DataArray with y/x coords, optionally dask-backed."""
+    """Build a DataArray with y/x coords, optionally dask/cupy-backed."""
     h, w = data.shape
     raster = xr.DataArray(
         data.astype(np.float64),
@@ -24,13 +26,24 @@ def _make_raster(data, backend='numpy', chunks=(3, 3)):
     raster['x'] = np.arange(w, dtype=np.float64)
     if 'dask' in backend and da is not None:
         raster.data = da.from_array(raster.data, chunks=chunks)
+    if 'cupy' in backend and has_cuda_and_cupy():
+        import cupy
+        if isinstance(raster.data, da.Array):
+            raster.data = raster.data.map_blocks(cupy.asarray)
+        else:
+            raster.data = cupy.asarray(raster.data)
     return raster
 
 
 def _compute(arr):
-    """Extract numpy data from DataArray (works for numpy or dask)."""
+    """Extract numpy data from DataArray (works for numpy, dask, or cupy)."""
     if da is not None and isinstance(arr.data, da.Array):
-        return arr.values
+        val = arr.data.compute()
+        if hasattr(val, 'get'):
+            return val.get()
+        return val
+    if hasattr(arr.data, 'get'):
+        return arr.data.get()
     return arr.data
 
 
@@ -400,3 +413,91 @@ def test_source_on_impassable_cell(backend):
 
     # Everything should be NaN — the only source is on impassable terrain
     assert np.all(np.isnan(out))
+
+
+# -----------------------------------------------------------------------
+# CuPy GPU spill-to-CPU tests
+# -----------------------------------------------------------------------
+
+@cuda_and_cupy_available
+def test_cupy_matches_numpy():
+    """CuPy (CPU fallback) path should produce identical results to numpy."""
+    np.random.seed(42)
+    source = np.zeros((7, 7))
+    source[3, 3] = 1.0
+
+    friction_data = np.random.uniform(0.5, 5.0, (7, 7))
+
+    result_np = _compute(cost_distance(
+        _make_raster(source, backend='numpy'),
+        _make_raster(friction_data, backend='numpy'),
+    ))
+    result_cupy = _compute(cost_distance(
+        _make_raster(source, backend='cupy'),
+        _make_raster(friction_data, backend='cupy'),
+    ))
+
+    np.testing.assert_allclose(result_cupy, result_np, equal_nan=True, atol=1e-5)
+
+
+@cuda_and_cupy_available
+def test_cupy_max_cost():
+    """CuPy path respects max_cost truncation."""
+    source = np.zeros((1, 10))
+    source[0, 0] = 1.0
+    friction_data = np.ones((1, 10))
+
+    result = _compute(cost_distance(
+        _make_raster(source, backend='cupy'),
+        _make_raster(friction_data, backend='cupy'),
+        max_cost=3.5,
+    ))
+
+    np.testing.assert_allclose(result[0, 3], 3.0, atol=1e-5)
+    assert np.isnan(result[0, 4])
+
+
+@cuda_and_cupy_available
+def test_cupy_returns_cupy_array():
+    """Result should be CuPy-backed when input is CuPy-backed."""
+    import cupy
+    from xrspatial.utils import is_cupy_array
+
+    source = np.zeros((3, 3))
+    source[1, 1] = 1.0
+    friction_data = np.ones((3, 3))
+
+    result = cost_distance(
+        _make_raster(source, backend='cupy'),
+        _make_raster(friction_data, backend='cupy'),
+    )
+    assert is_cupy_array(result.data)
+
+
+# -----------------------------------------------------------------------
+# Dask + CuPy GPU spill-to-CPU tests
+# -----------------------------------------------------------------------
+
+@cuda_and_cupy_available
+@pytest.mark.skipif(not has_dask_array(), reason="Requires dask.Array")
+def test_dask_cupy_matches_numpy():
+    """Dask+CuPy (CPU fallback) should produce identical results to numpy."""
+    np.random.seed(42)
+    source = np.zeros((10, 12))
+    source[2, 3] = 1.0
+    source[7, 9] = 2.0
+
+    friction_data = np.random.uniform(0.5, 5.0, (10, 12))
+
+    result_np = _compute(cost_distance(
+        _make_raster(source, backend='numpy'),
+        _make_raster(friction_data, backend='numpy'),
+        max_cost=20.0,
+    ))
+    result_dc = _compute(cost_distance(
+        _make_raster(source, backend='dask+cupy', chunks=(5, 6)),
+        _make_raster(friction_data, backend='dask+cupy', chunks=(5, 6)),
+        max_cost=20.0,
+    ))
+
+    np.testing.assert_allclose(result_dc, result_np, equal_nan=True, atol=1e-5)