refactoring aspect (#156)

* refactoring aspect

* fixed aspect gpu tests

* updated supported features in readme

Author: brendancol

README.md

@@ -126,7 +126,7 @@ In the GIS world, rasters are used for representing continuous phenomena (e.g. e
 
 | Name | NumPy xr.DataArray | Dask xr.DataArray | CuPy GPU xr.DataArray | Dask GPU xr.DataArray |
 |:----------:|:----------------------:|:--------------------:|:-------------------:|:------:|
-| [Aspect](xrspatial/aspect.py) | YES | NO | YES | NO |
+| [Aspect](xrspatial/aspect.py) | YES | YES | YES | NO |
 | [Curvature](xrspatial/curvature.py) | YES | NO | NO | NO |
 | [Hillshade](xrspatial/hillshade.py) | YES | YES  | NO | NO |
 | [Slope](xrspatial/slope.py) | YES  | YES  | YES | NO |

xrspatial/aspect.py

@@ -3,21 +3,33 @@
 import numpy as np
 import numba as nb
 
+from functools import partial 
+
 import dask.array as da
 
 from numba import cuda
 
+import xarray as xr
 from xarray import DataArray
 
 from xrspatial.utils import ngjit
 from xrspatial.utils import has_cuda
 from xrspatial.utils import cuda_args
+from xrspatial.utils import is_cupy_backed
+
+
+# 3rd-party
+try:
+    import cupy
+except ImportError:
+    class cupy(object):
+        ndarray = False
 
 RADIAN = 180 / np.pi
 
 
 @ngjit
-def _horn_aspect(data):
+def _cpu(data):
     out = np.zeros_like(data, dtype=np.float64)
     out[:] = np.nan
     rows, cols = data.shape
@@ -53,7 +65,7 @@ def _horn_aspect(data):
 
 
 @cuda.jit(device=True)
-def _gpu_aspect(arr):
+def _gpu(arr):
 
     a = arr[0, 0]
     b = arr[0, 1]
@@ -91,18 +103,66 @@ def _gpu_aspect(arr):
 
 
 @cuda.jit
-def _horn_aspect_cuda(arr, out):
-    minus_one = nb.float32(-1.)
+def _run_gpu(arr, out):
     i, j = cuda.grid(2)
     di = 1
     dj = 1
     if (i-di >= 1 and i+di < out.shape[0] - 1 and 
         j-dj >= 1 and j+dj < out.shape[1] - 1):
-        out[i, j] = _gpu_aspect(arr[i-di:i+di+1, j-dj:j+dj+1])
+        out[i, j] = _gpu(arr[i-di:i+di+1, j-dj:j+dj+1])
+
+
+def _run_cupy(data: cupy.ndarray) -> cupy.ndarray:
+
+    pad_rows = 3 // 2
+    pad_cols = 3 // 2
+    pad_width = ((pad_rows, pad_rows),
+                (pad_cols, pad_cols))
+
+    _data = np.pad(data, pad_width=pad_width, mode="reflect")
+
+    griddim, blockdim = cuda_args(_data.shape)
+    agg = cupy.empty(_data.shape, dtype='f4')
+    agg[:] = cupy.nan
+
+    _run_gpu[griddim, blockdim](_data, agg)
+    out = agg[pad_rows:-pad_rows, pad_cols:-pad_cols]
+    return out
+
+
+def _run_dask_cupy(data:da.Array) -> da.Array:
+
+    msg = 'Upstream bug in dask prevents cupy backed arrays'
+    raise NotImplementedError(msg)
+
+    # add any func args
+    # TODO: probably needs cellsize args
+    _func = partial(_run_cupy)
+
+    out = data.map_overlap(_func,
+                           depth=(1, 1),
+                           boundary=cupy.nan,
+                           dtype=cupy.float32,
+                           meta=cupy.array(()))
+    return out
+
+
+def _run_numpy(data:np.ndarray)-> np.ndarray:
+    out = _cpu(data)
+    return out
 
 
+def _run_dask_numpy(data:da.Array) -> da.Array:
+    _func = partial(_cpu)
 
-def aspect(agg, name='aspect', use_cuda=True, pad=True, use_cupy=True):
+    out = data.map_overlap(_func,
+                           depth=(1, 1),
+                           boundary=np.nan,
+                           meta=np.array(()))
+    return out
+
+
+def aspect(agg: xr.DataArray, name:str ='aspect'):
     """Returns downward slope direction in compass degrees (0 - 360) with 0 at 12 o'clock.
 
     Parameters
@@ -120,46 +180,27 @@ def aspect(agg, name='aspect', use_cuda=True, pad=True, use_cupy=True):
      - Burrough, P. A., and McDonell, R. A., 1998. Principles of Geographical Information Systems (Oxford University Press, New York), pp 406
     """
 
-    if not isinstance(agg, DataArray):
-        raise TypeError("agg must be instance of DataArray")
+    # numpy case
+    if isinstance(agg.data, np.ndarray):
+        out = _run_numpy(agg.data)
 
-    if has_cuda() and use_cuda:
-
-        if pad:
-            pad_rows = 3 // 2
-            pad_cols = 3 // 2
-            pad_width = ((pad_rows, pad_rows),
-                        (pad_cols, pad_cols))
-        else:
-            # If padding is not desired, set pads to 0
-            pad_rows = 0
-            pad_cols = 0
-            pad_width = 0
-
-        data = np.pad(agg.data, pad_width=pad_width, mode="reflect")
-
-        griddim, blockdim = cuda_args(data.shape)
-        out = np.empty(data.shape, dtype='f4')
-        out[:] = np.nan
-
-        if use_cupy:
-            import cupy
-            out = cupy.asarray(out)
-
-        _horn_aspect_cuda[griddim, blockdim](data, out)
-        if pad:
-            out = out[pad_rows:-pad_rows, pad_cols:-pad_cols]
+    # cupy case
+    elif has_cuda() and isinstance(agg.data, cupy.ndarray):
+        out = _run_cupy(agg.data)
 
+    # dask + cupy case
+    elif has_cuda() and isinstance(agg.data, da.Array) and is_cupy_backed(agg):
+        out = _run_dask_cupy(agg.data)
+    
+    # dask + numpy case
     elif isinstance(agg.data, da.Array):
-        out = agg.data.map_overlap(_horn_aspect,
-                                   depth=(1, 1),
-                                   boundary=np.nan,
-                                   meta=np.array(()))
+        out = _run_dask_numpy(agg.data)
+
     else:
-        out = _horn_aspect(agg.data)
+        raise TypeError('Unsupported Array Type: {}'.format(type(agg.data)))
 
-    return DataArray(out,
-                     name=name,
-                     dims=agg.dims,
-                     coords=agg.coords,
-                     attrs=agg.attrs)
+    return xr.DataArray(out,
+                        name=name,
+                        coords=agg.coords,
+                        dims=agg.dims,
+                        attrs=agg.attrs)