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

.claude/sweep-performance-state.csv

@@ -43,7 +43,7 @@ resample,2026-04-15T12:00:00Z,SAFE,compute-bound,0,false-positive,Downgraded. GP
 sieve,2026-04-14T12:00:00Z,WILL OOM,memory-bound,0,false-positive,False positive. Memory guards already in place on both dask paths. CCL is inherently global — documented limitation. CuPy CPU fallback is deliberate and documented.
 sky_view_factor,2026-03-31T18:00:00Z,SAFE,compute-bound,0,,
 slope,2026-03-31T18:00:00Z,SAFE,compute-bound,0,,
-surface_distance,2026-03-31T18:00:00Z,SAFE,memory-bound,0,1128,Memory guard added to dd_grid allocation.
+surface_distance,2026-08-16,RISKY,compute-bound,0,3723,"CRITICAL #3723/PR: _dijkstra and _dijkstra_geodesic sized the lazy-deletion heap height*width; peak occupancy hit 1.09x cap at 40% target density, _heap_push wrote OOB (SIGABRT 'double free or corruption'); fixed by doubling on demand rather than cost_distance's static (n_neighbors+1) bound, which would have raised the 80 B/px guard to ~272. MEDIUM (documented, not fixed): dask iterative path recomputes each source and elev block 6x (2 in _preprocess_tiles_sd, 4 across sweeps, 1 in _assemble_sd) - measured on 40x40/80x80/120x120; MEDIUM: cupy Bellman-Ford relaxation is 6-27x slower than the numpy Dijkstra (512^2: 0.315s vs 0.051s, 373 kernel launches each with a device sync), and DIRECTION mode round-trips to host for _vectorized_calc_direction; MEDIUM: map_overlap depth guard tests pad < max(chunks), so depth may approach chunk size (~9x redundant work) and exceed the smallest chunk when chunking is uneven. Memory guard is correctly chunk-scoped on dask (test_dask_path_bounded_per_chunk) - the morphology #3401 false-MemoryError hazard is NOT present. _sd_relax_kernel has ~8 float64 locals, no register pressure. OOM verdict RISKY: memory stays chunk-bounded but the whole iterative Dijkstra runs eagerly and single-threaded on the client at graph-build time."
 terrain,2026-03-31T18:00:00Z,RISKY,compute-bound,0,,
 terrain_metrics,2026-03-31T18:00:00Z,SAFE,memory-bound,0,,
 viewshed,2026-04-05T12:00:00Z,SAFE,memory-bound,0,fixed-in-tree,Tier B memory estimate tightened from 280 to 368 bytes/pixel (accounts for lexsort double-alloc + computed raster). astype copy=False avoids needless float64 copy.

benchmarks/benchmarks/surface_distance.py

@@ -1,4 +1,5 @@
 import numpy as np
+import xarray as xr
 
 from xrspatial.surface_distance import (
     surface_distance, surface_allocation, surface_direction,
@@ -24,3 +25,37 @@ def time_surface_allocation(self, nx, type):
 
     def time_surface_direction(self, nx, type):
         surface_direction(self.agg, self.elev)
+
+
+class SurfaceDistanceDenseTargets:
+    """Dijkstra with a mid-density target raster over rugged relief.
+
+    This is the regime where the lazy-deletion heap holds more than one
+    live entry per pixel (#3723).  The existing SurfaceDistance benchmark
+    misses it: its integer source raster makes nearly every pixel a
+    target, so almost no relaxation ever improves a distance.
+    """
+
+    params = ([200, 400], [0.05, 0.2, 0.4])
+    param_names = ("nx", "target_fraction")
+
+    def setup(self, nx, target_fraction):
+        ny = nx // 2
+        rng = np.random.default_rng(71942)
+        source = np.zeros((ny, nx), dtype=np.float64)
+        n_targets = max(1, int(ny * nx * target_fraction))
+        source.flat[rng.choice(ny * nx, size=n_targets, replace=False)] = 1.0
+        elev = rng.random((ny, nx)) * 200.0
+
+        coords = dict(y=np.arange(ny, dtype=np.float64),
+                      x=np.arange(nx, dtype=np.float64))
+        self.agg = xr.DataArray(source, coords=coords, dims=["y", "x"],
+                                attrs={"res": (1.0, 1.0)})
+        self.elev = xr.DataArray(elev, coords=coords, dims=["y", "x"],
+                                 attrs={"res": (1.0, 1.0)})
+
+    def time_surface_distance(self, nx, target_fraction):
+        surface_distance(self.agg, self.elev)
+
+    def peakmem_surface_distance(self, nx, target_fraction):
+        surface_distance(self.agg, self.elev)

xrspatial/surface_distance.py

@@ -170,6 +170,28 @@ def _check_gpu_memory(rows, cols):
 # ---------------------------------------------------------------------------
 
 
+@ngjit
+def _heap_grow(keys, rows, cols, size):
+    """Return copies of the heap arrays with twice the capacity.
+
+    The Dijkstra kernels below use a lazy-deletion min-heap: a pixel is
+    pushed again every time its tentative distance improves, and stale
+    entries linger until a pop skips them.  Live occupancy is therefore
+    bounded by the number of improving relaxations, not by the pixel
+    count, and a fixed ``height * width`` heap can overflow (#3723).
+    Growing on demand keeps the usual footprint at one entry per pixel
+    without capping the number of pushes.
+    """
+    cap = 2 * len(keys)
+    new_keys = np.empty(cap, dtype=np.float64)
+    new_rows = np.empty(cap, dtype=np.int64)
+    new_cols = np.empty(cap, dtype=np.int64)
+    new_keys[:size] = keys[:size]
+    new_rows[:size] = rows[:size]
+    new_cols[:size] = cols[:size]
+    return new_keys, new_rows, new_cols
+
+
 @ngjit
 def _seed_sources(source_data, elev_data, target_values,
                   dist, alloc, src_row, src_col):
@@ -211,7 +233,9 @@ def _dijkstra(elev_data, height, width, max_distance,
     """
     n_neighbors = len(dy)
 
-    max_heap = height * width
+    # Starting capacity: one entry per pixel, which the seeding loop below
+    # can never exceed.  Relaxations can, so _heap_grow doubles it there.
+    max_heap = max(height * width, 1)
     h_keys = np.empty(max_heap, dtype=np.float64)
     h_rows = np.empty(max_heap, dtype=np.int64)
     h_cols = np.empty(max_heap, dtype=np.int64)
@@ -260,6 +284,9 @@ def _dijkstra(elev_data, height, width, max_distance,
                 alloc[vr, vc] = alloc[ur, uc]
                 src_row[vr, vc] = src_row[ur, uc]
                 src_col[vr, vc] = src_col[ur, uc]
+                if h_size == len(h_keys):
+                    h_keys, h_rows, h_cols = _heap_grow(
+                        h_keys, h_rows, h_cols, h_size)
                 h_size = _heap_push(h_keys, h_rows, h_cols, h_size,
                                     new_cost, vr, vc)
 
@@ -274,7 +301,8 @@ def _dijkstra_geodesic(elev_data, height, width, max_distance,
     """
     n_neighbors = len(dy)
 
-    max_heap = height * width
+    # See _dijkstra for the heap-capacity rationale.
+    max_heap = max(height * width, 1)
     h_keys = np.empty(max_heap, dtype=np.float64)
     h_rows = np.empty(max_heap, dtype=np.int64)
     h_cols = np.empty(max_heap, dtype=np.int64)
@@ -322,6 +350,9 @@ def _dijkstra_geodesic(elev_data, height, width, max_distance,
                 alloc[vr, vc] = alloc[ur, uc]
                 src_row[vr, vc] = src_row[ur, uc]
                 src_col[vr, vc] = src_col[ur, uc]
+                if h_size == len(h_keys):
+                    h_keys, h_rows, h_cols = _heap_grow(
+                        h_keys, h_rows, h_cols, h_size)
                 h_size = _heap_push(h_keys, h_rows, h_cols, h_size,
                                     new_cost, vr, vc)
 

xrspatial/tests/test_surface_distance.py

@@ -792,3 +792,124 @@ def test_error_message_mentions_grid_size(self):
                 surface_distance(raster, elevation)
             with pytest.raises(MemoryError, match="dask"):
                 surface_distance(raster, elevation)
+
+
+# ---------------------------------------------------------------------------
+# Heap capacity regression (#3723)
+# ---------------------------------------------------------------------------
+
+
+def _reference_surface_distance(source, elev, connectivity=8, cellsize=1.0):
+    """Pure-Python multi-source Dijkstra reference for surface distance."""
+    import heapq
+
+    h, w = source.shape
+    diag = np.sqrt(2.0) * cellsize
+    if connectivity == 8:
+        nbrs = [(-1, -1, diag), (-1, 0, cellsize), (-1, 1, diag),
+                (0, -1, cellsize), (0, 1, cellsize),
+                (1, -1, diag), (1, 0, cellsize), (1, 1, diag)]
+    else:
+        nbrs = [(0, -1, cellsize), (-1, 0, cellsize),
+                (1, 0, cellsize), (0, 1, cellsize)]
+
+    dist = np.full((h, w), np.inf)
+    heap = []
+    for r in range(h):
+        for c in range(w):
+            if (source[r, c] != 0 and np.isfinite(source[r, c])
+                    and np.isfinite(elev[r, c])):
+                dist[r, c] = 0.0
+                heapq.heappush(heap, (0.0, r, c))
+
+    done = np.zeros((h, w), dtype=bool)
+    while heap:
+        d, r, c = heapq.heappop(heap)
+        if done[r, c]:
+            continue
+        done[r, c] = True
+        for dr, dc, hd in nbrs:
+            vr, vc = r + dr, c + dc
+            if not (0 <= vr < h and 0 <= vc < w) or done[vr, vc]:
+                continue
+            if not np.isfinite(elev[vr, vc]):
+                continue
+            dz = elev[vr, vc] - elev[r, c]
+            nd = d + np.sqrt(hd * hd + dz * dz)
+            if nd < dist[vr, vc]:
+                dist[vr, vc] = nd
+                heapq.heappush(heap, (nd, vr, vc))
+    return dist
+
+
+def _dense_target_scene(n=48, n_targets=921, relief=200.0, seed=1):
+    """Dense-target scene that overflowed the old height*width heap."""
+    rng = np.random.default_rng(seed)
+    source = np.zeros((n, n), dtype=np.float64)
+    source.flat[rng.choice(n * n, size=n_targets, replace=False)] = 1.0
+    elev = rng.random((n, n)) * relief
+    return source, elev
+
+
+def test_dense_targets_do_not_overflow_the_heap():
+    """A lazy-deletion heap can exceed height*width live entries.
+
+    Before #3723 the heap arrays were sized height*width, so this scene
+    made _heap_push write past the end of them (SIGABRT without bounds
+    checking, IndexError with NUMBA_BOUNDSCHECK=1).
+    """
+    source, elev = _dense_target_scene()
+    raster = _make_raster(source)
+    elevation = _make_raster(elev)
+
+    result = _compute(surface_distance(raster, elevation, connectivity=8))
+    expected = _reference_surface_distance(source, elev, connectivity=8)
+
+    assert np.all(np.isfinite(result))
+    np.testing.assert_allclose(result, expected.astype(np.float32),
+                               rtol=1e-5, atol=1e-4)
+
+
+def test_dense_targets_allocation_and_direction_do_not_overflow():
+    """The allocation and direction modes share the same Dijkstra kernel."""
+    source, elev = _dense_target_scene()
+    raster = _make_raster(source)
+    elevation = _make_raster(elev)
+
+    alloc = _compute(surface_allocation(raster, elevation, connectivity=8))
+    direction = _compute(surface_direction(raster, elevation, connectivity=8))
+
+    assert np.all(alloc == 1.0)
+    assert np.all(np.isfinite(direction))
+
+
+@pytest.mark.skipif(da is None, reason="dask not installed")
+def test_dense_targets_dask_iterative_does_not_overflow():
+    """The dask iterative path runs the same kernel per tile."""
+    source, elev = _dense_target_scene()
+    raster_np = _make_raster(source)
+    elev_np = _make_raster(elev)
+    raster = _make_raster(source, backend='dask+numpy', chunks=(48, 48))
+    elevation = _make_raster(elev, backend='dask+numpy', chunks=(48, 48))
+
+    np_result = _compute(surface_distance(raster_np, elev_np))
+    with pytest.warns(UserWarning, match="iterative"):
+        dask_result = _compute(surface_distance(raster, elevation))
+
+    np.testing.assert_allclose(dask_result, np_result, rtol=1e-5,
+                               equal_nan=True)
+
+
+def test_geodesic_dense_targets_do_not_overflow():
+    """_dijkstra_geodesic carries the same heap sizing."""
+    source, elev = _dense_target_scene(n=32, n_targets=410, relief=200.0)
+    h, w = source.shape
+    coords = {'y': np.linspace(10.0, 10.0 + 0.01 * (h - 1), h),
+              'x': np.linspace(20.0, 20.0 + 0.01 * (w - 1), w)}
+    raster = xr.DataArray(source, dims=['y', 'x'], coords=coords,
+                          attrs={'res': (0.01, 0.01)})
+    elevation = xr.DataArray(elev, dims=['y', 'x'], coords=coords,
+                             attrs={'res': (0.01, 0.01)})
+
+    result = _compute(surface_distance(raster, elevation, method='geodesic'))
+    assert np.all(np.isfinite(result))