spatialrust_search/
uniform_grid.rs

1use std::collections::HashMap;
2
3use spatialrust_core::{SpatialError, SpatialResult};
4
5/// Upper bound on dense grid cells; callers should fall back when exceeded.
6pub const MAX_UNIFORM_GRID_CELLS: u64 = 64_000_000;
7
8/// Returns grid origin (min corner) and cell counts for cell size `radius`.
9pub fn grid_bounds(
10    x: &[f32],
11    y: &[f32],
12    z: &[f32],
13    cell_size: f32,
14) -> SpatialResult<([f32; 3], [u32; 3])> {
15    let mut min = [f32::INFINITY; 3];
16    let mut max = [f32::NEG_INFINITY; 3];
17    for index in 0..x.len() {
18        for (axis, value) in [x[index], y[index], z[index]].into_iter().enumerate() {
19            min[axis] = min[axis].min(value);
20            max[axis] = max[axis].max(value);
21        }
22    }
23    let inv_cell = 1.0 / cell_size;
24    let mut dims = [0u32; 3];
25    for axis in 0..3 {
26        let span = ((max[axis] - min[axis]) * inv_cell).floor() as i64 + 1;
27        dims[axis] = span.max(1) as u32;
28    }
29    let cells = dims[0] as u64 * dims[1] as u64 * dims[2] as u64;
30    if cells > MAX_UNIFORM_GRID_CELLS {
31        return Err(SpatialError::InvalidArgument(format!(
32            "grid would need {cells} cells (cap {MAX_UNIFORM_GRID_CELLS}); use a larger radius or the CPU path"
33        )));
34    }
35    Ok((min, dims))
36}
37
38/// Returns whether a uniform grid with the given cell size fits within the cell cap.
39pub fn uniform_grid_fits(x: &[f32], y: &[f32], z: &[f32], cell_size: f32) -> bool {
40    grid_bounds(x, y, z, cell_size).is_ok()
41}
42
43/// Counting-sort points into grid cells, returning sorted indices and CSR offsets.
44pub fn build_grid(
45    x: &[f32],
46    y: &[f32],
47    z: &[f32],
48    origin: [f32; 3],
49    dims: [u32; 3],
50    cell_size: f32,
51) -> (Vec<u32>, Vec<u32>) {
52    let inv_cell = 1.0 / cell_size;
53    let n = x.len();
54    let num_cells = dims[0] as usize * dims[1] as usize * dims[2] as usize;
55
56    let cell_of = |index: usize| -> usize {
57        let cx = (((x[index] - origin[0]) * inv_cell).floor() as i64).clamp(0, dims[0] as i64 - 1)
58            as usize;
59        let cy = (((y[index] - origin[1]) * inv_cell).floor() as i64).clamp(0, dims[1] as i64 - 1)
60            as usize;
61        let cz = (((z[index] - origin[2]) * inv_cell).floor() as i64).clamp(0, dims[2] as i64 - 1)
62            as usize;
63        (cz * dims[1] as usize + cy) * dims[0] as usize + cx
64    };
65
66    let mut counts = vec![0u32; num_cells + 1];
67    for index in 0..n {
68        counts[cell_of(index)] += 1;
69    }
70    let mut acc = 0u32;
71    for slot in counts.iter_mut() {
72        let c = *slot;
73        *slot = acc;
74        acc += c;
75    }
76    let cell_start = counts;
77
78    let mut cursor = cell_start.clone();
79    let mut sorted = vec![0u32; n];
80    for index in 0..n {
81        let cell = cell_of(index);
82        let slot = cursor[cell];
83        sorted[slot as usize] = index as u32;
84        cursor[cell] = slot + 1;
85    }
86    (sorted, cell_start)
87}
88
89/// Connected-component roots via uniform-grid union-find (minimum index per component).
90pub fn euclidean_cluster_roots(
91    x: &[f32],
92    y: &[f32],
93    z: &[f32],
94    cluster_tolerance: f32,
95) -> SpatialResult<Vec<u32>> {
96    if x.len() != y.len() || x.len() != z.len() {
97        return Err(SpatialError::InvalidArgument("xyz arrays must have equal length".to_owned()));
98    }
99    let point_count = x.len();
100    if point_count == 0 {
101        return Ok(Vec::new());
102    }
103    if cluster_tolerance <= 0.0 || cluster_tolerance.is_nan() {
104        return Err(SpatialError::InvalidArgument("cluster_tolerance must be positive".to_owned()));
105    }
106
107    let (origin, dims) = grid_bounds(x, y, z, cluster_tolerance)?;
108    let (sorted, cell_start) = build_grid(x, y, z, origin, dims, cluster_tolerance);
109    let radius_sq = cluster_tolerance * cluster_tolerance;
110
111    #[cfg(feature = "parallel")]
112    if point_count >= PARALLEL_CLUSTER_MIN_POINTS {
113        return Ok(cluster_roots_parallel(
114            point_count,
115            x,
116            y,
117            z,
118            origin,
119            dims,
120            cluster_tolerance,
121            radius_sq,
122            &sorted,
123            &cell_start,
124        ));
125    }
126
127    Ok(cluster_roots_sequential(
128        point_count,
129        x,
130        y,
131        z,
132        origin,
133        dims,
134        cluster_tolerance,
135        radius_sq,
136        &sorted,
137        &cell_start,
138    ))
139}
140
141/// Computes Euclidean component roots from pre-built sparse grid segments.
142///
143/// This is the shared component-labeling phase for GPU-backed clustering: a
144/// backend may build and sort the sparse grid on an accelerator, then pass the
145/// compact segment metadata here for deterministic host-side union-find. The
146/// distance predicate and minimum-root semantics are identical to
147/// [`euclidean_cluster_roots`].
148pub fn euclidean_cluster_roots_from_segments(
149    x: &[f32],
150    y: &[f32],
151    z: &[f32],
152    cluster_tolerance: f32,
153    keys: &[(i64, i64, i64)],
154    point_indices: &[u32],
155    cell_starts: &[u32],
156    cell_counts: &[u32],
157) -> SpatialResult<Vec<u32>> {
158    if x.len() != y.len() || x.len() != z.len() {
159        return Err(SpatialError::InvalidArgument("xyz arrays must have equal length".to_owned()));
160    }
161    if cluster_tolerance <= 0.0 || cluster_tolerance.is_nan() {
162        return Err(SpatialError::InvalidArgument("cluster_tolerance must be positive".to_owned()));
163    }
164    if keys.len() != cell_starts.len() || keys.len() != cell_counts.len() {
165        return Err(SpatialError::BufferLengthMismatch {
166            expected: keys.len(),
167            found: cell_starts.len().max(cell_counts.len()),
168        });
169    }
170    if point_indices.len() != x.len() {
171        return Err(SpatialError::BufferLengthMismatch {
172            expected: x.len(),
173            found: point_indices.len(),
174        });
175    }
176    if x.is_empty() {
177        return Ok(Vec::new());
178    }
179
180    let mut cell_by_key = HashMap::with_capacity(keys.len());
181    for (cell, &key) in keys.iter().enumerate() {
182        if cell_by_key.insert(key, cell).is_some() {
183            return Err(SpatialError::InvalidArgument(
184                "sparse grid segments contain duplicate cell keys".to_owned(),
185            ));
186        }
187    }
188
189    let mut seen = vec![false; x.len()];
190    for cell in 0..keys.len() {
191        let start = cell_starts[cell] as usize;
192        let end = start.checked_add(cell_counts[cell] as usize).ok_or_else(|| {
193            SpatialError::InvalidArgument("sparse grid segment overflow".to_owned())
194        })?;
195        if end > point_indices.len() {
196            return Err(SpatialError::InvalidArgument(
197                "sparse grid segment exceeds point index buffer".to_owned(),
198            ));
199        }
200        for &point in &point_indices[start..end] {
201            let index = point as usize;
202            if index >= x.len() || std::mem::replace(&mut seen[index], true) {
203                return Err(SpatialError::InvalidArgument(
204                    "sparse grid point indices are not a permutation".to_owned(),
205                ));
206            }
207        }
208    }
209    if seen.iter().any(|present| !present) {
210        return Err(SpatialError::InvalidArgument(
211            "sparse grid point indices do not cover the input cloud".to_owned(),
212        ));
213    }
214
215    let radius_sq = cluster_tolerance * cluster_tolerance;
216    let mut parent: Vec<u32> = (0..x.len() as u32).collect();
217    for cell in 0..keys.len() {
218        let key = keys[cell];
219        let start = cell_starts[cell] as usize;
220        let end = start + cell_counts[cell] as usize;
221        for dz in -1_i64..=1 {
222            for dy in -1_i64..=1 {
223                for dx in -1_i64..=1 {
224                    let neighbor_key = (
225                        key.0.saturating_add(dx),
226                        key.1.saturating_add(dy),
227                        key.2.saturating_add(dz),
228                    );
229                    let Some(&neighbor_cell) = cell_by_key.get(&neighbor_key) else {
230                        continue;
231                    };
232                    let neighbor_start = cell_starts[neighbor_cell] as usize;
233                    let neighbor_end = neighbor_start + cell_counts[neighbor_cell] as usize;
234                    for &a in &point_indices[start..end] {
235                        let a = a as usize;
236                        for &b in &point_indices[neighbor_start..neighbor_end] {
237                            let b = b as usize;
238                            let dx = x[b] - x[a];
239                            let dy = y[b] - y[a];
240                            let dz = z[b] - z[a];
241                            if dx * dx + dy * dy + dz * dz <= radius_sq {
242                                union_min_root(&mut parent, a as u32, b as u32);
243                            }
244                        }
245                    }
246                }
247            }
248        }
249    }
250
251    Ok(compress_roots(&mut parent, x.len()))
252}
253
254/// Minimum point count before the `parallel` feature uses threaded union-find.
255#[cfg(feature = "parallel")]
256const PARALLEL_CLUSTER_MIN_POINTS: usize = 4_096;
257
258#[allow(clippy::too_many_arguments)]
259fn cluster_roots_sequential(
260    point_count: usize,
261    x: &[f32],
262    y: &[f32],
263    z: &[f32],
264    origin: [f32; 3],
265    dims: [u32; 3],
266    cluster_tolerance: f32,
267    radius_sq: f32,
268    sorted: &[u32],
269    cell_start: &[u32],
270) -> Vec<u32> {
271    let mut parent: Vec<u32> = (0..point_count as u32).collect();
272
273    for index in 0..point_count {
274        for neighbor in grid_radius_neighbors(
275            index,
276            x,
277            y,
278            z,
279            origin,
280            dims,
281            cluster_tolerance,
282            radius_sq,
283            sorted,
284            cell_start,
285        ) {
286            union_min_root(&mut parent, index as u32, neighbor as u32);
287        }
288    }
289
290    compress_roots(&mut parent, point_count)
291}
292
293#[cfg(feature = "parallel")]
294#[allow(clippy::too_many_arguments)]
295fn cluster_roots_parallel(
296    point_count: usize,
297    x: &[f32],
298    y: &[f32],
299    z: &[f32],
300    origin: [f32; 3],
301    dims: [u32; 3],
302    cluster_tolerance: f32,
303    radius_sq: f32,
304    sorted: &[u32],
305    cell_start: &[u32],
306) -> Vec<u32> {
307    use std::sync::atomic::{AtomicU32, Ordering};
308    use std::sync::Arc;
309
310    let parent: Arc<[AtomicU32]> =
311        (0..point_count as u32).map(AtomicU32::new).collect::<Vec<_>>().into();
312
313    let thread_count =
314        std::thread::available_parallelism().map_or(1, |count| count.get()).min(point_count).max(1);
315    let chunk = point_count.div_ceil(thread_count);
316
317    std::thread::scope(|scope| {
318        for thread in 0..thread_count {
319            let start = thread * chunk;
320            if start >= point_count {
321                break;
322            }
323            let end = (start + chunk).min(point_count);
324            let parent = Arc::clone(&parent);
325            scope.spawn(move || {
326                for index in start..end {
327                    for neighbor in grid_radius_neighbors(
328                        index,
329                        x,
330                        y,
331                        z,
332                        origin,
333                        dims,
334                        cluster_tolerance,
335                        radius_sq,
336                        sorted,
337                        cell_start,
338                    ) {
339                        atomic_union_min_root(&parent, index as u32, neighbor as u32);
340                    }
341                }
342            });
343        }
344    });
345
346    let mut parent_vec: Vec<u32> = parent.iter().map(|slot| slot.load(Ordering::Relaxed)).collect();
347    compress_roots(&mut parent_vec, point_count)
348}
349
350#[cfg(feature = "parallel")]
351fn atomic_union_min_root(parent: &[std::sync::atomic::AtomicU32], a: u32, b: u32) {
352    use std::sync::atomic::Ordering;
353
354    let mut ra = atomic_find_root(parent, a);
355    let mut rb = atomic_find_root(parent, b);
356    while ra != rb {
357        let (min_root, max_root) = if ra < rb { (ra, rb) } else { (rb, ra) };
358        match parent[max_root as usize].compare_exchange(
359            max_root,
360            min_root,
361            Ordering::Relaxed,
362            Ordering::Relaxed,
363        ) {
364            Ok(_) => return,
365            Err(current) => {
366                if current == min_root {
367                    return;
368                }
369                ra = atomic_find_root(parent, a);
370                rb = atomic_find_root(parent, b);
371            }
372        }
373    }
374}
375
376#[cfg(feature = "parallel")]
377fn atomic_find_root(parent: &[std::sync::atomic::AtomicU32], mut index: u32) -> u32 {
378    use std::sync::atomic::Ordering;
379
380    loop {
381        let parent_index = parent[index as usize].load(Ordering::Relaxed);
382        if parent_index == index {
383            return index;
384        }
385        index = parent_index;
386    }
387}
388
389fn compress_roots(parent: &mut [u32], point_count: usize) -> Vec<u32> {
390    let mut roots = vec![0u32; point_count];
391    for (index, root) in roots.iter_mut().enumerate() {
392        *root = find_root(parent, index as u32);
393    }
394    roots
395}
396
397fn find_root(parent: &mut [u32], mut index: u32) -> u32 {
398    let mut root = index;
399    while parent[root as usize] != root {
400        root = parent[root as usize];
401    }
402    while parent[index as usize] != root {
403        let next = parent[index as usize];
404        parent[index as usize] = root;
405        index = next;
406    }
407    root
408}
409
410fn union_min_root(parent: &mut [u32], a: u32, b: u32) {
411    let ra = find_root_readonly(parent, a);
412    let rb = find_root_readonly(parent, b);
413    if ra == rb {
414        return;
415    }
416    let (min_root, max_root) = if ra < rb { (ra, rb) } else { (rb, ra) };
417    parent[max_root as usize] = min_root;
418}
419
420fn find_root_readonly(parent: &[u32], mut index: u32) -> u32 {
421    while parent[index as usize] != index {
422        index = parent[index as usize];
423    }
424    index
425}
426
427#[allow(clippy::too_many_arguments)]
428fn grid_radius_neighbors<'a>(
429    index: usize,
430    x: &'a [f32],
431    y: &'a [f32],
432    z: &'a [f32],
433    origin: [f32; 3],
434    dims: [u32; 3],
435    tolerance: f32,
436    radius_sq: f32,
437    sorted: &'a [u32],
438    cell_start: &'a [u32],
439) -> GridRadiusNeighbors<'a> {
440    GridRadiusNeighbors {
441        index,
442        x,
443        y,
444        z,
445        dims,
446        radius_sq,
447        sorted,
448        cell_start,
449        cell: 0,
450        slot: 0,
451        end: 0,
452        dz: -1,
453        dy: -1,
454        dx: -1,
455        cx: cell_coord(x[index], origin[0], 1.0 / tolerance, dims[0]),
456        cy: cell_coord(y[index], origin[1], 1.0 / tolerance, dims[1]),
457        cz: cell_coord(z[index], origin[2], 1.0 / tolerance, dims[2]),
458        started: false,
459    }
460}
461
462struct GridRadiusNeighbors<'a> {
463    index: usize,
464    x: &'a [f32],
465    y: &'a [f32],
466    z: &'a [f32],
467    dims: [u32; 3],
468    radius_sq: f32,
469    sorted: &'a [u32],
470    cell_start: &'a [u32],
471    cell: usize,
472    slot: u32,
473    end: u32,
474    dz: i32,
475    dy: i32,
476    dx: i32,
477    cx: i32,
478    cy: i32,
479    cz: i32,
480    started: bool,
481}
482
483impl Iterator for GridRadiusNeighbors<'_> {
484    type Item = usize;
485
486    fn next(&mut self) -> Option<Self::Item> {
487        loop {
488            if !self.started {
489                self.started = true;
490                if !self.advance_cell() {
491                    return None;
492                }
493            } else if self.slot >= self.end {
494                if !self.advance_cell() {
495                    return None;
496                }
497            } else {
498                let neighbor = self.sorted[self.slot as usize] as usize;
499                self.slot += 1;
500                if neighbor == self.index {
501                    continue;
502                }
503                let dx = self.x[neighbor] - self.x[self.index];
504                let dy = self.y[neighbor] - self.y[self.index];
505                let dz = self.z[neighbor] - self.z[self.index];
506                if dx * dx + dy * dy + dz * dz <= self.radius_sq {
507                    return Some(neighbor);
508                }
509            }
510        }
511    }
512}
513
514impl GridRadiusNeighbors<'_> {
515    fn advance_cell(&mut self) -> bool {
516        let dimx = self.dims[0] as i32;
517        let dimy = self.dims[1] as i32;
518        let dimz = self.dims[2] as i32;
519
520        loop {
521            if self.dz > 1 {
522                return false;
523            }
524            let nx = self.cx + self.dx;
525            let ny = self.cy + self.dy;
526            let nz = self.cz + self.dz;
527            if nx >= 0 && ny >= 0 && nz >= 0 && nx < dimx && ny < dimy && nz < dimz {
528                self.cell = cell_index(nx, ny, nz, self.dims[0], self.dims[1]) as usize;
529                self.slot = self.cell_start[self.cell];
530                self.end = self.cell_start[self.cell + 1];
531                self.bump_offset();
532                return true;
533            }
534            self.bump_offset();
535        }
536    }
537
538    fn bump_offset(&mut self) {
539        self.dx += 1;
540        if self.dx > 1 {
541            self.dx = -1;
542            self.dy += 1;
543            if self.dy > 1 {
544                self.dy = -1;
545                self.dz += 1;
546            }
547        }
548    }
549}
550
551fn cell_coord(value: f32, origin: f32, inv_cell: f32, dim: u32) -> i32 {
552    let cell = ((value - origin) * inv_cell).floor() as i32;
553    cell.clamp(0, dim as i32 - 1)
554}
555
556fn cell_index(cx: i32, cy: i32, cz: i32, dimx: u32, dimy: u32) -> u32 {
557    (cz as u32 * dimy + cy as u32) * dimx + cx as u32
558}
559
560#[cfg(test)]
561mod tests {
562    use super::{euclidean_cluster_roots, euclidean_cluster_roots_from_segments};
563
564    #[test]
565    fn long_chain_is_one_component() {
566        let len = 300;
567        let spacing = 1.0;
568        let mut x = Vec::with_capacity(len);
569        for index in 0..len {
570            x.push(index as f32 * spacing);
571        }
572        let y = vec![0.0_f32; len];
573        let z = vec![0.0_f32; len];
574        let roots = euclidean_cluster_roots(&x, &y, &z, 1.5).unwrap();
575        assert!(roots.iter().all(|&root| root == 0));
576    }
577
578    #[test]
579    fn prebuilt_sparse_segments_match_direct_grid_roots() {
580        let x = [0.0_f32, 0.5, 2.0];
581        let y = [0.0_f32; 3];
582        let z = [0.0_f32; 3];
583        let expected = euclidean_cluster_roots(&x, &y, &z, 1.0).unwrap();
584        let actual = euclidean_cluster_roots_from_segments(
585            &x,
586            &y,
587            &z,
588            1.0,
589            &[(0, 0, 0), (2, 0, 0)],
590            &[0, 1, 2],
591            &[0, 2],
592            &[2, 1],
593        )
594        .unwrap();
595
596        assert_eq!(actual, expected);
597    }
598}
spatialrust_search/uniform_grid.rs

spatialrust_search/
uniform_grid.rs
