Grow the surface_distance Dijkstra heap on demand instead of capping it at height*width (#3723)

[brendancol]

Closes #3723

_dijkstra and _dijkstra_geodesic sized their lazy-deletion min-heap at height * width. That bound is wrong: a pixel is pushed again every time its tentative distance improves, and stale entries stay in the heap until a pop skips them, so live occupancy tracks the number of improving relaxations rather than the pixel count. _heap_push is @ngjit with no bounds checking, so passing the cap writes into whatever allocation follows h_keys / h_rows / h_cols.

On a 48x48 raster with 40% target pixels over 200 m of relief, peak occupancy hit 1.09x the cap and the call aborted the interpreter with double free or corruption (!prev). With NUMBA_BOUNDSCHECK=1 the same input raises IndexError: index is out of bounds from _surface_distance_numpy.

What changed

• The heap starts at one entry per pixel, which the seeding loop can never exceed, and _heap_grow doubles it when a relaxation would overflow.
• Regression tests for the numpy, dask iterative and geodesic paths, checking the dense-target scene against a pure-Python heapq Dijkstra reference.
• A SurfaceDistanceDenseTargets asv benchmark over 5%, 20% and 40% target density, with a peakmem variant. The existing SurfaceDistance benchmark cannot catch a regression here: its integer source raster makes nearly every pixel a target, so almost no relaxation ever improves a distance and the heap barely fills.

Why not the cost_distance approach

cost_distance fixed this defect at xrspatial/cost_distance.py:161 with a static height * width * (n_neighbors + 1) bound. That is also correct, but for 8-connectivity it multiplies the heap allocation by 9, from 24 to 216 bytes per pixel. This module documents an 80 bytes-per-pixel working set in _BYTES_PER_PIXEL and enforces it in _check_memory, so the constant would have to rise to roughly 272 and the guard would start refusing rasters about 3.4x smaller than it does today. Growing on demand keeps the documented footprint accurate and still removes the cap.

Backends

• numpy: fixed. _surface_distance_numpy calls both kernels directly.
• dask+numpy: fixed. The bounded map_overlap path calls _surface_distance_numpy per chunk and _run_tile calls _dijkstra per tile.
• cupy and dask+cupy bounded: not affected. They use the parallel relaxation kernel, which has no heap. The dask+cupy unbounded path converts to numpy and picks up the fix.

Test plan

• pytest xrspatial/tests/test_surface_distance.py -- 41 passed (37 before, 4 new)
• The #3723 reproduction under NUMBA_BOUNDSCHECK=1 returns a full finite grid instead of raising
• The same reproduction without bounds checking no longer aborts
• Dense-target result matches a pure-Python heapq Dijkstra to rtol=1e-5
• flake8 clean on the changed files (the one F841 in the test file predates this branch)
• SurfaceDistanceDenseTargets runs end to end at both grid sizes and all three densities

[brendancol]

PR Review: Grow the surface_distance Dijkstra heap on demand instead of capping it at height*width (#3723)

Self-review from the /sweep-performance agent that wrote the patch, posted as the audit trail. Read it with that bias in mind.

Blockers (must fix before merge)

None.

Suggestions (should fix, not blocking)

• xrspatial/surface_distance.py:80-92 -- the _BYTES_PER_PIXEL = 80 comment still itemizes h_keys / h_rows / h_cols at 8 bytes each, which was exact when the heap was a fixed height * width. It no longer is. After one doubling the heap is 48 bytes per pixel rather than 24, and _heap_grow allocates the replacement arrays before the originals go out of scope, so a growth event transiently holds three times the current heap. On the dense-target scenes this PR is about, _check_memory can therefore under-report peak by roughly 60%. The guard was advisory before and is not made worse by this change (the old code corrupted memory instead of over-allocating), but the constant should say the heap can grow so the next person reading it is not misled.

• xrspatial/surface_distance.py:246-251 and the geodesic equivalent -- the seeding loop pushes without a capacity check. The invariant holding it up is real: the loop visits each pixel once and pushes at most once per pixel, into a heap sized max(height * width, 1). It is stated in the comment at line 236. But it is exactly the kind of invariant that a later change to seeding would break silently, and the failure mode is the out-of-bounds write this PR exists to remove.

Nits (optional improvements)

• xrspatial/surface_distance.py:174-192 -- _heap_grow's docstring is longer than its body, and the part that matters (why the old bound was wrong) belongs where the capacity is chosen, in _dijkstra. Line 236 has a short version already; consider making that the canonical explanation and trimming the helper to one line.

• xrspatial/tests/test_surface_distance.py -- none of the four new tests assert that the heap actually grew. They assert the call completes and matches a heapq reference, which is what a user cares about, but if someone later changed the growth trigger to fire only above, say, four times capacity, the tests would still pass while the bug returned for denser inputs. Checking that the pre-fix code fails all four (verified: IndexError under NUMBA_BOUNDSCHECK=1) is decent evidence the tests bite today; it says nothing about tomorrow.

What looks good

• The bound in the issue is right, and the reasoning holds for both kernels. A lazy-deletion heap's live occupancy tracks improving relaxations, not pixels, and _heap_push has no bounds check because @ngjit does not add one.
• Rejecting cost_distance's static height * width * (n_neighbors + 1) bound is the right call here. It is correct but multiplies the heap by 9 for 8-connectivity, and this module, unlike cost_distance, enforces a documented per-pixel budget that would have to move from 80 to about 272.
• The regression tests were checked against the pre-fix source and all four fail there, including the geodesic one, so the geodesic kernel is genuinely covered rather than just touched.
• The new benchmark targets the regime the old one structurally could not reach. The existing SurfaceDistance class feeds get_xr_dataarray(is_int=True), which makes almost every pixel a target, so relaxations almost never improve anything and the heap never fills.
• Backend claims in the PR body check out: _surface_distance_cupy uses parallel relaxation with no heap, and the dask+cupy unbounded branch converts to numpy and inherits the fix.

Coverage the PR does not claim, and does not have

• cupy and dask+cupy are untested for this scenario because there is nothing to test; neither allocates a heap.
• The bounded map_overlap dask branch is not covered by a dense-target test. It calls _surface_distance_numpy per chunk, so it is fixed, but only the iterative branch has a test.

Checklist

• Algorithm matches reference/paper -- output verified against a pure-Python heapq multi-source Dijkstra
• All implemented backends produce consistent results -- dask iterative matches numpy at rtol=1e-5
• NaN handling is correct -- untouched by this change; barrier and unreachable tests still pass
• Edge cases are covered by tests -- dense targets at 40%, geodesic, dask tiles
• Dask chunk boundaries handled correctly -- _run_tile seeds boundaries before _dijkstra, so seeds stay within one entry per pixel
• No premature materialization or unnecessary copies -- the only new allocation is the heap doubling
• Benchmark exists -- SurfaceDistanceDenseTargets added
• README feature matrix updated -- not applicable, no new function and no backend change
• Docstrings present and accurate

[brendancol]

PR Review: follow-up pass after 4afb223

Second review pass over the same PR, covering the commit that responded to the first one.

Disposition of the first pass

• Fixed: the _BYTES_PER_PIXEL comment now says the heap terms are a starting capacity rather than a bound, so 80 reads as a floor.
• Fixed: both seeding loops now check capacity before pushing, so every push site in both kernels is guarded. The invariant that made the unguarded version safe was real but sat only in a comment, and the failure mode it protected was the out-of-bounds write this PR removes.
• Fixed: the capacity rationale moved from _heap_grow into _dijkstra at the point the capacity is chosen; _heap_grow is back to a one-line docstring.
• Fixed: test_dense_targets_dask_bounded_does_not_overflow covers the bounded map_overlap branch at max_distance=10.0 with 16x16 chunks, matching numpy at rtol=1e-5.
• Dismissed: asserting that the heap actually grew would mean reaching into @ngjit internals or exporting occupancy from the kernel, which costs more than it buys. The pre-fix check (all regression tests fail with IndexError under NUMBA_BOUNDSCHECK=1) stands as the evidence that the tests bite.

Blockers

None.

Suggestions

None outstanding.

Nits

• The seeding guard is dead code on every input that exists today, since a per-pixel loop cannot exceed a per-pixel capacity. It is one comparison inside a loop that already does a log n sift, so the cost is not the concern; someone reading it may just wonder why it is there. The comment above max_heap explains the class of bug, which should be enough.

Verification

• pytest xrspatial/tests/test_surface_distance.py -- 42 passed
• Same file under NUMBA_BOUNDSCHECK=1 -- 42 passed
• pytest xrspatial/tests/test_surface_distance.py xrspatial/tests/test_cost_distance.py -- 133 passed, so the shared _heap_push / _heap_pop are unaffected
• flake8 clean on xrspatial/surface_distance.py and benchmarks/benchmarks/surface_distance.py

[brendancol]

Heads-up for whoever reviews this: PR #3727 (issue #3722) fixes the same defect, found independently by the security sweep in the same /deep-sweep run over surface_distance. Both PRs identify the same root cause (lazy-deletion Dijkstra pushing more entries than height * width, _heap_push writing past the end with no bounds check) and both reproduce it as an interpreter abort.

The fixes are mutually exclusive and touch the same two kernels:

• Size the surface_distance Dijkstra heap to its real push bound (#3722) #3727 uses the static height * width * (n_neighbors + 1) bound, matching what cost_distance.py already carries.
• This PR grows the heap on demand, arguing the static bound raises the per-pixel working set from 80 to roughly 272 bytes, which would make _check_memory refuse rasters about 3.4x smaller than it does today.

Only one should land. Merging both will conflict.

[brendancol]

Follow-up to the cross-reference above, now that both sweeps have reported in full. The tradeoff between #3727 and #3728 is narrower than my first comment implied, and both fixes are internally consistent:

• Size the surface_distance Dijkstra heap to its real push bound (#3722) #3727 takes the static height * width * (n_neighbors + 1) bound and also raises _BYTES_PER_PIXEL from 80 to 272, so the Guard surface_distance() against unbounded eager allocations (#1303) #1305 memory guard keeps modelling the allocation it guards. Cost: _check_memory refuses rasters roughly 3.4x smaller than it does today. Matches the precedent already set in cost_distance.py.
• Grow the surface_distance Dijkstra heap on demand instead of capping it at height*width (#3723) #3728 grows the heap on demand, keeping the documented 80 bytes-per-pixel footprint and today's raster capacity, at the cost of diverging from that precedent and adding a _heap_grow path.

So it is a straight choice: consistency with cost_distance and a provably sufficient static bound (#3727), against preserved capacity and a smaller memory footprint (#3728). Both reproduce the abort and both verify the fix under NUMBA_BOUNDSCHECK=1.

One item worth carrying forward whichever way this lands: #3722 records that pathfinding.py:304 uses the same height * width heap sizing. The security sweep could not make it overflow because that path is single-source A*, but it should be re-checked rather than assumed safe.