proximity brute-force kernel: hoist metric dispatch, branchless argmin, parallel rows

[brendancol]

Reason or Problem

_process_numpy_bruteforce in xrspatial/proximity.py is the CPU kernel behind allocation(), direction(), every GREAT_CIRCLE call, and proximity() when scipy is missing. It is also the per-chunk function for the dask+numpy versions of those calls. The default Euclidean/Manhattan proximity() goes through cKDTree and is not affected.

The kernel does more work per target pair than it needs to:

• The inner loop calls _distance, which branches on the metric for every pixel/target pair. For GREAT_CIRCLE that means four range checks plus a sqrt and an asin per pair.
• Target coordinates are read back through xs[target_rows[k], target_cols[k]] on every pair instead of being gathered once.
• The argmin is a data-dependent if that the branch predictor cannot learn.
• The outer loop is written with prange, but the kernel is decorated with @ngjit, which has no parallel=True. So it runs on one core.

On a 300x600 raster with 1000 random targets on a 20-core box:

case	now
EUCLIDEAN / PROXIMITY	492 ms
EUCLIDEAN / ALLOCATION	505 ms
GREAT_CIRCLE / PROXIMITY	3567 ms

Proposal

Restructure the kernel without changing any output.

Gather target coordinates into flat float64 arrays once before the pixel loop. Hoist the metric dispatch out of the pair loop so there is one inner loop per metric. Each loop compares a monotone proxy of the distance (squared distance for Euclidean, |dx|+|dy| for Manhattan, the haversine term for great circle) and only evaluates the sqrt/asin and the float32 rounding when the proxy beats the current best. The strict < still runs on the float32 distance, so the documented float32 tie-break (lowest flat index wins on a float32 tie, #3689) stays bitwise identical. Use a select instead of a branch for the argmin update.

Then turn on parallel=True so the existing prange over rows does something. The launch has to be serialized behind a module-level threading.Lock, same pattern as convolution.py and terrain.py: the dask path calls the kernel per chunk from worker threads and numba's workqueue layer is not threadsafe across host threads (#3141).

For GREAT_CIRCLE, validate the coordinate grids once before the loop and raise the same ValueError messages the per-pair guards in great_circle_distance raise today. The public euclidean_distance, manhattan_distance, great_circle_distance and _distance functions do not change.

Usage: no API change.

Value: a spike measured about 30x on the two Euclidean cases and 50x on great circle with default threads, and 2.4x/3.2x single-threaded, with every output bitwise identical to the current kernel.

Stakeholders and Impacts

Anyone calling allocation, direction, or any op with distance_metric='GREAT_CIRCLE' on numpy or dask+numpy input. The cupy kernel and the cKDTree paths are untouched.

Drawbacks

A parallel kernel compiles more slowly on first call. The lock means concurrent dask chunks run one at a time, each internally parallel. Convolution and terrain already made that trade.

Alternatives

Comparing the raw float64 proxy and skipping the float32 rounding would be a little faster, but it changes which target wins on float32 near-ties and breaks test_tie_break_float32_precision_nonlattice_grid.

Additional Notes or Context

The asv classes Proximity, Allocation, Direction in benchmarks/benchmarks/proximity.py already time this path.