diff --git a/xrspatial/aspect.py b/xrspatial/aspect.py index 0409f429..aa1f46d5 100644 --- a/xrspatial/aspect.py +++ b/xrspatial/aspect.py @@ -121,6 +121,16 @@ def _run_dask_numpy(data: da.Array) -> da.Array: return out +def _run_dask_cupy(data: da.Array) -> da.Array: + data = data.astype(cupy.float32) + _func = partial(_run_cupy) + out = data.map_overlap(_func, + depth=(1, 1), + boundary=cupy.nan, + meta=cupy.array(())) + return out + + def aspect(agg: xr.DataArray, name: Optional[str] = 'aspect') -> xr.DataArray: """ @@ -249,8 +259,7 @@ def aspect(agg: xr.DataArray, numpy_func=_run_numpy, dask_func=_run_dask_numpy, cupy_func=_run_cupy, - dask_cupy_func=lambda *args: not_implemented_func( - *args, messages='aspect() does not support dask with cupy backed DataArray') # noqa + dask_cupy_func=_run_dask_cupy, ) out = mapper(agg)(agg.data) diff --git a/xrspatial/convolution.py b/xrspatial/convolution.py index 24bd3aa4..3bcc751f 100644 --- a/xrspatial/convolution.py +++ b/xrspatial/convolution.py @@ -373,15 +373,24 @@ def _convolve_2d_cupy(data, kernel): _convolve_2d_cuda[griddim, blockdim](data, kernel, cupy.asarray(out)) return out +def _convolve_2d_dask_cupy(data, kernel): + 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, + meta=cupy.array(())) + return out + def convolve_2d(data, kernel): mapper = ArrayTypeFunctionMapping( numpy_func=_convolve_2d_numpy, cupy_func=_convolve_2d_cupy, dask_func=_convolve_2d_dask_numpy, - dask_cupy_func=lambda *args: not_implemented_func( - *args, messages='convolution_2d() does not support dask with cupy backed xr.DataArray' # noqa - ) + dask_cupy_func=_convolve_2d_dask_cupy ) out = mapper(xr.DataArray(data))(data, kernel) return out diff --git a/xrspatial/curvature.py b/xrspatial/curvature.py index 9acc9f06..fcec5650 100644 --- a/xrspatial/curvature.py +++ b/xrspatial/curvature.py @@ -84,6 +84,16 @@ def _run_cupy(data: cupy.ndarray, return out +def _run_dask_cupy(data: da.Array, + cellsize: Union[int, float]) -> da.Array: + data = data.astype(cupy.float32) + _func = partial(_cpu, cellsize=cellsize) + out = data.map_overlap(_func, + depth=(1, 1), + boundary=cupy.nan, + meta=cupy.array(())) + return out + def curvature(agg: xr.DataArray, name: Optional[str] = 'curvature') -> xr.DataArray: @@ -209,8 +219,7 @@ def curvature(agg: xr.DataArray, numpy_func=_run_numpy, cupy_func=_run_cupy, dask_func=_run_dask_numpy, - dask_cupy_func=lambda *args: not_implemented_func( - *args, messages='curvature() does not support dask with cupy backed DataArray.'), # noqa + dask_cupy_func=_run_dask_cupy ) out = mapper(agg)(agg.data, cellsize) return xr.DataArray(out, diff --git a/xrspatial/slope.py b/xrspatial/slope.py index e337315a..3abef6fe 100644 --- a/xrspatial/slope.py +++ b/xrspatial/slope.py @@ -64,6 +64,21 @@ def _run_dask_numpy(data: da.Array, meta=np.array(())) return out +def _run_dask_cupy(data: da.Array, + cellsize_x: Union[int, float], + cellsize_y: Union[int, float]) -> da.Array: + data = data.astype(cupy.float32) + _func = partial(_run_cupy, + cellsize_x=cellsize_x, + cellsize_y=cellsize_y) + + out = data.map_overlap(_func, + depth=(1, 1), + boundary=cupy.nan, + meta=cupy.array(())) + return out + + @cuda.jit(device=True) def _gpu(arr, cellsize_x, cellsize_y): @@ -163,9 +178,7 @@ def slope(agg: xr.DataArray, numpy_func=_run_numpy, cupy_func=_run_cupy, dask_func=_run_dask_numpy, - dask_cupy_func=lambda *args: not_implemented_func( - *args, messages='slope() does not support dask with cupy backed DataArray' # noqa - ), + dask_cupy_func=_run_dask_cupy, ) out = mapper(agg)(agg.data, cellsize_x, cellsize_y)