From fc271121a963b8566f12f7511b0ef8d44657d759 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Fri, 4 Sep 2026 11:55:26 -0400 Subject: [PATCH 1/2] Speed up the proximity brute-force kernel (#3740) _process_numpy_bruteforce serves allocation(), direction(), every GREAT_CIRCLE call and proximity() without scipy, on numpy and per chunk on dask+numpy. It called _distance for every pixel/target pair, which branched on the metric and (for GREAT_CIRCLE) ran four range checks, a sqrt and an asin per pair, re-read the target coordinates through the row/col index arrays each time, and its prange over rows was serial because @ngjit has no parallel=True. Gather the target coordinates into flat arrays once, give each metric its own inner loop, and compile the pixel loop with parallel=True. The inner loops compare a monotone proxy of the distance (squared distance, |dx|+|dy|, the haversine term) and only take the sqrt/asin and the float32 rounding when the proxy beats the running best. The strict < still runs on the float32 distance, so the lowest-flat-index tie-break at float32 precision (#3689) is unchanged. GREAT_CIRCLE validates the coordinate grids once up front and raises the same messages the per-pair guards raised. The parallel kernel launch is serialized behind a module-level lock, same as convolution and terrain (#3141), because the dask path calls it from worker threads. 300x600 raster, 1000 random targets, 20-core host, median of 5: EUCLIDEAN/PROXIMITY 496 ms -> 87 ms (1 thread), 8.6 ms (20) EUCLIDEAN/ALLOCATION 501 ms -> 91 ms (1 thread), 10.2 ms (20) MANHATTAN/DIRECTION 488 ms -> 95 ms (1 thread), 15.5 ms (20) GREAT_CIRCLE/PROXIMITY 3563 ms -> 1136 ms (1 thread), 87.6 ms (20) Results are bit-identical to the previous kernel across all three metrics, all three modes, bounded and unbounded max_distance, explicit and default target_values, and NaN cells in the image. --- xrspatial/proximity.py | 244 ++++++++++++++++++++++++++---- xrspatial/tests/test_proximity.py | 108 +++++++++++++ 2 files changed, 322 insertions(+), 30 deletions(-) diff --git a/xrspatial/proximity.py b/xrspatial/proximity.py index bd54a5696..e86f6a689 100644 --- a/xrspatial/proximity.py +++ b/xrspatial/proximity.py @@ -1,3 +1,4 @@ +import threading import warnings from functools import partial @@ -15,7 +16,7 @@ import numpy as np import xarray as xr -from numba import cuda, prange +from numba import cuda, jit, prange try: import cupy @@ -493,33 +494,27 @@ def _is_target_value(v, target_values): return False -@ngjit -def _process_numpy_bruteforce( - img, xs, ys, target_values, max_distance, distance_metric, process_mode -): - """Exact nearest-target proximity / allocation / direction on the CPU. +# Numba parallel=True kernels must not be launched concurrently from multiple +# Python threads: the default 'workqueue' threading layer is not threadsafe and +# aborts the process (SIGABRT on macOS) when two host threads enter a parallel +# region at once. _process_dask maps _process_numpy over chunks under dask's +# threaded scheduler, and that reaches the brute-force kernel for GREAT_CIRCLE, +# ALLOCATION, DIRECTION and the no-scipy PROXIMITY fallback, so the kernel +# launch is serialized behind this lock. Same hazard and fix as the +# convolution, terrain and reproject kernels (#3141). +_PARALLEL_KERNEL_LOCK = threading.Lock() - For every pixel, scan all target pixels and keep the closest one under the - chosen distance metric. This is the same brute-force search the CUDA kernel - runs (see ``_proximity_cuda_kernel``). It covers what the cKDTree path - cannot: GREAT_CIRCLE (not a Minkowski metric), the tie-break-sensitive - ALLOCATION/DIRECTION modes, and PROXIMITY when scipy is missing. - ``xs`` and ``ys`` are the per-pixel 2D coordinate grids built by the caller. - """ +@ngjit +def _collect_targets(img, target_values): + """Row/col indices of every target pixel, in flat (row-major) order.""" height, width = img.shape - - # Collect target pixel rows/cols in flat arrays (two passes: count, fill). n_targets = 0 for line in range(height): for col in range(width): if _is_target_value(img[line, col], target_values): n_targets += 1 - output = np.full((height, width), np.nan, dtype=np.float32) - if n_targets == 0: - return output - target_rows = np.empty(n_targets, dtype=np.int64) target_cols = np.empty(n_targets, dtype=np.int64) t = 0 @@ -529,20 +524,165 @@ def _process_numpy_bruteforce( target_rows[t] = line target_cols[t] = col t += 1 + return target_rows, target_cols + + +# The three inner loops below each scan every target for one pixel and return +# (index, float32 distance) of the nearest one, or (-1, inf) with no targets. +# +# They compare a cheap proxy that is monotone in the distance (squared +# distance, |dx| + |dy|, the haversine term) and only evaluate the sqrt / +# arcsin and the float32 rounding when the proxy beats the running best. The +# strict ``<`` that decides the winner still runs on the float32 distance, the +# same value ``_distance`` returns, so the documented tie-break is unchanged: +# two targets whose float64 distances differ only past the float32 mantissa +# are a tie and the lowest flat index wins (issue #3689, and the matching +# comment in ``_proximity_cuda_kernel``). A candidate whose proxy does not +# beat the running best has a float32 distance >= the running best and could +# never have won under that rule, so skipping it changes nothing. + +@ngjit +def _nearest_euclidean(px, py, txs, tys): + best_proxy = np.inf + best_dist = np.float32(np.inf) + best_idx = -1 + for k in range(len(txs)): + dx = px - txs[k] + dy = py - tys[k] + proxy = dx * dx + dy * dy + if proxy < best_proxy: + d = np.float32(np.sqrt(proxy)) + better = d < best_dist + best_idx = k if better else best_idx + best_dist = d + best_proxy = proxy + return best_idx, best_dist + + +@ngjit +def _nearest_manhattan(px, py, txs, tys): + best_proxy = np.inf + best_dist = np.float32(np.inf) + best_idx = -1 + for k in range(len(txs)): + dx = px - txs[k] + dy = py - tys[k] + proxy = abs(dx) + abs(dy) + if proxy < best_proxy: + d = np.float32(proxy) + better = d < best_dist + best_idx = k if better else best_idx + best_dist = d + best_proxy = proxy + return best_idx, best_dist + + +@ngjit +def _nearest_great_circle(px, py, tlons, tlats, tcoslats): + # Same arithmetic, in the same order, as great_circle_distance with the + # default radius, so the float32 distance is bit-identical to _distance. + lat1 = np.radians(py) + lon1 = np.radians(px) + coslat1 = np.cos(lat1) + best_proxy = np.inf + best_dist = np.float32(np.inf) + best_idx = -1 + for k in range(len(tlons)): + dlon = tlons[k] - lon1 + dlat = tlats[k] - lat1 + proxy = np.sin(dlat / 2.0) ** 2 + \ + coslat1 * tcoslats[k] * np.sin(dlon / 2.0) ** 2 + if proxy < best_proxy: + d = np.float32(6378137 * 2 * np.arcsin(np.sqrt(proxy))) + better = d < best_dist + best_idx = k if better else best_idx + best_dist = d + best_proxy = proxy + return best_idx, best_dist + + +@ngjit +def _great_circle_target_terms(txs, tys): + # Precompute the per-target radians and cos(lat) with numba's np.radians / + # np.cos rather than numpy's, so they match what the per-pixel side of + # _nearest_great_circle (and great_circle_distance) computes bit for bit. + n = len(txs) + tlons = np.empty(n, dtype=np.float64) + tlats = np.empty(n, dtype=np.float64) + tcoslats = np.empty(n, dtype=np.float64) + for k in range(n): + tlons[k] = np.radians(txs[k]) + tlats[k] = np.radians(tys[k]) + tcoslats[k] = np.cos(tlats[k]) + return tlons, tlats, tcoslats + +@ngjit +def _great_circle_range_violation(xs, ys, txs, tys): + # Report the first out-of-range coordinate in the order the per-pair + # guards in great_circle_distance would have hit it when the pixel loop + # called it pair by pair: pixel (0, 0) against every target, then the + # remaining pixels. 0 = no violation, otherwise the guard number (1: x of + # the first point, 2: x of the second, 3: y of the first, 4: y of the + # second). NaN coordinates (dask halo padding) fail no comparison, as + # before. + px = xs[0, 0] + py = ys[0, 0] + if px > 180 or px < -180: + return 1 + if txs[0] > 180 or txs[0] < -180: + return 2 + if py > 90 or py < -90: + return 3 + if tys[0] > 90 or tys[0] < -90: + return 4 + for k in range(1, len(txs)): + if txs[k] > 180 or txs[k] < -180: + return 2 + if tys[k] > 90 or tys[k] < -90: + return 4 + height, width = xs.shape + for line in range(height): + for col in range(width): + if xs[line, col] > 180 or xs[line, col] < -180: + return 1 + if ys[line, col] > 90 or ys[line, col] < -90: + return 3 + return 0 + + +_GREAT_CIRCLE_RANGE_MESSAGES = { + 1: "Invalid x-coordinate of the first point." + "Must be in the range [-180, 180]", + 2: "Invalid x-coordinate of the second point." + "Must be in the range [-180, 180]", + 3: "Invalid y-coordinate of the first point." + "Must be in the range [-90, 90]", + 4: "Invalid y-coordinate of the second point." + "Must be in the range [-90, 90]", +} + + +@jit(nopython=True, nogil=True, parallel=True) +def _bruteforce_kernel( + img, xs, ys, target_rows, target_cols, txs, tys, tlons, tlats, tcoslats, + max_distance, distance_metric, process_mode, output +): + # The metric branch sits per pixel, outside the target loop, and the + # metric stays a runtime value so one compiled specialization serves all + # three. Rows are independent, so prange over them. + height, width = img.shape for line in prange(height): for col in range(width): px = xs[line, col] py = ys[line, col] - best_dist = np.float32(np.inf) - best_idx = -1 - for k in range(n_targets): - tx = xs[target_rows[k], target_cols[k]] - ty = ys[target_rows[k], target_cols[k]] - d = _distance(px, tx, py, ty, distance_metric) - if d < best_dist: - best_dist = d - best_idx = k + if distance_metric == EUCLIDEAN: + best_idx, best_dist = _nearest_euclidean(px, py, txs, tys) + elif distance_metric == GREAT_CIRCLE: + best_idx, best_dist = _nearest_great_circle( + px, py, tlons, tlats, tcoslats) + else: + best_idx, best_dist = _nearest_manhattan(px, py, txs, tys) if best_idx >= 0 and best_dist <= max_distance: if process_mode == PROXIMITY: output[line, col] = best_dist @@ -551,8 +691,52 @@ def _process_numpy_bruteforce( target_rows[best_idx], target_cols[best_idx]] else: output[line, col] = _calc_direction( - px, xs[target_rows[best_idx], target_cols[best_idx]], - py, ys[target_rows[best_idx], target_cols[best_idx]]) + px, txs[best_idx], py, tys[best_idx]) + + +def _process_numpy_bruteforce( + img, xs, ys, target_values, max_distance, distance_metric, process_mode +): + """Exact nearest-target proximity / allocation / direction on the CPU. + + For every pixel, scan all target pixels and keep the closest one under the + chosen distance metric. This is the same brute-force search the CUDA kernel + runs (see ``_proximity_cuda_kernel``). It covers what the cKDTree path + cannot: GREAT_CIRCLE (not a Minkowski metric), the tie-break-sensitive + ALLOCATION/DIRECTION modes, and PROXIMITY when scipy is missing. + + ``xs`` and ``ys`` are the per-pixel 2D coordinate grids built by the caller. + + The target coordinates are gathered into flat arrays once, the per-metric + inner loops live in ``_nearest_*`` and the pixel loop runs in parallel over + rows in ``_bruteforce_kernel``, serialized behind ``_PARALLEL_KERNEL_LOCK`` + because the dask path calls this per chunk from worker threads. + """ + target_rows, target_cols = _collect_targets(img, target_values) + + output = np.full(img.shape, np.nan, dtype=np.float32) + if len(target_rows) == 0: + return output + + txs = xs[target_rows, target_cols] + tys = ys[target_rows, target_cols] + + if distance_metric == GREAT_CIRCLE: + # The per-pair guards in great_circle_distance no longer run inside + # the loop; check the grids once up front and raise the same message. + violation = _great_circle_range_violation(xs, ys, txs, tys) + if violation: + raise ValueError(_GREAT_CIRCLE_RANGE_MESSAGES[violation]) + tlons, tlats, tcoslats = _great_circle_target_terms(txs, tys) + else: + tlons = tlats = tcoslats = np.empty(0, dtype=np.float64) + + with _PARALLEL_KERNEL_LOCK: + _bruteforce_kernel( + img, xs, ys, target_rows, target_cols, txs, tys, + tlons, tlats, tcoslats, + max_distance, distance_metric, process_mode, output, + ) return output diff --git a/xrspatial/tests/test_proximity.py b/xrspatial/tests/test_proximity.py index d065b4047..6b3ddc38e 100644 --- a/xrspatial/tests/test_proximity.py +++ b/xrspatial/tests/test_proximity.py @@ -729,6 +729,114 @@ def test_tie_break_float32_precision_nonlattice_grid( general_output_checks(raster, result, expected) +def _raster_on_backend(raster, backend, chunks=(2, 2)): + raster = raster.copy() + if has_cuda_and_cupy() and 'cupy' in backend: + import cupy + raster.data = cupy.asarray(raster.data) + if 'dask' in backend and da is not None: + raster.data = da.from_array(raster.data, chunks=chunks) + return raster + + +@pytest.fixture +def tie_break_proxy_gap_raster(): + # The brute-force kernel compares a cheap proxy (the squared distance) + # and only rounds to float32 when the proxy beats the running best. Here + # the proxies differ by a whole unit, 36000000 against 36000001, yet both + # distances round to float32(6000.0). The float64-closer target B must + # still lose to target A, the lower flat index (issue #3740). + data = np.zeros((3, 3), dtype=np.float64) + data[1, 2] = 1.0 # target A at (x=6000, y=1), flat index 5 + data[2, 1] = 2.0 # target B at (x=4800, y=3600), flat index 7 + raster = xr.DataArray(data, dims=['y', 'x']) + raster['x'] = np.array([0.0, 4800.0, 6000.0]) + raster['y'] = np.array([0.0, 1.0, 3600.0]) + return raster + + +@pytest.mark.parametrize("backend", ['numpy', 'dask+numpy', 'cupy', 'dask+cupy']) +@pytest.mark.parametrize("func", [allocation, direction]) +def test_tie_break_float32_beats_float64_proxy( + backend, func, tie_break_proxy_gap_raster): + numpy_raster = tie_break_proxy_gap_raster + d_a = np.sqrt(6000.0 * 6000.0 + 1.0 * 1.0) + d_b = np.sqrt(4800.0 * 4800.0 + 3600.0 * 3600.0) + assert d_b < d_a + assert np.float32(d_a) == np.float32(d_b) + + expected = func(numpy_raster).data + if func is allocation: + assert expected[0, 0] == 1.0 + else: + assert expected[0, 0] == _calc_direction(0.0, 6000.0, 0.0, 1.0) + + raster = _raster_on_backend(numpy_raster, backend) + result = func(raster) + general_output_checks(raster, result, expected) + + +@pytest.mark.parametrize("which, bad_x, bad_y, point", [ + ('pixel', -181.0, None, 'first'), + ('target', 181.0, None, 'second'), + ('pixel', None, 91.0, 'first'), + ('target', None, -91.0, 'second'), +]) +def test_bruteforce_great_circle_range_check_messages( + which, bad_x, bad_y, point): + # GREAT_CIRCLE validates the coordinate grids once before the pixel loop + # and must raise the same message the per-pair guard in + # great_circle_distance raises for that coordinate (issue #3740). + from xrspatial.proximity import GREAT_CIRCLE, PROXIMITY, _process_numpy_bruteforce + + img = np.zeros((2, 2)) + img[1, 1] = 1.0 + xs = np.tile(np.array([0.0, 1.0]), 2).reshape(2, 2) + ys = np.repeat(np.array([0.0, 1.0]), 2).reshape(2, 2) + row, col = (0, 0) if which == 'pixel' else (1, 1) + if bad_x is not None: + xs[row, col] = bad_x + kwargs = dict(x1=0.0, x2=0.0, y1=0.0, y2=0.0) + kwargs['x1' if point == 'first' else 'x2'] = bad_x + else: + ys[row, col] = bad_y + kwargs = dict(x1=0.0, x2=0.0, y1=0.0, y2=0.0) + kwargs['y1' if point == 'first' else 'y2'] = bad_y + with pytest.raises(ValueError) as reference: + great_circle_distance(**kwargs) + with pytest.raises(ValueError) as kernel: + _process_numpy_bruteforce( + img, xs, ys, np.array([]), np.float32(np.inf), + GREAT_CIRCLE, PROXIMITY) + assert str(kernel.value) == str(reference.value) + assert point in str(kernel.value) + + +def test_bruteforce_kernel_compiled_parallel(): + # prange over rows is a plain serial loop unless the kernel is compiled + # with parallel=True; guard against that regressing (issue #3740). + from xrspatial.proximity import _bruteforce_kernel + assert _bruteforce_kernel.targetoptions.get('parallel') is True + + +def test_bruteforce_concurrent_launches_match_serial(): + # The parallel kernel is launched per chunk from dask worker threads and + # is serialized behind _PARALLEL_KERNEL_LOCK; hammer it from a thread + # pool and check every caller gets the single-threaded answer. + from concurrent.futures import ThreadPoolExecutor + + raster = _simple_raster() + expected = allocation(raster, x='lon', y='lat').data + + def one(_): + return allocation(raster, x='lon', y='lat').data + + with ThreadPoolExecutor(max_workers=8) as pool: + results = list(pool.map(one, range(32))) + for result in results: + np.testing.assert_array_equal(result, expected) + + @pytest.mark.skipif(da is None, reason="dask is not installed") def test_proximity_dask_kdtree_no_targets(): """No target pixels found → result is all NaN.""" From 27d67714d72e6e5e7f07340c73b014c5b4bdc34f Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Fri, 4 Sep 2026 12:01:00 -0400 Subject: [PATCH 2/2] Address review: select on the running best, larger concurrency raster (#3740) Update best_dist only when the candidate wins so a NaN rounded distance (a haversine term a hair past 1.0) is skipped instead of poisoning every later comparison, as the old kernel did. Comment the placeholder great circle arrays on the non-great-circle call path. Hammer the lock test with a 40x40 raster so the concurrent launches overlap. --- xrspatial/proximity.py | 10 +++++++--- xrspatial/tests/test_proximity.py | 11 ++++++++--- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/xrspatial/proximity.py b/xrspatial/proximity.py index e86f6a689..d8d178596 100644 --- a/xrspatial/proximity.py +++ b/xrspatial/proximity.py @@ -540,6 +540,9 @@ def _collect_targets(img, target_values): # comment in ``_proximity_cuda_kernel``). A candidate whose proxy does not # beat the running best has a float32 distance >= the running best and could # never have won under that rule, so skipping it changes nothing. +# The running best is updated with a select rather than a branch, and only +# when the candidate wins, so a NaN distance (say a haversine term rounded a +# hair past 1.0) is skipped instead of poisoning every later comparison. @ngjit def _nearest_euclidean(px, py, txs, tys): @@ -554,7 +557,7 @@ def _nearest_euclidean(px, py, txs, tys): d = np.float32(np.sqrt(proxy)) better = d < best_dist best_idx = k if better else best_idx - best_dist = d + best_dist = d if better else best_dist best_proxy = proxy return best_idx, best_dist @@ -572,7 +575,7 @@ def _nearest_manhattan(px, py, txs, tys): d = np.float32(proxy) better = d < best_dist best_idx = k if better else best_idx - best_dist = d + best_dist = d if better else best_dist best_proxy = proxy return best_idx, best_dist @@ -596,7 +599,7 @@ def _nearest_great_circle(px, py, tlons, tlats, tcoslats): d = np.float32(6378137 * 2 * np.arcsin(np.sqrt(proxy))) better = d < best_dist best_idx = k if better else best_idx - best_dist = d + best_dist = d if better else best_dist best_proxy = proxy return best_idx, best_dist @@ -729,6 +732,7 @@ def _process_numpy_bruteforce( raise ValueError(_GREAT_CIRCLE_RANGE_MESSAGES[violation]) tlons, tlats, tcoslats = _great_circle_target_terms(txs, tys) else: + # Placeholders: the kernel only reads these under GREAT_CIRCLE. tlons = tlats = tcoslats = np.empty(0, dtype=np.float64) with _PARALLEL_KERNEL_LOCK: diff --git a/xrspatial/tests/test_proximity.py b/xrspatial/tests/test_proximity.py index 6b3ddc38e..b83a0f1e7 100644 --- a/xrspatial/tests/test_proximity.py +++ b/xrspatial/tests/test_proximity.py @@ -825,11 +825,16 @@ def test_bruteforce_concurrent_launches_match_serial(): # pool and check every caller gets the single-threaded answer. from concurrent.futures import ThreadPoolExecutor - raster = _simple_raster() - expected = allocation(raster, x='lon', y='lat').data + # Use a raster large enough that the launches actually overlap. + data = np.zeros((40, 40), dtype=np.float64) + rng = np.random.default_rng(3740) + data.flat[rng.choice(data.size, 50, replace=False)] = rng.integers( + 1, 6, 50) + raster = create_test_raster(data, backend='numpy') + expected = allocation(raster).data def one(_): - return allocation(raster, x='lon', y='lat').data + return allocation(raster).data with ThreadPoolExecutor(max_workers=8) as pool: results = list(pool.map(one, range(32)))