rigidity_spatial/
lib.rs

1//! Spatial indices.
2//!
3//! This crate is separate from the core because the core's direct
4//! dependencies are fixed and checked in CI: `kiddo` does not fit there,
5//! and file I/O is no place for a data structure.
6//!
7//! The core defines the [`NeighborSearch`] contract and a brute-force
8//! reference implementation; the working one lives here. Drawing the
9//! boundary at a trait rather than a type leaves room for a grid-based or
10//! GPU search later without touching ICP.
11
12use kiddo::{ImmutableKdTree, SquaredEuclidean};
13use nalgebra::Vector3;
14use rigidity_core::{Neighbor, NeighborSearch, PointCloud, neighbors::compare_neighbors};
15use std::num::NonZeroUsize;
16
17/// An error while building the index.
18#[derive(Debug, thiserror::Error)]
19pub enum SpatialError {
20    /// Kiddo failed to build the tree.
21    #[error("kd-tree construction failed: {0}")]
22    Construction(String),
23}
24
25/// A kd-tree built over a cloud.
26///
27/// The tree is built in **absolute** coordinates, the same ones
28/// [`PointCloud::point`] returns. That makes distances bit-for-bit
29/// comparable with the brute-force reference; otherwise the agreement test
30/// would need a tolerance and would stop catching indexing mistakes.
31pub struct KdTree {
32    tree: ImmutableKdTree<f64, 3>,
33    len: usize,
34}
35
36impl KdTree {
37    /// Builds a tree over a cloud.
38    pub fn build(cloud: &PointCloud) -> Result<Self, SpatialError> {
39        Self::build_observed(cloud, |_, _| {})
40    }
41
42    /// The same, reporting progress after each internal phase.
43    ///
44    /// `progress` receives `(completed, total)` in phases: the
45    /// coordinates are extracted, then the tree is built. Two steps is all
46    /// the resolution there is — `kiddo` builds the tree in one opaque
47    /// call, and on a large cloud that call is the larger half of the
48    /// wait. Pretending otherwise would mean a progress bar that lies.
49    pub fn build_observed<F>(cloud: &PointCloud, mut progress: F) -> Result<Self, SpatialError>
50    where
51        F: FnMut(usize, usize),
52    {
53        /// Coordinates, then construction.
54        const PHASES: usize = 2;
55
56        let entries: Vec<[f64; 3]> = (0..cloud.len())
57            .map(|i| {
58                let p = cloud.point(i);
59                [p.x, p.y, p.z]
60            })
61            .collect();
62        progress(1, PHASES);
63
64        let tree = ImmutableKdTree::new_from_slice(&entries)
65            .map_err(|e| SpatialError::Construction(e.to_string()))?;
66        progress(2, PHASES);
67
68        Ok(Self {
69            tree,
70            len: cloud.len(),
71        })
72    }
73
74    /// Number of indexed points.
75    pub fn len(&self) -> usize {
76        self.len
77    }
78
79    /// Whether the tree is empty.
80    pub fn is_empty(&self) -> bool {
81        self.len == 0
82    }
83}
84
85impl NeighborSearch for KdTree {
86    fn knn_into(&self, query: &Vector3<f64>, k: usize, out: &mut Vec<Neighbor>) {
87        out.clear();
88        let Some(count) = NonZeroUsize::new(k.min(self.len)) else {
89            return;
90        };
91        let point = [query.x, query.y, query.z];
92
93        // Fast path: ICP asks for exactly one neighbour, while `nearest_n`
94        // returns a `Vec` and therefore allocates per query. On a million
95        // points that is a million allocator calls per iteration.
96        // `nearest_one` returns a single value instead.
97        if count.get() == 1 {
98            let found = self
99                .tree
100                .query(&point)
101                .nearest_one::<SquaredEuclidean<f64>>()
102                .execute();
103            out.push(Neighbor {
104                index: found.item,
105                distance_squared: found.distance,
106            });
107            return;
108        }
109
110        let found = self
111            .tree
112            .query(&point)
113            .nearest_n::<SquaredEuclidean<f64>>(count)
114            .execute();
115        for item in found {
116            out.push(Neighbor {
117                index: item.item,
118                distance_squared: item.distance,
119            });
120        }
121        // kiddo makes no promise about breaking ties by index, so impose
122        // the contract here.
123        out.sort_unstable_by(compare_neighbors);
124    }
125
126    fn radius_into(&self, query: &Vector3<f64>, radius: f64, out: &mut Vec<Neighbor>) {
127        out.clear();
128        let point = [query.x, query.y, query.z];
129        // The threshold is in the metric's own units: for SquaredEuclidean
130        // that is a squared distance, not a distance.
131        let found = self
132            .tree
133            .query(&point)
134            .within::<SquaredEuclidean<f64>>(radius * radius)
135            .execute();
136        for item in found {
137            out.push(Neighbor {
138                index: item.item,
139                distance_squared: item.distance,
140            });
141        }
142        out.sort_unstable_by(compare_neighbors);
143    }
144}
rigidity_spatial/lib.rs

rigidity_spatial/
lib.rs
