surface_distance benchmarks miss cupy, dask+cupy, the bounded dask branch, and geodesic - and the source raster makes 99.95% of pixels sources

[brendancol]

benchmarks/benchmarks/surface_distance.py covers 1 of the 6 compute paths in xrspatial/surface_distance.py, and the one it does cover runs on source data that never triggers the Dijkstra relaxation.

The module is 1561 lines with a numba planar kernel, a separate numba geodesic kernel, a @cuda.jit relaxation kernel, a bounded dask map_overlap path, a ~570-line iterative dask tile path, and a dask+cupy wrapper. The benchmark file is 26 lines.

Coverage matrix

Compute path	Benchmarked
numpy planar (_dijkstra)	yes, but on degenerate input (see below)
numpy geodesic (_dijkstra_geodesic + _precompute_dd_grid)	no
cupy (_sd_relax_kernel)	no
dask bounded (_surface_distance_dask_bounded, map_overlap)	no
dask unbounded (_sd_dask_iterative)	yes
dask+cupy (_surface_distance_dask_cupy)	no

1. cupy and dask+cupy are never parameterized (HIGH)

params = ([100, 300, 1000], ["numpy", "dask"]). The module dispatches to _surface_distance_cupy and _surface_distance_dask_cupy, both reachable and both working on this host. A regression in the CUDA relaxation kernel would ship with nothing to catch it.

Both run fine at the suite's existing sizes (500x1000, sparse sources, this host):

nx=1000 cupy      unbounded  0.2496s
nx=1000 cupy      bounded    0.0138s
nx=1000 dask+cupy unbounded  1.3020s
nx=1000 dask+cupy bounded    0.1253s

benchmarks/benchmarks/pathfinding.py::AStarSearch already parameterizes all four types, so the pattern exists.

Worth noting while it is fresh: cupy unbounded at nx=1000 is slower than numpy (0.2496s vs 0.1127s). That is a separate question about the GPU relaxation kernel, not something this issue proposes to change.

2. Only the unbounded dask branch is measured, and it is never computed (HIGH)

_surface_distance_dask picks a branch on max_distance:

use_overlap = False
if np.isfinite(max_distance):
    pad = int(max_distance / min_cellsize) + 1
    chunks_y, chunks_x = source_da.chunks
    if pad < max(chunks_y) and pad < max(chunks_x):
        use_overlap = True

The benchmark calls surface_distance(self.agg, self.elev) with the default max_distance=np.inf, so it always lands on the iterative fallback. The call emits the warning every time:

WARNING: surface_distance: max_distance is infinite or the implied radius exceeds chunk dimensions; using iterative tile Dijkstra

The bounded branch is the one the docstring recommends ("A finite value enables efficient Dask parallelisation") and it has never been timed.

Compounding this: the benchmark never calls .compute(). Most dask benchmarks in the suite do (flood.py, twi.py, interpolate.py, flow_length.py, and others use if hasattr(result.data, 'compute'): result.data.compute()). Today the unbounded numbers are non-trivial only because _sd_dask_iterative happens to compute eagerly while building its graph. A bounded benchmark added without .compute() would measure 0.005s of graph construction and nothing else.

benchmarks/benchmarks/cost_distance.py already solved the same problem for the sibling module, with a time_cost_distance_bounded method that picks a max_cost whose pixel radius stays inside one chunk.

3. Source data makes almost every pixel a source (MEDIUM)

get_xr_dataarray(..., is_int=True) draws rng.integers(-nx, nx), and _seed_sources treats any non-zero finite value as a target. At nx=1000:

total pixels        : 500000
nonzero (= sources) : 499769 (99.95%)

So dist is 0 nearly everywhere before the main loop starts, new_cost < dist[vr, vc] is never true, and the relaxation body never executes. The benchmark measures pushing 500k entries onto the heap and popping them again.

current bench source pattern : 0.0158s   max distance found: 3.62
sparse (20 sources)          : 0.1005s   max distance found: 160.44   (6.4x slower)

A maximum distance of 3.6 map units on a 360x180 degree grid is one hop. The propagation this function exists to compute is not being timed.

4. method='geodesic' is never benchmarked (MEDIUM)

_dijkstra_geodesic is a separate 65-line kernel and _precompute_dd_grid builds an (n_neighbors, ny, nx) grid ahead of it. Neither is touched. It costs about 2x planar:

nx=300  geodesic 0.0165s vs planar 0.0066s (2.5x)
nx=1000 geodesic 0.2276s vs planar 0.1032s (2.2x)

5. Smallest size is 50x100 (LOW, not proposed for change)

nx=100 gives a 50x100 grid where setup and JIT dispatch dominate. There is a ladder up to nx=1000 so scaling is still visible; recording this rather than changing it.

Proposed fix

Benchmark-file changes only, no source changes:

• parameterize ["numpy", "cupy", "dask", "dask+cupy"], matching pathfinding.py
• replace the near-uniform source raster with scattered point sources so the search crosses the grid
• add a bounded (max_distance=) method that reaches map_overlap, following the cost_distance.py comment and radius arithmetic
• call .compute() on dask results, following the suite convention
• add a numpy-only geodesic class (the module raises NotImplementedError for geodesic on cupy and dask)

Changing the source pattern will show up as a one-time step change in the asv history for time_surface_distance and friends. That is the point: the old numbers were measuring heap drain.

Found by /sweep-benchmarks.