Skip to content

Fixes #878: extend KDTree path to allocation/direction to prevent OOM#893

Merged
brendancol merged 1 commit intomasterfrom
fix/878-allocation-direction-kdtree-oom
Feb 25, 2026
Merged

Fixes #878: extend KDTree path to allocation/direction to prevent OOM#893
brendancol merged 1 commit intomasterfrom
fix/878-allocation-direction-kdtree-oom

Conversation

@brendancol
Copy link
Contributor

Summary

  • Extends the memory-safe KDTree path (previously PROXIMITY-only) to handle ALLOCATION and DIRECTION modes, preventing OOM when max_distance=inf on large dask-backed rasters
  • cKDTree.query() returns both distances and nearest-target indices, which is enough to derive allocation values and compass direction angles
  • Adds a memory guard in the non-KDTree fallback path (GREAT_CIRCLE or missing scipy) that raises MemoryError with actionable guidance instead of silently OOMing

Key changes

  • _vectorized_calc_direction() — new array-based compass direction function matching the scalar _calc_direction() floating-point behaviour
  • _stream_target_counts() — now also caches target pixel values alongside coordinates
  • _collect_region_targets() — returns (coords, vals) tuple
  • _kdtree_chunk_fn() / _tiled_chunk_query() — extended with mode dispatch (PROXIMITY/ALLOCATION/DIRECTION)
  • KDTree gate widened — removed process_mode == PROXIMITY restriction
  • Memory guardMemoryError raised when single-chunk rechunk would exceed 80% available memory

Test plan

  • test_allocation_dask_kdtree_matches_numpy (EUCLIDEAN + MANHATTAN) — consistency check against proximity
  • test_direction_dask_kdtree_matches_numpy (EUCLIDEAN + MANHATTAN) — consistency check against allocation
  • test_allocation_dask_kdtree_with_target_values — target filtering through KDTree
  • test_direction_dask_kdtree_with_target_values — target filtering through KDTree
  • test_allocation_dask_kdtree_no_targets / test_direction_dask_kdtree_no_targets — all-NaN
  • test_allocation_dask_kdtree_max_distance / test_direction_dask_kdtree_max_distance — truncation
  • test_great_circle_dask_unbounded_memory_guard — MemoryError for GREAT_CIRCLE
  • test_no_scipy_dask_unbounded_memory_guard — MemoryError when scipy missing
  • All 51 existing + new tests pass, no regressions

The memory-safe KDTree path was previously gated to PROXIMITY mode only.
This extends it to handle ALLOCATION and DIRECTION modes, since
cKDTree.query() returns both distances and nearest-target indices —
enough to derive allocation values and direction angles.

Changes:
- Add _vectorized_calc_direction() for array-based compass directions
- Extend _stream_target_counts() to cache target pixel values
- Extend _collect_region_targets() to return target values alongside coords
- Extend _kdtree_chunk_fn() and _tiled_chunk_query() for all three modes
- Thread process_mode and values_cache through the full call chain
- Remove process_mode == PROXIMITY restriction from KDTree gate
- Add memory guard for non-KDTree fallback (GREAT_CIRCLE, no scipy)
- Update memory estimate from 48 to 52 bytes per target
@brendancol brendancol merged commit f801c5e into master Feb 25, 2026
10 checks passed
@brendancol brendancol deleted the fix/878-allocation-direction-kdtree-oom branch February 26, 2026 14:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant