From 2798eb9bbc489fa6112bde29e5d0bcdcf36c374c Mon Sep 17 00:00:00 2001 From: abstractqqq Date: Mon, 25 Dec 2023 21:09:29 -0500 Subject: [PATCH 1/3] added within_unsorted within_count --- benches/bench.rs | 108 +++++++++++++++++++++++++++++++++++++++++++++++ src/kdtree.rs | 40 ++++++++++++++++++ tests/kdtree.rs | 42 ++++++++++++++++++ 3 files changed, 190 insertions(+) diff --git a/benches/bench.rs b/benches/bench.rs index 0f9f40b..10d262f 100644 --- a/benches/bench.rs +++ b/benches/bench.rs @@ -40,3 +40,111 @@ fn bench_nearest_from_kdtree_with_1k_3d_points(b: &mut Bencher) { } b.iter(|| kdtree.nearest(&point.0, 8, &squared_euclidean).unwrap()); } + +#[bench] +fn bench_within_2k_data_01_radius(b: &mut Bencher) { + + let len = 2000usize; + let point = rand_data(); + let mut points = vec![]; + let mut kdtree = KdTree::with_capacity(3, 16); + for _ in 0..len { + points.push(rand_data()); + } + for i in 0..points.len() { + kdtree.add(&points[i].0, points[i].1).unwrap(); + } + + b.iter(|| kdtree.within(&point.0, 0.1, &squared_euclidean).unwrap()); + +} + +#[bench] +fn bench_within_2k_data_02_radius(b: &mut Bencher) { + + let len = 2000usize; + let point = rand_data(); + let mut points = vec![]; + let mut kdtree = KdTree::with_capacity(3, 16); + for _ in 0..len { + points.push(rand_data()); + } + for i in 0..points.len() { + kdtree.add(&points[i].0, points[i].1).unwrap(); + } + + b.iter(|| kdtree.within(&point.0, 0.2, &squared_euclidean).unwrap()); + +} + +#[bench] +fn bench_within_unsorted_2k_data_01_radius(b: &mut Bencher) { + + let len = 2000usize; + let point = rand_data(); + let mut points = vec![]; + let mut kdtree = KdTree::with_capacity(3, 16); + for _ in 0..len { + points.push(rand_data()); + } + for i in 0..points.len() { + kdtree.add(&points[i].0, points[i].1).unwrap(); + } + + b.iter(|| kdtree.within_unsorted(&point.0, 0.1, &squared_euclidean).unwrap()); + +} + +#[bench] +fn bench_within_unsorted_2k_data_02_radius(b: &mut Bencher) { + + let len = 2000usize; + let point = rand_data(); + let mut points = vec![]; + let mut kdtree = KdTree::with_capacity(3, 16); + for _ in 0..len { + points.push(rand_data()); + } + for i in 0..points.len() { + kdtree.add(&points[i].0, points[i].1).unwrap(); + } + + b.iter(|| kdtree.within_unsorted(&point.0, 0.2, &squared_euclidean).unwrap()); + +} + +#[bench] +fn bench_within_count_2k_data_01_radius(b: &mut Bencher) { + + let len = 2000usize; + let point = rand_data(); + let mut points = vec![]; + let mut kdtree = KdTree::with_capacity(3, 16); + for _ in 0..len { + points.push(rand_data()); + } + for i in 0..points.len() { + kdtree.add(&points[i].0, points[i].1).unwrap(); + } + + b.iter(|| kdtree.within_count(&point.0, 0.1, &squared_euclidean).unwrap()); + +} + +#[bench] +fn bench_within_count_2k_data_02_radius(b: &mut Bencher) { + + let len = 2000usize; + let point = rand_data(); + let mut points = vec![]; + let mut kdtree = KdTree::with_capacity(3, 16); + for _ in 0..len { + points.push(rand_data()); + } + for i in 0..points.len() { + kdtree.add(&points[i].0, points[i].1).unwrap(); + } + + b.iter(|| kdtree.within_count(&point.0, 0.2, &squared_euclidean).unwrap()); + +} \ No newline at end of file diff --git a/src/kdtree.rs b/src/kdtree.rs index b2823a3..c80c61c 100644 --- a/src/kdtree.rs +++ b/src/kdtree.rs @@ -113,6 +113,46 @@ impl + std::cmp::Pa Ok(evaluated.into_sorted_vec().into_iter().map(Into::into).collect()) } + pub fn within_unsorted(&self, point: &[A], radius: A, distance: &F) -> Result, ErrorKind> + where + F: Fn(&[A], &[A]) -> A, + { + self.check_point(point)?; + if self.size == 0 { + return Ok(vec![]); + } + let mut pending = BinaryHeap::new(); + let mut evaluated = BinaryHeap::>::new(); + pending.push(HeapElement { + distance: A::zero(), + element: self, + }); + while !pending.is_empty() && (-pending.peek().unwrap().distance <= radius) { + self.nearest_step(point, self.size, radius, distance, &mut pending, &mut evaluated); + } + Ok(evaluated.into_iter().map(Into::into).collect()) + } + + pub fn within_count(&self, point: &[A], radius: A, distance: &F) -> Result + where + F: Fn(&[A], &[A]) -> A, + { + self.check_point(point)?; + if self.size == 0 { + return Ok(0); + } + let mut pending = BinaryHeap::new(); + let mut evaluated = BinaryHeap::>::new(); + pending.push(HeapElement { + distance: A::zero(), + element: self, + }); + while !pending.is_empty() && (-pending.peek().unwrap().distance <= radius) { + self.nearest_step(point, self.size, radius, distance, &mut pending, &mut evaluated); + } + Ok(evaluated.len()) + } + fn nearest_step<'b, F>( &self, point: &[A], diff --git a/tests/kdtree.rs b/tests/kdtree.rs index 16a6984..98559f5 100644 --- a/tests/kdtree.rs +++ b/tests/kdtree.rs @@ -60,6 +60,48 @@ fn it_works() { vec![(0.0, &1), (2.0, &2), (2.0, &0)] ); + let unsorted1 = kdtree.within_unsorted(&POINT_A.0, 0.0, &squared_euclidean).unwrap(); + let ans1 = vec![(0.0, &0)]; + assert_eq!( + unsorted1.len(), + ans1.len() + ); + assert_eq!( + kdtree.within_count(&POINT_A.0, 0.0, &squared_euclidean).unwrap(), + ans1.len() + ); + for item in unsorted1 { + assert!(ans1.contains(&item)); + } + + let unsorted2 = kdtree.within_unsorted(&POINT_B.0, 1.0, &squared_euclidean).unwrap(); + let ans2 = vec![(0.0, &1)]; + assert_eq!( + unsorted2.len(), + ans2.len() + ); + assert_eq!( + kdtree.within_count(&POINT_B.0, 1.0, &squared_euclidean).unwrap(), + ans2.len() + ); + for item in unsorted2 { + assert!(ans2.contains(&item)); + } + + let unsorted3 = kdtree.within_unsorted(&POINT_B.0, 2.0, &squared_euclidean).unwrap(); + let ans3 = vec![(0.0, &1), (2.0, &2), (2.0, &0)]; + assert_eq!( + unsorted3.len(), + ans3.len() + ); + assert_eq!( + kdtree.within_count(&POINT_B.0, 2.0, &squared_euclidean).unwrap(), + ans3.len() + ); + for item in unsorted3 { + assert!(ans3.contains(&item)); + } + assert_eq!( kdtree .iter_nearest(&POINT_A.0, &squared_euclidean) From 59237b50780da2d006ad40da9810cd5e1c771c80 Mon Sep 17 00:00:00 2001 From: abstractqqq Date: Tue, 2 Jan 2024 10:52:36 -0500 Subject: [PATCH 2/3] did cargo fmt --- benches/bench.rs | 14 +------------- tests/kdtree.rs | 15 +++------------ 2 files changed, 4 insertions(+), 25 deletions(-) diff --git a/benches/bench.rs b/benches/bench.rs index 10d262f..aecdaf3 100644 --- a/benches/bench.rs +++ b/benches/bench.rs @@ -43,7 +43,6 @@ fn bench_nearest_from_kdtree_with_1k_3d_points(b: &mut Bencher) { #[bench] fn bench_within_2k_data_01_radius(b: &mut Bencher) { - let len = 2000usize; let point = rand_data(); let mut points = vec![]; @@ -56,12 +55,10 @@ fn bench_within_2k_data_01_radius(b: &mut Bencher) { } b.iter(|| kdtree.within(&point.0, 0.1, &squared_euclidean).unwrap()); - } #[bench] fn bench_within_2k_data_02_radius(b: &mut Bencher) { - let len = 2000usize; let point = rand_data(); let mut points = vec![]; @@ -74,12 +71,10 @@ fn bench_within_2k_data_02_radius(b: &mut Bencher) { } b.iter(|| kdtree.within(&point.0, 0.2, &squared_euclidean).unwrap()); - } #[bench] fn bench_within_unsorted_2k_data_01_radius(b: &mut Bencher) { - let len = 2000usize; let point = rand_data(); let mut points = vec![]; @@ -92,12 +87,10 @@ fn bench_within_unsorted_2k_data_01_radius(b: &mut Bencher) { } b.iter(|| kdtree.within_unsorted(&point.0, 0.1, &squared_euclidean).unwrap()); - } #[bench] fn bench_within_unsorted_2k_data_02_radius(b: &mut Bencher) { - let len = 2000usize; let point = rand_data(); let mut points = vec![]; @@ -110,12 +103,10 @@ fn bench_within_unsorted_2k_data_02_radius(b: &mut Bencher) { } b.iter(|| kdtree.within_unsorted(&point.0, 0.2, &squared_euclidean).unwrap()); - } #[bench] fn bench_within_count_2k_data_01_radius(b: &mut Bencher) { - let len = 2000usize; let point = rand_data(); let mut points = vec![]; @@ -128,12 +119,10 @@ fn bench_within_count_2k_data_01_radius(b: &mut Bencher) { } b.iter(|| kdtree.within_count(&point.0, 0.1, &squared_euclidean).unwrap()); - } #[bench] fn bench_within_count_2k_data_02_radius(b: &mut Bencher) { - let len = 2000usize; let point = rand_data(); let mut points = vec![]; @@ -146,5 +135,4 @@ fn bench_within_count_2k_data_02_radius(b: &mut Bencher) { } b.iter(|| kdtree.within_count(&point.0, 0.2, &squared_euclidean).unwrap()); - -} \ No newline at end of file +} diff --git a/tests/kdtree.rs b/tests/kdtree.rs index 98559f5..ee1f23b 100644 --- a/tests/kdtree.rs +++ b/tests/kdtree.rs @@ -62,10 +62,7 @@ fn it_works() { let unsorted1 = kdtree.within_unsorted(&POINT_A.0, 0.0, &squared_euclidean).unwrap(); let ans1 = vec![(0.0, &0)]; - assert_eq!( - unsorted1.len(), - ans1.len() - ); + assert_eq!(unsorted1.len(), ans1.len()); assert_eq!( kdtree.within_count(&POINT_A.0, 0.0, &squared_euclidean).unwrap(), ans1.len() @@ -76,10 +73,7 @@ fn it_works() { let unsorted2 = kdtree.within_unsorted(&POINT_B.0, 1.0, &squared_euclidean).unwrap(); let ans2 = vec![(0.0, &1)]; - assert_eq!( - unsorted2.len(), - ans2.len() - ); + assert_eq!(unsorted2.len(), ans2.len()); assert_eq!( kdtree.within_count(&POINT_B.0, 1.0, &squared_euclidean).unwrap(), ans2.len() @@ -90,10 +84,7 @@ fn it_works() { let unsorted3 = kdtree.within_unsorted(&POINT_B.0, 2.0, &squared_euclidean).unwrap(); let ans3 = vec![(0.0, &1), (2.0, &2), (2.0, &0)]; - assert_eq!( - unsorted3.len(), - ans3.len() - ); + assert_eq!(unsorted3.len(), ans3.len()); assert_eq!( kdtree.within_count(&POINT_B.0, 2.0, &squared_euclidean).unwrap(), ans3.len() From 1584872032408721d5bf50c364d44d801508735a Mon Sep 17 00:00:00 2001 From: abstractqqq Date: Thu, 1 Feb 2024 18:21:45 -0500 Subject: [PATCH 3/3] better org --- src/kdtree.rs | 39 ++++++++++++++++----------------------- 1 file changed, 16 insertions(+), 23 deletions(-) diff --git a/src/kdtree.rs b/src/kdtree.rs index c80c61c..1410b0c 100644 --- a/src/kdtree.rs +++ b/src/kdtree.rs @@ -93,14 +93,11 @@ impl + std::cmp::Pa .collect()) } - pub fn within(&self, point: &[A], radius: A, distance: &F) -> Result, ErrorKind> + #[inline(always)] + fn evaluated_heap(&self, point: &[A], radius: A, distance: &F) -> BinaryHeap> where F: Fn(&[A], &[A]) -> A, { - self.check_point(point)?; - if self.size == 0 { - return Ok(vec![]); - } let mut pending = BinaryHeap::new(); let mut evaluated = BinaryHeap::>::new(); pending.push(HeapElement { @@ -110,6 +107,18 @@ impl + std::cmp::Pa while !pending.is_empty() && (-pending.peek().unwrap().distance <= radius) { self.nearest_step(point, self.size, radius, distance, &mut pending, &mut evaluated); } + evaluated + } + + pub fn within(&self, point: &[A], radius: A, distance: &F) -> Result, ErrorKind> + where + F: Fn(&[A], &[A]) -> A, + { + self.check_point(point)?; + if self.size == 0 { + return Ok(vec![]); + } + let evaluated = self.evaluated_heap(point, radius, distance); Ok(evaluated.into_sorted_vec().into_iter().map(Into::into).collect()) } @@ -121,15 +130,7 @@ impl + std::cmp::Pa if self.size == 0 { return Ok(vec![]); } - let mut pending = BinaryHeap::new(); - let mut evaluated = BinaryHeap::>::new(); - pending.push(HeapElement { - distance: A::zero(), - element: self, - }); - while !pending.is_empty() && (-pending.peek().unwrap().distance <= radius) { - self.nearest_step(point, self.size, radius, distance, &mut pending, &mut evaluated); - } + let evaluated = self.evaluated_heap(point, radius, distance); Ok(evaluated.into_iter().map(Into::into).collect()) } @@ -141,15 +142,7 @@ impl + std::cmp::Pa if self.size == 0 { return Ok(0); } - let mut pending = BinaryHeap::new(); - let mut evaluated = BinaryHeap::>::new(); - pending.push(HeapElement { - distance: A::zero(), - element: self, - }); - while !pending.is_empty() && (-pending.peek().unwrap().distance <= radius) { - self.nearest_step(point, self.size, radius, distance, &mut pending, &mut evaluated); - } + let evaluated = self.evaluated_heap(point, radius, distance); Ok(evaluated.len()) }