qntz 0.1.2

Vector quantization primitives (RaBitQ, ternary, bit packing) for ANN systems.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

//! Entropy-coded quantization: RaBitQ + ANS.
//!
//! Demonstrates the pipeline: generate vectors -> fit RaBitQ -> quantize ->
//! count code frequencies -> build ANS table -> entropy-encode -> decode -> verify.
//!
//! After quantization, code distributions are non-uniform (some levels appear far
//! more often than others), so entropy coding compresses below the fixed-width
//! `total_bits * dim` rate.
//!
//! Run:
//!
//! ```sh
//! cargo run --example entropy_coded_quantization --features rabitq
//! ```

use ans::{decode, encode, FrequencyTable};
use qntz::rabitq::{RaBitQConfig, RaBitQQuantizer};

// Simple xorshift64 PRNG (no extra deps).
struct Xorshift64(u64);

impl Xorshift64 {
    fn new(seed: u64) -> Self {
        Self(seed.max(1))
    }

    fn next_u64(&mut self) -> u64 {
        let mut x = self.0;
        x ^= x << 13;
        x ^= x >> 7;
        x ^= x << 17;
        self.0 = x;
        x
    }

    /// Uniform f32 in (-1, 1).
    fn next_f32(&mut self) -> f32 {
        (self.next_u64() as f64 / u64::MAX as f64 * 2.0 - 1.0) as f32
    }
}

fn main() {
    let dim = 64;
    let num_vectors = 500;
    let total_bits: usize = 4; // 1 sign + 3 extended -> 16 possible code values
    let seed = 0xCAFE_BABE;

    // -- Generate random vectors --
    let mut rng = Xorshift64::new(seed);
    let mut flat_vectors = Vec::with_capacity(num_vectors * dim);
    for _ in 0..num_vectors * dim {
        flat_vectors.push(rng.next_f32());
    }

    // -- Fit quantizer --
    let config = RaBitQConfig {
        total_bits,
        t_const: None,
    };
    let mut quantizer = RaBitQQuantizer::with_config(dim, seed, config).unwrap();
    quantizer.fit(&flat_vectors, num_vectors).unwrap();

    // -- Quantize all vectors, collect codes --
    let alphabet_size = 1usize << total_bits; // 16 for 4-bit
    let mut all_codes: Vec<u32> = Vec::with_capacity(num_vectors * dim);
    let mut freq_counts = vec![0u32; alphabet_size];

    for i in 0..num_vectors {
        let vec_slice = &flat_vectors[i * dim..(i + 1) * dim];
        let qv = quantizer.quantize(vec_slice).unwrap();
        for &code in &qv.codes {
            let c = code as u32;
            all_codes.push(c);
            freq_counts[c as usize] += 1;
        }
    }

    let total_symbols = all_codes.len();

    // -- Print code distribution --
    println!("Code distribution ({total_symbols} symbols, {alphabet_size} alphabet):");
    let empirical_entropy = {
        let mut h = 0.0f64;
        let n = total_symbols as f64;
        for (sym, &count) in freq_counts.iter().enumerate() {
            let frac = count as f64 / n;
            if frac > 0.0 {
                h -= frac * frac.log2();
                println!("  sym {sym:>2}: count {count:>6}  ({:.1}%)", frac * 100.0);
            }
        }
        h
    };
    println!("  empirical entropy: {empirical_entropy:.3} bits/symbol");

    // -- Ensure no zero-frequency symbols (ANS requires nonzero for used symbols) --
    // We only need entries for symbols that appear. But FrequencyTable::from_counts
    // expects a contiguous alphabet 0..N, and we need to encode symbols by their
    // original value. So we keep the full alphabet and replace zeros with 1 (minimal
    // frequency) to avoid ZeroFrequency errors if any unused symbol is somehow encoded.
    let adjusted_counts: Vec<u32> = freq_counts
        .iter()
        .map(|&c| if c == 0 { 1 } else { c })
        .collect();

    // -- Build ANS frequency table --
    let precision_bits = 14;
    let table = FrequencyTable::from_counts(&adjusted_counts, precision_bits)
        .expect("failed to build frequency table");

    // -- Entropy-encode --
    let encoded_bytes = encode(&all_codes, &table).expect("ANS encode failed");

    // -- Decode and verify roundtrip --
    let decoded = decode(&encoded_bytes, &table, total_symbols).expect("ANS decode failed");
    assert_eq!(
        all_codes, decoded,
        "roundtrip mismatch: ANS decode differs from original codes"
    );
    println!("\nRoundtrip verified: all {total_symbols} symbols match.");

    // -- Size report --
    let raw_bytes = total_symbols * 2; // u16 per code
    let fixed_bits = total_symbols * total_bits;
    let fixed_bytes = fixed_bits.div_ceil(8);
    let ans_bytes = encoded_bytes.len();
    let theoretical_bytes = ((empirical_entropy * total_symbols as f64) / 8.0).ceil() as usize;

    println!("\nSize comparison:");
    println!("  raw codes (u16):        {raw_bytes:>8} bytes");
    println!("  fixed-width ({total_bits}-bit):    {fixed_bytes:>8} bytes");
    println!(
        "  ANS entropy-coded:      {:>8} bytes  ({:.2}x vs fixed-width)",
        ans_bytes,
        fixed_bytes as f64 / ans_bytes as f64
    );
    println!("  theoretical minimum:    {theoretical_bytes:>8} bytes  (Shannon entropy)");
    println!(
        "  ANS overhead vs theory: {:.1}%",
        (ans_bytes as f64 / theoretical_bytes as f64 - 1.0) * 100.0
    );
    println!(
        "  bits/symbol: fixed={total_bits:.1}, ANS={:.3}, entropy={empirical_entropy:.3}",
        (ans_bytes as f64 * 8.0) / total_symbols as f64
    );
}