ecp_core/
metric.rs

1//! Distance metrics and the distance computations built on them.
2
3use ndarray::{Array1, Array2, Axis};
4
5use crate::error::{EcpError, Result};
6
7/// How the distance between two vectors is measured. `as_str` and `FromStr`
8/// convert it to and from the name stored in an index's `info/metric`.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum Metric {
11    /// Euclidean distance; lower is closer.
12    L2,
13    /// Inner product; higher is closer.
14    IP,
15}
16
17impl Metric {
18    /// Returns the metric's name, as stored in `info/metric`.
19    pub fn as_str(self) -> &'static str {
20        match self {
21            Metric::L2 => "L2",
22            Metric::IP => "IP",
23        }
24    }
25}
26
27impl std::str::FromStr for Metric {
28    type Err = String;
29
30    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
31        match s {
32            "L2" => Ok(Metric::L2),
33            "IP" => Ok(Metric::IP),
34            other => Err(format!("unknown metric `{other}` (use \"L2\" or \"IP\")")),
35        }
36    }
37}
38
39/// Scores every row of `embeddings` against the query `q`, as L2 distances or,
40/// for IP, inner products. Set `is_normalized` only if every row is
41/// unit-length, which lets L2 skip computing the rows' norms.
42/// L2 distance for unit-length rows = `sqrt(1 − 2·e·q + ‖q‖²)`.
43pub fn calculate_distances(
44    embeddings: &Array2<f32>,
45    q: &Array1<f32>,
46    metric: &Metric,
47    is_normalized: bool,
48) -> Result<Array1<f32>> {
49    if embeddings.ncols() != q.len() {
50        return Err(EcpError::InvalidInput(format!(
51            "query has dimension {}, the index expects {}",
52            q.len(),
53            embeddings.ncols()
54        )));
55    }
56
57    Ok(match metric {
58        Metric::IP => embeddings.dot(q),
59        // Terms of similar size can cancel to a small negative where the true
60        // distance is zero, and the root would turn that into NaN
61        Metric::L2 if is_normalized => {
62            let dots = embeddings.dot(q);
63            let q_norm_sq = q.dot(q);
64            (1.0 - 2.0 * dots + q_norm_sq).mapv(|v| v.max(0.0).sqrt())
65        }
66        Metric::L2 => {
67            let q_2d = q.clone().insert_axis(Axis(0));
68            let neg_dist_sq = negative_squared_distances(embeddings, &q_2d);
69            neg_dist_sq.column(0).mapv(|v| (-v).max(0.0).sqrt())
70        }
71    })
72}
73
74/// Computes the squared L2 distance between every row of `a` and every row of
75/// `b`, negated so a higher score is closer, as an `(a.nrows(), b.nrows())` matrix.
76/// Squared distance = `‖a‖² − 2·a·b + ‖b‖²`.
77pub fn negative_squared_distances(a: &Array2<f32>, b: &Array2<f32>) -> Array2<f32> {
78    let a_norms_sq = a.map_axis(Axis(1), |row| row.dot(&row));
79    let b_norms_sq = b.map_axis(Axis(1), |row| row.dot(&row));
80    let cross = a.dot(&b.t());
81
82    let mut neg_dist_sq = 2.0 * cross;
83    neg_dist_sq -= &a_norms_sq.insert_axis(Axis(1));
84    neg_dist_sq -= &b_norms_sq.insert_axis(Axis(0));
85    neg_dist_sq
86}
87
88#[cfg(test)]
89#[path = "utests/metric.rs"]
90mod tests;
ecp_core/metric.rs

ecp_core/
metric.rs
