Make hotspots() lazy for Dask input#2802
Open
brendancol wants to merge 2 commits into
Open
Conversation
hotspots() used to call da.compute() on the global mean and std while building the graph, so calling it ran about 12 Dask tasks before you asked for any output. This keeps those reductions as lazy 0-d dask arrays that broadcast into the z-score. The convolution runs through map_overlap and the classification through map_blocks, so the Dask and Dask+CuPy paths build their graph without running a single task. test_dask_laziness.py now asserts zero tasks run on the call, via a dask task-count callback, instead of only checking the return type.
brendancol
commented
Jun 1, 2026
Contributor
Author
brendancol
left a comment
There was a problem hiding this comment.
PR Review: Make hotspots() lazy for Dask input
Blockers (must fix before merge)
None.
Suggestions (should fix, not blocking)
None.
Nits (optional improvements)
- focal.py: the convolve
map_overlappassesmeta=np.array((), dtype=np.float32), but for float64 input the convolved chunks are float64 (the up-front cast only runs for non-float input). The final result is still int8 from themap_blocksmeta, and I verified correctness against numpy on float64 input, so this is cosmetic. If you want the intermediate meta to be exact you could cast the data to float32 unconditionally.
What looks good
- The eager
da.compute()is gone from both dask paths. The global mean/std are 0-d dask reductions broadcast into the z-score, so nothing runs on the call. - Convolution still goes through
map_overlapwith depth = kernel radius and the boundary forwarded from the public API, so halo handling is unchanged. - The new test asserts zero tasks execute on the call via a dask task-count callback, which is what actually catches this regression. The old eager path ran 12 tasks.
- Dask output is bit-identical to numpy on float32, float64, integer, and ragged-chunk inputs.
- The dask+cupy path mirrors the numpy path and keeps each chunk on the device.
Checklist
- Algorithm matches reference (z-score classification unchanged)
- All implemented backends produce consistent results
- NaN handling is correct (boundary=nan halo unchanged)
- Edge cases covered (ragged chunks, integer/float input verified)
- Dask chunk boundaries handled correctly
- No premature materialization
- Benchmark exists or not needed (bug fix, no new API)
- README feature matrix updated (not applicable)
- Docstrings present and accurate
Address review nit: the convolve map_overlap declared a float32 meta but float64 input stayed float64 through the intermediate. Cast to float32 up front in both dask paths, matching the single-array numpy path.
brendancol
commented
Jun 1, 2026
Contributor
Author
brendancol
left a comment
There was a problem hiding this comment.
Follow-up review (after 29251b4)
The one nit from the previous pass is resolved. Both dask paths now cast the input to float32 up front, so the convolve map_overlap meta and the actual chunk dtype agree, and the dask paths line up with the single-array numpy path.
Rechecked after the change:
- hotspots and dask-laziness tests pass (33 passed).
- float64 input still produces a result bit-identical to numpy, lazy dtype int8.
- flake8 clean.
No remaining blockers, suggestions, or nits.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
hotspots()computed the global mean and std eagerly withda.compute()while building the graph for Dask input, so calling it ran about 12 Dask tasks before you asked for any output (xrspatial/focal.py around lines 1289-1293, and the Dask+CuPy path around line 1336).The Dask and Dask+CuPy paths are now lazy:
map_overlap, then classified viamap_blocks.ZeroDivisionErroris dropped on the dask path. The numpy and cupy single-array paths still raise it, and the zero-std test uses the numpy backend, so it stays green.Backend coverage
Test plan
test_dask_laziness.py::test_hotspots_no_eager_computeasserts zero dask tasks run on the call (task-count callback), not just the return type.pytest xrspatial/tests/test_focal.py xrspatial/tests/test_emerging_hotspots.py xrspatial/tests/test_dask_laziness.pypasses (218 passed).Closes #2772