use linfa::Float;
use ndarray::{aview1, ArrayBase, Data, Ix2};
#[cfg(feature = "serde")]
use serde_crate::{Deserialize, Serialize};
use crate::{
distance::Distance, BuildError, NearestNeighbour, NearestNeighbourBox, NearestNeighbourIndex,
NnError, Point,
};
#[derive(Debug)]
pub struct KdTreeIndex<'a, F: Float, D: Distance<F>>(
kdtree::KdTree<F, (Point<'a, F>, usize), &'a [F]>,
D,
);
impl<'a, F: Float, D: Distance<F>> KdTreeIndex<'a, F, D> {
pub fn new<DT: Data<Elem = F>>(
batch: &'a ArrayBase<DT, Ix2>,
leaf_size: usize,
dist_fn: D,
) -> Result<Self, BuildError> {
if leaf_size == 0 {
Err(BuildError::EmptyLeaf)
} else if batch.ncols() == 0 {
Err(BuildError::ZeroDimension)
} else {
let mut tree = kdtree::KdTree::with_capacity(batch.ncols().max(1), leaf_size);
for (i, point) in batch.rows().into_iter().enumerate() {
tree.add(
point.to_slice().expect("views should be contiguous"),
(point, i),
)
.unwrap();
}
Ok(Self(tree, dist_fn))
}
}
}
impl From<kdtree::ErrorKind> for NnError {
fn from(err: kdtree::ErrorKind) -> Self {
match err {
kdtree::ErrorKind::WrongDimension => NnError::WrongDimension,
kdtree::ErrorKind::NonFiniteCoordinate => panic!("infinite value found"),
_ => unreachable!(),
}
}
}
impl<'a, F: Float, D: Distance<F>> NearestNeighbourIndex<F> for KdTreeIndex<'a, F, D> {
fn k_nearest<'b>(
&self,
point: Point<'b, F>,
k: usize,
) -> Result<Vec<(Point<F>, usize)>, NnError> {
Ok(self
.0
.nearest(
point.to_slice().expect("views should be contiguous"),
k,
&|a, b| self.1.rdistance(aview1(a), aview1(b)),
)?
.into_iter()
.map(|(_, (pt, pos))| (pt.reborrow(), *pos))
.collect())
}
fn within_range<'b>(
&self,
point: Point<'b, F>,
range: F,
) -> Result<Vec<(Point<F>, usize)>, NnError> {
let range = self.1.dist_to_rdist(range);
Ok(self
.0
.within(
point.to_slice().expect("views should be contiguous"),
range,
&|a, b| self.1.rdistance(aview1(a), aview1(b)),
)?
.into_iter()
.map(|(_, (pt, pos))| (pt.reborrow(), *pos))
.collect())
}
}
#[derive(Default, Clone, Debug, PartialEq, Eq)]
#[cfg_attr(
feature = "serde",
derive(Serialize, Deserialize),
serde(crate = "serde_crate")
)]
pub struct KdTree;
impl KdTree {
pub fn new() -> Self {
Self
}
}
impl NearestNeighbour for KdTree {
fn from_batch_with_leaf_size<'a, F: Float, DT: Data<Elem = F>, D: 'a + Distance<F>>(
&self,
batch: &'a ArrayBase<DT, Ix2>,
leaf_size: usize,
dist_fn: D,
) -> Result<NearestNeighbourBox<'a, F>, BuildError> {
KdTreeIndex::new(batch, leaf_size, dist_fn).map(|v| Box::new(v) as NearestNeighbourBox<F>)
}
}