spatialrust_features/
neighborhood.rs1use spatialrust_core::{HasPositions3, SpatialError, SpatialResult};
2use spatialrust_search::{KdTree, NearestNeighborIndex, RadiusSearchIndex};
3
4pub trait NeighborhoodProvider {
6 fn query_k(&self, index: usize, k: usize) -> SpatialResult<Vec<usize>>;
8
9 fn query_radius(&self, index: usize, radius: f32) -> SpatialResult<Vec<usize>>;
11}
12
13#[derive(Clone, Debug)]
15pub struct KdTreeNeighborhood {
16 tree: KdTree,
17 x: Vec<f32>,
18 y: Vec<f32>,
19 z: Vec<f32>,
20}
21
22impl KdTreeNeighborhood {
23 pub fn from_point_cloud(cloud: &spatialrust_core::PointCloud) -> SpatialResult<Self> {
25 let (x, y, z) = cloud.positions3()?;
26 Ok(Self { tree: KdTree::from_slices(x, y, z), x: x.to_vec(), y: y.to_vec(), z: z.to_vec() })
27 }
28}
29
30impl NeighborhoodProvider for KdTreeNeighborhood {
31 fn query_k(&self, index: usize, k: usize) -> SpatialResult<Vec<usize>> {
32 if index >= self.x.len() {
33 return Err(SpatialError::InvalidArgument(format!(
34 "neighbor query index out of bounds: {index}"
35 )));
36 }
37 if k == 0 {
38 return Ok(Vec::new());
39 }
40
41 let neighbors =
42 self.tree.nearest_k(self.x[index], self.y[index], self.z[index], k.saturating_add(1));
43 Ok(neighbors
44 .into_iter()
45 .map(|neighbor| neighbor.index)
46 .filter(|&candidate| candidate != index)
47 .take(k)
48 .collect())
49 }
50
51 fn query_radius(&self, index: usize, radius: f32) -> SpatialResult<Vec<usize>> {
52 if index >= self.x.len() {
53 return Err(SpatialError::InvalidArgument(format!(
54 "neighbor query index out of bounds: {index}"
55 )));
56 }
57 if radius < 0.0 {
58 return Err(SpatialError::InvalidArgument("radius must be non-negative".to_owned()));
59 }
60
61 let neighbors =
62 self.tree.radius_search(self.x[index], self.y[index], self.z[index], radius);
63 Ok(neighbors
64 .into_iter()
65 .map(|neighbor| neighbor.index)
66 .filter(|&candidate| candidate != index)
67 .collect())
68 }
69}