Merge commit 'd803bfe2b1fe7f5e219e50ac20d6801a0a58ac75' as 'vendor/ruvector'

vendor/ruvector/crates/ruvector-dither/src/quantize.rs

@@ -0,0 +1,134 @@
+//! Drop-in quantization helpers that apply dither before rounding.
+
+use crate::DitherSource;
+
+/// Quantize a single value with deterministic dither.
+///
+/// # Arguments
+/// - `x`      – input activation in `[-1.0, 1.0]`
+/// - `bits`   – quantizer bit-width (e.g. 3, 5, 7, 8)
+/// - `eps`    – dither amplitude in LSB units (0.0 = no dither, 0.5 = half-LSB recommended)
+/// - `source` – stateful dither sequence
+///
+/// Returns the quantized value in `[-1.0, 1.0]`.
+///
+/// # Example
+/// ```
+/// use ruvector_dither::{GoldenRatioDither, quantize_dithered};
+/// let mut d = GoldenRatioDither::new(0.0);
+/// let q = quantize_dithered(0.314, 8, 0.5, &mut d);
+/// assert!(q >= -1.0 && q <= 1.0);
+/// ```
+#[inline]
+pub fn quantize_dithered(x: f32, bits: u32, eps: f32, source: &mut impl DitherSource) -> f32 {
+    assert!(bits >= 2 && bits <= 31, "bits must be in [2, 31]");
+    let qmax = ((1u32 << (bits - 1)) - 1) as f32;
+    let lsb = 1.0 / qmax;
+    let dither = source.next(eps * lsb);
+    let shifted = (x + dither) * qmax;
+    let rounded = shifted.round().clamp(-qmax, qmax);
+    rounded / qmax
+}
+
+/// Quantize a slice in-place with deterministic dither.
+///
+/// Each element gets an independent dither sample from `source`.
+///
+/// # Example
+/// ```
+/// use ruvector_dither::{GoldenRatioDither, quantize_slice_dithered};
+/// let mut vals = vec![0.1_f32, 0.5, -0.3, 0.9, -0.8];
+/// let mut d = GoldenRatioDither::new(0.0);
+/// quantize_slice_dithered(&mut vals, 5, 0.5, &mut d);
+/// for &v in &vals {
+///     assert!(v >= -1.0 && v <= 1.0);
+/// }
+/// ```
+pub fn quantize_slice_dithered(
+    xs: &mut [f32],
+    bits: u32,
+    eps: f32,
+    source: &mut impl DitherSource,
+) {
+    assert!(bits >= 2 && bits <= 31, "bits must be in [2, 31]");
+    let qmax = ((1u32 << (bits - 1)) - 1) as f32;
+    let lsb = 1.0 / qmax;
+    for x in xs.iter_mut() {
+        let dither = source.next(eps * lsb);
+        let shifted = (*x + dither) * qmax;
+        *x = shifted.round().clamp(-qmax, qmax) / qmax;
+    }
+}
+
+/// Quantize to a raw integer code (signed, in `[-(2^(bits-1)), 2^(bits-1)-1]`).
+///
+/// Useful when you need the integer representation rather than a re-scaled float.
+#[inline]
+pub fn quantize_to_code(x: f32, bits: u32, eps: f32, source: &mut impl DitherSource) -> i32 {
+    assert!(bits >= 2 && bits <= 31, "bits must be in [2, 31]");
+    let qmax = ((1u32 << (bits - 1)) - 1) as f32;
+    let lsb = 1.0 / qmax;
+    let dither = source.next(eps * lsb);
+    ((x + dither) * qmax).round().clamp(-qmax, qmax) as i32
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{GoldenRatioDither, PiDither};
+
+    #[test]
+    fn output_in_unit_range() {
+        let mut d = GoldenRatioDither::new(0.0);
+        for bits in [3u32, 5, 7, 8] {
+            for &x in &[-1.0_f32, -0.5, 0.0, 0.5, 1.0] {
+                let q = quantize_dithered(x, bits, 0.5, &mut d);
+                assert!(q >= -1.0 && q <= 1.0, "bits={bits}, x={x}, q={q}");
+            }
+        }
+    }
+
+    #[test]
+    fn dither_reduces_idle_tones() {
+        // A constant signal at exactly 0.5 * LSB without dither quantizes
+        // to the same code every time (idle tone).  With dither the code
+        // alternates, so the variance of codes should be > 0.
+        let bits = 5u32;
+        let qmax = ((1u32 << (bits - 1)) - 1) as f32;
+        let lsb = 1.0 / qmax;
+        let x = 0.5 * lsb; // exactly half an LSB
+
+        let mut codes_with: Vec<i32> = Vec::with_capacity(256);
+        let mut d = GoldenRatioDither::new(0.0);
+        for _ in 0..256 {
+            codes_with.push(quantize_to_code(x, bits, 0.5, &mut d));
+        }
+        let unique: std::collections::HashSet<i32> = codes_with.iter().copied().collect();
+        assert!(
+            unique.len() > 1,
+            "dithered signal must produce >1 unique code"
+        );
+    }
+
+    #[test]
+    fn slice_quantize_in_bounds() {
+        let mut vals: Vec<f32> = (-50..=50).map(|i| i as f32 * 0.02).collect();
+        let mut pi = PiDither::new(0);
+        quantize_slice_dithered(&mut vals, 7, 0.5, &mut pi);
+        for v in vals {
+            assert!(v >= -1.0 && v <= 1.0, "out of range: {v}");
+        }
+    }
+
+    #[test]
+    fn deterministic_with_same_seed() {
+        let input = vec![0.1_f32, 0.4, -0.7, 0.9];
+        let quantize = |input: &[f32]| {
+            let mut buf = input.to_vec();
+            let mut d = GoldenRatioDither::new(0.5);
+            quantize_slice_dithered(&mut buf, 8, 0.5, &mut d);
+            buf
+        };
+        assert_eq!(quantize(&input), quantize(&input));
+    }
+}