diff --git a/src/bench.rs b/src/bench.rs index f6f7ed9..529127f 100644 --- a/src/bench.rs +++ b/src/bench.rs @@ -26,7 +26,7 @@ fn bench_creating_1000_node_tree(b: &mut Bencher) { let points = generate_points(len); b.iter(|| { - kdtree::kdtree::Kdtree::new(&mut points.clone()); + kdtree::kdtree::KdTree::new(&mut points.clone()); }); } @@ -34,7 +34,7 @@ fn bench_single_loop_times_for_1000_node_tree(b: &mut Bencher) { let len = 1000usize; let points = generate_points(len); - let tree = kdtree::kdtree::Kdtree::new(&mut points.clone()); + let tree = kdtree::kdtree::KdTree::new(&mut points.clone()); b.iter(|| tree.nearest_search(&points[0])); @@ -46,14 +46,14 @@ fn bench_creating_1000_000_node_tree(b: &mut Bencher) { let points = generate_points(len); b.iter(|| { - kdtree::kdtree::Kdtree::new(&mut points.clone()); + kdtree::kdtree::KdTree::new(&mut points.clone()); }); } fn bench_adding_same_node_to_1000_tree(b: &mut Bencher) { let len = 1000usize; let mut points = generate_points(len); - let mut tree = kdtree::kdtree::Kdtree::new(&mut points); + let mut tree = kdtree::kdtree::KdTree::new(&mut points); let point = Point3WithId::new(-1 as i32, gen_random(), gen_random(), gen_random()); b.iter(|| { @@ -65,7 +65,7 @@ fn bench_incrementally_building_the_1000_tree(b: &mut Bencher) { b.iter(|| { let len = 1usize; let mut points = generate_points(len); - let mut tree = kdtree::kdtree::Kdtree::new(&mut points); + let mut tree = kdtree::kdtree::KdTree::new(&mut points); for _ in 0 .. 1000 { let point = Point3WithId::new(-1 as i32, gen_random(), gen_random(), gen_random()); tree.insert_node(point); diff --git a/src/kdtree/bounds.rs b/src/kdtree/bounds.rs index 3b870c9..7398088 100644 --- a/src/kdtree/bounds.rs +++ b/src/kdtree/bounds.rs @@ -1,5 +1,6 @@ use ::kdtree::*; +#[derive(Clone, Copy)] pub struct Bounds { pub bounds: [(f64, f64); 3], @@ -8,7 +9,7 @@ pub struct Bounds { } impl Bounds { - pub fn new_from_points(points: &[T]) -> Bounds { + pub fn new_from_points(points: &[T]) -> Bounds { let mut bounds = Bounds { bounds: [(0., 0.), (0., 0.), (0., 0.)], widest_dim: 0, @@ -45,6 +46,7 @@ impl Bounds { bounds: self.bounds.clone(), ..*self }; + cloned.bounds[dimension].1 = value; cloned.calculate_variables(); diff --git a/src/kdtree/mod.rs b/src/kdtree/mod.rs index 5427621..b5f76ce 100644 --- a/src/kdtree/mod.rs +++ b/src/kdtree/mod.rs @@ -9,27 +9,82 @@ use self::distance::*; use std::cmp; -pub trait KdtreePointTrait: Copy + PartialEq { +pub trait KdTreePoint: Copy + PartialEq { + fn dist_1d(left: f64, right: f64, dim: usize) -> f64 { + let diff = left - right; + + diff * diff + } + fn dims(&self) -> &[f64]; + fn dist(&self, other: &Self) -> f64 { + squared_euclidean(self.dims(), other.dims()) + } } -pub struct Kdtree { - nodes: Vec>, +pub struct NearestNeighboursIter<'a, 'b, T> { + range: f64, + kdtree: &'a KdTree, + ref_node: &'b T, + node_stack: Vec, +} + +impl<'a, 'b, T> Iterator for NearestNeighboursIter<'a, 'b, T> + where T: KdTreePoint +{ + type Item = &'a T; + + fn next(&mut self) -> Option { + let p = self.ref_node; + + loop { + let node_idx = self.node_stack.pop()?; + let node = &self.kdtree.nodes[node_idx]; + + let splitting_value = node.split_on; + let point_splitting_dim_value = p.dims()[node.dimension]; + let distance_on_single_dimension = T::dist_1d(splitting_value, point_splitting_dim_value, node.dimension); + + if distance_on_single_dimension <= self.range { + if let Some(idx) = node.left_node { + self.node_stack.push(idx); + } + + if let Some(idx) = node.right_node { + self.node_stack.push(idx); + } + + if p.dist(&node.point) <= self.range { + return Some(&node.point); + } + } else if point_splitting_dim_value <= splitting_value { + if let Some(idx) = node.left_node { + self.node_stack.push(idx); + } + } else { + if let Some(idx) = node.right_node { + self.node_stack.push(idx); + } + } + } + } +} + +pub struct KdTree { + nodes: Vec>, node_adding_dimension: usize, node_depth_during_last_rebuild: usize, current_node_depth: usize, } -impl Kdtree { - pub fn new(mut points: &mut [KdtreePoint]) -> Kdtree { +impl KdTree { + pub fn new(mut points: &mut [KP]) -> Self { if points.len() == 0 { panic!("empty vector point not allowed"); } - - - let mut tree = Kdtree { + let mut tree = KdTree { nodes: vec![], node_adding_dimension: 0, node_depth_during_last_rebuild: 0, @@ -41,8 +96,9 @@ impl Kdtree { tree } - pub fn rebuild_tree(&mut self, points : &mut [KdtreePoint]) { + pub fn rebuild_tree(&mut self, points: &mut [KP]) { self.nodes.clear(); + self.nodes.reserve(points.len()); self.node_depth_during_last_rebuild = 0; self.current_node_depth = 0; @@ -53,50 +109,60 @@ impl Kdtree { /// Can be used if you are sure that the tree is degenerated or if you will never again insert the nodes into the tree. pub fn gather_points_and_rebuild(&mut self) { - let mut points : Vec = vec![]; + let mut points : Vec = vec![]; self.gather_points(0,&mut points); self.rebuild_tree(&mut points); } - pub fn nearest_search(&self, node: &KdtreePoint) -> KdtreePoint - { + pub fn nearest_search(&self, node: &KP) -> KP { let mut nearest_neighbor = 0usize; - let mut best_distance = squared_euclidean(node.dims(), &self.nodes[0].point.dims()); + let mut best_distance = self.nodes[0].point.dist(&node); + self.nearest_search_impl(node, 0usize, &mut best_distance, &mut nearest_neighbor); self.nodes[nearest_neighbor].point } - pub fn has_neighbor_in_range(&self, node: &KdtreePoint, range: f64) -> bool { + pub fn nearest_search_dist<'a, 'b>(&'a self, node: &'b KP, dist: f64) -> NearestNeighboursIter<'a, 'b, KP> { + let mut node_stack = Vec::with_capacity(16); + node_stack.push(0); + + NearestNeighboursIter { + range: dist, + kdtree: &self, + ref_node: node, + node_stack, + } + } + + pub fn has_neighbor_in_range(&self, node: &KP, range: f64) -> bool { let squared_range = range * range; self.distance_squared_to_nearest(node) <= squared_range } - pub fn distance_squared_to_nearest(&self, node: &KdtreePoint) -> f64 { - squared_euclidean(&self.nearest_search(node).dims(), node.dims()) + pub fn distance_squared_to_nearest(&self, node: &KP) -> f64 { + self.nearest_search(node).dist(&node) } - pub fn insert_nodes_and_rebuild(&mut self, nodes_to_add : &mut [KdtreePoint]) { - let mut pts : Vec = vec![]; + pub fn insert_nodes_and_rebuild(&mut self, nodes_to_add : &mut [KP]) { + let mut pts : Vec = vec![]; self.gather_points(0, &mut pts); pts.extend(nodes_to_add.iter()); self.rebuild_tree(&mut pts); } - pub fn insert_node(&mut self, node_to_add : KdtreePoint) { - + pub fn insert_node(&mut self, node_to_add : KP) { let mut current_index = 0; let dimension = self.node_adding_dimension; - let index_of_new_node = self.add_node(node_to_add,dimension,node_to_add.dims()[dimension]); + let index_of_new_node = self.add_node(node_to_add, dimension,node_to_add.dims()[dimension]); self.node_adding_dimension = ( dimension + 1) % node_to_add.dims().len(); let mut should_pop_node = false; let mut depth = 0; loop { - depth +=1 ; let current_node = &mut self.nodes[current_index]; @@ -134,7 +200,7 @@ impl Kdtree { } } - fn nearest_search_impl(&self, p: &KdtreePoint, searched_index: usize, best_distance_squared: &mut f64, best_leaf_found: &mut usize) { + fn nearest_search_impl(&self, p: &KP, searched_index: usize, best_distance_squared: &mut f64, best_leaf_found: &mut usize) { let node = &self.nodes[searched_index]; let splitting_value = node.split_on; @@ -150,14 +216,14 @@ impl Kdtree { self.nearest_search_impl(p, closer_node, best_distance_squared, best_leaf_found); } - let distance = squared_euclidean(p.dims(), node.point.dims()); + let distance = p.dist(&node.point); if distance < *best_distance_squared { *best_distance_squared = distance; *best_leaf_found = searched_index; } if let Some(farther_node) = farther_node { - let distance_on_single_dimension = squared_euclidean(&[splitting_value], &[point_splitting_dim_value]); + let distance_on_single_dimension = KP::dist_1d(splitting_value, point_splitting_dim_value, node.dimension); if distance_on_single_dimension <= *best_distance_squared { self.nearest_search_impl(p, farther_node, best_distance_squared, best_leaf_found); @@ -165,14 +231,14 @@ impl Kdtree { } } - fn add_node(&mut self, p: KdtreePoint, dimension: usize, split_on: f64) -> usize { - let node = KdtreeNode::new(p, dimension, split_on); + fn add_node(&mut self, p: KP, dimension: usize, split_on: f64) -> usize { + let node = KdTreeNode::new(p, dimension, split_on); self.nodes.push(node); self.nodes.len() - 1 } - fn build_tree(&mut self, nodes: &mut [KdtreePoint], bounds: &Bounds, depth : usize) -> usize { + fn build_tree(&mut self, nodes: &mut [KP], bounds: &Bounds, depth : usize) -> usize { let splitting_index = partition::partition_sliding_midpoint(nodes, bounds.get_midvalue_of_widest_dim(), bounds.get_widest_dim()); let pivot_value = nodes[splitting_index].dims()[bounds.get_widest_dim()]; @@ -197,8 +263,9 @@ impl Kdtree { node_id } - fn gather_points(&self, current_index: usize, points : &mut Vec){ + fn gather_points(&self, current_index: usize, points : &mut Vec){ points.push(self.nodes[current_index].point); + if let Some(left_index) = self.nodes[current_index].left_node { self.gather_points(left_index, points); } @@ -209,7 +276,7 @@ impl Kdtree { } } -pub struct KdtreeNode { +pub struct KdTreeNode { left_node: Option, right_node: Option, @@ -218,9 +285,9 @@ pub struct KdtreeNode { split_on: f64 } -impl KdtreeNode { - fn new(p: T, splitting_dimension: usize, split_on_value: f64) -> KdtreeNode { - KdtreeNode { +impl KdTreeNode { + fn new(p: T, splitting_dimension: usize, split_on_value: f64) -> KdTreeNode { + KdTreeNode { left_node: None, right_node: None, @@ -242,7 +309,7 @@ mod tests { fn should_panic_given_empty_vector() { let mut empty_vec: Vec = vec![]; - Kdtree::new(&mut empty_vec); + KdTree::new(&mut empty_vec); } quickcheck! { @@ -257,7 +324,7 @@ mod tests { vec.push(p); } - let tree = Kdtree::new(&mut qc_value_vec_to_2d_points_vec(&xs)); + let tree = KdTree::new(&mut qc_value_vec_to_2d_points_vec(&xs)); let mut to_iterate : Vec = vec![]; to_iterate.push(0); @@ -284,7 +351,7 @@ mod tests { } let point_vec = qc_value_vec_to_2d_points_vec(&xs); - let tree = Kdtree::new(&mut point_vec.clone()); + let tree = KdTree::new(&mut point_vec.clone()); for p in &point_vec { let found_nn = tree.nearest_search(p); @@ -300,7 +367,7 @@ mod tests { fn has_neighbor_in_range() { let mut vec: Vec = vec![Point2WithId::new(0,2.,0.)]; - let tree = Kdtree::new(&mut vec); + let tree = KdTree::new(&mut vec); assert_eq!(false,tree.has_neighbor_in_range(&Point2WithId::new(0,0.,0.), 0.)); assert_eq!(false,tree.has_neighbor_in_range(&Point2WithId::new(0,0.,0.), 1.)); @@ -314,7 +381,7 @@ mod tests { let mut vec = vec![Point2WithId::new(0,0.,0.)]; - let mut tree = Kdtree::new(&mut vec); + let mut tree = KdTree::new(&mut vec); tree.insert_node(Point2WithId::new(0,1.,0.)); tree.insert_node(Point2WithId::new(0,-1.,0.)); @@ -333,7 +400,7 @@ mod tests { fn incremental_add_filters_duplicates() { let mut vec = vec![Point2WithId::new(0,0.,0.)]; - let mut tree = Kdtree::new(&mut vec); + let mut tree = KdTree::new(&mut vec); let node = Point2WithId::new(0,1.,0.); tree.insert_node(node); diff --git a/src/kdtree/partition.rs b/src/kdtree/partition.rs index 2791b1a..a456107 100644 --- a/src/kdtree/partition.rs +++ b/src/kdtree/partition.rs @@ -11,7 +11,7 @@ struct PartitionPointHelper { index_of_splitter: usize, } -fn partition_sliding_midpoint_helper(vec: &mut [T], midpoint_value: f64, partition_on_dimension: usize) -> PartitionPointHelper { +fn partition_sliding_midpoint_helper(vec: &mut [T], midpoint_value: f64, partition_on_dimension: usize) -> PartitionPointHelper { let mut closest_index = 0; let mut closest_distance = (vec[0].dims()[partition_on_dimension] - midpoint_value).abs(); @@ -48,7 +48,7 @@ fn partition_sliding_midpoint_helper(vec: &mut [T], midpoin } } -pub fn partition_sliding_midpoint(vec: &mut [T], midpoint_value: f64, partition_on_dimension: usize) -> usize { +pub fn partition_sliding_midpoint(vec: &mut [T], midpoint_value: f64, partition_on_dimension: usize) -> usize { let vec_len = vec.len(); debug_assert!(vec[0].dims().len() > partition_on_dimension); @@ -74,7 +74,7 @@ pub fn partition_sliding_midpoint(vec: &mut [T], midpoint_v } } -fn partition_kdtree(vec: &mut [T], index_of_splitting_point: usize, partition_on_dimension: usize) -> usize { +fn partition_kdtree(vec: &mut [T], index_of_splitting_point: usize, partition_on_dimension: usize) -> usize { if vec.len() == 1 { return 0; } diff --git a/src/kdtree/test_common.rs b/src/kdtree/test_common.rs index c37a598..38f91c3 100644 --- a/src/kdtree/test_common.rs +++ b/src/kdtree/test_common.rs @@ -1,6 +1,6 @@ -use super::KdtreePointTrait; +use super::KdTreePoint; -#[derive(Copy, Clone, PartialEq)] +#[derive(Copy, Clone, PartialEq, Debug)] pub struct Point3WithId { dims: [f64; 3], pub id: i32, @@ -15,14 +15,14 @@ impl Point3WithId { } } -impl KdtreePointTrait for Point3WithId { +impl KdTreePoint for Point3WithId { #[inline] fn dims(&self) -> &[f64] { return &self.dims; } } -#[derive(Copy, Clone, PartialEq)] +#[derive(Copy, Clone, PartialEq, Debug)] pub struct Point2WithId { dims: [f64; 2], pub id: i32, @@ -37,14 +37,14 @@ impl Point2WithId { } } -impl KdtreePointTrait for Point2WithId { +impl KdTreePoint for Point2WithId { #[inline] fn dims(&self) -> &[f64] { return &self.dims; } } -#[derive(Copy, Clone, PartialEq)] +#[derive(Copy, Clone, PartialEq, Debug)] pub struct Point1WithId { dims: [f64; 1], pub id: i32, @@ -59,7 +59,7 @@ impl Point1WithId { } } -impl KdtreePointTrait for Point1WithId { +impl KdTreePoint for Point1WithId { #[inline] fn dims(&self) -> &[f64] { return &self.dims; diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs index 68520f6..e7f6412 100644 --- a/tests/integration_tests.rs +++ b/tests/integration_tests.rs @@ -4,21 +4,19 @@ extern crate rand; use rand::Rng; use kdtree::kdtree::test_common::*; -use kdtree::kdtree::KdtreePointTrait; +use kdtree::kdtree::KdTreePoint; +use kdtree::kdtree::distance::squared_euclidean; fn gen_random() -> f64 { - rand::thread_rng().gen_range(0., 10000.) + rand::thread_rng().gen_range(0., 1000.) } fn find_nn_with_linear_search(points : &Vec, find_for : Point3WithId) -> &Point3WithId { - let distance_fun = kdtree::kdtree::distance::squared_euclidean; - - - let mut best_found_distance = distance_fun(find_for.dims(), points[0].dims()); + let mut best_found_distance = squared_euclidean(find_for.dims(), points[0].dims()); let mut closed_found_point = &points[0]; for p in points { - let dist = distance_fun(find_for.dims(), p.dims()); + let dist = squared_euclidean(find_for.dims(), p.dims()); if dist < best_found_distance { best_found_distance = dist; @@ -29,6 +27,20 @@ fn find_nn_with_linear_search(points : &Vec, find_for : Point3With closed_found_point } +fn find_neigbours_with_linear_search(points : &Vec, find_for : Point3WithId, dist: f64) -> Vec<&Point3WithId> { + let mut result = Vec::new(); + + for p in points { + let d = squared_euclidean(find_for.dims(), p.dims()); + + if d <= dist { + result.push(p); + } + } + + result +} + fn generate_points(point_count : usize) -> Vec { let mut points : Vec = vec![]; @@ -46,7 +58,7 @@ fn test_against_1000_random_points() { let points = generate_points(point_count); kdtree::kdtree::test_common::Point1WithId::new(0,0.); - let tree = kdtree::kdtree::Kdtree::new(&mut points.clone()); + let tree = kdtree::kdtree::KdTree::new(&mut points.clone()); //test points pushed into the tree, id should be equal. for i in 0 .. point_count { @@ -71,8 +83,8 @@ fn test_incrementally_build_tree_against_built_at_once() { let point_count = 2000usize; let mut points = generate_points(point_count); - let tree_built_at_once = kdtree::kdtree::Kdtree::new(&mut points.clone()); - let mut tree_built_incrementally = kdtree::kdtree::Kdtree::new(&mut points[0..1]); + let tree_built_at_once = kdtree::kdtree::KdTree::new(&mut points.clone()); + let mut tree_built_incrementally = kdtree::kdtree::KdTree::new(&mut points[0..1]); for i in 1 .. point_count { let p = &points[i]; @@ -95,3 +107,28 @@ fn test_incrementally_build_tree_against_built_at_once() { assert_eq!(tree_built_at_once.nearest_search(&p).id, tree_built_incrementally.nearest_search(&p).id); } } + + +#[test] +fn test_neighbour_search_with_distance() { + let point_count = 1000usize; + let points = generate_points(point_count); + let tree = kdtree::kdtree::KdTree::new(&mut points.clone()); + + for _ in 0 .. 500 { + let dist = 100.0; + let p = Point3WithId::new(0i32, gen_random(), gen_random(), gen_random()); + + let mut found_by_linear_search = find_neigbours_with_linear_search(&points, p, dist * dist); + let mut point_found_by_kdtree: Vec<_> = tree.nearest_search_dist(&p, dist * dist).collect(); + + assert_eq!(found_by_linear_search.len(), point_found_by_kdtree.len()); + + if point_found_by_kdtree.len() > 0 { + found_by_linear_search.sort_by(|a, b| a.id.cmp(&b.id)); + point_found_by_kdtree.sort_by(|a, b| a.id.cmp(&b.id)); + } + + assert_eq!(point_found_by_kdtree, found_by_linear_search); + } +} \ No newline at end of file