#[cfg(test)] mod tests { use std::sync::atomic::AtomicBool; use common::counter::hardware_counter::HardwareCounterCell; use quantization::encoded_storage::TestEncodedStorageBuilder; use quantization::encoded_vectors::{DistanceType, EncodedVectors, VectorParameters}; use quantization::encoded_vectors_tq::{self, EncodedVectorsTQ, ErrorCorrectionMetadata}; use quantization::turboquant::simd::{ CODEBOOK_SCALE_SQ_2BIT, CODEBOOK_SCALE_SQ_4BIT, score_1bit_internal_scalar, score_2bit_internal_scalar, score_2bit_internal_weighted_scalar, score_4bit_internal_scalar, score_4bit_internal_weighted_scalar, }; use quantization::turboquant::{TQBits, TQMode, TQRotation}; use rand::{RngExt, SeedableRng}; use crate::metrics::{dot_similarity, l1_similarity, l2_similarity}; const VECTORS_COUNT: usize = 513; const DIMS: &[usize] = &[16, 64, 65, 128, 384, 512]; const BITS: &[TQBits] = &[TQBits::Bits4, TQBits::Bits2, TQBits::Bits1_5, TQBits::Bits1]; /// Absolute tolerance for an approximate score: an empirical per-bit /// coefficient (≈ 1.8x observed max across VECTORS_COUNT trials) times /// the signal-std of the input data. The mean-error / signal-std ratio /// is shared by the Dot, Cosine, and L2 paths, so the same coefficient /// table is used; only the caller's `signal_std` differs by metric and /// data distribution. fn error(bits: TQBits, signal_std: f32) -> f32 { let coef = match bits { TQBits::Bits1 => 5.1, TQBits::Bits1_5 => 4.0, TQBits::Bits2 => 3.0, TQBits::Bits4 => 0.9, }; coef * signal_std } /// Signal-std of dot products of two independent U[-1, 1]^d vectors. fn dot_signal_std(dim: usize) -> f32 { (dim as f32 / 9.0).sqrt() } /// Signal-std of cosine of two independent unit-norm random vectors in d-dim. fn cosine_signal_std(dim: usize) -> f32 { 1.0 / (dim as f32).sqrt() } /// Signal-std of the L2-squared score for two independent U[-1, 1]^d /// vectors. The score is `‖q‖² + ‖v‖² − 2·`; the norms are stored /// exactly as extras and contribute no quantization noise, so the /// signal variance is dominated by the `−2·` term — i.e. twice /// the dot-product std. fn l2_signal_std(dim: usize) -> f32 { 2.0 * dot_signal_std(dim) } /// Tolerance for L1 scoring. Has its own per-bit coefficient table /// because the L1 path's noise/signal ratio doesn't track the dot/cosine /// pattern: it dequantizes both sides (full inverse rotation + Lloyd-Max /// noise per coord) before summing |a − b|. Per-coord errors enter the /// sum signed (sign(a_i − b_i)·δ_i), so cancellation makes the total /// scale as sqrt(dim). Coefficients are empirically calibrated /// (~1.8x observed max across both symmetric and asymmetric paths); /// the symmetric path dominates at low bits because both vectors carry /// dequantization noise. fn error_l1(dim: usize, bits: TQBits) -> f32 { let per_sqrt_dim = match bits { TQBits::Bits1 => 7.5, TQBits::Bits1_5 => 4.5, TQBits::Bits2 => 3.0, TQBits::Bits4 => 0.7, }; per_sqrt_dim * (dim as f32).sqrt() } /// Per-bits minimum dim for meaningful absolute-error testing. At low dim /// the Hadamard rotation has too few coordinates to Gaussianize well and /// the Lloyd-Max quantization error grows too large relative to the signal /// for the tolerance formulas above to be useful. fn should_test(dim: usize, bits: TQBits) -> bool { let min_dim = match bits { TQBits::Bits1 => 64, TQBits::Bits1_5 => 48, TQBits::Bits2 => 32, TQBits::Bits4 => 8, }; dim >= min_dim } fn l2_norm(v: &[f32]) -> f32 { v.iter().map(|x| x * x).sum::().sqrt() } fn normalize(v: &[f32]) -> Vec { let n = l2_norm(v); v.iter().map(|x| x / n).collect() } fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 { dot_similarity(a, b) / (l2_norm(a) * l2_norm(b)) } fn read_f32(src: &[u8], offset: usize) -> f32 { f32::from_le_bytes(src[offset..offset + size_of::()].try_into().unwrap()) } fn extra_len(distance: DistanceType, mode: TQMode) -> usize { let distance_extra_len = match distance { DistanceType::Dot | DistanceType::Cosine => size_of::(), DistanceType::L2 => 2 * size_of::(), DistanceType::L1 => unreachable!("L1 reference needs inverse rotation"), }; let mode_extra_len = match mode { TQMode::Normal => 0, TQMode::Plus => size_of::(), }; distance_extra_len + mode_extra_len } fn ec_correction(extra: &[u8]) -> f32 { read_f32(extra, extra.len() - size_of::()) } fn tq_plus_weight_scale(ec: &ErrorCorrectionMetadata) -> (Vec, f32) { let d_prime_sq_f32: Vec = ec .scale .iter() .map(|&s| { if s.abs() > f32::EPSILON { (s * s).recip() } else { 0.0 } }) .collect(); let max_d_prime_sq = d_prime_sq_f32.iter().copied().fold(0.0f32, f32::max); const QUANT_CAP: i16 = i16::MAX - 1; let weight_scale = if max_d_prime_sq > f32::EPSILON { f32::from(QUANT_CAP) / max_d_prime_sq } else { 1.0 }; let weights: Vec = d_prime_sq_f32 .iter() .map(|&x| (x * weight_scale).round().clamp(0.0, f32::from(QUANT_CAP)) as i16) .collect(); (weights, weight_scale) } fn score_scalar_reference( v1: &[u8], v2: &[u8], bits: TQBits, distance: DistanceType, mode: TQMode, ec: Option<&ErrorCorrectionMetadata>, ) -> f32 { let extra_len = extra_len(distance, mode); let (data_v1, extra_v1) = v1.split_at(v1.len() - extra_len); let (data_v2, extra_v2) = v2.split_at(v2.len() - extra_len); let raw_dot = match (mode, bits) { (TQMode::Plus, TQBits::Bits2) => { let ec = ec.expect("TQ+ reference requires error correction metadata"); let (weights, weight_scale) = tq_plus_weight_scale(ec); let raw_int = score_2bit_internal_weighted_scalar(data_v1, data_v2, &weights); let xm_a = ec_correction(extra_v1); let xm_b = ec_correction(extra_v2); let mm: f32 = ec.shift.iter().map(|&s| s * s).sum(); raw_int as f32 / (weight_scale * CODEBOOK_SCALE_SQ_2BIT) + xm_a + xm_b - mm } (TQMode::Plus, TQBits::Bits4) => { let ec = ec.expect("TQ+ reference requires error correction metadata"); let (weights, weight_scale) = tq_plus_weight_scale(ec); let raw_int = score_4bit_internal_weighted_scalar(data_v1, data_v2, &weights); let xm_a = ec_correction(extra_v1); let xm_b = ec_correction(extra_v2); let mm: f32 = ec.shift.iter().map(|&s| s * s).sum(); raw_int as f32 / (weight_scale * CODEBOOK_SCALE_SQ_4BIT) + xm_a + xm_b - mm } (_, TQBits::Bits1 | TQBits::Bits1_5) => score_1bit_internal_scalar(data_v1, data_v2), (_, TQBits::Bits2) => score_2bit_internal_scalar(data_v1, data_v2), (_, TQBits::Bits4) => score_4bit_internal_scalar(data_v1, data_v2), }; let v1_scale = read_f32(extra_v1, 0); let v2_scale = read_f32(extra_v2, 0); match distance { DistanceType::Dot | DistanceType::Cosine => raw_dot * v1_scale * v2_scale, DistanceType::L2 => { let l2_a = read_f32(extra_v1, size_of::()); let l2_b = read_f32(extra_v2, size_of::()); l2_a * l2_a + l2_b * l2_b - 2.0 * v1_scale * v2_scale * raw_dot } DistanceType::L1 => unreachable!("L1 reference needs inverse rotation"), } } #[test] fn test_tq_internal_score_matches_reference() { let dim = 128; let vectors_count = 32; let counter = HardwareCounterCell::new(); for &bits in BITS { for &distance in &[DistanceType::Dot, DistanceType::Cosine, DistanceType::L2] { let mut rng = rand::rngs::StdRng::seed_from_u64(42); let vector_data: Vec> = (0..vectors_count) .map(|_| { let vector: Vec = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect(); match distance { DistanceType::Cosine => normalize(&vector), DistanceType::Dot | DistanceType::L1 | DistanceType::L2 => vector, } }) .collect(); let vector_parameters = VectorParameters { dim, deprecated_count: None, distance_type: distance, invert: false, }; for &mode in &[TQMode::Normal, TQMode::Plus] { let quantized_vector_size = encoded_vectors_tq::get_quantized_vector_size( &vector_parameters, bits, mode, ); let encoded = EncodedVectorsTQ::encode( vector_data.iter(), TestEncodedStorageBuilder::new(None, quantized_vector_size), &vector_parameters, vectors_count, bits, mode, TQRotation::Padded, false, 1, None, &AtomicBool::new(false), ) .unwrap(); let ec = encoded.get_metadata().error_correction.as_ref(); for i in 1..vectors_count { let v1 = encoded.get_quantized_vector(0); let v2 = encoded.get_quantized_vector(i as u32); let optimized = encoded.score_internal(0, i as u32, &counter); let reference = score_scalar_reference(&v1, &v2, bits, distance, mode, ec); let tolerance = 1e-5 * reference.abs().max(1.0); assert!( (optimized - reference).abs() <= tolerance, "bits={bits:?}, mode={mode:?}, distance={distance:?}, i={i}: \ optimized={optimized}, reference={reference}, tolerance={tolerance}" ); } } } } } #[rstest::rstest] #[case::normal(TQMode::Normal, 1)] #[case::plus(TQMode::Plus, 1)] // Exercises the threaded TQ+ pre-pass — keeps the parallel coord-chunk // push covered by integration testing (single-threaded paths stay // covered by the cases above). #[case::plus_parallel(TQMode::Plus, 4)] fn test_tq_dot(#[case] mode: TQMode, #[case] num_threads: usize) { for &bits in BITS { for &dim in DIMS { if !should_test(dim, bits) { continue; } let error = error(bits, dot_signal_std(dim)); let mut rng = rand::rngs::StdRng::seed_from_u64(42); let mut vector_data: Vec> = vec![]; for _ in 0..VECTORS_COUNT { vector_data.push((0..dim).map(|_| rng.random_range(-1.0..1.0)).collect()); } let query: Vec = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect(); let vector_parameters = VectorParameters { dim, deprecated_count: None, distance_type: DistanceType::Dot, invert: false, }; let quantized_vector_size = encoded_vectors_tq::get_quantized_vector_size(&vector_parameters, bits, mode); let encoded = EncodedVectorsTQ::encode( vector_data.iter(), TestEncodedStorageBuilder::new(None, quantized_vector_size), &vector_parameters, VECTORS_COUNT, bits, mode, TQRotation::Padded, false, num_threads, None, &AtomicBool::new(false), ) .unwrap(); let query_u8 = encoded.encode_query(&query); let counter = HardwareCounterCell::new(); for (index, vector) in vector_data.iter().enumerate() { let score = encoded.score_point(&query_u8, index as u32, &counter); let original_score = dot_similarity(&query, vector); assert!( (score - original_score).abs() < error, "bits={bits:?}, dim={dim}, index={index}, score={score}, expected={original_score}" ); } } } } #[rstest::rstest] #[case::normal(TQMode::Normal)] #[case::plus(TQMode::Plus)] fn test_tq_cosine(#[case] mode: TQMode) { for &bits in BITS { for &dim in DIMS { if !should_test(dim, bits) { continue; } let error = error(bits, cosine_signal_std(dim)); let mut rng = rand::rngs::StdRng::seed_from_u64(42); let mut vector_data: Vec> = vec![]; for _ in 0..VECTORS_COUNT { let v: Vec = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect(); vector_data.push(normalize(&v)); } let raw_query: Vec = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect(); let query = normalize(&raw_query); let vector_parameters = VectorParameters { dim, deprecated_count: None, distance_type: DistanceType::Cosine, invert: false, }; let quantized_vector_size = encoded_vectors_tq::get_quantized_vector_size(&vector_parameters, bits, mode); let encoded = EncodedVectorsTQ::encode( vector_data.iter(), TestEncodedStorageBuilder::new(None, quantized_vector_size), &vector_parameters, VECTORS_COUNT, bits, mode, TQRotation::Padded, false, 1, None, &AtomicBool::new(false), ) .unwrap(); let query_u8 = encoded.encode_query(&query); let counter = HardwareCounterCell::new(); for (index, vector) in vector_data.iter().enumerate() { let score = encoded.score_point(&query_u8, index as u32, &counter); let original_score = cosine_similarity(&query, vector); assert!( (score - original_score).abs() < error, "bits={bits:?}, dim={dim}, index={index}, score={score}, expected={original_score}" ); } } } } #[rstest::rstest] #[case::normal(TQMode::Normal)] #[case::plus(TQMode::Plus)] fn test_tq_dot_internal(#[case] mode: TQMode) { for &bits in BITS { for &dim in DIMS { if !should_test(dim, bits) { continue; } let error = error(bits, dot_signal_std(dim)); let mut rng = rand::rngs::StdRng::seed_from_u64(42); let mut vector_data: Vec> = vec![]; for _ in 0..VECTORS_COUNT { vector_data.push((0..dim).map(|_| rng.random_range(-1.0..1.0)).collect()); } let vector_parameters = VectorParameters { dim, deprecated_count: None, distance_type: DistanceType::Dot, invert: false, }; let quantized_vector_size = encoded_vectors_tq::get_quantized_vector_size(&vector_parameters, bits, mode); let encoded = EncodedVectorsTQ::encode( vector_data.iter(), TestEncodedStorageBuilder::new(None, quantized_vector_size), &vector_parameters, VECTORS_COUNT, bits, mode, TQRotation::Padded, false, 1, None, &AtomicBool::new(false), ) .unwrap(); let counter = HardwareCounterCell::new(); for i in 1..VECTORS_COUNT { let score = encoded.score_internal(0, i as u32, &counter); let original_score = dot_similarity(&vector_data[0], &vector_data[i]); assert!( (score - original_score).abs() < error, "bits={bits:?}, dim={dim}, i={i}, score={score}, expected={original_score}" ); } } } } #[rstest::rstest] #[case::normal(TQMode::Normal)] #[case::plus(TQMode::Plus)] fn test_tq_cosine_internal(#[case] mode: TQMode) { for &bits in BITS { for &dim in DIMS { if !should_test(dim, bits) { continue; } let error = error(bits, cosine_signal_std(dim)); let mut rng = rand::rngs::StdRng::seed_from_u64(42); let mut vector_data: Vec> = vec![]; for _ in 0..VECTORS_COUNT { let v: Vec = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect(); vector_data.push(normalize(&v)); } let vector_parameters = VectorParameters { dim, deprecated_count: None, distance_type: DistanceType::Cosine, invert: false, }; let quantized_vector_size = encoded_vectors_tq::get_quantized_vector_size(&vector_parameters, bits, mode); let encoded = EncodedVectorsTQ::encode( vector_data.iter(), TestEncodedStorageBuilder::new(None, quantized_vector_size), &vector_parameters, VECTORS_COUNT, bits, mode, TQRotation::Padded, false, 1, None, &AtomicBool::new(false), ) .unwrap(); let counter = HardwareCounterCell::new(); for i in 1..VECTORS_COUNT { let score = encoded.score_internal(0, i as u32, &counter); let original_score = cosine_similarity(&vector_data[0], &vector_data[i]); assert!( (score - original_score).abs() < error, "bits={bits:?}, dim={dim}, i={i}, score={score}, expected={original_score}" ); } } } } #[rstest::rstest] #[case::normal(TQMode::Normal)] #[case::plus(TQMode::Plus)] fn test_tq_zero_vector_dot(#[case] mode: TQMode) { // A zero vector stored in the dataset, scored with Dot against a // non-zero query, should produce a score close to the true dot // product, which is exactly 0. for &bits in BITS { for &dim in DIMS { if !should_test(dim, bits) { continue; } let error = error(bits, dot_signal_std(dim)); let mut rng = rand::rngs::StdRng::seed_from_u64(42); // Place the zero vector at index 0; fill the rest with random // data so the encoded batch resembles a realistic input. let mut vector_data: Vec> = vec![vec![0.0f32; dim]]; for _ in 1..VECTORS_COUNT { vector_data.push((0..dim).map(|_| rng.random_range(-1.0..1.0)).collect()); } let query: Vec = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect(); let vector_parameters = VectorParameters { dim, deprecated_count: None, distance_type: DistanceType::Dot, invert: false, }; let quantized_vector_size = encoded_vectors_tq::get_quantized_vector_size(&vector_parameters, bits, mode); let encoded = EncodedVectorsTQ::encode( vector_data.iter(), TestEncodedStorageBuilder::new(None, quantized_vector_size), &vector_parameters, VECTORS_COUNT, bits, mode, TQRotation::Padded, false, 1, None, &AtomicBool::new(false), ) .unwrap(); let query_u8 = encoded.encode_query(&query); let counter = HardwareCounterCell::new(); let score = encoded.score_point(&query_u8, 0u32, &counter); assert!( score.abs() < error, "bits={bits:?}, dim={dim}, score={score} (expected ~0)" ); } } } #[rstest::rstest] #[case::normal(TQMode::Normal)] #[case::plus(TQMode::Plus)] fn test_tq_zero_query_dot(#[case] mode: TQMode) { // A zero query, scored with Dot against any encoded vector, should // produce a score close to the true dot product, which is exactly 0. for &bits in BITS { for &dim in DIMS { if !should_test(dim, bits) { continue; } let error = error(bits, dot_signal_std(dim)); let mut rng = rand::rngs::StdRng::seed_from_u64(42); let mut vector_data: Vec> = vec![]; for _ in 0..VECTORS_COUNT { vector_data.push((0..dim).map(|_| rng.random_range(-1.0..1.0)).collect()); } let query: Vec = vec![0.0f32; dim]; let vector_parameters = VectorParameters { dim, deprecated_count: None, distance_type: DistanceType::Dot, invert: false, }; let quantized_vector_size = encoded_vectors_tq::get_quantized_vector_size(&vector_parameters, bits, mode); let encoded = EncodedVectorsTQ::encode( vector_data.iter(), TestEncodedStorageBuilder::new(None, quantized_vector_size), &vector_parameters, VECTORS_COUNT, bits, mode, TQRotation::Padded, false, 1, None, &AtomicBool::new(false), ) .unwrap(); let query_u8 = encoded.encode_query(&query); let counter = HardwareCounterCell::new(); for index in 0..VECTORS_COUNT { let score = encoded.score_point(&query_u8, index as u32, &counter); assert!( score.abs() < error, "bits={bits:?}, dim={dim}, index={index}, score={score} (expected ~0)" ); } } } } #[rstest::rstest] #[case::normal(TQMode::Normal)] #[case::plus(TQMode::Plus)] fn test_tq_zero_vector_cosine(#[case] mode: TQMode) { // A zero vector stored in the dataset, scored with Cosine against a // unit-norm query, should produce a score close to 0. Cosine of a // zero vector is mathematically undefined; the implementation's // convention (preserve zero through preprocessing) yields a true // dot of zero post-rotation, so the encoded score should be ~0. for &bits in BITS { for &dim in DIMS { if !should_test(dim, bits) { continue; } let error = error(bits, cosine_signal_std(dim)); let mut rng = rand::rngs::StdRng::seed_from_u64(42); // Index 0 is the zero vector (not routed through `normalize`, // which would divide by zero); the rest are unit-norm random. let mut vector_data: Vec> = vec![vec![0.0f32; dim]]; for _ in 1..VECTORS_COUNT { let v: Vec = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect(); vector_data.push(normalize(&v)); } let raw_query: Vec = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect(); let query = normalize(&raw_query); let vector_parameters = VectorParameters { dim, deprecated_count: None, distance_type: DistanceType::Cosine, invert: false, }; let quantized_vector_size = encoded_vectors_tq::get_quantized_vector_size(&vector_parameters, bits, mode); let encoded = EncodedVectorsTQ::encode( vector_data.iter(), TestEncodedStorageBuilder::new(None, quantized_vector_size), &vector_parameters, VECTORS_COUNT, bits, mode, TQRotation::Padded, false, 1, None, &AtomicBool::new(false), ) .unwrap(); let query_u8 = encoded.encode_query(&query); let counter = HardwareCounterCell::new(); let score = encoded.score_point(&query_u8, 0u32, &counter); assert!( score.abs() < error, "bits={bits:?}, dim={dim}, score={score} (expected ~0)" ); } } } #[rstest::rstest] #[case::normal(TQMode::Normal)] #[case::plus(TQMode::Plus)] fn test_tq_zero_query_cosine(#[case] mode: TQMode) { // A zero query, scored with Cosine against any unit-norm vector, // should produce a score close to 0. Same convention as above: // zero is preserved through query preprocessing. for &bits in BITS { for &dim in DIMS { if !should_test(dim, bits) { continue; } let error = error(bits, cosine_signal_std(dim)); let mut rng = rand::rngs::StdRng::seed_from_u64(42); let mut vector_data: Vec> = vec![]; for _ in 0..VECTORS_COUNT { let v: Vec = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect(); vector_data.push(normalize(&v)); } let query: Vec = vec![0.0f32; dim]; let vector_parameters = VectorParameters { dim, deprecated_count: None, distance_type: DistanceType::Cosine, invert: false, }; let quantized_vector_size = encoded_vectors_tq::get_quantized_vector_size(&vector_parameters, bits, mode); let encoded = EncodedVectorsTQ::encode( vector_data.iter(), TestEncodedStorageBuilder::new(None, quantized_vector_size), &vector_parameters, VECTORS_COUNT, bits, mode, TQRotation::Padded, false, 1, None, &AtomicBool::new(false), ) .unwrap(); let query_u8 = encoded.encode_query(&query); let counter = HardwareCounterCell::new(); for index in 0..VECTORS_COUNT { let score = encoded.score_point(&query_u8, index as u32, &counter); assert!( score.abs() < error, "bits={bits:?}, dim={dim}, index={index}, score={score} (expected ~0)" ); } } } } #[rstest::rstest] #[case::normal(TQMode::Normal)] #[case::plus(TQMode::Plus)] fn test_tq_dim_one_dot(#[case] mode: TQMode) { // dim=1 is degenerate (no room for Hadamard Gaussianization, quant // noise dominates). Just verify the path doesn't panic and yields // finite, sanely bounded scores; accuracy bounds elsewhere don't apply. let dim = 1; for &bits in BITS { let mut rng = rand::rngs::StdRng::seed_from_u64(42); let mut vector_data: Vec> = vec![]; for _ in 0..VECTORS_COUNT { vector_data.push((0..dim).map(|_| rng.random_range(-1.0..1.0)).collect()); } let query: Vec = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect(); let vector_parameters = VectorParameters { dim, deprecated_count: None, distance_type: DistanceType::Dot, invert: false, }; let quantized_vector_size = encoded_vectors_tq::get_quantized_vector_size(&vector_parameters, bits, mode); let encoded = EncodedVectorsTQ::encode( vector_data.iter(), TestEncodedStorageBuilder::new(None, quantized_vector_size), &vector_parameters, VECTORS_COUNT, bits, mode, TQRotation::Padded, false, 1, None, &AtomicBool::new(false), ) .unwrap(); let query_u8 = encoded.encode_query(&query); let counter = HardwareCounterCell::new(); for index in 0..VECTORS_COUNT { let score = encoded.score_point(&query_u8, index as u32, &counter); assert!( score.is_finite() && score.abs() < 10.0, "bits={bits:?}, index={index}, score={score}" ); } } } // L2 (squared) and L1 score tests. Tolerance for L2 reuses the unified // `error()` with `l2_signal_std`; L1 has its own coefficient table // because its noise/signal scaling per bit doesn't match dot/cosine. #[rstest::rstest] #[case::normal(TQMode::Normal)] #[case::plus(TQMode::Plus)] fn test_tq_l2(#[case] mode: TQMode) { for &bits in BITS { for &dim in DIMS { if !should_test(dim, bits) { continue; } let error = error(bits, l2_signal_std(dim)); let mut rng = rand::rngs::StdRng::seed_from_u64(42); let mut vector_data: Vec> = vec![]; for _ in 0..VECTORS_COUNT { vector_data.push((0..dim).map(|_| rng.random_range(-1.0..1.0)).collect()); } let query: Vec = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect(); let vector_parameters = VectorParameters { dim, deprecated_count: None, distance_type: DistanceType::L2, invert: false, }; let quantized_vector_size = encoded_vectors_tq::get_quantized_vector_size(&vector_parameters, bits, mode); let encoded = EncodedVectorsTQ::encode( vector_data.iter(), TestEncodedStorageBuilder::new(None, quantized_vector_size), &vector_parameters, VECTORS_COUNT, bits, mode, TQRotation::Padded, false, 1, None, &AtomicBool::new(false), ) .unwrap(); let query_u8 = encoded.encode_query(&query); let counter = HardwareCounterCell::new(); for (index, vector) in vector_data.iter().enumerate() { let score = encoded.score_point(&query_u8, index as u32, &counter); let original_score = l2_similarity(&query, vector); assert!( (score - original_score).abs() < error, "bits={bits:?}, dim={dim}, index={index}, score={score}, expected={original_score}" ); } } } } #[rstest::rstest] #[case::normal(TQMode::Normal)] #[case::plus(TQMode::Plus)] fn test_tq_l2_internal(#[case] mode: TQMode) { for &bits in BITS { for &dim in DIMS { if !should_test(dim, bits) { continue; } let error = error(bits, l2_signal_std(dim)); let mut rng = rand::rngs::StdRng::seed_from_u64(42); let mut vector_data: Vec> = vec![]; for _ in 0..VECTORS_COUNT { vector_data.push((0..dim).map(|_| rng.random_range(-1.0..1.0)).collect()); } let vector_parameters = VectorParameters { dim, deprecated_count: None, distance_type: DistanceType::L2, invert: false, }; let quantized_vector_size = encoded_vectors_tq::get_quantized_vector_size(&vector_parameters, bits, mode); let encoded = EncodedVectorsTQ::encode( vector_data.iter(), TestEncodedStorageBuilder::new(None, quantized_vector_size), &vector_parameters, VECTORS_COUNT, bits, mode, TQRotation::Padded, false, 1, None, &AtomicBool::new(false), ) .unwrap(); let counter = HardwareCounterCell::new(); for i in 1..VECTORS_COUNT { let score = encoded.score_internal(0, i as u32, &counter); let original_score = l2_similarity(&vector_data[0], &vector_data[i]); assert!( (score - original_score).abs() < error, "bits={bits:?}, dim={dim}, i={i}, score={score}, expected={original_score}" ); } } } } #[rstest::rstest] #[case::normal(TQMode::Normal)] #[case::plus(TQMode::Plus)] fn test_tq_l1(#[case] mode: TQMode) { for &bits in BITS { for &dim in DIMS { if !should_test(dim, bits) { continue; } let error = error_l1(dim, bits); let mut rng = rand::rngs::StdRng::seed_from_u64(42); let mut vector_data: Vec> = vec![]; for _ in 0..VECTORS_COUNT { vector_data.push((0..dim).map(|_| rng.random_range(-1.0..1.0)).collect()); } let query: Vec = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect(); let vector_parameters = VectorParameters { dim, deprecated_count: None, distance_type: DistanceType::L1, invert: false, }; let quantized_vector_size = encoded_vectors_tq::get_quantized_vector_size(&vector_parameters, bits, mode); let encoded = EncodedVectorsTQ::encode( vector_data.iter(), TestEncodedStorageBuilder::new(None, quantized_vector_size), &vector_parameters, VECTORS_COUNT, bits, mode, TQRotation::Padded, false, 1, None, &AtomicBool::new(false), ) .unwrap(); let query_u8 = encoded.encode_query(&query); let counter = HardwareCounterCell::new(); for (index, vector) in vector_data.iter().enumerate() { let score = encoded.score_point(&query_u8, index as u32, &counter); let original_score = l1_similarity(&query, vector); assert!( (score - original_score).abs() < error, "bits={bits:?}, dim={dim}, index={index}, score={score}, expected={original_score}" ); } } } } #[rstest::rstest] #[case::normal(TQMode::Normal)] #[case::plus(TQMode::Plus)] fn test_tq_l1_internal(#[case] mode: TQMode) { for &bits in BITS { for &dim in DIMS { if !should_test(dim, bits) { continue; } let error = error_l1(dim, bits); let mut rng = rand::rngs::StdRng::seed_from_u64(42); let mut vector_data: Vec> = vec![]; for _ in 0..VECTORS_COUNT { vector_data.push((0..dim).map(|_| rng.random_range(-1.0..1.0)).collect()); } let vector_parameters = VectorParameters { dim, deprecated_count: None, distance_type: DistanceType::L1, invert: false, }; let quantized_vector_size = encoded_vectors_tq::get_quantized_vector_size(&vector_parameters, bits, mode); let encoded = EncodedVectorsTQ::encode( vector_data.iter(), TestEncodedStorageBuilder::new(None, quantized_vector_size), &vector_parameters, VECTORS_COUNT, bits, mode, TQRotation::Padded, false, 1, None, &AtomicBool::new(false), ) .unwrap(); let counter = HardwareCounterCell::new(); for i in 1..VECTORS_COUNT { let score = encoded.score_internal(0, i as u32, &counter); let original_score = l1_similarity(&vector_data[0], &vector_data[i]); assert!( (score - original_score).abs() < error, "bits={bits:?}, dim={dim}, i={i}, score={score}, expected={original_score}" ); } } } } /// Recall regression probe — non-uniform per-coordinate variance data /// (the case where TQ+ should HELP, not hurt). Asserts Plus-mode recall@k /// stays close to Normal-mode on data with a few high-variance "spike" /// directions. Catches the failure mode where renorm's `cn` drifts because /// `‖X+‖` has chi-squared spread across vectors — see /// [`TurboQuantizer::compute_centroid_norm`] for the EC-revert that fixes it. #[rstest::rstest] #[case::bits1(TQBits::Bits1)] #[case::bits2(TQBits::Bits2)] #[case::bits4(TQBits::Bits4)] fn recall_skewed_data(#[case] bits: TQBits) { use rand::prelude::StdRng; let dim = 256; let n = 1000; let n_queries = 50; let topk = 10; let mut rng = StdRng::seed_from_u64(42); // Pre-rotation per-coord scale: a few coords have huge variance, most // have small. Realistic embeddings have similar structure (a few // "spike" directions). let coord_scales: Vec = (0..dim).map(|i| if i < 8 { 100.0 } else { 1.0 }).collect(); let make_vec = |rng: &mut StdRng| -> Vec { coord_scales .iter() .map(|&s| s * (rng.random_range(-1.0..1.0))) .collect() }; let vectors: Vec> = (0..n).map(|_| make_vec(&mut rng)).collect(); let queries: Vec> = (0..n_queries).map(|_| make_vec(&mut rng)).collect(); let true_dot = |a: &[f32], b: &[f32]| -> f32 { a.iter().zip(b.iter()).map(|(&x, &y)| x * y).sum() }; let recall_for = |mode: TQMode| -> f32 { let vp = VectorParameters { dim, distance_type: DistanceType::Dot, invert: false, deprecated_count: None, }; let qsize = encoded_vectors_tq::get_quantized_vector_size(&vp, bits, mode); let encoded = EncodedVectorsTQ::encode( vectors.iter(), TestEncodedStorageBuilder::new(None, qsize), &vp, n, bits, mode, TQRotation::Padded, false, 1, None, &AtomicBool::new(false), ) .unwrap(); let counter = HardwareCounterCell::new(); let mut total = 0.0; for q in &queries { let mut truth: Vec<(usize, f32)> = vectors .iter() .enumerate() .map(|(i, v)| (i, true_dot(q, v))) .collect(); truth.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); let truth_top: Vec = truth.iter().take(topk).map(|x| x.0).collect(); let qq = encoded.encode_query(q); let mut q_scores: Vec<(usize, f32)> = (0..n) .map(|i| (i, encoded.score_point(&qq, i as u32, &counter))) .collect(); q_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); let q_top: Vec = q_scores.iter().take(topk).map(|x| x.0).collect(); let hits = truth_top.iter().filter(|i| q_top.contains(i)).count(); total += hits as f32 / topk as f32; } total / n_queries as f32 }; let normal = recall_for(TQMode::Normal); let plus = recall_for(TQMode::Plus); // Plus must not be meaningfully worse than Normal. We allow a small // 2% slack since the codebook is the same and EC's value comes from // distribution fit, which on this seed is marginal. assert!( plus >= normal - 0.02, "bits={bits:?}: Plus recall regressed (Normal={normal:.3}, Plus={plus:.3})" ); } #[test] fn test_tq_layout_size_is_multiple_of_alignment() { const LAYOUT_DIMS: &[usize] = &[1, 7, 33, 65, 100, 756, 768]; let vectors_count = 8; for &dim in LAYOUT_DIMS { for &bits in BITS { for &distance in &[ DistanceType::Dot, DistanceType::Cosine, DistanceType::L1, DistanceType::L2, ] { let mut rng = rand::rngs::StdRng::seed_from_u64(42); let vector_data: Vec> = (0..vectors_count) .map(|_| { let vector: Vec = (0..dim).map(|_| rng.random_range(-1.0..1.0)).collect(); match distance { DistanceType::Cosine => normalize(&vector), DistanceType::Dot | DistanceType::L1 | DistanceType::L2 => vector, } }) .collect(); let vector_parameters = VectorParameters { dim, deprecated_count: None, distance_type: distance, invert: false, }; for &mode in &[TQMode::Normal, TQMode::Plus] { let quantized_vector_size = encoded_vectors_tq::get_quantized_vector_size( &vector_parameters, bits, mode, ); let encoded = EncodedVectorsTQ::encode( vector_data.iter(), TestEncodedStorageBuilder::new(None, quantized_vector_size), &vector_parameters, vectors_count, bits, mode, TQRotation::Padded, false, 1, None, &AtomicBool::new(false), ) .unwrap(); let layout = encoded.layout(); assert_eq!(layout.size(), quantized_vector_size); assert_eq!( layout.size() % layout.align(), 0, "dim={dim} bits={bits:?} distance={distance:?} mode={mode:?}: \ layout size {} is not a multiple of alignment {}", layout.size(), layout.align(), ); } } } } } }