k-dimensional tree

Terada Yuichiro

Last update: Dec 2, 2022

Related tags

Command-line kd-tree

Overview

kd-tree

k-dimensional tree in Rust.

Fast, simple, and easy to use.

Usage

// construct kd-tree
let kdtree = kd_tree::KdTree::build_by_ordered_float(vec![
    [1.0, 2.0, 3.0],
    [3.0, 1.0, 2.0],
    [2.0, 3.0, 1.0],
]);

// search the nearest neighbor
let found = kdtree.nearest(&[3.1, 0.9, 2.1]).unwrap();
assert_eq!(found.item, &[3.0, 1.0, 2.0]);

// search k-nearest neighbors
let found = kdtree.nearests(&[1.5, 2.5, 1.8], 2);
assert_eq!(found[0].item, &[2.0, 3.0, 1.0]);
assert_eq!(found[1].item, &[1.0, 2.0, 3.0]);

// search points within a sphere
let found = kdtree.within_radius(&[2.0, 1.5, 2.5], 1.5);
assert_eq!(found.len(), 2);
assert!(found.iter().any(|&&p| p == [1.0, 2.0, 3.0]));
assert!(found.iter().any(|&&p| p == [3.0, 1.0, 2.0]));

With or without `KdPoint`

KdPoint trait represents k-dimensional point.

You can live with or without KdPoint.

With `KdPoint` explicitly

use kd_tree::{KdPoint, KdTree};

// define your own item type.
struct Item {
    point: [f64; 2],
    id: usize,
}

// implement `KdPoint` for your item type.
impl KdPoint for Item {
    type Scalar = f64;
    type Dim = typenum::U2; // 2 dimensional tree.
    fn at(&self, k: usize) -> f64 { self.point[k] }
}

// construct kd-tree from `Vec<Item>`.
// Note: you need to use `build_by_ordered_float()` because f64 doesn't implement `Ord` trait.
let kdtree: KdTree<Item> = KdTree::build_by_ordered_float(vec![
    Item { point: [1.0, 2.0], id: 111 },
    Item { point: [2.0, 3.0], id: 222 },
    Item { point: [3.0, 4.0], id: 333 },
]);

// search nearest item from [1.9, 3.1]
assert_eq!(kdtree.nearest(&[1.9, 3.1]).unwrap().item.id, 222);

With `KdPoint` implicitly

KdPoint trait is implemented for fixed-sized array of numerical types, such as [f64; 3] or [i32, 2] etc. So you can build kd-trees of those types without custom implementation of KdPoint.

let items: Vec<[i32; 3]> = vec![[1, 2, 3], [3, 1, 2], [2, 3, 1]];
let kdtree = kd_tree::KdTree::build(items);
assert_eq!(kdtree.nearest(&[3, 1, 2]).unwrap().item, &[3, 1, 2]);

KdPoint trait is also implemented for tuple of a KdPoint and an arbitrary type, like (P, T) where P: KdPoint. And a type alias named KdMap<P, T> is defined as KdTree<(P, T)>. So you can build a kd-tree from key-value pairs, as below:

let kdmap: kd_tree::KdMap<[isize; 3], &'static str> = kd_tree::KdMap::build(vec![
    ([1, 2, 3], "foo"),
    ([2, 3, 1], "bar"),
    ([3, 1, 2], "buzz"),
]);
assert_eq!(kdmap.nearest(&[3, 1, 2]).unwrap().item.1, "buzz");

`nalgebra` feature

KdPoint trait is implemented for nalgebra's vectors and points.

Enable nalgebra feature in your Cargo.toml:

kd-tree = { version = "...", features = ["nalgebra"] }

Then, you can use it as follows:

use nalgebra::Point3;
let items: Vec<Point3<i32>> = vec![
    Point3::new(1, 2, 3),
    Point3::new(3, 1, 2),
    Point3::new(2, 3, 1)
];
let kdtree = kd_tree::KdTree::build(items);
let query = Point3::new(3, 1, 2);
assert_eq!(kdtree.nearest(&query).unwrap().item, &query);

Without `KdPoint`

use std::collections::HashMap;
let items: HashMap<&'static str, [i32; 2]> = vec![
    ("a", [10, 20]),
    ("b", [20, 10]),
    ("c", [20, 20]),
].into_iter().collect();
let kdtree = kd_tree::KdTree2::build_by_key(items.keys().collect(), |key, k| items[*key][k]);
assert_eq!(kdtree.nearest_by(&[18, 21], |key, k| items[*key][k]).unwrap().item, &&"c");

To own, or not to own

KdSliceN<T, N> and KdTreeN<T, N> are similar to str and String, or Path and PathBuf.

KdSliceN<T, N> doesn't own its buffer, but KdTreeN<T, N>.
KdSliceN<T, N> is not Sized, so it must be dealed in reference mannar.
KdSliceN<T, N> implements Deref to [T].
KdTreeN<T, N> implements Deref to KdSliceN<T, N>.
Unlike PathBuf or String, which are mutable, KdTreeN<T, N> is immutable.

&KdSliceN<T, N> can be constructed directly, not via KdTreeN, as below:

let mut items: Vec<[i32; 3]> = vec![[1, 2, 3], [3, 1, 2], [2, 3, 1]];
let kdtree = kd_tree::KdSlice::sort(&mut items);
assert_eq!(kdtree.nearest(&[3, 1, 2]).unwrap().item, &[3, 1, 2]);

`KdIndexTreeN`

A KdIndexTreeN refers a slice of items, [T], and contains kd-tree of indices to the items, KdTreeN<usize, N>. Unlike [KdSlice::sort], [KdIndexTree::build] doesn't sort input items.

let items = vec![[1, 2, 3], [3, 1, 2], [2, 3, 1]];
let kdtree = kd_tree::KdIndexTree::build(&items);
assert_eq!(kdtree.nearest(&[3, 1, 2]).unwrap().item, &1); // nearest() returns an index of found item.

Semantic find-and-replace using tree-sitter-based macro expansion

15 Nov 10, 2022

lemmy-help is a emmylua parser as well as a CLI which takes that parsed tree and converts it into vim help docs.

117 Jan 3, 2023

Tree-based TUI chat client

cove Cove is a TUI client for euphoria.io, a threaded real-time chat platform. It runs on Linux, Windows and macOS. Manual installation This section c

8 Nov 16, 2022

Work in progress NCBI's Common Tree alternative in the terminal

Lomanai Usage lomanai --species 'Mus musculus' --species 'Homo sapiens' # Mammalia # ++Rodentia # | \-Mus musculus # \+Primates # \-Homo sapien

3 Dec 20, 2022

Mypyc DSL grammar for tree-sitter

tree-sitter-mypyc Mypyc DSL grammar for tree-sitter. Installing (Neovim) This is based on the Neovim Tree-sitter docs for adding new parsers. Basicall

3 Dec 30, 2022

tree-sitter meets Kakoune

kak-tree-sitter This is a binary server that interfaces tree-sitter with kakoune. Features Install Usage Design Credits Features Semantic highlighting

5 May 3, 2023

Build Abstract Syntax Trees and tree-walking models quickly in Rust.

astmaker Build Abstract Syntax Trees and tree-walking models quickly in Rust. Example This example creates an AST for simple math expressions, and an

100 Jun 5, 2023

rehype plugin to use tree-sitter to highlight code in pre code blocks

rehype-tree-sitter rehype plugin to use tree-sitter to highlight code in precode blocks Contents What is this? When should I use this? Install Use

5 Jul 25, 2023

A simple, fast, and easy to use Solidity test generator based on the Branching Tree Technique.

bulloak A simple, fast, and easy to use Solidity test generator based on the Branching Tree Technique. Installing cargo install bulloak Usage Basic Us

38 Aug 7, 2023

Comments

Add support for rayon-based sorting

Adds optional rayon feature and par_* variants of sort & build methods that employ Rayon to perform sorting in parallel.

Even on an 4-core machine, this provides almost 2x improvement for large lists:

Benchmarking construct/kd_tree (f64)/4: Warming up for 3.0000 s
Warning: Unable to complete 100 samples in 5.0s. You may wish to increase target time to 7.8s, enable flat sampling, or reduce sample count to 50.
construct/kd_tree (f64)/4                                                                             
                        time:   [1.5249 ms 1.5296 ms 1.5350 ms]
Found 7 outliers among 100 measurements (7.00%)
  1 (1.00%) low mild
  3 (3.00%) high mild
  3 (3.00%) high severe
Benchmarking construct/kd_tree (i32)/4: Warming up for 3.0000 s
Warning: Unable to complete 100 samples in 5.0s. You may wish to increase target time to 6.1s, enable flat sampling, or reduce sample count to 60.
construct/kd_tree (i32)/4                                                                             
                        time:   [1.2005 ms 1.2036 ms 1.2069 ms]
Found 3 outliers among 100 measurements (3.00%)
  1 (1.00%) low mild
  1 (1.00%) high mild
  1 (1.00%) high severe
construct/kd_tree + rayon (f64)/4                                                                            
                        time:   [823.38 us 829.92 us 837.19 us]
Found 5 outliers among 100 measurements (5.00%)
  4 (4.00%) high mild
  1 (1.00%) high severe
construct/kd_tree + rayon (i32)/4                                                                            
                        time:   [666.80 us 673.07 us 679.72 us]
Found 8 outliers among 100 measurements (8.00%)
  3 (3.00%) high mild
  5 (5.00%) high severe

opened by RReverser 3

Serde Support (behind a feature flag)

Hi, thanks for this great library. I successfully built my tree, and I would like to save it to disk using bincode. However, I'm getting this error

the trait bound `KdSliceN<EmbeddedBook, UInt<UInt<UTerm, B1>, B0>>: Serialize` is not satisfied
the following other types implement trait `Serialize`:
  &'a T
  &'a mut T
  ()
  (T0, T1)
  (T0, T1, T2)
  (T0, T1, T2, T3)
  (T0, T1, T2, T3, T4)
  (T0, T1, T2, T3, T4, T5)
and 183 others

let me know if this feature is out of scope for this library or not

this is my code for reference

use serde::{Deserialize, Serialize};
use serde_big_array::BigArray;
use kd_tree::{KdPoint, KdSlice2};

#[derive(Debug, Serialize, Deserialize)]
pub struct EmbeddedBook {
    pub title: String,

    pub author: String,

    pub summary: String,

    #[serde(with = "BigArray")]
    pub embeddings: [f32; 384],
}

impl KdPoint for EmbeddedBook {
    type Scalar = f32;
    type Dim = typenum::U2; // 2 dimensional tree.
    fn at(&self, k: usize) -> f32 {
        self.embeddings[k]
    }
}

#[derive(Serialize, Deserialize)]
struct BookSearcher<'a> {
    tree: &'a KdSlice2<EmbeddedBook>,
}

impl<'a> BookSearcher<'a> {
    fn new(embeddedbooks: &'a mut Vec<EmbeddedBook>) -> Self {
        let kdtree = kd_tree::KdSlice2::sort_by(embeddedbooks, |item1, item2, k| {
            item1.embeddings[k]
                .partial_cmp(&item2.embeddings[k])
                .unwrap()
        });
        BookSearcher { tree: kdtree }
    }
}

opened by sachaarbonel 2

nalgebra feature not working as advertised in README.md
When following your official docs, using the nalgebra feature does not work as advertised. The following error is given:

the trait bound `OPoint<i32, Const<3_usize>>: KdPoint` is not satisfied the trait `KdPoint` is not implemented for `OPoint<i32, Const<3_usize>>`

So I'm not sure if the feature is broken or the docs are wrong. I made an MRE to demonstrate (but it is basically the docs copied and pasted into a project).

Minimal reproducible example

https://github.com/DriesCruyskens/test-kd-tree
opened by DriesCruyskens 2
Make testing use numbers with a step
Instead of using completely random floating numbers, use stepped ones as this helps catch some interesting edge cases when those numbers start overlap.

E.g. with this change getting right away:

---- tests::test_nearests stdout ---- thread 'tests::test_nearests' panicked at 'assertion failed: `(left == right)` left: `8`, right: `5`', src\tests.rs:35:9 note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

(fixes are not included in this PR)
opened by RReverser 1