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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196
#![deny(missing_docs, missing_debug_implementations)]
//! A simple library implementing an immutable, flat representation of a [k-d tree](https://en.wikipedia.org/wiki/K-d_tree)
//!
//! The library supports arbitrary spatial queries via the [`Query`] trait and nearest neighbour search.
//! Its implementation is simple as the objects in the index are fixed after construction.
//! This also enables a flat and thereby cache-friendly memory layout.
//!
//! The library provides optional integration with [rayon] for parallel construction and queries and [serde] for (de-)serialization of the trees.
//!
//! # Example
//!
//! ```
//! use std::ops::ControlFlow;
//!
//! use sif_kdtree::{KdTree, Object, WithinDistance};
//!
//! struct Something(usize, [f64; 2]);
//!
//! impl Object for Something {
//! type Point = [f64; 2];
//!
//! fn position(&self) -> &Self::Point {
//! &self.1
//! }
//! }
//!
//! let index = KdTree::new(
//! [
//! Something(0, [-0.4, -3.3]),
//! Something(1, [-4.5, -1.8]),
//! Something(2, [0.7, 2.0]),
//! Something(3, [1.7, 1.5]),
//! Something(4, [-1.3, 2.3]),
//! Something(5, [2.2, 1.0]),
//! Something(6, [-3.7, 3.8]),
//! Something(7, [-3.2, -0.1]),
//! Something(8, [1.4, 2.7]),
//! Something(9, [3.1, -0.0]),
//! Something(10, [4.3, 0.8]),
//! Something(11, [3.9, -3.3]),
//! Something(12, [0.4, -3.2]),
//! ]
//! .into(),
//! );
//!
//! let mut close_by = Vec::new();
//!
//! index.look_up(&WithinDistance::new([0., 0.], 3.), |thing| {
//! close_by.push(thing.0);
//!
//! ControlFlow::Continue(())
//! });
//!
//! assert_eq!(close_by, [2, 4, 5, 3]);
//!
//! let closest = index.nearest(&[0., 0.]).unwrap().0;
//!
//! assert_eq!(closest, 2);
//! ```
mod look_up;
mod nearest;
mod sort;
pub use look_up::{Query, WithinBoundingBox, WithinDistance};
use std::ops::Deref;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
/// Defines a [finite-dimensional][Self::DIM] real space in terms of [coordinate values][Self::coord] along a chosen set of axes
pub trait Point {
/// The dimension of the underlying real space
const DIM: usize;
/// Access the coordinate value of the point along the given `axis`
fn coord(&self, axis: usize) -> f64;
/// Return the squared distance between `self` and `other`.
///
/// This is called during nearest neighbour search and hence only the relation between two distance values is required so that computing square roots can be avoided.
fn distance_2(&self, other: &Self) -> f64;
}
/// `N`-dimensional real space using [Euclidean distance](https://en.wikipedia.org/wiki/Euclidean_distance)
impl<const N: usize> Point for [f64; N] {
const DIM: usize = N;
fn coord(&self, axis: usize) -> f64 {
self[axis]
}
fn distance_2(&self, other: &Self) -> f64 {
(0..N).map(|axis| (self[axis] - other[axis]).powi(2)).sum()
}
}
/// Defines the objects which can be organized in a [`KdTree`] by positioning them in a real space defined via the [`Point`] trait
pub trait Object {
/// The [`Point`] implementation used to represent the [position][`Self::position`] of these objects.
type Point: Point;
/// Return the position associated with this object.
///
/// Note that calling this method is assumed to be cheap, returning a reference to a point stored in the interior of the object.
fn position(&self) -> &Self::Point;
}
/// An immutable, flat representation of a [k-d tree](https://en.wikipedia.org/wiki/K-d_tree)
///
/// Accelerates spatial queries and nearest neighbour search by sorting the objects according to the coordinate values of their positions.
///
/// Note that this tree dereferences to and deserializes as a slice of objects.
/// Modifying object positions through interior mutability or deserializing a modified sequence is safe but will lead to incorrect results.
#[derive(Debug, Default)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "serde", serde(transparent))]
pub struct KdTree<O> {
objects: Box<[O]>,
}
impl<O> Deref for KdTree<O> {
type Target = [O];
fn deref(&self) -> &Self::Target {
&self.objects
}
}
fn split<O>(objects: &[O]) -> (&[O], &O, &[O]) {
assert!(!objects.is_empty());
let mid = objects.len() / 2;
unsafe {
(
objects.get_unchecked(..mid),
objects.get_unchecked(mid),
objects.get_unchecked(mid + 1..),
)
}
}
fn contains<P>(aabb: &(P, P), position: &P) -> bool
where
P: Point,
{
(0..P::DIM).all(|axis| {
aabb.0.coord(axis) <= position.coord(axis) && position.coord(axis) <= aabb.1.coord(axis)
})
}
#[cfg(test)]
mod tests {
use super::*;
use std::cmp::Ordering;
use proptest::{collection::vec, strategy::Strategy};
pub fn random_points(len: usize) -> impl Strategy<Value = Vec<[f64; 2]>> {
(vec(0.0..=1.0, len), vec(0.0..=1.0, len))
.prop_map(|(x, y)| x.into_iter().zip(y).map(|(x, y)| [x, y]).collect())
}
#[derive(Debug, PartialEq)]
pub struct RandomObject(pub [f64; 2]);
impl Eq for RandomObject {}
impl PartialOrd for RandomObject {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for RandomObject {
fn cmp(&self, other: &Self) -> Ordering {
self.0.partial_cmp(&other.0).unwrap()
}
}
impl Object for RandomObject {
type Point = [f64; 2];
fn position(&self) -> &Self::Point {
&self.0
}
}
pub fn random_objects(len: usize) -> impl Strategy<Value = Box<[RandomObject]>> {
random_points(len).prop_map(|points| points.into_iter().map(RandomObject).collect())
}
}