Speed up the proximity brute-force kernel (#3740)

[brendancol]

Closes #3740

_process_numpy_bruteforce is the CPU kernel behind 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, so the metric branch ran per pair and great circle paid four range checks, a sqrt and an asin per pair. It re-read target coordinates through the row/col index arrays each time. Its prange over rows was serial because @ngjit has no parallel=True.

• Gather target coordinates into flat arrays once, give each metric its own inner loop (_nearest_euclidean, _nearest_manhattan, _nearest_great_circle), and compile the pixel loop with parallel=True.
• Each inner loop compares a monotone proxy of the distance (squared distance, |dx|+|dy|, the haversine term) and only takes 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 (allocation/direction: nearest-target tie evaluated at different float precision per backend, diverging on non-lattice grids #3689) is unchanged. A candidate whose proxy does not beat the running best has a float32 distance at or above it and could never have won.
• GREAT_CIRCLE validates the coordinate grids once before the loop and raises the same messages the per-pair guards in great_circle_distance raise, in the order they would have fired.
• The parallel launch is serialized behind a module-level _PARALLEL_KERNEL_LOCK, the same pattern as convolution.py and terrain.py (Streaming reproject thread pool aborts the process when numba parallel kernels run concurrently #3141), since _process_dask maps _process_numpy over chunks from worker threads.

The public euclidean_distance, manhattan_distance, great_circle_distance and _distance functions, the cupy kernel and the cKDTree paths are untouched.

Timings

300x600 raster, 1000 random targets, 20-core host, time.perf_counter, median of 5 after one warmup:

case	before	after, 1 thread	after, 20 threads
EUCLIDEAN / PROXIMITY	496 ms	87 ms (5.7x)	8.6 ms (58x)
EUCLIDEAN / ALLOCATION	501 ms	91 ms (5.5x)	10.2 ms (49x)
MANHATTAN / DIRECTION	488 ms	95 ms (5.1x)	15.5 ms (31x)
GREAT_CIRCLE / PROXIMITY	3563 ms	1136 ms (3.1x)	87.6 ms (41x)

Results are bit-identical to the kernel on main (np.array_equal(old, new, equal_nan=True)) on that raster across all 54 combinations of the three metrics, the three modes, max_distance at inf and at a finite value that leaves some pixels NaN, explicit and default target_values, and NaN cells in the image. The great-circle grid used lon in [-10, 10] and lat in [40, 45]. The single-thread gain is larger than the spike in #3740 predicted because the proxy comparison also skips the float32 rounding on most pairs.

The argmin is not a pure select on the proxy, which differs from the sketch in #3740. A raw float64 proxy comparison picks the float64-closer target on float32 near-ties, which breaks test_tie_break_float32_precision_nonlattice_grid and the CPU/GPU tie parity the CUDA kernel was aligned to. Gating the float32 rounding on the proxy keeps the old ordering exactly and only pays for it on record-breaking candidates.

Backends: numpy and dask+numpy run the new kernel. cupy and dask+cupy are unchanged.

Tests

• pytest xrspatial/tests/test_proximity.py (609 passed, including cupy and dask+cupy on this box)
• test_dask_task_names.py, test_accessor.py, test_dataset_support.py, test_balanced_allocation.py (178 passed)
• New: float32 tie where the squared distances differ by a whole unit (36000000 vs 36000001) but both round to float32(6000.0), on all four backends for allocation and direction
• New: the four GREAT_CIRCLE range messages match great_circle_distance byte for byte
• New: _bruteforce_kernel is compiled with parallel=True
• New: 32 concurrent allocation calls from 8 threads match the serial result

[brendancol]

PR Review: Speed up the proximity brute-force kernel (#3740)

Blockers (must fix before merge)

• none

Suggestions (should fix, not blocking)

• xrspatial/proximity.py:557, :575, :599: each inner loop assigns best_dist = d unconditionally after the proxy test, relying on d <= best_dist from monotonicity. That holds as long as d is a number. If the rounded distance ever comes out NaN (for great circle, a haversine term a hair above 1.0 sends arcsin(sqrt(a)) to NaN), best_dist becomes NaN and every later d < best_dist is False, so no later target can win. The old kernel skipped a NaN candidate and carried on. I could not construct an input that pushes a past 1.0 in float64, so this may be unreachable, but best_dist = d if better else best_dist costs nothing and removes the assumption. The proxy update can stay unconditional.

Nits (optional improvements)

• xrspatial/tests/test_proximity.py:822: the concurrency test hammers a 3x3 raster, so each launch finishes in microseconds and the eight threads rarely overlap inside the kernel. A raster on the order of the 40x40 tie-break fixture would make the launches actually contend for the lock.
• xrspatial/proximity.py:667: the kernel takes tlons, tlats, tcoslats as empty arrays for the Euclidean and Manhattan metrics. A one-line comment on the call site saying they are placeholders would save the next reader from looking for where they are filled.

What looks good

• The proxy-then-round argmin keeps the float32 tie-break exactly and the PR proves it with a case where the squared distances differ by a whole unit (36000000 vs 36000001) yet both round to float32(6000.0). The A/B against main is bitwise over 54 metric/mode/range/target/NaN combinations.
• _nearest_great_circle reproduces the arithmetic of great_circle_distance in the same evaluation order, and the target terms are precomputed with numba's np.radians/np.cos rather than numpy's, so no ulp drift between the two sides.
• The range check reproduces the order the per-pair guards fired in and the messages are compared byte for byte against great_circle_distance.
• The lock follows the convolution/terrain pattern and the dask path (_process_dask mapping _process_numpy) is covered because the lock lives inside _process_numpy_bruteforce.
• test_bruteforce_kernel_compiled_parallel guards the exact regression the issue describes (prange without parallel=True).
• Benchmarks Proximity, Allocation, Direction already cover this path; the performance label is on the PR.

Checklist

• Algorithm matches reference (bitwise A/B against the previous kernel)
• All implemented backends produce consistent results (numpy, dask+numpy, cupy, dask+cupy pass locally)
• NaN handling is correct (NaN image cells and NaN halo coordinates behave as before; see the suggestion on a NaN rounded distance)
• Edge cases are covered by tests
• Dask chunk boundaries handled correctly (unchanged; kernel is per chunk)
• No premature materialization or unnecessary copies
• Benchmark exists
• README feature matrix: not applicable, no new function
• Docstrings present and accurate

[brendancol]

PR Review: Speed up the proximity brute-force kernel (#3740), follow-up pass

Re-reviewed commit 27d6771 against the first pass.

Blockers (must fix before merge)

• none

Suggestions (should fix, not blocking)

• none

Nits (optional improvements)

• none

Disposition of the first pass

• Suggestion, NaN rounded distance poisoning the running best: fixed. All three inner loops now do best_dist = d if better else best_dist, and the shared comment block above them explains why. best_proxy still updates unconditionally, which is fine: any later candidate with a proxy at or above a NaN-producing one would produce NaN too and could not win under the old kernel either. A/B against main is still bitwise over the same 54 cases.
• Nit, tiny concurrency raster: fixed. The lock test now hammers a 40x40 raster with 50 random targets.
• Nit, placeholder great-circle arrays: fixed with a comment at the call site.

What looks good

• The follow-up touches only the three select lines, one comment, and the test; no behaviour change outside the NaN edge case. 609 tests pass locally including the cupy and dask+cupy backends.