fib_quant/
spherical_beta.rs

1use rand::Rng;
2use rand_distr::{Distribution, Gamma, StandardNormal};
3
4use crate::{beta_inv::beta_inv, FibQuantError, Result};
5
6/// Bennett-Gersho radial companding Beta shape `beta_{d,k}`.
7pub fn beta_d_k(d: usize, k: usize) -> Result<f64> {
8    validate_dk(d, k)?;
9    if k == d {
10        return Ok(1.0);
11    }
12    let dk_term = (d as f64 - k as f64 - 2.0) / 2.0;
13    let beta = (k as f64 / (k as f64 + 2.0)) * dk_term + 1.0;
14    if beta.is_finite() && beta > 0.0 {
15        Ok(beta)
16    } else {
17        Err(FibQuantError::NumericalFailure(format!(
18            "invalid beta_{{d,k}} for d={d}, k={k}: {beta}"
19        )))
20    }
21}
22
23/// Radius quantile for codeword `n` in `1..=n_total`.
24pub fn radius_quantile(d: usize, k: usize, n: usize, n_total: usize) -> Result<f64> {
25    validate_dk(d, k)?;
26    if n == 0 || n > n_total || n_total == 0 {
27        return Err(FibQuantError::InvalidCodebookSize(n_total));
28    }
29    if k == d {
30        return Ok(1.0);
31    }
32    let q = (n as f64 - 0.5) / n_total as f64;
33    if k == 2 {
34        return radius_quantile_k2_closed_form(d, q);
35    }
36    let alpha = k as f64 / 2.0;
37    let beta = beta_d_k(d, k)?;
38    Ok(beta_inv(q, alpha, beta)?.sqrt())
39}
40
41/// Paper closed-form radius path for `k = 2`.
42pub fn radius_quantile_k2_closed_form(d: usize, q: f64) -> Result<f64> {
43    if d <= 2 {
44        return Err(FibQuantError::InvalidBlockDim {
45            ambient_dim: d,
46            block_dim: 2,
47        });
48    }
49    if !(0.0..=1.0).contains(&q) || !q.is_finite() {
50        return Err(FibQuantError::NumericalFailure(format!(
51            "invalid k=2 radius quantile q={q}"
52        )));
53    }
54    Ok((1.0 - (1.0 - q).powf(4.0 / d as f64)).sqrt())
55}
56
57/// Sample the canonical spherical-Beta `k`-block source.
58pub fn sample_spherical_beta(d: usize, k: usize, rng: &mut impl Rng) -> Result<Vec<f64>> {
59    validate_dk(d, k)?;
60    if k == d {
61        return sample_unit_sphere(k, rng);
62    }
63    let r2 = sample_beta(k as f64 / 2.0, (d - k) as f64 / 2.0, rng)?;
64    let direction = sample_unit_sphere(k, rng)?;
65    let r = r2.sqrt();
66    Ok(direction.into_iter().map(|value| r * value).collect())
67}
68
69/// Sample by normalizing a Gaussian in `R^d` and projecting the first `k` coordinates.
70pub fn sample_reference_projection(d: usize, k: usize, rng: &mut impl Rng) -> Result<Vec<f64>> {
71    validate_dk(d, k)?;
72    let mut values = Vec::with_capacity(d);
73    let mut norm_sq = 0.0;
74    for _ in 0..d {
75        let value: f64 = StandardNormal.sample(rng);
76        norm_sq += value * value;
77        values.push(value);
78    }
79    if norm_sq == 0.0 || !norm_sq.is_finite() {
80        return Err(FibQuantError::NumericalFailure(
81            "zero/non-finite gaussian norm".into(),
82        ));
83    }
84    let norm = norm_sq.sqrt();
85    Ok(values
86        .into_iter()
87        .take(k)
88        .map(|value| value / norm)
89        .collect())
90}
91
92pub(crate) fn sample_unit_sphere(k: usize, rng: &mut impl Rng) -> Result<Vec<f64>> {
93    if k == 0 {
94        return Err(FibQuantError::InvalidBlockDim {
95            ambient_dim: 0,
96            block_dim: 0,
97        });
98    }
99    loop {
100        let mut values = Vec::with_capacity(k);
101        let mut norm_sq = 0.0;
102        for _ in 0..k {
103            let value: f64 = StandardNormal.sample(rng);
104            norm_sq += value * value;
105            values.push(value);
106        }
107        if norm_sq > 0.0 && norm_sq.is_finite() {
108            let norm = norm_sq.sqrt();
109            return Ok(values.into_iter().map(|value| value / norm).collect());
110        }
111    }
112}
113
114fn sample_beta(alpha: f64, beta: f64, rng: &mut impl Rng) -> Result<f64> {
115    let ga = Gamma::new(alpha, 1.0)
116        .map_err(|err| FibQuantError::NumericalFailure(format!("gamma alpha: {err}")))?;
117    let gb = Gamma::new(beta, 1.0)
118        .map_err(|err| FibQuantError::NumericalFailure(format!("gamma beta: {err}")))?;
119    let a: f64 = ga.sample(rng);
120    let b: f64 = gb.sample(rng);
121    let sum = a + b;
122    if sum <= 0.0 || !sum.is_finite() {
123        return Err(FibQuantError::NumericalFailure("beta sample sum".into()));
124    }
125    Ok(a / sum)
126}
127
128fn validate_dk(d: usize, k: usize) -> Result<()> {
129    if d == 0 {
130        return Err(FibQuantError::ZeroDimension);
131    }
132    if k == 0 || k > d {
133        return Err(FibQuantError::InvalidBlockDim {
134            ambient_dim: d,
135            block_dim: k,
136        });
137    }
138    Ok(())
139}
fib_quant/spherical_beta.rs

fib_quant/
spherical_beta.rs
