Skip to main content

ruvector_robotics/perception/
clustering.rs

1//! Shared spatial-hash clustering with union-find.
2//!
3//! Used by the obstacle detector, scene graph builder, and perception pipeline
4//! to avoid duplicating the same algorithm.
5//!
6//! ## Optimizations
7//!
8//! - Union-by-rank prevents tree degeneration, keeping `find` near O(α(n)).
9//! - Path halving in `find` for efficient path compression.
10//! - `#[inline]` on hot helpers to ensure inlining in tight loops.
11
12use crate::bridge::{Point3D, PointCloud};
13use std::collections::HashMap;
14
15/// Cluster a point cloud using spatial hashing and union-find over adjacent cells.
16///
17/// Points are binned into a 3-D grid with the given `cell_size`. Cells that
18/// share a face, edge, or corner (26-neighbourhood) are merged via union-find.
19/// Returns the resulting groups as separate point vectors.
20pub fn cluster_point_cloud(cloud: &PointCloud, cell_size: f64) -> Vec<Vec<Point3D>> {
21    if cloud.points.is_empty() || cell_size <= 0.0 {
22        return Vec::new();
23    }
24
25    // 1. Map each point to a grid cell.
26    let mut cell_map: HashMap<(i64, i64, i64), Vec<usize>> = HashMap::new();
27    for (idx, p) in cloud.points.iter().enumerate() {
28        let key = cell_key(p, cell_size);
29        cell_map.entry(key).or_default().push(idx);
30    }
31
32    // 2. Build union-find over cells (with rank for balanced merges).
33    let cells: Vec<(i64, i64, i64)> = cell_map.keys().copied().collect();
34    let cell_count = cells.len();
35    let cell_idx: HashMap<(i64, i64, i64), usize> = cells
36        .iter()
37        .enumerate()
38        .map(|(i, &k)| (k, i))
39        .collect();
40
41    let mut parent: Vec<usize> = (0..cell_count).collect();
42    let mut rank: Vec<u8> = vec![0; cell_count];
43
44    for &(cx, cy, cz) in &cells {
45        let a = cell_idx[&(cx, cy, cz)];
46        for dx in -1..=1_i64 {
47            for dy in -1..=1_i64 {
48                for dz in -1..=1_i64 {
49                    let neighbor = (cx + dx, cy + dy, cz + dz);
50                    if let Some(&b) = cell_idx.get(&neighbor) {
51                        uf_union(&mut parent, &mut rank, a, b);
52                    }
53                }
54            }
55        }
56    }
57
58    // 3. Group points by their root representative.
59    let mut groups: HashMap<usize, Vec<Point3D>> = HashMap::new();
60    for (key, point_indices) in &cell_map {
61        let ci = cell_idx[key];
62        let root = uf_find(&mut parent, ci);
63        let entry = groups.entry(root).or_default();
64        for &pi in point_indices {
65            entry.push(cloud.points[pi]);
66        }
67    }
68
69    groups.into_values().collect()
70}
71
72/// Compute the grid cell key for a point.
73#[inline]
74fn cell_key(p: &Point3D, cell_size: f64) -> (i64, i64, i64) {
75    (
76        (p.x as f64 / cell_size).floor() as i64,
77        (p.y as f64 / cell_size).floor() as i64,
78        (p.z as f64 / cell_size).floor() as i64,
79    )
80}
81
82/// Path-compressing find (path halving).
83#[inline]
84fn uf_find(parent: &mut [usize], mut i: usize) -> usize {
85    while parent[i] != i {
86        parent[i] = parent[parent[i]];
87        i = parent[i];
88    }
89    i
90}
91
92/// Union by rank: attaches the shorter tree under the taller root.
93#[inline]
94fn uf_union(parent: &mut [usize], rank: &mut [u8], a: usize, b: usize) {
95    let ra = uf_find(parent, a);
96    let rb = uf_find(parent, b);
97    if ra != rb {
98        match rank[ra].cmp(&rank[rb]) {
99            std::cmp::Ordering::Less => parent[ra] = rb,
100            std::cmp::Ordering::Greater => parent[rb] = ra,
101            std::cmp::Ordering::Equal => {
102                parent[rb] = ra;
103                rank[ra] += 1;
104            }
105        }
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    fn make_cloud(pts: &[[f32; 3]]) -> PointCloud {
114        let points: Vec<Point3D> = pts.iter().map(|a| Point3D::new(a[0], a[1], a[2])).collect();
115        PointCloud::new(points, 0)
116    }
117
118    #[test]
119    fn test_empty_cloud() {
120        let cloud = PointCloud::default();
121        let clusters = cluster_point_cloud(&cloud, 1.0);
122        assert!(clusters.is_empty());
123    }
124
125    #[test]
126    fn test_single_cluster() {
127        let cloud = make_cloud(&[
128            [1.0, 1.0, 0.0],
129            [1.1, 1.0, 0.0],
130            [1.0, 1.1, 0.0],
131        ]);
132        let clusters = cluster_point_cloud(&cloud, 0.5);
133        assert_eq!(clusters.len(), 1);
134        assert_eq!(clusters[0].len(), 3);
135    }
136
137    #[test]
138    fn test_two_clusters() {
139        let cloud = make_cloud(&[
140            [0.0, 0.0, 0.0],
141            [0.1, 0.0, 0.0],
142            [10.0, 10.0, 0.0],
143            [10.1, 10.0, 0.0],
144        ]);
145        let clusters = cluster_point_cloud(&cloud, 0.5);
146        assert_eq!(clusters.len(), 2);
147    }
148
149    #[test]
150    fn test_negative_coordinates() {
151        let cloud = make_cloud(&[
152            [-1.0, -1.0, 0.0],
153            [-0.9, -1.0, 0.0],
154            [1.0, 1.0, 0.0],
155        ]);
156        let clusters = cluster_point_cloud(&cloud, 0.5);
157        assert_eq!(clusters.len(), 2);
158    }
159}