Skip to main content

spatialrust_features/
neighborhood.rs

1use spatialrust_core::{HasPositions3, SpatialError, SpatialResult};
2use spatialrust_search::{KdTree, NearestNeighborIndex, RadiusSearchIndex};
3
4/// Neighborhood query abstraction for feature estimation.
5pub trait NeighborhoodProvider {
6    /// Returns up to `k` neighbor indices excluding the query point itself.
7    fn query_k(&self, index: usize, k: usize) -> SpatialResult<Vec<usize>>;
8
9    /// Returns all neighbor indices within `radius`, excluding the query point.
10    fn query_radius(&self, index: usize, radius: f32) -> SpatialResult<Vec<usize>>;
11}
12
13/// KD-tree backed neighborhood provider.
14#[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    /// Builds a neighborhood provider from a point cloud.
24    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}