use g_math::fixed_point::FixedPoint;
use crate::constants;
use crate::hyperbolic_geometry::HyperbolicPoint;
pub trait Metric<P> {
fn distance(&self, a: &P, b: &P) -> FixedPoint;
fn has_proxy(&self) -> bool {
false
}
fn proxy(&self, a: &P, b: &P) -> FixedPoint {
self.distance(a, b)
}
fn prune_left(&self, s_query: FixedPoint, median: FixedPoint, s_worst: FixedPoint) -> bool {
s_query - s_worst > median
}
fn prune_right(&self, s_query: FixedPoint, median: FixedPoint, s_worst: FixedPoint) -> bool {
s_query + s_worst < median
}
fn left_first(&self, s_query: FixedPoint, median: FixedPoint) -> bool {
s_query < median
}
}
pub struct EuclideanMetric;
impl Metric<Vec<FixedPoint>> for EuclideanMetric {
fn distance(&self, a: &Vec<FixedPoint>, b: &Vec<FixedPoint>) -> FixedPoint {
g_math::fixed_point::imperative::fused::euclidean_distance(a, b)
}
fn has_proxy(&self) -> bool {
true
}
fn proxy(&self, a: &Vec<FixedPoint>, b: &Vec<FixedPoint>) -> FixedPoint {
g_math::fixed_point::imperative::fused::euclidean_distance_squared(a, b)
}
fn prune_left(&self, s_query: FixedPoint, median: FixedPoint, s_worst: FixedPoint) -> bool {
if s_query <= median {
return false; }
if sq_product_may_saturate(s_query, median) {
return false;
}
let two = FixedPoint::from_int(2);
let x_ub = sqrt_upper_bound(s_query * median);
s_query + median - s_worst > two * x_ub
}
fn prune_right(&self, s_query: FixedPoint, median: FixedPoint, s_worst: FixedPoint) -> bool {
if median <= s_query {
return false;
}
if sq_product_may_saturate(s_query, median) {
return false;
}
let two = FixedPoint::from_int(2);
let x_ub = sqrt_upper_bound(s_query * median);
s_query + median - s_worst > two * x_ub
}
}
#[inline]
fn sq_product_may_saturate(a: FixedPoint, b: FixedPoint) -> bool {
const LIMIT: i128 = 1 << 95;
a.raw() >= LIMIT || b.raw() >= LIMIT
}
#[derive(Clone, Debug)]
pub struct CachedNormPoint {
pub point: HyperbolicPoint,
pub norm_sq: FixedPoint,
}
impl CachedNormPoint {
pub fn new(point: HyperbolicPoint) -> Self {
let norm_sq = point.coords().length_squared();
Self { point, norm_sq }
}
}
pub struct HyperbolicMetric;
fn near_boundary_sq() -> FixedPoint {
constants::near_boundary() * constants::near_boundary()
}
pub(crate) fn hyperbolic_ratio_sq(
a: &HyperbolicPoint,
a_norm_sq: FixedPoint,
b: &HyperbolicPoint,
b_norm_sq: FixedPoint,
) -> FixedPoint {
let zero = FixedPoint::from_int(0);
let one = FixedPoint::from_int(1);
let two = FixedPoint::from_int(2);
let cap = near_boundary_sq();
let eps_sq = constants::small_epsilon() * constants::small_epsilon();
if a_norm_sq < eps_sq {
return if b_norm_sq > cap { cap } else { b_norm_sq };
}
if b_norm_sq < eps_sq {
return if a_norm_sq > cap { cap } else { a_norm_sq };
}
let dot = a.coords().dot(b.coords());
let mut dist_sq = a_norm_sq + b_norm_sq - two * dot;
if dist_sq < zero {
dist_sq = zero; }
let den_sq = one - two * dot + a_norm_sq * b_norm_sq;
if den_sq < constants::min_safe_denominator() {
return cap;
}
let r_sq = dist_sq / den_sq;
if r_sq > cap {
cap
} else {
r_sq
}
}
pub(crate) fn sq_ratio_separation_exceeds(
s_hi: FixedPoint,
s_lo: FixedPoint,
s_tau: FixedPoint,
) -> bool {
let one = FixedPoint::from_int(1);
let two = FixedPoint::from_int(2);
let x_ub = sqrt_upper_bound(s_hi * s_lo);
let one_minus_x = one - x_ub;
s_hi + s_lo - two * x_ub > s_tau * (one_minus_x * one_minus_x)
}
impl Metric<CachedNormPoint> for HyperbolicMetric {
fn distance(&self, a: &CachedNormPoint, b: &CachedNormPoint) -> FixedPoint {
a.point.hyperbolic_distance(&b.point)
}
fn has_proxy(&self) -> bool {
true
}
fn proxy(&self, a: &CachedNormPoint, b: &CachedNormPoint) -> FixedPoint {
hyperbolic_ratio_sq(&a.point, a.norm_sq, &b.point, b.norm_sq)
}
fn prune_left(&self, s_query: FixedPoint, median: FixedPoint, s_worst: FixedPoint) -> bool {
s_query > median && sq_ratio_separation_exceeds(s_query, median, s_worst)
}
fn prune_right(&self, s_query: FixedPoint, median: FixedPoint, s_worst: FixedPoint) -> bool {
median > s_query && sq_ratio_separation_exceeds(median, s_query, s_worst)
}
fn left_first(&self, s_query: FixedPoint, median: FixedPoint) -> bool {
s_query < median
}
}
fn sqrt_upper_bound(p: FixedPoint) -> FixedPoint {
let raw = p.raw();
if raw <= 0 {
return FixedPoint::from_int(0);
}
let b = 128 - (raw as u128).leading_zeros();
let seed = FixedPoint::from_raw(1i128 << (b.div_ceil(2) + 32));
let x1 = FixedPoint::from_raw((seed.raw() + (p / seed).raw()) >> 1);
FixedPoint::from_raw((x1.raw() + (p / x1).raw()) >> 1)
}
struct TreeNode<P> {
unique_id: String,
point: P,
median: FixedPoint,
left: Option<Box<TreeNode<P>>>,
right: Option<Box<TreeNode<P>>>,
}
pub struct MetricVpTree<P> {
root: Option<Box<TreeNode<P>>>,
len: usize,
}
impl<P> MetricVpTree<P> {
pub fn build<M: Metric<P>>(mut entries: Vec<(String, P)>, metric: &M) -> Self {
entries.sort_by(|a, b| a.0.cmp(&b.0));
let len = entries.len();
Self {
root: Self::build_node(entries, metric),
len,
}
}
pub fn len(&self) -> usize {
self.len
}
pub fn is_empty(&self) -> bool {
self.len == 0
}
fn score<M: Metric<P>>(metric: &M, a: &P, b: &P) -> FixedPoint {
if metric.has_proxy() {
metric.proxy(a, b)
} else {
metric.distance(a, b)
}
}
fn build_node<M: Metric<P>>(
mut entries: Vec<(String, P)>,
metric: &M,
) -> Option<Box<TreeNode<P>>> {
if entries.is_empty() {
return None;
}
if entries.len() == 1 {
let (unique_id, point) = entries.remove(0);
return Some(Box::new(TreeNode {
unique_id,
point,
median: FixedPoint::from_int(0),
left: None,
right: None,
}));
}
let (vp_id, vp_point) = entries.swap_remove(0);
let mut with_scores: Vec<((String, P), FixedPoint)> = entries
.into_iter()
.map(|e| {
let s = Self::score(metric, &vp_point, &e.1);
(e, s)
})
.collect();
with_scores.sort_by(|a, b| a.1.cmp(&b.1).then_with(|| (a.0).0.cmp(&(b.0).0)));
let median = with_scores[with_scores.len() / 2].1;
let (left_vec, right_vec): (Vec<_>, Vec<_>) =
with_scores.into_iter().partition(|(_, s)| *s < median);
let left = Self::build_node(left_vec.into_iter().map(|(e, _)| e).collect(), metric);
let right = Self::build_node(right_vec.into_iter().map(|(e, _)| e).collect(), metric);
Some(Box::new(TreeNode {
unique_id: vp_id,
point: vp_point,
median,
left,
right,
}))
}
pub fn knn<M: Metric<P>>(&self, query: &P, k: usize, metric: &M) -> Vec<(String, FixedPoint)> {
if k == 0 {
return Vec::new();
}
let mut candidates: Vec<(FixedPoint, &TreeNode<P>)> = Vec::with_capacity(k + 1);
if let Some(ref root) = self.root {
Self::search_knn(root, query, k, metric, &mut candidates);
}
let mut results: Vec<(String, FixedPoint)> = candidates
.into_iter()
.map(|(score, node)| {
let d = if metric.has_proxy() {
metric.distance(query, &node.point)
} else {
score
};
(node.unique_id.clone(), d)
})
.collect();
results.sort_by(|a, b| a.1.cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
results
}
fn worst_score<P2>(candidates: &[(FixedPoint, P2)], k: usize) -> Option<FixedPoint> {
if candidates.len() < k {
return None;
}
Some(candidates.last().unwrap().0)
}
fn search_knn<'a, M: Metric<P>>(
node: &'a TreeNode<P>,
query: &P,
k: usize,
metric: &M,
candidates: &mut Vec<(FixedPoint, &'a TreeNode<P>)>,
) {
let s = Self::score(metric, query, &node.point);
let full = candidates.len() == k;
let admit = !full || {
let worst = candidates.last().unwrap();
(s, node.unique_id.as_str()) < (worst.0, worst.1.unique_id.as_str())
};
if admit {
candidates.push((s, node));
candidates.sort_by(|a, b| {
a.0.cmp(&b.0)
.then_with(|| a.1.unique_id.cmp(&b.1.unique_id))
});
if candidates.len() > k {
candidates.truncate(k);
}
}
let worst = Self::worst_score(candidates, k);
let descend_left = |worst: Option<FixedPoint>| {
worst.is_none_or(|w| !metric.prune_left(s, node.median, w))
};
let descend_right = |worst: Option<FixedPoint>| {
worst.is_none_or(|w| !metric.prune_right(s, node.median, w))
};
if metric.left_first(s, node.median) {
if let Some(ref left) = node.left {
if descend_left(worst) {
Self::search_knn(left, query, k, metric, candidates);
}
}
let worst = Self::worst_score(candidates, k);
if let Some(ref right) = node.right {
if descend_right(worst) {
Self::search_knn(right, query, k, metric, candidates);
}
}
} else {
if let Some(ref right) = node.right {
if descend_right(worst) {
Self::search_knn(right, query, k, metric, candidates);
}
}
let worst = Self::worst_score(candidates, k);
if let Some(ref left) = node.left {
if descend_left(worst) {
Self::search_knn(left, query, k, metric, candidates);
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn fp(v: i32) -> FixedPoint {
FixedPoint::from_int(v)
}
fn point(coords: &[i32]) -> Vec<FixedPoint> {
coords.iter().map(|&c| fp(c)).collect()
}
fn brute_knn(
entries: &[(String, Vec<FixedPoint>)],
query: &Vec<FixedPoint>,
k: usize,
) -> Vec<(String, FixedPoint)> {
let mut all: Vec<(FixedPoint, String)> = entries
.iter()
.map(|(uid, p)| (EuclideanMetric.distance(query, p), uid.clone()))
.collect();
all.sort_by(|a, b| a.0.cmp(&b.0).then_with(|| a.1.cmp(&b.1)));
all.truncate(k);
all.into_iter().map(|(d, uid)| (uid, d)).collect()
}
fn lcg_entries(n: usize, dims: usize, seed: u64) -> Vec<(String, Vec<FixedPoint>)> {
let mut state = seed;
let mut next = || {
state = state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
((state >> 33) % 2001) as i32 - 1000 };
(0..n)
.map(|i| {
let coords: Vec<FixedPoint> = (0..dims).map(|_| fp(next())).collect();
(format!("node_{:05}", i), coords)
})
.collect()
}
#[test]
fn knn_matches_brute_force() {
for &(n, dims, seed) in &[(50usize, 2usize, 7u64), (200, 4, 42), (500, 8, 1234)] {
let entries = lcg_entries(n, dims, seed);
let tree = MetricVpTree::build(entries.clone(), &EuclideanMetric);
let mut qstate = seed ^ 0xdead_beef;
let mut next = || {
qstate = qstate
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((qstate >> 33) % 2001) as i32 - 1000
};
for _ in 0..20 {
let query: Vec<FixedPoint> = (0..dims).map(|_| fp(next())).collect();
for &k in &[1usize, 5, 17, n + 10] {
let got = tree.knn(&query, k, &EuclideanMetric);
let want = brute_knn(&entries, &query, k);
assert_eq!(got, want, "n={} dims={} k={}", n, dims, k);
}
}
}
}
#[test]
fn ties_break_by_uid() {
let p = point(&[3, 4]);
let mut entries: Vec<(String, Vec<FixedPoint>)> = ["d", "b", "a", "c"]
.iter()
.map(|s| (s.to_string(), p.clone()))
.collect();
let tree1 = MetricVpTree::build(entries.clone(), &EuclideanMetric);
entries.reverse();
let tree2 = MetricVpTree::build(entries, &EuclideanMetric);
let query = point(&[0, 0]);
let r1 = tree1.knn(&query, 2, &EuclideanMetric);
let r2 = tree2.knn(&query, 2, &EuclideanMetric);
assert_eq!(r1, r2);
assert_eq!(r1[0].0, "a");
assert_eq!(r1[1].0, "b");
}
fn hpoint(x: f32, y: f32) -> CachedNormPoint {
CachedNormPoint::new(HyperbolicPoint::from_f32_slice(&[x, y]))
}
#[test]
fn hyperbolic_proxy_bounds_and_monotonicity() {
let m = HyperbolicMetric;
let cases = [
(hpoint(0.3, 0.4), hpoint(-0.2, 0.5)),
(hpoint(0.0, 0.0), hpoint(0.6, -0.3)),
(hpoint(0.55, 0.0), hpoint(0.0, 0.0)),
(hpoint(0.95, 0.0), hpoint(-0.95, 0.0)),
(hpoint(0.98, 0.01), hpoint(0.97, 0.02)),
(hpoint(0.1, 0.1), hpoint(0.1, 0.1)),
(hpoint(0.001, 0.0), hpoint(0.0, 0.001)),
];
let mut scored: Vec<(FixedPoint, FixedPoint)> = Vec::new();
let two = FixedPoint::from_int(2);
let one = FixedPoint::from_int(1);
for (a, b) in &cases {
let exact = m.distance(a, b);
let s = m.proxy(a, b);
assert!(two * s <= exact, "lower bound violated: s={:?} d={:?}", s, exact);
let denom = one - s;
let d_sq_ub = FixedPoint::from_int(4) * s / (denom * denom);
assert!(
exact * exact <= d_sq_ub + constants::epsilon(),
"upper bound violated: d²={:?} ub={:?}",
exact * exact, d_sq_ub
);
scored.push((s, exact));
}
let mut by_proxy = scored.clone();
by_proxy.sort_by(|a, b| a.0.cmp(&b.0));
let mut by_exact = scored;
by_exact.sort_by(|a, b| a.1.cmp(&b.1));
let d_order: Vec<_> = by_proxy.iter().map(|(_, d)| *d).collect();
let d_expected: Vec<_> = by_exact.iter().map(|(_, d)| *d).collect();
assert_eq!(d_order, d_expected, "proxy ordering diverged from distance ordering");
}
#[test]
fn hyperbolic_knn_matches_brute_force() {
let m = HyperbolicMetric;
let mut state = 7u64;
let mut next = || {
state = state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
(((state >> 33) % 1800) as f32 / 1000.0) - 0.9 };
let entries: Vec<(String, CachedNormPoint)> = (0..300)
.map(|i| {
let (x, y) = (next() * 0.7, next() * 0.7); (format!("p{:03}", i), hpoint(x, y))
})
.collect();
let tree = MetricVpTree::build(entries.clone(), &m);
for qi in [0usize, 111, 222] {
let query = entries[qi].1.clone();
for k in [1usize, 5, 20] {
let got = tree.knn(&query, k, &m);
let mut want: Vec<(FixedPoint, String)> = entries
.iter()
.map(|(uid, p)| (m.distance(&query, p), uid.clone()))
.collect();
want.sort_by(|a, b| a.0.cmp(&b.0).then_with(|| a.1.cmp(&b.1)));
want.truncate(k);
let want: Vec<(String, FixedPoint)> =
want.into_iter().map(|(d, uid)| (uid, d)).collect();
assert_eq!(got, want, "qi={} k={}", qi, k);
}
}
}
#[test]
fn empty_and_degenerate() {
let tree: MetricVpTree<Vec<FixedPoint>> =
MetricVpTree::build(Vec::new(), &EuclideanMetric);
assert!(tree.is_empty());
assert!(tree.knn(&point(&[0]), 5, &EuclideanMetric).is_empty());
let tree = MetricVpTree::build(vec![("only".to_string(), point(&[1]))], &EuclideanMetric);
assert_eq!(tree.len(), 1);
let r = tree.knn(&point(&[0]), 3, &EuclideanMetric);
assert_eq!(r.len(), 1);
assert_eq!(r[0].0, "only");
assert!(tree.knn(&point(&[0]), 0, &EuclideanMetric).is_empty());
}
}