rig_core/embeddings/
distance.rs

1//! Similarity and distance of embedding vectors.
2//!
3//! [`Embedding`](crate::embeddings::Embedding) reductions use fixed chunks and
4//! left-to-right summation for reproducible ordering.
5//!
6//! ```
7//! use rig_core::embeddings::{Embedding, distance::VectorDistance};
8//!
9//! let vector = Embedding { document: String::new(), vec: vec![1.0, 0.0] };
10//! assert_eq!(vector.cosine_similarity(&vector, true), 1.0);
11//! assert_eq!(vector.euclidean_distance(&vector), 0.0);
12//! ```
13
14/// Similarity and distance metrics of embedding vectors.
15/// Supply equal-length vectors; the embedding implementation pairs only their
16/// shared prefix. Unnormalized cosine requires nonzero magnitudes.
17pub trait VectorDistance {
18    /// The dot product.
19    fn dot_product(&self, other: &Self) -> f64;
20
21    /// Cosine similarity, from -1 to 1.
22    /// If `normalized` is true, the vectors are taken as unit length and the
23    /// dot product is returned.
24    fn cosine_similarity(&self, other: &Self, normalized: bool) -> f64;
25
26    /// The angle between the vectors as a fraction of pi, from 0 (same
27    /// direction) to 1 (opposite). `normalized` is as for
28    /// [`Self::cosine_similarity`].
29    fn angular_distance(&self, other: &Self, normalized: bool) -> f64 {
30        // Roundoff can push a valid cosine beyond the domain of acos.
31        self.cosine_similarity(other, normalized)
32            .clamp(-1.0, 1.0)
33            .acos()
34            / std::f64::consts::PI
35    }
36
37    /// The straight-line (L2) distance.
38    fn euclidean_distance(&self, other: &Self) -> f64;
39
40    /// The sum of the coordinates' absolute differences (L1).
41    fn manhattan_distance(&self, other: &Self) -> f64;
42
43    /// The largest absolute difference of a coordinate (L-infinity).
44    fn chebyshev_distance(&self, other: &Self) -> f64;
45}
46
47/// Reduction chunk size. Preserve left-to-right summation within and across
48/// chunks to keep floating-point results reproducible.
49const CHUNK: usize = 256;
50
51/// `term` of each pair of `a` and `b` summed in fixed chunks, left to right
52/// within and across them.
53fn pair_sum(a: &[f64], b: &[f64], term: impl Fn(f64, f64) -> f64) -> f64 {
54    a.chunks(CHUNK)
55        .zip(b.chunks(CHUNK))
56        .map(|(a, b)| a.iter().zip(b).map(|(x, y)| term(*x, *y)).sum::<f64>())
57        .sum()
58}
59
60impl VectorDistance for crate::embeddings::Embedding {
61    fn dot_product(&self, other: &Self) -> f64 {
62        dot(&self.vec, &other.vec)
63    }
64
65    fn cosine_similarity(&self, other: &Self, normalized: bool) -> f64 {
66        let dot_product = dot(&self.vec, &other.vec);
67        if normalized {
68            return dot_product;
69        }
70        dot_product / (magnitude(&self.vec) * magnitude(&other.vec))
71    }
72
73    fn euclidean_distance(&self, other: &Self) -> f64 {
74        pair_sum(&self.vec, &other.vec, |x, y| (x - y).powi(2)).sqrt()
75    }
76
77    fn manhattan_distance(&self, other: &Self) -> f64 {
78        pair_sum(&self.vec, &other.vec, |x, y| (x - y).abs())
79    }
80
81    fn chebyshev_distance(&self, other: &Self) -> f64 {
82        self.vec
83            .iter()
84            .zip(&other.vec)
85            .map(|(x, y)| (x - y).abs())
86            .fold(0.0, f64::max)
87    }
88}
89
90/// The dot product of `a` and `b`, summed as [`VectorDistance`] sums it.
91pub(crate) fn dot(a: &[f64], b: &[f64]) -> f64 {
92    pair_sum(a, b, |x, y| x * y)
93}
94
95/// The length of `vec`, summed as [`VectorDistance`] sums it.
96pub(crate) fn magnitude(vec: &[f64]) -> f64 {
97    pair_sum(vec, vec, |x, _| x.powi(2)).sqrt()
98}
99
100#[cfg(test)]
101mod tests;
rig_core/embeddings/distance.rs

rig_core/embeddings/
distance.rs
