usearch 2.21.4 → 2.25.0

package/numkong/rust/curved.rs

@@ -0,0 +1,391 @@
+//! Curved metric spaces: Bilinear forms and Mahalanobis distance.
+//!
+//! This module provides:
+//!
+//! - [`Bilinear`]: Bilinear form aᵀ × C × b with a metric tensor
+//! - [`Mahalanobis`]: Mahalanobis distance √((a−b)ᵀ × C × (a−b))
+
+use crate::types::{bf16, bf16c, f16, f16c, f32c, f64c, StorageElement};
+
+#[link(name = "numkong")]
+extern "C" {
+    fn nk_bilinear_f64(a: *const f64, b: *const f64, c: *const f64, n: usize, result: *mut f64);
+    fn nk_bilinear_f32(a: *const f32, b: *const f32, c: *const f32, n: usize, result: *mut f64);
+    fn nk_bilinear_f16(a: *const u16, b: *const u16, c: *const u16, n: usize, result: *mut f32);
+    fn nk_bilinear_bf16(a: *const u16, b: *const u16, c: *const u16, n: usize, result: *mut f32);
+    fn nk_bilinear_f64c(a: *const f64, b: *const f64, c: *const f64, n: usize, results: *mut f64);
+    fn nk_bilinear_f32c(a: *const f32, b: *const f32, c: *const f32, n: usize, results: *mut f64);
+    fn nk_bilinear_f16c(a: *const u16, b: *const u16, c: *const u16, n: usize, results: *mut f32);
+    fn nk_bilinear_bf16c(a: *const u16, b: *const u16, c: *const u16, n: usize, results: *mut f32);
+
+    // Mahalanobis distance
+    fn nk_mahalanobis_f64(a: *const f64, b: *const f64, c: *const f64, n: usize, result: *mut f64);
+    fn nk_mahalanobis_f32(a: *const f32, b: *const f32, c: *const f32, n: usize, result: *mut f64);
+    fn nk_mahalanobis_f16(a: *const u16, b: *const u16, c: *const u16, n: usize, result: *mut f32);
+    fn nk_mahalanobis_bf16(a: *const u16, b: *const u16, c: *const u16, n: usize, result: *mut f32);
+}
+
+/// Bilinear form computation: aᵀ × C × b where C is a metric tensor.
+///
+/// Computes the bilinear form of two vectors `a` and `b` with respect to
+/// a symmetric matrix `C` (given in row-major order as a flat slice of length n²).
+pub trait Bilinear: StorageElement {
+    /// Output type for results. `f32` widens to `f64`; f16/bf16 use f32.
+    type Output;
+
+    /// Computes the bilinear form aᵀ × C × b.
+    ///
+    /// # Arguments
+    /// * `a` - First vector of length n
+    /// * `b` - Second vector of length n
+    /// * `c` - Metric tensor (n×n matrix in row-major order, length n²)
+    ///
+    /// # Returns
+    /// `Some(result)` if inputs are valid, `None` if lengths are incompatible.
+    fn bilinear(a: &[Self], b: &[Self], c: &[Self]) -> Option<Self::Output>;
+}
+
+impl Bilinear for f64 {
+    type Output = f64;
+
+    fn bilinear(a: &[Self], b: &[Self], c: &[Self]) -> Option<Self::Output> {
+        let point_count = a.len();
+        if point_count == 0 || b.len() != point_count || c.len() != point_count * point_count {
+            return None;
+        }
+        let mut result: f64 = 0.0;
+        unsafe {
+            nk_bilinear_f64(a.as_ptr(), b.as_ptr(), c.as_ptr(), point_count, &mut result);
+        }
+        Some(result)
+    }
+}
+
+impl Bilinear for f32 {
+    type Output = f64;
+
+    fn bilinear(a: &[Self], b: &[Self], c: &[Self]) -> Option<Self::Output> {
+        let point_count = a.len();
+        if point_count == 0 || b.len() != point_count || c.len() != point_count * point_count {
+            return None;
+        }
+        let mut result: f64 = 0.0;
+        unsafe {
+            nk_bilinear_f32(a.as_ptr(), b.as_ptr(), c.as_ptr(), point_count, &mut result);
+        }
+        Some(result)
+    }
+}
+
+impl Bilinear for f16 {
+    type Output = f32;
+
+    fn bilinear(a: &[Self], b: &[Self], c: &[Self]) -> Option<Self::Output> {
+        let point_count = a.len();
+        if point_count == 0 || b.len() != point_count || c.len() != point_count * point_count {
+            return None;
+        }
+        let mut result: f32 = 0.0;
+        unsafe {
+            nk_bilinear_f16(
+                a.as_ptr() as *const u16,
+                b.as_ptr() as *const u16,
+                c.as_ptr() as *const u16,
+                point_count,
+                &mut result,
+            );
+        }
+        Some(result)
+    }
+}
+
+impl Bilinear for bf16 {
+    type Output = f32;
+
+    fn bilinear(a: &[Self], b: &[Self], c: &[Self]) -> Option<Self::Output> {
+        let point_count = a.len();
+        if point_count == 0 || b.len() != point_count || c.len() != point_count * point_count {
+            return None;
+        }
+        let mut result: f32 = 0.0;
+        unsafe {
+            nk_bilinear_bf16(
+                a.as_ptr() as *const u16,
+                b.as_ptr() as *const u16,
+                c.as_ptr() as *const u16,
+                point_count,
+                &mut result,
+            );
+        }
+        Some(result)
+    }
+}
+
+impl Bilinear for f64c {
+    type Output = f64c;
+
+    fn bilinear(a: &[Self], b: &[Self], c: &[Self]) -> Option<Self::Output> {
+        let point_count = a.len();
+        if point_count == 0 || b.len() != point_count || c.len() != point_count * point_count {
+            return None;
+        }
+        let mut result = [0.0f64; 2];
+        unsafe {
+            nk_bilinear_f64c(
+                a.as_ptr() as *const f64,
+                b.as_ptr() as *const f64,
+                c.as_ptr() as *const f64,
+                point_count,
+                result.as_mut_ptr(),
+            );
+        }
+        Some(f64c {
+            re: result[0],
+            im: result[1],
+        })
+    }
+}
+
+impl Bilinear for f32c {
+    type Output = f64c;
+
+    fn bilinear(a: &[Self], b: &[Self], c: &[Self]) -> Option<Self::Output> {
+        let point_count = a.len();
+        if point_count == 0 || b.len() != point_count || c.len() != point_count * point_count {
+            return None;
+        }
+        let mut result = [0.0f64; 2];
+        unsafe {
+            nk_bilinear_f32c(
+                a.as_ptr() as *const f32,
+                b.as_ptr() as *const f32,
+                c.as_ptr() as *const f32,
+                point_count,
+                result.as_mut_ptr(),
+            );
+        }
+        Some(f64c {
+            re: result[0],
+            im: result[1],
+        })
+    }
+}
+
+impl Bilinear for f16c {
+    type Output = f32c;
+
+    fn bilinear(a: &[Self], b: &[Self], c: &[Self]) -> Option<Self::Output> {
+        let point_count = a.len();
+        if point_count == 0 || b.len() != point_count || c.len() != point_count * point_count {
+            return None;
+        }
+        let mut result = [0.0f32; 2];
+        unsafe {
+            nk_bilinear_f16c(
+                a.as_ptr() as *const u16,
+                b.as_ptr() as *const u16,
+                c.as_ptr() as *const u16,
+                point_count,
+                result.as_mut_ptr(),
+            );
+        }
+        Some(f32c {
+            re: result[0],
+            im: result[1],
+        })
+    }
+}
+
+impl Bilinear for bf16c {
+    type Output = f32c;
+
+    fn bilinear(a: &[Self], b: &[Self], c: &[Self]) -> Option<Self::Output> {
+        let point_count = a.len();
+        if point_count == 0 || b.len() != point_count || c.len() != point_count * point_count {
+            return None;
+        }
+        let mut result = [0.0f32; 2];
+        unsafe {
+            nk_bilinear_bf16c(
+                a.as_ptr() as *const u16,
+                b.as_ptr() as *const u16,
+                c.as_ptr() as *const u16,
+                point_count,
+                result.as_mut_ptr(),
+            );
+        }
+        Some(f32c {
+            re: result[0],
+            im: result[1],
+        })
+    }
+}
+
+/// Mahalanobis distance: √((a−b)ᵀ × C × (a−b)).
+///
+/// Computes the Mahalanobis distance between two vectors `a` and `b` with respect
+/// to an inverse covariance matrix `C` (given in row-major order as a flat slice of length n²).
+pub trait Mahalanobis: StorageElement {
+    /// Output type for results. `f32` widens to `f64`; f16/bf16 use f32.
+    type Output;
+
+    /// Computes the Mahalanobis distance √((a−b)ᵀ × C × (a−b)).
+    ///
+    /// # Arguments
+    /// * `a` - First vector of length n
+    /// * `b` - Second vector of length n
+    /// * `c` - Inverse covariance matrix (n×n matrix in row-major order, length n²)
+    ///
+    /// # Returns
+    /// `Some(result)` if inputs are valid, `None` if lengths are incompatible.
+    fn mahalanobis(a: &[Self], b: &[Self], c: &[Self]) -> Option<Self::Output>;
+}
+
+impl Mahalanobis for f64 {
+    type Output = f64;
+
+    fn mahalanobis(a: &[Self], b: &[Self], c: &[Self]) -> Option<Self::Output> {
+        let point_count = a.len();
+        if point_count == 0 || b.len() != point_count || c.len() != point_count * point_count {
+            return None;
+        }
+        let mut result: f64 = 0.0;
+        unsafe {
+            nk_mahalanobis_f64(a.as_ptr(), b.as_ptr(), c.as_ptr(), point_count, &mut result);
+        }
+        Some(result)
+    }
+}
+
+impl Mahalanobis for f32 {
+    type Output = f64;
+
+    fn mahalanobis(a: &[Self], b: &[Self], c: &[Self]) -> Option<Self::Output> {
+        let point_count = a.len();
+        if point_count == 0 || b.len() != point_count || c.len() != point_count * point_count {
+            return None;
+        }
+        let mut result: f64 = 0.0;
+        unsafe {
+            nk_mahalanobis_f32(a.as_ptr(), b.as_ptr(), c.as_ptr(), point_count, &mut result);
+        }
+        Some(result)
+    }
+}
+
+impl Mahalanobis for f16 {
+    type Output = f32;
+
+    fn mahalanobis(a: &[Self], b: &[Self], c: &[Self]) -> Option<Self::Output> {
+        let point_count = a.len();
+        if point_count == 0 || b.len() != point_count || c.len() != point_count * point_count {
+            return None;
+        }
+        let mut result: f32 = 0.0;
+        unsafe {
+            nk_mahalanobis_f16(
+                a.as_ptr() as *const u16,
+                b.as_ptr() as *const u16,
+                c.as_ptr() as *const u16,
+                point_count,
+                &mut result,
+            );
+        }
+        Some(result)
+    }
+}
+
+impl Mahalanobis for bf16 {
+    type Output = f32;
+
+    fn mahalanobis(a: &[Self], b: &[Self], c: &[Self]) -> Option<Self::Output> {
+        let point_count = a.len();
+        if point_count == 0 || b.len() != point_count || c.len() != point_count * point_count {
+            return None;
+        }
+        let mut result: f32 = 0.0;
+        unsafe {
+            nk_mahalanobis_bf16(
+                a.as_ptr() as *const u16,
+                b.as_ptr() as *const u16,
+                c.as_ptr() as *const u16,
+                point_count,
+                &mut result,
+            );
+        }
+        Some(result)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::types::{assert_close, bf16, f16, FloatLike, NumberLike, TestableType};
+    /// Build an identity matrix of size point_count × point_count.
+    pub(crate) fn make_identity<Scalar: FloatLike>(point_count: usize) -> Vec<Scalar> {
+        let mut v = vec![Scalar::zero(); point_count * point_count];
+        for i in 0..point_count {
+            v[i * point_count + i] = Scalar::one();
+        }
+        v
+    }
+
+    fn check_bilinear<Scalar>(first_values: &[f32], second_values: &[f32], expected: f64)
+    where
+        Scalar: FloatLike + TestableType + Bilinear,
+        Scalar::Output: FloatLike,
+    {
+        let first: Vec<Scalar> = first_values.iter().map(|&v| Scalar::from_f32(v)).collect();
+        let second: Vec<Scalar> = second_values.iter().map(|&v| Scalar::from_f32(v)).collect();
+        let identity = make_identity::<Scalar>(first.len());
+        let result = Scalar::bilinear(&first, &second, &identity).unwrap();
+        assert_close(
+            result.to_f64(),
+            expected,
+            Scalar::atol(),
+            Scalar::rtol(),
+            &format!("bilinear<{}>", core::any::type_name::<Scalar>()),
+        );
+    }
+
+    #[test]
+    fn bilinear() {
+        // first=[1,2,3], second=[4,5,6], identity matrix → dot = 32
+        let first_values = &[1.0, 2.0, 3.0];
+        let second_values = &[4.0, 5.0, 6.0];
+        check_bilinear::<f64>(first_values, second_values, 32.0);
+        check_bilinear::<f32>(first_values, second_values, 32.0);
+        check_bilinear::<f16>(first_values, second_values, 32.0);
+        check_bilinear::<bf16>(first_values, second_values, 32.0);
+    }
+
+    fn check_mahalanobis<Scalar>(first_values: &[f32], second_values: &[f32], expected: f64)
+    where
+        Scalar: FloatLike + TestableType + Mahalanobis,
+        Scalar::Output: FloatLike,
+    {
+        let first: Vec<Scalar> = first_values.iter().map(|&v| Scalar::from_f32(v)).collect();
+        let second: Vec<Scalar> = second_values.iter().map(|&v| Scalar::from_f32(v)).collect();
+        let identity = make_identity::<Scalar>(first.len());
+        let result = Scalar::mahalanobis(&first, &second, &identity).unwrap();
+        assert_close(
+            result.to_f64(),
+            expected,
+            Scalar::atol(),
+            Scalar::rtol(),
+            &format!("mahalanobis<{}>", core::any::type_name::<Scalar>()),
+        );
+    }
+
+    #[test]
+    fn mahalanobis() {
+        // first=[1,2,3], second=[4,5,6], identity → sqrt(27)
+        let first_values = &[1.0, 2.0, 3.0];
+        let second_values = &[4.0, 5.0, 6.0];
+        let expected = (27.0_f64).sqrt();
+        check_mahalanobis::<f64>(first_values, second_values, expected);
+        check_mahalanobis::<f32>(first_values, second_values, expected);
+        check_mahalanobis::<f16>(first_values, second_values, expected);
+        check_mahalanobis::<bf16>(first_values, second_values, expected);
+    }
+}