Fixes #845: add configurable boundary mode to kernel-based spatial ops (#874)

Add a `boundary` parameter ('nan', 'nearest', 'reflect', 'wrap') to
slope, aspect, curvature, hillshade, convolve_2d, mean, apply,
focal_stats, and hotspots. Default is 'nan' (backward compatible).

Uses pad-compute-trim for numpy/cupy backends and passes boundary mode
through to dask map_overlap. Adds shared utilities (_validate_boundary,
_boundary_to_dask, _pad_array) in utils.py. Includes parametrized
cross-backend tests covering multiple array sizes, chunk sizes, and
boundary modes to verify numpy-vs-dask consistency.

Author: brendancol

xrspatial/aspect.py

@@ -16,7 +16,10 @@
 
 from xrspatial.utils import ArrayTypeFunctionMapping
 from xrspatial.utils import Z_UNITS
+from xrspatial.utils import _boundary_to_dask
 from xrspatial.utils import _extract_latlon_coords
+from xrspatial.utils import _pad_array
+from xrspatial.utils import _validate_boundary
 from xrspatial.utils import cuda_args
 from xrspatial.utils import ngjit
 from xrspatial.dataset_support import supports_dataset
@@ -54,7 +57,7 @@ class cupy(object):
 # =====================================================================
 
 @ngjit
-def _run_numpy(data: np.ndarray):
+def _cpu(data: np.ndarray):
     data = data.astype(np.float32)
     out = np.zeros_like(data, dtype=np.float32)
     out[:] = np.nan
@@ -90,6 +93,14 @@ def _run_numpy(data: np.ndarray):
     return out
 
 
+def _run_numpy(data: np.ndarray, boundary: str = 'nan') -> np.ndarray:
+    if boundary == 'nan':
+        return _cpu(data)
+    padded = _pad_array(data, 1, boundary)
+    result = _cpu(padded)
+    return result[1:-1, 1:-1]
+
+
 @cuda.jit(device=True)
 def _gpu(arr):
 
@@ -136,7 +147,12 @@ def _run_gpu(arr, out):
         out[i, j] = _gpu(arr[i-di:i+di+1, j-dj:j+dj+1])
 
 
-def _run_cupy(data: cupy.ndarray) -> cupy.ndarray:
+def _run_cupy(data: cupy.ndarray, boundary: str = 'nan') -> cupy.ndarray:
+    if boundary != 'nan':
+        padded = _pad_array(data, 1, boundary)
+        result = _run_cupy(padded)
+        return result[1:-1, 1:-1]
+
     data = data.astype(cupy.float32)
     griddim, blockdim = cuda_args(data.shape)
     out = cupy.empty(data.shape, dtype='f4')
@@ -145,22 +161,22 @@ def _run_cupy(data: cupy.ndarray) -> cupy.ndarray:
     return out
 
 
-def _run_dask_numpy(data: da.Array) -> da.Array:
+def _run_dask_numpy(data: da.Array, boundary: str = 'nan') -> da.Array:
     data = data.astype(np.float32)
-    _func = partial(_run_numpy)
+    _func = partial(_cpu)
     out = data.map_overlap(_func,
                            depth=(1, 1),
-                           boundary=np.nan,
+                           boundary=_boundary_to_dask(boundary),
                            meta=np.array(()))
     return out
 
 
-def _run_dask_cupy(data: da.Array) -> da.Array:
+def _run_dask_cupy(data: da.Array, boundary: str = 'nan') -> da.Array:
     data = data.astype(cupy.float32)
     _func = partial(_run_cupy)
     out = data.map_overlap(_func,
                            depth=(1, 1),
-                           boundary=cupy.nan,
+                           boundary=_boundary_to_dask(boundary, is_cupy=True),
                            meta=cupy.array(()))
     return out
 
@@ -169,7 +185,14 @@ def _run_dask_cupy(data: da.Array) -> da.Array:
 # Geodesic backend functions
 # =====================================================================
 
-def _run_numpy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor):
+def _run_numpy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor, boundary='nan'):
+    if boundary != 'nan':
+        data_p = _pad_array(data.astype(np.float64), 1, boundary)
+        lat_p = _pad_array(lat_2d, 1, boundary)
+        lon_p = _pad_array(lon_2d, 1, boundary)
+        stacked = np.stack([data_p, lat_p, lon_p], axis=0)
+        result = _cpu_geodesic_aspect(stacked, a2, b2, z_factor)
+        return result[1:-1, 1:-1]
     stacked = np.stack([
         data.astype(np.float64),
         lat_2d,
@@ -178,7 +201,22 @@ def _run_numpy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor):
     return _cpu_geodesic_aspect(stacked, a2, b2, z_factor)
 
 
-def _run_cupy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor):
+def _run_cupy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor, boundary='nan'):
+    if boundary != 'nan':
+        data_p = _pad_array(data.astype(cupy.float64), 1, boundary)
+        lat_p = _pad_array(cupy.asarray(lat_2d, dtype=cupy.float64), 1, boundary)
+        lon_p = _pad_array(cupy.asarray(lon_2d, dtype=cupy.float64), 1, boundary)
+        stacked = cupy.stack([data_p, lat_p, lon_p], axis=0)
+        H, W = data_p.shape
+        out = cupy.full((H, W), cupy.nan, dtype=cupy.float32)
+        a2_arr = cupy.array([a2], dtype=cupy.float64)
+        b2_arr = cupy.array([b2], dtype=cupy.float64)
+        zf_arr = cupy.array([z_factor], dtype=cupy.float64)
+        inv_2r_arr = cupy.array([INV_2R], dtype=cupy.float64)
+        griddim, blockdim = _geodesic_cuda_dims((H, W))
+        _run_gpu_geodesic_aspect[griddim, blockdim](stacked, a2_arr, b2_arr, zf_arr, inv_2r_arr, out)
+        return out[1:-1, 1:-1]
+
     lat_2d_gpu = cupy.asarray(lat_2d, dtype=cupy.float64)
     lon_2d_gpu = cupy.asarray(lon_2d, dtype=cupy.float64)
     stacked = cupy.stack([
@@ -227,7 +265,7 @@ def _dask_geodesic_aspect_chunk_cupy(stacked_chunk, a2, b2, z_factor):
     return out
 
 
-def _run_dask_numpy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor):
+def _run_dask_numpy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor, boundary='nan'):
     lat_dask = da.from_array(lat_2d, chunks=data.chunksize)
     lon_dask = da.from_array(lon_2d, chunks=data.chunksize)
     stacked = da.stack([
@@ -237,16 +275,17 @@ def _run_dask_numpy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor):
     ], axis=0).rechunk({0: 3})
 
     _func = partial(_dask_geodesic_aspect_chunk, a2=a2, b2=b2, z_factor=z_factor)
+    dask_bnd = _boundary_to_dask(boundary)
     out = stacked.map_overlap(
         _func,
         depth=(0, 1, 1),
-        boundary=np.nan,
+        boundary={0: np.nan, 1: dask_bnd, 2: dask_bnd},
         meta=np.array((), dtype=np.float32),
     )
     return out[0]
 
 
-def _run_dask_cupy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor):
+def _run_dask_cupy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor, boundary='nan'):
     lat_dask = da.from_array(cupy.asarray(lat_2d, dtype=cupy.float64),
                              chunks=data.chunksize)
     lon_dask = da.from_array(cupy.asarray(lon_2d, dtype=cupy.float64),
@@ -258,10 +297,11 @@ def _run_dask_cupy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor):
     ], axis=0).rechunk({0: 3})
 
     _func = partial(_dask_geodesic_aspect_chunk_cupy, a2=a2, b2=b2, z_factor=z_factor)
+    dask_bnd = _boundary_to_dask(boundary, is_cupy=True)
     out = stacked.map_overlap(
         _func,
         depth=(0, 1, 1),
-        boundary=cupy.nan,
+        boundary={0: cupy.nan, 1: dask_bnd, 2: dask_bnd},
         meta=cupy.array((), dtype=cupy.float32),
     )
     return out[0]
@@ -275,7 +315,8 @@ def _run_dask_cupy_geodesic(data, lat_2d, lon_2d, a2, b2, z_factor):
 def aspect(agg: xr.DataArray,
            name: Optional[str] = 'aspect',
            method: str = 'planar',
-           z_unit: str = 'meter') -> xr.DataArray:
+           z_unit: str = 'meter',
+           boundary: str = 'nan') -> xr.DataArray:
     """
     Calculates the aspect value of an elevation aggregate.
 
@@ -314,6 +355,12 @@ def aspect(agg: xr.DataArray,
         Unit of the elevation values.  Only used when ``method='geodesic'``.
         Accepted values: ``'meter'``, ``'foot'``, ``'kilometer'``, ``'mile'``
         (and common aliases).
+    boundary : str, default='nan'
+        How to handle edges where the kernel extends beyond the raster.
+        ``'nan'``     — fill missing neighbours with NaN (default).
+        ``'nearest'`` — repeat edge values.
+        ``'reflect'`` — mirror at boundary.
+        ``'wrap'``    — periodic / toroidal.
 
     Returns
     -------
@@ -353,6 +400,7 @@ def aspect(agg: xr.DataArray,
         raise ValueError(
             f"method must be 'planar' or 'geodesic', got {method!r}"
         )
+    _validate_boundary(boundary)
 
     if method == 'planar':
         mapper = ArrayTypeFunctionMapping(
@@ -361,7 +409,7 @@ def aspect(agg: xr.DataArray,
             cupy_func=_run_cupy,
             dask_cupy_func=_run_dask_cupy,
         )
-        out = mapper(agg)(agg.data)
+        out = mapper(agg)(agg.data, boundary=boundary)
 
     else:  # geodesic
         if z_unit not in Z_UNITS:
@@ -379,7 +427,7 @@ def aspect(agg: xr.DataArray,
             dask_func=_run_dask_numpy_geodesic,
             dask_cupy_func=_run_dask_cupy_geodesic,
         )
-        out = mapper(agg)(agg.data, lat_2d, lon_2d, WGS84_A2, WGS84_B2, z_factor)
+        out = mapper(agg)(agg.data, lat_2d, lon_2d, WGS84_A2, WGS84_B2, z_factor, boundary)
 
     return xr.DataArray(out,
                         name=name,

xrspatial/convolution.py

@@ -5,7 +5,8 @@
 import xarray as xr
 from numba import cuda, jit, prange
 
-from xrspatial.utils import (ArrayTypeFunctionMapping, cuda_args, get_dataarray_resolution,
+from xrspatial.utils import (ArrayTypeFunctionMapping, _boundary_to_dask, _pad_array,
+                             _validate_boundary, cuda_args, get_dataarray_resolution,
                              not_implemented_func)
 
 # 3rd-party
@@ -313,14 +314,28 @@ def _convolve_2d_numpy(data, kernel):
     return out
 
 
-def _convolve_2d_dask_numpy(data, kernel):
+def _convolve_2d_numpy_boundary(data, kernel, boundary='nan'):
+    if boundary == 'nan':
+        return _convolve_2d_numpy(data, kernel)
+    pad_h = kernel.shape[0] // 2
+    pad_w = kernel.shape[1] // 2
+    padded = _pad_array(data, (pad_h, pad_w), boundary)
+    result = _convolve_2d_numpy(padded, kernel)
+    r0 = pad_h if pad_h else None
+    r1 = -pad_h if pad_h else None
+    c0 = pad_w if pad_w else None
+    c1 = -pad_w if pad_w else None
+    return result[r0:r1, c0:c1]
+
+
+def _convolve_2d_dask_numpy(data, kernel, boundary='nan'):
     data = data.astype(np.float32)
     pad_h = kernel.shape[0] // 2
     pad_w = kernel.shape[1] // 2
     _func = partial(_convolve_2d_numpy, kernel=kernel)
     out = data.map_overlap(_func,
                            depth=(pad_h, pad_w),
-                           boundary=np.nan,
+                           boundary=_boundary_to_dask(boundary),
                            meta=np.array(()))
     return out
 
@@ -365,7 +380,18 @@ def _convolve_2d_cuda(data, kernel, out):
     out[i, j] = s
 
 
-def _convolve_2d_cupy(data, kernel):
+def _convolve_2d_cupy(data, kernel, boundary='nan'):
+    if boundary != 'nan':
+        pad_h = kernel.shape[0] // 2
+        pad_w = kernel.shape[1] // 2
+        padded = _pad_array(data, (pad_h, pad_w), boundary)
+        result = _convolve_2d_cupy(padded, kernel)
+        r0 = pad_h if pad_h else None
+        r1 = -pad_h if pad_h else None
+        c0 = pad_w if pad_w else None
+        c1 = -pad_w if pad_w else None
+        return result[r0:r1, c0:c1]
+
     data = data.astype(cupy.float32)
     out = cupy.empty(data.shape, dtype='f4')
     out[:, :] = cupy.nan
@@ -374,30 +400,31 @@ def _convolve_2d_cupy(data, kernel):
     return out
 
 
-def _convolve_2d_dask_cupy(data, kernel):
+def _convolve_2d_dask_cupy(data, kernel, boundary='nan'):
     data = data.astype(cupy.float32)
     pad_h = kernel.shape[0] // 2
     pad_w = kernel.shape[1] // 2
     _func = partial(_convolve_2d_cupy, kernel=kernel)
     out = data.map_overlap(_func,
                            depth=(pad_h, pad_w),
-                           boundary=cupy.nan,
+                           boundary=_boundary_to_dask(boundary, is_cupy=True),
                            meta=cupy.array(()))
     return out
 
 
-def convolve_2d(data, kernel):
+def convolve_2d(data, kernel, boundary='nan'):
+    _validate_boundary(boundary)
     mapper = ArrayTypeFunctionMapping(
-        numpy_func=_convolve_2d_numpy,
+        numpy_func=_convolve_2d_numpy_boundary,
         cupy_func=_convolve_2d_cupy,
         dask_func=_convolve_2d_dask_numpy,
         dask_cupy_func=_convolve_2d_dask_cupy
     )
-    out = mapper(xr.DataArray(data))(data, kernel)
+    out = mapper(xr.DataArray(data))(data, kernel, boundary)
     return out
 
 
-def convolution_2d(agg, kernel, name='convolution_2d'):
+def convolution_2d(agg, kernel, name='convolution_2d', boundary='nan'):
     """
     Calculates, for all inner cells of an array, the 2D convolution of
     each cell. Convolution is frequently used for image
@@ -413,6 +440,12 @@ def convolution_2d(agg, kernel, name='convolution_2d'):
     kernel : array-like object
         Impulse kernel, determines area to apply impulse function for
         each cell.
+    boundary : str, default='nan'
+        How to handle edges where the kernel extends beyond the raster.
+        ``'nan'``     -- fill missing neighbours with NaN (default).
+        ``'nearest'`` -- repeat edge values.
+        ``'reflect'`` -- mirror at boundary.
+        ``'wrap'``    -- periodic / toroidal.
 
     Returns
     -------
@@ -513,7 +546,7 @@ def convolution_2d(agg, kernel, name='convolution_2d'):
     """
 
     # wrapper of convolve_2d
-    out = convolve_2d(agg.data, kernel)
+    out = convolve_2d(agg.data, kernel, boundary)
     return xr.DataArray(out,
                         name=name,
                         coords=agg.coords,