1use statrs::distribution::{ContinuousCDF, Normal};
2
3use crate::{profile::DirectionMethod, FibQuantError, Result};
4
5const GOLDEN_RATIO: f64 = 1.618_033_988_749_895;
6
7pub 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
21pub 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
37pub 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}