ic_rig/embeddings/
distance.rs

1//! Vector similarity and distance metrics.
2//!
3//! Implemented on [`Embedding`] so you can compare vectors without a
4//! separate library.
5
6use super::Embedding;
7
8/// Which metric to use when scoring a query against a candidate embedding.
9///
10/// Pass this to any search function that accepts a `DistanceMetric`.
11/// For similarity metrics (higher = more similar) sort descending;
12/// for distance metrics (lower = closer) sort ascending.
13///
14/// # Example
15///
16/// ```rust,no_run
17/// use irig::embeddings::DistanceMetric;
18///
19/// let metric = DistanceMetric::Cosine { normalized: false };
20/// let score  = metric.score(&query_embedding, &candidate_embedding);
21/// ```
22#[derive(Debug, Clone, Copy, PartialEq)]
23pub enum DistanceMetric {
24    /// Cosine similarity in `[-1, 1]` (higher = more similar).
25    ///
26    /// Set `normalized = true` if both vectors are already unit-length
27    /// (e.g. `text-embedding-3-small` with `dimensions` set) — avoids the
28    /// magnitude computation and returns the raw dot product.
29    Cosine { normalized: bool },
30
31    /// Angular distance in `[0, 1]` (lower = closer in direction).
32    ///
33    /// Set `normalized = true` for the same reason as [`DistanceMetric::Cosine`].
34    Angular { normalized: bool },
35
36    /// Euclidean (L2) distance (lower = closer).
37    Euclidean,
38
39    /// Manhattan (L1) distance (lower = closer).
40    Manhattan,
41
42    /// Chebyshev (L∞) distance — maximum per-dimension difference (lower = closer).
43    Chebyshev,
44
45    /// Raw dot product (higher = more similar for unit vectors).
46    DotProduct,
47}
48
49impl DistanceMetric {
50    /// Score `a` against `b` using this metric.
51    pub fn score(&self, a: &Embedding, b: &Embedding) -> f64 {
52        match self {
53            Self::Cosine { normalized }  => a.cosine_similarity(b, *normalized),
54            Self::Angular { normalized } => a.angular_distance(b, *normalized),
55            Self::Euclidean              => a.euclidean_distance(b),
56            Self::Manhattan              => a.manhattan_distance(b),
57            Self::Chebyshev              => a.chebyshev_distance(b),
58            Self::DotProduct             => a.dot_product(b),
59        }
60    }
61
62    /// Returns `true` for similarity metrics (sort descending for best-first),
63    /// `false` for distance metrics (sort ascending for best-first).
64    pub fn higher_is_better(&self) -> bool {
65        matches!(self, Self::Cosine { .. } | Self::DotProduct)
66    }
67}
68
69/// Similarity and distance metrics for embedding vectors.
70pub trait VectorDistance {
71    /// Dot product of two vectors.
72    fn dot_product(&self, other: &Self) -> f64;
73
74    /// Cosine similarity in `[-1, 1]`.
75    ///
76    /// Pass `normalized = true` if both vectors are already unit-length
77    /// (e.g. models like `text-embedding-3-small` with `dimensions` set) —
78    /// this skips the magnitude computation and just returns the dot product.
79    fn cosine_similarity(&self, other: &Self, normalized: bool) -> f64;
80
81    /// Angular distance in `[0, 1]` (0 = identical direction).
82    fn angular_distance(&self, other: &Self, normalized: bool) -> f64;
83
84    /// Euclidean (L2) distance.
85    fn euclidean_distance(&self, other: &Self) -> f64;
86
87    /// Manhattan (L1) distance.
88    fn manhattan_distance(&self, other: &Self) -> f64;
89
90    /// Chebyshev (L∞) distance — maximum absolute difference across dimensions.
91    fn chebyshev_distance(&self, other: &Self) -> f64;
92}
93
94impl VectorDistance for Embedding {
95    fn dot_product(&self, other: &Self) -> f64 {
96        self.vec.iter().zip(&other.vec).map(|(a, b)| a * b).sum()
97    }
98
99    fn cosine_similarity(&self, other: &Self, normalized: bool) -> f64 {
100        let dot = self.dot_product(other);
101        if normalized {
102            dot
103        } else {
104            let mag_a: f64 = self.vec.iter().map(|x| x * x).sum::<f64>().sqrt();
105            let mag_b: f64 = other.vec.iter().map(|x| x * x).sum::<f64>().sqrt();
106            dot / (mag_a * mag_b)
107        }
108    }
109
110    fn angular_distance(&self, other: &Self, normalized: bool) -> f64 {
111        self.cosine_similarity(other, normalized).acos() / std::f64::consts::PI
112    }
113
114    fn euclidean_distance(&self, other: &Self) -> f64 {
115        self.vec
116            .iter()
117            .zip(&other.vec)
118            .map(|(a, b)| (a - b).powi(2))
119            .sum::<f64>()
120            .sqrt()
121    }
122
123    fn manhattan_distance(&self, other: &Self) -> f64 {
124        self.vec.iter().zip(&other.vec).map(|(a, b)| (a - b).abs()).sum()
125    }
126
127    fn chebyshev_distance(&self, other: &Self) -> f64 {
128        self.vec
129            .iter()
130            .zip(&other.vec)
131            .map(|(a, b)| (a - b).abs())
132            .fold(0.0_f64, f64::max)
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139
140    fn pair() -> (Embedding, Embedding) {
141        (
142            Embedding { document: "a".into(), vec: vec![1.0, 2.0, 3.0] },
143            Embedding { document: "b".into(), vec: vec![1.0, 5.0, 7.0] },
144        )
145    }
146
147    #[test]
148    fn dot_product() {
149        let (a, b) = pair();
150        assert_eq!(a.dot_product(&b), 32.0);
151    }
152
153    #[test]
154    fn cosine_similarity() {
155        let (a, b) = pair();
156        assert!((a.cosine_similarity(&b, false) - 0.9875414397573881).abs() < 1e-10);
157    }
158
159    #[test]
160    fn euclidean_distance() {
161        let (a, b) = pair();
162        assert_eq!(a.euclidean_distance(&b), 5.0);
163    }
164
165    #[test]
166    fn manhattan_distance() {
167        let (a, b) = pair();
168        assert_eq!(a.manhattan_distance(&b), 7.0);
169    }
170
171    #[test]
172    fn chebyshev_distance() {
173        let (a, b) = pair();
174        assert_eq!(a.chebyshev_distance(&b), 4.0);
175    }
176}
ic_rig/embeddings/distance.rs

ic_rig/embeddings/
distance.rs
