Let the GPU surface-distance relaxation run to convergence (#3732)

* Let the GPU surface-distance relaxation run to convergence (#3721)

_surface_distance_cupy capped its parallel relaxation at H + W passes.
One pass moves distance information about one pixel hop, so the number
of passes needed is the hop count of the longest shortest path. Those
two agree on open terrain, but NaN barriers make paths wind, and the
Bellman-Ford bound for a grid graph is the pixel count.

Hitting the cap exited the loop with no warning and no signal in the
output, so pixels the CPU reaches came back NaN, meaning unreachable.
On a 16x16 serpentine corridor the GPU found 37 of 136 reachable pixels;
at 24x24 it found 53 of 300.

The ceiling is now H * W and reaching it warns. The changed-flag check
still exits as soon as the solution settles, so open terrain runs the
same number of passes as before: 256x256 and 512x512 random elevation
both timed within noise of main (0.135 vs 0.136 s, 0.313 vs 0.326 s).

Also covers the bounded dask+cupy path, which runs this function per
chunk against the chunk's own H and W.

* Address review: cover DIRECTION in the long-path test, fix the warning stacklevel (#3721)

_surface_distance_cupy runs the same relaxation loop for all three
modes, and DIRECTION reads srow/scol, which a truncated loop leaves
unset just as it leaves dist at infinity. The long-path test now
parametrizes over surface_direction too.

The convergence warning used stacklevel=2, which lands on _compute.
_surface_distance_dask already uses 4 for the same call depth, so the
warning now points at the caller.

Author: brendancol

xrspatial/surface_distance.py

@@ -611,7 +611,15 @@ def _surface_distance_cupy(source_data, elev_data, cellsize_x, cellsize_y,
     changed = cp.zeros(1, dtype=cp.int32)
     griddim, blockdim = cuda_args((H, W))
 
-    max_iterations = H + W
+    # One pass moves distance information roughly one pixel hop, so the
+    # number of passes needed is the hop count of the longest shortest
+    # path.  On open terrain that is about H + W, but NaN barriers make
+    # paths wind, and the Bellman-Ford bound for a grid graph is the
+    # pixel count.  The `changed` check below exits as soon as the
+    # solution settles, so the ceiling only costs anything on the
+    # pathological inputs that need it.
+    max_iterations = H * W
+    converged = False
     for _ in range(max_iterations):
         changed[0] = 0
         _sd_relax_kernel[griddim, blockdim](
@@ -621,8 +629,19 @@ def _surface_distance_cupy(source_data, elev_data, cellsize_x, cellsize_y,
             np.float64(max_distance),
         )
         if int(changed[0]) == 0:
+            converged = True
             break
 
+    if not converged:
+        warnings.warn(
+            f"surface_distance: the GPU relaxation was still improving "
+            f"distances after {max_iterations} passes on a {H}x{W} raster "
+            f"and was stopped.  Some pixels may be reported as unreachable "
+            f"when a path exists.  Please report this raster upstream.",
+            UserWarning,
+            stacklevel=4,
+        )
+
     # Extract output
     if mode == DISTANCE:
         out = cp.where(cp.isinf(dist) | (dist > max_distance),

xrspatial/tests/test_surface_distance.py

@@ -748,6 +748,47 @@ def test_cupy_matches_numpy():
                                equal_nan=True)
 
 
+def _serpentine_elevation(n):
+    """NaN barriers everywhere except one winding open corridor.
+
+    The corridor from (0, 0) is about n*n/2 pixel hops long, far more than
+    the n+n the GPU relaxation used to allow itself.
+    """
+    elev = np.full((n, n), np.nan, dtype=np.float64)
+    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
+
+
+@pytest.mark.skipif(not has_cuda_and_cupy(), reason="cupy/cuda not available")
+@pytest.mark.parametrize(
+    "mode", [surface_distance, surface_allocation, surface_direction])
+def test_cupy_long_path_matches_numpy(mode):
+    """CuPy must relax until it converges, not for a fixed H+W (#3721).
+
+    A winding corridor needs far more relaxation passes than the raster is
+    wide plus tall. Stopping early makes reachable pixels come back NaN,
+    which reads as "no path exists".
+    """
+    n = 16
+    elev = _serpentine_elevation(n)
+    source = np.zeros((n, n), dtype=np.float64)
+    source[0, 0] = 1.0
+
+    np_result = _compute(mode(_make_raster(source), _make_raster(elev)))
+    cp_result = _compute(mode(_make_raster(source, backend='cupy'),
+                              _make_raster(elev, backend='cupy')))
+
+    # The corridor is genuinely reachable end to end on the CPU.
+    assert np.isfinite(np_result[n - 1, 0])
+    assert int(np.isfinite(np_result).sum()) > 4 * n
+
+    np.testing.assert_allclose(cp_result, np_result, rtol=1e-5,
+                               equal_nan=True)
+
+
 @pytest.mark.skipif(not has_cuda_and_cupy(), reason="cupy/cuda not available")
 def test_cupy_returns_cupy_array():
     """CuPy input should produce CuPy output."""