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

use crate::dist::DistanceMetric;
use crate::leaf_view::{LeafArena, LeafView, TlsLeafScratch};
use crate::results::result_collection::{FromLeafCandidate, ResultCollection};
use crate::{Axis, Content};

#[inline(always)]
pub(crate) fn nearest_n_within_with_query_wide_fallback<
    AX,
    T,
    D,
    E,
    R,
    const EXCLUSIVE: bool,
    const K: usize,
    const B: usize,
>(
    leaf: &LeafView<'_, AX, T, K, B>,
    query_wide: &[D::Output; K],
    dist: D::Output,
    results: &mut R,
) where
    AX: Axis<Coord = AX> + 'static,
    T: Content,
    D: DistanceMetric<AX>,
    D::Output: Axis<Coord = D::Output> + TlsLeafScratch + 'static,
    E: FromLeafCandidate<AX, T, D::Output, K>,
    R: ResultCollection<D::Output, E>,
{
    leaf.with_dists_for_slice_wide::<D, _>(query_wide, |dists| {
        leaf.update_nearest_dists::<_, E, _, EXCLUSIVE>(dists, leaf.items(), dist, results);
    });
}

#[inline(always)]
pub(crate) fn nearest_n_within_with_query_wide_arena_fallback<
    AX,
    T,
    D,
    E,
    R,
    const EXCLUSIVE: bool,
    const K: usize,
>(
    arena: &LeafArena<'_, AX, T, K>,
    query_wide: &[D::Output; K],
    dist: D::Output,
    results: &mut R,
) where
    AX: Axis<Coord = AX> + 'static,
    T: Content,
    D: DistanceMetric<AX>,
    D::Output: Axis<Coord = D::Output> + 'static,
    E: FromLeafCandidate<AX, T, D::Output, K>,
    R: ResultCollection<D::Output, E>,
{
    if arena.is_empty() {
        return;
    }

    arena.for_each_tiled_chunk(|tile| {
        for idx in 0..tile.len() {
            let mut candidate_dist = D::Output::zero();

            for dim in 0..K {
                let coord = unsafe { tile.point_unaligned(dim, idx) };
                D::combine_component(
                    &mut candidate_dist,
                    D::dist1(D::widen_coord(coord), unsafe {
                        *query_wide.get_unchecked(dim)
                    }),
                );
            }

            let is_within_dist = if EXCLUSIVE {
                candidate_dist < dist
            } else {
                candidate_dist <= dist
            };

            if is_within_dist {
                #[cfg(feature = "result_collection_stats")]
                crate::results::result_collection_stats::record_candidate_emitted();

                results.add(E::from_leaf_candidate(
                    tile.point(idx),
                    unsafe { tile.item_unaligned(idx) },
                    candidate_dist,
                ));
            }
        }
    });
}