Skip to main content

oxiphysics_core/spatial/
functions.rs

1//! Auto-generated module
2//!
3//! 🤖 Generated with [SplitRS](https://github.com/cool-japan/splitrs)
4
5use crate::math::Vec3;
6
7use super::types::{GridSpatialIndex, KdNode, KdTree, Octree, OctreeNode, SpatialIndexStats};
8
9/// Compute statistics for an `Octree`.
10pub fn octree_stats<T: Clone>(tree: &Octree<T>) -> SpatialIndexStats {
11    let mut stats = SpatialIndexStats {
12        total_items: tree.len(),
13        ..Default::default()
14    };
15    octree_node_stats(&tree.root, 0, &mut stats);
16    if stats.n_leaves > 0 {
17        stats.avg_items_per_leaf = stats.total_items as f64 / stats.n_leaves as f64;
18    }
19    stats
20}
21pub(super) fn octree_node_stats<T: Clone>(
22    node: &OctreeNode<T>,
23    depth: usize,
24    stats: &mut SpatialIndexStats,
25) {
26    if depth > stats.max_depth {
27        stats.max_depth = depth;
28    }
29    match node {
30        OctreeNode::Leaf { items } => {
31            stats.n_leaves += 1;
32            let _ = items;
33        }
34        OctreeNode::Internal { children, .. } => {
35            stats.n_internal += 1;
36            for child in children.iter() {
37                octree_node_stats(child, depth + 1, stats);
38            }
39        }
40    }
41}
42/// Compute statistics for a `KdTree`.
43pub fn kdtree_stats(tree: &KdTree) -> SpatialIndexStats {
44    SpatialIndexStats {
45        total_items: tree.len(),
46        max_depth: kdtree_max_depth(
47            &tree.nodes,
48            if tree.nodes.is_empty() { None } else { Some(0) },
49            0,
50        ),
51        n_leaves: tree
52            .nodes
53            .iter()
54            .filter(|n| n.left.is_none() && n.right.is_none())
55            .count(),
56        n_internal: tree
57            .nodes
58            .iter()
59            .filter(|n| n.left.is_some() || n.right.is_some())
60            .count(),
61        avg_items_per_leaf: 1.0,
62    }
63}
64pub(super) fn kdtree_max_depth(nodes: &[KdNode], idx: Option<usize>, depth: usize) -> usize {
65    match idx {
66        None => depth,
67        Some(i) => {
68            let l = kdtree_max_depth(nodes, nodes[i].left, depth + 1);
69            let r = kdtree_max_depth(nodes, nodes[i].right, depth + 1);
70            l.max(r)
71        }
72    }
73}
74/// Spatial join: find all pairs (i, j) from two point sets within distance `r`.
75///
76/// Returns unique pairs `(i, j)` where point `i` is from set A and `j` from set B.
77#[allow(dead_code)]
78pub fn spatial_join(points_a: &[Vec3], points_b: &[Vec3], r: f64) -> Vec<(usize, usize)> {
79    if points_a.is_empty() || points_b.is_empty() {
80        return Vec::new();
81    }
82    let mut grid = GridSpatialIndex::new(r.max(1e-15));
83    for &p in points_b {
84        grid.insert(p);
85    }
86    let mut pairs = Vec::new();
87    for (i, &pa) in points_a.iter().enumerate() {
88        let neighbors = grid.range_query(pa, r);
89        for j in neighbors {
90            pairs.push((i, j));
91        }
92    }
93    pairs
94}
95/// Spatial self-join: find all pairs (i, j) with i < j within distance `r`.
96#[allow(dead_code)]
97pub fn spatial_self_join(points: &[Vec3], r: f64) -> Vec<(usize, usize)> {
98    let n = points.len();
99    let mut grid = GridSpatialIndex::new(r.max(1e-15));
100    for &p in points {
101        grid.insert(p);
102    }
103    let r_sq = r * r;
104    let mut pairs = Vec::new();
105    for i in 0..n {
106        let neighbors = grid.range_query(points[i], r);
107        for j in neighbors {
108            if j > i {
109                let diff = points[i] - points[j];
110                if diff.norm_squared() <= r_sq {
111                    pairs.push((i, j));
112                }
113            }
114        }
115    }
116    pairs
117}
118#[cfg(test)]
119mod tests {
120    use super::*;
121    use crate::Vec3;
122
123    use crate::spatial::FlatRTree;
124    use crate::spatial::KdTree;
125    use crate::spatial::KdTreeWithDeletion;
126    use crate::spatial::LshIndex;
127
128    use crate::spatial::RTree;
129    use crate::spatial::RangeTree1D;
130    use crate::spatial::SpatialAabb;
131
132    fn unit_octree<T: Clone>(max_items: usize) -> Octree<T> {
133        Octree::new(
134            SpatialAabb::new(Vec3::new(0.0, 0.0, 0.0), Vec3::new(1.0, 1.0, 1.0)),
135            8,
136            max_items,
137        )
138    }
139    #[test]
140    fn test_octree_insert_query() {
141        let mut tree = unit_octree::<u32>(16);
142        for i in 0..100u32 {
143            let t = i as f64 / 100.0;
144            let p = Vec3::new(t, t * t, (1.0 - t) * t);
145            tree.insert(p, i);
146        }
147        assert_eq!(tree.len(), 100);
148        let hits = tree.query_sphere(Vec3::new(0.5, 0.25, 0.25), 0.1);
149        assert!(!hits.is_empty());
150        assert!(hits.len() < 100);
151        for (p, _) in &hits {
152            let d = (p - Vec3::new(0.5, 0.25, 0.25)).norm();
153            assert!(d <= 0.1 + 1e-10, "point outside sphere: d={d}");
154        }
155    }
156    #[test]
157    fn test_octree_nearest_neighbor() {
158        let mut tree = unit_octree::<usize>(8);
159        let pts = [
160            Vec3::new(0.1, 0.1, 0.1),
161            Vec3::new(0.9, 0.9, 0.9),
162            Vec3::new(0.5, 0.5, 0.5),
163            Vec3::new(0.2, 0.8, 0.3),
164            Vec3::new(0.7, 0.2, 0.6),
165            Vec3::new(0.4, 0.6, 0.1),
166            Vec3::new(0.3, 0.3, 0.7),
167            Vec3::new(0.8, 0.4, 0.2),
168            Vec3::new(0.1, 0.7, 0.9),
169            Vec3::new(0.6, 0.1, 0.4),
170        ];
171        for (i, &p) in pts.iter().enumerate() {
172            tree.insert(p, i);
173        }
174        let query = Vec3::new(0.45, 0.45, 0.45);
175        let (nearest_pos, _val, dist) = tree.nearest_neighbor(query).unwrap();
176        let (bf_pos, bf_dist) = pts
177            .iter()
178            .map(|&p| (p, (p - query).norm()))
179            .min_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
180            .unwrap();
181        assert!(
182            (dist - bf_dist).abs() < 1e-10,
183            "dist mismatch: {dist} vs {bf_dist}"
184        );
185        assert!((nearest_pos - bf_pos).norm() < 1e-10, "position mismatch");
186    }
187    #[test]
188    fn test_octree_subdivision() {
189        let mut tree: Octree<i32> = Octree::new(
190            SpatialAabb::new(Vec3::zeros(), Vec3::new(1.0, 1.0, 1.0)),
191            8,
192            2,
193        );
194        for i in 0..20i32 {
195            let f = i as f64 / 20.0;
196            tree.insert(Vec3::new(f, f, f), i);
197        }
198        assert_eq!(tree.len(), 20);
199        let all = tree.query_aabb(&SpatialAabb::new(Vec3::zeros(), Vec3::new(1.0, 1.0, 1.0)));
200        assert_eq!(all.len(), 20);
201    }
202    #[test]
203    fn test_kdtree_build_nearest() {
204        let pts = vec![
205            Vec3::new(1.0, 0.0, 0.0),
206            Vec3::new(0.0, 1.0, 0.0),
207            Vec3::new(0.0, 0.0, 1.0),
208            Vec3::new(2.0, 2.0, 2.0),
209            Vec3::new(0.1, 0.1, 0.1),
210        ];
211        let kd = KdTree::build(pts.clone());
212        assert_eq!(kd.len(), 5);
213        let query = Vec3::new(0.05, 0.05, 0.05);
214        let (idx, d2) = kd.nearest(query).unwrap();
215        let expected_idx = pts
216            .iter()
217            .enumerate()
218            .min_by(|a, b| {
219                (a.1 - query)
220                    .norm_squared()
221                    .partial_cmp(&(b.1 - query).norm_squared())
222                    .unwrap()
223            })
224            .map(|(i, _)| i)
225            .unwrap();
226        assert_eq!(idx, expected_idx, "nearest index mismatch");
227        let expected_d2 = (pts[expected_idx] - query).norm_squared();
228        assert!((d2 - expected_d2).abs() < 1e-10, "d² mismatch");
229    }
230    #[test]
231    fn test_kdtree_k_nearest() {
232        let pts = vec![
233            Vec3::new(0.0, 0.0, 0.0),
234            Vec3::new(1.0, 0.0, 0.0),
235            Vec3::new(2.0, 0.0, 0.0),
236            Vec3::new(3.0, 0.0, 0.0),
237            Vec3::new(4.0, 0.0, 0.0),
238        ];
239        let kd = KdTree::build(pts.clone());
240        let query = Vec3::new(1.5, 0.0, 0.0);
241        let k3 = kd.k_nearest(query, 3);
242        assert_eq!(k3.len(), 3, "expected 3 results");
243        for w in k3.windows(2) {
244            assert!(w[0].1 <= w[1].1 + 1e-10, "not sorted");
245        }
246        let first_d2 = k3[0].1;
247        assert!(
248            first_d2 < 0.26,
249            "closest should be within 0.5: d²={first_d2}"
250        );
251    }
252    #[test]
253    fn test_kdtree_range() {
254        let pts = vec![
255            Vec3::new(0.0, 0.0, 0.0),
256            Vec3::new(0.3, 0.0, 0.0),
257            Vec3::new(0.6, 0.0, 0.0),
258            Vec3::new(1.0, 0.0, 0.0),
259            Vec3::new(2.0, 0.0, 0.0),
260        ];
261        let kd = KdTree::build(pts.clone());
262        let query = Vec3::new(0.5, 0.0, 0.0);
263        let mut within = kd.range_search(query, 0.5);
264        within.sort_unstable();
265        assert_eq!(
266            within.len(),
267            4,
268            "expected 4 points within radius 0.5, got {}",
269            within.len()
270        );
271        assert!(
272            !within.contains(&4),
273            "point at distance 1.5 should not be in range"
274        );
275    }
276    #[test]
277    fn test_spatial_aabb_operations() {
278        let a = SpatialAabb::new(Vec3::zeros(), Vec3::new(2.0, 2.0, 2.0));
279        let c = a.center();
280        assert!((c - Vec3::new(1.0, 1.0, 1.0)).norm() < 1e-10);
281        let he = a.half_extents();
282        assert!((he - Vec3::new(1.0, 1.0, 1.0)).norm() < 1e-10);
283        assert!(a.contains_point(Vec3::new(1.0, 1.0, 1.0)));
284        assert!(a.contains_point(Vec3::new(0.0, 0.0, 0.0)));
285        assert!(!a.contains_point(Vec3::new(3.0, 1.0, 1.0)));
286        let b = SpatialAabb::new(Vec3::new(1.0, 1.0, 1.0), Vec3::new(3.0, 3.0, 3.0));
287        assert!(a.intersects(&b));
288        let c_box = SpatialAabb::new(Vec3::new(3.0, 3.0, 3.0), Vec3::new(5.0, 5.0, 5.0));
289        assert!(!a.intersects(&c_box));
290        let expanded = a.expand(1.0);
291        assert!((expanded.min - Vec3::new(-1.0, -1.0, -1.0)).norm() < 1e-10);
292        assert!((expanded.max - Vec3::new(3.0, 3.0, 3.0)).norm() < 1e-10);
293        assert!(expanded.contains_point(Vec3::new(2.0, 2.0, 2.0)));
294        assert!(expanded.contains_point(Vec3::new(-0.5, -0.5, -0.5)));
295    }
296    fn pt_aabb(x: f64, y: f64, z: f64) -> SpatialAabb {
297        SpatialAabb::new(
298            Vec3::new(x - 0.05, y - 0.05, z - 0.05),
299            Vec3::new(x + 0.05, y + 0.05, z + 0.05),
300        )
301    }
302    #[test]
303    fn test_rtree_build_and_len() {
304        let items: Vec<(SpatialAabb, u32)> = (0..20u32)
305            .map(|i| {
306                let f = i as f64 / 20.0;
307                (pt_aabb(f, f, f), i)
308            })
309            .collect();
310        let tree = RTree::build(items, 4);
311        assert_eq!(tree.len(), 20);
312    }
313    #[test]
314    fn test_rtree_query_finds_overlapping() {
315        let items: Vec<(SpatialAabb, u32)> = vec![
316            (pt_aabb(0.1, 0.1, 0.1), 0),
317            (pt_aabb(0.5, 0.5, 0.5), 1),
318            (pt_aabb(0.9, 0.9, 0.9), 2),
319        ];
320        let tree = RTree::build(items, 4);
321        let query = SpatialAabb::new(Vec3::new(0.0, 0.0, 0.0), Vec3::new(0.3, 0.3, 0.3));
322        let results = tree.query(&query);
323        assert!(!results.is_empty(), "Should find item near (0.1,0.1,0.1)");
324        assert!(results.contains(&&0u32), "Should find item 0");
325        assert!(!results.contains(&&2u32), "Should NOT find item 2");
326    }
327    #[test]
328    fn test_flat_rtree_empty_original() {
329        let tree: RTree<u32> = RTree::build(vec![], 4);
330        assert!(tree.is_empty());
331        let query = SpatialAabb::new(Vec3::zeros(), Vec3::new(1.0, 1.0, 1.0));
332        assert!(tree.query(&query).is_empty());
333    }
334    #[test]
335    fn test_rtree_query_no_overlap() {
336        let items = vec![(pt_aabb(0.1, 0.1, 0.1), 0u32)];
337        let tree = RTree::build(items, 4);
338        let query = SpatialAabb::new(Vec3::new(5.0, 5.0, 5.0), Vec3::new(6.0, 6.0, 6.0));
339        assert!(tree.query(&query).is_empty());
340    }
341    #[test]
342    fn test_octree_stats_empty() {
343        let tree: Octree<u32> = Octree::new(
344            SpatialAabb::new(Vec3::zeros(), Vec3::new(1.0, 1.0, 1.0)),
345            4,
346            8,
347        );
348        let stats = octree_stats(&tree);
349        assert_eq!(stats.total_items, 0);
350        assert_eq!(stats.n_leaves, 1);
351        assert_eq!(stats.n_internal, 0);
352    }
353    #[test]
354    fn test_octree_stats_after_inserts() {
355        let mut tree: Octree<u32> = Octree::new(
356            SpatialAabb::new(Vec3::zeros(), Vec3::new(1.0, 1.0, 1.0)),
357            6,
358            2,
359        );
360        for i in 0..12u32 {
361            let f = (i as f64 + 0.5) / 12.0;
362            tree.insert(Vec3::new(f, f, f), i);
363        }
364        let stats = octree_stats(&tree);
365        assert_eq!(stats.total_items, 12);
366        assert!(stats.n_leaves >= 1);
367    }
368    #[test]
369    fn test_kdtree_stats() {
370        let pts = vec![
371            Vec3::new(0.0, 0.0, 0.0),
372            Vec3::new(1.0, 0.0, 0.0),
373            Vec3::new(0.0, 1.0, 0.0),
374            Vec3::new(0.0, 0.0, 1.0),
375            Vec3::new(1.0, 1.0, 1.0),
376        ];
377        let kd = KdTree::build(pts);
378        let stats = kdtree_stats(&kd);
379        assert_eq!(stats.total_items, 5);
380        assert!(
381            stats.max_depth >= 2,
382            "Tree of 5 nodes should have depth >= 2"
383        );
384        assert!(stats.n_leaves >= 1);
385    }
386    #[test]
387    fn test_kdtree_deletion_n_active() {
388        let pts = vec![
389            Vec3::new(0.0, 0.0, 0.0),
390            Vec3::new(1.0, 0.0, 0.0),
391            Vec3::new(2.0, 0.0, 0.0),
392        ];
393        let mut kd = KdTreeWithDeletion::build(pts);
394        assert_eq!(kd.n_active(), 3);
395        kd.delete(1);
396        assert_eq!(kd.n_active(), 2);
397    }
398    #[test]
399    fn test_kdtree_deletion_nearest_active_skips_deleted() {
400        let pts = vec![
401            Vec3::new(0.0, 0.0, 0.0),
402            Vec3::new(0.1, 0.0, 0.0),
403            Vec3::new(1.0, 0.0, 0.0),
404        ];
405        let mut kd = KdTreeWithDeletion::build(pts);
406        let query = Vec3::new(0.05, 0.0, 0.0);
407        let before = kd.nearest_active(query).map(|(i, _)| i);
408        assert_eq!(before, Some(1), "Nearest should be index 1");
409        kd.delete(1);
410        let after = kd.nearest_active(query).map(|(i, _)| i);
411        assert!(after != Some(1), "Deleted index 1 should not be returned");
412    }
413    #[test]
414    fn test_kdtree_deletion_all_deleted_returns_none() {
415        let pts = vec![Vec3::new(0.0, 0.0, 0.0)];
416        let mut kd = KdTreeWithDeletion::build(pts);
417        kd.delete(0);
418        let result = kd.nearest_active(Vec3::zeros());
419        assert!(result.is_none(), "All deleted: should return None");
420    }
421    #[test]
422    fn test_range_tree_query_basic() {
423        let values = vec![1.0, 3.0, 5.0, 7.0, 9.0];
424        let tree = RangeTree1D::build(&values);
425        let result = tree.range_query(3.0, 7.0);
426        let mut sorted = result.clone();
427        sorted.sort_unstable();
428        assert_eq!(
429            sorted,
430            vec![1, 2, 3],
431            "Range [3,7] should return indices 1,2,3"
432        );
433    }
434    #[test]
435    fn test_range_tree_empty_range() {
436        let values = vec![1.0, 3.0, 5.0];
437        let tree = RangeTree1D::build(&values);
438        let result = tree.range_query(10.0, 20.0);
439        assert!(result.is_empty(), "No values in [10,20]");
440    }
441    #[test]
442    fn test_range_tree_all_in_range() {
443        let values = vec![1.0, 2.0, 3.0, 4.0, 5.0];
444        let tree = RangeTree1D::build(&values);
445        let result = tree.range_query(0.0, 10.0);
446        assert_eq!(result.len(), 5, "All 5 values should be in [0,10]");
447    }
448    #[test]
449    fn test_range_tree_min_max() {
450        let values = vec![5.0, 1.0, 9.0, 3.0];
451        let tree = RangeTree1D::build(&values);
452        assert!((tree.min_value().unwrap() - 1.0).abs() < 1e-12);
453        assert!((tree.max_value().unwrap() - 9.0).abs() < 1e-12);
454    }
455    #[test]
456    fn test_range_tree_empty() {
457        let tree = RangeTree1D::build(&[]);
458        assert!(tree.is_empty());
459        assert!(tree.min_value().is_none());
460        assert!(tree.max_value().is_none());
461    }
462    #[test]
463    fn test_range_tree_single() {
464        let tree = RangeTree1D::build(&[42.0]);
465        let result = tree.range_query(42.0, 42.0);
466        assert_eq!(result, vec![0]);
467    }
468    #[test]
469    fn test_flat_rtree_multiple_overlaps() {
470        let mut tree: FlatRTree<u32> = FlatRTree::new(4);
471        for i in 0..5_u32 {
472            let x = i as f64;
473            tree.insert(
474                SpatialAabb::new(Vec3::new(x, 0.0, 0.0), Vec3::new(x + 0.5, 0.5, 0.5)),
475                i,
476            );
477        }
478        let big_query = SpatialAabb::new(Vec3::new(-1.0, -1.0, -1.0), Vec3::new(10.0, 10.0, 10.0));
479        let results = tree.query_overlap(&big_query);
480        assert_eq!(results.len(), 5, "Large query should find all 5 items");
481    }
482    #[test]
483    fn test_flat_rtree_nearest_in_new() {
484        let mut tree: FlatRTree<usize> = FlatRTree::new(4);
485        for i in 0..5_usize {
486            let x = i as f64;
487            tree.insert(
488                SpatialAabb::new(Vec3::new(x, 0.0, 0.0), Vec3::new(x + 0.1, 0.1, 0.1)),
489                i,
490            );
491        }
492        let nn = tree.nearest(Vec3::new(2.05, 0.05, 0.05)).unwrap();
493        assert_eq!(*nn, 2, "Nearest should be index 2");
494    }
495    #[test]
496    fn test_grid_index_insert_query() {
497        let mut grid = GridSpatialIndex::new(1.0);
498        for i in 0..5_usize {
499            grid.insert(Vec3::new(i as f64, 0.0, 0.0));
500        }
501        let result = grid.range_query(Vec3::new(2.0, 0.0, 0.0), 1.5);
502        assert!(result.contains(&1));
503        assert!(result.contains(&2));
504        assert!(result.contains(&3));
505        assert!(!result.contains(&0), "x=0 is 2.0 away from query");
506        assert!(!result.contains(&4), "x=4 is 2.0 away from query");
507    }
508    #[test]
509    fn test_grid_index_empty() {
510        let grid = GridSpatialIndex::new(1.0);
511        let result = grid.range_query(Vec3::zeros(), 10.0);
512        assert!(result.is_empty());
513    }
514    #[test]
515    fn test_grid_index_self_query() {
516        let mut grid = GridSpatialIndex::new(1.0);
517        let idx = grid.insert(Vec3::new(0.5, 0.5, 0.5));
518        let result = grid.range_query(Vec3::new(0.5, 0.5, 0.5), 0.1);
519        assert!(result.contains(&idx), "Should find itself");
520    }
521    #[test]
522    fn test_spatial_join_basic() {
523        let a = vec![Vec3::new(0.0, 0.0, 0.0), Vec3::new(10.0, 0.0, 0.0)];
524        let b = vec![Vec3::new(0.5, 0.0, 0.0), Vec3::new(10.5, 0.0, 0.0)];
525        let pairs = spatial_join(&a, &b, 1.0);
526        assert!(pairs.contains(&(0, 0)), "Close pair (0,0) should be found");
527        assert!(pairs.contains(&(1, 1)), "Close pair (1,1) should be found");
528        assert!(
529            !pairs.contains(&(0, 1)),
530            "Far pair (0,1) should not be found"
531        );
532    }
533    #[test]
534    fn test_spatial_join_empty() {
535        let pairs = spatial_join(&[], &[Vec3::zeros()], 1.0);
536        assert!(pairs.is_empty());
537    }
538    #[test]
539    fn test_spatial_self_join() {
540        let points = vec![
541            Vec3::new(0.0, 0.0, 0.0),
542            Vec3::new(0.5, 0.0, 0.0),
543            Vec3::new(5.0, 0.0, 0.0),
544        ];
545        let pairs = spatial_self_join(&points, 1.0);
546        assert!(pairs.contains(&(0, 1)), "Close pair should be found");
547        let far_pairs: Vec<_> = pairs.iter().filter(|&&(i, j)| i == 0 && j == 2).collect();
548        assert!(far_pairs.is_empty(), "Far pair should not be found");
549    }
550    #[test]
551    fn test_lsh_insert_and_nn() {
552        let projections = vec![[1.0, 0.0, 0.0_f64], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
553        let mut lsh = LshIndex::new(3, 0.5, &projections);
554        for i in 0..10_usize {
555            lsh.insert(Vec3::new(i as f64, 0.0, 0.0));
556        }
557        assert_eq!(lsh.len(), 10);
558        let nn = lsh.query_approx_nn(Vec3::new(7.1, 0.0, 0.0)).unwrap();
559        assert_eq!(
560            lsh.points[nn].x.round() as usize,
561            7,
562            "NN should be at x=7, got x={}",
563            lsh.points[nn].x
564        );
565    }
566    #[test]
567    fn test_lsh_empty() {
568        let lsh = LshIndex::new(2, 1.0, &[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]);
569        assert!(lsh.is_empty());
570        assert!(lsh.query_approx_nn(Vec3::zeros()).is_none());
571    }
572    #[test]
573    fn test_lsh_single_point() {
574        let projections = vec![[1.0, 0.0, 0.0_f64]];
575        let mut lsh = LshIndex::new(1, 1.0, &projections);
576        lsh.insert(Vec3::new(3.0, 4.0, 5.0));
577        let nn = lsh.query_approx_nn(Vec3::new(3.1, 4.0, 5.0)).unwrap();
578        assert_eq!(nn, 0);
579    }
580}
581/// Compute the Morton code (Z-order curve index) for a 3D integer coordinate.
582///
583/// Interleaves the bits of (x, y, z) to form a single 64-bit Morton code.
584/// Each coordinate must be < 2^21 for 63-bit Morton codes.
585#[allow(dead_code)]
586pub fn morton_encode(x: u32, y: u32, z: u32) -> u64 {
587    let x = expand_bits(x as u64);
588    let y = expand_bits(y as u64);
589    let z = expand_bits(z as u64);
590    x | (y << 1) | (z << 2)
591}
592/// Expand bits by inserting two zero bits between each pair: abcd → a00b00c00d.
593pub(super) fn expand_bits(mut v: u64) -> u64 {
594    v &= 0x00000000001fffff;
595    v = (v | (v << 32)) & 0x001f00000000ffff;
596    v = (v | (v << 16)) & 0x001f0000ff0000ff;
597    v = (v | (v << 8)) & 0x100f00f00f00f00f;
598    v = (v | (v << 4)) & 0x10c30c30c30c30c3;
599    v = (v | (v << 2)) & 0x1249249249249249;
600    v
601}
602/// Decode a Morton code back to (x, y, z) integer coordinates.
603#[allow(dead_code)]
604pub fn morton_decode(code: u64) -> (u32, u32, u32) {
605    let x = compact_bits(code);
606    let y = compact_bits(code >> 1);
607    let z = compact_bits(code >> 2);
608    (x as u32, y as u32, z as u32)
609}
610/// Compact bits: reverse of `expand_bits`.
611pub(super) fn compact_bits(mut v: u64) -> u64 {
612    v &= 0x1249249249249249;
613    v = (v | (v >> 2)) & 0x10c30c30c30c30c3;
614    v = (v | (v >> 4)) & 0x100f00f00f00f00f;
615    v = (v | (v >> 8)) & 0x001f0000ff0000ff;
616    v = (v | (v >> 16)) & 0x001f00000000ffff;
617    v = (v | (v >> 32)) & 0x00000000001fffff;
618    v
619}
620/// Batch k-nearest-neighbor query.
621///
622/// For each query point, returns the k nearest points from the tree.
623/// More efficient than calling `k_nearest` repeatedly if many queries share
624/// the same tree.
625#[allow(dead_code)]
626pub fn kd_batch_knn(tree: &KdTree, queries: &[Vec3], k: usize) -> Vec<Vec<(usize, f64)>> {
627    queries.iter().map(|&q| tree.k_nearest(q, k)).collect()
628}
629/// Find all pairs of points within distance `r` between two KD-trees.
630///
631/// Returns `(index_from_tree_a, index_from_tree_b)` pairs.
632#[allow(dead_code)]
633pub fn kd_cross_match(tree_a: &KdTree, tree_b: &KdTree, r: f64) -> Vec<(usize, usize)> {
634    let mut pairs = Vec::new();
635    for (i, &qa) in tree_a.points.iter().enumerate() {
636        let near_b = tree_b.range_search(qa, r);
637        for j in near_b {
638            pairs.push((i, j));
639        }
640    }
641    pairs
642}
643/// Compute the k-nearest neighbors of `query` from `points` using brute force.
644///
645/// Useful for verification and small datasets.
646#[allow(dead_code)]
647pub fn brute_force_knn(points: &[Vec3], query: Vec3, k: usize) -> Vec<(usize, f64)> {
648    let mut dists: Vec<(usize, f64)> = points
649        .iter()
650        .enumerate()
651        .map(|(i, &p)| (i, (p - query).norm()))
652        .collect();
653    dists.sort_unstable_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
654    dists.truncate(k);
655    dists
656}
657/// Compute the k-nearest neighbors of each point in `queries` from `points`.
658#[allow(dead_code)]
659pub fn brute_force_batch_knn(
660    points: &[Vec3],
661    queries: &[Vec3],
662    k: usize,
663) -> Vec<Vec<(usize, f64)>> {
664    queries
665        .iter()
666        .map(|&q| brute_force_knn(points, q, k))
667        .collect()
668}
669/// Build a uniform-spacing grid of points in 3D.
670///
671/// Returns points on a `(nx × ny × nz)` grid with spacing `(dx, dy, dz)` starting at `origin`.
672#[allow(dead_code)]
673pub fn regular_grid_3d(
674    origin: Vec3,
675    nx: usize,
676    ny: usize,
677    nz: usize,
678    dx: f64,
679    dy: f64,
680    dz: f64,
681) -> Vec<Vec3> {
682    let mut pts = Vec::with_capacity(nx * ny * nz);
683    for ix in 0..nx {
684        for iy in 0..ny {
685            for iz in 0..nz {
686                pts.push(Vec3::new(
687                    origin.x + ix as f64 * dx,
688                    origin.y + iy as f64 * dy,
689                    origin.z + iz as f64 * dz,
690                ));
691            }
692        }
693    }
694    pts
695}
696#[cfg(test)]
697mod tests_new_spatial {
698
699    use crate::Vec3;
700    use crate::spatial::BallTree;
701
702    use crate::spatial::KdTree;
703
704    use crate::spatial::MortonSortedIndex;
705
706    use crate::spatial::VoxelGrid;
707    use crate::spatial::brute_force_batch_knn;
708    use crate::spatial::brute_force_knn;
709    use crate::spatial::kd_batch_knn;
710    use crate::spatial::kd_cross_match;
711    use crate::spatial::morton_decode;
712    use crate::spatial::morton_encode;
713    use crate::spatial::regular_grid_3d;
714    #[test]
715    fn test_ball_tree_nearest_single_point() {
716        let pts = vec![Vec3::new(1.0, 2.0, 3.0)];
717        let tree = BallTree::build(pts);
718        let nn = tree.nearest(Vec3::new(0.0, 0.0, 0.0));
719        assert!(nn.is_some());
720        assert_eq!(nn.unwrap().0, 0);
721    }
722    #[test]
723    fn test_ball_tree_nearest_empty() {
724        let tree = BallTree::build(vec![]);
725        assert!(tree.nearest(Vec3::zeros()).is_none());
726    }
727    #[test]
728    fn test_ball_tree_nearest_multiple_points() {
729        let pts: Vec<Vec3> = (0..10).map(|i| Vec3::new(i as f64, 0.0, 0.0)).collect();
730        let tree = BallTree::build(pts);
731        let nn = tree.nearest(Vec3::new(7.1, 0.0, 0.0));
732        assert!(nn.is_some());
733        let (idx, _) = nn.unwrap();
734        assert_eq!(idx, 7, "nearest to 7.1 should be at x=7, got idx={idx}");
735    }
736    #[test]
737    fn test_ball_tree_range_query() {
738        let pts: Vec<Vec3> = (0..20).map(|i| Vec3::new(i as f64, 0.0, 0.0)).collect();
739        let tree = BallTree::build(pts);
740        let result = tree.range_query(Vec3::new(10.0, 0.0, 0.0), 2.5);
741        assert!(result.contains(&8));
742        assert!(result.contains(&10));
743        assert!(result.contains(&12));
744        assert!(!result.contains(&6), "x=6 is 4 away, should not be found");
745    }
746    #[test]
747    fn test_ball_tree_range_empty() {
748        let pts = vec![Vec3::new(100.0, 0.0, 0.0)];
749        let tree = BallTree::build(pts);
750        let result = tree.range_query(Vec3::zeros(), 1.0);
751        assert!(result.is_empty(), "far point should not be in range");
752    }
753    #[test]
754    fn test_ball_tree_len() {
755        let pts: Vec<Vec3> = (0..5).map(|i| Vec3::new(i as f64, 0.0, 0.0)).collect();
756        let tree = BallTree::build(pts);
757        assert_eq!(tree.len(), 5);
758        assert!(!tree.is_empty());
759    }
760    #[test]
761    fn test_morton_encode_origin() {
762        assert_eq!(morton_encode(0, 0, 0), 0);
763    }
764    #[test]
765    fn test_morton_encode_decode_roundtrip() {
766        for &(x, y, z) in &[(1u32, 2u32, 3u32), (7, 15, 0), (100, 200, 50)] {
767            let code = morton_encode(x, y, z);
768            let (dx, dy, dz) = morton_decode(code);
769            assert_eq!(
770                (dx, dy, dz),
771                (x, y, z),
772                "Morton roundtrip failed for ({x},{y},{z}): got ({dx},{dy},{dz})"
773            );
774        }
775    }
776    #[test]
777    fn test_morton_encode_increasing() {
778        let m0 = morton_encode(0, 0, 0);
779        let m1 = morton_encode(1, 0, 0);
780        let m2 = morton_encode(2, 0, 0);
781        assert!(m1 > m0, "Morton(1,0,0) should be > Morton(0,0,0)");
782        assert!(m2 > m1, "Morton(2,0,0) should be > Morton(1,0,0)");
783    }
784    #[test]
785    fn test_morton_decode_identity() {
786        let (x, y, z) = morton_decode(0);
787        assert_eq!((x, y, z), (0, 0, 0));
788    }
789    #[test]
790    fn test_morton_sorted_index_len() {
791        let pts: Vec<Vec3> = (0..10).map(|i| Vec3::new(i as f64, 0.0, 0.0)).collect();
792        let idx = MortonSortedIndex::build(&pts, 1024);
793        assert_eq!(idx.len(), 10);
794        assert!(!idx.is_empty());
795    }
796    #[test]
797    fn test_morton_sorted_index_sorted_order() {
798        let pts: Vec<Vec3> = (0..8).map(|i| Vec3::new(i as f64, 0.0, 0.0)).collect();
799        let idx = MortonSortedIndex::build(&pts, 64);
800        let codes: Vec<u64> = idx.sorted.iter().map(|&(c, _)| c).collect();
801        for w in codes.windows(2) {
802            assert!(
803                w[0] <= w[1],
804                "Morton codes should be sorted: {} <= {}",
805                w[0],
806                w[1]
807            );
808        }
809    }
810    #[test]
811    fn test_morton_sorted_index_empty() {
812        let idx = MortonSortedIndex::build(&[], 16);
813        assert!(idx.is_empty());
814    }
815    #[test]
816    fn test_morton_sorted_index_all_indices_present() {
817        let pts: Vec<Vec3> = (0..5).map(|i| Vec3::new(i as f64, i as f64, 0.0)).collect();
818        let idx = MortonSortedIndex::build(&pts, 64);
819        let sorted = idx.sorted_indices();
820        let mut check = sorted.clone();
821        check.sort_unstable();
822        assert_eq!(check, vec![0, 1, 2, 3, 4], "all indices should be present");
823    }
824    #[test]
825    fn test_voxel_grid_insert_query() {
826        let mut grid = VoxelGrid::new(1.0);
827        for i in 0..5 {
828            grid.insert(Vec3::new(i as f64, 0.0, 0.0));
829        }
830        let result = grid.range_query(Vec3::new(2.0, 0.0, 0.0), 1.5);
831        assert!(result.contains(&1), "x=1 should be found");
832        assert!(result.contains(&2), "x=2 should be found");
833        assert!(result.contains(&3), "x=3 should be found");
834        assert!(!result.contains(&0), "x=0 is 2 away, should not be found");
835    }
836    #[test]
837    fn test_voxel_grid_nearest_linear() {
838        let mut grid = VoxelGrid::new(1.0);
839        for i in 0..5 {
840            grid.insert(Vec3::new(i as f64, 0.0, 0.0));
841        }
842        let nn = grid.nearest_linear(Vec3::new(3.9, 0.0, 0.0)).unwrap();
843        assert_eq!(nn.0, 4, "nearest to 3.9 should be at x=4, got {}", nn.0);
844    }
845    #[test]
846    fn test_voxel_grid_empty() {
847        let grid = VoxelGrid::new(1.0);
848        assert!(grid.is_empty());
849        let result = grid.range_query(Vec3::zeros(), 10.0);
850        assert!(result.is_empty());
851        assert!(grid.nearest_linear(Vec3::zeros()).is_none());
852    }
853    #[test]
854    fn test_voxel_grid_centroid_downsampling() {
855        let mut grid = VoxelGrid::new(1.0);
856        grid.insert(Vec3::new(0.0, 0.0, 0.0));
857        grid.insert(Vec3::new(0.5, 0.0, 0.0));
858        let centroids = grid.voxel_centroids();
859        assert_eq!(
860            centroids.len(),
861            1,
862            "both points in same voxel should give 1 centroid"
863        );
864        assert!(
865            (centroids[0].x - 0.25).abs() < 1e-12,
866            "centroid should be at 0.25, got {}",
867            centroids[0].x
868        );
869    }
870    #[test]
871    fn test_brute_force_knn_basic() {
872        let pts: Vec<Vec3> = (0..10).map(|i| Vec3::new(i as f64, 0.0, 0.0)).collect();
873        let result = brute_force_knn(&pts, Vec3::new(5.0, 0.0, 0.0), 3);
874        assert_eq!(result.len(), 3);
875        assert_eq!(result[0].0, 5, "nearest should be at x=5");
876    }
877    #[test]
878    fn test_brute_force_knn_empty() {
879        let result = brute_force_knn(&[], Vec3::zeros(), 5);
880        assert!(result.is_empty());
881    }
882    #[test]
883    fn test_brute_force_knn_sorted() {
884        let pts: Vec<Vec3> = (0..10).map(|i| Vec3::new(i as f64, 0.0, 0.0)).collect();
885        let result = brute_force_knn(&pts, Vec3::new(5.0, 0.0, 0.0), 5);
886        for w in result.windows(2) {
887            assert!(w[0].1 <= w[1].1, "KNN results should be sorted by distance");
888        }
889    }
890    #[test]
891    fn test_brute_force_batch_knn() {
892        let pts: Vec<Vec3> = (0..5).map(|i| Vec3::new(i as f64, 0.0, 0.0)).collect();
893        let queries = vec![Vec3::new(0.0, 0.0, 0.0), Vec3::new(4.0, 0.0, 0.0)];
894        let result = brute_force_batch_knn(&pts, &queries, 2);
895        assert_eq!(result.len(), 2);
896        assert_eq!(result[0][0].0, 0, "nearest to x=0 should be index 0");
897        assert_eq!(result[1][0].0, 4, "nearest to x=4 should be index 4");
898    }
899    #[test]
900    fn test_kd_batch_knn() {
901        let pts: Vec<Vec3> = (0..10).map(|i| Vec3::new(i as f64, 0.0, 0.0)).collect();
902        let tree = KdTree::build(pts);
903        let queries = vec![Vec3::new(3.0, 0.0, 0.0), Vec3::new(8.0, 0.0, 0.0)];
904        let results = kd_batch_knn(&tree, &queries, 2);
905        assert_eq!(results.len(), 2);
906        assert_eq!(results[0].len(), 2);
907        assert_eq!(results[1].len(), 2);
908    }
909    #[test]
910    fn test_kd_cross_match_close_pairs() {
911        let pts_a: Vec<Vec3> = vec![Vec3::new(0.0, 0.0, 0.0), Vec3::new(10.0, 0.0, 0.0)];
912        let pts_b: Vec<Vec3> = vec![Vec3::new(0.3, 0.0, 0.0), Vec3::new(10.3, 0.0, 0.0)];
913        let tree_a = KdTree::build(pts_a);
914        let tree_b = KdTree::build(pts_b);
915        let pairs = kd_cross_match(&tree_a, &tree_b, 0.5);
916        assert!(pairs.contains(&(0, 0)), "close pair (0,0) should be found");
917        assert!(pairs.contains(&(1, 1)), "close pair (1,1) should be found");
918    }
919    #[test]
920    fn test_kd_cross_match_no_pairs() {
921        let pts_a: Vec<Vec3> = vec![Vec3::new(0.0, 0.0, 0.0)];
922        let pts_b: Vec<Vec3> = vec![Vec3::new(100.0, 0.0, 0.0)];
923        let tree_a = KdTree::build(pts_a);
924        let tree_b = KdTree::build(pts_b);
925        let pairs = kd_cross_match(&tree_a, &tree_b, 1.0);
926        assert!(pairs.is_empty(), "far-apart points should have no pairs");
927    }
928    #[test]
929    fn test_regular_grid_3d_count() {
930        let pts = regular_grid_3d(Vec3::zeros(), 3, 4, 5, 1.0, 1.0, 1.0);
931        assert_eq!(pts.len(), 60, "3x4x5 grid should have 60 points");
932    }
933    #[test]
934    fn test_regular_grid_3d_origin() {
935        let pts = regular_grid_3d(Vec3::new(1.0, 2.0, 3.0), 2, 2, 2, 1.0, 1.0, 1.0);
936        assert_eq!(pts.len(), 8);
937        assert!((pts[0].x - 1.0).abs() < 1e-12);
938        assert!((pts[0].y - 2.0).abs() < 1e-12);
939        assert!((pts[0].z - 3.0).abs() < 1e-12);
940    }
941    #[test]
942    fn test_regular_grid_3d_spacing() {
943        let pts = regular_grid_3d(Vec3::zeros(), 3, 1, 1, 0.5, 1.0, 1.0);
944        assert_eq!(pts.len(), 3);
945        assert!(
946            (pts[1].x - 0.5).abs() < 1e-12,
947            "second point should be at x=0.5"
948        );
949        assert!(
950            (pts[2].x - 1.0).abs() < 1e-12,
951            "third point should be at x=1.0"
952        );
953    }
954}