cumulus3d_align/
kdtree.rs

1/*
2 * kdtree.rs
3 *
4 * Copyright (c) 2026 ParkSangWoo
5 *
6 * Author(s):
7 *
8 *      ParkSangWoo <dev.parksangwoo@gmail.com>
9 *
10 *
11 * This file is part of cumulus3d.
12 *
13 * Licensed under the Apache License, Version 2.0 (the "License");
14 * you may not use this file except in compliance with the License.
15 * You may obtain a copy of the License at
16 *
17 *      http://www.apache.org/licenses/LICENSE-2.0
18 *
19 * Unless required by applicable law or agreed to in writing, software
20 * distributed under the License is distributed on an "AS IS" BASIS,
21 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
22 * See the License for the specific language governing permissions and
23 * limitations under the License.
24 *
25 * SPDX-License-Identifier: Apache-2.0
26 */
27
28//! 3D KD-tree (own implementation). For proximity masking against the refined cloud. Queries are parallel with rayon.
29
30use rayon::prelude::*;
31
32const LEAF_SIZE: usize = 16;
33
34#[derive(Clone, Copy, Debug)]
35enum Node {
36    Leaf { start: u32, end: u32 },
37    Split { dim: u8, val: f64, left: u32, right: u32 },
38}
39
40/// Static KD-tree. Points are stored rearranged in tree order, together with their original indices.
41#[derive(Clone, Debug, Default)]
42pub struct KdTree {
43    pts: Vec<[f64; 3]>,
44    orig: Vec<u32>,
45    nodes: Vec<Node>,
46}
47
48#[inline]
49fn d2(a: &[f64; 3], b: &[f64; 3]) -> f64 {
50    let (x, y, z) = (a[0] - b[0], a[1] - b[1], a[2] - b[2]);
51    x * x + y * y + z * z
52}
53
54impl KdTree {
55    /// Creates from f64 points.
56    pub fn new(points: impl IntoIterator<Item = [f64; 3]>) -> Self {
57        let pts: Vec<[f64; 3]> = points.into_iter().collect();
58        assert!(pts.len() < u32::MAX as usize, "점이 너무 많음");
59        let mut idx: Vec<u32> = (0..pts.len() as u32).collect();
60        let mut nodes = Vec::new();
61        if !pts.is_empty() {
62            build(&pts, &mut idx, 0, &mut nodes);
63        }
64        let reordered = idx.iter().map(|&i| pts[i as usize]).collect();
65        Self { pts: reordered, orig: idx, nodes }
66    }
67    /// Creates from f32 points (PLY).
68    pub fn from_f32(points: &[[f32; 3]]) -> Self {
69        Self::new(points.iter().map(|p| [p[0] as f64, p[1] as f64, p[2] as f64]))
70    }
71    /// Number of points.
72    pub fn len(&self) -> usize {
73        self.pts.len()
74    }
75    /// True if there are no points.
76    pub fn is_empty(&self) -> bool {
77        self.pts.is_empty()
78    }
79
80    /// Nearest point (original index, distance²). None for an empty tree.
81    pub fn nearest(&self, q: &[f64; 3]) -> Option<(usize, f64)> {
82        if self.nodes.is_empty() {
83            return None;
84        }
85        let mut best = (usize::MAX, f64::INFINITY);
86        self.nearest_rec(0, q, &mut best);
87        Some((self.orig[best.0] as usize, best.1))
88    }
89
90    fn nearest_rec(&self, node: u32, q: &[f64; 3], best: &mut (usize, f64)) {
91        match self.nodes[node as usize] {
92            Node::Leaf { start, end } => {
93                for i in start as usize..end as usize {
94                    let d = d2(&self.pts[i], q);
95                    if d < best.1 {
96                        *best = (i, d);
97                    }
98                }
99            }
100            Node::Split { dim, val, left, right } => {
101                let diff = q[dim as usize] - val;
102                let (near, far) = if diff < 0.0 { (left, right) } else { (right, left) };
103                self.nearest_rec(near, q, best);
104                if diff * diff < best.1 {
105                    self.nearest_rec(far, q, best);
106                }
107            }
108        }
109    }
110
111    /// Whether any point lies within radius r (distance ≤ r) (early exit).
112    pub fn any_within(&self, q: &[f64; 3], r: f64) -> bool {
113        !self.nodes.is_empty() && self.any_rec(0, q, r * r)
114    }
115
116    fn any_rec(&self, node: u32, q: &[f64; 3], r2: f64) -> bool {
117        match self.nodes[node as usize] {
118            Node::Leaf { start, end } => self.pts[start as usize..end as usize].iter().any(|p| d2(p, q) <= r2),
119            Node::Split { dim, val, left, right } => {
120                let diff = q[dim as usize] - val;
121                let (near, far) = if diff < 0.0 { (left, right) } else { (right, left) };
122                self.any_rec(near, q, r2) || (diff * diff <= r2 && self.any_rec(far, q, r2))
123            }
124        }
125    }
126
127    /// All points within radius r (original index, order unspecified).
128    pub fn within_radius(&self, q: &[f64; 3], r: f64) -> Vec<usize> {
129        let mut out = Vec::new();
130        if !self.nodes.is_empty() {
131            self.radius_rec(0, q, r * r, &mut out);
132        }
133        out
134    }
135
136    fn radius_rec(&self, node: u32, q: &[f64; 3], r2: f64, out: &mut Vec<usize>) {
137        match self.nodes[node as usize] {
138            Node::Leaf { start, end } => {
139                for i in start as usize..end as usize {
140                    if d2(&self.pts[i], q) <= r2 {
141                        out.push(self.orig[i] as usize);
142                    }
143                }
144            }
145            Node::Split { dim, val, left, right } => {
146                let diff = q[dim as usize] - val;
147                if diff < 0.0 || diff * diff <= r2 {
148                    self.radius_rec(left, q, r2, out);
149                }
150                if diff >= 0.0 || diff * diff <= r2 {
151                    self.radius_rec(right, q, r2, out);
152                }
153            }
154        }
155    }
156
157    /// Parallel nearest-neighbor query.
158    pub fn nearest_many(&self, queries: &[[f64; 3]]) -> Vec<Option<(usize, f64)>> {
159        queries.par_iter().map(|q| self.nearest(q)).collect()
160    }
161
162    /// Parallel radius test (f32 queries). `result[i]` = whether query i has a point within r.
163    pub fn any_within_many_f32(&self, queries: &[[f32; 3]], r: f64) -> Vec<bool> {
164        queries.par_iter().map(|q| self.any_within(&[q[0] as f64, q[1] as f64, q[2] as f64], r)).collect()
165    }
166}
167
168/// Builds nodes for the idx range and returns the node index. `offset` = start position of the idx range within the whole.
169fn build(pts: &[[f64; 3]], idx: &mut [u32], offset: usize, nodes: &mut Vec<Node>) -> u32 {
170    let me = nodes.len() as u32;
171    let n = idx.len();
172    if n <= LEAF_SIZE {
173        nodes.push(Node::Leaf { start: offset as u32, end: (offset + n) as u32 });
174        return me;
175    }
176    // Split along the axis with the largest extent.
177    let (mut lo, mut hi) = ([f64::INFINITY; 3], [f64::NEG_INFINITY; 3]);
178    for &i in idx.iter() {
179        let p = &pts[i as usize];
180        for d in 0..3 {
181            lo[d] = lo[d].min(p[d]);
182            hi[d] = hi[d].max(p[d]);
183        }
184    }
185    let dim = (0..3).max_by(|&a, &b| (hi[a] - lo[a]).total_cmp(&(hi[b] - lo[b]))).unwrap_or(0);
186    if hi[dim] - lo[dim] <= 0.0 {
187        // All points at the same position: cannot split further.
188        nodes.push(Node::Leaf { start: offset as u32, end: (offset + n) as u32 });
189        return me;
190    }
191    let mid = n / 2;
192    idx.select_nth_unstable_by(mid, |a, b| pts[*a as usize][dim].total_cmp(&pts[*b as usize][dim]));
193    let val = pts[idx[mid] as usize][dim];
194    // Left ≤ val ≤ right. Pruning uses only the |q_d − val| lower bound, so it is exact even if equal values lie on both sides.
195    nodes.push(Node::Leaf { start: 0, end: 0 }); // reserve a slot
196    let (l, r) = idx.split_at_mut(mid);
197    let left = build(pts, l, offset, nodes);
198    let right = build(pts, r, offset + mid, nodes);
199    nodes[me as usize] = Node::Split { dim: dim as u8, val, left, right };
200    me
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206    use rand::RngExt;
207
208    #[test]
209    fn nearest_matches_brute_force() {
210        let mut rng = cumulus3d_core::ransac::make_rng(Some(21));
211        let pts: Vec<[f64; 3]> =
212            (0..10_000).map(|_| [rng.random_range(-50.0..50.0), rng.random_range(-50.0..50.0), rng.random_range(-5.0..5.0)]).collect();
213        let tree = KdTree::new(pts.iter().copied());
214        let qs: Vec<[f64; 3]> =
215            (0..2000).map(|_| [rng.random_range(-60.0..60.0), rng.random_range(-60.0..60.0), rng.random_range(-10.0..10.0)]).collect();
216        let got = tree.nearest_many(&qs);
217        for (q, g) in qs.iter().zip(&got) {
218            let (bi, bd) = pts.iter().enumerate().map(|(i, p)| (i, d2(p, q))).min_by(|a, b| a.1.total_cmp(&b.1)).unwrap();
219            let (gi, gd) = g.unwrap();
220            assert_eq!(gd, bd);
221            assert!(gi == bi || d2(&pts[gi], q) == bd);
222            // Radius queries also match brute force.
223            let r = 3.0;
224            let mut within = tree.within_radius(q, r);
225            within.sort_unstable();
226            let brute: Vec<usize> = pts.iter().enumerate().filter(|(_, p)| d2(p, q) <= r * r).map(|(i, _)| i).collect();
227            assert_eq!(within, brute);
228            assert_eq!(tree.any_within(q, r), !brute.is_empty());
229        }
230    }
231
232    #[test]
233    fn duplicates_and_empty() {
234        let tree = KdTree::new(std::iter::repeat_n([1.0, 2.0, 3.0], 100));
235        assert_eq!(tree.nearest(&[0.0, 0.0, 0.0]).unwrap().1, 14.0);
236        assert!(KdTree::new(std::iter::empty()).nearest(&[0.0; 3]).is_none());
237        assert!(!KdTree::default().any_within(&[0.0; 3], 1.0));
238    }
239}
cumulus3d_align/kdtree.rs

cumulus3d_align/
kdtree.rs
