Support for certain Dask+Cupy (#815)

Author: relativityhd

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)

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

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,

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)