added projected eigen distance function

Author: suhasjs

diskann-quantization/src/multi_vector/distance/fallback.rs

@@ -9,6 +9,7 @@ use diskann_vector::distance::InnerProduct;
 use diskann_vector::{DistanceFunctionMut, PureDistanceFunction};
 
 use super::max_sim::{Chamfer, MaxSim};
+use super::projected_eigen::ProjectedEigen;
 use crate::multi_vector::{MatRef, MaxSimError, Repr, Standard};
 
 /////////////////
@@ -100,6 +101,50 @@ impl FallbackKernel {
             f(i, min_dist);
         }
     }
+
+    /// Core kernel for computing per-query-vector projected-eigen scores.
+    ///
+    /// For each `query` vector, sums the negated squared inner product
+    /// against every document vector, then calls `f(index, score)` with the
+    /// result. If there are no vectors in the `doc`, the kernel returns
+    /// immediately.
+    ///
+    /// The callback can be used to aggregate scores as needed - as is the
+    /// case with [`ProjectedEigen`].
+    ///
+    /// # Arguments
+    ///
+    /// * `query` - The query multi-vector (wrapped as [`QueryMatRef`])
+    /// * `doc` - The document multi-vector
+    /// * `f` - Callback invoked with `(query_index, score)` for each query vector
+    #[inline]
+    pub(crate) fn projected_eigen_kernel<F, T: Copy>(
+        query: QueryMatRef<'_, Standard<T>>,
+        doc: MatRef<'_, Standard<T>>,
+        mut f: F,
+    ) where
+        F: FnMut(usize, f32),
+        InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
+    {
+        // Early exit if no doc vectors - callback should never be invoked
+        if doc.num_vectors() == 0 {
+            return;
+        }
+
+        for (i, q_vec) in query.rows().enumerate() {
+            let mut sum = 0.0f32;
+
+            for d_vec in doc.rows() {
+                // `InnerProduct::evaluate` returns the negated inner product;
+                // squaring discards the sign, so negate the squared value to
+                // obtain `-IP(q, d)²`.
+                let ip = InnerProduct::evaluate(q_vec, d_vec);
+                sum += -(ip * ip);
+            }
+
+            f(i, sum);
+        }
+    }
 }
 
 ////////////
@@ -159,6 +204,27 @@ where
     }
 }
 
+/////////////////////
+// ProjectedEigen //
+/////////////////////
+
+impl<T: Copy> PureDistanceFunction<QueryMatRef<'_, Standard<T>>, MatRef<'_, Standard<T>>, f32>
+    for ProjectedEigen
+where
+    InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
+{
+    #[inline(always)]
+    fn evaluate(query: QueryMatRef<'_, Standard<T>>, doc: MatRef<'_, Standard<T>>) -> f32 {
+        let mut sum = 0.0f32;
+
+        FallbackKernel::projected_eigen_kernel(query, doc, |_i, score| {
+            sum += score;
+        });
+
+        sum
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -185,6 +251,17 @@ mod tests {
             .fold(f32::MAX, f32::min)
     }
 
+    /// Naive implementation of projected-eigen for a single query vector
+    /// against all doc vectors: `\sum_{j} -IP(q, d_{j})^2`.
+    fn naive_projected_eigen_single(query_vec: &[f32], doc: &MatRef<'_, Standard<f32>>) -> f32 {
+        doc.rows()
+            .map(|d_vec| {
+                let ip: f32 = query_vec.iter().zip(d_vec.iter()).map(|(a, b)| a * b).sum();
+                -(ip * ip)
+            })
+            .sum()
+    }
+
     /// Generate deterministic test data.
     fn make_test_data(len: usize, ceil: usize, shift: usize) -> Vec<f32> {
         (0..len).map(|v| ((v + shift) % ceil) as f32).collect()
@@ -286,6 +363,22 @@ mod tests {
                     nd,
                     dim
                 );
+
+                // Test ProjectedEigen
+                let projected = ProjectedEigen::evaluate(query, doc);
+                let expected_projected: f32 = query
+                    .rows()
+                    .map(|q_vec| naive_projected_eigen_single(q_vec, &doc))
+                    .sum();
+
+                assert!(
+                    (projected - expected_projected).abs()
+                        < 1e-6 * expected_projected.abs().max(1.0),
+                    "ProjectedEigen mismatch for ({},{},{})",
+                    nq,
+                    nd,
+                    dim
+                );
             }
         }
 
@@ -303,5 +396,15 @@ mod tests {
 
             assert_eq!(result, 0.0);
         }
+
+        #[test]
+        fn projected_eigen_with_zero_docs_returns_zero() {
+            let query = make_query(&[1.0, 0.0, 0.0, 1.0], 2, 2);
+            let doc = make_doc(&[], 0, 2);
+
+            // No document vectors means no pairs contribute, so the sum is 0.
+            let result = ProjectedEigen::evaluate(query, doc);
+            assert_eq!(result, 0.0);
+        }
     }
 }

diskann-quantization/src/multi_vector/distance/mod.rs

@@ -47,9 +47,11 @@ mod isa;
 mod kernel;
 mod kernels;
 mod max_sim;
+mod projected_eigen;
 
 pub use factory::{MaxSimElement, build_max_sim};
 pub use fallback::QueryMatRef;
 pub use isa::{MaxSimIsa, NotSupported};
 pub use kernel::{BoxErase, Erase, MaxSimKernel};
 pub use max_sim::{Chamfer, MaxSim, MaxSimError};
+pub use projected_eigen::ProjectedEigen;

diskann-quantization/src/multi_vector/distance/projected_eigen.rs

@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+//! Projected-eigen distance type for multi-vector representations.
+
+/////////////////////
+// ProjectedEigen //
+/////////////////////
+
+/// Projected-eigen distance for multi-vector similarity.
+///
+/// Computes the negated sum of squared inner products over *all*
+/// query/document vector pairs:
+///
+/// ```text
+/// ProjectedEigen(Q, D) = \sum_{i} \sum_{j} -IP(q_i, d_j)²
+/// ```
+///
+/// Unlike [`Chamfer`](super::Chamfer), which keeps only the best-matching
+/// document vector per query vector, this accumulates a contribution from
+/// every pair, so the score reflects the full query–document interaction. As
+/// with the other multi-vector distances, lower is better.
+///
+/// Implements [`PureDistanceFunction`](diskann_vector::PureDistanceFunction)
+/// for matrix view types.
+#[derive(Debug, Clone, Copy)]
+pub struct ProjectedEigen;

diskann-quantization/src/multi_vector/mod.rs

@@ -58,7 +58,7 @@ pub(crate) mod matrix;
 pub use block_transposed::{BlockTransposed, BlockTransposedMut, BlockTransposedRef};
 pub use distance::{
     BoxErase, Chamfer, Erase, MaxSim, MaxSimElement, MaxSimError, MaxSimIsa, MaxSimKernel,
-    NotSupported, QueryMatRef, build_max_sim,
+    NotSupported, ProjectedEigen, QueryMatRef, build_max_sim,
 };
 pub use matrix::{
     Defaulted, LayoutError, Mat, MatMut, MatRef, NewCloned, NewMut, NewOwned, NewRef, Overflow,