kiddo 6.0.2

A high-performance, flexible, ergonomic k-d tree library. Ideal for geo- and astro- nearest-neighbour and k-nearest-neighbor queries
Documentation
use kiddo::dist::{
    DistanceMetric, DistanceMetricAvx2, DistanceMetricAvx512, DistanceMetricNeon,
    DistanceMetricScalar,
};
use kiddo::leaf_strategies::VecOfArenas;
use kiddo::stem_strategies::{DonnellySimdFull, Eytzinger};
use kiddo::KdTree;

struct AbsDistance;

impl DistanceMetricScalar<f64> for AbsDistance {
    type Output = f64;

    fn widen_coord(a: f64) -> Self::Output {
        a
    }

    fn dist1(a: Self::Output, b: Self::Output) -> Self::Output {
        (a - b).abs()
    }
}

impl DistanceMetricAvx512<f64> for AbsDistance {
    #[cfg(all(feature = "simd", target_feature = "avx512f"))]
    type Avx512F64Ops = kiddo::traits::dist::UnsupportedAvx512F64LeafOps;

    #[cfg(all(feature = "simd", target_feature = "avx512f"))]
    type Avx512F32Ops = kiddo::traits::dist::UnsupportedAvx512F32LeafOps;
}

impl DistanceMetricAvx2<f64> for AbsDistance {
    #[cfg(all(feature = "simd", target_arch = "x86_64", target_feature = "avx2"))]
    type Avx2F64Ops = kiddo::traits::dist::UnsupportedAvx2F64LeafOps;

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

impl DistanceMetricNeon<f64> for AbsDistance {
    #[cfg(all(feature = "simd", target_arch = "aarch64", target_feature = "neon"))]
    type NeonF64Ops = kiddo::traits::dist::UnsupportedNeonF64LeafOps;

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

type GeneralTree = KdTree<f64, u32, Eytzinger, VecOfArenas<f64, u32, 2, 32>, 2, 32>;
type BlockTree = KdTree<f64, u32, DonnellySimdFull<3>, VecOfArenas<f64, u32, 2, 32>, 2, 32>;

fn assert_distance_metric<M: DistanceMetric<f64>>() {}

#[test]
fn external_metric_works_for_general_and_block_strategies() {
    assert_distance_metric::<AbsDistance>();

    let points = [[0.0, 0.0], [1.0, 1.0], [2.0, 0.5], [4.0, 4.0]];
    let query = [1.9, 0.6];

    let general_tree: GeneralTree = KdTree::new_from_slice(&points).unwrap();
    let block_tree: BlockTree = KdTree::new_from_slice(&points).unwrap();

    let general_nearest = general_tree
        .query(&query)
        .nearest_one::<AbsDistance>()
        .execute();
    let block_nearest = block_tree
        .query(&query)
        .nearest_one::<AbsDistance>()
        .execute();

    assert_eq!(general_nearest.item, 2);
    assert_eq!(block_nearest.item, 2);
    assert!((general_nearest.distance - 0.2).abs() < 1e-12);
    assert!((block_nearest.distance - 0.2).abs() < 1e-12);
}