Skip to main content

fib_quant/
directions.rs

1use statrs::distribution::{ContinuousCDF, Normal};
2
3use crate::{profile::DirectionMethod, FibQuantError, Result};
4
5const GOLDEN_RATIO: f64 = 1.618_033_988_749_895;
6
7/// Planar Fibonacci spiral directions on `S^1`.
8pub fn fibonacci_spiral_2d(n: usize) -> Result<Vec<[f64; 2]>> {
9    if n == 0 {
10        return Err(FibQuantError::InvalidCodebookSize(n));
11    }
12    let theta_g = 1.0 - 1.0 / GOLDEN_RATIO;
13    Ok((0..n)
14        .map(|idx| {
15            let theta = 2.0 * std::f64::consts::PI * idx as f64 * theta_g;
16            [theta.cos(), theta.sin()]
17        })
18        .collect())
19}
20
21/// Fibonacci sphere directions on `S^2`.
22pub fn fibonacci_sphere_3d(n: usize) -> Result<Vec<[f64; 3]>> {
23    if n == 0 {
24        return Err(FibQuantError::InvalidCodebookSize(n));
25    }
26    let theta_g = 1.0 - 1.0 / GOLDEN_RATIO;
27    Ok((0..n)
28        .map(|idx| {
29            let z = 1.0 - (2.0 * (idx + 1) as f64 - 1.0) / n as f64;
30            let theta = 2.0 * std::f64::consts::PI * idx as f64 * theta_g;
31            let r = (1.0 - z * z).max(0.0).sqrt();
32            [r * theta.cos(), r * theta.sin(), z]
33        })
34        .collect())
35}
36
37/// Roberts-Kronecker sequence mapped through inverse normal and projected to `S^{k-1}`.
38pub fn roberts_kronecker(k: usize, n: usize) -> Result<Vec<Vec<f64>>> {
39    if k < 4 {
40        return Err(FibQuantError::InvalidBlockDim {
41            ambient_dim: k + 1,
42            block_dim: k,
43        });
44    }
45    if n == 0 {
46        return Err(FibQuantError::InvalidCodebookSize(n));
47    }
48    let phi = roberts_phi(k)?;
49    let normal = Normal::new(0.0, 1.0)
50        .map_err(|err| FibQuantError::NumericalFailure(format!("normal: {err}")))?;
51    let mut dirs = Vec::with_capacity(n);
52    for idx in 0..n {
53        let mut values = Vec::with_capacity(k);
54        let mut norm_sq = 0.0;
55        for j in 1..=k {
56            let step = phi.powi(-(j as i32));
57            let frac = (((idx as f64 + 0.5) * step).fract()).clamp(1.0e-12, 1.0 - 1.0e-12);
58            let value = normal.inverse_cdf(frac);
59            norm_sq += value * value;
60            values.push(value);
61        }
62        if norm_sq <= 0.0 || !norm_sq.is_finite() {
63            return Err(FibQuantError::NumericalFailure(
64                "Roberts-Kronecker projected zero/non-finite vector".into(),
65            ));
66        }
67        let norm = norm_sq.sqrt();
68        dirs.push(values.into_iter().map(|value| value / norm).collect());
69    }
70    Ok(dirs)
71}
72
73pub(crate) fn directions_for_method(
74    k: usize,
75    n: usize,
76    method: &DirectionMethod,
77) -> Result<Vec<Vec<f64>>> {
78    match k {
79        2 if method == &DirectionMethod::FibonacciSpiral => Ok(fibonacci_spiral_2d(n)?
80            .into_iter()
81            .map(|v| vec![v[0], v[1]])
82            .collect()),
83        3 if method == &DirectionMethod::FibonacciSphere => Ok(fibonacci_sphere_3d(n)?
84            .into_iter()
85            .map(|v| vec![v[0], v[1], v[2]])
86            .collect()),
87        _ if k >= 4 && method == &DirectionMethod::RobertsKronecker => roberts_kronecker(k, n),
88        _ => Err(FibQuantError::CorruptPayload(format!(
89            "direction method {method:?} is not supported for k={k}"
90        ))),
91    }
92}
93
94fn roberts_phi(k: usize) -> Result<f64> {
95    let mut lo = 1.0f64;
96    let mut hi = 2.0f64;
97    for _ in 0..128 {
98        let mid = (lo + hi) / 2.0;
99        let value = mid.powi((k + 1) as i32) - mid - 1.0;
100        if value > 0.0 {
101            hi = mid;
102        } else {
103            lo = mid;
104        }
105    }
106    let phi = (lo + hi) / 2.0;
107    if phi.is_finite() && phi > 1.0 {
108        Ok(phi)
109    } else {
110        Err(FibQuantError::NumericalFailure(format!(
111            "invalid Roberts phi for k={k}"
112        )))
113    }
114}