spatialrust_search/
chunked.rs

1//! Chunk-oriented query extensions for [`SpatialIndex`](crate::SpatialIndex) backends.
2//!
3//! These traits align query batching with [`SpatialTensor`](spatialrust_core::SpatialTensor)
4//! chunk iteration for parallel CPU/GPU algorithms.
5
6use std::ops::Range;
7
8use spatialrust_core::SpatialTensor;
9
10use crate::{Neighbor, RadiusSearchIndex};
11
12/// Range of query point indices (matches [`SpatialTensor::chunks`] ranges).
13pub type ChunkQueryRange = Range<usize>;
14
15/// Radius search over contiguous query index ranges.
16pub trait ChunkedRadiusSearchIndex: RadiusSearchIndex {
17    /// Appends `(query_index, neighbor)` for each point index in `chunk`.
18    fn radius_search_chunk_into(
19        &self,
20        x: &[f32],
21        y: &[f32],
22        z: &[f32],
23        chunk: ChunkQueryRange,
24        radius: f32,
25        out: &mut Vec<(usize, Neighbor)>,
26    );
27
28    /// Radius search for a single indexed query point.
29    fn radius_search_at(
30        &self,
31        x: &[f32],
32        y: &[f32],
33        z: &[f32],
34        index: usize,
35        radius: f32,
36    ) -> Vec<Neighbor> {
37        self.radius_search(x[index], y[index], z[index], radius)
38    }
39}
40
41/// k-NN search over contiguous query index ranges.
42pub trait ChunkedNearestNeighborIndex: crate::NearestNeighborIndex {
43    /// Appends `(query_index, neighbor)` for each point index in `chunk`.
44    fn nearest_k_chunk_into(
45        &self,
46        x: &[f32],
47        y: &[f32],
48        z: &[f32],
49        chunk: ChunkQueryRange,
50        k: usize,
51        out: &mut Vec<(usize, Neighbor)>,
52    );
53}
54
55/// Runs radius search for every chunk in a [`SpatialTensor`], appending tagged neighbors.
56pub fn radius_search_spatial_tensor<I: ChunkedRadiusSearchIndex>(
57    index: &I,
58    x: &[f32],
59    y: &[f32],
60    z: &[f32],
61    tensor: &SpatialTensor<'_>,
62    radius: f32,
63    out: &mut Vec<(usize, Neighbor)>,
64) {
65    for chunk in tensor.chunks() {
66        index.radius_search_chunk_into(x, y, z, chunk.range(), radius, out);
67    }
68}
69
70/// Runs k-NN search for every chunk in a [`SpatialTensor`], appending tagged neighbors.
71pub fn nearest_k_spatial_tensor<I: ChunkedNearestNeighborIndex>(
72    index: &I,
73    x: &[f32],
74    y: &[f32],
75    z: &[f32],
76    tensor: &SpatialTensor<'_>,
77    k: usize,
78    out: &mut Vec<(usize, Neighbor)>,
79) {
80    for chunk in tensor.chunks() {
81        index.nearest_k_chunk_into(x, y, z, chunk.range(), k, out);
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::{
88        nearest_k_spatial_tensor, radius_search_spatial_tensor, ChunkedNearestNeighborIndex,
89        ChunkedRadiusSearchIndex,
90    };
91    use crate::{brute::BruteForceIndex, kdtree::KdTree, RadiusSearchIndex};
92    use spatialrust_core::{PointCloudBuilder, StandardSchemas};
93
94    fn sample_cloud() -> (Vec<f32>, Vec<f32>, Vec<f32>) {
95        (
96            vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0],
97            vec![0.0, 0.0, 0.0, 1.0, 2.0, 0.0],
98            vec![0.0, 0.0, 0.0, 0.0, 0.0, 5.0],
99        )
100    }
101
102    fn chunked_matches_per_index<I>(index: &I, radius: f32)
103    where
104        I: ChunkedRadiusSearchIndex + RadiusSearchIndex,
105    {
106        let (x, y, z) = sample_cloud();
107        let mut chunked = Vec::new();
108        index.radius_search_chunk_into(&x, &y, &z, 1..4, radius, &mut chunked);
109
110        let mut expected = Vec::new();
111        for query in 1..4 {
112            for neighbor in index.radius_search_at(&x, &y, &z, query, radius) {
113                expected.push((query, neighbor));
114            }
115        }
116
117        chunked.sort_by_key(|a| (a.0, a.1.index));
118        expected.sort_by_key(|a| (a.0, a.1.index));
119        assert_eq!(chunked, expected);
120    }
121
122    #[test]
123    fn kdtree_chunked_radius_matches_per_index() {
124        let (x, y, z) = sample_cloud();
125        let tree = KdTree::from_slices(&x, &y, &z);
126        chunked_matches_per_index(&tree, 1.5);
127    }
128
129    #[test]
130    fn brute_chunked_radius_matches_per_index() {
131        let (x, y, z) = sample_cloud();
132        let index = BruteForceIndex::from_slices(&x, &y, &z);
133        chunked_matches_per_index(&index, 1.5);
134    }
135
136    #[test]
137    fn spatial_tensor_radius_matches_full_scan() {
138        let (x, y, z) = sample_cloud();
139        let tree = KdTree::from_slices(&x, &y, &z);
140
141        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
142        for index in 0..x.len() {
143            builder.push_point([x[index], y[index], z[index]]).unwrap();
144        }
145        let cloud = builder.build().unwrap();
146        let tensor = cloud.spatial_tensor_chunks(2).unwrap();
147
148        let mut tensor_out = Vec::new();
149        radius_search_spatial_tensor(&tree, &x, &y, &z, &tensor, 1.5, &mut tensor_out);
150
151        let mut full_out = Vec::new();
152        tree.radius_search_chunk_into(&x, &y, &z, 0..x.len(), 1.5, &mut full_out);
153
154        tensor_out.sort_by_key(|a| (a.0, a.1.index));
155        full_out.sort_by_key(|a| (a.0, a.1.index));
156        assert_eq!(tensor_out, full_out);
157    }
158
159    #[test]
160    fn spatial_tensor_nearest_k_matches_full_scan() {
161        let (x, y, z) = sample_cloud();
162        let tree = KdTree::from_slices(&x, &y, &z);
163
164        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
165        for index in 0..x.len() {
166            builder.push_point([x[index], y[index], z[index]]).unwrap();
167        }
168        let cloud = builder.build().unwrap();
169        let tensor = cloud.spatial_tensor_chunks(2).unwrap();
170
171        let mut tensor_out = Vec::new();
172        nearest_k_spatial_tensor(&tree, &x, &y, &z, &tensor, 2, &mut tensor_out);
173
174        let mut full_out = Vec::new();
175        tree.nearest_k_chunk_into(&x, &y, &z, 0..x.len(), 2, &mut full_out);
176
177        tensor_out.sort_by(|a, b| {
178            a.0.cmp(&b.0)
179                .then(
180                    a.1.distance_squared
181                        .partial_cmp(&b.1.distance_squared)
182                        .unwrap_or(std::cmp::Ordering::Equal),
183                )
184                .then(a.1.index.cmp(&b.1.index))
185        });
186        full_out.sort_by(|a, b| {
187            a.0.cmp(&b.0)
188                .then(
189                    a.1.distance_squared
190                        .partial_cmp(&b.1.distance_squared)
191                        .unwrap_or(std::cmp::Ordering::Equal),
192                )
193                .then(a.1.index.cmp(&b.1.index))
194        });
195        assert_eq!(tensor_out, full_out);
196    }
197}
spatialrust_search/chunked.rs

spatialrust_search/
chunked.rs
