Add GPU support for curvature (#150)

* add GPU support for curvature

* correct position check condition

* border edge effect, set to nan

* curvature: update test, cpu against gpu version

* refactor

* remove use_cuda and use_gpu params

* support dask numpy

* typo fixes

Author: thuydotm

xrspatial/curvature.py

@@ -1,74 +1,138 @@
+# std lib
+from functools import partial
+from typing import Union
+
+# 3rd-party
+try:
+    import cupy
+except ImportError:
+    class cupy(object):
+        ndarray = False
+
+import dask.array as da
+
+from numba import cuda
+
 import numpy as np
 import xarray as xr
-from numba import stencil, vectorize
 
+# local modules
+from xrspatial.utils import cuda_args
+from xrspatial.utils import get_dataarray_resolution
+from xrspatial.utils import has_cuda
+from xrspatial.utils import ngjit
+from xrspatial.utils import is_cupy_backed
+
+
+@ngjit
+def _cpu(data, cellsize):
+    out = np.empty(data.shape, np.float64)
+    out[:, :] = np.nan
+    rows, cols = data.shape
+    for y in range(1, rows - 1):
+        for x in range(1, cols - 1):
+            d = (data[y + 1, x] + data[y - 1, x]) / 2 - data[y, x]
+            e = (data[y, x + 1] + data[y, x - 1]) / 2 - data[y, x]
+            out[y, x] = -2 * (d + e) * 100 / (cellsize * cellsize)
+    return out
+
+
+def _run_numpy(data: np.ndarray,
+               cellsize: Union[int, float]) -> np.ndarray:
+    # TODO: handle border edge effect
+    out = _cpu(data, cellsize)
+    return out
+
+
+def _run_dask_numpy(data: da.Array,
+                    cellsize: Union[int, float]) -> da.Array:
+    _func = partial(_cpu,
+                    cellsize=cellsize)
+
+    out = data.map_overlap(_func,
+                           depth=(1, 1),
+                           boundary=np.nan,
+                           meta=np.array(()))
+    return out
+
+
+@cuda.jit(device=True)
+def _gpu(arr, cellsize):
+    d = (arr[1, 0] + arr[1, 2]) / 2 - arr[1, 1]
+    e = (arr[0, 1] + arr[2, 1]) / 2 - arr[1, 1]
+    curv = -2 * (d + e) * 100 / (cellsize[0] * cellsize[0])
+    return curv
 
-@stencil
-def kernel_D(arr):
-    return (arr[0, -1] + arr[0, 1]) / 2 - arr[0, 0]
 
+@cuda.jit
+def _run_gpu(arr, cellsize, out):
+    i, j = cuda.grid(2)
+    di = 1
+    dj = 1
+    if (i - di >= 0 and i + di <= out.shape[0] - 1 and
+            j - dj >= 0 and j + dj <= out.shape[1] - 1):
+        out[i, j] = _gpu(arr[i - di:i + di + 1, j - dj:j + dj + 1], cellsize)
 
-@stencil
-def kernel_E(arr):
-    return (arr[-1, 0] + arr[1, 0]) / 2 - arr[0, 0]
 
+def _run_cupy(data: cupy.ndarray,
+              cellsize: Union[int, float]) -> cupy.ndarray:
 
-@vectorize(["float64(float64, float64)"], nopython=True, target="parallel")
-def _curvature(matrix_D, matrix_E):
-    curv = -2 * (matrix_D + matrix_E) * 100
-    return curv
+    cellsize_arr = cupy.array([float(cellsize)], dtype='f4')
+
+    # TODO: add padding
+    griddim, blockdim = cuda_args(data.shape)
+    out = cupy.empty(data.shape, dtype='f4')
+    out[:] = cupy.nan
 
+    _run_gpu[griddim, blockdim](data, cellsize_arr, out)
 
-def curvature(raster):
-    """Compute the curvature (second derivatives) of a raster surface.
+    return out
+
+
+def _run_dask_cupy(data: da.Array,
+                   cellsize: Union[int, float]) -> da.Array:
+    msg = 'Upstream bug in dask prevents cupy backed arrays'
+    raise NotImplementedError(msg)
+
+
+def curvature(agg, name='curvature'):
+    """Compute the curvature (second derivatives) of a agg surface.
 
     Parameters
     ----------
-    raster: xarray.DataArray
-        2D input raster image with shape=(height, width)
+    agg: xarray.xr.DataArray
+        2D input agg image with shape=(height, width)
 
     Returns
     -------
-    curvature: xarray.DataArray
+    curvature: xarray.xr.DataArray
         Curvature image with shape=(height, width)
     """
 
-    if not isinstance(raster, xr.DataArray):
-        raise TypeError("`raster` must be instance of DataArray")
-
-    if raster.ndim != 2:
-        raise ValueError("`raster` must be 2D")
-
-    if not (issubclass(raster.values.dtype.type, np.integer) or
-            issubclass(raster.values.dtype.type, np.floating)):
-        raise ValueError(
-            "`raster` must be an array of integers or float")
-
-    raster_values = raster.values
+    cellsize_x, cellsize_y = get_dataarray_resolution(agg)
+    cellsize = (cellsize_x + cellsize_y) / 2
 
-    cell_size_x = 1
-    cell_size_y = 1
+    # numpy case
+    if isinstance(agg.data, np.ndarray):
+        out = _run_numpy(agg.data, cellsize)
 
-    # calculate cell size from input `raster`
-    for dim in raster.dims:
-        if (dim.lower().count('x')) > 0 or (dim.lower().count('lon')) > 0:
-            # dimension of x-coordinates
-            if len(raster[dim]) > 1:
-                cell_size_x = raster[dim].values[1] - raster[dim].values[0]
-        elif (dim.lower().count('y')) > 0 or (dim.lower().count('lat')) > 0:
-            # dimension of y-coordinates
-            if len(raster[dim]) > 1:
-                cell_size_y = raster[dim].values[1] - raster[dim].values[0]
+    # cupy case
+    elif has_cuda() and isinstance(agg.data, cupy.ndarray):
+        out = _run_cupy(agg.data, cellsize)
 
-    cell_size = (cell_size_x + cell_size_y) / 2
+    # dask + numpy case
+    elif isinstance(agg.data, da.Array):
+        out = _run_dask_numpy(agg.data, cellsize)
 
-    matrix_D = kernel_D(raster_values)
-    matrix_E = kernel_E(raster_values)
+    # dask + cupy case
+    elif has_cuda() and isinstance(agg.data, da.Array) and is_cupy_backed(agg):
+        out = _run_dask_cupy(agg.data, cellsize)
 
-    curvature_values = _curvature(matrix_D, matrix_E) / (cell_size * cell_size)
-    result = xr.DataArray(curvature_values,
-                          coords=raster.coords,
-                          dims=raster.dims,
-                          attrs=raster.attrs)
+    else:
+        raise TypeError('Unsupported Array Type: {}'.format(type(agg.data)))
 
-    return result
+    return xr.DataArray(out,
+                        name=name,
+                        coords=agg.coords,
+                        dims=agg.dims,
+                        attrs=agg.attrs)

xrspatial/slope.py

@@ -25,38 +25,39 @@ class cupy(object):
 from xrspatial.utils import ngjit
 from xrspatial.utils import is_cupy_backed
 
+
 @ngjit
 def _cpu(data, cellsize_x, cellsize_y):
     out = np.zeros_like(data)
     out[:] = np.nan
     rows, cols = data.shape
-    for y in range(1, rows-1):
-        for x in range(1, cols-1):
-            a = data[y+1, x-1]
-            b = data[y+1, x]
-            c = data[y+1, x+1]
-            d = data[y, x-1]
-            f = data[y, x+1]
-            g = data[y-1, x-1]
-            h = data[y-1, x]
-            i = data[y-1, x+1]
+    for y in range(1, rows - 1):
+        for x in range(1, cols - 1):
+            a = data[y + 1, x - 1]
+            b = data[y + 1, x]
+            c = data[y + 1, x + 1]
+            d = data[y, x - 1]
+            f = data[y, x + 1]
+            g = data[y - 1, x - 1]
+            h = data[y - 1, x]
+            i = data[y - 1, x + 1]
             dz_dx = ((c + 2 * f + i) - (a + 2 * d + g)) / (8 * cellsize_x)
             dz_dy = ((g + 2 * h + i) - (a + 2 * b + c)) / (8 * cellsize_y)
             p = (dz_dx * dz_dx + dz_dy * dz_dy) ** .5
             out[y, x] = np.arctan(p) * 57.29578
     return out
 
 
-def _run_numpy(data:np.ndarray,
-               cellsize_x:Union[int, float],
-               cellsize_y:Union[int, float]) -> np.ndarray:
+def _run_numpy(data: np.ndarray,
+               cellsize_x: Union[int, float],
+               cellsize_y: Union[int, float]) -> np.ndarray:
     out = _cpu(data, cellsize_x, cellsize_y)
     return out
 
 
-def _run_dask_numpy(data:da.Array,
-                   cellsize_x:Union[int, float],
-                   cellsize_y:Union[int, float]) -> da.Array:
+def _run_dask_numpy(data: da.Array,
+                    cellsize_x: Union[int, float],
+                    cellsize_y: Union[int, float]) -> da.Array:
     _func = partial(_cpu,
                     cellsize_x=cellsize_x,
                     cellsize_y=cellsize_y)
@@ -90,17 +91,16 @@ def _run_gpu(arr, cellsize_x_arr, cellsize_y_arr, out):
     i, j = cuda.grid(2)
     di = 1
     dj = 1
-    if (i-di >= 0 and i+di < out.shape[0] and
-        j-dj >= 0 and j+dj < out.shape[1]):
-        out[i, j] = _gpu(arr[i-di:i+di+1, j-dj:j+dj+1],
+    if (i - di >= 0 and i + di < out.shape[0] and
+            j - dj >= 0 and j + dj < out.shape[1]):
+        out[i, j] = _gpu(arr[i - di:i + di + 1, j - dj:j + dj + 1],
                          cellsize_x_arr,
                          cellsize_y_arr)
 
 
 def _run_cupy(data: cupy.ndarray,
               cellsize_x: Union[int, float],
               cellsize_y: Union[int, float]) -> cupy.ndarray:
-
     cellsize_x_arr = cupy.array([float(cellsize_x)], dtype='f4')
     cellsize_y_arr = cupy.array([float(cellsize_y)], dtype='f4')
 
@@ -115,10 +115,9 @@ def _run_cupy(data: cupy.ndarray,
     return out
 
 
-def _run_dask_cupy(data:da.Array,
-                   cellsize_x:Union[int, float],
-                   cellsize_y:Union[int, float]) -> da.Array:
-
+def _run_dask_cupy(data: da.Array,
+                   cellsize_x: Union[int, float],
+                   cellsize_y: Union[int, float]) -> da.Array:
     msg = 'Upstream bug in dask prevents cupy backed arrays'
     raise NotImplementedError(msg)
 
@@ -134,7 +133,7 @@ def _run_dask_cupy(data:da.Array,
     return out
 
 
-def slope(agg:xr.DataArray, name: str='slope') -> xr.DataArray:
+def slope(agg: xr.DataArray, name: str = 'slope') -> xr.DataArray:
     """Returns slope of input aggregate in degrees.
 
     Parameters
@@ -168,7 +167,7 @@ def slope(agg:xr.DataArray, name: str='slope') -> xr.DataArray:
     # dask + cupy case
     elif has_cuda() and isinstance(agg.data, da.Array) and is_cupy_backed(agg):
         out = _run_dask_cupy(agg.data, cellsize_x, cellsize_y)
-    
+
     # dask + numpy case
     elif isinstance(agg.data, da.Array):
         out = _run_dask_numpy(agg.data, cellsize_x, cellsize_y)
@@ -180,4 +179,4 @@ def slope(agg:xr.DataArray, name: str='slope') -> xr.DataArray:
                         name=name,
                         coords=agg.coords,
                         dims=agg.dims,
-                        attrs=agg.attrs)
+                        attrs=agg.attrs)