97 lines
2.9 KiB
Rust
97 lines
2.9 KiB
Rust
#[cfg(feature = "ndarray_15")]
|
|
use crate::ndarray_15_extra::Pow;
|
|
use ndarray::{ArrayBase, Ix1};
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
|
|
pub enum CosineSimilarityError {
|
|
#[error(
|
|
"Invalid vectors: Vectors must have the same length for similarity calculation. LHS: {lhs}, RHS: {rhs}"
|
|
)]
|
|
InvalidVectors { lhs: usize, rhs: usize },
|
|
}
|
|
pub trait CosineSimilarity<T, Rhs = Self> {
|
|
/// Computes the cosine similarity between two vectors.
|
|
///
|
|
/// A `Result` containing the cosine similarity as a `f64`, or an error if the vectors are invalid.
|
|
fn cosine_similarity(&self, rhs: Rhs) -> Result<T, CosineSimilarityError>;
|
|
}
|
|
|
|
impl<S1, S2, T> CosineSimilarity<T, ArrayBase<S2, Ix1>> for ArrayBase<S1, Ix1>
|
|
where
|
|
S1: ndarray::Data<Elem = T>,
|
|
S2: ndarray::Data<Elem = T>,
|
|
T: num::traits::Float + 'static,
|
|
{
|
|
fn cosine_similarity(&self, rhs: ArrayBase<S2, Ix1>) -> Result<T, CosineSimilarityError> {
|
|
if self.len() != rhs.len() {
|
|
return Err(CosineSimilarityError::InvalidVectors {
|
|
lhs: self.len(),
|
|
rhs: rhs.len(),
|
|
});
|
|
}
|
|
debug_assert!(
|
|
self.iter().all(|&x| x.is_finite()),
|
|
"LHS vector contains non-finite values"
|
|
);
|
|
debug_assert!(
|
|
rhs.iter().all(|&x| x.is_finite()),
|
|
"RHS vector contains non-finite values"
|
|
);
|
|
let numerator = self.dot(&rhs);
|
|
let denominator = self.powi(2).sum().sqrt() * rhs.powi(2).sum().sqrt();
|
|
Ok(numerator / denominator)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod cosine_tests {
|
|
use super::*;
|
|
use ndarray::*;
|
|
|
|
#[test]
|
|
fn test_same_vectors() {
|
|
let a = array![1.0, 2.0, 3.0];
|
|
let b = array![1.0, 2.0, 3.0];
|
|
assert_eq!(a.cosine_similarity(b).unwrap(), 1.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_orthogonal_vectors() {
|
|
let a = array![1.0, 0.0, 0.0];
|
|
let b = array![0.0, 1.0, 0.0];
|
|
assert_eq!(a.cosine_similarity(b).unwrap(), 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_opposite_vectors() {
|
|
let a = array![1.0, 2.0, 3.0];
|
|
let b = array![-1.0, -2.0, -3.0];
|
|
assert_eq!(a.cosine_similarity(b).unwrap(), -1.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_invalid_vectors() {
|
|
let a = array![1.0, 2.0];
|
|
let b = array![1.0, 2.0, 3.0];
|
|
assert!(matches!(
|
|
a.cosine_similarity(b),
|
|
Err(CosineSimilarityError::InvalidVectors { lhs: 2, rhs: 3 })
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn test_zero_vector() {
|
|
let a = array![0.0, 0.0, 0.0];
|
|
let b = array![1.0, 2.0, 3.0];
|
|
let similarity = a.cosine_similarity(b);
|
|
assert!(similarity.is_ok_and(|item: f64| item.is_nan()));
|
|
}
|
|
|
|
#[test]
|
|
fn test_different_ndarray_types() {
|
|
let a = array![1.0, 2.0, 3.0];
|
|
let b = array![1.0, 2.0, 3.0];
|
|
assert_eq!(a.cosine_similarity(b.view()).unwrap(), 1.0);
|
|
}
|
|
}
|