fib_quant/
profile.rs

1use serde::{Deserialize, Serialize};
2
3use crate::{digest::json_digest, rotation::ROTATION_ALGORITHM_VERSION, FibQuantError, Result};
4
5pub const PROFILE_SCHEMA: &str = "fib_quant_profile_v1";
6/// Maximum ambient dimension accepted by the alpha profile validator.
7pub const MAX_AMBIENT_DIM: usize = 16_384;
8/// Maximum block dimension accepted by the alpha profile validator.
9pub const MAX_BLOCK_DIM: usize = 256;
10/// Maximum codebook size accepted by the alpha profile validator.
11pub const MAX_CODEBOOK_SIZE: usize = 1 << 20;
12/// Maximum Lloyd training samples accepted by the alpha profile validator.
13pub const MAX_TRAINING_SAMPLES: u32 = 10_000_000;
14/// Maximum number of scalar values in a dense rotation matrix.
15pub const MAX_ROTATION_MATRIX_VALUES: usize = 16_777_216;
16/// Maximum number of scalar values in an `N x k` codebook.
17pub const MAX_CODEBOOK_VALUES: usize = 67_108_864;
18/// Maximum bits in a packed fixed-rate payload.
19pub const MAX_PACKED_INDEX_BITS: usize = 1 << 34;
20
21const RATE_TOLERANCE: f64 = 1.0e-12;
22const MAX_LLOYD_RESTARTS: u32 = 1_024;
23const MAX_LLOYD_ITERATIONS: u32 = 100_000;
24
25/// Norm payload representation.
26#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
27#[non_exhaustive]
28#[serde(rename_all = "snake_case")]
29pub enum NormFormat {
30    /// Paper path: fp16 scalar norm side header.
31    Fp16Paper,
32    /// Reference/test path: f32 scalar norm side header.
33    #[doc(hidden)]
34    F32Reference,
35}
36
37/// Source used for training samples.
38#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
39#[non_exhaustive]
40#[serde(rename_all = "snake_case")]
41pub enum SourceMode {
42    /// Direct spherical-Beta sampler.
43    CanonicalSphericalBeta,
44    /// Normalized Gaussian projection reference sampler.
45    #[doc(hidden)]
46    ReferenceGaussianProjection,
47}
48
49/// Radius initialization method.
50#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
51#[non_exhaustive]
52#[serde(rename_all = "snake_case")]
53pub enum RadiusMethod {
54    /// Bennett-Gersho Beta-quantile radii.
55    BetaQuantile,
56    /// Paper closed form for k=2.
57    K2ClosedForm,
58    /// Explicit large-d single-shell initialization.
59    #[doc(hidden)]
60    LargeDSingleShellExplicit,
61}
62
63/// Direction initialization method.
64#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
65#[non_exhaustive]
66#[serde(rename_all = "snake_case")]
67pub enum DirectionMethod {
68    /// Planar Fibonacci spiral.
69    FibonacciSpiral,
70    /// Fibonacci sphere.
71    FibonacciSphere,
72    /// Roberts-Kronecker rank-one sequence.
73    RobertsKronecker,
74}
75
76/// Empty-cell handling during Lloyd-Max refinement.
77#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
78#[non_exhaustive]
79#[serde(rename_all = "snake_case")]
80pub enum EmptyCellPolicy {
81    /// Split the occupied cell with highest distortion.
82    SplitHighestDistortion,
83    /// Fail if any cell is empty.
84    FailClosed,
85}
86
87/// Lloyd-Max refinement mode.
88#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
89#[non_exhaustive]
90#[serde(rename_all = "snake_case")]
91pub enum LloydMode {
92    /// Run multi-restart Lloyd-Max refinement (paper path).
93    Refine,
94    /// Skip refinement entirely — use the radial-angular initialization as-is.
95    /// Faster codebook construction (~6ms saved) at the cost of higher MSE.
96    /// The receipt will reflect this choice via `best_mse == init_mse`.
97    Skip,
98}
99
100/// Stable profile for paper-faithful FibQuant codebooks and payloads.
101#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
102pub struct FibQuantProfileV1 {
103    /// Stable schema marker.
104    pub schema_version: String,
105    /// Ambient vector dimension `d`.
106    pub ambient_dim: u32,
107    /// Block dimension `k`.
108    pub block_dim: u32,
109    /// Codebook size `N`.
110    pub codebook_size: u32,
111    /// Paper dense rate `log2(N) / k`.
112    pub paper_rate_bits_per_coord: f64,
113    /// Practical fixed-rate index width `ceil(log2(N))`.
114    pub wire_index_bits: u8,
115    /// Practical wire rate `wire_index_bits / k`.
116    pub wire_bits_per_coord: f64,
117    /// Norm header format.
118    pub norm_format: NormFormat,
119    /// Seed for ambient rotation.
120    pub rotation_seed: u64,
121    /// Rotation generation algorithm identity.
122    pub rotation_algorithm_version: String,
123    /// Seed for codebook construction and Lloyd training.
124    pub codebook_seed: u64,
125    /// Codebook algorithm/version string.
126    pub codebook_version: String,
127    /// Training source mode.
128    pub source_mode: SourceMode,
129    /// Radius method.
130    pub radius_method: RadiusMethod,
131    /// Direction method.
132    pub direction_method: DirectionMethod,
133    /// Number of Lloyd restarts.
134    pub lloyd_restarts: u32,
135    /// Number of Lloyd iterations per restart.
136    pub lloyd_iterations: u32,
137    /// Number of training samples.
138    pub training_samples: u32,
139    /// Empty-cell repair policy.
140    pub empty_cell_policy: EmptyCellPolicy,
141    /// Lloyd-Max refinement mode: Refine (paper path) or Skip (fast path).
142    pub lloyd_mode: LloydMode,
143}
144
145impl FibQuantProfileV1 {
146    /// Build a validated paper profile with method choices derived from `k`.
147    pub fn paper_default(
148        ambient_dim: usize,
149        block_dim: usize,
150        codebook_size: usize,
151        seed: u64,
152    ) -> Result<Self> {
153        validate_profile_parts(ambient_dim, block_dim, codebook_size)?;
154        let direction_method = match block_dim {
155            2 => DirectionMethod::FibonacciSpiral,
156            3 => DirectionMethod::FibonacciSphere,
157            _ => DirectionMethod::RobertsKronecker,
158        };
159        let radius_method = if block_dim == 2 {
160            RadiusMethod::K2ClosedForm
161        } else {
162            RadiusMethod::BetaQuantile
163        };
164        let wire_index_bits = wire_index_bits(codebook_size)?;
165        let profile = Self {
166            schema_version: PROFILE_SCHEMA.into(),
167            ambient_dim: ambient_dim as u32,
168            block_dim: block_dim as u32,
169            codebook_size: codebook_size as u32,
170            paper_rate_bits_per_coord: (codebook_size as f64).log2() / block_dim as f64,
171            wire_index_bits,
172            wire_bits_per_coord: f64::from(wire_index_bits) / block_dim as f64,
173            norm_format: NormFormat::Fp16Paper,
174            rotation_seed: seed,
175            rotation_algorithm_version: ROTATION_ALGORITHM_VERSION.into(),
176            codebook_seed: seed.wrapping_add(0x9e37_79b9_7f4a_7c15),
177            codebook_version: "fib-quant:paper-core-v1".into(),
178            source_mode: SourceMode::CanonicalSphericalBeta,
179            radius_method,
180            direction_method,
181            lloyd_restarts: 4,
182            lloyd_iterations: 25,
183            training_samples: default_training_samples(codebook_size)?,
184            empty_cell_policy: EmptyCellPolicy::SplitHighestDistortion,
185            lloyd_mode: LloydMode::Refine,
186        };
187        profile.validate()?;
188        Ok(profile)
189    }
190
191    /// Validate the complete profile.
192    pub fn validate(&self) -> Result<()> {
193        if self.schema_version != PROFILE_SCHEMA {
194            return Err(FibQuantError::CorruptPayload(format!(
195                "profile schema_version {}, expected {PROFILE_SCHEMA}",
196                self.schema_version
197            )));
198        }
199        validate_profile_parts(
200            self.ambient_dim as usize,
201            self.block_dim as usize,
202            self.codebook_size as usize,
203        )?;
204        validate_resource_bounds(
205            self.ambient_dim as usize,
206            self.block_dim as usize,
207            self.codebook_size as usize,
208            self.training_samples,
209            self.wire_index_bits,
210        )?;
211        if self.norm_format != NormFormat::Fp16Paper {
212            return Err(FibQuantError::CorruptPayload(
213                "paper profile requires fp16 norm side header".into(),
214            ));
215        }
216        if self.source_mode != SourceMode::CanonicalSphericalBeta {
217            return Err(FibQuantError::CorruptPayload(
218                "paper profile requires canonical spherical-Beta source mode".into(),
219            ));
220        }
221        if self.rotation_algorithm_version != ROTATION_ALGORITHM_VERSION {
222            return Err(FibQuantError::CorruptPayload(format!(
223                "rotation_algorithm_version {}, expected {ROTATION_ALGORITHM_VERSION}",
224                self.rotation_algorithm_version
225            )));
226        }
227        let expected_bits = wire_index_bits(self.codebook_size as usize)?;
228        if self.wire_index_bits != expected_bits {
229            return Err(FibQuantError::CorruptPayload(format!(
230                "wire_index_bits {} does not match ceil(log2(N)) {expected_bits}",
231                self.wire_index_bits
232            )));
233        }
234        let k = self.block_dim as usize;
235        let expected_paper_rate = (self.codebook_size as f64).log2() / k as f64;
236        validate_rate(
237            "paper_rate_bits_per_coord",
238            self.paper_rate_bits_per_coord,
239            expected_paper_rate,
240        )?;
241        let expected_wire_rate = f64::from(self.wire_index_bits) / k as f64;
242        validate_rate(
243            "wire_bits_per_coord",
244            self.wire_bits_per_coord,
245            expected_wire_rate,
246        )?;
247        validate_method_pair(k, &self.radius_method, &self.direction_method)?;
248        if self.lloyd_restarts == 0 || self.lloyd_restarts > MAX_LLOYD_RESTARTS {
249            return Err(FibQuantError::CorruptPayload(format!(
250                "lloyd_restarts {} outside supported range 1..={MAX_LLOYD_RESTARTS}",
251                self.lloyd_restarts
252            )));
253        }
254        if self.lloyd_iterations == 0 || self.lloyd_iterations > MAX_LLOYD_ITERATIONS {
255            return Err(FibQuantError::CorruptPayload(format!(
256                "lloyd_iterations {} outside supported range 1..={MAX_LLOYD_ITERATIONS}",
257                self.lloyd_iterations
258            )));
259        }
260        if self.training_samples < self.codebook_size
261            || self.training_samples > MAX_TRAINING_SAMPLES
262        {
263            return Err(FibQuantError::CorruptPayload(format!(
264                "training_samples {} outside supported range {}..={MAX_TRAINING_SAMPLES}",
265                self.training_samples, self.codebook_size
266            )));
267        }
268        Ok(())
269    }
270
271    /// Stable digest over all explicit profile fields.
272    pub fn digest(&self) -> Result<String> {
273        self.validate()?;
274        json_digest(PROFILE_SCHEMA, self)
275    }
276
277    /// Number of `k`-blocks per vector.
278    pub fn block_count(&self) -> u32 {
279        self.ambient_dim / self.block_dim
280    }
281}
282
283/// Return the fixed wire width for one index in `[0, N)`.
284pub fn wire_index_bits(codebook_size: usize) -> Result<u8> {
285    if codebook_size < 2 {
286        return Err(FibQuantError::InvalidCodebookSize(codebook_size));
287    }
288    let bits = usize::BITS - (codebook_size - 1).leading_zeros();
289    u8::try_from(bits).map_err(|_| FibQuantError::InvalidCodebookSize(codebook_size))
290}
291
292fn validate_profile_parts(
293    ambient_dim: usize,
294    block_dim: usize,
295    codebook_size: usize,
296) -> Result<()> {
297    if ambient_dim == 0 {
298        return Err(FibQuantError::ZeroDimension);
299    }
300    if block_dim == 0 || block_dim > ambient_dim {
301        return Err(FibQuantError::InvalidBlockDim {
302            ambient_dim,
303            block_dim,
304        });
305    }
306    if ambient_dim == block_dim {
307        return Err(FibQuantError::InvalidBlockDim {
308            ambient_dim,
309            block_dim,
310        });
311    }
312    if ambient_dim % block_dim != 0 {
313        return Err(FibQuantError::DimensionNotDivisible {
314            ambient_dim,
315            block_dim,
316        });
317    }
318    if ambient_dim > MAX_AMBIENT_DIM {
319        return Err(FibQuantError::ResourceLimitExceeded(format!(
320            "ambient_dim {ambient_dim} exceeds MAX_AMBIENT_DIM {MAX_AMBIENT_DIM}"
321        )));
322    }
323    if block_dim > MAX_BLOCK_DIM {
324        return Err(FibQuantError::ResourceLimitExceeded(format!(
325            "block_dim {block_dim} exceeds MAX_BLOCK_DIM {MAX_BLOCK_DIM}"
326        )));
327    }
328    if !(2..=MAX_CODEBOOK_SIZE).contains(&codebook_size) {
329        return Err(FibQuantError::InvalidCodebookSize(codebook_size));
330    }
331    Ok(())
332}
333
334fn default_training_samples(codebook_size: usize) -> Result<u32> {
335    let samples = 30usize
336        .checked_mul(codebook_size)
337        .ok_or_else(|| FibQuantError::ResourceLimitExceeded("30 * codebook_size overflow".into()))?
338        .max(256)
339        .min(MAX_TRAINING_SAMPLES as usize);
340    u32::try_from(samples)
341        .map_err(|_| FibQuantError::ResourceLimitExceeded("training sample count overflow".into()))
342}
343
344fn checked_profile_mul(lhs: usize, rhs: usize, label: &str) -> Result<usize> {
345    lhs.checked_mul(rhs)
346        .ok_or_else(|| FibQuantError::ResourceLimitExceeded(format!("{label} overflow")))
347}
348
349fn validate_resource_bounds(
350    ambient_dim: usize,
351    block_dim: usize,
352    codebook_size: usize,
353    training_samples: u32,
354    wire_index_bits: u8,
355) -> Result<()> {
356    let rotation_values =
357        checked_profile_mul(ambient_dim, ambient_dim, "ambient_dim * ambient_dim")?;
358    if rotation_values > MAX_ROTATION_MATRIX_VALUES {
359        return Err(FibQuantError::ResourceLimitExceeded(format!(
360            "rotation matrix values {rotation_values} exceed MAX_ROTATION_MATRIX_VALUES {MAX_ROTATION_MATRIX_VALUES}"
361        )));
362    }
363
364    let codebook_values =
365        checked_profile_mul(codebook_size, block_dim, "codebook_size * block_dim")?;
366    if codebook_values > MAX_CODEBOOK_VALUES {
367        return Err(FibQuantError::ResourceLimitExceeded(format!(
368            "codebook values {codebook_values} exceed MAX_CODEBOOK_VALUES {MAX_CODEBOOK_VALUES}"
369        )));
370    }
371
372    checked_profile_mul(
373        training_samples as usize,
374        block_dim,
375        "training_samples * block_dim",
376    )?;
377
378    let block_count = ambient_dim / block_dim;
379    let packed_bits = checked_profile_mul(
380        block_count,
381        wire_index_bits as usize,
382        "block_count * wire_index_bits",
383    )?;
384    if packed_bits > MAX_PACKED_INDEX_BITS {
385        return Err(FibQuantError::ResourceLimitExceeded(format!(
386            "packed index bits {packed_bits} exceed MAX_PACKED_INDEX_BITS {MAX_PACKED_INDEX_BITS}"
387        )));
388    }
389    Ok(())
390}
391
392fn validate_rate(name: &str, actual: f64, expected: f64) -> Result<()> {
393    if !actual.is_finite() || !expected.is_finite() || (actual - expected).abs() > RATE_TOLERANCE {
394        return Err(FibQuantError::CorruptPayload(format!(
395            "{name} {actual} does not match expected {expected}"
396        )));
397    }
398    Ok(())
399}
400
401fn validate_method_pair(
402    block_dim: usize,
403    radius: &RadiusMethod,
404    direction: &DirectionMethod,
405) -> Result<()> {
406    let valid = match block_dim {
407        2 => {
408            radius == &RadiusMethod::K2ClosedForm && direction == &DirectionMethod::FibonacciSpiral
409        }
410        3 => {
411            radius == &RadiusMethod::BetaQuantile && direction == &DirectionMethod::FibonacciSphere
412        }
413        _ => {
414            radius == &RadiusMethod::BetaQuantile && direction == &DirectionMethod::RobertsKronecker
415        }
416    };
417    if valid {
418        Ok(())
419    } else {
420        Err(FibQuantError::CorruptPayload(format!(
421            "unsupported radius/direction pair for k={block_dim}: {radius:?}/{direction:?}"
422        )))
423    }
424}
fib_quant/profile.rs

fib_quant/
profile.rs
