Fixes #901: add GPU (CuPy) backends for proximity, allocation, direction

brendancol · brendancol · commit 39b2fdcce10a · 2026-02-26T08:36:07.000-08:00
Add CUDA brute-force nearest-target kernel with device functions for
Euclidean, Manhattan, and great-circle distance metrics. Each thread
processes one pixel, scanning all targets to find the nearest. Supports
proximity (distance), allocation (target value), and direction modes.

Adds _process_cupy() and _process_dask_cupy() host functions with
dispatch wired into _process(). Tests parametrized over cupy backend.
diff --git a/README.md b/README.md
@@ -201,10 +201,10 @@ 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 | ✅️ | ✅ | | |
+| [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 | ✅️ | ✅ | 🔄 | 🔄 |
-| [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 | ✅️ | ✅ | | |
+| [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 | ✅️ | ✅ | ✅️ | ✅️ |
 
 --------
 
diff --git a/xrspatial/proximity.py b/xrspatial/proximity.py
@@ -12,12 +12,23 @@
 except ImportError:
     cKDTree = None
 
+import math as _math
+
 import numpy as np
 import xarray as xr
-from numba import prange
+from numba import cuda, prange
+
+try:
+    import cupy
+except ImportError:
+    class cupy(object):
+        ndarray = False
 
 from xrspatial.pathfinding import _available_memory_bytes
-from xrspatial.utils import get_dataarray_resolution, ngjit
+from xrspatial.utils import (
+    cuda_args, get_dataarray_resolution, has_cuda_and_cupy,
+    is_cupy_array, is_dask_cupy, ngjit,
+)
 from xrspatial.dataset_support import supports_dataset
 
 EUCLIDEAN = 0
@@ -281,6 +292,144 @@ def _vectorized_calc_direction(x1, x2, y1, y2):
     return result.astype(np.float32)
 
 
+# =====================================================================
+# GPU (CuPy / CUDA) backend
+# =====================================================================
+
+@cuda.jit(device=True)
+def _gpu_euclidean_distance(x1, x2, y1, y2):
+    dx = x1 - x2
+    dy = y1 - y2
+    return _math.sqrt(dx * dx + dy * dy)
+
+
+@cuda.jit(device=True)
+def _gpu_manhattan_distance(x1, x2, y1, y2):
+    return abs(x1 - x2) + abs(y1 - y2)
+
+
+@cuda.jit(device=True)
+def _gpu_great_circle_distance(x1, x2, y1, y2):
+    lat1 = y1 * 0.017453292519943295
+    lon1 = x1 * 0.017453292519943295
+    lat2 = y2 * 0.017453292519943295
+    lon2 = x2 * 0.017453292519943295
+    dlon = lon2 - lon1
+    dlat = lat2 - lat1
+    a = (_math.sin(dlat / 2.0) ** 2
+         + _math.cos(lat1) * _math.cos(lat2)
+         * _math.sin(dlon / 2.0) ** 2)
+    return 6378137.0 * 2.0 * _math.asin(_math.sqrt(a))
+
+
+@cuda.jit(device=True)
+def _gpu_distance(x1, x2, y1, y2, metric):
+    if metric == EUCLIDEAN:
+        return _gpu_euclidean_distance(x1, x2, y1, y2)
+    elif metric == GREAT_CIRCLE:
+        return _gpu_great_circle_distance(x1, x2, y1, y2)
+    else:
+        return _gpu_manhattan_distance(x1, x2, y1, y2)
+
+
+@cuda.jit(device=True)
+def _gpu_calc_direction(x1, x2, y1, y2):
+    if x1 == x2 and y1 == y2:
+        return 0.0
+    dx = x2 - x1
+    dy = y2 - y1
+    d = _math.atan2(-dy, dx) * 57.29578
+    if d < 0.0:
+        d = 90.0 - d
+    elif d > 90.0:
+        d = 360.0 - d + 90.0
+    else:
+        d = 90.0 - d
+    return d
+
+
+@cuda.jit
+def _proximity_cuda_kernel(target_xs, target_ys, target_vals, n_targets,
+                           y_coords, x_coords, max_distance,
+                           distance_metric, process_mode, out):
+    iy, ix = cuda.grid(2)
+    if iy >= out.shape[0] or ix >= out.shape[1]:
+        return
+
+    px = x_coords[ix]
+    py = y_coords[iy]
+
+    best_dist = 1.0e38
+    best_idx = -1
+
+    for k in range(n_targets):
+        d = _gpu_distance(px, target_xs[k], py, target_ys[k], distance_metric)
+        if d < best_dist:
+            best_dist = d
+            best_idx = k
+
+    if best_idx >= 0 and best_dist <= max_distance:
+        if process_mode == PROXIMITY:
+            out[iy, ix] = best_dist
+        elif process_mode == ALLOCATION:
+            out[iy, ix] = target_vals[best_idx]
+        else:
+            out[iy, ix] = _gpu_calc_direction(
+                px, target_xs[best_idx], py, target_ys[best_idx])
+
+
+def _process_cupy(raster_data, x_coords, y_coords, target_values,
+                  max_distance, distance_metric, process_mode):
+    """GPU proximity using CUDA brute-force nearest-target kernel."""
+    import cupy as cp
+
+    # Find target pixels on GPU
+    if len(target_values) == 0:
+        mask = cp.isfinite(raster_data) & (raster_data != 0)
+    else:
+        mask = cp.isin(raster_data, cp.asarray(target_values))
+        mask &= cp.isfinite(raster_data)
+
+    target_rows, target_cols = cp.where(mask)
+    n_targets = int(target_rows.shape[0])
+
+    if n_targets == 0:
+        return cp.full(raster_data.shape, cp.nan, dtype=cp.float32)
+
+    # Collect target world-coordinates and values
+    y_dev = cp.asarray(y_coords, dtype=cp.float64)
+    x_dev = cp.asarray(x_coords, dtype=cp.float64)
+    target_ys = y_dev[target_rows]
+    target_xs = x_dev[target_cols]
+    target_vals = raster_data[target_rows, target_cols].astype(cp.float32)
+
+    # Pre-fill output with NaN (pixels with no target within range stay NaN)
+    out = cp.full(raster_data.shape, cp.nan, dtype=cp.float32)
+
+    griddim, blockdim = cuda_args(raster_data.shape)
+    _proximity_cuda_kernel[griddim, blockdim](
+        target_xs, target_ys, target_vals, n_targets,
+        y_dev, x_dev,
+        np.float64(max_distance),
+        np.int32(distance_metric),
+        np.int32(process_mode),
+        out,
+    )
+
+    return out
+
+
+def _process_dask_cupy(raster, x_coords, y_coords, target_values,
+                       max_distance, distance_metric, process_mode):
+    """Dask+CuPy backend: compute to cupy, run GPU kernel."""
+    import cupy as cp
+
+    cp_data = raster.data.compute()
+    result = _process_cupy(cp_data, x_coords, y_coords, target_values,
+                           max_distance, distance_metric, process_mode)
+    return da.from_array(result, chunks=raster.data.chunks)
+
+
 @ngjit
 def _process_proximity_line(
     source_line,
@@ -1062,47 +1211,66 @@ def _process_dask(raster, xs, ys):
         ys = np.repeat(y_coords, raster.shape[1]).reshape(raster.shape)
         result = _process_numpy(raster.data, xs, ys)
 
-    elif da is not None and isinstance(raster.data, da.Array):
-        use_kdtree = (
-            cKDTree is not None
-            and distance_metric in (EUCLIDEAN, MANHATTAN)
-            and max_distance >= max_possible_distance
+    elif has_cuda_and_cupy() and is_cupy_array(raster.data):
+        result = _process_cupy(
+            raster.data, x_coords, y_coords,
+            target_values, max_distance, distance_metric, process_mode,
         )
-        if use_kdtree:
-            result = _process_dask_kdtree(
+
+    elif da is not None and isinstance(raster.data, da.Array):
+        if has_cuda_and_cupy() and is_dask_cupy(raster):
+            result = _process_dask_cupy(
                 raster, x_coords, y_coords,
-                target_values, max_distance, distance_metric,
-                process_mode,
+                target_values, max_distance, distance_metric, process_mode,
             )
         else:
-            # Memory guard: unbounded distance on large rasters can OOM
-            if max_distance >= max_possible_distance:
-                H, W = raster.shape
-                required = H * W * 4 * 3  # raster + xs + ys, float32
-                avail = _available_memory_bytes()
-                if required > 0.8 * avail:
-                    if cKDTree is None:
-                        raise MemoryError(
-                            "Raster too large for single-chunk processing "
-                            "and scipy is not installed for memory-safe "
-                            "KDTree path. Install scipy or set a finite "
-                            "max_distance."
-                        )
-                    else:  # must be GREAT_CIRCLE
-                        raise MemoryError(
-                            "GREAT_CIRCLE with unbounded max_distance on "
-                            "this raster would exceed available memory. "
-                            "Set a finite max_distance."
-                        )
+            use_kdtree = (
+                cKDTree is not None
+                and distance_metric in (EUCLIDEAN, MANHATTAN)
+                and max_distance >= max_possible_distance
+            )
+            if use_kdtree:
+                result = _process_dask_kdtree(
+                    raster, x_coords, y_coords,
+                    target_values, max_distance, distance_metric,
+                    process_mode,
+                )
+            else:
+                # Memory guard: unbounded distance on large rasters can OOM
+                if max_distance >= max_possible_distance:
+                    H, W = raster.shape
+                    required = H * W * 4 * 3  # raster + xs + ys, float32
+                    avail = _available_memory_bytes()
+                    if required > 0.8 * avail:
+                        if cKDTree is None:
+                            raise MemoryError(
+                                "Raster too large for single-chunk processing "
+                                "and scipy is not installed for memory-safe "
+                                "KDTree path. Install scipy or set a finite "
+                                "max_distance."
+                            )
+                        else:  # must be GREAT_CIRCLE
+                            raise MemoryError(
+                                "GREAT_CIRCLE with unbounded max_distance on "
+                                "this raster would exceed available memory. "
+                                "Set a finite max_distance."
+                            )
 
-            # Existing path: build 2D coordinate arrays as dask arrays
-            x_coords_da = da.from_array(x_coords, chunks=x_coords.shape[0])
-            y_coords_da = da.from_array(y_coords, chunks=y_coords.shape[0])
-            xs = da.tile(x_coords_da, (raster.shape[0], 1))
-            ys = da.repeat(y_coords_da, raster.shape[1]).reshape(raster.shape)
-            xs = xs.rechunk(raster.chunks)
-            ys = ys.rechunk(raster.chunks)
-            result = _process_dask(raster, xs, ys)
+                # Existing path: build 2D coordinate arrays as dask arrays
+                x_coords_da = da.from_array(x_coords, chunks=x_coords.shape[0])
+                y_coords_da = da.from_array(y_coords, chunks=y_coords.shape[0])
+                xs = da.tile(x_coords_da, (raster.shape[0], 1))
+                ys = da.repeat(y_coords_da, raster.shape[1]).reshape(
+                    raster.shape)
+                xs = xs.rechunk(raster.chunks)
+                ys = ys.rechunk(raster.chunks)
+                result = _process_dask(raster, xs, ys)
+
+    else:
+        raise TypeError(
+            f"Unsupported array type {type(raster.data).__name__} "
+            f"for proximity/allocation/direction"
+        )
 
     return result
 
diff --git a/xrspatial/tests/test_proximity.py b/xrspatial/tests/test_proximity.py
@@ -11,7 +11,9 @@
 
 from xrspatial import allocation, direction, euclidean_distance, great_circle_distance, proximity
 from xrspatial.proximity import _calc_direction
-from xrspatial.tests.general_checks import general_output_checks, create_test_raster
+from xrspatial.tests.general_checks import (
+    general_output_checks, create_test_raster, has_cuda_and_cupy,
+)
 
 
 def test_great_circle_distance():
@@ -38,8 +40,11 @@ def test_raster(backend):
     raster = xr.DataArray(data, dims=['lat', 'lon'])
     raster['lon'] = _lon
     raster['lat'] = _lat
+    if has_cuda_and_cupy() and 'cupy' in backend:
+        import cupy
+        raster.data = cupy.asarray(data)
     if 'dask' in backend and da is not None:
-        raster.data = da.from_array(data, chunks=(4, 3))
+        raster.data = da.from_array(raster.data, chunks=(4, 3))
     return raster
 
 
@@ -167,43 +172,43 @@ def qgis_proximity_distance_target_values():
     return target_values, qgis_result
 
 
-@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy'])
+@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy', 'cupy'])
 def test_default_proximity(test_raster, result_default_proximity):
     default_prox = proximity(test_raster, x='lon', y='lat')
     general_output_checks(test_raster, default_prox, result_default_proximity, verify_dtype=True)
 
 
-@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy'])
+@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy', 'cupy'])
 def test_target_proximity(test_raster, result_target_proximity):
     target_values, expected_result = result_target_proximity
     target_prox = proximity(test_raster, x='lon', y='lat', target_values=target_values)
     general_output_checks(test_raster, target_prox, expected_result, verify_dtype=True)
 
 
-@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy'])
+@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy', 'cupy'])
 def test_manhattan_proximity(test_raster, result_manhattan_proximity):
     manhattan_prox = proximity(test_raster, x='lon', y='lat', distance_metric='MANHATTAN')
     general_output_checks(
         test_raster, manhattan_prox, result_manhattan_proximity, verify_dtype=True
     )
 
 
-@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy'])
+@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy', 'cupy'])
 def test_great_circle_proximity(test_raster, result_great_circle_proximity):
     great_circle_prox = proximity(test_raster, x='lon', y='lat', distance_metric='GREAT_CIRCLE')
     general_output_checks(
         test_raster, great_circle_prox, result_great_circle_proximity, verify_dtype=True
     )
 
 
-@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy'])
+@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy', 'cupy'])
 def test_max_distance_proximity(test_raster, result_max_distance_proximity):
     max_distance, expected_result = result_max_distance_proximity
     max_distance_prox = proximity(test_raster, x='lon', y='lat', max_distance=max_distance)
     general_output_checks(test_raster, max_distance_prox, expected_result, verify_dtype=True)
 
 
-@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy'])
+@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy', 'cupy'])
 def test_default_allocation(test_raster, result_default_allocation):
     allocation_agg = allocation(test_raster, x='lon', y='lat')
     general_output_checks(test_raster, allocation_agg, result_default_allocation, verify_dtype=True)
@@ -224,7 +229,7 @@ def test_default_allocation_against_proximity(test_raster, result_default_proxim
             np.testing.assert_allclose(result_default_proximity[y, x], d)
 
 
-@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy'])
+@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy', 'cupy'])
 def test_max_distance_allocation(test_raster, result_max_distance_allocation):
     max_distance, expected_result = result_max_distance_allocation
     max_distance_alloc = allocation(test_raster, x='lon', y='lat', max_distance=max_distance)
@@ -247,7 +252,7 @@ def test_calc_direction():
     assert (abs(output-expected_output) <= tolerance).all()
 
 
-@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy'])
+@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy', 'cupy'])
 def test_default_direction(test_raster, result_default_direction):
     direction_agg = direction(test_raster, x='lon', y='lat')
     general_output_checks(test_raster, direction_agg, result_default_direction)
@@ -267,7 +272,7 @@ def test_default_direction_against_allocation(test_raster, result_default_alloca
             np.testing.assert_allclose(direction_agg.data[y, x], d)
 
 
-@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy'])
+@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy', 'cupy'])
 def test_max_distance_direction(test_raster, result_max_distance_direction):
     max_distance, expected_result = result_max_distance_direction
     max_distance_direction = direction(test_raster, x='lon', y='lat', max_distance=max_distance)
