rigidity_core/
neighbors.rs

1//! Neighbour search: the contract, and a brute-force reference implementation.
2
3use nalgebra::Vector3;
4
5use crate::cloud::PointCloud;
6
7/// A neighbour that was found.
8#[derive(Debug, Clone, Copy, PartialEq)]
9pub struct Neighbor {
10    /// Index of the point in the source cloud.
11    pub index: u32,
12    /// Squared distance to the query. Squared rather than plain: the
13    /// square root is never needed for ranking neighbours, and taking it
14    /// costs cycles and adds rounding.
15    pub distance_squared: f64,
16}
17
18/// The order neighbours are returned in: by increasing distance, ties
19/// broken by increasing index.
20///
21/// Breaking ties by index is mandatory. Without it two equally correct
22/// implementations return different orders for equidistant points, and
23/// further up the stack that turns into a different order of summation.
24pub fn compare_neighbors(a: &Neighbor, b: &Neighbor) -> std::cmp::Ordering {
25    a.distance_squared
26        .total_cmp(&b.distance_squared)
27        .then(a.index.cmp(&b.index))
28}
29
30/// Nearest-neighbour search.
31///
32/// Implementations must honour the order defined by [`compare_neighbors`].
33pub trait NeighborSearch {
34    /// The `k` nearest points to `query`, written into `out`, whose
35    /// previous contents are discarded.
36    ///
37    /// The buffer is supplied by the caller so that the ICP hot loop does
38    /// not allocate once per point.
39    fn knn_into(&self, query: &Vector3<f64>, k: usize, out: &mut Vec<Neighbor>);
40
41    /// Every point inside the ball of radius `radius`, written into `out`.
42    fn radius_into(&self, query: &Vector3<f64>, radius: f64, out: &mut Vec<Neighbor>);
43
44    /// Allocating convenience wrapper over [`knn_into`](Self::knn_into).
45    fn knn(&self, query: &Vector3<f64>, k: usize) -> Vec<Neighbor> {
46        let mut out = Vec::with_capacity(k);
47        self.knn_into(query, k, &mut out);
48        out
49    }
50
51    /// Allocating convenience wrapper over [`radius_into`](Self::radius_into).
52    fn radius(&self, query: &Vector3<f64>, radius: f64) -> Vec<Neighbor> {
53        let mut out = Vec::new();
54        self.radius_into(query, radius, &mut out);
55        out
56    }
57}
58
59/// Brute force.
60///
61/// It exists as the reference the kd-tree is checked against. There is no
62/// reason to optimise it; it never sits in the hot path.
63pub struct BruteForce<'a> {
64    cloud: &'a PointCloud,
65}
66
67impl<'a> BruteForce<'a> {
68    /// Wraps a cloud.
69    pub fn new(cloud: &'a PointCloud) -> Self {
70        Self { cloud }
71    }
72
73    fn all_distances(&self, query: &Vector3<f64>, out: &mut Vec<Neighbor>) {
74        out.clear();
75        out.reserve(self.cloud.len());
76        for i in 0..self.cloud.len() {
77            out.push(Neighbor {
78                index: i as u32,
79                distance_squared: (self.cloud.point(i) - query).norm_squared(),
80            });
81        }
82    }
83}
84
85impl NeighborSearch for BruteForce<'_> {
86    fn knn_into(&self, query: &Vector3<f64>, k: usize, out: &mut Vec<Neighbor>) {
87        self.all_distances(query, out);
88        out.sort_unstable_by(compare_neighbors);
89        out.truncate(k);
90    }
91
92    fn radius_into(&self, query: &Vector3<f64>, radius: f64, out: &mut Vec<Neighbor>) {
93        let limit = radius * radius;
94        self.all_distances(query, out);
95        out.retain(|n| n.distance_squared <= limit);
96        out.sort_unstable_by(compare_neighbors);
97    }
98}
rigidity_core/neighbors.rs

rigidity_core/
neighbors.rs
