From 12b145c12becc53931fc6343495d9782ecb1eaa4 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Sat, 27 Jun 2026 07:58:44 -0700 Subject: [PATCH 1/2] Fix dask EUCLIDEAN/MANHATTAN proximity NaN at float32-ulp max_distance (#3392) The cKDTree backends (eager numpy PROXIMITY and the dask paths for every mode) compared the float64 distance against max_distance, while the brute-force and CUDA kernels round the distance to float32 first. A max_distance landing inside the float32 ulp gap just below a target's distance dropped that target to NaN on the cKDTree paths but kept it on the brute-force backends. _inclusive_upper_bound now widens the bound to the midpoint between the largest float32 <= max_distance and the next float32 up, matching the kernels' float32 keep test, while preserving the absolute floor that keeps the p=2 boundary inclusive at small max_distance. --- xrspatial/proximity.py | 58 ++++++++++++++---------- xrspatial/tests/test_proximity.py | 75 +++++++++++++++++++++++++++++++ 2 files changed, 110 insertions(+), 23 deletions(-) diff --git a/xrspatial/proximity.py b/xrspatial/proximity.py index 7ac0d24c5..c560b776d 100644 --- a/xrspatial/proximity.py +++ b/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, diff --git a/xrspatial/tests/test_proximity.py b/xrspatial/tests/test_proximity.py index 537612429..687884c24 100644 --- a/xrspatial/tests/test_proximity.py +++ b/xrspatial/tests/test_proximity.py @@ -2636,6 +2636,81 @@ 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. +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(da is None, reason="dask is not installed") +@pytest.mark.parametrize("metric", ['EUCLIDEAN', 'MANHATTAN']) +@pytest.mark.parametrize("func", [proximity, allocation, direction]) +def test_float32_ulp_max_distance_dask_matches_numpy(func, metric): + # 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 the eager numpy and dask+numpy backends. + 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 + p = 2 if metric == 'EUCLIDEAN' else 1 + d64 = float(np.linalg.norm([lon[j] - lon[0]], ord=p)) + 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) + numpy_raster = xr.DataArray( + data, dims=['lat', 'lon'], coords={'lon': lon, 'lat': lat}) + numpy_val = func(numpy_raster, **kwargs).data[0, j] + + dask_raster = xr.DataArray( + data.copy(), dims=['lat', 'lon'], coords={'lon': lon, 'lat': lat}) + dask_raster.data = da.from_array(dask_raster.data, chunks=(1, 2)) + dask_val = func(dask_raster, **kwargs).data.compute()[0, j] + + # The target qualifies (its float32 distance is <= max_distance), so + # neither backend NaNs it out and both return the same value. + assert not np.isnan(numpy_val) + assert np.isnan(numpy_val) == np.isnan(dask_val) + np.testing.assert_allclose(numpy_val, dask_val, rtol=1e-6) + + +@pytest.mark.skipif(da is None, reason="dask is not installed") +@pytest.mark.parametrize("func", [proximity, allocation, direction]) +def test_float32_ulp_max_distance_just_below_excludes(func): + # 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. + 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) + numpy_raster = xr.DataArray( + data, dims=['lat', 'lon'], coords={'lon': lon, 'lat': lat}) + numpy_val = func(numpy_raster, **kwargs).data[0, j] + + dask_raster = xr.DataArray( + data.copy(), dims=['lat', 'lon'], coords={'lon': lon, 'lat': lat}) + dask_raster.data = da.from_array(dask_raster.data, chunks=(1, 2)) + dask_val = func(dask_raster, **kwargs).data.compute()[0, j] + + assert np.isnan(numpy_val) + assert np.isnan(dask_val) + + # --------------------------------------------------------------------------- # Regression tests for issue #3442: EUCLIDEAN proximity/allocation/direction # dropped a target sitting exactly at max_distance to NaN on the numpy and From 917dc906fa7d2bc8bf7305ffed62eba9944c6a2d Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Sat, 27 Jun 2026 08:03:14 -0700 Subject: [PATCH 2/2] Address review: cover all four backends at the float32-ulp boundary (#3392) Parametrize the new keep/exclude tests across numpy, dask+numpy, cupy, and dask+cupy so the brute-force GPU backends are pinned to the cKDTree backends at the boundary and cannot drift apart. Replace the one-element np.linalg.norm with a plain longitude difference. --- xrspatial/tests/test_proximity.py | 92 ++++++++++++++++++++----------- 1 file changed, 61 insertions(+), 31 deletions(-) diff --git a/xrspatial/tests/test_proximity.py b/xrspatial/tests/test_proximity.py index 687884c24..3bb31502f 100644 --- a/xrspatial/tests/test_proximity.py +++ b/xrspatial/tests/test_proximity.py @@ -2647,47 +2647,86 @@ def test_max_distance_float32_boundary_allocation_gpu_matches_cpu(): # --------------------------------------------------------------------------- -@pytest.mark.skipif(da is None, reason="dask is not installed") +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_dask_matches_numpy(func, metric): +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 the eager numpy and dask+numpy backends. + # 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 - p = 2 if metric == 'EUCLIDEAN' else 1 - d64 = float(np.linalg.norm([lon[j] - lon[0]], ord=p)) + # 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) - numpy_raster = xr.DataArray( - data, dims=['lat', 'lon'], coords={'lon': lon, 'lat': lat}) - numpy_val = func(numpy_raster, **kwargs).data[0, j] + ref = func(_ulp_raster('numpy', data, lon, lat), **kwargs).data[0, j] + val = _ulp_value(func(_ulp_raster(backend, data, lon, lat), **kwargs), j) - dask_raster = xr.DataArray( - data.copy(), dims=['lat', 'lon'], coords={'lon': lon, 'lat': lat}) - dask_raster.data = da.from_array(dask_raster.data, chunks=(1, 2)) - dask_val = func(dask_raster, **kwargs).data.compute()[0, 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) - # The target qualifies (its float32 distance is <= max_distance), so - # neither backend NaNs it out and both return the same value. - assert not np.isnan(numpy_val) - assert np.isnan(numpy_val) == np.isnan(dask_val) - np.testing.assert_allclose(numpy_val, dask_val, rtol=1e-6) - -@pytest.mark.skipif(da is None, reason="dask is not installed") +@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): +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) @@ -2698,17 +2737,8 @@ def test_float32_ulp_max_distance_just_below_excludes(func): assert md < float(d32) kwargs = dict(x='lon', y='lat', max_distance=md) - numpy_raster = xr.DataArray( - data, dims=['lat', 'lon'], coords={'lon': lon, 'lat': lat}) - numpy_val = func(numpy_raster, **kwargs).data[0, j] - - dask_raster = xr.DataArray( - data.copy(), dims=['lat', 'lon'], coords={'lon': lon, 'lat': lat}) - dask_raster.data = da.from_array(dask_raster.data, chunks=(1, 2)) - dask_val = func(dask_raster, **kwargs).data.compute()[0, j] - - assert np.isnan(numpy_val) - assert np.isnan(dask_val) + val = _ulp_value(func(_ulp_raster(backend, data, lon, lat), **kwargs), j) + assert np.isnan(val) # ---------------------------------------------------------------------------