kiddo 6.0.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

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;

/// Chebyshev / L-infinity distance metric, parameterized by output type `R`.
pub struct Chebyshev<R>(core::marker::PhantomData<R>);

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

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

    #[inline(always)]
    fn dist1(a: R, b: R) -> R {
        R::saturating_dist(a, b)
    }

    #[inline(always)]
    fn combine_component(acc: &mut Self::Output, component: Self::Output) {
        if component > *acc {
            *acc = component;
        }
    }

    #[inline(always)]
    fn rect_dist_after_update<const K: usize>(
        _rd: Self::Output,
        off: &[Self::Output; K],
        dim: usize,
        new_off: Self::Output,
    ) -> Self::Output {
        let mut acc = Self::Output::zero();

        for axis in 0..K {
            let off_val = if axis == dim { new_off } else { off[axis] };
            Self::combine_component(&mut acc, off_val);
        }

        acc
    }

    #[inline(always)]
    fn rect_dist_after_axis_update(rd: R, old_off: R, new_off: R) -> R {
        // `rd` is the largest offset across all axes. When this axis was not the largest,
        // the other axes still account for `rd` and the update is exact. When it was (or
        // tied for) the largest, the other axes' maximum is not recoverable from `rd`
        // alone, so fall back to this axis alone: an under-estimate that keeps the sibling
        // in the search rather than pruning it incorrectly.
        if R::cmp(old_off, rd) == core::cmp::Ordering::Less {
            R::max(rd, new_off)
        } else {
            new_off
        }
    }
}

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

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

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

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

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

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