Add 3D multi-band support to focal functions (#924)

* Add 3D multi-band support to focal functions (mean, apply, focal_stats, hotspots)

Slices along the first dimension, applies the 2D focal op per band,
and concatenates results. Fixes _apply_per_band parameter name
collision where `func` was passed both positionally and as a keyword
argument. Removes stale validation tests that expected 3D rejection.

* Mark GPU CPU-fallback backends in feature matrix

Flag 17 cells across classify, polygonize, and hydrology sections
where CuPy or Dask+CuPy backends accept GPU input but convert to
numpy internally rather than running native GPU kernels.

Author: brendancol

README.md

@@ -139,13 +139,13 @@ 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 |
 |:----------:|:------------|:----------------------:|:--------------------:|:-------------------:|:------:|
-| [Box Plot](xrspatial/classify.py) | Classifies values into bins based on box plot quartile boundaries | ✅️ |✅ | ✅ |✅ |
+| [Box Plot](xrspatial/classify.py) | Classifies values into bins based on box plot quartile boundaries | ✅️ |✅ | ✅ | 🔄 |
 | [Equal Interval](xrspatial/classify.py) | Divides the value range into equal-width bins | ✅️ |✅ | ✅ |✅ |
-| [Head/Tail Breaks](xrspatial/classify.py) | Classifies heavy-tailed distributions using recursive mean splitting | ✅️ |✅ | ✅ |✅ |
-| [Maximum Breaks](xrspatial/classify.py) | Finds natural groupings by maximizing differences between sorted values | ✅️ |✅ | ✅ |✅ |
-| [Natural Breaks](xrspatial/classify.py) | Optimizes class boundaries to minimize within-class variance (Jenks) | ✅️ |✅ | ✅ |✅ |
-| [Percentiles](xrspatial/classify.py) | Assigns classes based on user-defined percentile breakpoints | ✅️ |✅ | ✅ |✅ |
-| [Quantile](xrspatial/classify.py) | Distributes values into classes with equal observation counts | ✅️ |✅ | ✅ |✅ |
+| [Head/Tail Breaks](xrspatial/classify.py) | Classifies heavy-tailed distributions using recursive mean splitting | ✅️ |✅ | 🔄 | 🔄 |
+| [Maximum Breaks](xrspatial/classify.py) | Finds natural groupings by maximizing differences between sorted values | ✅️ |✅ | 🔄 | 🔄 |
+| [Natural Breaks](xrspatial/classify.py) | Optimizes class boundaries to minimize within-class variance (Jenks) | ✅️ |✅ | 🔄 | 🔄 |
+| [Percentiles](xrspatial/classify.py) | Assigns classes based on user-defined percentile breakpoints | ✅️ |✅ | ✅ | 🔄 |
+| [Quantile](xrspatial/classify.py) | Distributes values into classes with equal observation counts | ✅️ |✅ | ✅ | 🔄 |
 | [Reclassify](xrspatial/classify.py) | Remaps pixel values to new classes using a user-defined lookup | ✅️ |✅ | ✅ |✅ |
 | [Std Mean](xrspatial/classify.py) | Classifies values by standard deviation intervals from the mean | ✅️ |✅ | ✅ |✅ |
 
@@ -216,7 +216,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 |
 |:-----|:------------|:------------------:|:-----------------:|:---------------------:|:---------------------:|
-| [Polygonize](xrspatial/polygonize.py) | Converts contiguous regions of equal value into vector polygons | ✅️ | ✅️ | ✅️ | ✅️ |
+| [Polygonize](xrspatial/polygonize.py) | Converts contiguous regions of equal value into vector polygons | ✅️ | ✅️ | ✅️ | 🔄 |
 
 --------
 
@@ -245,13 +245,13 @@ In the GIS world, rasters are used for representing continuous phenomena (e.g. e
 |:----------:|:------------|:----------------------:|:--------------------:|:-------------------:|:------:|
 | [Flow Direction (D8)](xrspatial/flow_direction.py) | Computes D8 flow direction from each cell toward the steepest downhill neighbor | ✅️ | ✅️ | ✅️ | ✅️ |
 | [Flow Direction (Dinf)](xrspatial/flow_direction_dinf.py) | Computes D-infinity flow direction as a continuous angle toward the steepest downslope facet | ✅️ | ✅️ | ✅️ | ✅️ |
-| [Flow Accumulation (D8)](xrspatial/flow_accumulation.py) | Counts upstream cells draining through each cell in a D8 flow direction grid | ✅️ | ✅️ | ✅️ | ✅️ |
-| [Watershed](xrspatial/watershed.py) | Labels each cell with the pour point it drains to via D8 flow direction | ✅️ | ✅️ | ✅️ | ✅️ |
-| [Basins](xrspatial/watershed.py) | Delineates drainage basins by labeling each cell with its outlet ID | ✅️ | ✅️ | ✅️ | ✅️ |
-| [Stream Order](xrspatial/stream_order.py) | Assigns Strahler or Shreve stream order to cells in a drainage network | ✅️ | ✅️ | ✅️ | ✅️ |
-| [Stream Link](xrspatial/stream_link.py) | Assigns unique IDs to each stream segment between junctions | ✅️ | ✅️ | ✅️ | ✅️ |
-| [Snap Pour Point](xrspatial/snap_pour_point.py) | Snaps pour points to the highest-accumulation cell within a search radius | ✅️ | ✅️ | ✅️ | ✅️ |
-| [Flow Path](xrspatial/flow_path.py) | Traces downstream flow paths from start points through a D8 direction grid | ✅️ | ✅️ | ✅️ | ✅️ |
+| [Flow Accumulation (D8)](xrspatial/flow_accumulation.py) | Counts upstream cells draining through each cell in a D8 flow direction grid | ✅️ | ✅️ | ✅️ | 🔄 |
+| [Watershed](xrspatial/watershed.py) | Labels each cell with the pour point it drains to via D8 flow direction | ✅️ | ✅️ | ✅️ | 🔄 |
+| [Basins](xrspatial/watershed.py) | Delineates drainage basins by labeling each cell with its outlet ID | ✅️ | ✅️ | ✅️ | 🔄 |
+| [Stream Order](xrspatial/stream_order.py) | Assigns Strahler or Shreve stream order to cells in a drainage network | ✅️ | ✅️ | ✅️ | 🔄 |
+| [Stream Link](xrspatial/stream_link.py) | Assigns unique IDs to each stream segment between junctions | ✅️ | ✅️ | ✅️ | 🔄 |
+| [Snap Pour Point](xrspatial/snap_pour_point.py) | Snaps pour points to the highest-accumulation cell within a search radius | ✅️ | ✅️ | 🔄 | 🔄 |
+| [Flow Path](xrspatial/flow_path.py) | Traces downstream flow paths from start points through a D8 direction grid | ✅️ | ✅️ | 🔄 | 🔄 |
 
 -----------
 

xrspatial/focal.py

@@ -35,6 +35,21 @@ class cupy(object):
                              cuda_args, ngjit, not_implemented_func)
 from xrspatial.dataset_support import supports_dataset
 
+
+def _apply_per_band(band_func, agg, *args, **kwargs):
+    """Apply a 2D focal function independently to each band of a 3D array.
+
+    Slices along the first dimension, calls *band_func* on each 2D slice,
+    and stacks the results back together.
+    """
+    band_dim = agg.dims[0]
+    slices = []
+    for i in range(agg.sizes[band_dim]):
+        band = agg.isel({band_dim: i})
+        slices.append(band_func(band, *args, **kwargs))
+    return xr.concat(slices, dim=band_dim)
+
+
 # TODO: Make convolution more generic with numba first-class functions.
 
 
@@ -281,9 +296,14 @@ def mean(agg, passes=1, excludes=[np.nan], name='mean', boundary='nan'):
         Dimensions without coordinates: dim_0, dim_1
     """
 
-    _validate_raster(agg, func_name='mean', name='agg')
+    _validate_raster(agg, func_name='mean', name='agg', ndim=(2, 3))
     _validate_scalar(passes, func_name='mean', name='passes', dtype=int, min_val=1)
     _validate_boundary(boundary)
+
+    if agg.ndim == 3:
+        return _apply_per_band(mean, agg, passes=passes, excludes=excludes,
+                               name=name, boundary=boundary)
+
     out = agg.data.astype(float)
     for i in range(passes):
         out = _mean(out, tuple(excludes), boundary)
@@ -512,7 +532,11 @@ def apply(raster, kernel, func=_calc_mean, name='focal_apply', boundary='nan'):
            [2. , 2. , 2. , 1.5]])
     Dimensions without coordinates: y, x
     """
-    _validate_raster(raster, func_name='apply', name='raster')
+    _validate_raster(raster, func_name='apply', name='raster', ndim=(2, 3))
+
+    if raster.ndim == 3:
+        return _apply_per_band(apply, raster, kernel=kernel, func=func,
+                               name=name, boundary=boundary)
 
     # Validate the kernel
     kernel = custom_kernel(kernel)
@@ -957,7 +981,11 @@ def focal_stats(agg,
           * stats    (stats) object 'min' 'sum'
         Dimensions without coordinates: dim_0, dim_1
     """
-    _validate_raster(agg, func_name='focal_stats', name='agg')
+    _validate_raster(agg, func_name='focal_stats', name='agg', ndim=(2, 3))
+
+    if agg.ndim == 3:
+        return _apply_per_band(focal_stats, agg, kernel=kernel,
+                               stats_funcs=stats_funcs, boundary=boundary)
 
     # Validate the kernel
     kernel = custom_kernel(kernel)
@@ -1237,7 +1265,11 @@ def hotspots(raster, kernel, boundary='nan'):
         Dimensions without coordinates: dim_0, dim_1
     """
 
-    _validate_raster(raster, func_name='hotspots', name='raster')
+    _validate_raster(raster, func_name='hotspots', name='raster', ndim=(2, 3))
+
+    if raster.ndim == 3:
+        return _apply_per_band(hotspots, raster, kernel=kernel,
+                               boundary=boundary)
 
     _validate_boundary(boundary)
 

xrspatial/tests/test_focal.py

@@ -756,3 +756,119 @@ def test_convolution_2d_boundary_no_nan(boundary):
     assert not np.any(np.isnan(da_result.data.compute()))
     np.testing.assert_allclose(
         np_result.data, da_result.data.compute(), equal_nan=True, rtol=1e-5)
+
+
+# --- 3D (multi-band) focal tests ---
+
+
+@pytest.fixture
+def rgb_data():
+    rng = np.random.default_rng(123)
+    return rng.random((3, 12, 14)).astype(np.float64)
+
+
+def test_mean_3d_numpy(rgb_data):
+    agg = xr.DataArray(rgb_data, dims=['band', 'y', 'x'])
+    result = mean(agg)
+    assert result.shape == (3, 12, 14)
+    assert result.dims == ('band', 'y', 'x')
+    for i in range(3):
+        band_result = mean(agg.isel(band=i))
+        np.testing.assert_allclose(result.isel(band=i).data, band_result.data)
+
+
+@dask_array_available
+def test_mean_3d_dask(rgb_data):
+    dask_data = da.from_array(rgb_data, chunks=(1, 6, 7))
+    agg = xr.DataArray(dask_data, dims=['band', 'y', 'x'])
+    result = mean(agg)
+    assert result.shape == (3, 12, 14)
+    # compare against numpy per-band
+    numpy_agg = xr.DataArray(rgb_data, dims=['band', 'y', 'x'])
+    numpy_result = mean(numpy_agg)
+    np.testing.assert_allclose(
+        result.data.compute(), numpy_result.data, equal_nan=True, rtol=1e-5)
+
+
+def test_apply_3d_numpy(rgb_data):
+    kernel = np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]])
+    agg = xr.DataArray(rgb_data, dims=['band', 'y', 'x'])
+    result = apply(agg, kernel)
+    assert result.shape == (3, 12, 14)
+    assert result.dims == ('band', 'y', 'x')
+    for i in range(3):
+        band_result = apply(agg.isel(band=i), kernel)
+        np.testing.assert_allclose(result.isel(band=i).data, band_result.data)
+
+
+@dask_array_available
+def test_apply_3d_dask(rgb_data):
+    kernel = np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]])
+    dask_data = da.from_array(rgb_data, chunks=(1, 6, 7))
+    agg = xr.DataArray(dask_data, dims=['band', 'y', 'x'])
+    result = apply(agg, kernel)
+    assert result.shape == (3, 12, 14)
+    numpy_agg = xr.DataArray(rgb_data, dims=['band', 'y', 'x'])
+    numpy_result = apply(numpy_agg, kernel)
+    np.testing.assert_allclose(
+        result.data.compute(), numpy_result.data, equal_nan=True, rtol=1e-5)
+
+
+def test_focal_stats_3d_numpy(rgb_data):
+    kernel = custom_kernel(np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]]))
+    stats = ['mean', 'max']
+    agg = xr.DataArray(rgb_data, dims=['band', 'y', 'x'])
+    result = focal_stats(agg, kernel, stats_funcs=stats)
+    # 3D input -> 4D output: (band, stats, y, x)
+    assert result.shape == (3, 2, 12, 14)
+    for i in range(3):
+        band_result = focal_stats(agg.isel(band=i), kernel, stats_funcs=stats)
+        np.testing.assert_allclose(
+            result.isel(band=i).data, band_result.data, equal_nan=True)
+
+
+@dask_array_available
+def test_focal_stats_3d_dask(rgb_data):
+    kernel = custom_kernel(np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]]))
+    stats = ['mean', 'max']
+    dask_data = da.from_array(rgb_data, chunks=(1, 6, 7))
+    agg = xr.DataArray(dask_data, dims=['band', 'y', 'x'])
+    result = focal_stats(agg, kernel, stats_funcs=stats)
+    assert result.shape == (3, 2, 12, 14)
+    numpy_agg = xr.DataArray(rgb_data, dims=['band', 'y', 'x'])
+    numpy_result = focal_stats(numpy_agg, kernel, stats_funcs=stats)
+    np.testing.assert_allclose(
+        result.data.compute(), numpy_result.data, equal_nan=True, rtol=1e-5)
+
+
+def test_hotspots_3d_numpy():
+    rng = np.random.default_rng(42)
+    data_2d = rng.standard_normal((10, 12)).astype(np.float64)
+    # stack 3 copies with different scales to avoid zero-std bands
+    data_3d = np.stack([data_2d, data_2d * 2, data_2d * 0.5])
+    kernel = np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]], dtype=np.float64)
+    agg = xr.DataArray(data_3d, dims=['band', 'y', 'x'])
+    result = hotspots(agg, kernel)
+    assert result.shape == (3, 10, 12)
+    assert result.dims == ('band', 'y', 'x')
+    for i in range(3):
+        band_result = hotspots(agg.isel(band=i), kernel)
+        np.testing.assert_array_equal(result.isel(band=i).data, band_result.data)
+
+
+@dask_array_available
+def test_hotspots_3d_dask():
+    rng = np.random.default_rng(42)
+    data_2d = rng.standard_normal((10, 12)).astype(np.float64)
+    data_3d = np.stack([data_2d, data_2d * 2, data_2d * 0.5])
+    kernel = np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]], dtype=np.float64)
+    # numpy reference
+    numpy_agg = xr.DataArray(data_3d, dims=['band', 'y', 'x'])
+    numpy_result = hotspots(numpy_agg, kernel)
+    # dask
+    dask_data = da.from_array(data_3d, chunks=(1, 5, 6))
+    dask_agg = xr.DataArray(dask_data, dims=['band', 'y', 'x'])
+    dask_result = hotspots(dask_agg, kernel)
+    assert dask_result.shape == (3, 10, 12)
+    np.testing.assert_array_equal(
+        dask_result.data.compute(), numpy_result.data)

xrspatial/tests/test_validation.py

@@ -224,14 +224,6 @@ def test_surface_rejects_3d(self, func, args):
         with pytest.raises(ValueError, match="2D"):
             func(self._agg_3d, *args)
 
-    def test_mean_rejects_3d(self):
-        with pytest.raises(ValueError, match="2D"):
-            mean(self._agg_3d)
-
-    def test_focal_apply_rejects_3d(self):
-        with pytest.raises(ValueError, match="2D"):
-            focal_apply(self._agg_3d, _kernel_3x3)
-
     @pytest.mark.parametrize('func', [proximity, allocation, direction])
     def test_proximity_rejects_3d(self, func):
         with pytest.raises(ValueError, match="2D"):