Fixes #906: add dask+cupy backends for perlin, terrain, crosstab

Add missing dask+cupy backends for perlin() and generate_terrain()
using per-chunk GPU kernel launches via da.map_blocks with block_info
coordinate mapping. Add cupy and dask+cupy backends for crosstab()
by converting to numpy/dask+numpy and delegating. Fix _find_cats()
to handle pure cupy arrays without wrapping in @delayed.

Author: brendancol

xrspatial/perlin.py

@@ -192,6 +192,37 @@ def _perlin_cupy(data: cupy.ndarray,
     return data
 
 
+def _perlin_dask_cupy(data: da.Array,
+                      freq: tuple,
+                      seed: int) -> da.Array:
+    np.random.seed(seed)
+    p = cupy.asarray(np.random.permutation(2**20))
+    p = cupy.append(p, p)
+
+    height, width = data.shape
+
+    def _chunk_perlin(block, block_info=None):
+        info = block_info[0]
+        y_start, y_end = info['array-location'][0]
+        x_start, x_end = info['array-location'][1]
+        x0 = freq[0] * x_start / width
+        x1 = freq[0] * x_end / width
+        y0 = freq[1] * y_start / height
+        y1 = freq[1] * y_end / height
+
+        out = cupy.empty(block.shape, dtype=cupy.float32)
+        griddim, blockdim = cuda_args(block.shape)
+        _perlin_gpu[griddim, blockdim](p, x0, x1, y0, y1, 1.0, out)
+        return out
+
+    data = da.map_blocks(_chunk_perlin, data, dtype=cupy.float32,
+                         meta=cupy.array((), dtype=cupy.float32))
+
+    min_val, max_val = dask.compute(da.min(data), da.max(data))
+    data = (data - min_val) / (max_val - min_val)
+    return data
+
+
 def perlin(agg: xr.DataArray,
            freq: tuple = (1, 1),
            seed: int = 5,
@@ -246,9 +277,7 @@ def perlin(agg: xr.DataArray,
         numpy_func=_perlin_numpy,
         cupy_func=_perlin_cupy,
         dask_func=_perlin_dask_numpy,
-        dask_cupy_func=lambda *args: not_implemented_func(
-            *args, messages='perlin() does not support dask with cupy backed DataArray',  # noqa
-        )
+        dask_cupy_func=_perlin_dask_cupy
     )
     out = mapper(agg)(agg.data, freq, seed)
     result = xr.DataArray(out,

xrspatial/terrain.py

@@ -162,6 +162,55 @@ def _terrain_gpu(height_map, seed, x_range=(0, 1), y_range=(0, 1)):
     return height_map
 
 
+def _terrain_dask_cupy(data: da.Array,
+                       seed: int,
+                       x_range_scaled: tuple,
+                       y_range_scaled: tuple,
+                       zfactor: int) -> da.Array:
+    data = data * 0
+    height, width = data.shape
+
+    NOISE_LAYERS = ((1 / 2**i, (2**i, 2**i)) for i in range(16))
+
+    nrange = np.arange(2**20, dtype=np.int32)
+
+    for i, (m, (xfreq, yfreq)) in enumerate(NOISE_LAYERS):
+        np.random.seed(seed + i)
+        p = cupy.asarray(np.random.permutation(nrange))
+        p = cupy.append(p, p)
+
+        def _chunk_noise(block, block_info=None, _p=p, _m=m,
+                         _xr=x_range_scaled, _yr=y_range_scaled,
+                         _xf=xfreq, _yf=yfreq):
+            info = block_info[0]
+            y_start, y_end = info['array-location'][0]
+            x_start, x_end = info['array-location'][1]
+            x0 = _xr[0] + (_xr[1] - _xr[0]) * x_start / width
+            x1 = _xr[0] + (_xr[1] - _xr[0]) * x_end / width
+            y0 = _yr[0] + (_yr[1] - _yr[0]) * y_start / height
+            y1 = _yr[0] + (_yr[1] - _yr[0]) * y_end / height
+
+            out = cupy.empty(block.shape, dtype=cupy.float32)
+            griddim, blockdim = cuda_args(block.shape)
+            _perlin_gpu[griddim, blockdim](
+                _p, x0 * _xf, x1 * _xf, y0 * _yf, y1 * _yf, _m, out
+            )
+            return out
+
+        noise = da.map_blocks(_chunk_noise, data, dtype=cupy.float32,
+                              meta=cupy.array((), dtype=cupy.float32))
+        data = data + noise
+
+    data /= (1.00 + 0.50 + 0.25 + 0.13 + 0.06 + 0.03)
+    data = data ** 3
+
+    min_val, max_val = dask.compute(da.min(data), da.max(data))
+    data = (data - min_val) / (max_val - min_val)
+    data = da.where(data < 0.3, 0, data)
+    data *= zfactor
+    return data
+
+
 def _terrain_cupy(data: cupy.ndarray,
                   seed: int,
                   x_range_scaled: tuple,
@@ -263,9 +312,7 @@ def generate_terrain(agg: xr.DataArray,
         numpy_func=_terrain_numpy,
         cupy_func=_terrain_cupy,
         dask_func=_terrain_dask_numpy,
-        dask_cupy_func=lambda *args: not_implemented_func(
-            *args, messages='generate_terrain() does not support dask with cupy backed DataArray'  # noqa
-        )
+        dask_cupy_func=_terrain_dask_cupy
     )
     out = mapper(agg)(agg.data, seed, x_range_scaled, y_range_scaled, zfactor)
     canvas = ds.Canvas(

xrspatial/tests/test_dask_cupy_gaps.py

@@ -0,0 +1,158 @@
+"""Tests for dask+cupy backends: perlin, terrain, crosstab."""
+
+import numpy as np
+import xarray as xr
+
+from xrspatial import generate_terrain, perlin
+from xrspatial.tests.general_checks import cuda_and_cupy_available, dask_array_available
+from xrspatial.utils import has_cuda_and_cupy
+
+
+def _make_raster(shape=(50, 50), backend='numpy', chunks=(10, 10)):
+    data = np.zeros(shape, dtype=np.float32)
+    raster = xr.DataArray(data, dims=['y', 'x'])
+    if has_cuda_and_cupy() and 'cupy' in backend:
+        import cupy
+        raster.data = cupy.asarray(raster.data)
+    if 'dask' in backend:
+        import dask.array as da
+        raster.data = da.from_array(raster.data, chunks=chunks)
+    return raster
+
+
+# ---- perlin dask+cupy ----
+
+@cuda_and_cupy_available
+@dask_array_available
+def test_perlin_dask_cupy():
+    import cupy
+    import dask.array as da
+
+    cupy_raster = _make_raster(backend='cupy')
+    result_cupy = perlin(cupy_raster)
+
+    dask_cupy_raster = _make_raster(backend='dask+cupy')
+    result_dask_cupy = perlin(dask_cupy_raster)
+
+    assert isinstance(result_dask_cupy.data, da.Array)
+    assert result_dask_cupy.shape == cupy_raster.shape
+
+    computed = result_dask_cupy.data.compute()
+    assert isinstance(computed, cupy.ndarray)
+
+    vals = computed.get()
+    assert vals.min() >= 0.0
+    assert vals.max() <= 1.0
+
+    np.testing.assert_allclose(
+        result_cupy.data.get(), vals, rtol=1e-4, atol=1e-4,
+    )
+
+
+# ---- terrain dask+cupy ----
+
+@cuda_and_cupy_available
+@dask_array_available
+def test_terrain_dask_cupy():
+    import cupy
+    import dask.array as da
+
+    cupy_raster = _make_raster(backend='cupy')
+    terrain_cupy = generate_terrain(cupy_raster)
+
+    dask_cupy_raster = _make_raster(backend='dask+cupy')
+    terrain_dask_cupy = generate_terrain(dask_cupy_raster)
+
+    assert isinstance(terrain_dask_cupy.data, da.Array)
+    assert terrain_dask_cupy.shape == cupy_raster.shape
+
+    computed = terrain_dask_cupy.data.compute()
+    assert isinstance(computed, cupy.ndarray)
+
+    vals = computed.get()
+    assert vals.min() >= 0.0
+
+    np.testing.assert_allclose(
+        terrain_cupy.data.get(), vals, rtol=1e-4, atol=1e-4,
+    )
+
+
+# ---- crosstab cupy ----
+
+@cuda_and_cupy_available
+def test_crosstab_cupy():
+    import cupy
+    from xrspatial.zonal import crosstab
+
+    zones_np = np.array([
+        [1, 1, 2, 2],
+        [1, 1, 2, 2],
+        [3, 3, 4, 4],
+        [3, 3, 4, 4],
+    ], dtype=np.float64)
+    values_np = np.array([
+        [10, 10, 20, 20],
+        [10, 20, 20, 30],
+        [30, 30, 10, 10],
+        [30, 10, 20, 20],
+    ], dtype=np.float64)
+
+    zones_xr = xr.DataArray(zones_np, dims=['y', 'x'])
+    values_xr = xr.DataArray(values_np, dims=['y', 'x'])
+    df_numpy = crosstab(zones_xr, values_xr)
+
+    zones_cupy = xr.DataArray(cupy.asarray(zones_np), dims=['y', 'x'])
+    values_cupy = xr.DataArray(cupy.asarray(values_np), dims=['y', 'x'])
+    df_cupy = crosstab(zones_cupy, values_cupy)
+
+    # Both should be pandas DataFrames with identical content
+    assert list(df_numpy.columns) == list(df_cupy.columns)
+    np.testing.assert_array_equal(df_numpy.values, df_cupy.values)
+
+
+# ---- crosstab dask+cupy ----
+
+@cuda_and_cupy_available
+@dask_array_available
+def test_crosstab_dask_cupy():
+    import cupy
+    import dask.array as da
+    from xrspatial.zonal import crosstab
+
+    zones_np = np.array([
+        [1, 1, 2, 2],
+        [1, 1, 2, 2],
+        [3, 3, 4, 4],
+        [3, 3, 4, 4],
+    ], dtype=np.float64)
+    values_np = np.array([
+        [10, 10, 20, 20],
+        [10, 20, 20, 30],
+        [30, 30, 10, 10],
+        [30, 10, 20, 20],
+    ], dtype=np.float64)
+
+    zones_xr = xr.DataArray(zones_np, dims=['y', 'x'])
+    values_xr = xr.DataArray(values_np, dims=['y', 'x'])
+    df_numpy = crosstab(zones_xr, values_xr)
+
+    zones_gpu = cupy.asarray(zones_np)
+    values_gpu = cupy.asarray(values_np)
+    zones_dask = xr.DataArray(
+        da.from_array(zones_gpu, chunks=(2, 2)), dims=['y', 'x']
+    )
+    values_dask = xr.DataArray(
+        da.from_array(values_gpu, chunks=(2, 2)), dims=['y', 'x']
+    )
+    df_dask_cupy = crosstab(zones_dask, values_dask)
+
+    # dask case returns dask DataFrame
+    df_computed = df_dask_cupy.compute()
+    # sort both by zone for stable comparison
+    df_numpy_sorted = df_numpy.sort_values('zone').reset_index(drop=True)
+    df_computed_sorted = df_computed.sort_values('zone').reset_index(drop=True)
+
+    assert list(df_numpy_sorted.columns) == list(df_computed_sorted.columns)
+    np.testing.assert_array_equal(
+        df_numpy_sorted.values, df_computed_sorted.values,
+    )

xrspatial/zonal.py

@@ -721,7 +721,7 @@ def _find_cats(values, cat_ids, nodata_values):
     if cat_ids is None:
         cat_ids = unique_cats
     else:
-        if isinstance(values.data, np.ndarray):
+        if isinstance(values.data, np.ndarray) or is_cupy_array(values.data):
             # remove cats that do not exist in `values` raster
             cat_ids = [c for c in cat_ids if c in unique_cats]
         else:
@@ -960,6 +960,52 @@ def _crosstab_dask_numpy(
     return dd.from_delayed(crosstab_df)
 
 
+def _crosstab_cupy(
+    zones: np.ndarray,
+    values: np.ndarray,
+    zone_ids,
+    unique_cats,
+    cat_ids,
+    nodata_values,
+    agg: str,
+):
+    # unique_cats / cat_ids may be cupy arrays from _find_cats
+    if is_cupy_array(unique_cats):
+        unique_cats = cupy.asnumpy(unique_cats)
+    if is_cupy_array(cat_ids):
+        cat_ids = cupy.asnumpy(cat_ids)
+    return _crosstab_numpy(
+        cupy.asnumpy(zones), cupy.asnumpy(values),
+        zone_ids, unique_cats, cat_ids, nodata_values, agg,
+    )
+
+
+def _crosstab_dask_cupy(
+    zones,
+    values,
+    zone_ids,
+    unique_cats,
+    cat_ids,
+    nodata_values,
+    agg: str,
+):
+    zones_cpu = zones.map_blocks(
+        lambda x: x.get(), dtype=zones.dtype, meta=np.array(()),
+    )
+    values_cpu = values.map_blocks(
+        lambda x: x.get(), dtype=values.dtype, meta=np.array(()),
+    )
+    # unique_cats / cat_ids may be cupy arrays from _find_cats
+    if is_cupy_array(unique_cats):
+        unique_cats = cupy.asnumpy(unique_cats)
+    if is_cupy_array(cat_ids):
+        cat_ids = cupy.asnumpy(cat_ids)
+    return _crosstab_dask_numpy(
+        zones_cpu, values_cpu, zone_ids, unique_cats, cat_ids,
+        nodata_values, agg,
+    )
+
+
 def crosstab(
     zones: xr.DataArray,
     values: xr.DataArray,
@@ -1161,12 +1207,8 @@ def crosstab(
     mapper = ArrayTypeFunctionMapping(
         numpy_func=_crosstab_numpy,
         dask_func=_crosstab_dask_numpy,
-        cupy_func=lambda *args: not_implemented_func(
-            *args, messages='crosstab() does not support cupy backed DataArray'
-        ),
-        dask_cupy_func=lambda *args: not_implemented_func(
-            *args, messages='crosstab() does not support dask with cupy backed DataArray'  # noqa
-        ),
+        cupy_func=_crosstab_cupy,
+        dask_cupy_func=_crosstab_dask_cupy,
     )
     crosstab_df = mapper(values)(
         zones.data, values.data,