Match cKDTree max_distance check to the brute-force float32 convention (#3392)

xrspatial/proximity.py

@@ -750,33 +750,45 @@ def _chunk_func(data_chunk, xs_chunk, ys_chunk):
 
 
 def _inclusive_upper_bound(max_distance):
-    """Exclusive ``distance_upper_bound`` that cKDTree treats as inclusive.
-
-    ``cKDTree.query``'s ``distance_upper_bound`` is exclusive, while the
-    brute-force and CUDA kernels keep ``dist <= max_distance``. To turn the
-    exclusive check into the inclusive one the bound has to be widened just
-    past ``max_distance``.
-
-    A one-ulp widen (``np.nextafter``) is not enough for the EUCLIDEAN query
-    (``p=2``): cKDTree compares *squared* distances internally, and the square
-    of ``nextafter(0.0, inf)`` (the smallest subnormal) underflows back to
-    0.0, so the bound collapses to exactly ``max_distance`` again and a
-    target sitting exactly at the bound -- most visibly a target pixel itself
-    at ``max_distance=0`` -- is dropped to NaN. numpy/dask then disagreed with
-    the cupy/dask+cupy brute-force backends, which include it.
-
-    Widen by a relative epsilon with an absolute floor instead. The bump stays
-    large enough to survive squaring for ``p=2`` (and is harmless for ``p=1``)
-    while staying far below the next representable distance, so nothing beyond
-    ``max_distance`` is pulled in. Because the bump is relative, a target at a
-    true distance within roughly ``8 * eps * max_distance`` of the bound could
-    qualify, but at that magnitude the two distances are not distinguishable in
-    float64 anyway, so the result is still correct.
+    """Exclusive cKDTree ``distance_upper_bound`` matching the float32 keep test.
+
+    The brute-force and CUDA kernels round each candidate distance to float32
+    and keep it when ``np.float32(dist) <= max_distance`` (see
+    ``_process_numpy_bruteforce`` and ``_proximity_cuda_kernel``). cKDTree
+    instead compares *float64* distances against an *exclusive* bound, so to
+    reproduce the kernels' keep decision the bound has to account for both
+    differences.
+
+    * Float32 rounding: a target whose true float64 distance rounds *down*
+      across a float32 step to a value ``<= max_distance`` is kept by the
+      kernels but, against a float64-exact bound, dropped by cKDTree to NaN.
+      This surfaces when ``max_distance`` sits in the float32 ulp gap just
+      below a target's distance (issue #3392). The bound is therefore widened
+      to the midpoint between the largest float32 ``<= max_distance`` and the
+      next float32 up: the supremum of float64 distances that still round to a
+      float32 value ``<= max_distance``. Nothing beyond that midpoint is pulled
+      in, so no target the kernels exclude is admitted.
+
+    * Inclusive boundary for ``p=2``: cKDTree compares *squared* distances
+      internally, and the square of ``nextafter(0.0, inf)`` (the smallest
+      subnormal) underflows back to 0.0, collapsing the bound to exactly
+      ``max_distance`` and dropping a target sitting on the bound -- most
+      visibly a target pixel itself at ``max_distance=0``. A relative bump with
+      an absolute floor stays large enough to survive squaring for ``p=2`` (and
+      is harmless for ``p=1``). At small ``max_distance`` the float32 midpoint
+      underflows toward 0, so this floor is what keeps the boundary inclusive.
     """
     if not np.isfinite(max_distance):
         return np.inf
+    # Largest float32 value <= max_distance (np.float32 rounds to nearest, so it
+    # can land just above; step down one float32 ulp when it does).
+    f = np.float32(max_distance)
+    if f > max_distance:
+        f = np.nextafter(f, np.float32(-np.inf))
+    nf = np.nextafter(f, np.float32(np.inf))
+    f32_bound = (float(f) + float(nf)) / 2.0
     bump = 4.0 * np.finfo(np.float64).eps * max(abs(max_distance), 1.0)
-    return max_distance + bump
+    return max(f32_bound, max_distance + bump)
 
 
 def _process_numpy_kdtree(img, xs, ys, target_values, max_distance, p,

xrspatial/tests/test_proximity.py

@@ -2636,6 +2636,111 @@ def test_max_distance_float32_boundary_allocation_gpu_matches_cpu():
     np.testing.assert_allclose(numpy_val, cupy_val, rtol=1e-6)
 
 
+# ---------------------------------------------------------------------------
+# Issue #3392: split off from #3389. The cKDTree backends (eager numpy
+# PROXIMITY, and the dask paths for every mode) compared the float64 distance
+# against max_distance, while the brute-force/CUDA kernels round the distance
+# to float32 first. A max_distance landing inside the float32 ulp gap just
+# below a target's distance then dropped the target to NaN on the cKDTree
+# backends while the brute-force backends kept it. _inclusive_upper_bound now
+# widens the bound by half a float32 ulp so the cKDTree decision matches.
+# ---------------------------------------------------------------------------
+
+
+def _ulp_raster(backend, data, lon, lat):
+    """Build a (1, 3) raster on *backend* with explicit lon/lat coords.
+
+    create_test_raster lays its own coords, but these tests need the exact
+    0.0/0.1/0.7 longitudes that produce a float32 ulp gap, so build the
+    backend arrays directly.
+    """
+    raster = xr.DataArray(
+        data.copy(), dims=['lat', 'lon'], coords={'lon': lon, 'lat': lat})
+    if backend == 'numpy':
+        return raster
+    if backend == 'dask+numpy':
+        raster.data = da.from_array(data.copy(), chunks=(1, 2))
+        return raster
+    import cupy
+    if backend == 'cupy':
+        raster.data = cupy.asarray(data)
+        return raster
+    raster.data = da.from_array(cupy.asarray(data), chunks=(1, 2))  # dask+cupy
+    return raster
+
+
+def _ulp_value(agg, j):
+    out = agg.data
+    if da is not None and isinstance(out, da.Array):
+        out = out.compute()
+    if hasattr(out, 'get'):  # cupy
+        out = out.get()
+    return out[0, j]
+
+
+@pytest.mark.parametrize(
+    "backend", ['numpy', 'dask+numpy', 'cupy', 'dask+cupy'])
+@pytest.mark.parametrize("metric", ['EUCLIDEAN', 'MANHATTAN'])
+@pytest.mark.parametrize("func", [proximity, allocation, direction])
+def test_float32_ulp_max_distance_backends_keep_target(func, metric, backend):
+    # A target whose true float64 distance sits one float32-ulp above
+    # max_distance, but whose float32-rounded distance is <= max_distance, must
+    # be kept identically on every backend. The eager numpy result is the
+    # reference; each backend has to agree with it.
+    if 'cupy' in backend and not has_cuda_and_cupy():
+        pytest.skip("Requires CUDA and CuPy")
+    if 'dask' in backend and da is None:
+        pytest.skip("dask is not installed")
+
+    lon = np.array([0.0, 0.1, 0.7])
+    lat = np.array([0.0])
+    data = np.zeros((1, 3), dtype=np.float64)
+    data[0, 0] = 9.0
+    j = 2
+    # lat is constant, so the target is one axis away: the distance is just the
+    # longitude separation for both EUCLIDEAN (p=2) and MANHATTAN (p=1).
+    d64 = abs(float(lon[j] - lon[0]))
+    md = (float(np.float32(d64)) + d64) / 2.0
+    assert float(np.float32(d64)) < md < d64, "needs a real float32 ulp gap"
+
+    kwargs = dict(x='lon', y='lat', max_distance=md, distance_metric=metric)
+    ref = func(_ulp_raster('numpy', data, lon, lat), **kwargs).data[0, j]
+    val = _ulp_value(func(_ulp_raster(backend, data, lon, lat), **kwargs), j)
+
+    # The target qualifies (its float32 distance is <= max_distance), so no
+    # backend NaNs it out and all return the same value.
+    assert not np.isnan(ref)
+    assert np.isnan(ref) == np.isnan(val)
+    np.testing.assert_allclose(ref, val, rtol=1e-6)
+
+
+@pytest.mark.parametrize(
+    "backend", ['numpy', 'dask+numpy', 'cupy', 'dask+cupy'])
+@pytest.mark.parametrize("func", [proximity, allocation, direction])
+def test_float32_ulp_max_distance_just_below_excludes(func, backend):
+    # The mirror case: when max_distance is below the target's float32-rounded
+    # distance, the target is genuinely out of range and every backend must
+    # NaN it out. Guards against the bound widening so far it admits targets
+    # the brute-force kernels exclude.
+    if 'cupy' in backend and not has_cuda_and_cupy():
+        pytest.skip("Requires CUDA and CuPy")
+    if 'dask' in backend and da is None:
+        pytest.skip("dask is not installed")
+
+    lon = np.array([0.0, 0.1, 0.7])
+    lat = np.array([0.0])
+    data = np.zeros((1, 3), dtype=np.float64)
+    data[0, 0] = 9.0
+    j = 2
+    d32 = np.float32(0.7)
+    md = float(np.nextafter(d32, np.float32(-np.inf)))  # one float32 ulp below
+    assert md < float(d32)
+
+    kwargs = dict(x='lon', y='lat', max_distance=md)
+    val = _ulp_value(func(_ulp_raster(backend, data, lon, lat), **kwargs), j)
+    assert np.isnan(val)
+
+
 # ---------------------------------------------------------------------------
 # Regression tests for issue #3442: EUCLIDEAN proximity/allocation/direction
 # dropped a target sitting exactly at max_distance to NaN on the numpy and