Skip to content

Commit 52f9fb5

Browse files
brendancolbrendancol
authored
Fixes #906: fill remaining dask+cupy gaps (terrain, perlin, crosstab) (#913)
* 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. * Update README feature matrix for new and existing GPU backends - Terrain Generation, Perlin Noise: mark dask+cupy as ✅ (native) - Zonal Cross Tabulate: mark cupy and dask+cupy as 🔄 (CPU fallback) - Zonal Statistics: mark cupy as ✅ (native), dask+cupy as 🔄 (fallback) - Apply: mark cupy and dask+cupy as 🔄 (CPU fallback) * Add dask/cupy/dask+cupy support for trim() and crop() Both functions only need numpy for the boundary scan; the output is a slice of the original DataArray, so the backend is preserved. Convert to numpy for the scan, then slice the original raster. Update README feature matrix for trim and crop. * Add native GPU kernel for zonal.apply() with automatic CPU fallback JIT-compile the user's scalar function as a CUDA device function and run it inside a CUDA kernel, avoiding the GPU→CPU→GPU round-trip. Non-compilable functions (e.g. using dict, str) automatically fall back to the existing CPU path via exception handling.
1 parent f5ce56a commit 52f9fb5

5 files changed

Lines changed: 664 additions & 35 deletions

File tree

‎README.md‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -224,10 +224,10 @@ In the GIS world, rasters are used for representing continuous phenomena (e.g. e
224224
| [Curvature](xrspatial/curvature.py) | Measures rate of slope change (concavity/convexity) at each cell | ✅️ |✅️ |✅️ | ✅️ |
225225
| [Hillshade](xrspatial/hillshade.py) | Simulates terrain illumination from a given sun angle and azimuth | ✅️ | ✅️ | ✅️ | ✅️ |
226226
| [Slope](xrspatial/slope.py) | Computes terrain gradient steepness at each cell in degrees | ✅️ | ✅️ | ✅️ | ✅️ |
227-
| [Terrain Generation](xrspatial/terrain.py) | Generates synthetic terrain elevation using fractal noise | ✅️ | ✅️ | ✅️ | |
227+
| [Terrain Generation](xrspatial/terrain.py) | Generates synthetic terrain elevation using fractal noise | ✅️ | ✅️ | ✅️ | ✅️ |
228228
| [Viewshed](xrspatial/viewshed.py) | Determines visible cells from a given observer point on terrain | ✅️ | ✅️ | ✅️ | ✅️ |
229229
| [Min Observable Height](xrspatial/experimental/min_observable_height.py) | Finds the minimum observer height needed to see each cell *(experimental)* | ✅️ | | | |
230-
| [Perlin Noise](xrspatial/perlin.py) | Generates smooth continuous random noise for procedural textures | ✅️ | ✅️ | ✅️ | |
230+
| [Perlin Noise](xrspatial/perlin.py) | Generates smooth continuous random noise for procedural textures | ✅️ | ✅️ | ✅️ | ✅️ |
231231
| [Bump Mapping](xrspatial/bump.py) | Adds randomized bump features to simulate natural terrain variation | ✅️ | | | |
232232

233233
-----------
@@ -236,12 +236,12 @@ In the GIS world, rasters are used for representing continuous phenomena (e.g. e
236236

237237
| Name | Description | NumPy xr.DataArray | Dask xr.DataArray | CuPy GPU xr.DataArray | Dask GPU xr.DataArray |
238238
|:----------:|:------------|:----------------------:|:--------------------:|:-------------------:|:------:|
239-
| [Apply](xrspatial/zonal.py) | Applies a custom function to each zone in a classified raster | ✅️ | ✅️ | | |
240-
| [Crop](xrspatial/zonal.py) | Extracts the bounding rectangle of a specific zone | ✅️ | | | |
239+
| [Apply](xrspatial/zonal.py) | Applies a custom function to each zone in a classified raster | ✅️ | ✅️ | ✅️ | ✅️ |
240+
| [Crop](xrspatial/zonal.py) | Extracts the bounding rectangle of a specific zone | ✅️ | ✅️ | ✅️ | ✅️ |
241241
| [Regions](xrspatial/zonal.py) | Identifies connected regions of non-zero cells | ✅️ | ✅️ | ✅️ | ✅️ |
242-
| [Trim](xrspatial/zonal.py) | Removes nodata border rows and columns from a raster | ✅️ | | | |
243-
| [Zonal Statistics](xrspatial/zonal.py) | Computes summary statistics for a value raster within each zone | ✅️ | ✅️| | |
244-
| [Zonal Cross Tabulate](xrspatial/zonal.py) | Cross-tabulates agreement between two categorical rasters | ✅️ | ✅️| | |
242+
| [Trim](xrspatial/zonal.py) | Removes nodata border rows and columns from a raster | ✅️ | ✅️ | ✅️ | ✅️ |
243+
| [Zonal Statistics](xrspatial/zonal.py) | Computes summary statistics for a value raster within each zone | ✅️ | ✅️| ✅️ | 🔄 |
244+
| [Zonal Cross Tabulate](xrspatial/zonal.py) | Cross-tabulates agreement between two categorical rasters | ✅️ | ✅️| 🔄 | 🔄 |
245245

246246
#### Usage
247247

‎xrspatial/perlin.py‎

Lines changed: 32 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -192,6 +192,37 @@ def _perlin_cupy(data: cupy.ndarray,
192192
return data
193193

194194

195+
def _perlin_dask_cupy(data: da.Array,
196+
freq: tuple,
197+
seed: int) -> da.Array:
198+
np.random.seed(seed)
199+
p = cupy.asarray(np.random.permutation(2**20))
200+
p = cupy.append(p, p)
201+
202+
height, width = data.shape
203+
204+
def _chunk_perlin(block, block_info=None):
205+
info = block_info[0]
206+
y_start, y_end = info['array-location'][0]
207+
x_start, x_end = info['array-location'][1]
208+
x0 = freq[0] * x_start / width
209+
x1 = freq[0] * x_end / width
210+
y0 = freq[1] * y_start / height
211+
y1 = freq[1] * y_end / height
212+
213+
out = cupy.empty(block.shape, dtype=cupy.float32)
214+
griddim, blockdim = cuda_args(block.shape)
215+
_perlin_gpu[griddim, blockdim](p, x0, x1, y0, y1, 1.0, out)
216+
return out
217+
218+
data = da.map_blocks(_chunk_perlin, data, dtype=cupy.float32,
219+
meta=cupy.array((), dtype=cupy.float32))
220+
221+
min_val, max_val = dask.compute(da.min(data), da.max(data))
222+
data = (data - min_val) / (max_val - min_val)
223+
return data
224+
225+
195226
def perlin(agg: xr.DataArray,
196227
freq: tuple = (1, 1),
197228
seed: int = 5,
@@ -246,9 +277,7 @@ def perlin(agg: xr.DataArray,
246277
numpy_func=_perlin_numpy,
247278
cupy_func=_perlin_cupy,
248279
dask_func=_perlin_dask_numpy,
249-
dask_cupy_func=lambda *args: not_implemented_func(
250-
*args, messages='perlin() does not support dask with cupy backed DataArray', # noqa
251-
)
280+
dask_cupy_func=_perlin_dask_cupy
252281
)
253282
out = mapper(agg)(agg.data, freq, seed)
254283
result = xr.DataArray(out,

‎xrspatial/terrain.py‎

Lines changed: 50 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,55 @@ def _terrain_gpu(height_map, seed, x_range=(0, 1), y_range=(0, 1)):
162162
return height_map
163163

164164

165+
def _terrain_dask_cupy(data: da.Array,
166+
seed: int,
167+
x_range_scaled: tuple,
168+
y_range_scaled: tuple,
169+
zfactor: int) -> da.Array:
170+
data = data * 0
171+
height, width = data.shape
172+
173+
NOISE_LAYERS = ((1 / 2**i, (2**i, 2**i)) for i in range(16))
174+
175+
nrange = np.arange(2**20, dtype=np.int32)
176+
177+
for i, (m, (xfreq, yfreq)) in enumerate(NOISE_LAYERS):
178+
np.random.seed(seed + i)
179+
p = cupy.asarray(np.random.permutation(nrange))
180+
p = cupy.append(p, p)
181+
182+
def _chunk_noise(block, block_info=None, _p=p, _m=m,
183+
_xr=x_range_scaled, _yr=y_range_scaled,
184+
_xf=xfreq, _yf=yfreq):
185+
info = block_info[0]
186+
y_start, y_end = info['array-location'][0]
187+
x_start, x_end = info['array-location'][1]
188+
x0 = _xr[0] + (_xr[1] - _xr[0]) * x_start / width
189+
x1 = _xr[0] + (_xr[1] - _xr[0]) * x_end / width
190+
y0 = _yr[0] + (_yr[1] - _yr[0]) * y_start / height
191+
y1 = _yr[0] + (_yr[1] - _yr[0]) * y_end / height
192+
193+
out = cupy.empty(block.shape, dtype=cupy.float32)
194+
griddim, blockdim = cuda_args(block.shape)
195+
_perlin_gpu[griddim, blockdim](
196+
_p, x0 * _xf, x1 * _xf, y0 * _yf, y1 * _yf, _m, out
197+
)
198+
return out
199+
200+
noise = da.map_blocks(_chunk_noise, data, dtype=cupy.float32,
201+
meta=cupy.array((), dtype=cupy.float32))
202+
data = data + noise
203+
204+
data /= (1.00 + 0.50 + 0.25 + 0.13 + 0.06 + 0.03)
205+
data = data ** 3
206+
207+
min_val, max_val = dask.compute(da.min(data), da.max(data))
208+
data = (data - min_val) / (max_val - min_val)
209+
data = da.where(data < 0.3, 0, data)
210+
data *= zfactor
211+
return data
212+
213+
165214
def _terrain_cupy(data: cupy.ndarray,
166215
seed: int,
167216
x_range_scaled: tuple,
@@ -263,9 +312,7 @@ def generate_terrain(agg: xr.DataArray,
263312
numpy_func=_terrain_numpy,
264313
cupy_func=_terrain_cupy,
265314
dask_func=_terrain_dask_numpy,
266-
dask_cupy_func=lambda *args: not_implemented_func(
267-
*args, messages='generate_terrain() does not support dask with cupy backed DataArray' # noqa
268-
)
315+
dask_cupy_func=_terrain_dask_cupy
269316
)
270317
out = mapper(agg)(agg.data, seed, x_range_scaled, y_range_scaled, zfactor)
271318
canvas = ds.Canvas(

0 commit comments

Comments
 (0)