use oxirs_core::simd::SimdOps;
use std::cmp::Ordering;
#[derive(Debug, Clone)]
pub struct TreeIndexConfig {
pub tree_type: TreeType,
pub max_leaf_size: usize,
pub random_seed: Option<u64>,
pub parallel_construction: bool,
pub distance_metric: DistanceMetric,
}
impl Default for TreeIndexConfig {
fn default() -> Self {
Self {
tree_type: TreeType::BallTree,
max_leaf_size: 16, random_seed: None,
parallel_construction: true,
distance_metric: DistanceMetric::Euclidean,
}
}
}
#[derive(Debug, Clone, Copy)]
pub enum TreeType {
BallTree,
KdTree,
VpTree,
CoverTree,
RandomProjectionTree,
}
#[derive(Debug, Clone, Copy)]
pub enum DistanceMetric {
Euclidean,
Manhattan,
Cosine,
Minkowski(f32),
}
impl DistanceMetric {
pub(crate) fn distance(&self, a: &[f32], b: &[f32]) -> f32 {
match self {
DistanceMetric::Euclidean => f32::euclidean_distance(a, b),
DistanceMetric::Manhattan => f32::manhattan_distance(a, b),
DistanceMetric::Cosine => f32::cosine_distance(a, b),
DistanceMetric::Minkowski(p) => a
.iter()
.zip(b.iter())
.map(|(x, y)| (x - y).abs().powf(*p))
.sum::<f32>()
.powf(1.0 / p),
}
}
}
#[derive(Debug, Clone)]
pub(crate) struct SearchResult {
pub(crate) index: usize,
pub(crate) distance: f32,
}
impl PartialEq for SearchResult {
fn eq(&self, other: &Self) -> bool {
self.distance == other.distance
}
}
impl Eq for SearchResult {}
impl PartialOrd for SearchResult {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for SearchResult {
fn cmp(&self, other: &Self) -> Ordering {
self.partial_cmp(other).unwrap_or(Ordering::Equal)
}
}