Skip to content

Fix hypsometric_integral dask+cupy backend (#2525)#2529

Merged
brendancol merged 2 commits into
mainfrom
issue-2525
May 28, 2026
Merged

Fix hypsometric_integral dask+cupy backend (#2525)#2529
brendancol merged 2 commits into
mainfrom
issue-2525

Conversation

@brendancol
Copy link
Copy Markdown
Contributor

Closes #2525.

Summary

  • _hi_dask_cupy previously converted cupy chunks to numpy and then called _hi_dask_numpy, returning a dask+numpy result for dask+cupy input. Downstream ArrayTypeFunctionMapping dispatch silently routed through the numpy path.
  • After this change, the painted output is re-wrapped through map_blocks(cupy.asarray, ...) so each chunk lands back as a cupy array. The dask graph stays lazy and is_dask_cupy(result) is True for dask+cupy inputs.
  • The numpy reduce in the middle is kept as-is. The per-zone HI lookup is a Python dict so the reduce must pass through host memory; only the painted output needs to match the input backend.

Backend coverage

  • numpy: unchanged, returns numpy ndarray
  • cupy: unchanged, returns cupy ndarray
  • dask+numpy: unchanged, returns dask Array with numpy chunks
  • dask+cupy: now returns dask Array with cupy chunks (was: numpy chunks)

Test plan

  • New parametrized regression test test_hi_preserves_backend_2525 asserts the output backend matches the input across all four backends.
  • Existing test_hypsometric_integral.py suite still passes (33 tests).
  • Existing test_zonal.py suite still passes (125 tests).
  • Verified the new test fails with the fix reverted (confirms it catches the regression).

Audit trail

Surfaced by /sweep-metadata on the zonal module, 2026-05-27. Cat 5 (backend-inconsistent metadata), severity MEDIUM.

_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
@github-actions github-actions Bot added the performance PR touches performance-sensitive code label May 27, 2026
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.
@brendancol
Copy link
Copy Markdown
Contributor Author

Review summary

The fix is correct and minimal. _hi_dask_cupy now wraps the painted output via map_blocks(cupy.asarray, ..., meta=cupy.array(())), mirroring the pattern in _apply_dask_cupy (zonal.py:1720, 1731). All four backends remain wired in ArrayTypeFunctionMapping, and the new parametrized test verifies both DataArray-level and chunk-level backend identity.

Blockers

None.

Suggestions

None blocking. The extra map_blocks(cupy.asarray, ...) adds one host-to-device copy per chunk, but this is unavoidable given that the per-zone lookup is a Python dict and must go through host memory anyway.

Nits applied in follow-up commit

  • Replaced the locally redeclared _has_cuda_and_cupy() with the shared has_cuda_and_cupy from xrspatial.utils, and swapped da is None for has_dask_array() to match the rest of the test suite.
  • Dropped the em-dash in the _hi_dask_cupy docstring.

What looks good

  • Test asserts chunk-level type via .blocks[0, 0].compute(), not just the DataArray wrapper.
  • PR body confirms the test fails when the fix is reverted.
  • Docstring updated to explain why the round-trip through host memory exists.

Pre-existing items (not from this PR)

  • xrspatial/zonal.py:39 F401 not_implemented_func unused and xrspatial/zonal.py:455 E501 both exist on the merge base.

@brendancol brendancol merged commit 29b8283 into main May 28, 2026
6 of 7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

performance PR touches performance-sensitive code

Projects

None yet

Development

Successfully merging this pull request may close these issues.

hypsometric_integral dask+cupy backend silently returns dask+numpy result

1 participant