1use crate::{check_input, Result};
4
5pub fn autocovariance(x: &[f64], max_lag: usize) -> Result<Vec<f64>> {
14 check_input(x, max_lag + 2)?;
15 let n = x.len();
16 let mean = x.iter().sum::<f64>() / n as f64;
17 let xc: Vec<f64> = x.iter().map(|v| v - mean).collect();
18 if max_lag > 64 {
19 Ok(acov_fft(&xc, max_lag))
20 } else {
21 Ok(acov_direct(&xc, max_lag))
22 }
23}
24
25fn acov_direct(xc: &[f64], max_lag: usize) -> Vec<f64> {
26 let n = xc.len();
27 let mut r = Vec::with_capacity(max_lag + 1);
28 for k in 0..=max_lag {
29 let mut s = 0.0;
30 for t in 0..n - k {
31 s += xc[t] * xc[t + k];
32 }
33 r.push(s / n as f64);
34 }
35 r
36}
37
38fn acov_fft(xc: &[f64], max_lag: usize) -> Vec<f64> {
41 let n = xc.len();
42 let m = (n + max_lag + 1).next_power_of_two();
43 let mut re = vec![0.0; m];
44 let mut im = vec![0.0; m];
45 re[..n].copy_from_slice(xc);
46 fft(&mut re, &mut im, false);
47 for i in 0..m {
48 re[i] = re[i] * re[i] + im[i] * im[i];
49 im[i] = 0.0;
50 }
51 fft(&mut re, &mut im, true);
52 let scale = 1.0 / (m as f64 * n as f64); (0..=max_lag).map(|k| re[k] * scale).collect()
54}
55
56fn fft(re: &mut [f64], im: &mut [f64], inverse: bool) {
60 let n = re.len();
61 debug_assert!(n.is_power_of_two());
62 let mut j = 0usize;
64 for i in 1..n {
65 let mut bit = n >> 1;
66 while j & bit != 0 {
67 j ^= bit;
68 bit >>= 1;
69 }
70 j |= bit;
71 if i < j {
72 re.swap(i, j);
73 im.swap(i, j);
74 }
75 }
76 let sign = if inverse { 1.0 } else { -1.0 };
77 let mut len = 2;
78 while len <= n {
79 let half = len / 2;
80 let base = sign * 2.0 * core::f64::consts::PI / len as f64;
81 let mut i = 0;
82 while i < n {
83 for k in 0..half {
84 let ang = base * k as f64;
85 let (wi, wr) = ang.sin_cos();
86 let (ur, ui) = (re[i + k], im[i + k]);
87 let (xr, xi) = (re[i + k + half], im[i + k + half]);
88 let vr = xr * wr - xi * wi;
89 let vi = xr * wi + xi * wr;
90 re[i + k] = ur + vr;
91 im[i + k] = ui + vi;
92 re[i + k + half] = ur - vr;
93 im[i + k + half] = ui - vi;
94 }
95 i += len;
96 }
97 len <<= 1;
98 }
99}
100
101pub fn inverse_acov_from_ar(phi: &[f64]) -> Vec<f64> {
108 let p = phi.len();
109 let mut vi = Vec::with_capacity(p + 1);
110 vi.push(1.0 + phi.iter().map(|v| v * v).sum::<f64>());
111 for k in 1..=p {
112 let mut s = -phi[k - 1];
113 for i in 1..=p - k {
114 s += phi[i - 1] * phi[i + k - 1];
115 }
116 vi.push(s);
117 }
118 vi
119}
120
121pub fn acf(x: &[f64], max_lag: usize) -> Result<Vec<f64>> {
123 let r = autocovariance(x, max_lag)?;
124 let r0 = r[0];
125 if r0 <= 0.0 {
126 return Err(crate::Error::Singular);
127 }
128 Ok(r.iter().map(|v| v / r0).collect())
129}
130
131pub fn bartlett_se(acf: &[f64], n: usize) -> Vec<f64> {
136 let mut out = Vec::with_capacity(acf.len());
137 let mut cum = 0.0;
138 for (k, &rk) in acf.iter().enumerate() {
139 if k == 0 {
140 out.push(0.0);
141 } else {
142 out.push(((1.0 + 2.0 * cum) / n as f64).sqrt());
143 }
144 if k >= 1 {
145 cum += rk * rk;
146 }
147 }
148 out
149}
150
151#[cfg(test)]
152mod tests {
153 use super::*;
154
155 #[test]
156 fn white_noise_acf_small() {
157 let x: Vec<f64> = (0..2000)
158 .map(|i| ((i * 7919) % 1000) as f64 / 1000.0 - 0.5)
159 .collect();
160 let r = acf(&x, 5).unwrap();
161 assert!((r[0] - 1.0).abs() < 1e-12);
162 }
163
164 #[test]
165 fn fft_path_matches_direct_sums() {
166 let mut s = 1u64;
168 let x: Vec<f64> = (0..500)
169 .map(|_| {
170 s = s
171 .wrapping_mul(6364136223846793005)
172 .wrapping_add(1442695040888963407);
173 (s >> 11) as f64 / (1u64 << 53) as f64 - 0.5
174 })
175 .collect();
176 let mean = x.iter().sum::<f64>() / x.len() as f64;
177 let xc: Vec<f64> = x.iter().map(|v| v - mean).collect();
178 let direct = acov_direct(&xc, 120);
179 let fast = acov_fft(&xc, 120);
180 for (k, (a, b)) in direct.iter().zip(&fast).enumerate() {
181 assert!(
182 (a - b).abs() < 1e-12 * direct[0].abs(),
183 "lag {k}: {a} vs {b}"
184 );
185 }
186 let public = autocovariance(&x, 120).unwrap();
188 assert!((public[120] - direct[120]).abs() < 1e-12 * direct[0].abs());
189 }
190
191 #[test]
192 fn inverse_acov_matches_cleveland_example_1() {
193 let phi = [-0.7, 0.0, 0.0, 0.0, 0.0, -0.4];
197 let vi = inverse_acov_from_ar(&phi);
198 let ri: Vec<f64> = vi.iter().map(|v| v / vi[0]).collect();
199 assert!((ri[1] - 0.7 / 1.65).abs() < 1e-12); assert!((ri[5] - 0.28 / 1.65).abs() < 1e-12); assert!((ri[6] - 0.4 / 1.65).abs() < 1e-12); for k in [2, 3, 4] {
203 assert!(ri[k].abs() < 1e-12, "ri[{k}] = {}", ri[k]);
204 }
205 }
206
207 #[test]
208 fn bartlett_lag1_is_one_over_sqrt_n() {
209 let se = bartlett_se(&[1.0, 0.5, 0.2], 400);
210 assert!((se[1] - 0.05).abs() < 1e-12);
211 assert!(se[2] > se[1]);
212 }
213}