Add Drop-in Dask Support

xrspatial/aspect.py

@@ -3,6 +3,8 @@
 import numpy as np
 import numba as nb
 
+import dask.array as da
+
 from numba import cuda
 
 from xarray import DataArray
@@ -147,7 +149,12 @@ def aspect(agg, name='aspect', use_cuda=True, pad=True, use_cupy=True):
         _horn_aspect_cuda[griddim, blockdim](data, out)
         if pad:
             out = out[pad_rows:-pad_rows, pad_cols:-pad_cols]
-
+
+    elif isinstance(agg.data, da.Array):
+        out = agg.data.map_overlap(_horn_aspect,
+                                   depth=(1, 1),
+                                   boundary=np.nan,
+                                   meta=np.array(()))
     else:
         out = _horn_aspect(agg.data)
 

xrspatial/hillshade.py

@@ -1,9 +1,25 @@
-from __future__ import division, absolute_import
+from functools import partial
 
 import numpy as np
 
 from xarray import DataArray
 
+import dask.array as da
+
+
+def _hillshade(data, azimuth=225, angle_altitude=25):
+    azimuth = 360.0 - azimuth
+    x, y = np.gradient(data)
+    slope = np.pi/2. - np.arctan(np.sqrt(x*x + y*y))
+    aspect = np.arctan2(-x, y)
+    azimuthrad = azimuth*np.pi/180.
+    altituderad = angle_altitude*np.pi/180.
+    shaded = np.sin(altituderad) * np.sin(slope) + np.cos(altituderad) * np.cos(slope)*np.cos((azimuthrad - np.pi/2.) - aspect)
+    result = (shaded + 1) / 2
+    result[(0, -1), :] = np.nan
+    result[:, (0, -1)] = np.nan
+    return data
+
 
 def hillshade(agg, azimuth=225, angle_altitude=25, name='hillshade'):
     """Illuminates 2D DataArray from specific azimuth and altitude.
@@ -34,13 +50,15 @@ def hillshade(agg, azimuth=225, angle_altitude=25, name='hillshade'):
     Algorithm References:
      - http://geoexamples.blogspot.com/2014/03/shaded-relief-images-using-gdal-python.html
     """
-    azimuth = 360.0 - azimuth
-    x, y = np.gradient(agg.data)
-    slope = np.pi/2. - np.arctan(np.sqrt(x*x + y*y))
-    aspect = np.arctan2(-x, y)
-    azimuthrad = azimuth*np.pi/180.
-    altituderad = angle_altitude*np.pi/180.
-    shaded = np.sin(altituderad) * np.sin(slope) + np.cos(altituderad) * np.cos(slope)*np.cos((azimuthrad - np.pi/2.) - aspect)
-    data = (shaded + 1) / 2
-    return DataArray(data, name=name, dims=agg.dims,
+
+    if isinstance(agg.data, da.Array):
+        _func = partial(_hillshade, azimuth=azimuth, angle_altitude=angle_altitude)
+        out = agg.data.map_overlap(_func,
+                                   depth=(1, 1),
+                                   boundary=np.nan,
+                                   meta=np.array(()))
+    else:
+        out = _hillshade(agg.data, azimuth, angle_altitude)
+
+    return DataArray(out, name=name, dims=agg.dims,
                      coords=agg.coords, attrs=agg.attrs)

xrspatial/slope.py

@@ -1,18 +1,32 @@
+# std lib
+from functools import partial
 from math import atan
-import numpy as np
-import numba as nb
+from typing import Union
+
+# 3rd-party
+try:
+    import cupy
+except ImportError:
+    class cupy(object):
+        ndarray = False
 
+import dask.array as da
+
+import numba as nb
 from numba import cuda
 
-from xarray import DataArray
+import numpy as np
+import xarray as xr
 
-from xrspatial.utils import ngjit
-from xrspatial.utils import has_cuda
+# 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 _horn_slope(data, cellsize_x, cellsize_y):
+def _cpu(data, cellsize_x, cellsize_y):
     out = np.zeros_like(data)
     out[:] = np.nan
     rows, cols = data.shape
@@ -33,8 +47,29 @@ def _horn_slope(data, cellsize_x, cellsize_y):
     return out
 
 
+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:
+    _func = partial(_cpu,
+                    cellsize_x=cellsize_x,
+                    cellsize_y=cellsize_y)
+
+    out = data.map_overlap(_func,
+                           depth=(1, 1),
+                           boundary=np.nan,
+                           meta=np.array(()))
+    return out
+
+
 @cuda.jit(device=True)
-def _gpu_slope(arr, cellsize_x, cellsize_y):
+def _gpu(arr, cellsize_x, cellsize_y):
     a = arr[2, 0]
     b = arr[2, 1]
     c = arr[2, 2]
@@ -53,28 +88,73 @@ def _gpu_slope(arr, cellsize_x, cellsize_y):
 
 
 @cuda.jit
-def _horn_slope_cuda(arr, cellsize_x_arr, cellsize_y_arr, out):
+def _run_gpu(arr, cellsize_x_arr, cellsize_y_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_slope(arr[i-di:i+di+1, j-dj:j+dj+1],
-                               cellsize_x_arr,
-                               cellsize_y_arr)
+        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:
 
-def slope(agg, name='slope', use_cuda=True, pad=True, use_cupy=True):
+    cellsize_x_arr = cupy.array([float(cellsize_x)], dtype='f4')
+    cellsize_y_arr = cupy.array([float(cellsize_y)], dtype='f4')
+
+    pad_rows = 3 // 2
+    pad_cols = 3 // 2
+    pad_width = ((pad_rows, pad_rows),
+                (pad_cols, pad_cols))
+
+    slope_data = np.pad(data, pad_width=pad_width, mode="reflect")
+
+    griddim, blockdim = cuda_args(slope_data.shape)
+    slope_agg = cupy.empty(slope_data.shape, dtype='f4')
+    slope_agg[:] = cupy.nan
+
+    _run_gpu[griddim, blockdim](slope_data,
+                                cellsize_x_arr,
+                                cellsize_y_arr,
+                                slope_agg)
+    out = slope_agg[pad_rows:-pad_rows, pad_cols:-pad_cols]
+    return out
+
+
+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)
+
+    _func = partial(_run_cupy,
+                    cellsize_x=cellsize_x,
+                    cellsize_y=cellsize_y)
+
+    out = data.map_overlap(_func,
+                           depth=(1, 1),
+                           boundary=cupy.nan,
+                           dtype=cupy.float32,
+                           meta=cupy.array(()))
+    return out
+
+
+def slope(agg:xr.DataArray, name: str='slope') -> xr.DataArray:
     """Returns slope of input aggregate in degrees.
+
     Parameters
     ----------
-    agg : DataArray
+    agg : xr.DataArray
     name : str - name property of output xr.DataArray
-    use_cuda : bool - use CUDA device if available
 
     Returns
     -------
-    data: DataArray
+    data: xr.DataArray
 
     Notes:
     ------
@@ -85,61 +165,29 @@ def slope(agg, name='slope', use_cuda=True, pad=True, use_cupy=True):
       (Oxford University Press, New York), pp 406
     """
 
-    if not isinstance(agg, DataArray):
-        raise TypeError("agg must be instance of DataArray")
-
-    if not agg.attrs.get('res'):
-        raise ValueError('input xarray must have `res` attr.')
-
-    # get cellsize out from 'res' attribute
-    cellsize = agg.attrs.get('res')
-    if isinstance(cellsize, tuple) and len(cellsize) == 2 \
-            and isinstance(cellsize[0], (int, float)) \
-            and isinstance(cellsize[1], (int, float)):
-        cellsize_x, cellsize_y = cellsize
-    elif isinstance(cellsize, (int, float)):
-        cellsize_x = cellsize
-        cellsize_y = cellsize
-    else:
-        raise ValueError('`res` attr of input xarray must be a numeric'
-                         ' or a tuple of numeric values.')
+    cellsize_x, cellsize_y = get_dataarray_resolution(agg)
+
+    # numpy case
+    if isinstance(agg.data, np.ndarray):
+        out = _run_numpy(agg.data, cellsize_x, cellsize_y)
+
+    # cupy case
+    elif has_cuda() and isinstance(agg.data, cupy.ndarray):
+        out = _run_cupy(agg.data, cellsize_x, cellsize_y)
+
+    # 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)
 
-    if has_cuda() and use_cuda:
-        cellsize_x_arr = np.array([float(cellsize_x)], dtype='f4')
-        cellsize_y_arr = np.array([float(cellsize_y)], dtype='f4')
-
-        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
-
-        slope_data = np.pad(agg.data, pad_width=pad_width, mode="reflect")
-
-        griddim, blockdim = cuda_args(slope_data.shape)
-        slope_agg = np.empty(slope_data.shape, dtype='f4')
-        slope_agg[:] = np.nan
-
-        if use_cupy:
-            import cupy
-            slope_agg = cupy.asarray(slope_agg)
-
-        _horn_slope_cuda[griddim, blockdim](slope_data,
-                                            cellsize_x_arr,
-                                            cellsize_y_arr,
-                                            slope_agg)
-        if pad:
-            slope_agg = slope_agg[pad_rows:-pad_rows, pad_cols:-pad_cols]
+    # dask + numpy case
+    elif isinstance(agg.data, da.Array):
+        out = _run_dask_numpy(agg.data, cellsize_x, cellsize_y)
+
     else:
-        slope_agg = _horn_slope(agg.data, cellsize_x, cellsize_y)
+        raise TypeError('Unsupported Array Type: {}'.format(type(agg.data)))
 
-    return DataArray(slope_agg,
-                     name=name,
-                     coords=agg.coords,
-                     dims=agg.dims,
-                     attrs=agg.attrs)
+    return xr.DataArray(out,
+                        name=name,
+                        coords=agg.coords,
+                        dims=agg.dims,
+                        attrs=agg.attrs)

xrspatial/tests/test_aspect.py

@@ -1,6 +1,7 @@
-import xarray as xr
+import dask.array as da
 import numpy as np
 import pytest
+import xarray as xr
 
 from xrspatial import aspect
 from xrspatial.utils import doesnt_have_cuda
@@ -325,4 +326,27 @@ def test_aspect_against_qgis_gpu():
     assert np.isclose(xrspatial_vals, qgis_vals, equal_nan=True).all()
 
     #assert (np.isnan(xrspatial_vals) | (
-    #        (0. <= xrspatial_vals) & (xrspatial_vals <= 360.))).all()
+    #        (0. <= xrspatial_vals) & (xrspatial_vals <= 360.))).all()
+
+
+def test_numpy_equals_dask():
+
+    # input data
+    data = np.asarray([[np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
+                       [1584.8767, 1584.8767, 1585.0546, 1585.2324, 1585.2324, 1585.2324],
+                       [1585.0546, 1585.0546, 1585.2324, 1585.588, 1585.588, 1585.588],
+                       [1585.2324, 1585.4102, 1585.588, 1585.588, 1585.588, 1585.588],
+                       [1585.588, 1585.588, 1585.7659, 1585.7659, 1585.7659, 1585.7659],
+                       [1585.7659, 1585.9437, 1585.7659, 1585.7659, 1585.7659, 1585.7659],
+                       [1585.9437, 1585.9437, 1585.9437, 1585.7659, 1585.7659, 1585.7659]],
+                      dtype=np.float32)
+
+    small_numpy_based_data_array = xr.DataArray(data, attrs={'res': (10.0, 10.0)})
+    small_das_based_data_array = xr.DataArray(da.from_array(data, chunks=(2, 2)),
+                                              attrs={'res': (10.0, 10.0)})
+
+    numpy_result = aspect(small_numpy_based_data_array, name='numpy_slope', use_cuda=False)
+    dask_result = aspect(small_das_based_data_array, name='dask_slope', use_cuda=False)
+    dask_result.data = dask_result.data.compute()
+
+    assert np.isclose(numpy_result, dask_result, equal_nan=True).all()