Skip to content

focal_stats(): gpu case #709

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
May 12, 2022
Merged

focal_stats(): gpu case #709

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4c0c6e1
focal stats gpu case
thuydotm May 12, 2022
326b441
flake8
thuydotm May 12, 2022
49d4778
add tests
thuydotm May 12, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
205 changes: 189 additions & 16 deletions xrspatial/focal.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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=[
Expand Down Expand Up @@ -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
Expand Down
69 changes: 39 additions & 30 deletions xrspatial/tests/test_focal.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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([
Expand Down