Add GPU support for hillshade (#151)

* correct docs

* add gpu support for hillshade

* curvature: correct order of if else statement

* hillshade: add gpu cupy support

Author: thuydotm

xrspatial/curvature.py

@@ -120,14 +120,14 @@ def curvature(agg, name='curvature'):
     elif has_cuda() and isinstance(agg.data, cupy.ndarray):
         out = _run_cupy(agg.data, cellsize)
 
-    # dask + numpy case
-    elif isinstance(agg.data, da.Array):
-        out = _run_dask_numpy(agg.data, cellsize)
-
     # dask + cupy case
     elif has_cuda() and isinstance(agg.data, da.Array) and is_cupy_backed(agg):
         out = _run_dask_cupy(agg.data, cellsize)
 
+    # dask + numpy case
+    elif isinstance(agg.data, da.Array):
+        out = _run_dask_numpy(agg.data, cellsize)
+
     else:
         raise TypeError('Unsupported Array Type: {}'.format(type(agg.data)))
 

xrspatial/hillshade.py

@@ -1,13 +1,29 @@
+# std lib
 from functools import partial
+from math import sqrt
+from typing import Union
 
-import numpy as np
-
-from xarray import DataArray
+# 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
 
-def _hillshade(data, azimuth=225, angle_altitude=25):
+# local modules
+from xrspatial.utils import cuda_args
+from xrspatial.utils import has_cuda
+from xrspatial.utils import is_cupy_backed
+
+
+def _run_numpy(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))
@@ -18,7 +34,71 @@ def _hillshade(data, azimuth=225, angle_altitude=25):
     result = (shaded + 1) / 2
     result[(0, -1), :] = np.nan
     result[:, (0, -1)] = np.nan
-    return data
+    return result
+
+
+def _run_dask_numpy(data, azimuth, angle_altitude):
+    _func = partial(_run_numpy, azimuth=azimuth, angle_altitude=angle_altitude)
+    out = data.map_overlap(_func,
+                           depth=(1, 1),
+                           boundary=np.nan,
+                           meta=np.array(()))
+    return out
+
+
+@cuda.jit
+def _gpu_calc(x, y, out):
+    i, j = cuda.grid(2)
+    if i < out.shape[0] and j < out.shape[1]:
+        out[i, j] = sqrt(x[i, j] * x[i, j] + y[i, j] * y[i, j])
+
+
+@cuda.jit
+def _gpu_cos_part(cos_altituderad, cos_slope, cos_aspect, out):
+    i, j = cuda.grid(2)
+    if i < out.shape[0] and j < out.shape[1]:
+        out[i, j] = cos_altituderad * cos_slope[i, j] * cos_aspect[i, j]
+
+
+def _run_cupy(data, azimuth, angle_altitude):
+    x, y = np.gradient(data.get())
+    x = cupy.asarray(x, dtype=x.dtype)
+    y = cupy.asarray(y, dtype=y.dtype)
+
+    altituderad = angle_altitude * np.pi / 180.
+    sin_altituderad = np.sin(altituderad)
+    cos_altituderad = np.cos(altituderad)
+
+    griddim, blockdim = cuda_args(data.shape)
+    arctan_part = cupy.empty(data.shape, dtype='f4')
+    _gpu_calc[griddim, blockdim](x, y, arctan_part)
+
+    slope = np.pi / 2. - np.arctan(arctan_part)
+    sin_slope = np.sin(slope)
+    sin_part = sin_altituderad * sin_slope
+
+    azimuthrad = (360.0 - azimuth) * np.pi / 180.
+    aspect = (azimuthrad - np.pi / 2.) - np.arctan2(-x, y)
+    cos_aspect = np.cos(aspect)
+    cos_slope = np.cos(slope)
+
+    cos_part = cupy.empty(data.shape, dtype='f4')
+    _gpu_cos_part[griddim, blockdim](cos_altituderad, cos_slope,
+                                     cos_aspect, cos_part)
+    shaded = sin_part + cos_part
+    out = (shaded + 1) / 2
+
+    out[0, :] = cupy.nan
+    out[-1, :] = cupy.nan
+    out[:, 0] = cupy.nan
+    out[:, -1] = cupy.nan
+
+    return out
+
+
+def _run_dask_cupy(data, azimuth, angle_altitude):
+    msg = 'Upstream bug in dask prevents cupy backed arrays'
+    raise NotImplementedError(msg)
 
 
 def hillshade(agg, azimuth=225, angle_altitude=25, name='hillshade'):
@@ -32,14 +112,6 @@ def hillshade(agg, azimuth=225, angle_altitude=25, name='hillshade'):
     azimuth : int, optional (default: 315)
         The angle between the north vector and the perpendicular projection
         of the light source down onto the horizon specified in degrees.
-    cmap : list of colors or matplotlib.colors.Colormap, optional
-        The colormap to use. Can be either a list of colors (in any of the
-        formats described above), or a matplotlib colormap object.
-        Default is `["lightgray", "black"]`
-    alpha : int, optional
-        Value between 0 - 255 representing the alpha value of pixels which contain
-        data (i.e. non-nan values). Regardless of this value, `NaN` values are
-        set to fully transparent.
 
     Returns
     -------
@@ -51,14 +123,27 @@ def hillshade(agg, azimuth=225, angle_altitude=25, name='hillshade'):
      - http://geoexamples.blogspot.com/2014/03/shaded-relief-images-using-gdal-python.html
     """
 
-    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(()))
+    # numpy case
+    if isinstance(agg.data, np.ndarray):
+        out = _run_numpy(agg.data, azimuth, angle_altitude)
+
+    # cupy case
+    elif has_cuda() and isinstance(agg.data, cupy.ndarray):
+        out = _run_cupy(agg.data, azimuth, angle_altitude)
+
+    # dask + cupy case
+    elif has_cuda() and isinstance(agg.data, da.Array) and is_cupy_backed(agg):
+        out = _run_dask_cupy(agg.data, azimuth, angle_altitude)
+
+    # dask + numpy case
+    elif isinstance(agg.data, da.Array):
+        out = _run_dask_numpy(agg.data, azimuth, angle_altitude)
+
     else:
-        out = _hillshade(agg.data, azimuth, angle_altitude)
+        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)

xrspatial/tests/test_hillshade.py

@@ -1,9 +1,22 @@
+import pytest
 import xarray as xr
 import numpy as np
 
+import dask.array as da
+
 from xrspatial import hillshade
+from xrspatial.utils import doesnt_have_cuda
 
-import dask.array as da
+
+elevation = np.asarray(
+    [[1432.6542, 1432.4764, 1432.4764, 1432.1207, 1431.9429, np.nan],
+     [1432.6542, 1432.6542, 1432.4764, 1432.2986, 1432.1207, np.nan],
+     [1432.832, 1432.6542, 1432.4764, 1432.2986, 1432.1207, np.nan],
+     [1432.832, 1432.6542, 1432.4764, 1432.4764, 1432.1207, np.nan],
+     [1432.832, 1432.6542, 1432.6542, 1432.4764, 1432.2986, np.nan],
+     [1432.832, 1432.6542, 1432.6542, 1432.4764, 1432.2986, np.nan],
+     [1432.832, 1432.832, 1432.6542, 1432.4764, 1432.4764, np.nan]],
+    dtype=np.float32)
 
 
 def _do_sparse_array(data_array):
@@ -53,26 +66,30 @@ def test_hillshade():
 
 
 def test_numpy_equals_dask():
-
-    # input data
-    data = np.asarray(
-        [[1432.6542, 1432.4764, 1432.4764, 1432.1207, 1431.9429, np.nan],
-         [1432.6542, 1432.6542, 1432.4764, 1432.2986, 1432.1207, np.nan],
-         [1432.832, 1432.6542, 1432.4764, 1432.2986, 1432.1207, np.nan],
-         [1432.832, 1432.6542, 1432.4764, 1432.4764, 1432.1207, np.nan],
-         [1432.832, 1432.6542, 1432.6542, 1432.4764, 1432.2986, np.nan],
-         [1432.832, 1432.6542, 1432.6542, 1432.4764, 1432.2986, np.nan],
-         [1432.832, 1432.832, 1432.6542, 1432.4764, 1432.4764, np.nan]],
-        dtype=np.float32)
-
     attrs = {'res': (10.0, 10.0)}
 
-    small_numpy_based_data_array = xr.DataArray(data, attrs=attrs)
-    dask_data = da.from_array(data, chunks=(3, 3))
+    small_numpy_based_data_array = xr.DataArray(elevation, attrs=attrs)
+    dask_data = da.from_array(elevation, chunks=(3, 3))
     small_das_based_data_array = xr.DataArray(dask_data, attrs=attrs)
 
     numpy_result = hillshade(small_numpy_based_data_array, name='numpy')
     dask_result = hillshade(small_das_based_data_array, name='dask')
     dask_result.data = dask_result.data.compute()
 
     assert np.isclose(numpy_result, dask_result, equal_nan=True).all()
+
+
+@pytest.mark.skipif(doesnt_have_cuda(), reason="CUDA Device not Available")
+def test_hillshade_gpu_equals_cpu():
+
+    import cupy
+
+    small_da = xr.DataArray(elevation, attrs={'res': (10.0, 10.0)})
+    cpu = hillshade(small_da, name='numpy_result')
+
+    small_da_cupy = xr.DataArray(cupy.asarray(elevation), attrs={'res': (10.0, 10.0)})
+    gpu = hillshade(small_da_cupy, name='cupy_result')
+
+    assert isinstance(gpu.data, cupy.ndarray)
+
+    assert np.isclose(cpu, gpu, equal_nan=True).all()