Let the GPU surface-distance relaxation run to convergence

xrspatial/surface_distance.py

@@ -580,7 +580,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](
@@ -590,8 +598,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

@@ -543,6 +543,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."""