ic-rig 0.2.0

//! Vector similarity and distance metrics.
//!
//! Implemented on [`Embedding`] so you can compare vectors without a
//! separate library.

use super::Embedding;

/// Which metric to use when scoring a query against a candidate embedding.
///
/// Pass this to any search function that accepts a `DistanceMetric`.
/// For similarity metrics (higher = more similar) sort descending;
/// for distance metrics (lower = closer) sort ascending.
///
/// # Example
///
/// ```rust,no_run
/// use irig::embeddings::DistanceMetric;
///
/// let metric = DistanceMetric::Cosine { normalized: false };
/// let score  = metric.score(&query_embedding, &candidate_embedding);
/// ```
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum DistanceMetric {
    /// Cosine similarity in `[-1, 1]` (higher = more similar).
    ///
    /// Set `normalized = true` if both vectors are already unit-length
    /// (e.g. `text-embedding-3-small` with `dimensions` set) — avoids the
    /// magnitude computation and returns the raw dot product.
    Cosine { normalized: bool },

    /// Angular distance in `[0, 1]` (lower = closer in direction).
    ///
    /// Set `normalized = true` for the same reason as [`DistanceMetric::Cosine`].
    Angular { normalized: bool },

    /// Euclidean (L2) distance (lower = closer).
    Euclidean,

    /// Manhattan (L1) distance (lower = closer).
    Manhattan,

    /// Chebyshev (L∞) distance — maximum per-dimension difference (lower = closer).
    Chebyshev,

    /// Raw dot product (higher = more similar for unit vectors).
    DotProduct,
}

impl DistanceMetric {
    /// Score `a` against `b` using this metric.
    pub fn score(&self, a: &Embedding, b: &Embedding) -> f64 {
        match self {
            Self::Cosine { normalized }  => a.cosine_similarity(b, *normalized),
            Self::Angular { normalized } => a.angular_distance(b, *normalized),
            Self::Euclidean              => a.euclidean_distance(b),
            Self::Manhattan              => a.manhattan_distance(b),
            Self::Chebyshev              => a.chebyshev_distance(b),
            Self::DotProduct             => a.dot_product(b),
        }
    }

    /// Returns `true` for similarity metrics (sort descending for best-first),
    /// `false` for distance metrics (sort ascending for best-first).
    pub fn higher_is_better(&self) -> bool {
        matches!(self, Self::Cosine { .. } | Self::DotProduct)
    }
}

/// Similarity and distance metrics for embedding vectors.
pub trait VectorDistance {
    /// Dot product of two vectors.
    fn dot_product(&self, other: &Self) -> f64;

    /// Cosine similarity in `[-1, 1]`.
    ///
    /// Pass `normalized = true` if both vectors are already unit-length
    /// (e.g. models like `text-embedding-3-small` with `dimensions` set) —
    /// this skips the magnitude computation and just returns the dot product.
    fn cosine_similarity(&self, other: &Self, normalized: bool) -> f64;

    /// Angular distance in `[0, 1]` (0 = identical direction).
    fn angular_distance(&self, other: &Self, normalized: bool) -> f64;

    /// Euclidean (L2) distance.
    fn euclidean_distance(&self, other: &Self) -> f64;

    /// Manhattan (L1) distance.
    fn manhattan_distance(&self, other: &Self) -> f64;

    /// Chebyshev (L∞) distance — maximum absolute difference across dimensions.
    fn chebyshev_distance(&self, other: &Self) -> f64;
}

impl VectorDistance for Embedding {
    fn dot_product(&self, other: &Self) -> f64 {
        self.vec.iter().zip(&other.vec).map(|(a, b)| a * b).sum()
    }

    fn cosine_similarity(&self, other: &Self, normalized: bool) -> f64 {
        let dot = self.dot_product(other);
        if normalized {
            dot
        } else {
            let mag_a: f64 = self.vec.iter().map(|x| x * x).sum::<f64>().sqrt();
            let mag_b: f64 = other.vec.iter().map(|x| x * x).sum::<f64>().sqrt();
            dot / (mag_a * mag_b)
        }
    }

    fn angular_distance(&self, other: &Self, normalized: bool) -> f64 {
        self.cosine_similarity(other, normalized).acos() / std::f64::consts::PI
    }

    fn euclidean_distance(&self, other: &Self) -> f64 {
        self.vec
            .iter()
            .zip(&other.vec)
            .map(|(a, b)| (a - b).powi(2))
            .sum::<f64>()
            .sqrt()
    }

    fn manhattan_distance(&self, other: &Self) -> f64 {
        self.vec.iter().zip(&other.vec).map(|(a, b)| (a - b).abs()).sum()
    }

    fn chebyshev_distance(&self, other: &Self) -> f64 {
        self.vec
            .iter()
            .zip(&other.vec)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0_f64, f64::max)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn pair() -> (Embedding, Embedding) {
        (
            Embedding { document: "a".into(), vec: vec![1.0, 2.0, 3.0] },
            Embedding { document: "b".into(), vec: vec![1.0, 5.0, 7.0] },
        )
    }

    #[test]
    fn dot_product() {
        let (a, b) = pair();
        assert_eq!(a.dot_product(&b), 32.0);
    }

    #[test]
    fn cosine_similarity() {
        let (a, b) = pair();
        assert!((a.cosine_similarity(&b, false) - 0.9875414397573881).abs() < 1e-10);
    }

    #[test]
    fn euclidean_distance() {
        let (a, b) = pair();
        assert_eq!(a.euclidean_distance(&b), 5.0);
    }

    #[test]
    fn manhattan_distance() {
        let (a, b) = pair();
        assert_eq!(a.manhattan_distance(&b), 7.0);
    }

    #[test]
    fn chebyshev_distance() {
        let (a, b) = pair();
        assert_eq!(a.chebyshev_distance(&b), 4.0);
    }
}