cupy surface_distance stops relaxing after H+W passes and reports reachable pixels as NaN

[brendancol]

Description

The cupy backend of surface_distance() solves the shortest-path problem by
repeated parallel relaxation (_sd_relax_kernel), and caps the number of
passes at H + W:

# xrspatial/surface_distance.py:583
max_iterations = H + W
for _ in range(max_iterations):
    changed[0] = 0
    _sd_relax_kernel[griddim, blockdim](...)
    if int(changed[0]) == 0:
        break

One pass propagates distance information roughly one pixel hop. The number of
passes needed is therefore the hop count of the longest shortest path, not the
raster's width plus height. Those two numbers agree on open terrain, but they
come apart as soon as NaN barriers force paths to wind, and the worst case for
a grid graph is H * W, not H + W.

When the cap is reached the loop exits with no warning and no signal in the
output. Pixels the CPU backend reaches are reported as NaN, meaning
unreachable. A user reading the result cannot tell the difference between "no
path exists" and "we stopped looking."

This affects all three public functions on the cupy backend
(surface_distance, surface_allocation, surface_direction) and the
bounded dask+cupy path, where _surface_distance_cupy runs per chunk and the
cap is the chunk's H + W rather than the raster's.

Reproduction

Ran on this host with CUDA available.

import numpy as np, xarray as xr, cupy
from xrspatial.surface_distance import surface_distance

def serpentine(n):
    """NaN everywhere except one winding open corridor from (0, 0)."""
    elev = np.full((n, n), np.nan)
    for r in range(0, n, 2):
        elev[r, :] = 0.0
    for r in range(1, n, 2):
        elev[r, n - 1 if (r // 2) % 2 == 0 else 0] = 0.0
    return elev

def make(data, gpu=False):
    a = xr.DataArray(
        data.astype(float), dims=['y', 'x'],
        coords={'y': np.arange(data.shape[0], dtype=float),
                'x': np.arange(data.shape[1], dtype=float)},
        attrs={'res': (1.0, 1.0)})
    if gpu:
        a.data = cupy.asarray(a.data)
    return a

for N in (8, 12, 16, 24):
    elev = serpentine(N)
    src = np.zeros((N, N)); src[0, 0] = 1.0
    cpu = np.asarray(surface_distance(make(src), make(elev)).data)
    gpu = surface_distance(make(src, True), make(elev, True)).data.get()
    print(f"N={N:3d} cap={2*N:3d}  reachable cpu={int(np.isfinite(cpu).sum()):4d}"
          f" gpu={int(np.isfinite(gpu).sum()):4d}"
          f"  wrongly NaN on gpu={int((np.isfinite(cpu) & ~np.isfinite(gpu)).sum()):4d}")

N=  8 cap= 16  reachable cpu=  36 gpu=  21  wrongly NaN on gpu=  15
N= 12 cap= 24  reachable cpu=  78 gpu=  29  wrongly NaN on gpu=  49
N= 16 cap= 32  reachable cpu= 136 gpu=  37  wrongly NaN on gpu=  99
N= 24 cap= 48  reachable cpu= 300 gpu=  53  wrongly NaN on gpu= 247

At N=16 the GPU finds 37 of the 136 reachable pixels and calls the other 99
unreachable. The last corridor row is the clearest case:

row 15, numpy: [126.21   nan   nan ...  nan]
row 15, cupy : [   nan   nan   nan ...  nan]

Distances that both backends do report agree exactly (max abs diff 0.0), so
this is purely truncated propagation rather than arithmetic drift.

A serpentine grid is an easy way to make the failure obvious, but nothing
about it is special. Any NaN-masked terrain whose shortest paths wind more
than H + W hops (a river network, a canyon system, a coastline with inlets)
hits the same cap.

The numpy and dask+numpy paths run an exact Dijkstra and are correct here. I
also checked the dask iterative sweep cap (max(n_tile_y, n_tile_x) + 10)
against the same mazes at several chunk sizes and could not make it produce a
wrong answer, so this report is about the GPU cap only.

Expected behaviour

The cupy backend should relax until nothing changes, and its result should
match numpy. If some bound must remain as a safety net, it should be the
Bellman-Ford bound for the graph (H * W) rather than H + W, and reaching
it should warn rather than pass off a truncated answer as a finished one.

The early-exit check on changed already means well-behaved terrain costs the
same number of passes as it does today, so raising the ceiling is not a
throughput change for normal input.

Environment

xarray-spatial main @ 069fc13, Python 3.14, cupy with CUDA available.

[brendancol]

Independent confirmation from the reference-validation sweep, plus a
second failure mode that I think raises the severity here.

I validated surface_distance against scipy.sparse.csgraph.dijkstra
(scipy 1.16.1) run over an explicitly built grid graph with
sqrt(horizontal**2 + dz**2) edge weights. On the serpentine-barrier
scene the numpy backend matches the reference bit for bit and cupy does
not, exactly as described above. Driving _sd_relax_kernel by hand on a
32x32 raster, convergence took 342 passes against a budget of
H + W = 64; at 342 the cupy result matched scipy exactly
(max_abs = 0).

The part worth adding: NaN barriers are not required, and the output is
not always NaN. A raster with no impassable pixels at all, only steep
ridges, makes the truncated relaxation return finite distances that are
silently far too large:

n = 40
elev = np.zeros((n, n))
for i, r in enumerate(range(3, n - 1, 3)):
    elev[r, :] = 3000.0                          # tall ridge, still passable
    elev[r, -1 if i % 2 == 0 else 0] = 0.0       # a cheap way round the end
src = np.zeros((n, n)); src[0, 0] = 1.0

Crossing a ridge is short in hops but expensive; going around the end is
cheap but many hops long. Within the H + W budget the GPU only ever
finds the expensive crossings.

numpy vs scipy: MATCH   max_abs=0        max_rel=0
cupy  vs scipy: DIVERGE max_abs=23724.5  max_rel=1.74
finite cupy pixels over-estimated: 1314 of 1600
worst over-estimate: 23724.5 m (173.8% too large)
cupy NaN where scipy finite: 0 px

So the failure is not confined to barrier mazes and does not always
announce itself as NaN. Here the cupy backend returns a fully finite,
plausible-looking distance field in which 82% of pixels are wrong, some by
a factor of 2.7. Nothing downstream can detect it. That reads as CRITICAL
by the severity rubric rather than HIGH.

Whatever bound replaces H + W, it would be worth warning (or raising)
when the loop exits on the cap instead of on changed == 0, so a
truncated result is never returned silently.