ogeom_mesh/
point_index.rs

1//! A k-d tree over a point cloud: the nearest point, the `k` nearest, and
2//! every point within a ball.
3//!
4//! The same contract as [`Bvh`](crate::bvh::Bvh): built once, immutable,
5//! shared across threads, and every answer the one brute force gives over
6//! the same points with the same arithmetic, ties going to the lowest point
7//! index. A cell is entered only if its box's computed squared distance is no
8//! more than the best found so far, and since every axis gap to a box is at
9//! most the gap to a point in it, rounded the same way, no point a query
10//! would accept lies in a cell it skips.
11
12use ogeom_core::{OgeomResult, ogeom_bail};
13use ogeom_math::Point;
14
15use crate::bvh::{distance2, gap};
16
17type P = [f64; 3];
18
19/// A point a query found, by its index in the cloud.
20#[derive(Debug, Clone, Copy, PartialEq)]
21pub struct Neighbour {
22    /// The point's index in the cloud the index was built from.
23    pub index: u32,
24    /// Its distance from the query point.
25    pub distance: f64,
26}
27
28#[derive(Debug, Clone, Copy)]
29struct Cell {
30    low: P,
31    high: P,
32    /// A leaf's first slot, or an interior cell's upper half. Its lower
33    /// half is the next cell.
34    first: u32,
35    /// How many points a leaf holds; zero for an interior cell.
36    count: u32,
37}
38
39/// A k-d tree over points, split at the median of each cell's widest axis.
40///
41/// Holds its own copy of the points.
42#[derive(Debug, Clone)]
43pub struct PointIndex {
44    cells: Vec<Cell>,
45    /// The points, slots in leaf order.
46    points: Vec<P>,
47    /// The point in each slot.
48    order: Vec<u32>,
49}
50
51/// The most points a leaf holds.
52const LEAF_MOST: usize = 8;
53
54impl PointIndex {
55    /// Build the tree over a cloud.
56    ///
57    /// # Errors
58    ///
59    /// [`OgeomError::Numeric`](ogeom_core::OgeomError::Numeric) if a point
60    /// is not finite, and
61    /// [`OgeomError::Range`](ogeom_core::OgeomError::Range) for a cloud of
62    /// more than four billion points.
63    pub fn build(points: &[Point]) -> OgeomResult<Self> {
64        if u32::try_from(points.len()).is_err() {
65            ogeom_bail!(Range, "a cloud of more than four billion points");
66        }
67        let mut held = Vec::with_capacity(points.len());
68        for (i, p) in points.iter().enumerate() {
69            if !(p.x.is_finite() && p.y.is_finite() && p.z.is_finite()) {
70                ogeom_bail!(Numeric, "point {i} is not finite");
71            }
72            held.push([p.x, p.y, p.z]);
73        }
74        let mut order: Vec<u32> = (0..held.len())
75            .map(|i| u32::try_from(i).unwrap_or(u32::MAX))
76            .collect();
77        let cells = build_cells(&held, &mut order);
78        Ok(Self {
79            cells,
80            points: order.iter().map(|&i| held[i as usize]).collect(),
81            order,
82        })
83    }
84
85    /// How many points the tree holds.
86    #[must_use]
87    pub fn len(&self) -> usize {
88        self.order.len()
89    }
90
91    /// Whether the tree holds no point.
92    #[must_use]
93    pub fn is_empty(&self) -> bool {
94        self.order.is_empty()
95    }
96
97    /// The points the tree was built over, in their cloud order.
98    #[must_use]
99    pub fn points(&self) -> Vec<Point> {
100        let mut out = vec![Point::ORIGIN; self.order.len()];
101        for (p, &i) in self.points.iter().zip(&self.order) {
102            out[i as usize] = Point::new(p[0], p[1], p[2]);
103        }
104        out
105    }
106
107    /// The point nearest `p`, no farther than `max`; the lowest index among
108    /// equally near ones. `max` may be infinite.
109    #[must_use]
110    pub fn nearest(&self, p: Point, max: f64) -> Option<Neighbour> {
111        self.k_nearest_with(p, 1, max, &mut Vec::new()).pop()
112    }
113
114    /// [`PointIndex::nearest`] for each point, in parallel, answers in point
115    /// order.
116    #[must_use]
117    pub fn nearest_many(&self, points: &[Point], max: f64) -> Vec<Option<Neighbour>> {
118        crate::bvh::batched(
119            points,
120            |&p| p,
121            |&p, stack| self.k_nearest_with(p, 1, max, stack).pop(),
122        )
123    }
124
125    /// The `k` points nearest `p` no farther than `max`, nearest first, ties
126    /// in index order: fewer where fewer are that near.
127    #[must_use]
128    pub fn k_nearest(&self, p: Point, k: usize, max: f64) -> Vec<Neighbour> {
129        self.k_nearest_with(p, k, max, &mut Vec::new())
130    }
131
132    fn k_nearest_with(
133        &self,
134        p: Point,
135        k: usize,
136        max: f64,
137        stack: &mut Vec<(u32, f64)>,
138    ) -> Vec<Neighbour> {
139        let p = [p.x, p.y, p.z];
140        let reach = max * max;
141        if k == 0
142            || self.cells.is_empty()
143            || !p.iter().all(|c| c.is_finite())
144            || max.is_nan()
145            || max < 0.0
146        {
147            return Vec::new();
148        }
149        // The best so far as (squared distance, index), sorted, at most `k`.
150        let mut best: Vec<(f64, u32)> = Vec::with_capacity(k.min(self.order.len()) + 1);
151        let bound = |best: &Vec<(f64, u32)>| {
152            if best.len() == k {
153                best[k - 1].0
154            } else {
155                reach
156            }
157        };
158        stack.clear();
159        stack.push((0, cell_gap(&self.cells[0], p)));
160        while let Some((at, g)) = stack.pop() {
161            if g > bound(&best) {
162                continue;
163            }
164            let cell = &self.cells[at as usize];
165            if cell.count > 0 {
166                let first = cell.first as usize;
167                for slot in first..first + cell.count as usize {
168                    let d2 = distance2(self.points[slot], p);
169                    let index = self.order[slot];
170                    let full = best.len() == k;
171                    if d2 > reach || (full && (d2, index) >= best[k - 1]) {
172                        continue;
173                    }
174                    let at = best.partition_point(|&held| held < (d2, index));
175                    best.insert(at, (d2, index));
176                    best.truncate(k);
177                }
178                continue;
179            }
180            let (lower, upper) = (at + 1, cell.first);
181            let gl = cell_gap(&self.cells[lower as usize], p);
182            let gu = cell_gap(&self.cells[upper as usize], p);
183            let (near, far) = if gl <= gu {
184                ((lower, gl), (upper, gu))
185            } else {
186                ((upper, gu), (lower, gl))
187            };
188            let limit = bound(&best);
189            if far.1 <= limit {
190                stack.push(far);
191            }
192            if near.1 <= limit {
193                stack.push(near);
194            }
195        }
196        best.into_iter()
197            .map(|(d2, index)| Neighbour {
198                index,
199                distance: d2.sqrt(),
200            })
201            .collect()
202    }
203
204    /// Every point within `radius` of `centre`, in index order.
205    #[must_use]
206    pub fn within_ball(&self, centre: Point, radius: f64) -> Vec<u32> {
207        let c = [centre.x, centre.y, centre.z];
208        let reach = radius * radius;
209        let mut found = Vec::new();
210        if self.cells.is_empty()
211            || !c.iter().all(|v| v.is_finite())
212            || radius.is_nan()
213            || radius < 0.0
214        {
215            return found;
216        }
217        let mut stack = vec![0_u32];
218        while let Some(at) = stack.pop() {
219            let cell = &self.cells[at as usize];
220            if cell_gap(cell, c) > reach {
221                continue;
222            }
223            if cell.count > 0 {
224                let first = cell.first as usize;
225                for slot in first..first + cell.count as usize {
226                    if distance2(self.points[slot], c) <= reach {
227                        found.push(self.order[slot]);
228                    }
229                }
230            } else {
231                stack.push(cell.first);
232                stack.push(at + 1);
233            }
234        }
235        found.sort_unstable();
236        found
237    }
238}
239
240fn cell_gap(cell: &Cell, p: P) -> f64 {
241    gap(&cell.low, &cell.high, p)
242}
243
244/// Build the cells, reordering `order` into leaf order: depth first, each
245/// lower half right after its parent.
246fn build_cells(points: &[P], order: &mut [u32]) -> Vec<Cell> {
247    let mut cells: Vec<Cell> = Vec::new();
248    if order.is_empty() {
249        return cells;
250    }
251    let mut pending: Vec<(usize, usize, Option<usize>)> = vec![(0, order.len(), None)];
252    while let Some((start, end, parent)) = pending.pop() {
253        let index = cells.len();
254        if let Some(parent) = parent {
255            cells[parent].first = u32::try_from(index).unwrap_or(u32::MAX);
256        }
257        let range = &mut order[start..end];
258        let (mut low, mut high) = ([f64::INFINITY; 3], [f64::NEG_INFINITY; 3]);
259        for &i in range.iter() {
260            let p = points[i as usize];
261            for k in 0..3 {
262                low[k] = low[k].min(p[k]);
263                high[k] = high[k].max(p[k]);
264            }
265        }
266        let count = end - start;
267        let cell = Cell {
268            low,
269            high,
270            first: u32::try_from(start).unwrap_or(u32::MAX),
271            count: u32::try_from(count).unwrap_or(u32::MAX),
272        };
273        if count <= LEAF_MOST {
274            cells.push(cell);
275            continue;
276        }
277        let spread = [0, 1, 2].map(|k| high[k] - low[k]);
278        let axis = if spread[0] >= spread[1] && spread[0] >= spread[2] {
279            0
280        } else if spread[1] >= spread[2] {
281            1
282        } else {
283            2
284        };
285        let half = count / 2;
286        range.select_nth_unstable_by(half, |&a, &b| {
287            points[a as usize][axis]
288                .total_cmp(&points[b as usize][axis])
289                .then(a.cmp(&b))
290        });
291        cells.push(Cell {
292            first: 0,
293            count: 0,
294            ..cell
295        });
296        pending.push((start + half, end, Some(index)));
297        pending.push((start, start + half, None));
298    }
299    cells
300}
301
302#[cfg(test)]
303#[allow(clippy::unwrap_used, clippy::expect_used)]
304mod tests {
305    use super::*;
306    use crate::bvh::tests::Seeded;
307
308    fn brute(points: &[Point], p: Point, max: f64) -> Vec<Neighbour> {
309        let q = [p.x, p.y, p.z];
310        let mut all: Vec<(f64, u32)> = points
311            .iter()
312            .enumerate()
313            .map(|(i, x)| (distance2([x.x, x.y, x.z], q), u32::try_from(i).unwrap()))
314            .filter(|&(d2, _)| d2 <= max * max)
315            .collect();
316        all.sort_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
317        all.into_iter()
318            .map(|(d2, index)| Neighbour {
319                index,
320                distance: d2.sqrt(),
321            })
322            .collect()
323    }
324
325    /// A cloud with repeats and points on a coarse lattice, so ties in
326    /// distance are common, plus uniform noise.
327    fn cloud(count: usize, rng: &mut Seeded) -> Vec<Point> {
328        let mut points = Vec::with_capacity(count);
329        for i in 0..count {
330            let p = match i % 4 {
331                0 if i > 0 => points[rng.below(i)],
332                1 => {
333                    let at =
334                        |rng: &mut Seeded| f64::from(u32::try_from(rng.below(5)).unwrap()) * 0.5;
335                    Point::new(at(rng), at(rng), at(rng))
336                }
337                _ => rng.point(0.0, 2.0),
338            };
339            points.push(p);
340        }
341        points
342    }
343
344    /// Nearest, k nearest and ball queries agree with brute force, bit for
345    /// bit and tie for tie: 1e5 of each.
346    #[test]
347    fn every_query_equals_brute_force_on_random_clouds() {
348        let mut rng = Seeded(0x0009_0177);
349        let mut found = [0_usize; 3];
350        for round in 0..10 {
351            let points = cloud(100 + 40 * round, &mut rng);
352            let index = PointIndex::build(&points).unwrap();
353            assert_eq!(index.len(), points.len());
354            for q in 0..10_000 {
355                let p = if q % 3 == 0 {
356                    points[rng.below(points.len())]
357                } else {
358                    rng.point(-0.5, 2.5)
359                };
360                let max = if q % 2 == 0 {
361                    f64::INFINITY
362                } else {
363                    rng.range(0.0, 0.5)
364                };
365                let all = brute(&points, p, max);
366                let near = index.nearest(p, max);
367                assert_eq!(near, all.first().copied(), "{p:?} {max}");
368                found[0] += usize::from(near.is_some());
369                let k = 1 + rng.below(12);
370                let some = index.k_nearest(p, k, max);
371                assert_eq!(some, all[..k.min(all.len())], "{p:?} {k} {max}");
372                found[1] += usize::from(some.len() == k);
373                let r = rng.range(0.0, 0.5);
374                let mut ball: Vec<u32> = brute(&points, p, r).iter().map(|n| n.index).collect();
375                ball.sort_unstable();
376                assert_eq!(index.within_ball(p, r), ball, "{p:?} {r}");
377                found[2] += usize::from(!ball.is_empty());
378            }
379        }
380        assert!(found.iter().all(|&n| n > 25_000), "{found:?}");
381    }
382
383    /// A query on a repeated point goes to the lowest of its copies.
384    #[test]
385    fn ties_go_to_the_lowest_index() {
386        let p = Point::new(1.0, 2.0, 3.0);
387        let mut points = vec![Point::new(9.0, 9.0, 9.0); 20];
388        points[7] = p;
389        points[13] = p;
390        points[3] = Point::new(1.0, 2.0, 4.0);
391        points[17] = Point::new(1.0, 2.0, 2.0);
392        let index = PointIndex::build(&points).unwrap();
393        assert_eq!(
394            index.nearest(p, 0.0),
395            Some(Neighbour {
396                index: 7,
397                distance: 0.0
398            })
399        );
400        let three: Vec<u32> = index.k_nearest(p, 3, 1.0).iter().map(|n| n.index).collect();
401        assert_eq!(three, vec![7, 13, 3]);
402        assert_eq!(index.k_nearest(p, 50, 1.0).len(), 4);
403        assert_eq!(index.within_ball(p, 1.0), vec![3, 7, 13, 17]);
404        ogeom_core::parallel::set_threads(4);
405        let many = index.nearest_many(&vec![p; 600], 0.5);
406        ogeom_core::parallel::set_threads(0);
407        assert!(many.iter().all(|n| *n
408            == Some(Neighbour {
409                index: 7,
410                distance: 0.0
411            })));
412    }
413
414    #[test]
415    fn an_empty_cloud_answers_nothing_and_a_bad_point_is_refused() {
416        let index = PointIndex::build(&[]).unwrap();
417        assert!(index.is_empty());
418        let p = Point::new(0.0, 0.0, 0.0);
419        assert!(index.nearest(p, f64::INFINITY).is_none());
420        assert!(index.k_nearest(p, 3, f64::INFINITY).is_empty());
421        assert!(index.within_ball(p, 1.0).is_empty());
422        assert!(PointIndex::build(&[p, Point::new(0.0, f64::INFINITY, 0.0)]).is_err());
423        let one = PointIndex::build(&[p]).unwrap();
424        assert!(one.k_nearest(p, 0, 1.0).is_empty());
425        assert!(one.nearest(Point::new(2.0, 0.0, 0.0), 1.0).is_none());
426        // A negative reach is no reach, whatever its square.
427        assert!(one.nearest(Point::new(1.0, 0.0, 0.0), -1.0).is_none());
428        assert!(one.within_ball(Point::new(1.0, 0.0, 0.0), -1.0).is_empty());
429    }
430}
ogeom_mesh/point_index.rs

ogeom_mesh/
point_index.rs
