kiddo 5.3.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
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176

use crate::kiddo::Axis;

// pub (crate) struct BoundsChecker<'a, F, A, const K: usize> where
//     F: Fn(&[A; K], &[A; K]) -> A,
//     A: Axis,
// {
//     pub(crate) min_bound: [A; K],
//     pub(crate) max_bound: [A; K],
//     pub(crate) distance_fn: &'a F,
// }
//
// impl<'a, 'b, F, A: Axis, const K: usize> BoundsChecker<'a, F, A, K> where
//     F: Fn(&[A; K], &[A; K]) -> A,
//     A: Axis,
// {
//     #[inline(never)]
//     pub(crate) fn new(
//         min_bound: [A; K],
//         max_bound: [A; K],
//         distance_fn: &'a F,
//     ) -> BoundsChecker<'a,F, A, K>
//     {
//         BoundsChecker {
//             min_bound,
//             max_bound,
//             distance_fn
//         }
//     }
// }

//
// impl<'a, 'b, F, A, const K: usize> BoundsChecker<'a, F, A, K> where
//     F: Fn(&[A; K], &[A; K]) -> A,
//     A: Axis
// {
//     #[inline(never)]
//     pub(crate) fn dist_to_bound(&self, p1: [A; K]) -> A
//     where
//     F: Fn(&[A; K], &[A; K]) -> A,
//     A: Axis,
//     {
//         let mut p2 = [A::nan(); K];
//
//         //// SLOWER
//         // let mut p2 = [A::nan(); K];
//         // for i in 0..K {
//         //     p2[i] = clamp(p1[i], bounds[i]);
//         // }
//
//         p1.into_iter()
//             .zip(self.min_bound)
//             .zip(&self.max_bound)
//             .map(|((v, min_bound), &max_bound)| clamp(v, min_bound, max_bound))
//             .zip(p2.iter_mut())
//             .for_each(|(clamped_val, p2_coord)| *p2_coord = clamped_val);
//
//         (self.distance_fn)(&p1, &p2)
//     }
// }

// #[inline(never)]
pub(crate) fn distance_to_bounds<F, A, const K: usize>(
    p1: &[A; K],
    min_bound: &[A; K],
    max_bound: &[A; K],
    distance: &F,
) -> A
where
    F: Fn(&[A; K], &[A; K]) -> A,
    A: Axis,
{
    let mut p2 = [A::nan(); K];

    //// SLOWER
    // for i in 0..K {
    //     p2[i] = clamp(p1[i], bounds[i]);
    // }

    p1.iter()
        .zip(min_bound)
        .zip(max_bound)
        .map(|((&v, &min_bound), &max_bound)| clamp(v, min_bound, max_bound))
        .zip(p2.iter_mut())
        .for_each(|(clamped_val, p2_coord)| *p2_coord = clamped_val);

    distance(p1, &p2)
}

// #[inline(never)]
#[allow(dead_code)]
pub(crate) fn distance_to_bounds_simd<F, A, const K: usize>(
    p1: &[A; K],
    min_bound: &[A; K],
    max_bound: &[A; K],
    distance: &F,
) -> A
where
    F: Fn(&[A; K], &[A; K]) -> A,
    A: Axis,
{
    let mut p2 = [A::nan(); K];

    //// SLOWER
    // for i in 0..K {
    //     p2[i] = clamp(p1[i], bounds[i]);
    // }

    p1.iter()
        .zip(min_bound)
        .zip(max_bound)
        .map(|((&v, &min_bound), &max_bound)| clamp(v, min_bound, max_bound))
        .zip(p2.iter_mut())
        .for_each(|(clamped_val, p2_coord)| *p2_coord = clamped_val);

    distance(p1, &p2)
}

// #[inline(never)]
pub(crate) fn clamp<A: Axis>(val: A, min_bound: A, max_bound: A) -> A {
    if val < min_bound {
        min_bound
    } else if val > max_bound {
        max_bound
    } else {
        val
    }
}

// #[inline(never)]
pub(crate) fn extend<A: Axis, const K: usize>(
    min_bound: &mut [A; K],
    max_bound: &mut [A; K],
    point: &[A; K],
) {
    min_bound.iter_mut().enumerate().for_each(|(dim, bound)| {
        if point[dim] < *bound {
            *bound = point[dim];
        }
    });

    max_bound.iter_mut().enumerate().for_each(|(dim, bound)| {
        if point[dim] > *bound {
            *bound = point[dim];
        }
    });
}

//// SLOWER
// #[inline(never)]
// pub(crate) fn extend<A: Axis, const K: usize>(bounds: &mut [(A, A); K], point: &[A; K]) {
//     point.iter().zip(bounds.iter_mut()).for_each(|(&point_coord, bound)| {
//         if point_coord < bound.0 {
//             bound.0 = point_coord;
//         }
//         if point_coord > bound.1 {
//             bound.1 = point_coord;
//         }
//     });
// }

//// SLOWER
// #[inline(never)]
// pub(crate) fn extend<A: Axis, const K: usize>(bounds: &mut [(A, A); K], point: &[A; K]) {
//     point.iter().zip(bounds.iter_mut()).for_each(|(&point_coord, bound)| {
//         bound.0 = bound.0.min(point_coord);
//         bound.1 = bound.1.max(point_coord);
//     });
// }

//// DOESN'T WORK
// #[inline(never)]
// pub(crate) fn extend<A: Axis, const K: usize>(bounds: &mut [(A, A); K], point: &[A; K]) {
//     bounds = point.iter().zip(bounds.iter()).map(|(&point_coord, (&lo, &hi))| {
//         (lo.min(point_coord), hi.max(point_coord))
//     }).collect();
// }