spatialrust_search/
brute.rs

1use crate::{
2    chunked::{ChunkQueryRange, ChunkedNearestNeighborIndex, ChunkedRadiusSearchIndex},
3    NearestNeighborIndex, Neighbor, RadiusSearchIndex, SpatialIndex,
4};
5
6/// Reference index using brute-force search for correctness tests.
7#[derive(Clone, Debug)]
8pub struct BruteForceIndex {
9    x: Vec<f32>,
10    y: Vec<f32>,
11    z: Vec<f32>,
12}
13
14impl BruteForceIndex {
15    /// Builds a brute-force index from coordinate slices.
16    #[must_use]
17    pub fn from_slices(x: &[f32], y: &[f32], z: &[f32]) -> Self {
18        assert_eq!(x.len(), y.len());
19        assert_eq!(x.len(), z.len());
20        Self { x: x.to_vec(), y: y.to_vec(), z: z.to_vec() }
21    }
22}
23
24impl SpatialIndex for BruteForceIndex {
25    fn len(&self) -> usize {
26        self.x.len()
27    }
28}
29
30impl NearestNeighborIndex for BruteForceIndex {
31    fn nearest_one(&self, x: f32, y: f32, z: f32) -> Option<Neighbor> {
32        brute_force_knn(&self.x, &self.y, &self.z, x, y, z, 1).into_iter().next()
33    }
34
35    fn nearest_k(&self, x: f32, y: f32, z: f32, k: usize) -> Vec<Neighbor> {
36        brute_force_knn(&self.x, &self.y, &self.z, x, y, z, k)
37    }
38}
39
40impl RadiusSearchIndex for BruteForceIndex {
41    fn radius_search(&self, x: f32, y: f32, z: f32, radius: f32) -> Vec<Neighbor> {
42        brute_force_radius(&self.x, &self.y, &self.z, x, y, z, radius)
43    }
44}
45
46impl ChunkedRadiusSearchIndex for BruteForceIndex {
47    fn radius_search_chunk_into(
48        &self,
49        x: &[f32],
50        y: &[f32],
51        z: &[f32],
52        chunk: ChunkQueryRange,
53        radius: f32,
54        out: &mut Vec<(usize, Neighbor)>,
55    ) {
56        for index in chunk {
57            for neighbor in self.radius_search_at(x, y, z, index, radius) {
58                out.push((index, neighbor));
59            }
60        }
61    }
62}
63
64impl ChunkedNearestNeighborIndex for BruteForceIndex {
65    fn nearest_k_chunk_into(
66        &self,
67        x: &[f32],
68        y: &[f32],
69        z: &[f32],
70        chunk: ChunkQueryRange,
71        k: usize,
72        out: &mut Vec<(usize, Neighbor)>,
73    ) {
74        for index in chunk {
75            for neighbor in self.nearest_k(x[index], y[index], z[index], k) {
76                out.push((index, neighbor));
77            }
78        }
79    }
80}
81
82/// Finds up to `k` nearest neighbors by brute force.
83#[must_use]
84pub fn brute_force_knn(
85    x: &[f32],
86    y: &[f32],
87    z: &[f32],
88    qx: f32,
89    qy: f32,
90    qz: f32,
91    k: usize,
92) -> Vec<Neighbor> {
93    if k == 0 || x.is_empty() {
94        return Vec::new();
95    }
96
97    let mut neighbors: Vec<Neighbor> = x
98        .iter()
99        .enumerate()
100        .map(|(index, &px)| Neighbor {
101            index,
102            distance_squared: squared_distance(px, y[index], z[index], qx, qy, qz),
103        })
104        .collect();
105    neighbors.sort_by(|a, b| {
106        a.distance_squared.partial_cmp(&b.distance_squared).unwrap_or(std::cmp::Ordering::Equal)
107    });
108    neighbors.truncate(k);
109    neighbors
110}
111
112/// Finds all neighbors within `radius` by brute force.
113#[must_use]
114pub fn brute_force_radius(
115    x: &[f32],
116    y: &[f32],
117    z: &[f32],
118    qx: f32,
119    qy: f32,
120    qz: f32,
121    radius: f32,
122) -> Vec<Neighbor> {
123    let radius_sq = radius * radius;
124    let mut neighbors = Vec::new();
125    for (index, &px) in x.iter().enumerate() {
126        let dist_sq = squared_distance(px, y[index], z[index], qx, qy, qz);
127        if dist_sq <= radius_sq {
128            neighbors.push(Neighbor { index, distance_squared: dist_sq });
129        }
130    }
131    neighbors.sort_by(|a, b| {
132        a.distance_squared.partial_cmp(&b.distance_squared).unwrap_or(std::cmp::Ordering::Equal)
133    });
134    neighbors
135}
136
137fn squared_distance(px: f32, py: f32, pz: f32, qx: f32, qy: f32, qz: f32) -> f32 {
138    let dx = px - qx;
139    let dy = py - qy;
140    let dz = pz - qz;
141    dx * dx + dy * dy + dz * dz
142}
spatialrust_search/brute.rs

spatialrust_search/
brute.rs
