Skip to main content

chronix_analytics/forecast/
arima.rs

1//! ARIMA(p,d,q) and SARIMA(p,d,q)(P,D,Q)m forecast models.
2//!
3//! # Estimation
4//!
5//! Both models are estimated by **conditional sum of squares** (CSS), the
6//! standard fast alternative to full maximum likelihood:
7//!
8//! 1. The series is differenced — seasonally first (`D` times at lag `m`),
9//!    then regularly (`d` times).
10//! 2. The autoregressive part is **seeded** by the Burg algorithm, which
11//!    always returns a stable AR polynomial and behaves well on short series.
12//! 3. Every block is then refined **together** by minimising the sum of
13//!    squared one-step innovations — φ and θ for `ArimaModel`, φ, Φ, θ and Θ
14//!    for `SarimaModel`. Jointly, because the blocks interact: the seasonal
15//!    and non-seasonal AR parts multiply, and an MA term drags the lag-1
16//!    autocorrelation away from φ, so a Burg AR fitted to mixed data answers
17//!    a different number and an MA optimised against it fits the residue of
18//!    the wrong model. `ArimaModel` used to do exactly that and returned
19//!    φ = 0.22, θ = -0.02 on 2 000 points of an ARMA(1,1) whose maximum
20//!    likelihood is 0.571 and -0.367.
21//! 4. The search runs over **unconstrained** reals, mapped to coefficients by
22//!    the Jones (1980) reparameterisation: `tanh` gives partial
23//!    autocorrelations in `(-1, 1)`, and the Levinson–Durbin recursion turns
24//!    those into a polynomial whose roots lie outside the unit circle for any
25//!    order. Stationarity and invertibility therefore hold **by
26//!    construction** rather than by a box on each coefficient, which is
27//!    neither sufficient (`φ = (0.99, 0.012)` has a root at 0.9985) nor
28//!    necessary (`φ = (1.2, -0.4)` is an ordinary stationary AR(2) outside
29//!    it).
30//!
31//! The first `max(deg φ*, deg θ*)` innovations are **conditioned away**: they
32//! are computed from pre-sample history the model does not have, and for
33//! `ARIMA(p,0,q)` the first of them is the raw first observation, so counting
34//! them would make the residual variance measure the series' level rather than
35//! its noise. That is what "conditional" in CSS means.
36//!
37//! # Seasonality
38//!
39//! `SarimaModel` expands the multiplicative form
40//!
41//! ```text
42//! φ(B)·Φ(Bᵐ)·(1-B)ᵈ·(1-Bᵐ)ᴰ·xₜ = θ(B)·Θ(Bᵐ)·εₜ
43//! ```
44//!
45//! into a single AR polynomial of degree `p + P·m` and a single MA polynomial
46//! of degree `q + Q·m`, then runs the ordinary ARMA recursion over them. One
47//! recursion and one forecast path, so the seasonal orders reach the forecast
48//! by construction.
49
50use 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
55// ── Shared ARMA core ────────────────────────────────────────────────────
56//
57// One definition of each operation, used by both models. `ArimaModel` is the
58// `P = D = Q = 0` case of `SarimaModel`; keeping the recursions here is what
59// makes that true rather than merely claimed.
60
61/// Difference a series `d` times at lag 1.
62pub(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
73/// Difference a series `d` times at lag `m` (seasonal differencing).
74pub(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
90/// Reconstruct the original scale from lag-`m` differenced predictions.
91///
92/// `anchor` is the tail of the *previous* difference level, and must hold at
93/// least `m` values: prediction `h` adds back the value `m` steps earlier,
94/// which for `h < m` comes from history and for `h >= m` from an already
95/// reconstructed prediction.
96///
97/// Kahan compensated summation bounds the rounding error that otherwise
98/// accumulates across a long horizon, where every value is the sum of every
99/// prediction before it.
100fn 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
114/// Estimate AR coefficients with the Burg algorithm.
115///
116/// Burg minimises the sum of forward and backward prediction errors. Unlike
117/// Yule–Walker it always produces a stable AR polynomial and behaves better on
118/// short series, and it runs in O(N·p) through the Levinson recursion.
119///
120/// Returns `(coefficients, effective_order)`; the effective order is below `p`
121/// when the recursion stops early on near-constant data.
122///
123/// # Errors
124///
125/// [`ForecastError::InsufficientData`] when fewer than `p + 1` observations
126/// are supplied.
127pub(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            // Near-constant data: the reflection coefficient is 0/0. Stop and
156            // report the order actually reached rather than dividing.
157            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
185/// Multiply two polynomials given as coefficient vectors (index = power of B).
186fn 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
199/// Expand `φ(B)·Φ(Bᵐ)` into one AR coefficient vector of degree `p + P·m`.
200///
201/// Both inputs and the output use the recursion convention
202/// `xₜ = Σ φᵢ·xₜ₋ᵢ + …`, i.e. the AR *polynomial* is `1 − Σ φᵢBⁱ`; the sign
203/// flip happens here so no caller has to remember it.
204pub(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
221/// Expand `θ(B)·Θ(Bᵐ)` into one MA coefficient vector of degree `q + Q·m`.
222///
223/// Convention `xₜ = … + Σ θⱼ·εₜ₋ⱼ + εₜ`, i.e. the MA polynomial is
224/// `1 + Σ θⱼBʲ`, so no sign flip is needed — but stating it is what stops the
225/// two conventions being mixed, which is the classic way a seasonal MA term
226/// ends up with the wrong sign.
227pub(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
242/// One-step innovations `εₜ` of `xₜ = Σφᵢxₜ₋ᵢ + Σθⱼεₜ₋ⱼ + εₜ`, conditioned on
243/// zero pre-sample history.
244///
245/// The returned vector is aligned with `values`; entries before
246/// [`warmup`](conditioning_warmup) are conditioned on history that does not
247/// exist and must not be used as residuals.
248pub(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/// Whether the model carries a constant, by the Hyndman–Khandakar rule
268/// `d + D <= 1`.
269///
270/// The constant is the mean of the *fully differenced* series, so what it
271/// means depends on how much differencing precedes it: with no differencing
272/// it is the **level** the series reverts to, and with one it is the **drift**
273/// per step. Without it, `ARIMA(1,0,0)` on a series sitting at 100 has to
274/// explain that level with a near-unit root, and `ARIMA(0,1,0)` on a trend
275/// forecasts a flat line.
276///
277/// Above one difference a constant implies a polynomial trend that keeps
278/// accelerating, which is essentially never what is wanted from a metric —
279/// which is why `auto.arima` draws the line in the same place.
280#[inline]
281pub(crate) fn includes_constant(d: usize, seasonal_d: usize) -> bool {
282    d + seasonal_d <= 1
283}
284
285/// Number of leading innovations that are conditioned on absent history.
286#[inline]
287pub(crate) fn conditioning_warmup(ar_len: usize, ma_len: usize) -> usize {
288    ar_len.max(ma_len)
289}
290
291/// Conditional sum of squares: the CSS objective.
292/// Map unconstrained reals onto a **stationary** AR coefficient vector.
293///
294/// Jones (1980) / Monahan (1984): `tanh` puts each value in `(-1, 1)`, where
295/// it is a partial autocorrelation, and the Levinson–Durbin recursion turns a
296/// sequence of PACFs in that open interval into a polynomial `1 - Σφⱼzʲ`
297/// whose every root lies strictly outside the unit circle — for any order,
298/// by construction. It is what `statsmodels`' `enforce_stationarity` does.
299///
300/// The alternative, and what stood here, is a box: bound every coefficient to
301/// ±0.99 and hope. That is neither sufficient nor necessary. Not sufficient,
302/// because `φ = (0.99, 0.012)` is inside the box and has a root at 0.9985, so
303/// the forecast recursion **diverges** — 21 of 80 fits on a random walk
304/// landed there. Not necessary, because a perfectly ordinary stationary AR(2)
305/// such as `φ = (1.2, -0.4)` is outside it and cannot be reached at all.
306fn 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
321/// The inverse of [`constrain_stationary`], used to start the optimiser from
322/// a Burg estimate rather than from zero.
323///
324/// Burg always returns a stationary polynomial, so the peel below stays
325/// inside `(-1, 1)`; the clamp is for the boundary case where it lands
326/// exactly on it.
327fn 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
346/// Map unconstrained reals onto an **invertible** MA coefficient vector.
347///
348/// The same recursion, negated. `constrain_stationary` produces the
349/// coefficients of `1 - Σψⱼzʲ`; this crate's MA convention is
350/// `xₜ = … + Σθⱼεₜ₋ⱼ + εₜ`, whose polynomial is `1 + Σθⱼzʲ`, so `θ = -ψ`
351/// gives the identical polynomial and therefore the identical roots.
352fn constrain_invertible(unconstrained: &[f64]) -> Vec<f64> {
353    constrain_stationary(unconstrained)
354        .iter()
355        .map(|v| -v)
356        .collect()
357}
358
359/// Minimise the conditional sum of squares over **every** AR and MA block at
360/// once, and return the parameter vector.
361///
362/// `expand` maps a parameter vector to the `(AR, MA)` polynomials the
363/// innovations recursion runs over: the identity for ARIMA, the
364/// multiplicative expansion for SARIMA. Brent below two parameters,
365/// Nelder–Mead above, over the unconstrained parameters.
366///
367/// **Jointly is the whole point.** `ArimaModel` used to fix the AR part at its
368/// Burg estimate and optimise only the MA coefficients against it, and a Burg
369/// AR fitted to mixed ARMA data is biased — the moving-average term drags the
370/// lag-1 autocorrelation away from φ, so Burg answers a different number and
371/// the MA search then fits the residue of the wrong model. On 2 000 points of
372/// `x = 0.6x[t-1] + e - 0.4e[t-1]` it returned φ = 0.22, θ = -0.02 where
373/// `SarimaModel` — the sibling, which had always optimised its blocks
374/// together — returned 0.571 and -0.367, matching statsmodels' maximum
375/// likelihood to four decimals. Pure AR and pure MA were unaffected, which is
376/// why the defect survived: the two orders anybody tests by hand are the two
377/// that have nothing to interact with.
378fn 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    // The search runs in the **unconstrained** space: `expand` puts every
391    // candidate through `constrain_stationary` / `constrain_invertible`
392    // first, so there is no region to keep the optimiser out of and no box to
393    // choose. `tanh` saturates long before ±12, which is only here so the
394    // simplex cannot wander to infinity on a flat objective.
395    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
416/// ψ-weights of the MA(∞) representation `xₜ = Σⱼ ψⱼ·εₜ₋ⱼ`, `ψ₀ = 1`.
417///
418/// The forecast variance at horizon `h` is `σ²·Σⱼ₌₀^{h-1} ψⱼ²`. Recursion:
419/// `ψⱼ = Σᵢ₌₁ᵖ φᵢ·ψⱼ₋ᵢ + θⱼ`, with `θⱼ = 0` for `j > q`.
420pub(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
438/// Extend a differenced series `horizon` steps under the ARMA recursion,
439/// returning only the new values. Future innovations are zero, which is the
440/// conditional expectation.
441fn 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    // Align the innovation history with the value history: a shorter residual
452    // vector (after streaming updates trimmed it) must not shift the MA lags.
453    // Both adjustments happen at the *front*: `truncate` dropped the newest
454    // innovations, which are exactly the ones the MA lags read, so an
455    // over-long buffer shifted every lag by the excess.
456    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
483/// Residual standard deviation over the innovations that are not conditioned
484/// on absent history, with a degrees-of-freedom correction for the estimated
485/// parameters.
486fn 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
497/// The **full** autoregressive operator of an ARIMA/SARIMA model,
498/// `φ(B)·Φ(Bᵐ)·(1-B)ᵈ·(1-Bᵐ)ᴰ`, expanded into a single coefficient vector in
499/// the same convention as `ar` (`xₜ = Σᵢ cᵢ·xₜ₋ᵢ + …`).
500///
501/// # Why the differencing has to be in here
502///
503/// The ψ-weights of the *stationary* ARMA part describe the error of the
504/// differenced series. The forecast is issued on the original scale, and
505/// integration accumulates those errors: for `ARIMA(0,1,0)` every ψⱼ is 1, so
506/// `Var[e_h] = h·σ²` and the interval widens as `√h`. Taking ψ from the ARMA
507/// part alone gives `ψ = [1]` and a **constant-width** interval for a random
508/// walk — the textbook example of an interval that must widen.
509pub(crate) fn integrated_ar(ar: &[f64], d: usize, seasonal_d: usize, m: usize) -> Vec<f64> {
510    // φ(B) = 1 - φ₁B - … - φₚBᵖ
511    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
528/// Build symmetric 95 % prediction intervals from the ψ-weights of the
529/// **integrated** model — see [`integrated_ar`].
530fn 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
554/// Forecast timestamps at `last_ts + h·interval`, saturating rather than
555/// wrapping on overflow.
556fn 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// ── ARIMA ───────────────────────────────────────────────────────────────
563
564/// ARIMA(p,d,q) — non-seasonal, estimated by Burg + conditional sum of
565/// squares.
566#[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    /// Mean of the differenced series; the level for `d = 0`, the drift for
575    /// `d = 1`, and zero above that.
576    constant: f64,
577    /// Leading innovations conditioned on absent history; excluded from every
578    /// residual statistic.
579    warmup: usize,
580    last_ts: i64,
581    interval_ns: i64,
582    fitted: bool,
583}
584
585impl ArimaModel {
586    /// Creates a new ARIMA model with the given orders.
587    #[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    /// Free parameters, counting the innovation variance and the constant if
614    /// the model carries one. This is the `k` of every information criterion
615    /// below; counting a constant that is not there is a quiet way to make
616    /// two model families incomparable.
617    fn n_params(&self) -> usize {
618        self.p + self.q + usize::from(includes_constant(self.d, 0)) + 1
619    }
620
621    /// Residual sum of squares over the innovations that are not conditioned
622    /// on absent history, and how many there are.
623    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    /// Akaike Information Criterion, `n·ln(RSS/n) + 2k`.
633    ///
634    /// `n` counts only the innovations the conditional likelihood is defined
635    /// over, which is what makes the value comparable between two models of
636    /// the same series at the same differencing order. It is *not* comparable
637    /// across differencing orders — a differenced series is a different
638    /// series — which is why [`auto_arima`] chooses `d` by a unit-root test
639    /// rather than by this number.
640    #[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    /// Bayesian Information Criterion, `n·ln(RSS/n) + k·ln(n)`.
650    #[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    /// Corrected Akaike Information Criterion (AICc).
660    #[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    /// Comprehensive model diagnostics.
670    ///
671    /// Requires both the original training values and the in-sample
672    /// predictions (fitted values) for the model.
673    #[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        // Burg seeds the AR block; the optimiser then moves **both** blocks
715        // together. Holding the AR part at Burg and fitting only the MA
716        // against it is biased for any mixed model — see `minimise_css`.
717        let (seed_ar, effective_ar_order) = estimate_ar_burg(&centered, 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(&centered, &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(&centered, &ar_coeffs, &ma_coeffs);
730
731        // Emit the effective order so a dashboard can see silent degradation.
732        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(&centered, &self.residuals, ar_coeffs, ma_coeffs, horizon)
782                .into_iter()
783                .map(|v| v + constant)
784                .collect();
785
786        // Undifference outward: anchor `i` is the last value of the series
787        // differenced `i` times, so index 0 is the original scale.
788        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        // Bound the buffers so a long-running stream does not grow without
856        // limit, keeping enough context for the model orders.
857        //
858        // `residuals` is aligned one-for-one with `difference(history, d)`,
859        // which is the only reason `predict` may read them at the same index.
860        // Trimming the two to *independent* caps broke that alignment by `d`
861        // after the first trim, silently shifting every MA lag from then on.
862        // One cap, and the residual buffer follows the history buffer.
863        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// ── SARIMA ──────────────────────────────────────────────────────────────
885
886/// Configuration for constructing a [`SarimaModel`].
887///
888/// A config struct rather than a seven-parameter constructor: `(1, 1, 1, 0,
889/// 1, 1, 24)` is unreadable at a call site, and swapping two of those numbers
890/// changes the model silently.
891#[derive(Debug, Clone, Copy, PartialEq, Eq)]
892pub struct SarimaConfig {
893    /// Non-seasonal autoregressive order `p`.
894    pub p: usize,
895    /// Non-seasonal differencing order `d`.
896    pub d: usize,
897    /// Non-seasonal moving-average order `q`.
898    pub q: usize,
899    /// Seasonal autoregressive order `P`.
900    pub sp: usize,
901    /// Seasonal differencing order `D`.
902    pub sd: usize,
903    /// Seasonal moving-average order `Q`.
904    pub sq: usize,
905    /// Seasonal period `m` (e.g. 24 for hourly data with daily seasonality).
906    pub m: usize,
907}
908
909/// SARIMA(p,d,q)(P,D,Q)m.
910///
911/// The multiplicative form is expanded into a single AR polynomial of degree
912/// `p + P·m` and a single MA polynomial of degree `q + Q·m`, so the seasonal
913/// orders participate in the forecast recursion rather than only in the
914/// stored parameters. All four coefficient blocks are estimated jointly by
915/// conditional sum of squares, seeded from a Burg fit of the non-seasonal AR
916/// part — the seasonal and non-seasonal AR factors multiply, so estimating
917/// one while pretending the other is absent biases both.
918#[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    /// Original-scale training values.
929    history: Vec<f64>,
930    /// Innovations aligned with the fully differenced series.
931    residuals: Vec<f64>,
932    /// Expanded AR polynomial `φ*` of degree `p + P·m`.
933    ar_expanded: Vec<f64>,
934    /// Expanded MA polynomial `θ*` of degree `q + Q·m`.
935    ma_expanded: Vec<f64>,
936    /// Mean of the fully differenced series; see [`includes_constant`].
937    constant: f64,
938    warmup: usize,
939    last_ts: i64,
940    interval_ns: i64,
941    fitted: bool,
942}
943
944impl SarimaModel {
945    /// Creates a new SARIMA model from a [`SarimaConfig`].
946    #[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    /// Creates a new SARIMA model with explicit order parameters.
952    ///
953    /// A period `m` of 0 or 1 leaves no seasonal structure to model; the
954    /// seasonal orders are then ignored and the model degenerates to
955    /// ARIMA(p,d,q), which is what those periods mean.
956    #[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    /// Fully difference the series: seasonally `D` times at lag `m`, then
996    /// regularly `d` times.
997    fn transform(&self, values: &[f64]) -> Vec<f64> {
998        difference(&seasonal_difference(values, self.m, self.sd), self.cfg_d)
999    }
1000
1001    /// The fully differenced series with the constant removed — the scale the
1002    /// ARMA recursion is defined on.
1003    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    /// Free parameters, counting the innovation variance and the constant.
1011    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    /// The expanded polynomials implied by a packed parameter vector
1021    /// `[φ₁…φ_p, Φ₁…Φ_P, θ₁…θ_q, Θ₁…Θ_Q]`.
1022    /// Split an **unconstrained** parameter vector into the four blocks, put
1023    /// each through its stationarity or invertibility transformation, and
1024    /// return the four constrained blocks.
1025    ///
1026    /// Each block is constrained on its own, which is what the multiplicative
1027    /// form asks for: `φ(B)·Φ(Bᵐ)` is stable exactly when both factors are.
1028    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    /// The expanded `(AR, MA)` polynomials for an unconstrained vector.
1039    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        // The expanded polynomials reach back `p + P·m` steps, and each round
1055        // of differencing costs observations up front. Ten usable innovations
1056        // is the floor below which the CSS objective is noise.
1057        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        // Undifference in the reverse of the order differencing was applied:
1155        // regular first (innermost), then seasonal. Each level's anchor is the
1156        // tail of the series differenced one step less than the predictions.
1157        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        // The interval is computed from the ψ-weights of the *integrated*
1189        // operator: differencing has unit lead coefficient, so the one-step
1190        // error is the same on both scales, but beyond one step integration
1191        // accumulates the earlier errors and the interval on the original
1192        // scale is strictly wider than the one on the differenced scale.
1193        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        // Keep enough original-scale history to re-derive the differenced
1240        // series the recursion reads, plus a margin — and keep `residuals`
1241        // aligned one-for-one with that differenced series, which is the
1242        // invariant `predict` depends on. See `ArimaModel::update`.
1243        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
1265// ── Unit-root testing: how many times to difference ─────────────────────
1266
1267/// Critical value of the KPSS level-stationarity statistic at 5 %.
1268///
1269/// Kwiatkowski, Phillips, Schmidt & Shin (1992), Table 1, level-stationary
1270/// case: 10 % → 0.347, **5 % → 0.463**, 2.5 % → 0.574, 1 % → 0.739.
1271pub const KPSS_CRITICAL_5PCT: f64 = 0.463;
1272
1273/// KPSS test statistic for level stationarity.
1274///
1275/// The null hypothesis is that the series is stationary around a constant, so
1276/// a statistic **above** the critical value is evidence that it is not and
1277/// that a difference is called for. That direction is the opposite of an
1278/// augmented Dickey–Fuller test and is the usual way this test is misread.
1279///
1280/// The long-run variance uses a Bartlett kernel with the Schwert short
1281/// truncation lag `⌊4·(n/100)^¼⌋`, which is the default in both R's
1282/// `tseries::kpss.test` and `urca::ur.kpss`.
1283///
1284/// Returns 0 for a series with fewer than three observations or no variance,
1285/// both of which are stationary by any reading.
1286#[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    // Partial sums of the demeaned series.
1296    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    // Bartlett-weighted long-run variance.
1311    #[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/// Choose the differencing order `d` by successive KPSS tests.
1331///
1332/// This is the Hyndman–Khandakar rule used by R's `auto.arima`: difference
1333/// while the series still rejects level stationarity, up to `max_d`.
1334///
1335/// It exists because **AIC cannot choose `d`**. Differencing changes the
1336/// series the likelihood is computed over — a different number of
1337/// observations, on a different scale — so the value is not comparable across
1338/// orders, and comparing it anyway systematically over-differences: on
1339/// stationary noise the `d = 1` fits win every time, and the resulting model
1340/// forecasts a random walk.
1341#[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(&current) <= KPSS_CRITICAL_5PCT {
1347            break;
1348        }
1349        current = difference(&current, 1);
1350        d += 1;
1351    }
1352    d
1353}
1354
1355/// Seasonal strength above which a series is treated as having a seasonal
1356/// unit root, from Wang, Smith & Hyndman (2006).
1357///
1358/// R's `forecast::nsdiffs` uses this exact threshold as its **default** test,
1359/// with the comment "Threshold chosen based on seasonal M3 auto.arima
1360/// accuracy" — it was fitted by minimising MASE over the M3 and M4
1361/// collections rather than derived, which is why it is a bare number and why
1362/// copying it is better than inventing one.
1363pub const SEASONAL_STRENGTH_THRESHOLD: f64 = 0.64;
1364
1365/// How many seasonal differences a series needs, from its seasonal strength.
1366///
1367/// The measure is `max(0, min(1, 1 − Var(remainder) / Var(remainder +
1368/// seasonal)))` over an STL decomposition: the share of the seasonal-plus-noise
1369/// variance that the seasonal component explains. One seasonal difference is
1370/// taken when it exceeds [`SEASONAL_STRENGTH_THRESHOLD`].
1371///
1372/// This is the default test in R's `forecast::nsdiffs`, in preference to OCSB
1373/// — which that package keeps as an option rather than a default. Reading the
1374/// two guards out of its source mattered as much as the formula: a **constant**
1375/// series needs no differencing however the variance ratio comes out (both
1376/// variances are zero), and a series shorter than two periods cannot be
1377/// decomposed at all, let alone differenced.
1378///
1379/// # Why this is not cosmetic
1380///
1381/// The alternative is what stood here: a fixed `D = 1`. Every seasonal
1382/// candidate was seasonally differenced, so a series with strong but
1383/// *stationary* seasonality — a daily load curve that repeats rather than
1384/// drifting — was always over-differenced, and the model that fits it,
1385/// `SARIMA(p,d,q)(P,0,Q)[m]`, could not be proposed at all. Over-differencing
1386/// does not fail; it inflates the forecast variance and widens every interval
1387/// built on it.
1388#[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    // Fewer than two full periods cannot be decomposed, and `nsdiffs`
1394    // refuses for the same reason.
1395    if values.len() < 2 * period {
1396        return 0;
1397    }
1398    // A constant series has zero variance in both terms, so the ratio is
1399    // meaningless; R checks this before the heuristic and so does this.
1400    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(&current, period) else {
1409            break;
1410        };
1411        if strength <= SEASONAL_STRENGTH_THRESHOLD {
1412            break;
1413        }
1414        current = seasonal_difference(&current, period, 1);
1415        d += 1;
1416        if current.len() < 2 * period {
1417            break;
1418        }
1419    }
1420    d
1421}
1422
1423/// `max(0, min(1, 1 − Var(remainder) / Var(remainder + seasonal)))`, or
1424/// `None` if the series cannot be decomposed.
1425fn 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
1445/// Sample variance, or `None` for fewer than two observations.
1446fn 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// ── auto_arima ──────────────────────────────────────────────────────────
1456
1457/// Result of automatic ARIMA model selection.
1458#[derive(Debug, Clone)]
1459pub struct AutoArimaResult {
1460    /// Best `(p, d, q)` order.
1461    pub order: (usize, usize, usize),
1462    /// AICc score of the best model.
1463    pub aic: f64,
1464    /// Every evaluated candidate, sorted by score ascending.
1465    pub candidates: Vec<(usize, usize, usize, f64)>,
1466    /// Number of model fits performed.
1467    pub evals: usize,
1468    /// Differencing order chosen by the unit-root test, before `max_d` was
1469    /// applied. Exposed so a caller can see that the cap bound the answer.
1470    pub kpss_order: usize,
1471}
1472
1473/// Search strategy for [`auto_arima`].
1474#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
1475pub enum SearchStrategy {
1476    /// Exhaustive grid over `p ∈ [0, max_p]`, `q ∈ [0, max_q]`.
1477    #[default]
1478    Exhaustive,
1479    /// Hyndman–Khandakar stepwise search — O(P+Q) fits instead of O(P×Q), at
1480    /// the risk of a local optimum.
1481    Stepwise,
1482}
1483
1484/// Budget controls for [`auto_arima`].
1485#[derive(Debug, Clone, Copy, Default)]
1486pub struct AutoArimaOptions {
1487    /// Hard cap on the number of model fits. `None` means unlimited.
1488    pub max_evals: Option<usize>,
1489    /// Search strategy.
1490    pub strategy: SearchStrategy,
1491}
1492
1493/// Automatically select an ARIMA(p,d,q) order.
1494///
1495/// # How the order is chosen
1496///
1497/// `d` comes from a **KPSS unit-root test**
1498/// ([`select_differencing_order`]), capped at `max_d`; `p` and `q` are then
1499/// chosen by **AICc** at that fixed `d`. Splitting the decision this way is
1500/// not a refinement — it is what makes the information criterion meaningful,
1501/// because AIC/AICc compare models of the *same* data and two differencing
1502/// orders are two different datasets.
1503///
1504/// AICc rather than AIC: the correction matters at exactly the series lengths
1505/// a gateway has, and it converges to AIC as `n` grows.
1506///
1507/// # Strategies
1508///
1509/// - [`SearchStrategy::Exhaustive`] (default) fits every `(p, q)` in the grid
1510///   in parallel and returns the global AICc minimum within it. With `d`
1511///   fixed the grid is `(max_p+1)·(max_q+1)`, which is small enough that this
1512///   is usually the right choice.
1513/// - [`SearchStrategy::Stepwise`] follows Hyndman–Khandakar: seed with
1514///   `(2,q=2)`, `(0,0)`, `(1,0)`, `(0,1)`, then walk to improving neighbours.
1515///
1516/// Use [`AutoArimaOptions::max_evals`] to cap the cost of either.
1517///
1518/// # Errors
1519///
1520/// [`ForecastError::InsufficientData`] below 30 observations, and
1521/// [`ForecastError::NumericalInstability`] when no candidate fits.
1522pub 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/// Hyndman–Khandakar stepwise walk over `(p, q)` at a fixed `d`.
1591#[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    /// Deterministic uniform noise on `[-amp, amp]`, so every test in this
1697    /// module is reproducible without a RNG dependency.
1698    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    // ── Core recursions ─────────────────────────────────────────────
1715
1716    #[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            // Undifference the tail back to the original scale.
1722            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        // (1 - 0.5B)(1 - 0.3B⁴) = 1 - 0.5B - 0.3B⁴ + 0.15B⁵
1762        // In the "plus" convention that is φ* = [0.5, 0, 0, 0.3, -0.15].
1763        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        // (1 + 0.4B)(1 + 0.2B³) = 1 + 0.4B + 0.2B³ + 0.08B⁴
1770        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        // Roots of z² − φ₁z − φ₂ must lie inside the unit circle.
1816        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    // ── ARIMA ───────────────────────────────────────────────────────
1829
1830    /// A random walk is `ARIMA(0,1,0)`: every ψⱼ is 1, so `Var[e_h] = h·σ²`
1831    /// and the variance ratio is **exactly** `h`. Before the integration term
1832    /// entered the ψ-weights this was a constant-width interval.
1833    #[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    /// `ARIMA(1,1,0)`: `φ*(B) = (1-φB)(1-B)`, so
1855    /// `ψⱼ = 1 + φ + … + φʲ = (1-φ^{j+1})/(1-φ)`.
1856    #[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    /// `integrated_ar` must expand `φ(B)(1-B)^d(1-B^m)^D` exactly.
1896    #[test]
1897    fn integrated_ar_expands_the_differencing_operator() {
1898        // (1-B) → xₜ = xₜ₋₁
1899        assert_eq!(integrated_ar(&[], 1, 0, 0), vec![1.0]);
1900        // (1-B)² = 1 - 2B + B² → xₜ = 2xₜ₋₁ - xₜ₋₂
1901        let d2 = integrated_ar(&[], 2, 0, 0);
1902        assert!((d2[0] - 2.0).abs() < 1e-12 && (d2[1] + 1.0).abs() < 1e-12);
1903        // (1-0.5B)(1-B) = 1 - 1.5B + 0.5B²
1904        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        // (1-B⁴): coefficient 1 at lag 4, zero elsewhere.
1907        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        // No differencing → unchanged.
1912        assert_eq!(integrated_ar(&[0.3, -0.2], 0, 0, 12), vec![0.3, -0.2]);
1913    }
1914
1915    /// `arma_forecast` must read the **newest** innovations. With a residual
1916    /// buffer longer than the value buffer, `θ₁ = 0.5` and last innovation
1917    /// 5.0, the one-step forecast is exactly 2.5; truncating from the tail
1918    /// gave 1.5.
1919    #[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    /// `predict` reads `residuals[i]` alongside `difference(history, d)[i]`,
1928    /// so after the streaming buffers are trimmed the two lengths must still
1929    /// differ by exactly `d`.
1930    #[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        // …and the MA term of the one-step forecast must use the newest
1949        // innovation: ŷ = last + c + θ₁·e_last for ARIMA(p,1,1) with p's
1950        // contribution written out from the differenced history.
1951        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    /// The same alignment invariant for SARIMA.
1974    #[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        // d = 1, so the model carries a drift: the mean first difference is
2012        // exactly 2, and the forecast continues the line.
2013        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    /// Above one difference there is no constant, so `ARIMA(0,2,0)` is the
2020    /// pure `2x_{t-1} - x_{t-2}` extrapolation — exact on a line, and short of
2021    /// a quadratic by its (unmodelled) curvature. Both are the right answer
2022    /// for the model that was asked for.
2023    #[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        // No constant at d = 2, so the forecast is the linear continuation of
2044        // the last first difference: x_{n+h} = x_{n-1} + (h+1)·(x_{n-1} −
2045        // x_{n-2}). It falls short of the true quadratic by the curvature the
2046        // model does not carry, and it must do so *exactly*.
2047        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    /// The conditioning warm-up is the whole reason ARIMA(p,0,q) intervals are
2069    /// usable: without it the first innovation is the raw first observation,
2070    /// so a series sitting at 100 with noise of 0.1 gets an interval two
2071    /// orders of magnitude too wide.
2072    #[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    // ── KPSS and differencing order ─────────────────────────────────
2135
2136    #[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    /// Strong seasonality takes a seasonal difference.
2176    ///
2177    /// This is the M3/M4-tuned behaviour of the measure, and it is worth
2178    /// being explicit that it is a **strength** test, not a unit-root test:
2179    /// a perfectly repeating pattern around a stable level scores near 1 and
2180    /// is differenced. R's `nsdiffs` does the same, deliberately — the
2181    /// threshold was fitted by minimising forecast error, not derived from
2182    /// stationarity theory.
2183    #[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    /// Faint seasonality buried in noise takes none — and this is the case
2196    /// the fixed `D = 1` got wrong.
2197    ///
2198    /// `detect_period` will happily return a period for a series whose
2199    /// seasonal term explains almost none of its variance. Differencing it
2200    /// anyway spends `m` observations, adds a moving-average term the data
2201    /// does not support, and widens every interval built on the result.
2202    #[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    /// The measure's blind spot, pinned rather than left to be rediscovered.
2226    ///
2227    /// A **seasonal random walk** — each season's level drifts from cycle to
2228    /// cycle — is the textbook case for a seasonal difference, and the
2229    /// strength heuristic scores it *low*: STL's cycle-subseries smoother
2230    /// cannot fit a seasonal shape that keeps moving, so the variation lands
2231    /// in the remainder and the ratio collapses. This is exactly what a real
2232    /// seasonal unit-root test (OCSB, Canova–Hansen) is for, and why R keeps
2233    /// them as options beside this default. Recorded as a test so the limit
2234    /// is a known quantity rather than a surprise.
2235    /// Smallest |root| of `1 + θ₁z + … + θ_qz^q`, by Durand–Kerner.
2236    ///
2237    /// The MA polynomial is invertible iff every root lies strictly outside
2238    /// the unit circle, so the answer is compared against 1.
2239    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        // `1 + 0.99z - 0.99z²` has a root at ≈ -0.6225: inside the unit
2304        // circle, so this θ is non-invertible *and* inside the ±0.99 box
2305        // every coefficient is bounded to. The box is not the constraint.
2306        let m = min_ma_root_modulus(&[0.99, -0.99]);
2307        assert!((m - 0.6225).abs() < 1e-3, "min |root| was {m}");
2308        // An invertible one, for the other direction.
2309        assert!(min_ma_root_modulus(&[0.5]) > 1.0);
2310    }
2311
2312    #[test]
2313    fn a_fitted_ma_polynomial_is_invertible() {
2314        // `estimate_ma` bounds each coefficient to ±0.99, which does not
2315        // imply invertibility for q > 1 — see the test above. What keeps
2316        // fits inside the invertible region is the **objective**: the CSS
2317        // error recursion diverges outside it, so the optimiser has no
2318        // reason to go there. That is a property of the estimator, not of
2319        // the bounds, so it is worth checking rather than assuming — if the
2320        // estimator is ever replaced (an exact likelihood, say), this is the
2321        // guarantee that quietly changes.
2322        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        // Over-differenced noise is the textbook way to manufacture a
2371        // non-invertible MA(1), so it belongs in the sample rather than
2372        // beside it.
2373        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        // Each season's level is a random walk across cycles, so the series
2388        // has a seasonal unit root and `D = 1` is the textbook answer.
2389        //
2390        // This used to assert `0`, under the name "the strength measure is
2391        // blind to a stochastic seasonal level", with a note asking whoever
2392        // saw it change to check which test was being used before calling it
2393        // an improvement. It was not the measure that changed: the STL under
2394        // it padded the low-pass filter by repeating the first and last
2395        // cycle, which is right only for a series that is already periodic,
2396        // and a seasonal level that wanders is exactly the series that is
2397        // not.
2398        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        // The other direction, and the one that matters here: a monotonic
2420        // counter is the commonest series in a metrics database and has no
2421        // seasonal component at all. The seasonal strength measure is a
2422        // ratio of the remainder's variance to the seasonal-plus-remainder
2423        // variance, so it says nothing useful unless the decomposition
2424        // underneath it is sound — and a straight line came out of the old
2425        // STL with a seasonal swing of 278 on values running 1_000 to
2426        // 11_619, which read as strength 0.81 against a 0.64 threshold.
2427        //
2428        // Asserted for a clean ramp and for two noise levels, because the
2429        // clean one alone would also pass under a guard that merely refuses
2430        // to divide by something small.
2431        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        // Both variances are zero, so the ratio says nothing; the guard has
2462        // to come first. R's `nsdiffs` checks `is.constant` before the
2463        // heuristic for exactly this reason.
2464        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    // ── auto_arima ──────────────────────────────────────────────────
2496
2497    #[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    /// Why `d` comes from KPSS: scoring it by an information criterion picks
2508    /// `d = 1` on stationary noise, because differencing shrinks the residual
2509    /// sum of squares the criterion is comparing.
2510    #[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        // Noisy, or every candidate fits exactly, RSS is zero and no
2588        // information criterion is defined for any of them.
2589        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    // ── SARIMA ──────────────────────────────────────────────────────
2614
2615    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    /// A seasonal model must beat a non-seasonal one on seasonal data — the
2637    /// property that fails the moment the seasonal coefficients stop reaching
2638    /// the forecast recursion.
2639    #[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        // A pure repeating pattern with a linear drift: (0,0,0)(0,1,0)m
2673        // forecasts x_{t} = x_{t-m} + drift, which reproduces the pattern.
2674        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    /// A prediction interval widens with the horizon, strictly.
2719    ///
2720    /// The series carries **noise** on purpose. A noiseless seasonal sine is
2721    /// fitted so well that `residual_std` is ~1e-10, every width is ~1e-9,
2722    /// and "strictly wider" becomes a question about the last bits of a
2723    /// float rather than about the ψ-weights: the assertion passed only
2724    /// while the fit was poor enough to leave something to measure. It
2725    /// started failing the moment the estimator got better, which is the
2726    /// wrong direction for a test to point.
2727    #[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        // Strict growth, with the ψ-weight closed form: `>=` also passes for
2745        // a constant-width interval, which is exactly the bug it missed.
2746        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}