1use crate::error::{SpatialError, SpatialResult};
16use scirs2_core::ndarray::{Array1, Array2, ArrayView1, ArrayView2};
17use std::cmp::Ordering;
18use std::collections::{BinaryHeap, VecDeque};
19
20const MAX_POINTS_PER_NODE: usize = 8;
22const MAX_DEPTH: usize = 20;
24
25#[derive(Debug, Clone)]
27pub struct BoundingBox2D {
28 pub min: Array1<f64>,
30 pub max: Array1<f64>,
32}
33
34impl BoundingBox2D {
35 pub fn new(min: &ArrayView1<f64>, max: &ArrayView1<f64>) -> SpatialResult<Self> {
51 if min.len() != 2 || max.len() != 2 {
52 return Err(SpatialError::DimensionError(format!(
53 "Min and max must have 2 elements, got {} and {}",
54 min.len(),
55 max.len()
56 )));
57 }
58
59 for i in 0..2 {
61 if min[i] > max[i] {
62 return Err(SpatialError::ValueError(format!(
63 "Min must be <= max for all dimensions, got min[{}]={} > max[{}]={}",
64 i, min[i], i, max[i]
65 )));
66 }
67 }
68
69 Ok(BoundingBox2D {
70 min: min.to_owned(),
71 max: max.to_owned(),
72 })
73 }
74
75 pub fn from_points(points: &ArrayView2<'_, f64>) -> SpatialResult<Self> {
89 if points.is_empty() {
90 return Err(SpatialError::ValueError(
91 "Cannot create bounding box from empty point set".into(),
92 ));
93 }
94
95 if points.ncols() != 2 {
96 return Err(SpatialError::DimensionError(format!(
97 "Points must have 2 columns, got {}",
98 points.ncols()
99 )));
100 }
101
102 let mut min = Array1::from_vec(vec![f64::INFINITY, f64::INFINITY]);
104 let mut max = Array1::from_vec(vec![f64::NEG_INFINITY, f64::NEG_INFINITY]);
105
106 for row in points.rows() {
107 for d in 0..2 {
108 if row[d] < min[d] {
109 min[d] = row[d];
110 }
111 if row[d] > max[d] {
112 max[d] = row[d];
113 }
114 }
115 }
116
117 Ok(BoundingBox2D { min, max })
118 }
119
120 pub fn contains(&self, point: &ArrayView1<f64>) -> SpatialResult<bool> {
134 if point.len() != 2 {
135 return Err(SpatialError::DimensionError(format!(
136 "Point must have 2 elements, got {}",
137 point.len()
138 )));
139 }
140
141 for d in 0..2 {
142 if point[d] < self.min[d] || point[d] > self.max[d] {
143 return Ok(false);
144 }
145 }
146
147 Ok(true)
148 }
149
150 pub fn center(&self) -> Array1<f64> {
156 let mut center = Array1::zeros(2);
157 for d in 0..2 {
158 center[d] = (self.min[d] + self.max[d]) / 2.0;
159 }
160 center
161 }
162
163 pub fn dimensions(&self) -> Array1<f64> {
169 let mut dims = Array1::zeros(2);
170 for d in 0..2 {
171 dims[d] = self.max[d] - self.min[d];
172 }
173 dims
174 }
175
176 pub fn overlaps(&self, other: &BoundingBox2D) -> bool {
186 for d in 0..2 {
187 if self.max[d] < other.min[d] || self.min[d] > other.max[d] {
188 return false;
189 }
190 }
191 true
192 }
193
194 pub fn squared_distance_to_point(&self, point: &ArrayView1<f64>) -> SpatialResult<f64> {
208 if point.len() != 2 {
209 return Err(SpatialError::DimensionError(format!(
210 "Point must have 2 elements, got {}",
211 point.len()
212 )));
213 }
214
215 let mut squared_dist = 0.0;
216
217 for d in 0..2 {
218 let v = point[d];
219
220 if v < self.min[d] {
221 squared_dist += (v - self.min[d]) * (v - self.min[d]);
223 } else if v > self.max[d] {
224 squared_dist += (v - self.max[d]) * (v - self.max[d]);
226 }
227 }
229
230 Ok(squared_dist)
231 }
232
233 pub fn split_into_quadrants(&self) -> [BoundingBox2D; 4] {
239 let center = self.center();
240
241 [
248 BoundingBox2D {
250 min: self.min.clone(),
251 max: center.clone(),
252 },
253 BoundingBox2D {
255 min: Array1::from_vec(vec![center[0], self.min[1]]),
256 max: Array1::from_vec(vec![self.max[0], center[1]]),
257 },
258 BoundingBox2D {
260 min: Array1::from_vec(vec![self.min[0], center[1]]),
261 max: Array1::from_vec(vec![center[0], self.max[1]]),
262 },
263 BoundingBox2D {
265 min: center,
266 max: self.max.clone(),
267 },
268 ]
269 }
270}
271
272#[derive(Debug)]
274enum QuadtreeNode {
275 Internal {
277 bounds: BoundingBox2D,
279 children: Box<[Option<QuadtreeNode>; 4]>,
281 },
282 Leaf {
284 bounds: BoundingBox2D,
286 points: Vec<usize>,
288 point_data: Array2<f64>,
290 },
291}
292
293#[derive(Debug, Clone, PartialEq)]
295struct DistancePoint {
296 index: usize,
298 distance_sq: f64,
300}
301
302impl Ord for DistancePoint {
313 fn cmp(&self, other: &Self) -> Ordering {
314 self.distance_sq
315 .partial_cmp(&other.distance_sq)
316 .unwrap_or(Ordering::Equal)
317 }
318}
319
320impl PartialOrd for DistancePoint {
321 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
322 Some(self.cmp(other))
323 }
324}
325
326impl Eq for DistancePoint {}
327
328#[derive(Debug, Clone, PartialEq)]
330struct DistanceNode {
331 node: *const QuadtreeNode,
333 min_distance_sq: f64,
335}
336
337impl Ord for DistanceNode {
340 fn cmp(&self, other: &Self) -> Ordering {
341 other
342 .min_distance_sq
343 .partial_cmp(&self.min_distance_sq)
344 .unwrap_or(Ordering::Equal)
345 }
346}
347
348impl PartialOrd for DistanceNode {
349 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
350 Some(self.cmp(other))
351 }
352}
353
354impl Eq for DistanceNode {}
355
356#[derive(Debug)]
358pub struct Quadtree {
359 root: Option<QuadtreeNode>,
361 size: usize,
363 points: Array2<f64>,
365}
366
367impl Quadtree {
368 pub fn new(points: &ArrayView2<'_, f64>) -> SpatialResult<Self> {
382 if points.is_empty() {
383 return Err(SpatialError::ValueError(
384 "Cannot create quadtree from empty point set".into(),
385 ));
386 }
387
388 if points.ncols() != 2 {
389 return Err(SpatialError::DimensionError(format!(
390 "Points must have 2 columns, got {}",
391 points.ncols()
392 )));
393 }
394
395 let size = points.nrows();
396 let bounds = BoundingBox2D::from_points(points)?;
397 let points_owned = points.to_owned();
398
399 let indices: Vec<usize> = (0..size).collect();
401
402 let root = Some(Self::build_tree(indices, bounds, &points_owned, 0)?);
404
405 Ok(Quadtree {
406 root,
407 size,
408 points: points_owned,
409 })
410 }
411
412 fn build_tree(
414 indices: Vec<usize>,
415 bounds: BoundingBox2D,
416 points: &Array2<f64>,
417 depth: usize,
418 ) -> SpatialResult<QuadtreeNode> {
419 if depth >= MAX_DEPTH || indices.len() <= MAX_POINTS_PER_NODE {
421 return Ok(QuadtreeNode::Leaf {
422 bounds,
423 points: indices,
424 point_data: points.to_owned(),
425 });
426 }
427
428 let quadrants = bounds.split_into_quadrants();
430
431 let mut quadrant_points: [Vec<usize>; 4] = Default::default();
433
434 for &idx in &indices {
436 let point = points.row(idx);
437 let center = bounds.center();
438
439 let mut quadrant_idx = 0;
441 if point[0] >= center[0] {
442 quadrant_idx |= 1;
443 } if point[1] >= center[1] {
445 quadrant_idx |= 2;
446 } quadrant_points[quadrant_idx].push(idx);
449 }
450
451 let mut children: [Option<QuadtreeNode>; 4] = Default::default();
453
454 for i in 0..4 {
455 if !quadrant_points[i].is_empty() {
456 children[i] = Some(Self::build_tree(
457 quadrant_points[i].clone(),
458 quadrants[i].clone(),
459 points,
460 depth + 1,
461 )?);
462 }
463 }
464
465 Ok(QuadtreeNode::Internal {
466 bounds,
467 children: Box::new(children),
468 })
469 }
470
471 pub fn query_nearest(
488 &self,
489 query: &ArrayView1<f64>,
490 k: usize,
491 ) -> SpatialResult<(Vec<usize>, Vec<f64>)> {
492 if query.len() != 2 {
493 return Err(SpatialError::DimensionError(format!(
494 "Query point must have 2 dimensions, got {}",
495 query.len()
496 )));
497 }
498
499 if k == 0 {
500 return Err(SpatialError::ValueError("k must be > 0".into()));
501 }
502
503 if self.root.is_none() {
504 return Ok((Vec::new(), Vec::new()));
505 }
506
507 let mut node_queue = BinaryHeap::new();
509
510 let mut result_queue = BinaryHeap::new();
512 let mut worst_dist = f64::INFINITY;
513
514 let root_ref = self.root.as_ref().expect("Operation failed") as *const QuadtreeNode;
516 let root_dist = match self.root.as_ref().expect("Operation failed") {
517 QuadtreeNode::Internal { bounds, .. } => bounds.squared_distance_to_point(query)?,
518 QuadtreeNode::Leaf { bounds, .. } => bounds.squared_distance_to_point(query)?,
519 };
520
521 node_queue.push(DistanceNode {
522 node: root_ref,
523 min_distance_sq: root_dist,
524 });
525
526 while let Some(dist_node) = node_queue.pop() {
528 if dist_node.min_distance_sq > worst_dist && result_queue.len() >= k {
530 continue;
531 }
532
533 let node = unsafe { &*dist_node.node };
536
537 match node {
538 QuadtreeNode::Leaf {
539 points, point_data, ..
540 } => {
541 for &idx in points {
543 let point = point_data.row(idx);
544 let dist_sq = squared_distance(query, &point);
545
546 if result_queue.len() < k || dist_sq < worst_dist {
548 result_queue.push(DistancePoint {
549 index: idx,
550 distance_sq: dist_sq,
551 });
552
553 if result_queue.len() > k {
555 result_queue.pop();
556 if let Some(worst) = result_queue.peek() {
558 worst_dist = worst.distance_sq;
559 }
560 }
561 }
562 }
563 }
564 QuadtreeNode::Internal { children, .. } => {
565 for child in children.iter().flatten() {
567 let child_ref = child as *const QuadtreeNode;
568
569 let min_dist = match child {
570 QuadtreeNode::Internal { bounds, .. } => {
571 bounds.squared_distance_to_point(query)?
572 }
573 QuadtreeNode::Leaf { bounds, .. } => {
574 bounds.squared_distance_to_point(query)?
575 }
576 };
577
578 node_queue.push(DistanceNode {
579 node: child_ref,
580 min_distance_sq: min_dist,
581 });
582 }
583 }
584 }
585 }
586
587 let mut result_indices = Vec::with_capacity(result_queue.len());
589 let mut result_distances = Vec::with_capacity(result_queue.len());
590
591 let mut temp_results = Vec::new();
593 while let Some(result) = result_queue.pop() {
594 temp_results.push(result);
595 }
596
597 for result in temp_results.iter().rev() {
599 result_indices.push(result.index);
600 result_distances.push(result.distance_sq);
601 }
602
603 Ok((result_indices, result_distances))
604 }
605
606 pub fn query_radius(
623 &self,
624 query: &ArrayView1<f64>,
625 radius: f64,
626 ) -> SpatialResult<(Vec<usize>, Vec<f64>)> {
627 if query.len() != 2 {
628 return Err(SpatialError::DimensionError(format!(
629 "Query point must have 2 dimensions, got {}",
630 query.len()
631 )));
632 }
633
634 if radius < 0.0 {
635 return Err(SpatialError::ValueError(
636 "Radius must be non-negative".into(),
637 ));
638 }
639
640 let radius_sq = radius * radius;
641
642 if self.root.is_none() {
643 return Ok((Vec::new(), Vec::new()));
644 }
645
646 let mut result_indices = Vec::new();
647 let mut result_distances = Vec::new();
648
649 let mut node_queue = VecDeque::new();
651 node_queue.push_back(self.root.as_ref().expect("Operation failed"));
652
653 while let Some(node) = node_queue.pop_front() {
654 match node {
655 QuadtreeNode::Leaf {
656 points,
657 point_data,
658 bounds,
659 ..
660 } => {
661 if bounds.squared_distance_to_point(query)? > radius_sq {
663 continue;
664 }
665
666 for &idx in points {
668 let point = point_data.row(idx);
669 let dist_sq = squared_distance(query, &point);
670
671 if dist_sq <= radius_sq {
672 result_indices.push(idx);
673 result_distances.push(dist_sq);
674 }
675 }
676 }
677 QuadtreeNode::Internal {
678 children, bounds, ..
679 } => {
680 if bounds.squared_distance_to_point(query)? > radius_sq {
682 continue;
683 }
684
685 for child in children.iter().flatten() {
687 node_queue.push_back(child);
688 }
689 }
690 }
691 }
692
693 Ok((result_indices, result_distances))
694 }
695
696 pub fn points_in_region(&self, region: &BoundingBox2D) -> bool {
706 if self.root.is_none() {
707 return false;
708 }
709
710 let mut node_stack = Vec::new();
712 node_stack.push(self.root.as_ref().expect("Operation failed"));
713
714 while let Some(node) = node_stack.pop() {
715 match node {
716 QuadtreeNode::Leaf {
717 points,
718 point_data,
719 bounds,
720 ..
721 } => {
722 if !bounds.overlaps(region) {
724 continue;
725 }
726
727 for &idx in points {
729 let point = point_data.row(idx);
730 let point_in_region = region.contains(&point.view()).unwrap_or(false);
731
732 if point_in_region {
733 return true;
734 }
735 }
736 }
737 QuadtreeNode::Internal {
738 children, bounds, ..
739 } => {
740 if !bounds.overlaps(region) {
742 continue;
743 }
744
745 for child in children.iter().flatten() {
747 node_stack.push(child);
748 }
749 }
750 }
751 }
752
753 false
754 }
755
756 pub fn get_points_in_region(&self, region: &BoundingBox2D) -> Vec<usize> {
766 if self.root.is_none() {
767 return Vec::new();
768 }
769
770 let mut result_indices = Vec::new();
771
772 let mut node_stack = Vec::new();
774 node_stack.push(self.root.as_ref().expect("Operation failed"));
775
776 while let Some(node) = node_stack.pop() {
777 match node {
778 QuadtreeNode::Leaf {
779 points,
780 point_data,
781 bounds,
782 ..
783 } => {
784 if !bounds.overlaps(region) {
786 continue;
787 }
788
789 for &idx in points {
791 let point = point_data.row(idx);
792 let point_in_region = region.contains(&point.view()).unwrap_or(false);
793
794 if point_in_region {
795 result_indices.push(idx);
796 }
797 }
798 }
799 QuadtreeNode::Internal {
800 children, bounds, ..
801 } => {
802 if !bounds.overlaps(region) {
804 continue;
805 }
806
807 for child in children.iter().flatten() {
809 node_stack.push(child);
810 }
811 }
812 }
813 }
814
815 result_indices
816 }
817
818 pub fn get_point(&self, index: usize) -> Option<Array1<f64>> {
828 if index < self.size {
829 Some(self.points.row(index).to_owned())
830 } else {
831 None
832 }
833 }
834
835 pub fn size(&self) -> usize {
841 self.size
842 }
843
844 pub fn bounds(&self) -> Option<BoundingBox2D> {
850 match &self.root {
851 Some(QuadtreeNode::Internal { bounds, .. }) => Some(bounds.clone()),
852 Some(QuadtreeNode::Leaf { bounds, .. }) => Some(bounds.clone()),
853 None => None,
854 }
855 }
856
857 pub fn max_depth(&self) -> usize {
863 Quadtree::compute_max_depth(self.root.as_ref())
864 }
865
866 #[allow(clippy::only_used_in_recursion)]
868 fn compute_max_depth(node: Option<&QuadtreeNode>) -> usize {
869 match node {
870 None => 0,
871 Some(QuadtreeNode::Leaf { .. }) => 1,
872 Some(QuadtreeNode::Internal { children, .. }) => {
873 let mut max_child_depth = 0;
874 for child in children.iter().flatten() {
875 let child_depth = Self::compute_max_depth(Some(child));
876 max_child_depth = max_child_depth.max(child_depth);
877 }
878 1 + max_child_depth
879 }
880 }
881 }
882}
883
884#[allow(dead_code)]
895fn squared_distance(p1: &ArrayView1<f64>, p2: &ArrayView1<f64>) -> f64 {
896 let mut sum_sq = 0.0;
897 for i in 0..p1.len().min(p2.len()) {
898 let diff = p1[i] - p2[i];
899 sum_sq += diff * diff;
900 }
901 sum_sq
902}
903
904#[cfg(test)]
905mod tests {
906 use super::*;
907 use scirs2_core::ndarray::array;
908
909 #[test]
910 fn test_bounding_box_creation() {
911 let min = array![0.0, 0.0];
913 let max = array![1.0, 1.0];
914 let bbox = BoundingBox2D::new(&min.view(), &max.view()).expect("Operation failed");
915
916 assert_eq!(bbox.min, min);
917 assert_eq!(bbox.max, max);
918
919 let points = array![[0.0, 0.0], [1.0, 1.0], [0.5, 0.5],];
921 let bbox = BoundingBox2D::from_points(&points.view()).expect("Operation failed");
922
923 assert_eq!(bbox.min, min);
924 assert_eq!(bbox.max, max);
925
926 let bad_min = array![0.0];
928 let result = BoundingBox2D::new(&bad_min.view(), &max.view());
929 assert!(result.is_err());
930
931 let bad_minmax = array![2.0, 0.0];
932 let result = BoundingBox2D::new(&bad_minmax.view(), &max.view());
933 assert!(result.is_err());
934 }
935
936 #[test]
937 fn test_bounding_box_operations() {
938 let min = array![0.0, 0.0];
939 let max = array![2.0, 4.0];
940 let bbox = BoundingBox2D::new(&min.view(), &max.view()).expect("Operation failed");
941
942 let center = bbox.center();
944 assert_eq!(center, array![1.0, 2.0]);
945
946 let dims = bbox.dimensions();
948 assert_eq!(dims, array![2.0, 4.0]);
949
950 let inside_point = array![1.0, 1.0];
952 assert!(bbox
953 .contains(&inside_point.view())
954 .expect("Operation failed"));
955
956 let outside_point = array![3.0, 3.0];
957 assert!(!bbox
958 .contains(&outside_point.view())
959 .expect("Operation failed"));
960
961 let edge_point = array![0.0, 4.0];
962 assert!(bbox.contains(&edge_point.view()).expect("Operation failed"));
963
964 let overlapping_box =
966 BoundingBox2D::new(&array![1.0, 1.0].view(), &array![3.0, 3.0].view())
967 .expect("Operation failed");
968 assert!(bbox.overlaps(&overlapping_box));
969
970 let non_overlapping_box =
971 BoundingBox2D::new(&array![3.0, 5.0].view(), &array![4.0, 6.0].view())
972 .expect("Operation failed");
973 assert!(!bbox.overlaps(&non_overlapping_box));
974
975 let inside_dist = bbox
977 .squared_distance_to_point(&inside_point.view())
978 .expect("Operation failed");
979 assert_eq!(inside_dist, 0.0);
980
981 let outside_dist = bbox
982 .squared_distance_to_point(&array![3.0, 5.0].view())
983 .expect("Operation failed");
984 assert_eq!(outside_dist, 1.0 + 1.0); }
986
987 #[test]
988 fn test_quadtree_creation() {
989 let points = array![[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [0.5, 0.5],];
991
992 let quadtree = Quadtree::new(&points.view()).expect("Operation failed");
993
994 assert_eq!(quadtree.size(), 5);
996
997 let bounds = quadtree.bounds().expect("Operation failed");
998 assert_eq!(bounds.min, array![0.0, 0.0]);
999 assert_eq!(bounds.max, array![1.0, 1.0]);
1000
1001 assert!(quadtree.max_depth() > 0);
1003 }
1004
1005 #[test]
1006 fn test_nearest_neighbor_search() {
1007 let points = array![
1009 [0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [0.5, 0.5], [2.0, 2.0], ];
1016
1017 let quadtree = Quadtree::new(&points.view()).expect("Operation failed");
1018
1019 let query = array![0.1, 0.1];
1021 let (indices, distances) = quadtree
1022 .query_nearest(&query.view(), 1)
1023 .expect("Operation failed");
1024
1025 assert_eq!(indices.len(), 1);
1026 assert_eq!(indices[0], 0); assert!(distances[0] >= 0.0);
1028
1029 let (indices, distances) = quadtree
1031 .query_nearest(&query.view(), 3)
1032 .expect("Operation failed");
1033
1034 assert!(!indices.is_empty());
1036
1037 for d in distances.iter() {
1039 assert!(*d >= 0.0);
1040 }
1041
1042 let (indices, distances) = quadtree
1044 .query_nearest(&query.view(), 10)
1045 .expect("Operation failed");
1046
1047 assert_eq!(indices.len(), 6); assert_eq!(distances.len(), 6);
1049 }
1050
1051 #[test]
1052 fn test_radius_search() {
1053 let points = array![
1055 [0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [0.5, 0.5], [2.0, 2.0], ];
1062
1063 let quadtree = Quadtree::new(&points.view()).expect("Operation failed");
1064
1065 let query = array![0.0, 0.0];
1067 let radius = 0.5;
1068 let (indices, distances) = quadtree
1069 .query_radius(&query.view(), radius)
1070 .expect("Operation failed");
1071
1072 assert_eq!(indices.len(), 1);
1073 assert_eq!(indices[0], 0); let radius = 1.5;
1077 let (indices, distances) = quadtree
1078 .query_radius(&query.view(), radius)
1079 .expect("Operation failed");
1080
1081 assert!(indices.len() >= 4); for &dist in &distances {
1085 assert!(dist <= radius * radius);
1086 }
1087
1088 let radius = 4.0;
1090 let (indices, distances) = quadtree
1091 .query_radius(&query.view(), radius)
1092 .expect("Operation failed");
1093
1094 assert_eq!(indices.len(), 6); }
1096
1097 #[test]
1098 fn test_region_queries() {
1099 let points = array![
1101 [0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [0.5, 0.5], [2.0, 2.0], ];
1108
1109 let quadtree = Quadtree::new(&points.view()).expect("Operation failed");
1110
1111 let region = BoundingBox2D::new(&array![0.25, 0.25].view(), &array![0.75, 0.75].view())
1113 .expect("Operation failed");
1114
1115 assert!(quadtree.points_in_region(®ion));
1117
1118 let indices = quadtree.get_points_in_region(®ion);
1120 assert_eq!(indices.len(), 1);
1121 assert_eq!(indices[0], 4); let large_region = BoundingBox2D::new(&array![0.0, 0.0].view(), &array![1.0, 1.0].view())
1125 .expect("Operation failed");
1126
1127 let indices = quadtree.get_points_in_region(&large_region);
1128 assert_eq!(indices.len(), 5); let empty_region = BoundingBox2D::new(&array![1.5, 1.5].view(), &array![1.9, 1.9].view())
1132 .expect("Operation failed");
1133
1134 assert!(!quadtree.points_in_region(&empty_region));
1135 let indices = quadtree.get_points_in_region(&empty_region);
1136 assert_eq!(indices.len(), 0);
1137 }
1138}