From 2c411427afbc2c8301b083e3c6d0e4a1e056fa59 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Mon, 29 Jun 2026 09:44:35 -0700 Subject: [PATCH 1/2] generate_terrain: map dask backends over empty skeleton (#3574) The dask backends regenerate every cell from coordinates and never read the template's values, but they mapped the per-chunk generator over the template array itself. For a from_template() grid that meant computing a da.full(shape, nan) and immediately discarding it. Map over a da.empty_like skeleton instead so the template drops out of the graph. Also removes the leftover data * 0 in the dask+cupy path, which built a NaN array the per-chunk cupy.zeros discarded. --- xrspatial/terrain.py | 21 ++++++++++--- xrspatial/tests/test_terrain.py | 55 +++++++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+), 5 deletions(-) diff --git a/xrspatial/terrain.py b/xrspatial/terrain.py index ebd680378..b6f4a0ce2 100644 --- a/xrspatial/terrain.py +++ b/xrspatial/terrain.py @@ -378,6 +378,13 @@ def _terrain_dask_numpy(data, seed, x_range_scaled, y_range_scaled, zfactor, """ height, width = data.shape + # The chunk functions below regenerate every cell from coordinates and + # never read the incoming block values. Map over an uninitialised + # skeleton with the same shape/chunking so the (NaN-filled) template the + # caller passed in — e.g. from_template()'s da.full — drops out of the + # graph instead of being memset and then thrown away. + skeleton = da.empty_like(data, dtype=np.float32) + # --- worley pre-pass: get global min/max to avoid per-chunk seams --- w_norm_range = None if worley_blend > 0: @@ -417,7 +424,7 @@ def _chunk_worley_raw(block, block_info=None): return _worley_numpy_xy(w_p, x_arr, y_arr) raw_worley = da.map_blocks( - _chunk_worley_raw, data, dtype=np.float32, + _chunk_worley_raw, skeleton, dtype=np.float32, meta=np.array((), dtype=np.float32), ) (raw_worley,) = dask.persist(raw_worley) @@ -452,7 +459,7 @@ def _chunk_terrain(block, block_info=None): out *= zfactor return out - return da.map_blocks(_chunk_terrain, data, dtype=np.float32, + return da.map_blocks(_chunk_terrain, skeleton, dtype=np.float32, meta=np.array((), dtype=np.float32), **_dask_task_name_kwargs('xrspatial.terrain')) @@ -634,9 +641,13 @@ def _terrain_dask_cupy(data, seed, x_range_scaled, y_range_scaled, zfactor, computes global worley noise range so normalization is consistent across chunk boundaries. """ - data = data * 0 height, width = data.shape + # See _terrain_dask_numpy: the chunk functions regenerate every cell from + # coordinates and never read block values, so map over an uninitialised + # skeleton instead of materialising (and discarding) the template's fill. + skeleton = da.empty_like(data, dtype=cupy.float32) + # --- worley pre-pass: get global min/max to avoid per-chunk seams --- w_norm_range = None if worley_blend > 0: @@ -691,7 +702,7 @@ def _chunk_worley_raw(block, block_info=None): return w_noise raw_worley = da.map_blocks( - _chunk_worley_raw, data, dtype=cupy.float32, + _chunk_worley_raw, skeleton, dtype=cupy.float32, meta=cupy.array((), dtype=cupy.float32), ) (raw_worley,) = dask.persist(raw_worley) @@ -719,7 +730,7 @@ def _chunk_terrain(block, block_info=None): ) return out - data = da.map_blocks(_chunk_terrain, data, dtype=cupy.float32, + data = da.map_blocks(_chunk_terrain, skeleton, dtype=cupy.float32, meta=cupy.array((), dtype=cupy.float32), **_dask_task_name_kwargs('xrspatial.terrain')) diff --git a/xrspatial/tests/test_terrain.py b/xrspatial/tests/test_terrain.py index 96f297dda..00e1a33c8 100644 --- a/xrspatial/tests/test_terrain.py +++ b/xrspatial/tests/test_terrain.py @@ -530,6 +530,61 @@ def test_terrain_all_nan_template_dask_cupy_matches_numpy(): rtol=1e-4, atol=1e-4) +# ===================================================================== +# Issue #3574: the dask backends regenerate every cell from coordinates +# and must not materialize the template's values. Mapping over a +# da.empty_like skeleton drops the template (e.g. from_template's +# da.full(nan)) out of the graph entirely. +# ===================================================================== + +def _poisoned_template(backend='dask'): + """A dask template whose blocks raise if they are ever computed. + + generate_terrain regenerates terrain from coordinates and never reads + these values, so its result must compute without tripping the poison. + """ + def _boom(block): + raise AssertionError("template values were materialized") + + if has_cuda_and_cupy() and 'cupy' in backend: + import cupy + base = da.zeros((50, 50), chunks=(10, 10), dtype=cupy.float32, + meta=cupy.array((), dtype=cupy.float32)) + poisoned = da.map_blocks(_boom, base, dtype=cupy.float32, + meta=cupy.array((), dtype=cupy.float32)) + else: + base = da.zeros((50, 50), chunks=(10, 10), dtype=np.float32) + poisoned = da.map_blocks(_boom, base, dtype=np.float32, + meta=np.array((), dtype=np.float32)) + return xr.DataArray(poisoned, dims=['y', 'x']) + + +@dask_array_available +@pytest.mark.parametrize("worley_blend", [0.0, 0.2]) +def test_terrain_dask_does_not_materialize_template(worley_blend): + """The template's values must never be read; mapping over an empty + skeleton keeps them out of the graph (both map_blocks call sites).""" + terrain = generate_terrain(_poisoned_template(), worley_blend=worley_blend) + result = terrain.compute() # raises if the template were materialized + assert np.isfinite(result.data).all() + + +@dask_array_available +def test_terrain_dask_skeleton_matches_numpy(): + """Swapping the template for an empty skeleton must not change output.""" + expected = generate_terrain(create_test_arr()).data + actual = generate_terrain(_poisoned_template()).compute().data + np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-7) + + +@cuda_and_cupy_available +@dask_array_available +def test_terrain_dask_cupy_does_not_materialize_template(): + terrain = generate_terrain(_poisoned_template(backend='dask+cupy')) + result = terrain.compute() + assert np.isfinite(result.data.get()).all() + + # --------------------------------------------------------------------------- # Fused numpy fast path (worley off) -- independent correctness anchor # --------------------------------------------------------------------------- From 8a9b1ebcf27a10b6dcfbd6e4f07a9d82bb3fa6c5 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Mon, 29 Jun 2026 09:46:44 -0700 Subject: [PATCH 2/2] Address review: exercise worley pre-pass on dask+cupy too (#3574) --- xrspatial/tests/test_terrain.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/xrspatial/tests/test_terrain.py b/xrspatial/tests/test_terrain.py index 00e1a33c8..379eef406 100644 --- a/xrspatial/tests/test_terrain.py +++ b/xrspatial/tests/test_terrain.py @@ -579,8 +579,10 @@ def test_terrain_dask_skeleton_matches_numpy(): @cuda_and_cupy_available @dask_array_available -def test_terrain_dask_cupy_does_not_materialize_template(): - terrain = generate_terrain(_poisoned_template(backend='dask+cupy')) +@pytest.mark.parametrize("worley_blend", [0.0, 0.2]) +def test_terrain_dask_cupy_does_not_materialize_template(worley_blend): + terrain = generate_terrain(_poisoned_template(backend='dask+cupy'), + worley_blend=worley_blend) result = terrain.compute() assert np.isfinite(result.data.get()).all()