From 968e4e4e014a2b15de884843768108d55f6f7c73 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Wed, 27 May 2026 11:01:44 -0700 Subject: [PATCH 1/2] Fix hypsometric_integral dask+cupy backend (#2525) _hi_dask_cupy converted cupy chunks to numpy then delegated to _hi_dask_numpy, returning a dask+numpy array even when the input was dask+cupy. Wrap the painted output through map_blocks(cupy.asarray, ...) so chunks land back as cupy arrays. is_dask_cupy(result) is now True for dask+cupy inputs, matching _regions_dask_cupy's behavior. Add test_hi_preserves_backend_2525 covering all four backends (numpy / cupy / dask+numpy / dask+cupy). Closes #2525 --- xrspatial/tests/test_hypsometric_integral.py | 63 ++++++++++++++++++++ xrspatial/zonal.py | 14 ++++- 2 files changed, 75 insertions(+), 2 deletions(-) diff --git a/xrspatial/tests/test_hypsometric_integral.py b/xrspatial/tests/test_hypsometric_integral.py index 0debcc338..351b7dd92 100644 --- a/xrspatial/tests/test_hypsometric_integral.py +++ b/xrspatial/tests/test_hypsometric_integral.py @@ -254,3 +254,66 @@ def test_hypsometric_integral_list_of_pairs_zones(): result = hypsometric_integral(zones_pairs, values, nodata=0) assert isinstance(result, xr.DataArray) assert result.shape == values.shape + + +# --- backend-consistency regression (#2525) --------------------------------- + +def _has_cuda_and_cupy(): + try: + from numba import cuda + import cupy # noqa: F401 + return cuda.is_available() + except Exception: + return False + + +@pytest.mark.parametrize("backend", ['numpy', 'cupy', 'dask+numpy', 'dask+cupy']) +def test_hi_preserves_backend_2525(backend): + """Regression: output backend must match input backend. + + Before the fix, `_hi_dask_cupy` converted chunks to numpy and called + `_hi_dask_numpy`, returning a dask+numpy array even when the input was + dask+cupy. Downstream dispatch via ArrayTypeFunctionMapping then routed + through the numpy path silently. See issue #2525. + """ + if 'cupy' in backend and not _has_cuda_and_cupy(): + pytest.skip("Requires CUDA and CuPy") + if 'dask' in backend and da is None: + pytest.skip("Requires Dask") + + from xrspatial.zonal import hypsometric_integral + from xrspatial.utils import ( + is_cupy_array, is_dask_cupy, has_dask_array, + ) + + zones_np = np.array([[1, 1, 2, 2], + [1, 1, 2, 2]], dtype=np.int32) + values_np = np.array([[10.0, 20.0, 30.0, 40.0], + [15.0, 25.0, 35.0, 45.0]]) + + zones = create_test_raster(zones_np, backend, dims=['y', 'x'], + chunks=(2, 2)) + values = create_test_raster(values_np, backend, dims=['y', 'x'], + chunks=(2, 2)) + + result = hypsometric_integral(zones, values) + assert isinstance(result, xr.DataArray) + assert result.shape == values.shape + + if backend == 'numpy': + assert isinstance(result.data, np.ndarray) + elif backend == 'cupy': + assert is_cupy_array(result.data) + elif backend == 'dask+numpy': + assert has_dask_array() and isinstance(result.data, da.Array) + assert not is_dask_cupy(result) + # confirm chunks are numpy + sample = result.data.blocks[0, 0].compute() + assert isinstance(sample, np.ndarray) + elif backend == 'dask+cupy': + assert is_dask_cupy(result), ( + "dask+cupy input must produce dask+cupy output (#2525)" + ) + # confirm chunks are cupy + sample = result.data.blocks[0, 0].compute() + assert is_cupy_array(sample) diff --git a/xrspatial/zonal.py b/xrspatial/zonal.py index 9db5e3d6f..a34238281 100644 --- a/xrspatial/zonal.py +++ b/xrspatial/zonal.py @@ -1538,14 +1538,24 @@ def _paint(zones_chunk, values_chunk, hi_map): def _hi_dask_cupy(zones_data, values_data, nodata): - """Dask+cupy backend: convert chunks to numpy, delegate.""" + """Dask+cupy backend: convert chunks to numpy, delegate, then re-wrap as cupy. + + The per-zone HI lookup table is a Python dict so the reduce step has to + pass through host memory. The painted output is wrapped back as cupy + chunks so the returned dask graph yields cupy arrays — keeping the + backend consistent with the dask+cupy input (issue #2525). + """ zones_cpu = zones_data.map_blocks( lambda x: x.get(), dtype=zones_data.dtype, meta=np.array(()), ) values_cpu = values_data.map_blocks( lambda x: x.get(), dtype=values_data.dtype, meta=np.array(()), ) - return _hi_dask_numpy(zones_cpu, values_cpu, nodata) + result_cpu = _hi_dask_numpy(zones_cpu, values_cpu, nodata) + # Re-wrap each numpy chunk as cupy so dispatch downstream stays on GPU. + return result_cpu.map_blocks( + cupy.asarray, dtype=result_cpu.dtype, meta=cupy.array(()), + ) def hypsometric_integral( From 06cb48d3826adebc480d66561b8c5b6a25c9640f Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Thu, 28 May 2026 07:17:32 -0700 Subject: [PATCH 2/2] Use shared backend helpers in HI regression test Replace locally redeclared _has_cuda_and_cupy with has_cuda_and_cupy from xrspatial.utils, and use has_dask_array() instead of `da is None`. Same pattern other tests follow. Drop the em-dash in the _hi_dask_cupy docstring while in the area. --- xrspatial/tests/test_hypsometric_integral.py | 21 +++++++------------- xrspatial/zonal.py | 2 +- 2 files changed, 8 insertions(+), 15 deletions(-) diff --git a/xrspatial/tests/test_hypsometric_integral.py b/xrspatial/tests/test_hypsometric_integral.py index 351b7dd92..abe1267a7 100644 --- a/xrspatial/tests/test_hypsometric_integral.py +++ b/xrspatial/tests/test_hypsometric_integral.py @@ -258,15 +258,6 @@ def test_hypsometric_integral_list_of_pairs_zones(): # --- backend-consistency regression (#2525) --------------------------------- -def _has_cuda_and_cupy(): - try: - from numba import cuda - import cupy # noqa: F401 - return cuda.is_available() - except Exception: - return False - - @pytest.mark.parametrize("backend", ['numpy', 'cupy', 'dask+numpy', 'dask+cupy']) def test_hi_preserves_backend_2525(backend): """Regression: output backend must match input backend. @@ -276,15 +267,17 @@ def test_hi_preserves_backend_2525(backend): dask+cupy. Downstream dispatch via ArrayTypeFunctionMapping then routed through the numpy path silently. See issue #2525. """ - if 'cupy' in backend and not _has_cuda_and_cupy(): + from xrspatial.utils import ( + has_cuda_and_cupy, has_dask_array, + is_cupy_array, is_dask_cupy, + ) + + if 'cupy' in backend and not has_cuda_and_cupy(): pytest.skip("Requires CUDA and CuPy") - if 'dask' in backend and da is None: + if 'dask' in backend and not has_dask_array(): pytest.skip("Requires Dask") from xrspatial.zonal import hypsometric_integral - from xrspatial.utils import ( - is_cupy_array, is_dask_cupy, has_dask_array, - ) zones_np = np.array([[1, 1, 2, 2], [1, 1, 2, 2]], dtype=np.int32) diff --git a/xrspatial/zonal.py b/xrspatial/zonal.py index a34238281..b24c6cf64 100644 --- a/xrspatial/zonal.py +++ b/xrspatial/zonal.py @@ -1542,7 +1542,7 @@ def _hi_dask_cupy(zones_data, values_data, nodata): The per-zone HI lookup table is a Python dict so the reduce step has to pass through host memory. The painted output is wrapped back as cupy - chunks so the returned dask graph yields cupy arrays — keeping the + chunks so the returned dask graph yields cupy arrays, keeping the backend consistent with the dask+cupy input (issue #2525). """ zones_cpu = zones_data.map_blocks(