1use crate::accumulator::BinnedAccumulatorF64;
5
6fn binned_sum(data: &[f64]) -> f64 {
11 let mut acc = BinnedAccumulatorF64::new();
12 acc.add_slice(data);
13 acc.finalize()
14}
15
16fn binned_mean(data: &[f64]) -> f64 {
21 if data.is_empty() {
22 return 0.0;
23 }
24 binned_sum(data) / data.len() as f64
25}
26
27pub fn acf(data: &[f64], max_lag: usize) -> Vec<f64> {
37 let n = data.len();
38 if n == 0 {
39 return vec![f64::NAN; max_lag + 1];
40 }
41
42 let mean = binned_mean(data);
43
44 let centered: Vec<f64> = data.iter().map(|&x| x - mean).collect();
46
47 let sq: Vec<f64> = centered.iter().map(|&x| x * x).collect();
49 let gamma0 = binned_sum(&sq);
50
51 if gamma0 == 0.0 {
52 let mut result = vec![0.0; max_lag + 1];
54 result[0] = 1.0;
55 return result;
56 }
57
58 let mut result = Vec::with_capacity(max_lag + 1);
59 for k in 0..=max_lag {
60 if k >= n {
61 result.push(f64::NAN);
62 continue;
63 }
64 let prods: Vec<f64> = (0..n - k).map(|t| centered[t] * centered[t + k]).collect();
65 let gamma_k = binned_sum(&prods);
66 result.push(gamma_k / gamma0);
67 }
68
69 result
70}
71
72pub fn pacf(data: &[f64], max_lag: usize) -> Vec<f64> {
80 let n = data.len();
81 if n == 0 || max_lag == 0 {
82 return vec![1.0];
83 }
84
85 let r = acf(data, max_lag);
87
88 let mut result = vec![0.0; max_lag + 1];
89 result[0] = 1.0;
90
91 if max_lag >= n {
92 for i in n..=max_lag {
94 result[i] = f64::NAN;
95 }
96 }
97
98 let effective_max = max_lag.min(n - 1);
102
103 let mut phi_prev = vec![0.0; effective_max + 1];
104 phi_prev[1] = r[1];
106 result[1] = r[1];
107
108 for m in 2..=effective_max {
109 let num_terms: Vec<f64> = (1..m).map(|j| phi_prev[j] * r[m - j]).collect();
112 let den_terms: Vec<f64> = (1..m).map(|j| phi_prev[j] * r[j]).collect();
113
114 let num = r[m] - binned_sum(&num_terms);
115 let den = 1.0 - binned_sum(&den_terms);
116
117 if den.abs() < 1e-15 {
118 result[m] = f64::NAN;
119 break;
120 }
121
122 let phi_mm = num / den;
123 result[m] = phi_mm;
124
125 let mut phi_new = vec![0.0; effective_max + 1];
127 for j in 1..m {
128 phi_new[j] = phi_prev[j] - phi_mm * phi_prev[m - j];
129 }
130 phi_new[m] = phi_mm;
131 phi_prev = phi_new;
132 }
133
134 result
135}
136
137pub fn ewma(data: &[f64], alpha: f64) -> Vec<f64> {
147 if data.is_empty() {
148 return Vec::new();
149 }
150
151 let mut result = Vec::with_capacity(data.len());
152 result.push(data[0]);
153
154 for i in 1..data.len() {
155 let prev = result[i - 1];
156 result.push(alpha * data[i] + (1.0 - alpha) * prev);
157 }
158
159 result
160}
161
162pub fn ema(data: &[f64], span: usize) -> Vec<f64> {
171 let alpha = 2.0 / (span as f64 + 1.0);
172 ewma(data, alpha)
173}
174
175pub fn seasonal_decompose(
187 data: &[f64],
188 period: usize,
189 model: &str,
190) -> Result<(Vec<f64>, Vec<f64>, Vec<f64>), String> {
191 let n = data.len();
192
193 if period < 2 {
194 return Err("seasonal_decompose: period must be >= 2".into());
195 }
196 if n < 2 * period {
197 return Err(format!(
198 "seasonal_decompose: need at least {} observations for period {}, got {}",
199 2 * period,
200 period,
201 n
202 ));
203 }
204 if model != "additive" && model != "multiplicative" {
205 return Err(format!(
206 "seasonal_decompose: model must be \"additive\" or \"multiplicative\", got \"{}\"",
207 model
208 ));
209 }
210
211 let is_mult = model == "multiplicative";
212
213 if is_mult {
215 for &v in data {
216 if v <= 0.0 {
217 return Err(
218 "seasonal_decompose: multiplicative model requires all positive data".into(),
219 );
220 }
221 }
222 }
223
224 let trend = centered_moving_average(data, period);
226
227 let mut detrended = vec![f64::NAN; n];
229 for i in 0..n {
230 if trend[i].is_nan() {
231 continue;
232 }
233 if is_mult {
234 if trend[i] != 0.0 {
235 detrended[i] = data[i] / trend[i];
236 }
237 } else {
238 detrended[i] = data[i] - trend[i];
239 }
240 }
241
242 let mut seasonal = vec![0.0; n];
244 let mut period_avgs = vec![0.0; period];
245
246 for p in 0..period {
247 let mut vals = Vec::new();
248 let mut idx = p;
249 while idx < n {
250 if !detrended[idx].is_nan() {
251 vals.push(detrended[idx]);
252 }
253 idx += period;
254 }
255 if !vals.is_empty() {
256 period_avgs[p] = binned_mean(&vals);
257 }
258 }
259
260 if is_mult {
262 let avg = binned_mean(&period_avgs);
263 if avg != 0.0 {
264 for v in &mut period_avgs {
265 *v /= avg;
266 }
267 }
268 } else {
269 let avg = binned_mean(&period_avgs);
270 for v in &mut period_avgs {
271 *v -= avg;
272 }
273 }
274
275 for i in 0..n {
277 seasonal[i] = period_avgs[i % period];
278 }
279
280 let mut residual = vec![f64::NAN; n];
282 for i in 0..n {
283 if trend[i].is_nan() {
284 continue;
285 }
286 if is_mult {
287 if seasonal[i] != 0.0 {
288 residual[i] = data[i] / (trend[i] * seasonal[i]);
289 }
290 } else {
291 residual[i] = data[i] - trend[i] - seasonal[i];
292 }
293 }
294
295 Ok((trend, seasonal, residual))
296}
297
298fn centered_moving_average(data: &[f64], period: usize) -> Vec<f64> {
301 let n = data.len();
302 let mut result = vec![f64::NAN; n];
303
304 if period % 2 == 1 {
305 let half = period / 2;
307 for i in half..n.saturating_sub(half) {
308 let window = &data[i - half..=i + half];
309 result[i] = binned_mean(window);
310 }
311 } else {
312 let mut ma = vec![f64::NAN; n];
314 let half = period / 2;
315 for i in (period - 1)..n {
317 let window = &data[i + 1 - period..=i];
318 ma[i] = binned_mean(window);
319 }
320 for i in half..n.saturating_sub(half) {
322 let left_idx = i + half - 1; let right_idx = left_idx + 1;
324 if left_idx < n && right_idx < n && !ma[left_idx].is_nan() && !ma[right_idx].is_nan()
325 {
326 result[i] = (ma[left_idx] + ma[right_idx]) / 2.0;
327 }
328 }
329 }
330
331 result
332}
333
334pub fn diff(data: &[f64], periods: usize) -> Vec<f64> {
342 if periods >= data.len() {
343 return Vec::new();
344 }
345 (periods..data.len())
346 .map(|i| data[i] - data[i - periods])
347 .collect()
348}
349
350pub fn arima_diff(data: &[f64], d: usize) -> Vec<f64> {
361 let mut current = data.to_vec();
362 for _ in 0..d {
363 if current.len() <= 1 {
364 return Vec::new();
365 }
366 current = diff(¤t, 1);
367 }
368 current
369}
370
371pub fn ar_fit(data: &[f64], p: usize) -> Result<Vec<f64>, String> {
381 if p == 0 {
382 return Err("ar_fit: p must be > 0".into());
383 }
384 if data.len() <= p {
385 return Err(format!(
386 "ar_fit: need at least {} observations for AR({}), got {}",
387 p + 1,
388 p,
389 data.len()
390 ));
391 }
392
393 let r = acf(data, p);
394
395 let mut mat_data = vec![0.0f64; p * p];
397 for i in 0..p {
398 for j in 0..p {
399 let lag = if i >= j { i - j } else { j - i };
400 mat_data[i * p + j] = r[lag];
401 }
402 }
403
404 let rhs: Vec<f64> = (1..=p).map(|k| r[k]).collect();
406
407 use crate::tensor::Tensor;
408 let r_matrix =
409 Tensor::from_vec(mat_data, &[p, p]).map_err(|e| format!("ar_fit: {e}"))?;
410 let r_vec = Tensor::from_vec(rhs, &[p]).map_err(|e| format!("ar_fit: {e}"))?;
411 let phi_tensor = r_matrix.solve(&r_vec).map_err(|e| format!("ar_fit: {e}"))?;
412
413 Ok(phi_tensor.to_vec())
414}
415
416pub fn ar_forecast(coeffs: &[f64], history: &[f64], steps: usize) -> Result<Vec<f64>, String> {
426 let p = coeffs.len();
427 if p == 0 {
428 return Err("ar_forecast: need at least one coefficient".into());
429 }
430 if history.len() < p {
431 return Err(format!(
432 "ar_forecast: need at least {} history values for AR({}), got {}",
433 p,
434 p,
435 history.len()
436 ));
437 }
438
439 let mut buf: Vec<f64> = history.to_vec();
441 let mut predictions = Vec::with_capacity(steps);
442
443 for _ in 0..steps {
444 let n = buf.len();
445 let mut acc = BinnedAccumulatorF64::new();
446 for i in 0..p {
447 acc.add(coeffs[i] * buf[n - 1 - i]);
448 }
449 let val = acc.finalize();
450 predictions.push(val);
451 buf.push(val);
452 }
453
454 Ok(predictions)
455}
456
457#[cfg(test)]
462mod tests {
463 use super::*;
464
465 #[test]
468 fn test_acf_constant_series() {
469 let data = vec![5.0; 100];
470 let result = acf(&data, 5);
471 assert_eq!(result.len(), 6);
472 assert_eq!(result[0], 1.0);
473 for k in 1..=5 {
474 assert_eq!(result[k], 0.0, "ACF at lag {} should be 0 for constant series", k);
475 }
476 }
477
478 #[test]
479 fn test_acf_lag_zero_is_one() {
480 let data: Vec<f64> = (0..50).map(|i| (i as f64).sin()).collect();
481 let result = acf(&data, 10);
482 assert!((result[0] - 1.0).abs() < 1e-12);
483 }
484
485 #[test]
486 fn test_acf_sinusoidal_periodicity() {
487 let data: Vec<f64> = (0..200)
489 .map(|i| (2.0 * std::f64::consts::PI * i as f64 / 20.0).sin())
490 .collect();
491 let result = acf(&data, 25);
492 assert!(
494 result[20] > 0.9,
495 "ACF at lag=period should be high, got {}",
496 result[20]
497 );
498 assert!(
500 result[10] < -0.9,
501 "ACF at lag=period/2 should be negative, got {}",
502 result[10]
503 );
504 }
505
506 #[test]
507 fn test_acf_empty() {
508 let result = acf(&[], 3);
509 assert_eq!(result.len(), 4);
510 assert!(result[0].is_nan());
511 }
512
513 #[test]
514 fn test_acf_determinism() {
515 let data: Vec<f64> = (0..100).map(|i| (i as f64) * 0.1 + (i as f64).sin()).collect();
516 let r1 = acf(&data, 10);
517 let r2 = acf(&data, 10);
518 for (a, b) in r1.iter().zip(r2.iter()) {
519 assert_eq!(a.to_bits(), b.to_bits());
520 }
521 }
522
523 #[test]
526 fn test_pacf_lag_zero_is_one() {
527 let data: Vec<f64> = (0..100).map(|i| (i as f64) * 0.3).collect();
528 let result = pacf(&data, 5);
529 assert!((result[0] - 1.0).abs() < 1e-12);
530 }
531
532 #[test]
533 fn test_pacf_ar1_process() {
534 let mut data = vec![0.0; 500];
537 let mut rng = cjc_repro::Rng::seeded(42);
538 for t in 1..500 {
539 data[t] = 0.8 * data[t - 1] + (rng.next_f64() - 0.5) * 0.1;
540 }
541 let result = pacf(&data, 5);
542 assert!(
543 result[1].abs() > 0.5,
544 "PACF at lag 1 should be significant for AR(1), got {}",
545 result[1]
546 );
547 for k in 2..=5 {
549 assert!(
550 result[k].abs() < 0.3,
551 "PACF at lag {} should be small for AR(1), got {}",
552 k,
553 result[k]
554 );
555 }
556 }
557
558 #[test]
559 fn test_pacf_determinism() {
560 let data: Vec<f64> = (0..100).map(|i| (i as f64).cos()).collect();
561 let r1 = pacf(&data, 5);
562 let r2 = pacf(&data, 5);
563 for (a, b) in r1.iter().zip(r2.iter()) {
564 assert_eq!(a.to_bits(), b.to_bits());
565 }
566 }
567
568 #[test]
571 fn test_ewma_alpha_one_returns_original() {
572 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0];
573 let result = ewma(&data, 1.0);
574 assert_eq!(result, data);
575 }
576
577 #[test]
578 fn test_ewma_alpha_zero_returns_first() {
579 let data = vec![10.0, 20.0, 30.0, 40.0];
580 let result = ewma(&data, 0.0);
581 for &v in &result {
583 assert_eq!(v, 10.0);
584 }
585 }
586
587 #[test]
588 fn test_ewma_length() {
589 let data = vec![1.0, 2.0, 3.0];
590 let result = ewma(&data, 0.5);
591 assert_eq!(result.len(), data.len());
592 }
593
594 #[test]
595 fn test_ewma_empty() {
596 let result = ewma(&[], 0.5);
597 assert!(result.is_empty());
598 }
599
600 #[test]
601 fn test_ewma_smoothing() {
602 let data = vec![1.0, 3.0, 5.0];
604 let result = ewma(&data, 0.5);
605 assert_eq!(result[0], 1.0);
606 assert!((result[1] - 2.0).abs() < 1e-12);
607 assert!((result[2] - 3.5).abs() < 1e-12);
608 }
609
610 #[test]
613 fn test_ema_span_relationship() {
614 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0];
615 let span = 3;
616 let alpha = 2.0 / (span as f64 + 1.0); let ema_result = ema(&data, span);
618 let ewma_result = ewma(&data, alpha);
619 for (a, b) in ema_result.iter().zip(ewma_result.iter()) {
620 assert_eq!(a.to_bits(), b.to_bits());
621 }
622 }
623
624 #[test]
627 fn test_diff_first_differences() {
628 let data = vec![1.0, 3.0, 6.0, 10.0, 15.0];
629 let result = diff(&data, 1);
630 assert_eq!(result, vec![2.0, 3.0, 4.0, 5.0]);
631 }
632
633 #[test]
634 fn test_diff_periods_two() {
635 let data = vec![1.0, 2.0, 4.0, 7.0, 11.0];
636 let result = diff(&data, 2);
637 assert_eq!(result, vec![3.0, 5.0, 7.0]);
639 }
640
641 #[test]
642 fn test_diff_empty_when_periods_too_large() {
643 let data = vec![1.0, 2.0];
644 let result = diff(&data, 5);
645 assert!(result.is_empty());
646 }
647
648 #[test]
649 fn test_diff_length() {
650 let data = vec![1.0; 10];
651 let result = diff(&data, 3);
652 assert_eq!(result.len(), 7);
653 }
654
655 #[test]
658 fn test_seasonal_decompose_additive_reconstruction() {
659 let period = 4;
661 let n = 40;
662 let mut data = vec![0.0; n];
663 for i in 0..n {
664 let trend = 10.0 + 0.5 * i as f64;
665 let seasonal = [2.0, -1.0, 0.5, -1.5][i % period];
666 data[i] = trend + seasonal;
667 }
668
669 let (trend, seasonal, residual) =
670 seasonal_decompose(&data, period, "additive").unwrap();
671
672 for i in 0..n {
674 if trend[i].is_nan() || residual[i].is_nan() {
675 continue;
676 }
677 let reconstructed = trend[i] + seasonal[i] + residual[i];
678 assert!(
679 (reconstructed - data[i]).abs() < 1e-10,
680 "Reconstruction failed at i={}: {} vs {}",
681 i,
682 reconstructed,
683 data[i]
684 );
685 }
686 }
687
688 #[test]
689 fn test_seasonal_decompose_multiplicative_reconstruction() {
690 let period = 4;
691 let n = 40;
692 let mut data = vec![0.0; n];
693 for i in 0..n {
694 let trend = 100.0 + 2.0 * i as f64;
695 let seasonal = [1.1, 0.9, 1.05, 0.95][i % period];
696 data[i] = trend * seasonal;
697 }
698
699 let (trend, seasonal, residual) =
700 seasonal_decompose(&data, period, "multiplicative").unwrap();
701
702 for i in 0..n {
703 if trend[i].is_nan() || residual[i].is_nan() {
704 continue;
705 }
706 let reconstructed = trend[i] * seasonal[i] * residual[i];
707 assert!(
708 (reconstructed - data[i]).abs() < 1e-6,
709 "Multiplicative reconstruction failed at i={}: {} vs {}",
710 i,
711 reconstructed,
712 data[i]
713 );
714 }
715 }
716
717 #[test]
718 fn test_seasonal_decompose_invalid_period() {
719 let data = vec![1.0; 20];
720 assert!(seasonal_decompose(&data, 1, "additive").is_err());
721 }
722
723 #[test]
724 fn test_seasonal_decompose_too_short() {
725 let data = vec![1.0; 5];
726 assert!(seasonal_decompose(&data, 4, "additive").is_err());
727 }
728
729 #[test]
730 fn test_seasonal_decompose_invalid_model() {
731 let data = vec![1.0; 20];
732 assert!(seasonal_decompose(&data, 4, "invalid").is_err());
733 }
734
735 #[test]
736 fn test_seasonal_decompose_multiplicative_negative_data() {
737 let data = vec![1.0, -1.0, 2.0, 3.0, 1.0, -1.0, 2.0, 3.0];
738 assert!(seasonal_decompose(&data, 4, "multiplicative").is_err());
739 }
740
741 #[test]
742 fn test_seasonal_decompose_seasonal_component_sums_to_zero() {
743 let period = 4;
744 let n = 40;
745 let mut data = vec![0.0; n];
746 for i in 0..n {
747 data[i] = 10.0 + 0.5 * i as f64 + [2.0, -1.0, 0.5, -1.5][i % period];
748 }
749
750 let (_, seasonal, _) = seasonal_decompose(&data, period, "additive").unwrap();
751
752 let one_period: Vec<f64> = (0..period).map(|i| seasonal[i]).collect();
754 let period_sum = binned_sum(&one_period);
755 assert!(
756 period_sum.abs() < 1e-10,
757 "Seasonal component should sum to ~0 over one period, got {}",
758 period_sum
759 );
760 }
761
762 #[test]
763 fn test_seasonal_decompose_determinism() {
764 let period = 4;
765 let n = 40;
766 let data: Vec<f64> = (0..n).map(|i| 10.0 + (i as f64).sin()).collect();
767
768 let (t1, s1, r1) = seasonal_decompose(&data, period, "additive").unwrap();
769 let (t2, s2, r2) = seasonal_decompose(&data, period, "additive").unwrap();
770
771 for i in 0..n {
772 assert_eq!(t1[i].to_bits(), t2[i].to_bits());
773 assert_eq!(s1[i].to_bits(), s2[i].to_bits());
774 assert_eq!(r1[i].to_bits(), r2[i].to_bits());
775 }
776 }
777}