mod nalgebra;
mod nearest;
mod nearests;
mod sort;
mod tests;
mod within;
use nearest::*;
use nearests::*;
use sort::*;
use std::cmp::Ordering;
use std::marker::PhantomData;
use typenum::Unsigned;
use within::*;
pub trait KdPoint {
type Scalar: num_traits::NumAssign + Copy + PartialOrd;
type Dim: Unsigned;
fn dim() -> usize {
<Self::Dim as Unsigned>::to_usize()
}
fn at(&self, i: usize) -> Self::Scalar;
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ItemAndDistance<'a, T, Scalar> {
pub item: &'a T,
pub squared_distance: Scalar,
}
#[derive(Debug, PartialEq, Eq)]
pub struct KdSliceN<T, N: Unsigned>(PhantomData<N>, [T]);
pub type KdSlice<T> = KdSliceN<T, <T as KdPoint>::Dim>;
impl<T, N: Unsigned> std::ops::Deref for KdSliceN<T, N> {
type Target = [T];
fn deref(&self) -> &[T] {
&self.1
}
}
impl<T: Clone, N: Unsigned> std::borrow::ToOwned for KdSliceN<T, N> {
type Owned = KdTreeN<T, N>;
fn to_owned(&self) -> Self::Owned {
KdTreeN(PhantomData, self.1.to_vec())
}
}
impl<T, N: Unsigned> KdSliceN<T, N> {
pub fn items(&self) -> &[T] {
&self.1
}
unsafe fn new_unchecked(items: &[T]) -> &Self {
&*(items as *const _ as *const Self)
}
pub fn sort_by<F>(items: &mut [T], compare: F) -> &Self
where
F: Fn(&T, &T, usize) -> Ordering + Copy,
{
kd_sort_by(items, N::to_usize(), compare);
unsafe { Self::new_unchecked(items) }
}
pub fn sort_by_key<Key: Ord, F>(items: &mut [T], kd_key: F) -> &Self
where
F: Fn(&T, usize) -> Key + Copy,
{
Self::sort_by(items, |item1, item2, k| {
kd_key(item1, k).cmp(&kd_key(item2, k))
})
}
pub fn sort_by_ordered_float(points: &mut [T]) -> &Self
where
T: KdPoint<Dim = N>,
T::Scalar: ordered_float::FloatCore,
{
Self::sort_by_key(points, |item, k| ordered_float::OrderedFloat(item.at(k)))
}
pub fn sort(points: &mut [T]) -> &Self
where
T: KdPoint<Dim = N>,
T::Scalar: Ord,
{
Self::sort_by_key(points, |item, k| item.at(k))
}
pub fn nearest_by<Q: KdPoint<Dim = N>>(
&self,
query: &Q,
coord: impl Fn(&T, usize) -> Q::Scalar + Copy,
) -> Option<ItemAndDistance<'_, T, Q::Scalar>> {
if self.is_empty() {
None
} else {
Some(kd_nearest_by(self.items(), query, coord))
}
}
pub fn nearest(
&self,
query: &impl KdPoint<Scalar = T::Scalar, Dim = N>,
) -> Option<ItemAndDistance<'_, T, T::Scalar>>
where
T: KdPoint<Dim = N>,
{
if self.is_empty() {
None
} else {
Some(kd_nearest(self.items(), query))
}
}
pub fn nearests_by<Q: KdPoint<Dim = N>>(
&self,
query: &Q,
num: usize,
coord: impl Fn(&T, usize) -> Q::Scalar + Copy,
) -> Vec<ItemAndDistance<'_, T, Q::Scalar>> {
kd_nearests_by(self.items(), query, num, coord)
}
pub fn nearests(
&self,
query: &impl KdPoint<Scalar = T::Scalar, Dim = N>,
num: usize,
) -> Vec<ItemAndDistance<'_, T, T::Scalar>>
where
T: KdPoint<Dim = N>,
{
kd_nearests(self.items(), query, num)
}
pub fn within_by_cmp(&self, compare: impl Fn(&T, usize) -> Ordering + Copy) -> Vec<&T> {
kd_within_by_cmp(self, N::to_usize(), compare)
}
pub fn within_by<Q: KdPoint<Dim = N>>(
&self,
query: &[Q; 2],
coord: impl Fn(&T, usize) -> Q::Scalar + Copy,
) -> Vec<&T> {
assert!((0..Q::dim()).all(|k| query[0].at(k) <= query[1].at(k)));
self.within_by_cmp(|item, k| {
let a = coord(item, k);
if a < query[0].at(k) {
Ordering::Less
} else if a > query[1].at(k) {
Ordering::Greater
} else {
Ordering::Equal
}
})
}
pub fn within(&self, query: &[impl KdPoint<Scalar = T::Scalar, Dim = N>; 2]) -> Vec<&T>
where
T: KdPoint<Dim = N>,
{
self.within_by(query, |item, k| item.at(k))
}
pub fn within_radius_by<Q: KdPoint<Dim = N>>(
&self,
query: &Q,
radius: Q::Scalar,
coord: impl Fn(&T, usize) -> Q::Scalar + Copy,
) -> Vec<&T> {
let mut results = self.within_by_cmp(|item, k| {
let coord = coord(item, k);
if coord < query.at(k) - radius {
Ordering::Less
} else if coord > query.at(k) + radius {
Ordering::Greater
} else {
Ordering::Equal
}
});
results.retain(|item| {
let mut distance = <Q::Scalar as num_traits::Zero>::zero();
for k in 0..N::to_usize() {
let diff = coord(item, k) - query.at(k);
distance += diff * diff;
}
distance < radius * radius
});
results
}
pub fn within_radius(
&self,
query: &impl KdPoint<Scalar = T::Scalar, Dim = N>,
radius: T::Scalar,
) -> Vec<&T>
where
T: KdPoint<Dim = N>,
{
self.within_radius_by(query, radius, |item, k| item.at(k))
}
}
#[cfg(feature = "rayon")]
impl<T: Send, N: Unsigned> KdSliceN<T, N> {
pub fn par_sort_by<F>(items: &mut [T], compare: F) -> &Self
where
F: Fn(&T, &T, usize) -> Ordering + Copy + Send,
{
kd_par_sort_by(items, N::to_usize(), compare);
unsafe { Self::new_unchecked(items) }
}
pub fn par_sort_by_key<Key: Ord, F>(items: &mut [T], kd_key: F) -> &Self
where
F: Fn(&T, usize) -> Key + Copy + Send,
{
Self::par_sort_by(items, move |item1, item2, k| {
kd_key(item1, k).cmp(&kd_key(item2, k))
})
}
pub fn par_sort_by_ordered_float(points: &mut [T]) -> &Self
where
T: KdPoint<Dim = N>,
T::Scalar: ordered_float::FloatCore,
{
Self::par_sort_by_key(points, |item, k| ordered_float::OrderedFloat(item.at(k)))
}
pub fn par_sort(points: &mut [T]) -> &Self
where
T: KdPoint<Dim = N>,
T::Scalar: Ord,
{
Self::par_sort_by_key(points, |item, k| item.at(k))
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct KdTreeN<T, N: Unsigned>(PhantomData<N>, Vec<T>);
pub type KdTree<T> = KdTreeN<T, <T as KdPoint>::Dim>;
impl<T, N: Unsigned> std::ops::Deref for KdTreeN<T, N> {
type Target = KdSliceN<T, N>;
fn deref(&self) -> &Self::Target {
unsafe { KdSliceN::new_unchecked(&self.1) }
}
}
impl<T, N: Unsigned> AsRef<KdSliceN<T, N>> for KdTreeN<T, N> {
fn as_ref(&self) -> &KdSliceN<T, N> {
self
}
}
impl<T, N: Unsigned> std::borrow::Borrow<KdSliceN<T, N>> for KdTreeN<T, N> {
fn borrow(&self) -> &KdSliceN<T, N> {
self
}
}
impl<T, N: Unsigned> From<KdTreeN<T, N>> for Vec<T> {
fn from(src: KdTreeN<T, N>) -> Self {
src.1
}
}
impl<T, N: Unsigned> KdTreeN<T, N> {
pub fn into_vec(self) -> Vec<T> {
self.1
}
pub fn build_by<F>(mut items: Vec<T>, compare: F) -> Self
where
F: Fn(&T, &T, usize) -> Ordering + Copy,
{
kd_sort_by(&mut items, N::to_usize(), compare);
Self(PhantomData, items)
}
pub fn build_by_key<Key, F>(items: Vec<T>, kd_key: F) -> Self
where
Key: Ord,
F: Fn(&T, usize) -> Key + Copy,
{
Self::build_by(items, |item1, item2, k| {
kd_key(item1, k).cmp(&kd_key(item2, k))
})
}
pub fn build_by_ordered_float(points: Vec<T>) -> Self
where
T: KdPoint<Dim = N>,
T::Scalar: ordered_float::FloatCore,
{
Self::build_by_key(points, |item, k| ordered_float::OrderedFloat(item.at(k)))
}
pub fn build(points: Vec<T>) -> Self
where
T: KdPoint<Dim = N>,
T::Scalar: Ord,
{
Self::build_by_key(points, |item, k| item.at(k))
}
}
#[cfg(feature = "serde")]
mod impl_serde {
use super::{KdTreeN, PhantomData, Unsigned};
impl<T: serde::Serialize, N: Unsigned> serde::Serialize for KdTreeN<T, N> {
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
self.1.serialize(serializer)
}
}
impl<'de, T: serde::Deserialize<'de>, N: Unsigned> serde::Deserialize<'de> for KdTreeN<T, N> {
fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
Vec::<T>::deserialize(deserializer).map(|items| Self(PhantomData, items))
}
}
}
#[cfg(feature = "rayon")]
impl<T: Send, N: Unsigned> KdTreeN<T, N> {
pub fn par_build_by<F>(mut items: Vec<T>, compare: F) -> Self
where
F: Fn(&T, &T, usize) -> Ordering + Copy + Send,
{
kd_par_sort_by(&mut items, N::to_usize(), compare);
Self(PhantomData, items)
}
pub fn par_build_by_key<Key, F>(items: Vec<T>, kd_key: F) -> Self
where
Key: Ord,
F: Fn(&T, usize) -> Key + Copy + Send,
{
Self::par_build_by(items, move |item1, item2, k| {
kd_key(item1, k).cmp(&kd_key(item2, k))
})
}
pub fn par_build_by_ordered_float(points: Vec<T>) -> Self
where
T: KdPoint<Dim = N>,
T::Scalar: ordered_float::FloatCore,
{
Self::par_build_by_key(points, |item, k| ordered_float::OrderedFloat(item.at(k)))
}
pub fn par_build(points: Vec<T>) -> Self
where
T: KdPoint<Dim = N>,
T::Scalar: Ord,
{
Self::par_build_by_key(points, |item, k| item.at(k))
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KdIndexTreeN<'a, T, N: Unsigned> {
source: &'a [T],
kdtree: KdTreeN<usize, N>,
}
pub type KdIndexTree<'a, T> = KdIndexTreeN<'a, T, <T as KdPoint>::Dim>;
impl<'a, T, N: Unsigned> KdIndexTreeN<'a, T, N> {
pub fn source(&self) -> &'a [T] {
self.source
}
pub fn indices(&self) -> &KdSliceN<usize, N> {
&self.kdtree
}
pub fn item(&self, i: usize) -> &'a T {
&self.source[i]
}
pub fn build_by<F>(source: &'a [T], compare: F) -> Self
where
F: Fn(&T, &T, usize) -> Ordering + Copy,
{
Self {
source,
kdtree: KdTreeN::build_by((0..source.len()).collect(), |i1, i2, k| {
compare(&source[*i1], &source[*i2], k)
}),
}
}
pub fn build_by_key<Key, F>(source: &'a [T], kd_key: F) -> Self
where
Key: Ord,
F: Fn(&T, usize) -> Key + Copy,
{
Self::build_by(source, |item1, item2, k| {
kd_key(item1, k).cmp(&kd_key(item2, k))
})
}
pub fn build_by_ordered_float(points: &'a [T]) -> Self
where
T: KdPoint<Dim = N>,
T::Scalar: ordered_float::FloatCore,
{
Self::build_by_key(points, |item, k| ordered_float::OrderedFloat(item.at(k)))
}
pub fn build(points: &'a [T]) -> Self
where
T: KdPoint<Dim = N>,
T::Scalar: Ord,
{
Self::build_by_key(points, |item, k| item.at(k))
}
pub fn nearest_by<Q: KdPoint<Dim = N>>(
&self,
query: &Q,
coord: impl Fn(&T, usize) -> Q::Scalar + Copy,
) -> Option<ItemAndDistance<'_, usize, Q::Scalar>> {
self.kdtree
.nearest_by(query, |&index, k| coord(&self.source[index], k))
}
pub fn nearest(
&self,
query: &impl KdPoint<Scalar = T::Scalar, Dim = N>,
) -> Option<ItemAndDistance<'_, usize, T::Scalar>>
where
T: KdPoint<Dim = N>,
{
self.nearest_by(query, |item, k| item.at(k))
}
pub fn nearests_by<Q: KdPoint<Dim = N>>(
&self,
query: &Q,
num: usize,
coord: impl Fn(&T, usize) -> Q::Scalar + Copy,
) -> Vec<ItemAndDistance<'_, usize, Q::Scalar>> {
self.kdtree
.nearests_by(query, num, |&index, k| coord(&self.source[index], k))
}
pub fn nearests(
&self,
query: &impl KdPoint<Scalar = T::Scalar, Dim = N>,
num: usize,
) -> Vec<ItemAndDistance<'_, usize, T::Scalar>>
where
T: KdPoint<Dim = N>,
{
self.nearests_by(query, num, |item, k| item.at(k))
}
pub fn within_by_cmp(&self, compare: impl Fn(&T, usize) -> Ordering + Copy) -> Vec<&usize> {
self.kdtree
.within_by_cmp(|&index, k| compare(&self.source[index], k))
}
pub fn within_by<Q: KdPoint<Dim = N>>(
&self,
query: &[Q; 2],
coord: impl Fn(&T, usize) -> Q::Scalar + Copy,
) -> Vec<&usize> {
self.kdtree
.within_by(query, |&index, k| coord(&self.source[index], k))
}
pub fn within(&self, query: &[impl KdPoint<Scalar = T::Scalar, Dim = N>; 2]) -> Vec<&usize>
where
T: KdPoint<Dim = N>,
{
self.within_by(query, |item, k| item.at(k))
}
pub fn within_radius_by<Q: KdPoint<Dim = N>>(
&self,
query: &Q,
radius: Q::Scalar,
coord: impl Fn(&T, usize) -> Q::Scalar + Copy,
) -> Vec<&usize> {
self.kdtree
.within_radius_by(query, radius, |&index, k| coord(&self.source[index], k))
}
pub fn within_radius(
&self,
query: &impl KdPoint<Scalar = T::Scalar, Dim = N>,
radius: T::Scalar,
) -> Vec<&usize>
where
T: KdPoint<Dim = N>,
{
self.within_radius_by(query, radius, |item, k| item.at(k))
}
}
#[cfg(feature = "rayon")]
impl<'a, T: Sync, N: Unsigned> KdIndexTreeN<'a, T, N> {
pub fn par_build_by<F>(source: &'a [T], compare: F) -> Self
where
F: Fn(&T, &T, usize) -> Ordering + Copy + Send,
{
Self {
source,
kdtree: KdTreeN::par_build_by((0..source.len()).collect(), move |i1, i2, k| {
compare(&source[*i1], &source[*i2], k)
}),
}
}
pub fn par_build_by_key<Key, F>(source: &'a [T], kd_key: F) -> Self
where
Key: Ord,
F: Fn(&T, usize) -> Key + Copy + Send,
{
Self::par_build_by(source, move |item1, item2, k| {
kd_key(item1, k).cmp(&kd_key(item2, k))
})
}
pub fn par_build_by_ordered_float(points: &'a [T]) -> Self
where
T: KdPoint<Dim = N>,
T::Scalar: ordered_float::FloatCore,
{
Self::par_build_by_key(points, |item, k| ordered_float::OrderedFloat(item.at(k)))
}
pub fn par_build(points: &'a [T]) -> Self
where
T: KdPoint<Dim = N>,
T::Scalar: Ord,
{
Self::par_build_by_key(points, |item, k| item.at(k))
}
}
macro_rules! define_kdtree_aliases {
($($dim:literal),*) => {
$(
paste::paste! {
pub type [<KdSlice $dim>]<T> = KdSliceN<T, typenum::[<U $dim>]>;
pub type [<KdTree $dim>]<T> = KdTreeN<T, typenum::[<U $dim>]>;
pub type [<KdIndexTree $dim>]<'a, T> = KdIndexTreeN<'a, T, typenum::[<U $dim>]>;
}
)*
};
}
define_kdtree_aliases!(1, 2, 3, 4, 5, 6, 7, 8);
macro_rules! impl_kd_points {
($($len:literal),*) => {
$(
paste::paste!{
impl<T: num_traits::NumAssign + Copy + PartialOrd> KdPoint for [T; $len] {
type Scalar = T;
type Dim = typenum::[<U $len>];
fn at(&self, i: usize) -> T { self[i] }
}
}
)*
};
}
impl_kd_points!(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16);
impl<P: KdPoint, T> KdPoint for (P, T) {
type Scalar = P::Scalar;
type Dim = P::Dim;
fn at(&self, k: usize) -> Self::Scalar {
self.0.at(k)
}
}
pub type KdMap<P, T> = KdTree<(P, T)>;
pub type KdMapSlice<P, T> = KdSlice<(P, T)>;