kiddo 6.0.0

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

use std::ops::{Add, Mul};

use fixed::traits::LossyFrom;

use crate::Axis;

use crate::dist::{
    DistanceMetricAvx2, DistanceMetricAvx512, DistanceMetricNeon, DistanceMetricScalar,
};

#[cfg(all(feature = "simd", target_arch = "x86_64", target_feature = "avx2"))]
mod avx2;

#[cfg(all(feature = "simd", target_feature = "avx512f"))]
mod avx512;

#[cfg(all(feature = "simd", target_arch = "aarch64", target_feature = "neon"))]
mod neon;

/// Squared Euclidean distance metric, parameterized by output type `R`.
pub struct SquaredEuclidean<R>(core::marker::PhantomData<R>);

impl<A, R> DistanceMetricScalar<A> for SquaredEuclidean<R>
where
    A: Copy,
    R: Axis<Coord = R> + LossyFrom<A> + Mul<Output = R> + Add<Output = R>,
{
    type Output = R;

    #[inline(always)]
    fn widen_coord(a: A) -> R {
        R::lossy_from(a)
    }

    #[inline(always)]
    fn dist1(a: R, b: R) -> R {
        let d = if R::IS_SIGNED || a >= b { a - b } else { b - a };
        d * d
    }
}

impl<A, R> DistanceMetricAvx512<A> for SquaredEuclidean<R>
where
    A: Copy,
    R: Axis<Coord = R> + LossyFrom<A> + Mul<Output = R> + Add<Output = R>,
{
    #[cfg(all(feature = "simd", target_feature = "avx512f"))]
    type Avx512F64Ops = avx512::SquaredEuclideanAvx512F64LeafOps;

    #[cfg(all(feature = "simd", target_feature = "avx512f"))]
    type Avx512F32Ops = avx512::SquaredEuclideanAvx512F32LeafOps;
}

impl<A, R> DistanceMetricAvx2<A> for SquaredEuclidean<R>
where
    A: Copy,
    R: Axis<Coord = R> + LossyFrom<A> + Mul<Output = R> + Add<Output = R>,
{
    #[cfg(all(feature = "simd", target_arch = "x86_64", target_feature = "avx2"))]
    type Avx2F64Ops = avx2::SquaredEuclideanAvx2F64LeafOps;

    #[cfg(all(feature = "simd", target_arch = "x86_64", target_feature = "avx2"))]
    type Avx2F32Ops = avx2::SquaredEuclideanAvx2F32LeafOps;
}

impl<A, R> DistanceMetricNeon<A> for SquaredEuclidean<R>
where
    A: Copy,
    R: Axis<Coord = R> + LossyFrom<A> + Mul<Output = R> + Add<Output = R>,
{
    #[cfg(all(feature = "simd", target_arch = "aarch64", target_feature = "neon"))]
    type NeonF64Ops = neon::SquaredEuclideanNeonF64LeafOps;

    #[cfg(all(feature = "simd", target_arch = "aarch64", target_feature = "neon"))]
    type NeonF32Ops = neon::SquaredEuclideanNeonF32LeafOps;
}

#[cfg(feature = "cargo_asm")]
pub mod cargo_asm {
    use super::SquaredEuclidean;
    use crate::dist::DistanceMetricScalar;

    #[inline(never)]
    #[unsafe(no_mangle)]
    pub fn v6_squared_euclidean_dist1_f64_cargo_asm_hook(a: f64, b: f64) -> f64 {
        <SquaredEuclidean<f64> as DistanceMetricScalar<f64>>::dist1(a, b)
    }

    #[inline(never)]
    #[unsafe(no_mangle)]
    pub fn v6_squared_euclidean_dist1_u16_cargo_asm_hook(a: u16, b: u16) -> u16 {
        <SquaredEuclidean<u16> as DistanceMetricScalar<u16>>::dist1(a, b)
    }
}

#[cfg(test)]
mod tests {
    use super::SquaredEuclidean;
    use crate::dist::DistanceMetricScalar;

    #[test]
    fn signed_dist1_squares_direct_delta_in_both_directions() {
        type Metric = SquaredEuclidean<f64>;

        assert_eq!(<Metric as DistanceMetricScalar<f64>>::dist1(5.0, 2.0), 9.0);
        assert_eq!(<Metric as DistanceMetricScalar<f64>>::dist1(2.0, 5.0), 9.0);
    }

    #[test]
    fn unsigned_dist1_uses_ordered_subtraction() {
        type Metric = SquaredEuclidean<u16>;

        assert_eq!(<Metric as DistanceMetricScalar<u16>>::dist1(5, 2), 9);
        assert_eq!(<Metric as DistanceMetricScalar<u16>>::dist1(2, 5), 9);
    }

    #[cfg(feature = "fixed")]
    #[test]
    fn signed_fixed_dist1_squares_direct_delta() {
        use fixed::{types::extra::U16, FixedI32};

        type Output = FixedI32<U16>;
        type Metric = SquaredEuclidean<Output>;
        let a = Output::from_num(-1.5);
        let b = Output::from_num(0.5);

        assert_eq!(
            <Metric as DistanceMetricScalar<Output>>::dist1(a, b),
            Output::from_num(4)
        );
        assert_eq!(
            <Metric as DistanceMetricScalar<Output>>::dist1(b, a),
            Output::from_num(4)
        );
    }
}