kiddo 6.0.0-alpha.1

A high-performance, flexible, ergonomic k-d tree library. Ideal for geo- and astro- nearest-neighbour and k-nearest-neighbor queries
Documentation
#![cfg_attr(coverage_nightly, feature(coverage_attribute))]
#![cfg_attr(coverage_nightly, coverage(off))]

use kiddo::kd_tree::KdTree;
use kiddo::leaf_strategy::FlatVec;
use kiddo::Eytzinger;
use kiddo::SquaredEuclidean;
use rand::{RngExt, SeedableRng};
use rand_chacha::ChaCha8Rng;
use std::hint::black_box;
use std::time::Instant;

const DEFAULT_POINT_COUNT: usize = 10_000_000;
const DEFAULT_QUERY_COUNT: usize = 1_000;
const DEFAULT_QUERY_BATCH_REPEATS: usize = 2_000;
const POINT_SEED: u64 = 0x5eed_0000_0000_0001;
const QUERY_SEED: u64 = 0x5eed_0000_0000_0002;

type Tree = KdTree<f64, u32, Eytzinger, FlatVec<f64, u32, 3, 32>, 3, 32>;

fn read_usize_env(var: &str, default: usize) -> usize {
    std::env::var(var)
        .ok()
        .and_then(|value| value.parse::<usize>().ok())
        .unwrap_or(default)
}

fn main() {
    let point_count = read_usize_env("KIDDO_PROFILE_POINTS", DEFAULT_POINT_COUNT);
    let query_count = read_usize_env("KIDDO_PROFILE_QUERIES", DEFAULT_QUERY_COUNT);
    let query_batch_repeats = read_usize_env(
        "KIDDO_PROFILE_QUERY_BATCH_REPEATS",
        DEFAULT_QUERY_BATCH_REPEATS,
    );

    eprintln!(
        "profiling v6 nearest_one: points={} queries={} query_batch_repeats={} point_seed={} query_seed={}",
        point_count, query_count, query_batch_repeats, POINT_SEED, QUERY_SEED
    );

    let mut point_rng = ChaCha8Rng::seed_from_u64(POINT_SEED);
    let points: Vec<[f64; 3]> = (0..point_count)
        .map(|_| point_rng.random::<[f64; 3]>())
        .collect();

    let build_start = Instant::now();
    let tree: Tree = KdTree::new_from_slice(&points).unwrap();
    let build_elapsed = build_start.elapsed();

    let mut query_rng = ChaCha8Rng::seed_from_u64(QUERY_SEED);
    let queries: Vec<[f64; 3]> = (0..query_count)
        .map(|_| query_rng.random::<[f64; 3]>())
        .collect();

    let query_start = Instant::now();
    let mut checksum_dist = 0.0f64;
    let mut checksum_item = 0u64;

    for _ in 0..query_batch_repeats {
        for query in &queries {
            let result = tree
                .query(black_box(query))
                .nearest_one::<SquaredEuclidean<f64>>()
                .execute();
            checksum_dist += result.distance;
            checksum_item = checksum_item.wrapping_add(result.item as u64);
        }
    }

    let query_elapsed = query_start.elapsed();

    println!(
        "v6 build={:?} query_total={:?} query_per_batch={:?} checksum_dist={:.17e} checksum_item={}",
        build_elapsed,
        query_elapsed,
        query_elapsed / (query_batch_repeats as u32),
        checksum_dist,
        checksum_item
    );
}