kiddo 6.0.0-alpha.1

A high-performance, flexible, ergonomic k-d tree library. Ideal for geo- and astro- nearest-neighbour and k-nearest-neighbor queries
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158

use std::fmt::{self, Display};
use std::num::NonZero;

use kiddo::kd_tree::KdTree;
use kiddo::leaf_strategy::VecOfArrays;
use kiddo::{Eytzinger, SquaredEuclidean};

type TestTree<T> = KdTree<f64, T, Eytzinger, VecOfArrays<f64, T, 2, 32>, 2, 32>;

#[derive(Clone, Copy, Debug, PartialEq, Default, Eq, Ord, PartialOrd)]
struct MyFixedString {
    bytes: [u8; 32],
}

impl MyFixedString {
    fn from_str(s: &str) -> Self {
        let mut bytes = [0u8; 32];
        let copy_length = s.len().min(32);
        bytes[..copy_length].copy_from_slice(&s.as_bytes()[..copy_length]);
        Self { bytes }
    }

    fn as_str(&self) -> &str {
        let len = self.bytes.iter().position(|&b| b == 0).unwrap_or(32);
        std::str::from_utf8(&self.bytes[..len]).unwrap()
    }
}

impl Display for MyFixedString {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

#[test]
fn test_kdtree_with_numeric_id() {
    // Create a new KdTree with 2D points (K=2) and u32 as the data type
    let mut tree: TestTree<u32> = KdTree::default();

    // Add some points with associated numeric IDs
    tree.add(&[0.0, 0.0], 1001).unwrap();
    tree.add(&[1.0, 1.0], 1002).unwrap();
    tree.add(&[2.0, 2.0], 1003).unwrap();
    tree.add(&[-1.0, -1.0], 1004).unwrap();

    // Test nearest neighbor query
    let query_point = [0.5, 0.25];
    let nearest = tree
        .query(&query_point)
        .nearest_one::<SquaredEuclidean<f64>>()
        .execute();

    // The closest point should be [0.0, 0.0] with ID 1001
    assert_eq!(nearest.item, 1001);
    assert!((nearest.distance - (0.5f64 * 0.5 + 0.25 * 0.25)).abs() < f64::EPSILON);

    // Test k-nearest neighbors
    let k_nearest = tree
        .query(&query_point)
        .nearest_n::<SquaredEuclidean<f64>>(NonZero::new(2).unwrap())
        .execute();
    assert_eq!(k_nearest.len(), 2);
    assert_eq!(k_nearest[0].item, 1001);
    assert_eq!(k_nearest[1].item, 1002);

    // Test within radius
    let radius = 2.0;
    let within_results = tree
        .query(&query_point)
        .within::<SquaredEuclidean<f64>>(radius)
        .execute();
    assert_eq!(within_results.len(), 2); // Should find two points within radius
    assert!(within_results.iter().any(|r| r.item == 1001));
    assert!(within_results.iter().any(|r| r.item == 1002));
}

#[test]
/// Test case with an esoteric T type
fn test_kdtree_with_fixed_string() {
    // Create a new KdTree with 2D points (K=2) and MyFixedString as the data type
    let mut tree: TestTree<MyFixedString> = KdTree::default();

    // Add some points with associated MyFixedString data
    tree.add(&[0.0, 0.0], MyFixedString::from_str("Origin"))
        .unwrap();
    tree.add(&[1.0, 1.0], MyFixedString::from_str("Point A"))
        .unwrap();
    tree.add(&[2.0, 2.0], MyFixedString::from_str("Point B"))
        .unwrap();
    tree.add(&[-1.0, -1.0], MyFixedString::from_str("Point C"))
        .unwrap();

    // Test nearest neighbor query
    let query_point = [0.5, 0.25];
    let nearest = tree
        .query(&query_point)
        .nearest_one::<SquaredEuclidean<f64>>()
        .execute();

    // The closest point should be [0.0, 0.0] with data "Origin"
    assert_eq!(nearest.item.as_str(), "Origin");
    assert!((nearest.distance - (0.5f64 * 0.5 + 0.25 * 0.25)).abs() < f64::EPSILON);

    // Test k-nearest neighbors
    let k_nearest = tree
        .query(&query_point)
        .nearest_n::<SquaredEuclidean<f64>>(NonZero::new(2).unwrap())
        .execute();
    assert_eq!(k_nearest.len(), 2);
    assert_eq!(k_nearest[0].item.as_str(), "Origin");
    assert_eq!(k_nearest[1].item.as_str(), "Point A");

    // Test within radius
    let radius = 2.0;
    let within_results = tree
        .query(&query_point)
        .within::<SquaredEuclidean<f64>>(radius)
        .execute();
    assert_eq!(within_results.len(), 2); // Should find "Origin" and "Point A"
}

#[test]
fn test_kdtree_with_empty_type_as_content() {
    // Create a new KdTree with 2D points (K=2) and () as the data type
    let mut tree: TestTree<()> = KdTree::default();

    // Add some points with associated numeric IDs
    tree.add(&[0.0, 0.0], ()).unwrap();
    tree.add(&[1.0, 1.0], ()).unwrap();
    tree.add(&[2.0, 2.0], ()).unwrap();
    tree.add(&[-1.0, -1.0], ()).unwrap();

    // Test nearest neighbor query
    let query_point = [0.5, 0.25];
    let nearest = tree
        .query(&query_point)
        .nearest_one::<SquaredEuclidean<f64>>()
        .execute();

    // The closest point should be [0.0, 0.0]
    assert!((nearest.distance - (0.5f64 * 0.5 + 0.25 * 0.25)).abs() < f64::EPSILON);
    assert_eq!(nearest.item, ());

    // Test k-nearest neighbors
    let k_nearest = tree
        .query(&query_point)
        .nearest_n::<SquaredEuclidean<f64>>(NonZero::new(2).unwrap())
        .execute();
    assert_eq!(k_nearest.len(), 2);

    // Test within radius
    let radius = 2.0;
    let within_results = tree
        .query(&query_point)
        .within::<SquaredEuclidean<f64>>(radius)
        .execute();
    assert_eq!(within_results.len(), 2); // Should find two points within radius
}