fib_quant/
codebook.rs

1use serde::{Deserialize, Serialize};
2
3use crate::{
4    digest::json_digest,
5    directions::directions_for_method,
6    lloyd::{refine_codebook, LloydReportV1},
7    profile::{FibQuantProfileV1, LloydMode, RadiusMethod},
8    rotation::StoredRotation,
9    spherical_beta::{radius_quantile, radius_quantile_k2_closed_form},
10    FibQuantError, Result,
11};
12
13pub const CODEBOOK_SCHEMA: &str = "fib_codebook_v1";
14
15/// Persisted FibQuant codebook in row-major codeword order.
16#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
17pub struct FibCodebookV1 {
18    /// Stable schema marker.
19    pub schema_version: String,
20    /// Profile used to create this codebook.
21    pub profile: FibQuantProfileV1,
22    /// Digest of `profile`.
23    pub profile_digest: String,
24    /// Digest of profile, report, and row-major codewords.
25    pub codebook_digest: String,
26    /// Digest of deterministic ambient rotation identity and matrix.
27    pub rotation_digest: String,
28    /// Row-major `N x k` f32 codeword storage.
29    pub codewords: Vec<f32>,
30    /// Initial codebook training MSE before Lloyd-Max.
31    pub init_mse: f64,
32    /// Best training MSE after Lloyd-Max.
33    pub training_mse: f64,
34    /// Lloyd-Max refinement receipt.
35    pub refinement_report: LloydReportV1,
36}
37
38impl FibCodebookV1 {
39    /// Build a full refined codebook for a profile.
40    pub fn build(profile: FibQuantProfileV1) -> Result<Self> {
41        profile.validate()?;
42        let profile_digest = profile.digest()?;
43        let rotation_digest =
44            StoredRotation::new(profile.ambient_dim as usize, profile.rotation_seed)?.digest()?;
45        let initial = build_initial_codebook(&profile)?;
46        let refined = match profile.lloyd_mode {
47            LloydMode::Refine => refine_codebook(&profile, &initial)?,
48            LloydMode::Skip => crate::lloyd::skip_refinement(&profile, &initial)?,
49        };
50        let mut codebook = Self {
51            schema_version: CODEBOOK_SCHEMA.into(),
52            profile,
53            profile_digest,
54            codebook_digest: String::new(),
55            rotation_digest,
56            codewords: refined.codewords,
57            init_mse: refined.init_mse,
58            training_mse: refined.training_mse,
59            refinement_report: refined.report,
60        };
61        codebook.codebook_digest = codebook.compute_digest()?;
62        Ok(codebook)
63    }
64
65    /// Validate shape and digest fields.
66    pub fn validate(&self) -> Result<()> {
67        if self.schema_version != CODEBOOK_SCHEMA {
68            return Err(FibQuantError::CorruptPayload(format!(
69                "codebook schema_version {}, expected {CODEBOOK_SCHEMA}",
70                self.schema_version
71            )));
72        }
73        self.profile.validate()?;
74        self.refinement_report
75            .validate_against_profile(&self.profile)?;
76        let expected_profile = self.profile.digest()?;
77        if self.profile_digest != expected_profile {
78            return Err(FibQuantError::ProfileDigestMismatch {
79                expected: expected_profile,
80                actual: self.profile_digest.clone(),
81            });
82        }
83        let expected_rotation = StoredRotation::new(
84            self.profile.ambient_dim as usize,
85            self.profile.rotation_seed,
86        )?
87        .digest()?;
88        if self.rotation_digest != expected_rotation {
89            return Err(FibQuantError::RotationDigestMismatch {
90                expected: expected_rotation,
91                actual: self.rotation_digest.clone(),
92            });
93        }
94        let expected_codebook = self.compute_digest()?;
95        if self.codebook_digest != expected_codebook {
96            return Err(FibQuantError::CodebookDigestMismatch {
97                expected: expected_codebook,
98                actual: self.codebook_digest.clone(),
99            });
100        }
101        let expected_len = (self.profile.codebook_size as usize)
102            .checked_mul(self.profile.block_dim as usize)
103            .ok_or_else(|| {
104                FibQuantError::ResourceLimitExceeded("codebook value count overflow".into())
105            })?;
106        if self.codewords.len() != expected_len {
107            return Err(FibQuantError::CorruptPayload(format!(
108                "codebook has {} values, expected {expected_len}",
109                self.codewords.len()
110            )));
111        }
112        if self.codewords.iter().any(|value| !value.is_finite()) {
113            return Err(FibQuantError::CorruptPayload(
114                "codebook contains non-finite value".into(),
115            ));
116        }
117        Ok(())
118    }
119
120    /// Return codeword `index` as f64 values.
121    pub fn codeword(&self, index: usize) -> Result<Vec<f64>> {
122        let n = self.profile.codebook_size as usize;
123        let k = self.profile.block_dim as usize;
124        if index >= n {
125            return Err(FibQuantError::IndexOutOfRange {
126                index: index as u32,
127                codebook_size: n as u32,
128            });
129        }
130        Ok(self.codewords[index * k..(index + 1) * k]
131            .iter()
132            .map(|value| f64::from(*value))
133            .collect())
134    }
135
136    /// Deterministic codebook digest.
137    pub fn compute_digest(&self) -> Result<String> {
138        #[derive(Serialize)]
139        struct DigestView<'a> {
140            schema_version: &'a str,
141            profile_digest: &'a str,
142            rotation_digest: &'a str,
143            codewords: &'a [f32],
144            init_mse: f64,
145            training_mse: f64,
146            refinement_report: &'a LloydReportV1,
147        }
148        json_digest(
149            CODEBOOK_SCHEMA,
150            &DigestView {
151                schema_version: &self.schema_version,
152                profile_digest: &self.profile_digest,
153                rotation_digest: &self.rotation_digest,
154                codewords: &self.codewords,
155                init_mse: self.init_mse,
156                training_mse: self.training_mse,
157                refinement_report: &self.refinement_report,
158            },
159        )
160    }
161}
162
163/// Build deterministic radial-angular initialization.
164pub fn build_initial_codebook(profile: &FibQuantProfileV1) -> Result<Vec<f64>> {
165    profile.validate()?;
166    let d = profile.ambient_dim as usize;
167    let k = profile.block_dim as usize;
168    let n = profile.codebook_size as usize;
169    let directions = directions_for_method(k, n, &profile.direction_method)?;
170    let value_count = n.checked_mul(k).ok_or_else(|| {
171        FibQuantError::ResourceLimitExceeded("codebook value count overflow".into())
172    })?;
173    let mut codewords = Vec::with_capacity(value_count);
174    for (idx, direction) in directions.iter().enumerate() {
175        let radius = radius_for_method(profile, d, k, idx + 1, n)?;
176        for value in direction {
177            let code = radius * value;
178            if !code.is_finite() {
179                return Err(FibQuantError::NumericalFailure(
180                    "non-finite initialized codeword".into(),
181                ));
182            }
183            codewords.push(code);
184        }
185    }
186    Ok(codewords)
187}
188
189fn radius_for_method(
190    profile: &FibQuantProfileV1,
191    d: usize,
192    k: usize,
193    idx: usize,
194    n: usize,
195) -> Result<f64> {
196    match profile.radius_method {
197        RadiusMethod::K2ClosedForm if k == 2 => {
198            let q = (idx as f64 - 0.5) / n as f64;
199            radius_quantile_k2_closed_form(d, q)
200        }
201        RadiusMethod::BetaQuantile if k >= 3 => radius_quantile(d, k, idx, n),
202        _ => Err(FibQuantError::CorruptPayload(format!(
203            "radius method {:?} is not supported for k={k}",
204            profile.radius_method
205        ))),
206    }
207}
fib_quant/codebook.rs

fib_quant/
codebook.rs
