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.

Author: brendancol

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"):