focal_stats(): gpu case (#709)

* focal stats gpu case

* flake8

* add tests

Author: thuydotm

xrspatial/focal.py

@@ -1,6 +1,6 @@
 import copy
 from functools import partial
-from math import isnan
+from math import isnan, sqrt
 
 import dask.array as da
 import numba as nb
@@ -420,6 +420,178 @@ def apply(raster, kernel, func=_calc_mean, name='focal_apply'):
     return result
 
 
+@cuda.jit
+def _focal_min_cuda(data, kernel, out):
+    i, j = cuda.grid(2)
+
+    delta_rows = kernel.shape[0] // 2
+    delta_cols = kernel.shape[1] // 2
+
+    data_rows, data_cols = data.shape
+
+    if i < delta_rows or i >= data_rows - delta_rows or \
+            j < delta_cols or j >= data_cols - delta_cols:
+        return
+
+    s = data[i, j]
+    for k in range(kernel.shape[0]):
+        for h in range(kernel.shape[1]):
+            i_k = i + k - delta_rows
+            j_h = j + h - delta_cols
+            if (i_k >= 0) and (i_k < data_rows) and (j_h >= 0) and (j_h < data_cols):
+                if (kernel[k, h] != 0) and s > data[i_k, j_h]:
+                    s = data[i_k, j_h]
+    out[i, j] = s
+
+
+@cuda.jit
+def _focal_max_cuda(data, kernel, out):
+    i, j = cuda.grid(2)
+
+    delta_rows = kernel.shape[0] // 2
+    delta_cols = kernel.shape[1] // 2
+
+    data_rows, data_cols = data.shape
+
+    if i < delta_rows or i >= data_rows - delta_rows or \
+            j < delta_cols or j >= data_cols - delta_cols:
+        return
+
+    s = data[i, j]
+    for k in range(kernel.shape[0]):
+        for h in range(kernel.shape[1]):
+            i_k = i + k - delta_rows
+            j_h = j + h - delta_cols
+            if (i_k >= 0) and (i_k < data_rows) and (j_h >= 0) and (j_h < data_cols):
+                if (kernel[k, h] != 0) and s < data[i_k, j_h]:
+                    s = data[i_k, j_h]
+    out[i, j] = s
+
+
+def _focal_range_cupy(data, kernel):
+    focal_min = _focal_stats_func_cupy(data, kernel, _focal_min_cuda)
+    focal_max = _focal_stats_func_cupy(data, kernel, _focal_max_cuda)
+    out = focal_max - focal_min
+    return out
+
+
+@cuda.jit
+def _focal_std_cuda(data, kernel, out):
+    i, j = cuda.grid(2)
+
+    delta_rows = kernel.shape[0] // 2
+    delta_cols = kernel.shape[1] // 2
+
+    data_rows, data_cols = data.shape
+
+    if i < delta_rows or i >= data_rows - delta_rows or \
+            j < delta_cols or j >= data_cols - delta_cols:
+        return
+
+    sum_squares = 0
+    sum = 0
+    count = 0
+    for k in range(kernel.shape[0]):
+        for h in range(kernel.shape[1]):
+            i_k = i + k - delta_rows
+            j_h = j + h - delta_cols
+            if (i_k >= 0) and (i_k < data_rows) and (j_h >= 0) and (j_h < data_cols):
+                sum_squares += (kernel[k, h]*data[i_k, j_h])**2
+                sum += kernel[k, h]*data[i_k, j_h]
+                count += kernel[k, h]
+    squared_sum = sum**2
+    out[i, j] = sqrt((sum_squares - squared_sum/count) / count)
+
+
+@cuda.jit
+def _focal_var_cuda(data, kernel, out):
+    i, j = cuda.grid(2)
+
+    delta_rows = kernel.shape[0] // 2
+    delta_cols = kernel.shape[1] // 2
+
+    data_rows, data_cols = data.shape
+
+    if i < delta_rows or i >= data_rows - delta_rows or \
+            j < delta_cols or j >= data_cols - delta_cols:
+        return
+
+    sum_squares = 0
+    sum = 0
+    count = 0
+    for k in range(kernel.shape[0]):
+        for h in range(kernel.shape[1]):
+            i_k = i + k - delta_rows
+            j_h = j + h - delta_cols
+            if (i_k >= 0) and (i_k < data_rows) and (j_h >= 0) and (j_h < data_cols):
+                sum_squares += (kernel[k, h]*data[i_k, j_h])**2
+                sum += kernel[k, h]*data[i_k, j_h]
+                count += kernel[k, h]
+    squared_sum = sum**2
+    out[i, j] = (sum_squares - squared_sum/count) / count
+
+
+def _focal_mean_cupy(data, kernel):
+    out = convolve_2d(data, kernel / kernel.sum())
+    return out
+
+
+def _focal_sum_cupy(data, kernel):
+    out = convolve_2d(data, kernel)
+    return out
+
+
+def _focal_stats_func_cupy(data, kernel, func=_focal_max_cuda):
+    out = cupy.empty(data.shape, dtype='f4')
+    out[:, :] = cupy.nan
+    griddim, blockdim = cuda_args(data.shape)
+    func[griddim, blockdim](data, kernel, cupy.asarray(out))
+    return out
+
+
+def _focal_stats_cupy(agg, kernel, stats_funcs):
+    _stats_cupy_mapper = dict(
+        mean=_focal_mean_cupy,
+        sum=_focal_sum_cupy,
+        range=_focal_range_cupy,
+        max=lambda *args: _focal_stats_func_cupy(*args, func=_focal_max_cuda),
+        min=lambda *args: _focal_stats_func_cupy(*args, func=_focal_min_cuda),
+        std=lambda *args: _focal_stats_func_cupy(*args, func=_focal_std_cuda),
+        var=lambda *args: _focal_stats_func_cupy(*args, func=_focal_var_cuda),
+    )
+    stats_aggs = []
+    for stats in stats_funcs:
+        data = agg.data.astype(cupy.float32)
+        stats_data = _stats_cupy_mapper[stats](data, kernel)
+        stats_agg = xr.DataArray(
+            stats_data,
+            dims=agg.dims,
+            coords=agg.coords,
+            attrs=agg.attrs
+          )
+        stats_aggs.append(stats_agg)
+    stats = xr.concat(stats_aggs, pd.Index(stats_funcs, name='stats'))
+    return stats
+
+
+def _focal_stats_cpu(agg, kernel, stats_funcs):
+    _function_mapping = {
+        'mean': _calc_mean,
+        'max': _calc_max,
+        'min': _calc_min,
+        'range': _calc_range,
+        'std': _calc_std,
+        'var': _calc_var,
+        'sum': _calc_sum
+    }
+    stats_aggs = []
+    for stats in stats_funcs:
+        stats_agg = apply(agg, kernel, func=_function_mapping[stats])
+        stats_aggs.append(stats_agg)
+    stats = xr.concat(stats_aggs, pd.Index(stats_funcs, name='stats'))
+    return stats
+
+
 def focal_stats(agg,
                 kernel,
                 stats_funcs=[
@@ -480,24 +652,25 @@ def focal_stats(agg,
           * stats    (stats) object 'min' 'sum'
         Dimensions without coordinates: dim_0, dim_1
     """
+    # validate raster
+    if not isinstance(agg, DataArray):
+        raise TypeError("`agg` must be instance of DataArray")
 
-    _function_mapping = {
-        'mean': _calc_mean,
-        'max': _calc_max,
-        'min': _calc_min,
-        'range': _calc_range,
-        'std': _calc_std,
-        'var': _calc_var,
-        'sum': _calc_sum
-    }
+    if agg.ndim != 2:
+        raise ValueError("`agg` must be 2D")
 
-    stats_aggs = []
-    for stats in stats_funcs:
-        stats_agg = apply(agg, kernel, func=_function_mapping[stats])
-        stats_aggs.append(stats_agg)
+    # Validate the kernel
+    kernel = custom_kernel(kernel)
 
-    stats = xr.concat(stats_aggs, pd.Index(stats_funcs, name='stats'))
-    return stats
+    mapper = ArrayTypeFunctionMapping(
+        numpy_func=_focal_stats_cpu,
+        cupy_func=_focal_stats_cupy,
+        dask_func=_focal_stats_cpu,
+        dask_cupy_func=lambda *args: not_implemented_func(
+            *args, messages='focal_stats() does not support dask with cupy backed DataArray.'),
+    )
+    result = mapper(agg)(agg, kernel, stats_funcs)
+    return result
 
 
 @ngjit

xrspatial/tests/test_focal.py

@@ -322,44 +322,43 @@ def test_apply_dask_numpy(data_apply):
 @pytest.fixture
 def data_focal_stats():
     data = np.arange(16).reshape(4, 4)
-    cellsize = (1, 1)
-    kernel = circle_kernel(*cellsize, 1.5)
+    kernel = custom_kernel(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 0]]))
     expected_result = np.asarray([
         # mean
-        [[1.66666667, 2., 3., 4.],
-         [4.25, 5., 6., 6.75],
-         [8.25, 9., 10., 10.75],
-         [11., 12., 13., 13.33333333]],
+        [[0, 1, 2, 3.],
+         [4, 2.5,  3.5,  4.5],
+         [8, 6.5,  7.5,  8.5],
+         [12, 10.5,  11.5,  12.5]],
         # max
-        [[4., 5., 6., 7.],
-         [8., 9., 10., 11.],
-         [12., 13., 14., 15.],
-         [13., 14., 15., 15.]],
+        [[0, 1, 2, 3.],
+         [4, 5, 6, 7.],
+         [8, 9, 10, 11.],
+         [12, 13, 14, 15.]],
         # min
-        [[0., 0., 1., 2.],
-         [0., 1., 2., 3.],
-         [4., 5., 6., 7.],
-         [8., 9., 10., 11.]],
+        [[0, 1, 2, 3.],
+         [4, 0, 1, 2.],
+         [8, 4, 5, 6.],
+         [12, 8, 9, 10.]],
         # range
-        [[4., 5., 5., 5.],
-         [8., 8., 8., 8.],
-         [8., 8., 8., 8.],
-         [5., 5., 5., 4.]],
+        [[0, 0, 0, 0.],
+         [0, 5, 5, 5.],
+         [0, 5, 5, 5.],
+         [0, 5, 5, 5.]],
         # std
-        [[1.69967317, 1.87082869, 1.87082869, 2.1602469],
-         [2.86138079, 2.60768096, 2.60768096, 2.86138079],
-         [2.86138079, 2.60768096, 2.60768096, 2.86138079],
-         [2.1602469, 1.87082869, 1.87082869, 1.69967317]],
+        [[0, 0, 0, 0.],
+         [0, 2.5,  2.5,  2.5],
+         [0, 2.5,  2.5,  2.5],
+         [0, 2.5,  2.5,  2.5]],
         # var
-        [[2.88888889, 3.5, 3.5, 4.66666667],
-         [8.1875, 6.8, 6.8, 8.1875],
-         [8.1875, 6.8, 6.8, 8.1875],
-         [4.66666667, 3.5, 3.5, 2.88888889]],
+        [[0, 0, 0, 0.],
+         [0, 6.25, 6.25, 6.25],
+         [0, 6.25, 6.25, 6.25],
+         [0, 6.25, 6.25, 6.25]],
         # sum
-        [[5., 8., 12., 12.],
-         [17., 25., 30., 27.],
-         [33., 45., 50., 43.],
-         [33., 48., 52., 40.]]
+        [[0, 1, 2, 3.],
+         [4, 5, 7, 9.],
+         [8, 13, 15, 17.],
+         [12, 21, 23, 25.]]
     ])
     return data, kernel, expected_result
 
@@ -383,6 +382,16 @@ def test_focal_stats_dask_numpy(data_focal_stats):
     )
 
 
+@cuda_and_cupy_available
+def test_focal_stats_gpu(data_focal_stats):
+    data, kernel, expected_result = data_focal_stats
+    cupy_agg = create_test_raster(data, backend='cupy')
+    cupy_focalstats = focal_stats(cupy_agg, kernel)
+    general_output_checks(
+        cupy_agg, cupy_focalstats, verify_attrs=False, expected_results=expected_result
+    )
+
+
 @pytest.fixture
 def data_hotspots():
     data = np.asarray([