vibespatial.spatial.spatial_index_knn_device¶

Device-side k-nearest-neighbor spatial query.

Provides spatial_index_knn_device, a GPU-accelerated k-NN spatial query that replaces the CPU STRtree nearest path in sjoin_nearest.

Pipeline:

Expand query bounds by max_distance (or progressive expansion)
Generate bbox candidate pairs via spatial_index_device_query
Compute exact distances for candidate pairs (reuses nearest.py strategies)
Per-query top-k selection via CCCL segmented_sort
Output: device-resident (query_idx, target_idx, distance) triples

ADR-0002: METRIC kernel class for distance computation – fp64 required. ADR-0033 tier classification:

Tier 2 (CuPy): element-wise distance filtering, gather/scatter

Tier 3a (CCCL): segmented_sort (per-query ranking), exclusive_sum, lower_bound, upper_bound, compact_indices

Attributes¶

`cp`
`logger`

Classes¶

DeviceKnnResult

Device-resident k-NN query result.

Functions¶

spatial_index_knn_device(→ DeviceKnnResult | None)

GPU-accelerated k-nearest-neighbor spatial query.

Module Contents¶

vibespatial.spatial.spatial_index_knn_device.cp = None¶

vibespatial.spatial.spatial_index_knn_device.logger¶

class vibespatial.spatial.spatial_index_knn_device.DeviceKnnResult¶

Device-resident k-NN query result.

All arrays are CuPy device arrays to avoid D->H transfers when the result feeds directly into the next GPU pipeline stage (e.g., sjoin attribute assembly).

Attributes¶

d_query_idxdevice int32 array: Query geometry indices (one per result pair).
d_target_idxdevice int32 array: Target geometry indices (one per result pair).
d_distancesdevice float64 array: Exact distances for each (query, target) pair.
total_pairsint: Number of result pairs.
kint: Requested k value.

d_query_idx: Any¶

d_target_idx: Any¶

d_distances: Any¶

total_pairs: int¶

k: int¶

to_host() → tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]¶: Copy result to host as numpy arrays.

vibespatial.spatial.spatial_index_knn_device.spatial_index_knn_device(query_owned: vibespatial.geometry.owned.OwnedGeometryArray, tree_owned: vibespatial.geometry.owned.OwnedGeometryArray, query_bounds: numpy.ndarray, tree_bounds: numpy.ndarray, *, k: int = 1, max_distance: float | None = None, exclusive: bool = False, return_all: bool = True, precision: vibespatial.runtime.precision.PrecisionMode | str = PrecisionMode.AUTO) → DeviceKnnResult | None¶

GPU-accelerated k-nearest-neighbor spatial query.

Replaces the CPU STRtree nearest path with a fully device-resident pipeline: candidate generation -> exact distance -> per-query top-k.

Parameters¶

query_ownedOwnedGeometryArray: Query geometries (source of the nearest search).
tree_ownedOwnedGeometryArray: Target geometries (the “tree” to search against).
query_boundsnp.ndarray, shape (Q, 4): Pre-computed query bounding boxes.
tree_boundsnp.ndarray, shape (M, 4): Pre-computed target bounding boxes.
kint: Number of nearest neighbours per query. k=1 is the most common case.
max_distancefloat or None: Maximum search distance. Candidates beyond this are pruned. When None, an effective distance is computed from the data extent.
exclusivebool: If True, exclude identical geometries from results.
return_allbool: If True, return all k-nearest ties. If False, return exactly one per query (the first nearest).
precisionPrecisionMode: Precision mode for distance computation. METRIC class requires fp64 per ADR-0002 on all devices.

Returns¶

DeviceKnnResult or None: Device-resident result with (query_idx, target_idx, distance) triples, or None if the GPU path is not applicable.