1use crate::forecast::error::ForecastError;
51use crate::forecast::result::{ForecastResult, ModelParams, ModelType};
52use crate::forecast::traits::ForecastModel;
53use crate::forecast::util::{median_interval, validate_input};
54
55pub(crate) fn difference(values: &[f64], d: usize) -> Vec<f64> {
63 let mut current = values.to_vec();
64 for _ in 0..d {
65 if current.len() < 2 {
66 return Vec::new();
67 }
68 current = current.windows(2).map(|w| w[1] - w[0]).collect();
69 }
70 current
71}
72
73pub(crate) fn seasonal_difference(values: &[f64], m: usize, d: usize) -> Vec<f64> {
75 if m == 0 {
76 return values.to_vec();
77 }
78 let mut current = values.to_vec();
79 for _ in 0..d {
80 if current.len() <= m {
81 return Vec::new();
82 }
83 current = (m..current.len())
84 .map(|i| current[i] - current[i - m])
85 .collect();
86 }
87 current
88}
89
90fn undifference_lag(predictions: &[f64], anchor: &[f64], m: usize) -> Vec<f64> {
101 let mut buf: Vec<f64> = anchor.to_vec();
102 let mut comp = vec![0.0_f64; m.max(1)];
103 for (h, &v) in predictions.iter().enumerate() {
104 let base = buf[buf.len() - m];
105 let slot = h % m.max(1);
106 let y = v - comp[slot];
107 let t = base + y;
108 comp[slot] = (t - base) - y;
109 buf.push(t);
110 }
111 buf[anchor.len()..].to_vec()
112}
113
114pub(crate) fn estimate_ar_burg(
128 values: &[f64],
129 p: usize,
130) -> Result<(Vec<f64>, usize), ForecastError> {
131 if values.len() < p + 1 {
132 return Err(ForecastError::InsufficientData {
133 min: p + 1,
134 got: values.len(),
135 });
136 }
137 if p == 0 {
138 return Ok((Vec::new(), 0));
139 }
140
141 let n = values.len();
142 let mut ef: Vec<f64> = values.to_vec();
143 let mut eb: Vec<f64> = values.to_vec();
144 let mut ar = vec![0.0; p];
145 let mut effective_order = p;
146
147 for m in 0..p {
148 let mut num = 0.0;
149 let mut den = 0.0;
150 for t in (m + 1)..n {
151 num += ef[t] * eb[t - 1];
152 den += ef[t] * ef[t] + eb[t - 1] * eb[t - 1];
153 }
154 if den.abs() < 1e-15 {
155 effective_order = m;
158 tracing::warn!(
159 requested_order = p,
160 effective_order = m,
161 "Burg AR estimation terminated early due to numerical instability"
162 );
163 break;
164 }
165 let k = 2.0 * num / den;
166
167 let mut ar_new = vec![0.0; p];
168 for j in 0..m {
169 ar_new[j] = ar[j] - k * ar[m - 1 - j];
170 }
171 ar_new[m] = k;
172 ar = ar_new;
173
174 let ef_prev = ef.clone();
175 let eb_prev = eb.clone();
176 for t in (m + 1)..n {
177 ef[t] = ef_prev[t] - k * eb_prev[t - 1];
178 eb[t] = eb_prev[t - 1] - k * ef_prev[t];
179 }
180 }
181
182 Ok((ar, effective_order))
183}
184
185fn poly_mul(a: &[f64], b: &[f64]) -> Vec<f64> {
187 let mut out = vec![0.0; a.len() + b.len() - 1];
188 for (i, &ai) in a.iter().enumerate() {
189 if ai == 0.0 {
190 continue;
191 }
192 for (j, &bj) in b.iter().enumerate() {
193 out[i + j] += ai * bj;
194 }
195 }
196 out
197}
198
199pub(crate) fn expand_ar(phi: &[f64], seasonal_phi: &[f64], m: usize) -> Vec<f64> {
205 if seasonal_phi.is_empty() || m == 0 {
206 return phi.to_vec();
207 }
208 let mut a = vec![0.0; phi.len() + 1];
209 a[0] = 1.0;
210 for (i, &c) in phi.iter().enumerate() {
211 a[i + 1] = -c;
212 }
213 let mut s = vec![0.0; seasonal_phi.len() * m + 1];
214 s[0] = 1.0;
215 for (k, &c) in seasonal_phi.iter().enumerate() {
216 s[(k + 1) * m] = -c;
217 }
218 poly_mul(&a, &s)[1..].iter().map(|c| -c).collect()
219}
220
221pub(crate) fn expand_ma(theta: &[f64], seasonal_theta: &[f64], m: usize) -> Vec<f64> {
228 if seasonal_theta.is_empty() || m == 0 {
229 return theta.to_vec();
230 }
231 let mut b = vec![0.0; theta.len() + 1];
232 b[0] = 1.0;
233 b[1..].copy_from_slice(theta);
234 let mut t = vec![0.0; seasonal_theta.len() * m + 1];
235 t[0] = 1.0;
236 for (k, &c) in seasonal_theta.iter().enumerate() {
237 t[(k + 1) * m] = c;
238 }
239 poly_mul(&b, &t)[1..].to_vec()
240}
241
242pub(crate) fn innovations(values: &[f64], ar: &[f64], ma: &[f64]) -> Vec<f64> {
249 let mut eps = vec![0.0; values.len()];
250 for t in 0..values.len() {
251 let mut pred = 0.0;
252 for (j, &c) in ar.iter().enumerate() {
253 if t > j {
254 pred += c * values[t - j - 1];
255 }
256 }
257 for (j, &c) in ma.iter().enumerate() {
258 if t > j {
259 pred += c * eps[t - j - 1];
260 }
261 }
262 eps[t] = values[t] - pred;
263 }
264 eps
265}
266
267#[inline]
281pub(crate) fn includes_constant(d: usize, seasonal_d: usize) -> bool {
282 d + seasonal_d <= 1
283}
284
285#[inline]
287pub(crate) fn conditioning_warmup(ar_len: usize, ma_len: usize) -> usize {
288 ar_len.max(ma_len)
289}
290
291fn constrain_stationary(unconstrained: &[f64]) -> Vec<f64> {
307 let p = unconstrained.len();
308 let mut phi = vec![0.0; p];
309 let mut next = vec![0.0; p];
310 for k in 0..p {
311 let r = unconstrained[k].tanh();
312 next[k] = r;
313 for j in 0..k {
314 next[j] = phi[j] - r * phi[k - 1 - j];
315 }
316 phi[..=k].copy_from_slice(&next[..=k]);
317 }
318 phi
319}
320
321fn unconstrain_stationary(phi: &[f64]) -> Vec<f64> {
328 let p = phi.len();
329 let mut cur = phi.to_vec();
330 let mut out = vec![0.0; p];
331 for k in (0..p).rev() {
332 let r = cur[k].clamp(-0.999_999, 0.999_999);
333 out[k] = r.atanh();
334 if k == 0 {
335 break;
336 }
337 let denom = 1.0 - r * r;
338 let prev: Vec<f64> = (0..k)
339 .map(|j| r.mul_add(cur[k - 1 - j], cur[j]) / denom)
340 .collect();
341 cur[..k].copy_from_slice(&prev);
342 }
343 out
344}
345
346fn constrain_invertible(unconstrained: &[f64]) -> Vec<f64> {
353 constrain_stationary(unconstrained)
354 .iter()
355 .map(|v| -v)
356 .collect()
357}
358
359fn minimise_css<F>(w: &[f64], x0: &[f64], expand: F) -> Vec<f64>
379where
380 F: Fn(&[f64]) -> (Vec<f64>, Vec<f64>),
381{
382 let n_params = x0.len();
383 if n_params == 0 {
384 return Vec::new();
385 }
386 let objective = |x: &[f64]| {
387 let (ar, ma) = expand(x);
388 css(w, &ar, &ma)
389 };
390 if n_params == 1 {
396 let opt =
397 crate::forecast::optimizer::minimize_brent(|v| objective(&[v]), -12.0, 12.0, 1e-9, 300);
398 vec![opt.params[0]]
399 } else {
400 let bounds: Vec<crate::forecast::optimizer::Bound> = (0..n_params)
401 .map(|_| crate::forecast::optimizer::Bound::new(-12.0, 12.0))
402 .collect();
403 crate::forecast::optimizer::minimize_nelder_mead(objective, x0, &bounds, 2000, 1e-10).params
404 }
405}
406
407fn css(values: &[f64], ar: &[f64], ma: &[f64]) -> f64 {
408 let warmup = conditioning_warmup(ar.len(), ma.len());
409 if values.len() <= warmup {
410 return f64::INFINITY;
411 }
412 let eps = innovations(values, ar, ma);
413 eps[warmup..].iter().map(|e| e * e).sum()
414}
415
416pub(crate) fn psi_weights(ar: &[f64], ma: &[f64], horizon: usize) -> Vec<f64> {
421 let mut psi = Vec::with_capacity(horizon.max(1));
422 psi.push(1.0);
423 for j in 1..horizon {
424 let mut val = 0.0;
425 for (i, &c) in ar.iter().enumerate() {
426 if j > i {
427 val += c * psi[j - i - 1];
428 }
429 }
430 if j <= ma.len() {
431 val += ma[j - 1];
432 }
433 psi.push(val);
434 }
435 psi
436}
437
438fn arma_forecast(
442 values: &[f64],
443 residuals: &[f64],
444 ar: &[f64],
445 ma: &[f64],
446 horizon: usize,
447) -> Vec<f64> {
448 let n = values.len();
449 let mut extended = values.to_vec();
450 let mut eps = residuals.to_vec();
451 if eps.len() < n {
457 let mut padded = vec![0.0; n - eps.len()];
458 padded.extend_from_slice(&eps);
459 eps = padded;
460 } else if eps.len() > n {
461 eps.drain(..eps.len() - n);
462 }
463
464 for _ in 0..horizon {
465 let idx = extended.len();
466 let mut pred = 0.0;
467 for (j, &c) in ar.iter().enumerate() {
468 if idx > j {
469 pred += c * extended[idx - j - 1];
470 }
471 }
472 for (j, &c) in ma.iter().enumerate() {
473 if idx > j {
474 pred += c * eps[idx - j - 1];
475 }
476 }
477 extended.push(pred);
478 eps.push(0.0);
479 }
480 extended[n..].to_vec()
481}
482
483fn residual_std(residuals: &[f64], warmup: usize, n_params: usize) -> f64 {
487 let usable = residuals.get(warmup..).unwrap_or(&[]);
488 if usable.len() < 2 {
489 return 0.0;
490 }
491 let mean: f64 = usable.iter().sum::<f64>() / usable.len() as f64;
492 let dof = usable.len().saturating_sub(n_params).max(1);
493 let var: f64 = usable.iter().map(|&r| (r - mean).powi(2)).sum::<f64>() / dof as f64;
494 var.sqrt()
495}
496
497pub(crate) fn integrated_ar(ar: &[f64], d: usize, seasonal_d: usize, m: usize) -> Vec<f64> {
510 let mut poly = Vec::with_capacity(ar.len() + 1);
512 poly.push(1.0);
513 poly.extend(ar.iter().map(|c| -c));
514 for _ in 0..d {
515 poly = poly_mul(&poly, &[1.0, -1.0]);
516 }
517 if m > 0 && seasonal_d > 0 {
518 let mut sd = vec![0.0; m + 1];
519 sd[0] = 1.0;
520 sd[m] = -1.0;
521 for _ in 0..seasonal_d {
522 poly = poly_mul(&poly, &sd);
523 }
524 }
525 poly[1..].iter().map(|c| -c).collect()
526}
527
528fn prediction_intervals(
531 values: &[f64],
532 ar: &[f64],
533 ma: &[f64],
534 sigma: f64,
535 horizon: usize,
536 d: usize,
537 seasonal_d: usize,
538 m: usize,
539) -> (Vec<f64>, Vec<f64>) {
540 const Z95: f64 = 1.959_963_984_540_054;
541 let phi_star = integrated_ar(ar, d, seasonal_d, m);
542 let psi = psi_weights(&phi_star, ma, horizon);
543 let mut lower = Vec::with_capacity(horizon);
544 let mut upper = Vec::with_capacity(horizon);
545 for h in 1..=horizon {
546 let scale = psi[..h].iter().map(|w| w * w).sum::<f64>().sqrt();
547 let width = Z95 * sigma * scale;
548 lower.push(values[h - 1] - width);
549 upper.push(values[h - 1] + width);
550 }
551 (lower, upper)
552}
553
554fn forecast_timestamps(last_ts: i64, interval_ns: i64, horizon: usize) -> Vec<i64> {
557 (1..=horizon as i64)
558 .map(|h| last_ts.saturating_add(h.saturating_mul(interval_ns)))
559 .collect()
560}
561
562#[derive(serde::Serialize, serde::Deserialize)]
567pub struct ArimaModel {
568 p: usize,
569 d: usize,
570 q: usize,
571 params: ModelParams,
572 history: Vec<f64>,
573 residuals: Vec<f64>,
574 constant: f64,
577 warmup: usize,
580 last_ts: i64,
581 interval_ns: i64,
582 fitted: bool,
583}
584
585impl ArimaModel {
586 #[must_use]
588 pub fn new(p: usize, d: usize, q: usize) -> Self {
589 Self {
590 p,
591 d,
592 q,
593 params: ModelParams::Arima {
594 p,
595 d,
596 q,
597 ar_coeffs: Vec::new(),
598 ma_coeffs: Vec::new(),
599 constant: 0.0,
600 residual_std: 0.0,
601 effective_ar_order: None,
602 },
603 history: Vec::new(),
604 residuals: Vec::new(),
605 constant: 0.0,
606 warmup: 0,
607 last_ts: 0,
608 interval_ns: 0,
609 fitted: false,
610 }
611 }
612
613 fn n_params(&self) -> usize {
618 self.p + self.q + usize::from(includes_constant(self.d, 0)) + 1
619 }
620
621 fn usable_rss(&self) -> Option<(f64, usize)> {
624 let warmup = self.warmup.min(self.residuals.len());
625 let usable = &self.residuals[warmup..];
626 if usable.is_empty() {
627 return None;
628 }
629 Some((usable.iter().map(|r| r * r).sum(), usable.len()))
630 }
631
632 #[must_use]
641 pub fn aic(&self) -> Option<f64> {
642 if !self.fitted {
643 return None;
644 }
645 let (rss, n) = self.usable_rss()?;
646 crate::forecast::diagnostics::aic(rss, n, self.n_params())
647 }
648
649 #[must_use]
651 pub fn bic(&self) -> Option<f64> {
652 if !self.fitted {
653 return None;
654 }
655 let (rss, n) = self.usable_rss()?;
656 crate::forecast::diagnostics::bic(rss, n, self.n_params())
657 }
658
659 #[must_use]
661 pub fn aicc(&self) -> Option<f64> {
662 if !self.fitted {
663 return None;
664 }
665 let (rss, n) = self.usable_rss()?;
666 crate::forecast::diagnostics::aicc(rss, n, self.n_params())
667 }
668
669 #[must_use]
674 pub fn diagnostics(
675 &self,
676 actual: &[f64],
677 predicted: &[f64],
678 ) -> Option<crate::forecast::diagnostics::ModelDiagnostics> {
679 if !self.fitted {
680 return None;
681 }
682 crate::forecast::diagnostics::compute_diagnostics(actual, predicted, self.n_params(), None)
683 }
684}
685
686impl ForecastModel for ArimaModel {
687 #[tracing::instrument(skip_all, level = "debug")]
688 fn fit(&mut self, timestamps: &[i64], values: &[f64]) -> Result<(), ForecastError> {
689 let _start = std::time::Instant::now();
690 validate_input(timestamps, values)?;
691 let min_required = self.p.max(self.q) + 2 * self.d + 10;
692 if values.len() < min_required {
693 return Err(ForecastError::InsufficientData {
694 min: min_required,
695 got: values.len(),
696 });
697 }
698
699 self.history = values.to_vec();
700 let diffed = difference(values, self.d);
701 if diffed.is_empty() {
702 return Err(ForecastError::InsufficientData {
703 min: min_required,
704 got: values.len(),
705 });
706 }
707 self.constant = if includes_constant(self.d, 0) {
708 diffed.iter().sum::<f64>() / diffed.len() as f64
709 } else {
710 0.0
711 };
712 let centered: Vec<f64> = diffed.iter().map(|v| v - self.constant).collect();
713
714 let (seed_ar, effective_ar_order) = estimate_ar_burg(¢ered, self.p)?;
718 let ar_len = seed_ar.len();
719 let mut x0 = vec![0.0; ar_len + self.q];
720 x0[..ar_len].copy_from_slice(&unconstrain_stationary(&seed_ar));
721 let best = minimise_css(¢ered, &x0, |x| {
722 (
723 constrain_stationary(&x[..ar_len]),
724 constrain_invertible(&x[ar_len..]),
725 )
726 });
727 let ar_coeffs = constrain_stationary(&best[..ar_len]);
728 let ma_coeffs = constrain_invertible(&best[ar_len..]);
729 let residuals = innovations(¢ered, &ar_coeffs, &ma_coeffs);
730
731 metrics::gauge!("chronix_forecast_arima_effective_order",
733 "requested" => format!("{}", self.p))
734 .set(effective_ar_order as f64);
735
736 self.warmup = conditioning_warmup(ar_coeffs.len(), ma_coeffs.len());
737 let sigma = residual_std(&residuals, self.warmup, self.n_params());
738
739 self.residuals = residuals;
740 self.interval_ns = median_interval(timestamps);
741 self.last_ts = timestamps[timestamps.len() - 1];
742
743 self.params = ModelParams::Arima {
744 p: self.p,
745 d: self.d,
746 q: self.q,
747 ar_coeffs,
748 ma_coeffs,
749 constant: self.constant,
750 residual_std: sigma,
751 effective_ar_order: (effective_ar_order < self.p).then_some(effective_ar_order),
752 };
753 self.fitted = true;
754 metrics::histogram!("chronix_forecast_fit_duration_seconds", "model_type" => "arima")
755 .record(_start.elapsed().as_secs_f64());
756 Ok(())
757 }
758
759 #[tracing::instrument(skip_all, level = "debug")]
760 fn predict(&self, horizon: usize) -> Result<ForecastResult, ForecastError> {
761 let _start = std::time::Instant::now();
762 if !self.fitted {
763 return Err(ForecastError::NotFitted);
764 }
765 let ModelParams::Arima {
766 ar_coeffs,
767 ma_coeffs,
768 constant,
769 residual_std: sigma,
770 ..
771 } = &self.params
772 else {
773 return Err(ForecastError::InvalidInput(
774 "unexpected model params variant".into(),
775 ));
776 };
777
778 let diffed = difference(&self.history, self.d);
779 let centered: Vec<f64> = diffed.iter().map(|v| v - constant).collect();
780 let diff_preds: Vec<f64> =
781 arma_forecast(¢ered, &self.residuals, ar_coeffs, ma_coeffs, horizon)
782 .into_iter()
783 .map(|v| v + constant)
784 .collect();
785
786 let mut anchors = Vec::with_capacity(self.d);
789 let mut level = self.history.clone();
790 for _ in 0..self.d {
791 anchors.push(
792 *level
793 .last()
794 .ok_or(ForecastError::InsufficientData { min: 1, got: 0 })?,
795 );
796 level = difference(&level, 1);
797 }
798 let mut values = diff_preds;
799 for anchor in anchors.into_iter().rev() {
800 values = undifference_lag(&values, &[anchor], 1);
801 }
802
803 let (lower, upper) =
804 prediction_intervals(&values, ar_coeffs, ma_coeffs, *sigma, horizon, self.d, 0, 0);
805
806 metrics::histogram!("chronix_forecast_predict_duration_seconds", "model_type" => "arima")
807 .record(_start.elapsed().as_secs_f64());
808 Ok(ForecastResult {
809 values,
810 timestamps: forecast_timestamps(self.last_ts, self.interval_ns, horizon),
811 confidence_lower: lower,
812 confidence_upper: upper,
813 confidence_level: 0.95,
814 })
815 }
816
817 fn update(&mut self, timestamp: i64, value: f64) -> Result<(), ForecastError> {
818 if !self.fitted {
819 return Err(ForecastError::NotFitted);
820 }
821 self.history.push(value);
822 self.last_ts = timestamp;
823
824 let ModelParams::Arima {
825 ar_coeffs,
826 ma_coeffs,
827 constant,
828 ..
829 } = &self.params
830 else {
831 return Err(ForecastError::InvalidInput(
832 "unexpected model params variant".into(),
833 ));
834 };
835 let diffed: Vec<f64> = difference(&self.history, self.d)
836 .into_iter()
837 .map(|v| v - constant)
838 .collect();
839 let Some(idx) = diffed.len().checked_sub(1) else {
840 return Ok(());
841 };
842 let mut pred = 0.0;
843 for (j, &c) in ar_coeffs.iter().enumerate() {
844 if idx > j {
845 pred += c * diffed[idx - j - 1];
846 }
847 }
848 for (j, &c) in ma_coeffs.iter().enumerate() {
849 if let Some(e) = idx.checked_sub(j + 1).and_then(|k| self.residuals.get(k)) {
850 pred += c * e;
851 }
852 }
853 self.residuals.push(diffed[idx] - pred);
854
855 let max_history = (self.p + self.d + self.q + 1).max(64) * 4;
864 if self.history.len() > max_history {
865 self.history.drain(..self.history.len() - max_history);
866 }
867 let want_residuals = self.history.len().saturating_sub(self.d);
868 if self.residuals.len() > want_residuals {
869 self.residuals
870 .drain(..self.residuals.len() - want_residuals);
871 }
872 Ok(())
873 }
874
875 fn model_type(&self) -> ModelType {
876 ModelType::Arima
877 }
878
879 fn params(&self) -> &ModelParams {
880 &self.params
881 }
882}
883
884#[derive(Debug, Clone, Copy, PartialEq, Eq)]
892pub struct SarimaConfig {
893 pub p: usize,
895 pub d: usize,
897 pub q: usize,
899 pub sp: usize,
901 pub sd: usize,
903 pub sq: usize,
905 pub m: usize,
907}
908
909#[derive(serde::Serialize, serde::Deserialize)]
919pub struct SarimaModel {
920 cfg_p: usize,
921 cfg_d: usize,
922 cfg_q: usize,
923 sp: usize,
924 sd: usize,
925 sq: usize,
926 m: usize,
927 params: ModelParams,
928 history: Vec<f64>,
930 residuals: Vec<f64>,
932 ar_expanded: Vec<f64>,
934 ma_expanded: Vec<f64>,
936 constant: f64,
938 warmup: usize,
939 last_ts: i64,
940 interval_ns: i64,
941 fitted: bool,
942}
943
944impl SarimaModel {
945 #[must_use]
947 pub fn from_config(cfg: SarimaConfig) -> Self {
948 Self::new(cfg.p, cfg.d, cfg.q, cfg.sp, cfg.sd, cfg.sq, cfg.m)
949 }
950
951 #[must_use]
957 pub fn new(p: usize, d: usize, q: usize, sp: usize, sd: usize, sq: usize, m: usize) -> Self {
958 let seasonal = m > 1;
959 let (sp, sd, sq) = if seasonal { (sp, sd, sq) } else { (0, 0, 0) };
960 Self {
961 cfg_p: p,
962 cfg_d: d,
963 cfg_q: q,
964 sp,
965 sd,
966 sq,
967 m,
968 params: ModelParams::Sarima {
969 p,
970 d,
971 q,
972 sp,
973 sd,
974 sq,
975 m,
976 ar_coeffs: Vec::new(),
977 ma_coeffs: Vec::new(),
978 sar_coeffs: Vec::new(),
979 sma_coeffs: Vec::new(),
980 constant: 0.0,
981 residual_std: 0.0,
982 },
983 history: Vec::new(),
984 residuals: Vec::new(),
985 ar_expanded: Vec::new(),
986 ma_expanded: Vec::new(),
987 constant: 0.0,
988 warmup: 0,
989 last_ts: 0,
990 interval_ns: 0,
991 fitted: false,
992 }
993 }
994
995 fn transform(&self, values: &[f64]) -> Vec<f64> {
998 difference(&seasonal_difference(values, self.m, self.sd), self.cfg_d)
999 }
1000
1001 fn centered(&self, values: &[f64]) -> Vec<f64> {
1004 self.transform(values)
1005 .into_iter()
1006 .map(|v| v - self.constant)
1007 .collect()
1008 }
1009
1010 fn n_free(&self) -> usize {
1012 self.cfg_p
1013 + self.sp
1014 + self.cfg_q
1015 + self.sq
1016 + usize::from(includes_constant(self.cfg_d, self.sd))
1017 + 1
1018 }
1019
1020 fn blocks(&self, x: &[f64]) -> (Vec<f64>, Vec<f64>, Vec<f64>, Vec<f64>) {
1029 let (p, sp, q) = (self.cfg_p, self.sp, self.cfg_q);
1030 (
1031 constrain_stationary(&x[..p]),
1032 constrain_stationary(&x[p..p + sp]),
1033 constrain_invertible(&x[p + sp..p + sp + q]),
1034 constrain_invertible(&x[p + sp + q..]),
1035 )
1036 }
1037
1038 fn expand(&self, x: &[f64]) -> (Vec<f64>, Vec<f64>) {
1040 let (phi, sphi, theta, stheta) = self.blocks(x);
1041 (
1042 expand_ar(&phi, &sphi, self.m),
1043 expand_ma(&theta, &stheta, self.m),
1044 )
1045 }
1046}
1047
1048impl ForecastModel for SarimaModel {
1049 #[tracing::instrument(skip_all, level = "debug")]
1050 fn fit(&mut self, timestamps: &[i64], values: &[f64]) -> Result<(), ForecastError> {
1051 let _start = std::time::Instant::now();
1052 validate_input(timestamps, values)?;
1053
1054 let reach = (self.cfg_p + self.sp * self.m).max(self.cfg_q + self.sq * self.m);
1058 let min_required = reach + 2 * self.cfg_d + 2 * self.sd * self.m + 10;
1059 if values.len() < min_required {
1060 return Err(ForecastError::InsufficientData {
1061 min: min_required,
1062 got: values.len(),
1063 });
1064 }
1065
1066 self.history = values.to_vec();
1067 let differenced = self.transform(values);
1068 if differenced.is_empty() {
1069 return Err(ForecastError::InsufficientData {
1070 min: min_required,
1071 got: values.len(),
1072 });
1073 }
1074 self.constant = if includes_constant(self.cfg_d, self.sd) {
1075 differenced.iter().sum::<f64>() / differenced.len() as f64
1076 } else {
1077 0.0
1078 };
1079 let w: Vec<f64> = differenced.iter().map(|v| v - self.constant).collect();
1080
1081 let n_params = self.cfg_p + self.sp + self.cfg_q + self.sq;
1082 let mut x0 = vec![0.0; n_params];
1083 if self.cfg_p > 0 {
1084 let (phi, _) = estimate_ar_burg(&w, self.cfg_p)?;
1085 if phi.len() == self.cfg_p {
1086 x0[..self.cfg_p].copy_from_slice(&unconstrain_stationary(&phi));
1087 }
1088 }
1089
1090 let best = minimise_css(&w, &x0, |x| self.expand(x));
1091
1092 let (phi, sphi, theta, stheta) = self.blocks(&best);
1093 let (ar_expanded, ma_expanded) = self.expand(&best);
1094 let residuals = innovations(&w, &ar_expanded, &ma_expanded);
1095 self.warmup = conditioning_warmup(ar_expanded.len(), ma_expanded.len());
1096 let sigma = residual_std(&residuals, self.warmup, self.n_free());
1097
1098 let (p, sp, q) = (self.cfg_p, self.sp, self.cfg_q);
1099 self.params = ModelParams::Sarima {
1100 p,
1101 d: self.cfg_d,
1102 q,
1103 sp,
1104 sd: self.sd,
1105 sq: self.sq,
1106 m: self.m,
1107 ar_coeffs: phi,
1108 ma_coeffs: theta,
1109 sar_coeffs: sphi,
1110 sma_coeffs: stheta,
1111 constant: self.constant,
1112 residual_std: sigma,
1113 };
1114 self.ar_expanded = ar_expanded;
1115 self.ma_expanded = ma_expanded;
1116 self.residuals = residuals;
1117 self.interval_ns = median_interval(timestamps);
1118 self.last_ts = timestamps[timestamps.len() - 1];
1119 self.fitted = true;
1120
1121 metrics::histogram!("chronix_forecast_fit_duration_seconds", "model_type" => "sarima")
1122 .record(_start.elapsed().as_secs_f64());
1123 Ok(())
1124 }
1125
1126 #[tracing::instrument(skip_all, level = "debug")]
1127 fn predict(&self, horizon: usize) -> Result<ForecastResult, ForecastError> {
1128 let _start = std::time::Instant::now();
1129 if !self.fitted {
1130 return Err(ForecastError::NotFitted);
1131 }
1132 let ModelParams::Sarima {
1133 residual_std: sigma,
1134 ..
1135 } = &self.params
1136 else {
1137 return Err(ForecastError::InvalidInput(
1138 "unexpected model params variant".into(),
1139 ));
1140 };
1141
1142 let w = self.centered(&self.history);
1143 let mut values: Vec<f64> = arma_forecast(
1144 &w,
1145 &self.residuals,
1146 &self.ar_expanded,
1147 &self.ma_expanded,
1148 horizon,
1149 )
1150 .into_iter()
1151 .map(|v| v + self.constant)
1152 .collect();
1153
1154 let seasonal_base = seasonal_difference(&self.history, self.m, self.sd);
1158 let mut regular_levels = Vec::with_capacity(self.cfg_d);
1159 let mut level = seasonal_base.clone();
1160 for _ in 0..self.cfg_d {
1161 regular_levels.push(level.clone());
1162 level = difference(&level, 1);
1163 }
1164 for hist in regular_levels.into_iter().rev() {
1165 let anchor = *hist
1166 .last()
1167 .ok_or(ForecastError::InsufficientData { min: 1, got: 0 })?;
1168 values = undifference_lag(&values, &[anchor], 1);
1169 }
1170
1171 let mut seasonal_levels = Vec::with_capacity(self.sd);
1172 let mut level = self.history.clone();
1173 for _ in 0..self.sd {
1174 seasonal_levels.push(level.clone());
1175 level = seasonal_difference(&level, self.m, 1);
1176 }
1177 for hist in seasonal_levels.into_iter().rev() {
1178 if hist.len() < self.m {
1179 return Err(ForecastError::InsufficientData {
1180 min: self.m,
1181 got: hist.len(),
1182 });
1183 }
1184 let anchor = &hist[hist.len() - self.m..];
1185 values = undifference_lag(&values, anchor, self.m);
1186 }
1187
1188 let (lower, upper) = prediction_intervals(
1194 &values,
1195 &self.ar_expanded,
1196 &self.ma_expanded,
1197 *sigma,
1198 horizon,
1199 self.cfg_d,
1200 self.sd,
1201 self.m,
1202 );
1203
1204 metrics::histogram!("chronix_forecast_predict_duration_seconds", "model_type" => "sarima")
1205 .record(_start.elapsed().as_secs_f64());
1206 Ok(ForecastResult {
1207 values,
1208 timestamps: forecast_timestamps(self.last_ts, self.interval_ns, horizon),
1209 confidence_lower: lower,
1210 confidence_upper: upper,
1211 confidence_level: 0.95,
1212 })
1213 }
1214
1215 fn update(&mut self, timestamp: i64, value: f64) -> Result<(), ForecastError> {
1216 if !self.fitted {
1217 return Err(ForecastError::NotFitted);
1218 }
1219 self.history.push(value);
1220 self.last_ts = timestamp;
1221
1222 let w = self.centered(&self.history);
1223 let Some(idx) = w.len().checked_sub(1) else {
1224 return Ok(());
1225 };
1226 let mut pred = 0.0;
1227 for (j, &c) in self.ar_expanded.iter().enumerate() {
1228 if idx > j {
1229 pred += c * w[idx - j - 1];
1230 }
1231 }
1232 for (j, &c) in self.ma_expanded.iter().enumerate() {
1233 if let Some(e) = idx.checked_sub(j + 1).and_then(|k| self.residuals.get(k)) {
1234 pred += c * e;
1235 }
1236 }
1237 self.residuals.push(w[idx] - pred);
1238
1239 let reach = self.ar_expanded.len().max(self.ma_expanded.len());
1244 let max_history = (reach + self.cfg_d + self.sd * self.m + 1).max(64) * 4;
1245 if self.history.len() > max_history {
1246 self.history.drain(..self.history.len() - max_history);
1247 }
1248 let want_residuals = self.centered(&self.history).len();
1249 if self.residuals.len() > want_residuals {
1250 self.residuals
1251 .drain(..self.residuals.len() - want_residuals);
1252 }
1253 Ok(())
1254 }
1255
1256 fn model_type(&self) -> ModelType {
1257 ModelType::Sarima
1258 }
1259
1260 fn params(&self) -> &ModelParams {
1261 &self.params
1262 }
1263}
1264
1265pub const KPSS_CRITICAL_5PCT: f64 = 0.463;
1272
1273#[must_use]
1287pub fn kpss_statistic(values: &[f64]) -> f64 {
1288 let n = values.len();
1289 if n < 3 {
1290 return 0.0;
1291 }
1292 let n_f = n as f64;
1293 let mean = values.iter().sum::<f64>() / n_f;
1294
1295 let mut partial = 0.0;
1297 let mut sum_sq_partial = 0.0;
1298 let mut gamma0 = 0.0;
1299 let resid: Vec<f64> = values.iter().map(|v| v - mean).collect();
1300 for &e in &resid {
1301 partial += e;
1302 sum_sq_partial += partial * partial;
1303 gamma0 += e * e;
1304 }
1305 gamma0 /= n_f;
1306 if gamma0 <= 0.0 {
1307 return 0.0;
1308 }
1309
1310 #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
1312 let lag = (4.0 * (n_f / 100.0).powf(0.25)).floor() as usize;
1313 let mut long_run = gamma0;
1314 for j in 1..=lag.min(n - 1) {
1315 let gamma_j: f64 = resid[j..]
1316 .iter()
1317 .zip(&resid[..n - j])
1318 .map(|(a, b)| a * b)
1319 .sum::<f64>()
1320 / n_f;
1321 long_run += 2.0 * (1.0 - j as f64 / (lag + 1) as f64) * gamma_j;
1322 }
1323 if long_run <= 0.0 {
1324 return 0.0;
1325 }
1326
1327 sum_sq_partial / (n_f * n_f * long_run)
1328}
1329
1330#[must_use]
1342pub fn select_differencing_order(values: &[f64], max_d: usize) -> usize {
1343 let mut current = values.to_vec();
1344 let mut d = 0;
1345 while d < max_d {
1346 if current.len() < 4 || kpss_statistic(¤t) <= KPSS_CRITICAL_5PCT {
1347 break;
1348 }
1349 current = difference(¤t, 1);
1350 d += 1;
1351 }
1352 d
1353}
1354
1355pub const SEASONAL_STRENGTH_THRESHOLD: f64 = 0.64;
1364
1365#[must_use]
1389pub fn select_seasonal_differencing_order(values: &[f64], period: usize, max_d: usize) -> usize {
1390 if period < 2 || max_d == 0 {
1391 return 0;
1392 }
1393 if values.len() < 2 * period {
1396 return 0;
1397 }
1398 let first = values.first().copied().unwrap_or(0.0);
1401 if values.iter().all(|v| (v - first).abs() < f64::EPSILON) {
1402 return 0;
1403 }
1404
1405 let mut current = values.to_vec();
1406 let mut d = 0usize;
1407 while d < max_d {
1408 let Some(strength) = seasonal_strength(¤t, period) else {
1409 break;
1410 };
1411 if strength <= SEASONAL_STRENGTH_THRESHOLD {
1412 break;
1413 }
1414 current = seasonal_difference(¤t, period, 1);
1415 d += 1;
1416 if current.len() < 2 * period {
1417 break;
1418 }
1419 }
1420 d
1421}
1422
1423fn seasonal_strength(values: &[f64], period: usize) -> Option<f64> {
1426 let decomposition = crate::preprocess::decomposition::stl_decompose(
1427 values,
1428 &crate::preprocess::decomposition::StlConfig::new(period),
1429 )
1430 .ok()?;
1431 let remainder_var = variance(&decomposition.residual)?;
1432 let combined: Vec<f64> = decomposition
1433 .residual
1434 .iter()
1435 .zip(&decomposition.seasonal)
1436 .map(|(r, s)| r + s)
1437 .collect();
1438 let combined_var = variance(&combined)?;
1439 if combined_var <= 0.0 {
1440 return Some(0.0);
1441 }
1442 Some((1.0 - remainder_var / combined_var).clamp(0.0, 1.0))
1443}
1444
1445fn variance(values: &[f64]) -> Option<f64> {
1447 if values.len() < 2 {
1448 return None;
1449 }
1450 let n = values.len() as f64;
1451 let mean = values.iter().sum::<f64>() / n;
1452 Some(values.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / (n - 1.0))
1453}
1454
1455#[derive(Debug, Clone)]
1459pub struct AutoArimaResult {
1460 pub order: (usize, usize, usize),
1462 pub aic: f64,
1464 pub candidates: Vec<(usize, usize, usize, f64)>,
1466 pub evals: usize,
1468 pub kpss_order: usize,
1471}
1472
1473#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
1475pub enum SearchStrategy {
1476 #[default]
1478 Exhaustive,
1479 Stepwise,
1482}
1483
1484#[derive(Debug, Clone, Copy, Default)]
1486pub struct AutoArimaOptions {
1487 pub max_evals: Option<usize>,
1489 pub strategy: SearchStrategy,
1491}
1492
1493pub fn auto_arima(
1523 timestamps: &[i64],
1524 values: &[f64],
1525 max_p: usize,
1526 max_d: usize,
1527 max_q: usize,
1528 options: Option<&AutoArimaOptions>,
1529) -> Result<AutoArimaResult, ForecastError> {
1530 validate_input(timestamps, values)?;
1531 if values.len() < 30 {
1532 return Err(ForecastError::InsufficientData {
1533 min: 30,
1534 got: values.len(),
1535 });
1536 }
1537
1538 let kpss_order = select_differencing_order(values, max_d.max(2));
1539 let d = kpss_order.min(max_d);
1540 let opts = options.copied().unwrap_or_default();
1541 let max_evals = opts.max_evals.unwrap_or(usize::MAX);
1542
1543 let score = |p: usize, q: usize| -> Option<f64> {
1544 let min_required = p.max(q) + 2 * d + 10;
1545 if values.len() < min_required {
1546 return None;
1547 }
1548 let mut model = ArimaModel::new(p, d, q);
1549 model.fit(timestamps, values).ok()?;
1550 model.aicc().filter(|a| a.is_finite())
1551 };
1552
1553 let (order, best, mut candidates, evals) = match opts.strategy {
1554 SearchStrategy::Exhaustive => {
1555 use rayon::prelude::*;
1556 let grid: Vec<(usize, usize)> = (0..=max_p)
1557 .flat_map(|p| (0..=max_q).map(move |q| (p, q)))
1558 .filter(|&(p, q)| p + q > 0)
1559 .take(max_evals)
1560 .collect();
1561 let found: Vec<(usize, usize, usize, f64)> = grid
1562 .par_iter()
1563 .filter_map(|&(p, q)| score(p, q).map(|a| (p, d, q, a)))
1564 .collect();
1565 let evals = found.len();
1566 let best = found
1567 .iter()
1568 .copied()
1569 .min_by(|a, b| a.3.total_cmp(&b.3))
1570 .ok_or_else(|| {
1571 ForecastError::NumericalInstability(
1572 "no ARIMA model in the grid could be fitted".into(),
1573 )
1574 })?;
1575 ((best.0, best.1, best.2), best.3, found, evals)
1576 }
1577 SearchStrategy::Stepwise => stepwise(max_p, max_q, d, max_evals, &score)?,
1578 };
1579
1580 candidates.sort_by(|a, b| a.3.total_cmp(&b.3));
1581 Ok(AutoArimaResult {
1582 order,
1583 aic: best,
1584 candidates,
1585 evals,
1586 kpss_order,
1587 })
1588}
1589
1590#[allow(clippy::type_complexity)]
1592fn stepwise(
1593 max_p: usize,
1594 max_q: usize,
1595 d: usize,
1596 max_evals: usize,
1597 score: &dyn Fn(usize, usize) -> Option<f64>,
1598) -> Result<
1599 (
1600 (usize, usize, usize),
1601 f64,
1602 Vec<(usize, usize, usize, f64)>,
1603 usize,
1604 ),
1605 ForecastError,
1606> {
1607 let mut candidates: Vec<(usize, usize, usize, f64)> = Vec::new();
1608 let mut visited: std::collections::HashSet<(usize, usize)> = std::collections::HashSet::new();
1609 let mut evals = 0usize;
1610
1611 let mut best_score = f64::INFINITY;
1612 let mut best = (0usize, 0usize);
1613
1614 let mut try_model = |p: usize,
1615 q: usize,
1616 candidates: &mut Vec<(usize, usize, usize, f64)>,
1617 evals: &mut usize|
1618 -> Option<f64> {
1619 if p > max_p || q > max_q || (p == 0 && q == 0) {
1620 return None;
1621 }
1622 if !visited.insert((p, q)) {
1623 return candidates
1624 .iter()
1625 .find(|c| c.0 == p && c.2 == q)
1626 .map(|c| c.3);
1627 }
1628 if *evals >= max_evals {
1629 return None;
1630 }
1631 let s = score(p, q)?;
1632 *evals += 1;
1633 candidates.push((p, d, q, s));
1634 Some(s)
1635 };
1636
1637 for &(p, q) in &[(2, 2), (0, 0), (1, 0), (0, 1)] {
1638 let (p, q) = (p.min(max_p), q.min(max_q));
1639 if let Some(s) = try_model(p, q, &mut candidates, &mut evals)
1640 && s < best_score
1641 {
1642 best_score = s;
1643 best = (p, q);
1644 }
1645 }
1646 if best_score.is_infinite() {
1647 return Err(ForecastError::NumericalInstability(
1648 "no ARIMA model in the stepwise seed set could be fitted".into(),
1649 ));
1650 }
1651
1652 const DELTAS: [(i8, i8); 8] = [
1653 (1, 0),
1654 (-1, 0),
1655 (0, 1),
1656 (0, -1),
1657 (1, -1),
1658 (-1, 1),
1659 (1, 1),
1660 (-1, -1),
1661 ];
1662 loop {
1663 let (cp, cq) = best;
1664 let mut improved = false;
1665 for (dp, dq) in DELTAS {
1666 let (Ok(np), Ok(nq)) = (
1667 usize::try_from(cp as i64 + i64::from(dp)),
1668 usize::try_from(cq as i64 + i64::from(dq)),
1669 ) else {
1670 continue;
1671 };
1672 if let Some(s) = try_model(np, nq, &mut candidates, &mut evals)
1673 && s < best_score
1674 {
1675 best_score = s;
1676 best = (np, nq);
1677 improved = true;
1678 }
1679 }
1680 if !improved || evals >= max_evals {
1681 break;
1682 }
1683 }
1684
1685 Ok(((best.0, d, best.1), best_score, candidates, evals))
1686}
1687
1688impl crate::forecast::storage::ModelStore for ArimaModel {}
1689impl crate::forecast::storage::ModelStore for SarimaModel {}
1690
1691#[cfg(test)]
1692#[allow(clippy::unwrap_used, clippy::expect_used)]
1693mod tests {
1694 use super::*;
1695
1696 fn noise(seed: u64, n: usize, amp: f64) -> Vec<f64> {
1699 let mut s = seed;
1700 (0..n)
1701 .map(|_| {
1702 s = s
1703 .wrapping_mul(6_364_136_223_846_793_005)
1704 .wrapping_add(1_442_695_040_888_963_407);
1705 ((s >> 33) as f64 / (1u64 << 31) as f64 - 0.5) * 2.0 * amp
1706 })
1707 .collect()
1708 }
1709
1710 fn ts(n: usize) -> Vec<i64> {
1711 (0..n as i64).map(|i| i * 1_000_000_000).collect()
1712 }
1713
1714 #[test]
1717 fn difference_and_undifference_round_trip() {
1718 let series: Vec<f64> = (0..40).map(|i| (i as f64).powi(2) + 3.0).collect();
1719 for d in 1..=3 {
1720 let diffed = difference(&series, d);
1721 let split = 30;
1723 let mut preds = diffed[diffed.len() - (40 - split)..].to_vec();
1724 let mut anchors = Vec::new();
1725 let mut level = series[..split].to_vec();
1726 for _ in 0..d {
1727 anchors.push(*level.last().unwrap());
1728 level = difference(&level, 1);
1729 }
1730 for a in anchors.into_iter().rev() {
1731 preds = undifference_lag(&preds, &[a], 1);
1732 }
1733 for (i, v) in preds.iter().enumerate() {
1734 assert!(
1735 (v - series[split + i]).abs() < 1e-6,
1736 "d={d} i={i}: {v} vs {}",
1737 series[split + i]
1738 );
1739 }
1740 }
1741 }
1742
1743 #[test]
1744 fn seasonal_difference_and_undifference_round_trip() {
1745 let m = 4;
1746 let series: Vec<f64> = (0..40)
1747 .map(|i| (i % m) as f64 * 2.0 + i as f64 * 0.1)
1748 .collect();
1749 let diffed = seasonal_difference(&series, m, 1);
1750 let split = 30;
1751 let preds = diffed[diffed.len() - (40 - split)..].to_vec();
1752 let anchor = &series[split - m..split];
1753 let restored = undifference_lag(&preds, anchor, m);
1754 for (i, v) in restored.iter().enumerate() {
1755 assert!((v - series[split + i]).abs() < 1e-9, "i={i}: {v}");
1756 }
1757 }
1758
1759 #[test]
1760 fn expanded_polynomials_match_the_multiplicative_form() {
1761 let ar = expand_ar(&[0.5], &[0.3], 4);
1764 let expected = [0.5, 0.0, 0.0, 0.3, -0.15];
1765 assert_eq!(ar.len(), expected.len());
1766 for (a, b) in ar.iter().zip(expected) {
1767 assert!((a - b).abs() < 1e-12, "{ar:?}");
1768 }
1769 let ma = expand_ma(&[0.4], &[0.2], 3);
1771 let expected = [0.4, 0.0, 0.2, 0.08];
1772 for (a, b) in ma.iter().zip(expected) {
1773 assert!((a - b).abs() < 1e-12, "{ma:?}");
1774 }
1775 }
1776
1777 #[test]
1778 fn expanded_polynomials_degenerate_without_seasonal_terms() {
1779 assert_eq!(expand_ar(&[0.5, -0.2], &[], 12), vec![0.5, -0.2]);
1780 assert_eq!(expand_ma(&[0.4], &[], 12), vec![0.4]);
1781 }
1782
1783 #[test]
1784 fn psi_weights_degenerate_to_sqrt_h() {
1785 let psi = psi_weights(&[], &[], 5);
1786 assert_eq!(psi.len(), 5);
1787 assert!((psi[0] - 1.0).abs() < 1e-12);
1788 for (j, w) in psi.iter().enumerate().skip(1) {
1789 assert!(w.abs() < 1e-12, "ψ_{j} should be 0 for pure noise");
1790 }
1791
1792 let psi_ar = psi_weights(&[0.5], &[], 4);
1793 for (j, expected) in [1.0, 0.5, 0.25, 0.125].into_iter().enumerate() {
1794 assert!((psi_ar[j] - expected).abs() < 1e-12);
1795 }
1796
1797 let psi_ma = psi_weights(&[], &[0.8], 4);
1798 assert!((psi_ma[0] - 1.0).abs() < 1e-12);
1799 assert!((psi_ma[1] - 0.8).abs() < 1e-12);
1800 assert!(psi_ma[2].abs() < 1e-12);
1801 assert!(psi_ma[3].abs() < 1e-12);
1802 }
1803
1804 #[test]
1805 fn burg_recovers_a_known_ar2_and_stays_stable() {
1806 let n = 1000;
1807 let eps = noise(12345, n, 0.25);
1808 let mut vals = vec![0.0; n];
1809 for t in 2..n {
1810 vals[t] = 0.75 * vals[t - 1] - 0.5 * vals[t - 2] + eps[t];
1811 }
1812 let (coeffs, _) = estimate_ar_burg(&vals[200..], 2).unwrap();
1813 assert!((coeffs[0] - 0.75).abs() < 0.15, "{coeffs:?}");
1814 assert!((coeffs[1] + 0.5).abs() < 0.15, "{coeffs:?}");
1815 let (b, c) = (-coeffs[0], -coeffs[1]);
1817 let disc = b * b - 4.0 * c;
1818 if disc >= 0.0 {
1819 let r1 = (-b + disc.sqrt()) / 2.0;
1820 let r2 = (-b - disc.sqrt()) / 2.0;
1821 assert!(
1822 r1.abs() < 1.0 && r2.abs() < 1.0,
1823 "unstable roots {r1}, {r2}"
1824 );
1825 }
1826 }
1827
1828 #[test]
1834 fn random_walk_interval_variance_grows_exactly_linearly() {
1835 let e = noise(7, 200, 1.0);
1836 let mut vals = vec![100.0];
1837 for v in e.iter().take(199) {
1838 vals.push(vals.last().unwrap() + v);
1839 }
1840 let mut model = ArimaModel::new(0, 1, 0);
1841 model.fit(&ts(200), &vals).unwrap();
1842 let r = model.predict(12).unwrap();
1843 let half = |h: usize| r.confidence_upper[h] - r.values[h];
1844 assert!(half(0) > 0.0);
1845 for h in 1..=12 {
1846 let ratio = (half(h - 1) / half(0)).powi(2);
1847 assert!(
1848 (ratio - h as f64).abs() < 1e-9,
1849 "var({h})/var(1) = {ratio}, want {h}"
1850 );
1851 }
1852 }
1853
1854 #[test]
1857 fn arima_110_interval_matches_the_psi_closed_form() {
1858 let e = noise(11, 240, 1.0);
1859 let mut vals = vec![50.0];
1860 let mut prev_d = 0.0;
1861 for v in e.iter().take(239) {
1862 let d = 0.6 * prev_d + v;
1863 vals.push(vals.last().unwrap() + d);
1864 prev_d = d;
1865 }
1866 let mut model = ArimaModel::new(1, 1, 0);
1867 model.fit(&ts(240), &vals).unwrap();
1868 let ModelParams::Arima {
1869 ar_coeffs,
1870 residual_std,
1871 ..
1872 } = model.params()
1873 else {
1874 panic!()
1875 };
1876 let phi = ar_coeffs[0];
1877 assert!(phi.abs() > 0.1, "need a non-trivial φ, got {phi}");
1878 let r = model.predict(10).unwrap();
1879 for h in 1..=10 {
1880 let factor: f64 = (0..h)
1881 .map(|j| {
1882 let psi = (1.0 - phi.powi(j as i32 + 1)) / (1.0 - phi);
1883 psi * psi
1884 })
1885 .sum::<f64>();
1886 let expected = 1.959_963_984_540_054 * residual_std * factor.sqrt();
1887 let half = r.confidence_upper[h - 1] - r.values[h - 1];
1888 assert!(
1889 (half - expected).abs() < 1e-9,
1890 "h={h}: half-width {half} != {expected}"
1891 );
1892 }
1893 }
1894
1895 #[test]
1897 fn integrated_ar_expands_the_differencing_operator() {
1898 assert_eq!(integrated_ar(&[], 1, 0, 0), vec![1.0]);
1900 let d2 = integrated_ar(&[], 2, 0, 0);
1902 assert!((d2[0] - 2.0).abs() < 1e-12 && (d2[1] + 1.0).abs() < 1e-12);
1903 let m = integrated_ar(&[0.5], 1, 0, 0);
1905 assert!((m[0] - 1.5).abs() < 1e-12 && (m[1] + 0.5).abs() < 1e-12);
1906 let sd = integrated_ar(&[], 0, 1, 4);
1908 assert_eq!(sd.len(), 4);
1909 assert!((sd[3] - 1.0).abs() < 1e-12);
1910 assert!(sd[..3].iter().all(|c| c.abs() < 1e-12));
1911 assert_eq!(integrated_ar(&[0.3, -0.2], 0, 0, 12), vec![0.3, -0.2]);
1913 }
1914
1915 #[test]
1920 fn arma_forecast_reads_the_newest_innovations() {
1921 let values = [0.0, 0.0, 0.0];
1922 let residuals = [1.0, 2.0, 3.0, 4.0, 5.0];
1923 let out = arma_forecast(&values, &residuals, &[], &[0.5], 1);
1924 assert!((out[0] - 2.5).abs() < 1e-12, "got {}", out[0]);
1925 }
1926
1927 #[test]
1931 fn update_keeps_history_and_residuals_aligned() {
1932 let vals: Vec<f64> = (0..100)
1933 .map(|i| 10.0 + i as f64 * 0.3 + ((i * 17) % 7) as f64 * 0.4)
1934 .collect();
1935 let mut model = ArimaModel::new(1, 1, 1);
1936 model.fit(&ts(100), &vals).unwrap();
1937 for i in 0..300 {
1938 let v = 40.0 + ((i * 13) % 11) as f64 * 0.5;
1939 model.update((100 + i) as i64 * 1_000_000_000, v).unwrap();
1940 }
1941 assert!(model.history.len() < 300, "history was never trimmed");
1942 assert_eq!(
1943 model.residuals.len(),
1944 model.history.len() - model.d,
1945 "residuals must stay aligned with difference(history, d)"
1946 );
1947
1948 let ModelParams::Arima {
1952 ar_coeffs,
1953 ma_coeffs,
1954 constant,
1955 ..
1956 } = model.params().clone()
1957 else {
1958 panic!()
1959 };
1960 let diffed: Vec<f64> = difference(&model.history, 1)
1961 .into_iter()
1962 .map(|v| v - constant)
1963 .collect();
1964 let n = diffed.len();
1965 let expected_diff = constant
1966 + ar_coeffs[0] * diffed[n - 1]
1967 + ma_coeffs[0] * model.residuals[model.residuals.len() - 1];
1968 let expected = model.history[model.history.len() - 1] + expected_diff;
1969 let got = model.predict(1).unwrap().values[0];
1970 assert!((got - expected).abs() < 1e-9, "{got} != {expected}");
1971 }
1972
1973 #[test]
1975 fn sarima_update_keeps_history_and_residuals_aligned() {
1976 let vals = {
1977 (0..120)
1978 .map(|i| 50.0 + 5.0 * (std::f64::consts::TAU * (i % 12) as f64 / 12.0).sin())
1979 .collect::<Vec<f64>>()
1980 };
1981 let mut model = SarimaModel::new(1, 1, 1, 0, 1, 0, 12);
1982 model.fit(&ts(120), &vals).unwrap();
1983 for i in 0..400 {
1984 let v = 50.0 + 5.0 * (std::f64::consts::TAU * (i % 12) as f64 / 12.0).sin();
1985 model.update((120 + i) as i64 * 1_000_000_000, v).unwrap();
1986 }
1987 assert!(model.history.len() < 400, "history was never trimmed");
1988 assert_eq!(model.residuals.len(), model.centered(&model.history).len());
1989 }
1990
1991 #[test]
1992 fn arima_random_walk() {
1993 let e = noise(42, 200, 1.0);
1994 let mut vals = vec![100.0];
1995 for v in e.iter().take(199) {
1996 vals.push(vals.last().unwrap() + v);
1997 }
1998 let mut model = ArimaModel::new(1, 1, 0);
1999 model.fit(&ts(200), &vals).unwrap();
2000 let result = model.predict(10).unwrap();
2001 assert_eq!(result.values.len(), 10);
2002 assert!((result.values[0] - vals[199]).abs() < 10.0);
2003 }
2004
2005 #[test]
2006 fn arima_linear_trend() {
2007 let vals: Vec<f64> = (0..100).map(|i| f64::from(i) * 2.0 + 10.0).collect();
2008 let mut model = ArimaModel::new(0, 1, 0);
2009 model.fit(&ts(100), &vals).unwrap();
2010 let result = model.predict(5).unwrap();
2011 for (i, v) in result.values.iter().enumerate() {
2014 let expected = 208.0 + (i + 1) as f64 * 2.0;
2015 assert!((v - expected).abs() < 1e-6, "h={i}: {v} vs {expected}");
2016 }
2017 }
2018
2019 #[test]
2024 fn arima_d2_extrapolates_a_line_exactly() {
2025 let n = 60usize;
2026 let vals: Vec<f64> = (0..n).map(|i| 3.0 * i as f64 - 7.0).collect();
2027 let mut model = ArimaModel::new(0, 2, 0);
2028 model.fit(&ts(n), &vals).unwrap();
2029 let result = model.predict(5).unwrap();
2030 for (h, v) in result.values.iter().enumerate() {
2031 let expected = 3.0 * (n + h) as f64 - 7.0;
2032 assert!((v - expected).abs() < 1e-6, "h={h}: {v} vs {expected}");
2033 }
2034 }
2035
2036 #[test]
2037 fn arima_d2_undifference_correctness() {
2038 let n = 100usize;
2039 let vals: Vec<f64> = (0..n).map(|i| (i as f64).powi(2)).collect();
2040 let mut model = ArimaModel::new(0, 2, 0);
2041 model.fit(&ts(n), &vals).unwrap();
2042 let result = model.predict(5).unwrap();
2043 let slope = vals[n - 1] - vals[n - 2];
2048 for (h, v) in result.values.iter().enumerate() {
2049 let expected = vals[n - 1] + (h + 1) as f64 * slope;
2050 assert!((v - expected).abs() < 1e-6, "h={h}: {v} vs {expected}");
2051 }
2052 }
2053
2054 #[test]
2055 fn arima_confidence_intervals_bracket_the_forecast() {
2056 let vals: Vec<f64> = (0..100)
2057 .map(|i| (f64::from(i)).sin() * 5.0 + 50.0)
2058 .collect();
2059 let mut model = ArimaModel::new(2, 0, 1);
2060 model.fit(&ts(100), &vals).unwrap();
2061 let result = model.predict(10).unwrap();
2062 for i in 0..10 {
2063 assert!(result.confidence_lower[i] <= result.values[i]);
2064 assert!(result.confidence_upper[i] >= result.values[i]);
2065 }
2066 }
2067
2068 #[test]
2073 fn arima_intervals_track_the_noise_not_the_level() {
2074 let n = 200;
2075 let e = noise(7, n, 0.1);
2076 let mut vals = Vec::with_capacity(n);
2077 let mut x = 100.0_f64;
2078 for v in &e {
2079 x = 100.0 + 0.5 * (x - 100.0) + v;
2080 vals.push(x);
2081 }
2082 let mut model = ArimaModel::new(1, 0, 0);
2083 model.fit(&ts(n), &vals).unwrap();
2084 let r = model.predict(3).unwrap();
2085 let half_width = r.confidence_upper[0] - r.values[0];
2086 assert!(
2087 half_width < 0.5,
2088 "one-step 95% half-width {half_width} exceeds the noise scale (0.1)"
2089 );
2090 assert!(
2091 half_width > 0.05,
2092 "interval collapsed to nothing: {half_width}"
2093 );
2094 }
2095
2096 #[test]
2097 fn arima_insufficient_data() {
2098 let mut model = ArimaModel::new(2, 1, 1);
2099 assert!(model.fit(&[0, 1, 2], &[1.0, 2.0, 3.0]).is_err());
2100 }
2101
2102 #[test]
2103 fn arima_update() {
2104 let vals: Vec<f64> = (0..50).map(f64::from).collect();
2105 let mut model = ArimaModel::new(1, 1, 0);
2106 model.fit(&ts(50), &vals).unwrap();
2107 assert!(model.update(50_000_000_000, 50.0).is_ok());
2108 assert!(model.predict(3).unwrap().values.len() == 3);
2109 }
2110
2111 #[test]
2112 fn css_recovers_a_known_arma11() {
2113 let n = 1000;
2114 let eps = noise(99999, n, 0.2);
2115 let mut vals = vec![0.0; n];
2116 for t in 1..n {
2117 vals[t] = 0.6 * vals[t - 1] + eps[t] + 0.4 * eps[t - 1];
2118 }
2119 let vals = &vals[200..];
2120 let mut model = ArimaModel::new(1, 0, 1);
2121 model.fit(&ts(vals.len()), vals).unwrap();
2122 let ModelParams::Arima {
2123 ma_coeffs,
2124 ar_coeffs,
2125 ..
2126 } = model.params()
2127 else {
2128 panic!("wrong variant");
2129 };
2130 assert!((ma_coeffs[0] - 0.4).abs() < 0.3, "θ = {}", ma_coeffs[0]);
2131 assert!((ar_coeffs[0] - 0.6).abs() < 0.2, "φ = {}", ar_coeffs[0]);
2132 }
2133
2134 #[test]
2137 fn kpss_accepts_stationary_and_rejects_a_random_walk() {
2138 let white = noise(3, 300, 1.0);
2139 assert!(
2140 kpss_statistic(&white) < KPSS_CRITICAL_5PCT,
2141 "white noise rejected: {}",
2142 kpss_statistic(&white)
2143 );
2144
2145 let mut walk = vec![0.0];
2146 for v in white.iter().take(299) {
2147 walk.push(walk.last().unwrap() + v);
2148 }
2149 assert!(
2150 kpss_statistic(&walk) > KPSS_CRITICAL_5PCT,
2151 "random walk accepted: {}",
2152 kpss_statistic(&walk)
2153 );
2154 assert_eq!(select_differencing_order(&white, 2), 0);
2155 assert_eq!(select_differencing_order(&walk, 2), 1);
2156 }
2157
2158 #[test]
2159 fn kpss_is_scale_invariant() {
2160 let v = noise(11, 200, 1.0);
2161 let scaled: Vec<f64> = v.iter().map(|x| x * 1000.0 + 5.0).collect();
2162 let a = kpss_statistic(&v);
2163 let b = kpss_statistic(&scaled);
2164 assert!((a - b).abs() < 1e-9, "{a} vs {b}");
2165 }
2166
2167 #[test]
2168 fn kpss_handles_degenerate_input() {
2169 assert_eq!(kpss_statistic(&[]), 0.0);
2170 assert_eq!(kpss_statistic(&[1.0, 1.0]), 0.0);
2171 assert_eq!(kpss_statistic(&[7.0; 50]), 0.0);
2172 assert_eq!(select_differencing_order(&[7.0; 50], 2), 0);
2173 }
2174
2175 #[test]
2184 fn strong_seasonality_takes_a_seasonal_difference() {
2185 let m = 24usize;
2186 let values: Vec<f64> = (0..(m * 14))
2187 .map(|i| {
2188 let phase = (i % m) as f64 / m as f64 * std::f64::consts::TAU;
2189 100.0 + 20.0 * phase.sin() + ((i * 37) % 11) as f64 * 0.05
2190 })
2191 .collect();
2192 assert_eq!(select_seasonal_differencing_order(&values, m, 1), 1);
2193 }
2194
2195 #[test]
2203 fn weak_seasonality_takes_no_seasonal_difference() {
2204 let m = 24usize;
2205 let values: Vec<f64> = (0..(m * 14))
2206 .map(|i| {
2207 let phase = (i % m) as f64 / m as f64 * std::f64::consts::TAU;
2208 let noise = (((i * 2_654_435_761usize) % 1000) as f64 / 1000.0 - 0.5) * 40.0;
2209 100.0 + 0.5 * phase.sin() + noise
2210 })
2211 .collect();
2212 assert_eq!(select_seasonal_differencing_order(&values, m, 1), 0);
2213 }
2214
2215 #[test]
2216 fn pure_noise_takes_no_seasonal_difference() {
2217 let m = 24usize;
2218 let values: Vec<f64> = (0..(m * 14))
2219 .map(|i| 100.0 + (((i * 2_654_435_761usize) % 1000) as f64 / 1000.0 - 0.5) * 40.0)
2220 .collect();
2221 assert_eq!(select_seasonal_differencing_order(&values, m, 0), 0);
2222 assert_eq!(select_seasonal_differencing_order(&values, m, 1), 0);
2223 }
2224
2225 fn min_ma_root_modulus(theta: &[f64]) -> f64 {
2240 let q = theta.len();
2241 if q == 0 {
2242 return f64::INFINITY;
2243 }
2244 let mut c: Vec<(f64, f64)> = std::iter::once(1.0)
2245 .chain(theta.iter().copied())
2246 .map(|v| (v, 0.0))
2247 .collect();
2248 let lead = c[q].0;
2249 if lead.abs() < 1e-12 {
2250 return min_ma_root_modulus(&theta[..q - 1]);
2251 }
2252 for v in &mut c {
2253 v.0 /= lead;
2254 }
2255 let mul = |a: (f64, f64), b: (f64, f64)| (a.0 * b.0 - a.1 * b.1, a.0 * b.1 + a.1 * b.0);
2256 let sub = |a: (f64, f64), b: (f64, f64)| (a.0 - b.0, a.1 - b.1);
2257 let div = |a: (f64, f64), b: (f64, f64)| {
2258 let d = b.0 * b.0 + b.1 * b.1;
2259 ((a.0 * b.0 + a.1 * b.1) / d, (a.1 * b.0 - a.0 * b.1) / d)
2260 };
2261 let eval = |z: (f64, f64)| {
2262 let mut acc = (0.0, 0.0);
2263 for k in (0..=q).rev() {
2264 acc = mul(acc, z);
2265 acc = (acc.0 + c[k].0, acc.1 + c[k].1);
2266 }
2267 acc
2268 };
2269 let mut roots: Vec<(f64, f64)> = (0..q)
2270 .map(|k| {
2271 let a = 0.4 + 0.9 * k as f64;
2272 (a.cos() * 0.9, a.sin() * 0.9)
2273 })
2274 .collect();
2275 for _ in 0..500 {
2276 let mut moved = 0.0f64;
2277 for i in 0..q {
2278 let mut denom = (1.0, 0.0);
2279 for j in 0..q {
2280 if i != j {
2281 denom = mul(denom, sub(roots[i], roots[j]));
2282 }
2283 }
2284 if denom.0.abs() + denom.1.abs() < 1e-300 {
2285 continue;
2286 }
2287 let delta = div(eval(roots[i]), denom);
2288 roots[i] = sub(roots[i], delta);
2289 moved = moved.max(delta.0.abs() + delta.1.abs());
2290 }
2291 if moved < 1e-14 {
2292 break;
2293 }
2294 }
2295 roots
2296 .iter()
2297 .map(|r| (r.0 * r.0 + r.1 * r.1).sqrt())
2298 .fold(f64::INFINITY, f64::min)
2299 }
2300
2301 #[test]
2302 fn the_root_finder_agrees_with_the_quadratic_formula() {
2303 let m = min_ma_root_modulus(&[0.99, -0.99]);
2307 assert!((m - 0.6225).abs() < 1e-3, "min |root| was {m}");
2308 assert!(min_ma_root_modulus(&[0.5]) > 1.0);
2310 }
2311
2312 #[test]
2313 fn a_fitted_ma_polynomial_is_invertible() {
2314 let mut state = 7u64;
2323 let mut next = move || {
2324 state = state
2325 .wrapping_mul(6_364_136_223_846_793_005)
2326 .wrapping_add(1);
2327 ((state >> 33) as f64 / f64::from(u32::MAX >> 1)) - 1.0
2328 };
2329
2330 let mut checked = 0usize;
2331 let mut worst = f64::INFINITY;
2332 for q in 2..=3usize {
2333 for trial in 0..40 {
2334 let n = 160usize;
2335 let mut v = Vec::with_capacity(n);
2336 let (mut prev0, mut prev1, mut walk) = (0.0f64, 0.0f64, 0.0f64);
2337 for i in 0..n {
2338 let e = next();
2339 v.push(match trial % 4 {
2340 0 => e,
2341 1 => e + 0.95 * prev0 + 0.9 * prev1,
2342 2 => {
2343 walk += e;
2344 walk
2345 }
2346 _ => i as f64 * 0.05 + e,
2347 });
2348 prev1 = prev0;
2349 prev0 = e;
2350 }
2351 let ts: Vec<i64> = (0..n as i64).collect();
2352 let mut model = ArimaModel::new(1, 0, q);
2353 if model.fit(&ts, &v).is_err() {
2354 continue;
2355 }
2356 let ModelParams::Arima { ma_coeffs, .. } = model.params() else {
2357 continue;
2358 };
2359 let r = min_ma_root_modulus(ma_coeffs);
2360 worst = worst.min(r);
2361 assert!(
2362 r > 1.0,
2363 "q={q} trial={trial}: fitted θ={ma_coeffs:?} is non-invertible \
2364 (min |root| = {r})"
2365 );
2366 checked += 1;
2367 }
2368 }
2369 assert!(checked >= 60, "only {checked} fits were checked");
2370 let raw: Vec<f64> = (0..121).map(|_| next()).collect();
2374 let over: Vec<f64> = raw.windows(2).map(|w| w[1] - w[0]).collect();
2375 let ts: Vec<i64> = (0..over.len() as i64).collect();
2376 let mut model = ArimaModel::new(1, 0, 2);
2377 if model.fit(&ts, &over).is_ok()
2378 && let ModelParams::Arima { ma_coeffs, .. } = model.params()
2379 {
2380 let r = min_ma_root_modulus(ma_coeffs);
2381 assert!(r > 1.0, "over-differenced fit is non-invertible: {r}");
2382 }
2383 }
2384
2385 #[test]
2386 fn a_stochastic_seasonal_level_is_differenced() {
2387 let m = 24usize;
2399 let mut values = vec![0.0f64; 12 * m];
2400 let mut season_level = vec![0.0f64; m];
2401 let mut rng_state = 42u64;
2402 let mut next = || {
2403 rng_state = rng_state
2404 .wrapping_mul(6_364_136_223_846_793_005)
2405 .wrapping_add(1);
2406 ((rng_state >> 33) as f64 / f64::from(u32::MAX >> 1)) - 1.0
2407 };
2408 for cycle in 0..12 {
2409 for season in 0..m {
2410 season_level[season] += next() * 3.0;
2411 values[cycle * m + season] = season_level[season] + next() * 0.1;
2412 }
2413 }
2414 assert_eq!(select_seasonal_differencing_order(&values, m, 1), 1);
2415 }
2416
2417 #[test]
2418 fn a_counter_is_never_seasonally_differenced() {
2419 let m = 24usize;
2432 let n = 12 * m;
2433 let clean: Vec<f64> = (0..n).map(|i| 1000.0 + i as f64 * 37.0).collect();
2434 assert_eq!(select_seasonal_differencing_order(&clean, m, 1), 0);
2435
2436 let mut rng_state = 3u64;
2437 let mut next = || {
2438 rng_state = rng_state
2439 .wrapping_mul(6_364_136_223_846_793_005)
2440 .wrapping_add(1);
2441 ((rng_state >> 33) as f64 / f64::from(u32::MAX >> 1)) - 1.0
2442 };
2443 for jitter in [0.5f64, 15.0] {
2444 let mut acc = 1000.0;
2445 let noisy: Vec<f64> = (0..n)
2446 .map(|_| {
2447 acc += 37.0 + next() * jitter;
2448 acc
2449 })
2450 .collect();
2451 assert_eq!(
2452 select_seasonal_differencing_order(&noisy, m, 1),
2453 0,
2454 "a counter with jitter {jitter} was seasonally differenced"
2455 );
2456 }
2457 }
2458
2459 #[test]
2460 fn a_constant_series_is_never_seasonally_differenced() {
2461 let values = vec![7.0f64; 200];
2465 assert_eq!(select_seasonal_differencing_order(&values, 24, 1), 0);
2466 }
2467
2468 #[test]
2469 fn a_series_shorter_than_two_periods_is_never_seasonally_differenced() {
2470 let values: Vec<f64> = (0..30).map(|i| f64::from(i % 24)).collect();
2471 assert_eq!(
2472 select_seasonal_differencing_order(&values, 24, 1),
2473 0,
2474 "STL needs two full periods; there is nothing to measure",
2475 );
2476 }
2477
2478 #[test]
2479 fn a_period_below_two_is_not_seasonal() {
2480 let values: Vec<f64> = (0..100).map(f64::from).collect();
2481 assert_eq!(select_seasonal_differencing_order(&values, 1, 1), 0);
2482 assert_eq!(select_seasonal_differencing_order(&values, 0, 1), 0);
2483 }
2484
2485 #[test]
2486 fn select_differencing_order_respects_the_cap() {
2487 let mut quad: Vec<f64> = (0..200).map(|i| (f64::from(i)).powi(2)).collect();
2488 quad.iter_mut()
2489 .zip(noise(5, 200, 1.0))
2490 .for_each(|(v, e)| *v += e);
2491 assert_eq!(select_differencing_order(&quad, 1), 1);
2492 assert!(select_differencing_order(&quad, 3) >= 2);
2493 }
2494
2495 #[test]
2498 fn auto_arima_differences_a_trend() {
2499 let vals: Vec<f64> = (0..100).map(|i| f64::from(i) * 2.0 + 10.0).collect();
2500 let result = auto_arima(&ts(100), &vals, 2, 2, 2, None).unwrap();
2501 assert!(result.order.1 >= 1, "d should be >= 1 for trending data");
2502 assert!(!result.candidates.is_empty());
2503 assert!(result.aic.is_finite());
2504 assert!(result.evals > 0);
2505 }
2506
2507 #[test]
2511 fn auto_arima_does_not_overdifference_stationary_noise() {
2512 let vals: Vec<f64> = noise(2024, 200, 0.5).iter().map(|v| v + 10.0).collect();
2513 let result = auto_arima(&ts(200), &vals, 2, 2, 2, None).unwrap();
2514 assert_eq!(result.order.1, 0, "stationary series was differenced");
2515 assert_eq!(result.kpss_order, 0);
2516 }
2517
2518 #[test]
2519 fn auto_arima_reports_when_max_d_bound_the_answer() {
2520 let vals: Vec<f64> = (0..200)
2521 .map(|i| (f64::from(i)).powi(2))
2522 .zip(noise(6, 200, 1.0))
2523 .map(|(a, b)| a + b)
2524 .collect();
2525 let result = auto_arima(&ts(200), &vals, 1, 1, 1, None).unwrap();
2526 assert_eq!(result.order.1, 1, "capped at max_d");
2527 assert!(
2528 result.kpss_order >= 2,
2529 "kpss wanted more: {}",
2530 result.kpss_order
2531 );
2532 }
2533
2534 #[test]
2535 fn auto_arima_returns_sorted_candidates_at_one_differencing_order() {
2536 let vals: Vec<f64> = (0..80)
2537 .map(|i| (f64::from(i) * 0.1).sin() * 5.0 + 50.0)
2538 .collect();
2539 let result = auto_arima(&ts(80), &vals, 2, 1, 2, None).unwrap();
2540 for w in result.candidates.windows(2) {
2541 assert!(w[0].3 <= w[1].3, "candidates not sorted");
2542 }
2543 assert!(
2544 result.candidates.iter().all(|c| c.1 == result.order.1),
2545 "candidates must share one differencing order to be comparable"
2546 );
2547 }
2548
2549 #[test]
2550 fn auto_arima_insufficient_data() {
2551 let vals: Vec<f64> = (0..10).map(f64::from).collect();
2552 assert!(auto_arima(&ts(10), &vals, 3, 2, 3, None).is_err());
2553 }
2554
2555 #[test]
2556 fn auto_arima_max_evals_budget() {
2557 let vals: Vec<f64> = (0..100).map(|i| f64::from(i) * 2.0 + 10.0).collect();
2558 let opts = AutoArimaOptions {
2559 max_evals: Some(3),
2560 ..Default::default()
2561 };
2562 let result = auto_arima(&ts(100), &vals, 3, 2, 3, Some(&opts)).unwrap();
2563 assert!(result.evals <= 3, "got {} evals", result.evals);
2564 assert!(result.aic.is_finite());
2565 }
2566
2567 #[test]
2568 fn auto_arima_stepwise_agrees_with_exhaustive_on_a_small_grid() {
2569 let vals: Vec<f64> = (0..100).map(|i| f64::from(i) * 2.0 + 10.0).collect();
2570 let exhaustive = auto_arima(&ts(100), &vals, 2, 1, 2, None).unwrap();
2571 let opts = AutoArimaOptions {
2572 strategy: SearchStrategy::Stepwise,
2573 ..Default::default()
2574 };
2575 let stepwise = auto_arima(&ts(100), &vals, 2, 1, 2, Some(&opts)).unwrap();
2576 let denom = exhaustive.aic.abs().max(1e-6);
2577 assert!(
2578 (stepwise.aic - exhaustive.aic).abs() / denom < 0.10,
2579 "stepwise {} vs exhaustive {}",
2580 stepwise.aic,
2581 exhaustive.aic
2582 );
2583 }
2584
2585 #[test]
2586 fn auto_arima_stepwise_uses_fewer_fits() {
2587 let vals: Vec<f64> = (0..150)
2590 .map(|i| f64::from(i) * 2.0 + 10.0)
2591 .zip(noise(31, 150, 1.0))
2592 .map(|(a, b)| a + b)
2593 .collect();
2594 let exhaustive = auto_arima(&ts(150), &vals, 5, 2, 5, None).unwrap();
2595 assert!(
2596 exhaustive.evals > 10,
2597 "grid collapsed: {}",
2598 exhaustive.evals
2599 );
2600 let opts = AutoArimaOptions {
2601 strategy: SearchStrategy::Stepwise,
2602 ..Default::default()
2603 };
2604 let stepwise = auto_arima(&ts(150), &vals, 5, 2, 5, Some(&opts)).unwrap();
2605 assert!(
2606 stepwise.evals < exhaustive.evals,
2607 "stepwise {} vs exhaustive {}",
2608 stepwise.evals,
2609 exhaustive.evals
2610 );
2611 }
2612
2613 fn seasonal(n: usize, m: usize, amp: f64, level: f64) -> Vec<f64> {
2616 (0..n)
2617 .map(|i| level + amp * (2.0 * std::f64::consts::PI * (i % m) as f64 / m as f64).sin())
2618 .collect()
2619 }
2620
2621 #[test]
2622 fn sarima_seasonal_orders_change_the_forecast() {
2623 let vals = seasonal(240, 12, 5.0, 50.0);
2624 let mut plain = SarimaModel::new(1, 0, 0, 0, 0, 0, 12);
2625 plain.fit(&ts(240), &vals).unwrap();
2626 let mut seas = SarimaModel::new(1, 0, 0, 1, 0, 1, 12);
2627 seas.fit(&ts(240), &vals).unwrap();
2628 let a = plain.predict(12).unwrap();
2629 let b = seas.predict(12).unwrap();
2630 assert_ne!(
2631 a.values, b.values,
2632 "seasonal orders must reach the forecast recursion"
2633 );
2634 }
2635
2636 #[test]
2640 fn sarima_forecasts_the_season() {
2641 let m = 12;
2642 let n = 240;
2643 let vals = seasonal(n, m, 5.0, 50.0);
2644 let truth = seasonal(n + m, m, 5.0, 50.0);
2645 let expected = &truth[n..];
2646
2647 let mut seas = SarimaModel::new(0, 0, 0, 1, 0, 1, m);
2648 seas.fit(&ts(n), &vals).unwrap();
2649 let got = seas.predict(m).unwrap();
2650
2651 let mut plain = ArimaModel::new(1, 0, 0);
2652 plain.fit(&ts(n), &vals).unwrap();
2653 let baseline = plain.predict(m).unwrap();
2654
2655 let err = |p: &[f64]| -> f64 {
2656 p.iter()
2657 .zip(expected)
2658 .map(|(a, b)| (a - b).abs())
2659 .sum::<f64>()
2660 / m as f64
2661 };
2662 let seasonal_err = err(&got.values);
2663 let baseline_err = err(&baseline.values);
2664 assert!(
2665 seasonal_err < baseline_err * 0.5,
2666 "SARIMA MAE {seasonal_err:.4} should beat ARIMA MAE {baseline_err:.4}"
2667 );
2668 }
2669
2670 #[test]
2671 fn sarima_seasonal_differencing_round_trips_the_level() {
2672 let m = 12;
2675 let n = 240;
2676 let vals: Vec<f64> = (0..n)
2677 .map(|i| 50.0 + 5.0 * ((i % m) as f64) + 0.25 * i as f64)
2678 .collect();
2679 let mut model = SarimaModel::new(0, 0, 0, 0, 1, 0, m);
2680 model.fit(&ts(n), &vals).unwrap();
2681 let r = model.predict(m).unwrap();
2682 for (h, v) in r.values.iter().enumerate() {
2683 let expected = 50.0 + 5.0 * ((n + h) % m) as f64 + 0.25 * (n + h) as f64;
2684 assert!((v - expected).abs() < 1.0, "h={h}: {v} vs {expected}");
2685 }
2686 }
2687
2688 #[test]
2689 fn sarima_ignores_seasonal_orders_when_the_period_is_degenerate() {
2690 let vals: Vec<f64> = noise(8, 120, 1.0).iter().map(|v| v + 5.0).collect();
2691 let mut model = SarimaModel::new(1, 0, 0, 2, 1, 2, 1);
2692 model.fit(&ts(120), &vals).unwrap();
2693 let ModelParams::Sarima { sp, sd, sq, .. } = model.params() else {
2694 panic!("wrong variant");
2695 };
2696 assert_eq!((*sp, *sd, *sq), (0, 0, 0));
2697 }
2698
2699 #[test]
2700 fn sarima_insufficient_data() {
2701 let vals = seasonal(20, 12, 5.0, 50.0);
2702 let mut model = SarimaModel::new(1, 0, 0, 1, 1, 0, 12);
2703 assert!(model.fit(&ts(20), &vals).is_err());
2704 }
2705
2706 #[test]
2707 fn sarima_update_advances_the_forecast() {
2708 let vals = seasonal(240, 12, 5.0, 50.0);
2709 let mut model = SarimaModel::new(1, 0, 0, 1, 0, 0, 12);
2710 model.fit(&ts(240), &vals).unwrap();
2711 let before = model.predict(1).unwrap();
2712 model.update(240_000_000_000, 60.0).unwrap();
2713 let after = model.predict(1).unwrap();
2714 assert_ne!(before.values[0], after.values[0]);
2715 assert!(after.timestamps[0] > before.timestamps[0]);
2716 }
2717
2718 #[test]
2728 fn sarima_intervals_widen_with_horizon() {
2729 let mut state = 31u64;
2730 let mut noise = move || {
2731 state = state
2732 .wrapping_mul(6_364_136_223_846_793_005)
2733 .wrapping_add(1);
2734 (((state >> 33) as f64 / f64::from(u32::MAX >> 1)) - 1.0) * 0.5
2735 };
2736 let vals: Vec<f64> = seasonal(240, 12, 5.0, 50.0)
2737 .into_iter()
2738 .map(|v| v + noise())
2739 .collect();
2740 let mut model = SarimaModel::new(1, 0, 1, 1, 0, 1, 12);
2741 model.fit(&ts(240), &vals).unwrap();
2742 let r = model.predict(24).unwrap();
2743 let width = |h: usize| r.confidence_upper[h] - r.confidence_lower[h];
2744 for h in 1..24 {
2747 assert!(
2748 width(h) > width(h - 1),
2749 "width must grow strictly: h={h} {} vs {}",
2750 width(h),
2751 width(h - 1)
2752 );
2753 }
2754 let ModelParams::Sarima { residual_std, .. } = model.params() else {
2755 panic!()
2756 };
2757 let phi_star = integrated_ar(&model.ar_expanded, 0, 0, 12);
2758 let psi = psi_weights(&phi_star, &model.ma_expanded, 24);
2759 for h in 1..=24 {
2760 let scale: f64 = psi[..h].iter().map(|w| w * w).sum::<f64>().sqrt();
2761 let expected = 2.0 * 1.959_963_984_540_054 * residual_std * scale;
2762 assert!(
2763 (width(h - 1) - expected).abs() < 1e-9,
2764 "h={h}: width {} != {expected}",
2765 width(h - 1)
2766 );
2767 }
2768 for h in 0..24 {
2769 assert!(r.confidence_lower[h] <= r.values[h]);
2770 assert!(r.confidence_upper[h] >= r.values[h]);
2771 }
2772 }
2773}