fib_quant/
lattice.rs

1/// Experimental lattice quantization module.
2///
3/// Provides Z1 (integer lattice) and A2-shaped (hexagonal lattice placeholder)
4/// quantization. A2 operates on pairs using scalar rounding with pair-preserving
5/// padding — this is NOT a full hexagonal nearest-lattice implementation.
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum LatticeKind {
8    Z1,
9    A2,
10}
11
12#[derive(Debug, Clone, PartialEq)]
13pub struct LatticeQuantizationResult {
14    pub kind: LatticeKind,
15    pub codes: Vec<i16>,
16    pub reconstructed: Vec<f32>,
17    pub mse: f32,
18}
19
20fn sanitize_scale(scale: f32) -> f32 {
21    if scale.is_finite() && scale > 0.0 {
22        scale
23    } else {
24        1.0
25    }
26}
27
28fn clamp_to_i16(v: f32) -> i16 {
29    let rounded = v.round();
30    if rounded >= i16::MAX as f32 {
31        i16::MAX
32    } else if rounded <= i16::MIN as f32 {
33        i16::MIN
34    } else {
35        rounded as i16
36    }
37}
38
39fn mse(original: &[f32], reconstructed: &[f32]) -> f32 {
40    let n = original.len();
41    if n == 0 {
42        return 0.0;
43    }
44    let sum: f32 = original
45        .iter()
46        .zip(reconstructed.iter())
47        .map(|(a, b)| (a - b).powi(2))
48        .sum();
49    sum / n as f32
50}
51
52/// Quantize via Z1 (integer lattice): round(value * scale), reconstruct = code / scale.
53pub fn quantize_z1(values: &[f32], scale: f32) -> LatticeQuantizationResult {
54    let scale = sanitize_scale(scale);
55    let codes: Vec<i16> = values.iter().map(|&v| clamp_to_i16(v * scale)).collect();
56    let reconstructed: Vec<f32> = codes.iter().map(|&c| c as f32 / scale).collect();
57    let mse = mse(values, &reconstructed);
58    LatticeQuantizationResult {
59        kind: LatticeKind::Z1,
60        codes,
61        reconstructed,
62        mse,
63    }
64}
65
66/// Quantize via A2-shaped pair quantization (placeholder, not full hexagonal nearest-lattice).
67///
68/// Operates on consecutive pairs. If the input length is odd, the trailing element
69/// is quantized independently via Z1 behavior and appended.
70pub fn quantize_a2_pairs(values: &[f32], scale: f32) -> LatticeQuantizationResult {
71    let scale = sanitize_scale(scale);
72    let n = values.len();
73    let mut codes = Vec::with_capacity(n);
74    let mut reconstructed = Vec::with_capacity(n);
75
76    let pairs = n / 2;
77    for i in 0..pairs {
78        let a = values[2 * i];
79        let b = values[2 * i + 1];
80        let ca = clamp_to_i16(a * scale);
81        let cb = clamp_to_i16(b * scale);
82        codes.push(ca);
83        codes.push(cb);
84        reconstructed.push(ca as f32 / scale);
85        reconstructed.push(cb as f32 / scale);
86    }
87
88    // Odd trailing element: Z1 scalar fallback.
89    if n % 2 == 1 {
90        let v = values[n - 1];
91        let c = clamp_to_i16(v * scale);
92        codes.push(c);
93        reconstructed.push(c as f32 / scale);
94    }
95
96    let mse = mse(values, &reconstructed);
97    LatticeQuantizationResult {
98        kind: LatticeKind::A2,
99        codes,
100        reconstructed,
101        mse,
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108
109    #[test]
110    fn z1_deterministic_and_same_length() {
111        let values = vec![0.1, -0.5, 1.2, -3.7, 0.0];
112        let r1 = quantize_z1(&values, 10.0);
113        let r2 = quantize_z1(&values, 10.0);
114        assert_eq!(r1.codes, r2.codes);
115        assert_eq!(r1.codes.len(), values.len());
116        assert_eq!(r1.reconstructed.len(), values.len());
117    }
118
119    #[test]
120    fn z1_mse_finite() {
121        let values: Vec<f32> = (0..64).map(|i| i as f32 * 0.1 - 3.2).collect();
122        let r = quantize_z1(&values, 4.0);
123        assert!(r.mse.is_finite());
124    }
125
126    #[test]
127    fn a2_preserves_length_even() {
128        let values = vec![1.0, 2.0, 3.0, 4.0];
129        let r = quantize_a2_pairs(&values, 1.0);
130        assert_eq!(r.codes.len(), 4);
131        assert_eq!(r.reconstructed.len(), 4);
132    }
133
134    #[test]
135    fn a2_preserves_length_odd() {
136        let values = vec![1.0, 2.0, 3.0];
137        let r = quantize_a2_pairs(&values, 1.0);
138        assert_eq!(r.codes.len(), 3);
139        assert_eq!(r.reconstructed.len(), 3);
140    }
141
142    #[test]
143    fn nonpositive_scale_falls_back_to_1() {
144        let values = vec![2.0, -2.0];
145        let r0 = quantize_z1(&values, 0.0);
146        let rn = quantize_z1(&values, -5.0);
147        let rf = quantize_z1(&values, f32::NAN);
148        let r1 = quantize_z1(&values, 1.0);
149        assert_eq!(r0.codes, r1.codes);
150        assert_eq!(rn.codes, r1.codes);
151        assert_eq!(rf.codes, r1.codes);
152    }
153
154    #[test]
155    fn large_values_clamp_to_i16_bounds() {
156        let values = vec![1e10_f32, -1e10_f32];
157        let r = quantize_z1(&values, 1.0);
158        assert_eq!(r.codes[0], i16::MAX);
159        assert_eq!(r.codes[1], i16::MIN);
160    }
161}
fib_quant/lattice.rs

fib_quant/
lattice.rs
