Skip to main content

fdars_core/
detrend.rs

1//! Detrending and decomposition functions for non-stationary functional data.
2//!
3//! This module provides methods for removing trends from functional data
4//! to enable more accurate seasonal analysis. It includes:
5//! - Linear detrending (least squares)
6//! - Polynomial detrending (QR decomposition)
7//! - Differencing (first and second order)
8//! - LOESS detrending (local polynomial regression)
9//! - Spline detrending (P-splines)
10//! - Automatic method selection via AIC
11
12use crate::iter_maybe_parallel;
13use crate::matrix::FdMatrix;
14use crate::smoothing::local_polynomial;
15use nalgebra::{DMatrix, DVector, Dyn, SVD};
16#[cfg(feature = "parallel")]
17use rayon::iter::ParallelIterator;
18use std::borrow::Cow;
19
20/// Result of detrending operation.
21#[derive(Debug, Clone)]
22pub struct TrendResult {
23    /// Estimated trend values (n x m)
24    pub trend: FdMatrix,
25    /// Detrended data (n x m)
26    pub detrended: FdMatrix,
27    /// Method used for detrending
28    pub method: Cow<'static, str>,
29    /// Polynomial coefficients (for polynomial methods, per sample)
30    /// For n samples with polynomial degree d: n x (d+1)
31    pub coefficients: Option<FdMatrix>,
32    /// Residual sum of squares for each sample
33    pub rss: Vec<f64>,
34    /// Number of parameters (for AIC calculation)
35    pub n_params: usize,
36}
37
38impl TrendResult {
39    /// Construct a no-op TrendResult (zero trend, data copied to detrended).
40    fn empty(data: &FdMatrix, n: usize, m: usize, method: Cow<'static, str>, n_params: usize) -> Self {
41        TrendResult {
42            trend: FdMatrix::zeros(n, m),
43            detrended: FdMatrix::from_slice(data.as_slice(), n, m)
44                .unwrap_or_else(|| FdMatrix::zeros(n, m)),
45            method,
46            coefficients: None,
47            rss: vec![0.0; n],
48            n_params,
49        }
50    }
51}
52
53/// Result of seasonal decomposition.
54#[derive(Debug, Clone)]
55pub struct DecomposeResult {
56    /// Trend component (n x m)
57    pub trend: FdMatrix,
58    /// Seasonal component (n x m)
59    pub seasonal: FdMatrix,
60    /// Remainder/residual component (n x m)
61    pub remainder: FdMatrix,
62    /// Period used for decomposition
63    pub period: f64,
64    /// Decomposition method ("additive" or "multiplicative")
65    pub method: Cow<'static, str>,
66}
67
68/// Remove linear trend from functional data using least squares.
69///
70/// # Arguments
71/// * `data` - Matrix (n x m): n samples, m evaluation points
72/// * `argvals` - Time/argument values of length m
73///
74/// # Returns
75/// TrendResult with trend, detrended data, and coefficients (intercept, slope)
76pub fn detrend_linear(data: &FdMatrix, argvals: &[f64]) -> TrendResult {
77    let (n, m) = data.shape();
78    if n == 0 || m < 2 || argvals.len() != m {
79        return TrendResult::empty(data, n, m, Cow::Borrowed("linear"), 2);
80    }
81
82    let mean_t: f64 = argvals.iter().sum::<f64>() / m as f64;
83    let ss_t: f64 = argvals.iter().map(|&t| (t - mean_t).powi(2)).sum();
84
85    // Compute only scalar coefficients in parallel (no Vec allocations per sample)
86    let scalars: Vec<(f64, f64, f64)> = iter_maybe_parallel!(0..n)
87        .map(|i| {
88            let mut sum_y = 0.0;
89            for j in 0..m {
90                sum_y += data[(i, j)];
91            }
92            let mean_y = sum_y / m as f64;
93            let mut sp = 0.0;
94            for j in 0..m {
95                sp += (argvals[j] - mean_t) * (data[(i, j)] - mean_y);
96            }
97            let slope = if ss_t.abs() > 1e-15 { sp / ss_t } else { 0.0 };
98            let intercept = mean_y - slope * mean_t;
99            let mut rss = 0.0;
100            for j in 0..m {
101                let residual = data[(i, j)] - (intercept + slope * argvals[j]);
102                rss += residual * residual;
103            }
104            (intercept, slope, rss)
105        })
106        .collect();
107
108    // Write directly into pre-allocated output matrices
109    let mut trend = FdMatrix::zeros(n, m);
110    let mut detrended = FdMatrix::zeros(n, m);
111    let mut coefficients = FdMatrix::zeros(n, 2);
112    let mut rss = vec![0.0; n];
113
114    for (i, &(intercept, slope, r)) in scalars.iter().enumerate() {
115        coefficients[(i, 0)] = intercept;
116        coefficients[(i, 1)] = slope;
117        rss[i] = r;
118        for j in 0..m {
119            let trend_val = intercept + slope * argvals[j];
120            trend[(i, j)] = trend_val;
121            detrended[(i, j)] = data[(i, j)] - trend_val;
122        }
123    }
124
125    TrendResult {
126        trend,
127        detrended,
128        method: Cow::Borrowed("linear"),
129        coefficients: Some(coefficients),
130        rss,
131        n_params: 2,
132    }
133}
134
135fn build_vandermonde_matrix(t_norm: &[f64], m: usize, n_coef: usize) -> DMatrix<f64> {
136    let mut design = DMatrix::zeros(m, n_coef);
137    for j in 0..m {
138        let t = t_norm[j];
139        let mut power = 1.0;
140        for k in 0..n_coef {
141            design[(j, k)] = power;
142            power *= t;
143        }
144    }
145    design
146}
147
148fn fit_polynomial_single_curve(
149    curve: &[f64],
150    svd: &SVD<f64, Dyn, Dyn>,
151    design: &DMatrix<f64>,
152    n_coef: usize,
153    m: usize,
154) -> (Vec<f64>, Vec<f64>, Vec<f64>, f64) {
155    let y = DVector::from_row_slice(curve);
156    let beta = svd
157        .solve(&y, 1e-10)
158        .unwrap_or_else(|_| DVector::zeros(n_coef));
159    let fitted = design * &beta;
160    let mut trend = vec![0.0; m];
161    let mut detrended = vec![0.0; m];
162    let mut rss = 0.0;
163    for j in 0..m {
164        trend[j] = fitted[j];
165        detrended[j] = curve[j] - fitted[j];
166        rss += detrended[j].powi(2);
167    }
168    let coefs: Vec<f64> = beta.iter().cloned().collect();
169    (trend, detrended, coefs, rss)
170}
171
172fn diff_single_curve(curve: &[f64], m: usize, order: usize) -> (Vec<f64>, Vec<f64>, Vec<f64>, f64) {
173    let diff1: Vec<f64> = (0..m - 1).map(|j| curve[j + 1] - curve[j]).collect();
174    let detrended = if order == 2 {
175        (0..diff1.len() - 1)
176            .map(|j| diff1[j + 1] - diff1[j])
177            .collect()
178    } else {
179        diff1.clone()
180    };
181    let initial_values = if order == 2 {
182        vec![curve[0], curve[1]]
183    } else {
184        vec![curve[0]]
185    };
186    let rss: f64 = detrended.iter().map(|&x| x.powi(2)).sum();
187    let new_m = m - order;
188    let mut trend = vec![0.0; m];
189    trend[0] = curve[0];
190    if order == 1 {
191        for j in 1..m {
192            trend[j] = curve[j] - if j <= new_m { detrended[j - 1] } else { 0.0 };
193        }
194    } else {
195        trend = curve.to_vec();
196    }
197    let mut det_full = vec![0.0; m];
198    det_full[..new_m].copy_from_slice(&detrended[..new_m]);
199    (trend, det_full, initial_values, rss)
200}
201
202fn reassemble_polynomial_results(
203    results: Vec<(Vec<f64>, Vec<f64>, Vec<f64>, f64)>,
204    n: usize,
205    m: usize,
206    n_coef: usize,
207) -> (FdMatrix, FdMatrix, FdMatrix, Vec<f64>) {
208    let mut trend = FdMatrix::zeros(n, m);
209    let mut detrended = FdMatrix::zeros(n, m);
210    let mut coefficients = FdMatrix::zeros(n, n_coef);
211    let mut rss = vec![0.0; n];
212    for (i, (t, d, coefs, r)) in results.into_iter().enumerate() {
213        for j in 0..m {
214            trend[(i, j)] = t[j];
215            detrended[(i, j)] = d[j];
216        }
217        for k in 0..n_coef {
218            coefficients[(i, k)] = coefs[k];
219        }
220        rss[i] = r;
221    }
222    (trend, detrended, coefficients, rss)
223}
224
225/// Reassemble per-curve (trend, detrended, rss) results into FdMatrix outputs.
226fn reassemble_trend_results(
227    results: Vec<(Vec<f64>, Vec<f64>, f64)>,
228    n: usize,
229    m: usize,
230) -> (FdMatrix, FdMatrix, Vec<f64>) {
231    let mut trend = FdMatrix::zeros(n, m);
232    let mut detrended = FdMatrix::zeros(n, m);
233    let mut rss = vec![0.0; n];
234    for (i, (t, d, r)) in results.into_iter().enumerate() {
235        for j in 0..m {
236            trend[(i, j)] = t[j];
237            detrended[(i, j)] = d[j];
238        }
239        rss[i] = r;
240    }
241    (trend, detrended, rss)
242}
243
244/// Remove polynomial trend from functional data using QR decomposition.
245pub fn detrend_polynomial(data: &FdMatrix, argvals: &[f64], degree: usize) -> TrendResult {
246    let (n, m) = data.shape();
247    if n == 0 || m < degree + 1 || argvals.len() != m || degree == 0 {
248        return TrendResult::empty(data, n, m, Cow::Owned(format!("polynomial({})", degree)), degree + 1);
249    }
250    if degree == 1 {
251        let mut result = detrend_linear(data, argvals);
252        result.method = Cow::Borrowed("polynomial(1)");
253        return result;
254    }
255    let n_coef = degree + 1;
256    let t_min = argvals.iter().cloned().fold(f64::INFINITY, f64::min);
257    let t_max = argvals.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
258    let t_range = if (t_max - t_min).abs() > 1e-15 {
259        t_max - t_min
260    } else {
261        1.0
262    };
263    let t_norm: Vec<f64> = argvals.iter().map(|&t| (t - t_min) / t_range).collect();
264    let design = build_vandermonde_matrix(&t_norm, m, n_coef);
265    let svd = design.clone().svd(true, true);
266    let results: Vec<(Vec<f64>, Vec<f64>, Vec<f64>, f64)> = iter_maybe_parallel!(0..n)
267        .map(|i| {
268            let curve: Vec<f64> = (0..m).map(|j| data[(i, j)]).collect();
269            fit_polynomial_single_curve(&curve, &svd, &design, n_coef, m)
270        })
271        .collect();
272    let (trend, detrended, coefficients, rss) =
273        reassemble_polynomial_results(results, n, m, n_coef);
274    TrendResult {
275        trend,
276        detrended,
277        method: Cow::Owned(format!("polynomial({})", degree)),
278        coefficients: Some(coefficients),
279        rss,
280        n_params: n_coef,
281    }
282}
283
284/// Remove trend by differencing.
285pub fn detrend_diff(data: &FdMatrix, order: usize) -> TrendResult {
286    let (n, m) = data.shape();
287    if n == 0 || m <= order || order == 0 || order > 2 {
288        return TrendResult::empty(data, n, m, Cow::Owned(format!("diff{}", order)), order);
289    }
290    let results: Vec<(Vec<f64>, Vec<f64>, Vec<f64>, f64)> = iter_maybe_parallel!(0..n)
291        .map(|i| {
292            let curve: Vec<f64> = (0..m).map(|j| data[(i, j)]).collect();
293            diff_single_curve(&curve, m, order)
294        })
295        .collect();
296    let (trend, detrended, coefficients, rss) =
297        reassemble_polynomial_results(results, n, m, order);
298    TrendResult {
299        trend,
300        detrended,
301        method: Cow::Owned(format!("diff{}", order)),
302        coefficients: Some(coefficients),
303        rss,
304        n_params: order,
305    }
306}
307
308/// Remove trend using LOESS (local polynomial regression).
309pub fn detrend_loess(
310    data: &FdMatrix,
311    argvals: &[f64],
312    bandwidth: f64,
313    degree: usize,
314) -> TrendResult {
315    let (n, m) = data.shape();
316    if n == 0 || m < 3 || argvals.len() != m || bandwidth <= 0.0 {
317        return TrendResult::empty(data, n, m, Cow::Borrowed("loess"), (m as f64 * bandwidth.max(0.0)).ceil() as usize);
318    }
319    let t_min = argvals.iter().cloned().fold(f64::INFINITY, f64::min);
320    let t_max = argvals.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
321    let abs_bandwidth = (t_max - t_min) * bandwidth;
322    let results: Vec<(Vec<f64>, Vec<f64>, f64)> = iter_maybe_parallel!(0..n)
323        .map(|i| {
324            let curve: Vec<f64> = (0..m).map(|j| data[(i, j)]).collect();
325            let trend =
326                local_polynomial(argvals, &curve, argvals, abs_bandwidth, degree, "gaussian");
327            let mut detrended = vec![0.0; m];
328            let mut rss = 0.0;
329            for j in 0..m {
330                detrended[j] = curve[j] - trend[j];
331                rss += detrended[j].powi(2);
332            }
333            (trend, detrended, rss)
334        })
335        .collect();
336    let (trend, detrended, rss) = reassemble_trend_results(results, n, m);
337    let n_params = (m as f64 * bandwidth).ceil() as usize;
338    TrendResult {
339        trend,
340        detrended,
341        method: Cow::Borrowed("loess"),
342        coefficients: None,
343        rss,
344        n_params,
345    }
346}
347
348/// Automatically select the best detrending method using AIC.
349pub fn auto_detrend(data: &FdMatrix, argvals: &[f64]) -> TrendResult {
350    let (n, m) = data.shape();
351    if n == 0 || m < 4 || argvals.len() != m {
352        return TrendResult::empty(data, n, m, Cow::Borrowed("auto(none)"), 0);
353    }
354    let compute_aic = |result: &TrendResult| -> f64 {
355        let mut total_aic = 0.0;
356        for i in 0..n {
357            let rss = result.rss[i];
358            let k = result.n_params as f64;
359            let aic = if rss > 1e-15 {
360                m as f64 * (rss / m as f64).ln() + 2.0 * k
361            } else {
362                f64::NEG_INFINITY
363            };
364            total_aic += aic;
365        }
366        total_aic / n as f64
367    };
368    let linear = detrend_linear(data, argvals);
369    let poly2 = detrend_polynomial(data, argvals, 2);
370    let poly3 = detrend_polynomial(data, argvals, 3);
371    let loess = detrend_loess(data, argvals, 0.3, 2);
372    let aic_linear = compute_aic(&linear);
373    let aic_poly2 = compute_aic(&poly2);
374    let aic_poly3 = compute_aic(&poly3);
375    let aic_loess = compute_aic(&loess);
376    let methods = [
377        (aic_linear, "linear", linear),
378        (aic_poly2, "polynomial(2)", poly2),
379        (aic_poly3, "polynomial(3)", poly3),
380        (aic_loess, "loess", loess),
381    ];
382    let (_, best_name, mut best_result) = methods
383        .into_iter()
384        .min_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal))
385        .unwrap();
386    best_result.method = Cow::Owned(format!("auto({})", best_name));
387    best_result
388}
389
390fn fit_fourier_seasonal(
391    detrended_i: &[f64],
392    argvals: &[f64],
393    omega: f64,
394    n_harm: usize,
395    m: usize,
396) -> (Vec<f64>, Vec<f64>) {
397    let n_coef = 2 * n_harm;
398    let mut design = DMatrix::zeros(m, n_coef);
399    for j in 0..m {
400        let t = argvals[j];
401        for k in 0..n_harm {
402            let freq = (k + 1) as f64 * omega;
403            design[(j, 2 * k)] = (freq * t).cos();
404            design[(j, 2 * k + 1)] = (freq * t).sin();
405        }
406    }
407    let y = DVector::from_row_slice(detrended_i);
408    let svd = design.clone().svd(true, true);
409    let coef = svd
410        .solve(&y, 1e-10)
411        .unwrap_or_else(|_| DVector::zeros(n_coef));
412    let fitted = &design * &coef;
413    let seasonal: Vec<f64> = fitted.iter().cloned().collect();
414    let remainder: Vec<f64> = (0..m).map(|j| detrended_i[j] - seasonal[j]).collect();
415    (seasonal, remainder)
416}
417
418/// Additive seasonal decomposition: data = trend + seasonal + remainder
419pub fn decompose_additive(
420    data: &FdMatrix,
421    argvals: &[f64],
422    period: f64,
423    trend_method: &str,
424    bandwidth: f64,
425    n_harmonics: usize,
426) -> DecomposeResult {
427    let (n, m) = data.shape();
428    if n == 0 || m < 4 || argvals.len() != m || period <= 0.0 {
429        return DecomposeResult {
430            trend: FdMatrix::zeros(n, m),
431            seasonal: FdMatrix::zeros(n, m),
432            remainder: FdMatrix::from_slice(data.as_slice(), n, m)
433                .unwrap_or_else(|| FdMatrix::zeros(n, m)),
434            period,
435            method: Cow::Borrowed("additive"),
436        };
437    }
438    let _ = trend_method;
439    let trend_result = detrend_loess(data, argvals, bandwidth.max(0.3), 2);
440    let n_harm = n_harmonics.max(1).min(m / 4);
441    let omega = 2.0 * std::f64::consts::PI / period;
442    let results: Vec<(Vec<f64>, Vec<f64>, Vec<f64>)> = iter_maybe_parallel!(0..n)
443        .map(|i| {
444            let trend_i: Vec<f64> = (0..m).map(|j| trend_result.trend[(i, j)]).collect();
445            let detrended_i: Vec<f64> = (0..m).map(|j| trend_result.detrended[(i, j)]).collect();
446            let (seasonal, remainder) =
447                fit_fourier_seasonal(&detrended_i, argvals, omega, n_harm, m);
448            (trend_i, seasonal, remainder)
449        })
450        .collect();
451    let mut trend = FdMatrix::zeros(n, m);
452    let mut seasonal = FdMatrix::zeros(n, m);
453    let mut remainder = FdMatrix::zeros(n, m);
454    for (i, (t, s, r)) in results.into_iter().enumerate() {
455        for j in 0..m {
456            trend[(i, j)] = t[j];
457            seasonal[(i, j)] = s[j];
458            remainder[(i, j)] = r[j];
459        }
460    }
461    DecomposeResult {
462        trend,
463        seasonal,
464        remainder,
465        period,
466        method: Cow::Borrowed("additive"),
467    }
468}
469
470/// Multiplicative seasonal decomposition: data = trend * seasonal * remainder
471pub fn decompose_multiplicative(
472    data: &FdMatrix,
473    argvals: &[f64],
474    period: f64,
475    trend_method: &str,
476    bandwidth: f64,
477    n_harmonics: usize,
478) -> DecomposeResult {
479    let (n, m) = data.shape();
480    if n == 0 || m < 4 || argvals.len() != m || period <= 0.0 {
481        return DecomposeResult {
482            trend: FdMatrix::zeros(n, m),
483            seasonal: FdMatrix::zeros(n, m),
484            remainder: FdMatrix::from_slice(data.as_slice(), n, m)
485                .unwrap_or_else(|| FdMatrix::zeros(n, m)),
486            period,
487            method: Cow::Borrowed("multiplicative"),
488        };
489    }
490    let min_val = data
491        .as_slice()
492        .iter()
493        .cloned()
494        .fold(f64::INFINITY, f64::min);
495    let shift = if min_val <= 0.0 { -min_val + 1.0 } else { 0.0 };
496    let log_data_vec: Vec<f64> = data.as_slice().iter().map(|&x| (x + shift).ln()).collect();
497    let log_data = FdMatrix::from_column_major(log_data_vec, n, m).unwrap();
498    let additive_result = decompose_additive(
499        &log_data,
500        argvals,
501        period,
502        trend_method,
503        bandwidth,
504        n_harmonics,
505    );
506    let mut trend = FdMatrix::zeros(n, m);
507    let mut seasonal = FdMatrix::zeros(n, m);
508    let mut remainder = FdMatrix::zeros(n, m);
509    for i in 0..n {
510        for j in 0..m {
511            trend[(i, j)] = additive_result.trend[(i, j)].exp() - shift;
512            seasonal[(i, j)] = additive_result.seasonal[(i, j)].exp();
513            remainder[(i, j)] = additive_result.remainder[(i, j)].exp();
514        }
515    }
516    DecomposeResult {
517        trend,
518        seasonal,
519        remainder,
520        period,
521        method: Cow::Borrowed("multiplicative"),
522    }
523}
524
525// ============================================================================
526// STL Decomposition (Cleveland et al., 1990)
527// ============================================================================
528
529/// Result of STL decomposition including robustness weights.
530#[derive(Debug, Clone)]
531pub struct StlResult {
532    /// Trend component (n x m)
533    pub trend: FdMatrix,
534    /// Seasonal component (n x m)
535    pub seasonal: FdMatrix,
536    /// Remainder/residual component (n x m)
537    pub remainder: FdMatrix,
538    /// Robustness weights per point (n x m)
539    pub weights: FdMatrix,
540    /// Period used for decomposition
541    pub period: usize,
542    /// Seasonal smoothing window
543    pub s_window: usize,
544    /// Trend smoothing window
545    pub t_window: usize,
546    /// Number of inner loop iterations performed
547    pub inner_iterations: usize,
548    /// Number of outer loop iterations performed
549    pub outer_iterations: usize,
550}
551
552/// STL Decomposition: Seasonal and Trend decomposition using LOESS
553pub fn stl_decompose(
554    data: &FdMatrix,
555    period: usize,
556    s_window: Option<usize>,
557    t_window: Option<usize>,
558    l_window: Option<usize>,
559    robust: bool,
560    inner_iterations: Option<usize>,
561    outer_iterations: Option<usize>,
562) -> StlResult {
563    let (n, m) = data.shape();
564    if n == 0 || m < 2 * period || period < 2 {
565        return StlResult {
566            trend: FdMatrix::zeros(n, m),
567            seasonal: FdMatrix::zeros(n, m),
568            remainder: FdMatrix::from_slice(data.as_slice(), n, m)
569                .unwrap_or_else(|| FdMatrix::zeros(n, m)),
570            weights: FdMatrix::from_column_major(vec![1.0; n * m], n, m)
571                .unwrap_or_else(|| FdMatrix::zeros(n, m)),
572            period,
573            s_window: 0,
574            t_window: 0,
575            inner_iterations: 0,
576            outer_iterations: 0,
577        };
578    }
579    let s_win = s_window.unwrap_or(7).max(3) | 1;
580    let t_win = t_window.unwrap_or_else(|| {
581        let ratio = 1.5 * period as f64 / (1.0 - 1.5 / s_win as f64);
582        let val = ratio.ceil() as usize;
583        val.max(3) | 1
584    });
585    let l_win = l_window.unwrap_or(period) | 1;
586    let n_inner = inner_iterations.unwrap_or(2);
587    let n_outer = outer_iterations.unwrap_or(if robust { 15 } else { 1 });
588    let results: Vec<(Vec<f64>, Vec<f64>, Vec<f64>, Vec<f64>)> = iter_maybe_parallel!(0..n)
589        .map(|i| {
590            let curve: Vec<f64> = (0..m).map(|j| data[(i, j)]).collect();
591            stl_single_series(
592                &curve, period, s_win, t_win, l_win, robust, n_inner, n_outer,
593            )
594        })
595        .collect();
596    let mut trend = FdMatrix::zeros(n, m);
597    let mut seasonal = FdMatrix::zeros(n, m);
598    let mut remainder = FdMatrix::zeros(n, m);
599    let mut weights = FdMatrix::from_column_major(vec![1.0; n * m], n, m).unwrap();
600    for (i, (t, s, r, w)) in results.into_iter().enumerate() {
601        for j in 0..m {
602            trend[(i, j)] = t[j];
603            seasonal[(i, j)] = s[j];
604            remainder[(i, j)] = r[j];
605            weights[(i, j)] = w[j];
606        }
607    }
608    StlResult {
609        trend,
610        seasonal,
611        remainder,
612        weights,
613        period,
614        s_window: s_win,
615        t_window: t_win,
616        inner_iterations: n_inner,
617        outer_iterations: n_outer,
618    }
619}
620
621fn stl_single_series(
622    data: &[f64],
623    period: usize,
624    s_window: usize,
625    t_window: usize,
626    l_window: usize,
627    robust: bool,
628    n_inner: usize,
629    n_outer: usize,
630) -> (Vec<f64>, Vec<f64>, Vec<f64>, Vec<f64>) {
631    let m = data.len();
632    let mut trend = vec![0.0; m];
633    let mut seasonal = vec![0.0; m];
634    let mut weights = vec![1.0; m];
635    for _outer in 0..n_outer {
636        for _inner in 0..n_inner {
637            let detrended: Vec<f64> = data
638                .iter()
639                .zip(trend.iter())
640                .map(|(&y, &t)| y - t)
641                .collect();
642            let cycle_smoothed = smooth_cycle_subseries(&detrended, period, s_window, &weights);
643            let low_pass = stl_lowpass_filter(&cycle_smoothed, period, l_window);
644            seasonal = cycle_smoothed
645                .iter()
646                .zip(low_pass.iter())
647                .map(|(&c, &l)| c - l)
648                .collect();
649            let deseasonalized: Vec<f64> = data
650                .iter()
651                .zip(seasonal.iter())
652                .map(|(&y, &s)| y - s)
653                .collect();
654            trend = weighted_loess(&deseasonalized, t_window, &weights);
655        }
656        if robust && _outer < n_outer - 1 {
657            let remainder: Vec<f64> = data
658                .iter()
659                .zip(trend.iter())
660                .zip(seasonal.iter())
661                .map(|((&y, &t), &s)| y - t - s)
662                .collect();
663            weights = compute_robustness_weights(&remainder);
664        }
665    }
666    let remainder: Vec<f64> = data
667        .iter()
668        .zip(trend.iter())
669        .zip(seasonal.iter())
670        .map(|((&y, &t), &s)| y - t - s)
671        .collect();
672    (trend, seasonal, remainder, weights)
673}
674
675fn smooth_cycle_subseries(
676    data: &[f64],
677    period: usize,
678    s_window: usize,
679    weights: &[f64],
680) -> Vec<f64> {
681    let m = data.len();
682    let n_cycles = m.div_ceil(period);
683    let mut result = vec![0.0; m];
684    for pos in 0..period {
685        let mut subseries_idx: Vec<usize> = Vec::new();
686        let mut subseries_vals: Vec<f64> = Vec::new();
687        let mut subseries_weights: Vec<f64> = Vec::new();
688        for cycle in 0..n_cycles {
689            let idx = cycle * period + pos;
690            if idx < m {
691                subseries_idx.push(idx);
692                subseries_vals.push(data[idx]);
693                subseries_weights.push(weights[idx]);
694            }
695        }
696        if subseries_vals.is_empty() {
697            continue;
698        }
699        let smoothed = weighted_loess(&subseries_vals, s_window, &subseries_weights);
700        for (i, &idx) in subseries_idx.iter().enumerate() {
701            result[idx] = smoothed[i];
702        }
703    }
704    result
705}
706
707fn stl_lowpass_filter(data: &[f64], period: usize, _l_window: usize) -> Vec<f64> {
708    let ma1 = moving_average(data, period);
709    let ma2 = moving_average(&ma1, period);
710    moving_average(&ma2, 3)
711}
712
713fn moving_average(data: &[f64], window: usize) -> Vec<f64> {
714    let m = data.len();
715    if m == 0 || window == 0 {
716        return data.to_vec();
717    }
718    let half = window / 2;
719    let mut result = vec![0.0; m];
720    for i in 0..m {
721        let start = i.saturating_sub(half);
722        let end = (i + half + 1).min(m);
723        let sum: f64 = data[start..end].iter().sum();
724        let count = (end - start) as f64;
725        result[i] = sum / count;
726    }
727    result
728}
729
730fn weighted_loess(data: &[f64], window: usize, weights: &[f64]) -> Vec<f64> {
731    let m = data.len();
732    if m == 0 {
733        return vec![];
734    }
735    let half = window / 2;
736    let mut result = vec![0.0; m];
737    for i in 0..m {
738        let start = i.saturating_sub(half);
739        let end = (i + half + 1).min(m);
740        let mut sum_w = 0.0;
741        let mut sum_wx = 0.0;
742        let mut sum_wy = 0.0;
743        let mut sum_wxx = 0.0;
744        let mut sum_wxy = 0.0;
745        for j in start..end {
746            let dist = (j as f64 - i as f64).abs() / (half.max(1) as f64);
747            let tricube = if dist < 1.0 {
748                (1.0 - dist.powi(3)).powi(3)
749            } else {
750                0.0
751            };
752            let w = tricube * weights[j];
753            let x = j as f64;
754            let y = data[j];
755            sum_w += w;
756            sum_wx += w * x;
757            sum_wy += w * y;
758            sum_wxx += w * x * x;
759            sum_wxy += w * x * y;
760        }
761        if sum_w > 1e-10 {
762            let denom = sum_w * sum_wxx - sum_wx * sum_wx;
763            if denom.abs() > 1e-10 {
764                let intercept = (sum_wxx * sum_wy - sum_wx * sum_wxy) / denom;
765                let slope = (sum_w * sum_wxy - sum_wx * sum_wy) / denom;
766                result[i] = intercept + slope * i as f64;
767            } else {
768                result[i] = sum_wy / sum_w;
769            }
770        } else {
771            result[i] = data[i];
772        }
773    }
774    result
775}
776
777fn compute_robustness_weights(residuals: &[f64]) -> Vec<f64> {
778    let m = residuals.len();
779    if m == 0 {
780        return vec![];
781    }
782    let mut abs_residuals: Vec<f64> = residuals.iter().map(|&r| r.abs()).collect();
783    abs_residuals.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
784    let median_idx = m / 2;
785    let mad = if m % 2 == 0 {
786        (abs_residuals[median_idx - 1] + abs_residuals[median_idx]) / 2.0
787    } else {
788        abs_residuals[median_idx]
789    };
790    let h = 6.0 * mad.max(1e-10);
791    residuals
792        .iter()
793        .map(|&r| {
794            let u = r.abs() / h;
795            if u < 1.0 {
796                (1.0 - u * u).powi(2)
797            } else {
798                0.0
799            }
800        })
801        .collect()
802}
803
804/// Wrapper function for functional data STL decomposition.
805pub fn stl_fdata(
806    data: &FdMatrix,
807    _argvals: &[f64],
808    period: usize,
809    s_window: Option<usize>,
810    t_window: Option<usize>,
811    robust: bool,
812) -> StlResult {
813    stl_decompose(data, period, s_window, t_window, None, robust, None, None)
814}
815
816#[cfg(test)]
817mod tests {
818    use super::*;
819    use std::f64::consts::PI;
820
821    #[test]
822    fn test_detrend_linear_removes_linear_trend() {
823        let m = 100;
824        let argvals: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64 * 10.0).collect();
825        let data_vec: Vec<f64> = argvals
826            .iter()
827            .map(|&t| 2.0 + 0.5 * t + (2.0 * PI * t / 2.0).sin())
828            .collect();
829        let data = FdMatrix::from_column_major(data_vec, 1, m).unwrap();
830        let result = detrend_linear(&data, &argvals);
831        let expected: Vec<f64> = argvals
832            .iter()
833            .map(|&t| (2.0 * PI * t / 2.0).sin())
834            .collect();
835        let mut max_diff = 0.0f64;
836        for j in 0..m {
837            let diff = (result.detrended[(0, j)] - expected[j]).abs();
838            max_diff = max_diff.max(diff);
839        }
840        assert!(max_diff < 0.2, "Max difference: {}", max_diff);
841    }
842
843    #[test]
844    fn test_detrend_polynomial_removes_quadratic_trend() {
845        let m = 100;
846        let argvals: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64 * 10.0).collect();
847        let data_vec: Vec<f64> = argvals
848            .iter()
849            .map(|&t| 1.0 + 0.5 * t - 0.1 * t * t + (2.0 * PI * t / 2.0).sin())
850            .collect();
851        let data = FdMatrix::from_column_major(data_vec, 1, m).unwrap();
852        let result = detrend_polynomial(&data, &argvals, 2);
853        let expected: Vec<f64> = argvals
854            .iter()
855            .map(|&t| (2.0 * PI * t / 2.0).sin())
856            .collect();
857        let detrended_vec: Vec<f64> = (0..m).map(|j| result.detrended[(0, j)]).collect();
858        let mean_det: f64 = detrended_vec.iter().sum::<f64>() / m as f64;
859        let mean_exp: f64 = expected.iter().sum::<f64>() / m as f64;
860        let mut num = 0.0;
861        let mut den_det = 0.0;
862        let mut den_exp = 0.0;
863        for j in 0..m {
864            num += (detrended_vec[j] - mean_det) * (expected[j] - mean_exp);
865            den_det += (detrended_vec[j] - mean_det).powi(2);
866            den_exp += (expected[j] - mean_exp).powi(2);
867        }
868        let corr = num / (den_det.sqrt() * den_exp.sqrt());
869        assert!(corr > 0.95, "Correlation: {}", corr);
870    }
871
872    #[test]
873    fn test_detrend_diff1() {
874        let m = 100;
875        let data_vec: Vec<f64> = {
876            let mut v = vec![0.0; m];
877            v[0] = 1.0;
878            for i in 1..m {
879                v[i] = v[i - 1] + 0.1 * (i as f64).sin();
880            }
881            v
882        };
883        let data = FdMatrix::from_column_major(data_vec.clone(), 1, m).unwrap();
884        let result = detrend_diff(&data, 1);
885        for j in 0..m - 1 {
886            let expected = data_vec[j + 1] - data_vec[j];
887            assert!(
888                (result.detrended[(0, j)] - expected).abs() < 1e-10,
889                "Mismatch at {}: {} vs {}",
890                j,
891                result.detrended[(0, j)],
892                expected
893            );
894        }
895    }
896
897    #[test]
898    fn test_auto_detrend_selects_linear_for_linear_data() {
899        let m = 100;
900        let argvals: Vec<f64> = (0..m).map(|i| i as f64).collect();
901        let data_vec: Vec<f64> = argvals.iter().map(|&t| 2.0 + 0.5 * t).collect();
902        let data = FdMatrix::from_column_major(data_vec, 1, m).unwrap();
903        let result = auto_detrend(&data, &argvals);
904        assert!(
905            result.method.contains("linear") || result.method.contains("polynomial"),
906            "Method: {}",
907            result.method
908        );
909    }
910
911    #[test]
912    fn test_detrend_loess_removes_linear_trend() {
913        let m = 100;
914        let argvals: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64 * 10.0).collect();
915        let data_vec: Vec<f64> = argvals
916            .iter()
917            .map(|&t| 2.0 + 0.5 * t + (2.0 * PI * t / 2.0).sin())
918            .collect();
919        let data = FdMatrix::from_column_major(data_vec, 1, m).unwrap();
920        let result = detrend_loess(&data, &argvals, 0.3, 1);
921        let expected: Vec<f64> = argvals
922            .iter()
923            .map(|&t| (2.0 * PI * t / 2.0).sin())
924            .collect();
925        let detrended_vec: Vec<f64> = (0..m).map(|j| result.detrended[(0, j)]).collect();
926        let mean_det: f64 = detrended_vec.iter().sum::<f64>() / m as f64;
927        let mean_exp: f64 = expected.iter().sum::<f64>() / m as f64;
928        let mut num = 0.0;
929        let mut den_det = 0.0;
930        let mut den_exp = 0.0;
931        for j in 0..m {
932            num += (detrended_vec[j] - mean_det) * (expected[j] - mean_exp);
933            den_det += (detrended_vec[j] - mean_det).powi(2);
934            den_exp += (expected[j] - mean_exp).powi(2);
935        }
936        let corr = num / (den_det.sqrt() * den_exp.sqrt());
937        assert!(corr > 0.9, "Correlation: {}", corr);
938        assert_eq!(result.method, "loess");
939    }
940
941    #[test]
942    fn test_detrend_loess_removes_quadratic_trend() {
943        let m = 100;
944        let argvals: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64 * 10.0).collect();
945        let data_vec: Vec<f64> = argvals
946            .iter()
947            .map(|&t| 1.0 + 0.3 * t - 0.05 * t * t + (2.0 * PI * t / 2.0).sin())
948            .collect();
949        let data = FdMatrix::from_column_major(data_vec, 1, m).unwrap();
950        let result = detrend_loess(&data, &argvals, 0.3, 2);
951        assert_eq!(result.trend.ncols(), m);
952        assert_eq!(result.detrended.ncols(), m);
953        assert!(result.rss[0] > 0.0);
954    }
955
956    #[test]
957    fn test_detrend_loess_different_bandwidths() {
958        let m = 100;
959        let argvals: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64 * 10.0).collect();
960        let data_vec: Vec<f64> = argvals
961            .iter()
962            .enumerate()
963            .map(|(i, &t)| (2.0 * PI * t / 2.0).sin() + 0.1 * ((i * 17) % 100) as f64 / 100.0)
964            .collect();
965        let data = FdMatrix::from_column_major(data_vec, 1, m).unwrap();
966        let result_small = detrend_loess(&data, &argvals, 0.1, 1);
967        let result_large = detrend_loess(&data, &argvals, 0.5, 1);
968        assert_eq!(result_small.trend.ncols(), m);
969        assert_eq!(result_large.trend.ncols(), m);
970        assert!(result_large.n_params >= result_small.n_params);
971    }
972
973    #[test]
974    fn test_detrend_loess_short_series() {
975        let m = 10;
976        let argvals: Vec<f64> = (0..m).map(|i| i as f64).collect();
977        let data_vec: Vec<f64> = argvals.iter().map(|&t| t * 2.0).collect();
978        let data = FdMatrix::from_column_major(data_vec, 1, m).unwrap();
979        let result = detrend_loess(&data, &argvals, 0.3, 1);
980        assert_eq!(result.trend.ncols(), m);
981        assert_eq!(result.detrended.ncols(), m);
982    }
983
984    #[test]
985    fn test_decompose_additive_separates_components() {
986        let m = 200;
987        let period = 2.0;
988        let argvals: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64 * 10.0).collect();
989        let data_vec: Vec<f64> = argvals
990            .iter()
991            .map(|&t| 2.0 + 0.5 * t + (2.0 * PI * t / period).sin())
992            .collect();
993        let data = FdMatrix::from_column_major(data_vec.clone(), 1, m).unwrap();
994        let result = decompose_additive(&data, &argvals, period, "loess", 0.3, 3);
995        assert_eq!(result.trend.ncols(), m);
996        assert_eq!(result.seasonal.ncols(), m);
997        assert_eq!(result.remainder.ncols(), m);
998        assert_eq!(result.method, "additive");
999        assert_eq!(result.period, period);
1000        for j in 0..m {
1001            let reconstructed =
1002                result.trend[(0, j)] + result.seasonal[(0, j)] + result.remainder[(0, j)];
1003            assert!(
1004                (reconstructed - data_vec[j]).abs() < 0.5,
1005                "Reconstruction error at {}: {} vs {}",
1006                j,
1007                reconstructed,
1008                data_vec[j]
1009            );
1010        }
1011    }
1012
1013    #[test]
1014    fn test_decompose_additive_different_harmonics() {
1015        let m = 200;
1016        let period = 2.0;
1017        let argvals: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64 * 10.0).collect();
1018        let data_vec: Vec<f64> = argvals
1019            .iter()
1020            .map(|&t| 1.0 + (2.0 * PI * t / period).sin())
1021            .collect();
1022        let data = FdMatrix::from_column_major(data_vec, 1, m).unwrap();
1023        let result1 = decompose_additive(&data, &argvals, period, "loess", 0.3, 1);
1024        let result5 = decompose_additive(&data, &argvals, period, "loess", 0.3, 5);
1025        assert_eq!(result1.seasonal.ncols(), m);
1026        assert_eq!(result5.seasonal.ncols(), m);
1027    }
1028
1029    #[test]
1030    fn test_decompose_additive_residual_properties() {
1031        let m = 200;
1032        let period = 2.0;
1033        let argvals: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64 * 10.0).collect();
1034        let data_vec: Vec<f64> = argvals
1035            .iter()
1036            .map(|&t| 2.0 + 0.3 * t + (2.0 * PI * t / period).sin())
1037            .collect();
1038        let data = FdMatrix::from_column_major(data_vec.clone(), 1, m).unwrap();
1039        let result = decompose_additive(&data, &argvals, period, "loess", 0.3, 3);
1040        let remainder_vec: Vec<f64> = (0..m).map(|j| result.remainder[(0, j)]).collect();
1041        let mean_rem: f64 = remainder_vec.iter().sum::<f64>() / m as f64;
1042        assert!(mean_rem.abs() < 0.5, "Remainder mean: {}", mean_rem);
1043        let data_mean: f64 = data_vec.iter().sum::<f64>() / m as f64;
1044        let var_data: f64 = data_vec
1045            .iter()
1046            .map(|&x| (x - data_mean).powi(2))
1047            .sum::<f64>()
1048            / m as f64;
1049        let var_rem: f64 = remainder_vec
1050            .iter()
1051            .map(|&x| (x - mean_rem).powi(2))
1052            .sum::<f64>()
1053            / m as f64;
1054        assert!(
1055            var_rem < var_data,
1056            "Remainder variance {} should be < data variance {}",
1057            var_rem,
1058            var_data
1059        );
1060    }
1061
1062    #[test]
1063    fn test_decompose_additive_multi_sample() {
1064        let n = 3;
1065        let m = 100;
1066        let period = 2.0;
1067        let argvals: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64 * 10.0).collect();
1068        let mut data = FdMatrix::zeros(n, m);
1069        for i in 0..n {
1070            let amp = (i + 1) as f64;
1071            for j in 0..m {
1072                data[(i, j)] =
1073                    1.0 + 0.1 * argvals[j] + amp * (2.0 * PI * argvals[j] / period).sin();
1074            }
1075        }
1076        let result = decompose_additive(&data, &argvals, period, "loess", 0.3, 2);
1077        assert_eq!(result.trend.shape(), (n, m));
1078        assert_eq!(result.seasonal.shape(), (n, m));
1079        assert_eq!(result.remainder.shape(), (n, m));
1080    }
1081
1082    #[test]
1083    fn test_decompose_multiplicative_basic() {
1084        let m = 200;
1085        let period = 2.0;
1086        let argvals: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64 * 10.0).collect();
1087        let data_vec: Vec<f64> = argvals
1088            .iter()
1089            .map(|&t| (2.0 + 0.1 * t) * (1.0 + 0.3 * (2.0 * PI * t / period).sin()))
1090            .collect();
1091        let data = FdMatrix::from_column_major(data_vec, 1, m).unwrap();
1092        let result = decompose_multiplicative(&data, &argvals, period, "loess", 0.3, 3);
1093        assert_eq!(result.trend.ncols(), m);
1094        assert_eq!(result.seasonal.ncols(), m);
1095        assert_eq!(result.remainder.ncols(), m);
1096        assert_eq!(result.method, "multiplicative");
1097        let seasonal_vec: Vec<f64> = (0..m).map(|j| result.seasonal[(0, j)]).collect();
1098        let mean_seasonal: f64 = seasonal_vec.iter().sum::<f64>() / m as f64;
1099        assert!(
1100            (mean_seasonal - 1.0).abs() < 0.5,
1101            "Mean seasonal factor: {}",
1102            mean_seasonal
1103        );
1104    }
1105
1106    #[test]
1107    fn test_decompose_multiplicative_non_positive_data() {
1108        let m = 100;
1109        let period = 2.0;
1110        let argvals: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64 * 10.0).collect();
1111        let data_vec: Vec<f64> = argvals
1112            .iter()
1113            .map(|&t| -1.0 + (2.0 * PI * t / period).sin())
1114            .collect();
1115        let data = FdMatrix::from_column_major(data_vec, 1, m).unwrap();
1116        let result = decompose_multiplicative(&data, &argvals, period, "loess", 0.3, 2);
1117        assert_eq!(result.trend.ncols(), m);
1118        assert_eq!(result.seasonal.ncols(), m);
1119        for j in 0..m {
1120            let s = result.seasonal[(0, j)];
1121            assert!(s.is_finite(), "Seasonal should be finite");
1122        }
1123    }
1124
1125    #[test]
1126    fn test_decompose_multiplicative_vs_additive() {
1127        let m = 200;
1128        let period = 2.0;
1129        let argvals: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64 * 10.0).collect();
1130        let data_vec: Vec<f64> = argvals
1131            .iter()
1132            .map(|&t| 5.0 + (2.0 * PI * t / period).sin())
1133            .collect();
1134        let data = FdMatrix::from_column_major(data_vec, 1, m).unwrap();
1135        let add_result = decompose_additive(&data, &argvals, period, "loess", 0.3, 3);
1136        let mult_result = decompose_multiplicative(&data, &argvals, period, "loess", 0.3, 3);
1137        assert_eq!(add_result.seasonal.ncols(), m);
1138        assert_eq!(mult_result.seasonal.ncols(), m);
1139        let add_seasonal_vec: Vec<f64> = (0..m).map(|j| add_result.seasonal[(0, j)]).collect();
1140        let add_mean: f64 = add_seasonal_vec.iter().sum::<f64>() / m as f64;
1141        let mult_seasonal_vec: Vec<f64> = (0..m).map(|j| mult_result.seasonal[(0, j)]).collect();
1142        let mult_mean: f64 = mult_seasonal_vec.iter().sum::<f64>() / m as f64;
1143        assert!(
1144            add_mean.abs() < mult_mean,
1145            "Additive mean {} vs mult mean {}",
1146            add_mean,
1147            mult_mean
1148        );
1149    }
1150
1151    #[test]
1152    fn test_decompose_multiplicative_edge_cases() {
1153        let empty = FdMatrix::zeros(0, 0);
1154        let result = decompose_multiplicative(&empty, &[], 2.0, "loess", 0.3, 2);
1155        assert_eq!(result.trend.len(), 0);
1156        let m = 5;
1157        let argvals: Vec<f64> = (0..m).map(|i| i as f64).collect();
1158        let data = FdMatrix::from_column_major(vec![1.0, 2.0, 3.0, 2.0, 1.0], 1, m).unwrap();
1159        let result = decompose_multiplicative(&data, &argvals, 2.0, "loess", 0.3, 1);
1160        assert_eq!(result.remainder.ncols(), m);
1161    }
1162
1163    #[test]
1164    fn test_stl_decompose_basic() {
1165        let period = 12;
1166        let n_cycles = 10;
1167        let m = period * n_cycles;
1168        let data_vec: Vec<f64> = (0..m)
1169            .map(|i| {
1170                let t = i as f64;
1171                0.01 * t + (2.0 * PI * t / period as f64).sin()
1172            })
1173            .collect();
1174        let data = FdMatrix::from_column_major(data_vec.clone(), 1, m).unwrap();
1175        let result = stl_decompose(&data, period, None, None, None, false, None, None);
1176        assert_eq!(result.trend.ncols(), m);
1177        assert_eq!(result.seasonal.ncols(), m);
1178        assert_eq!(result.remainder.ncols(), m);
1179        assert_eq!(result.period, period);
1180        for j in 0..m {
1181            let reconstructed =
1182                result.trend[(0, j)] + result.seasonal[(0, j)] + result.remainder[(0, j)];
1183            assert!(
1184                (reconstructed - data_vec[j]).abs() < 1e-8,
1185                "Reconstruction error at {}: {} vs {}",
1186                j,
1187                reconstructed,
1188                data_vec[j]
1189            );
1190        }
1191    }
1192
1193    #[test]
1194    fn test_stl_decompose_robust() {
1195        let period = 12;
1196        let n_cycles = 10;
1197        let m = period * n_cycles;
1198        let mut data_vec: Vec<f64> = (0..m)
1199            .map(|i| {
1200                let t = i as f64;
1201                0.01 * t + (2.0 * PI * t / period as f64).sin()
1202            })
1203            .collect();
1204        data_vec[30] += 10.0;
1205        data_vec[60] += 10.0;
1206        let data = FdMatrix::from_column_major(data_vec, 1, m).unwrap();
1207        let result = stl_decompose(&data, period, None, None, None, true, None, Some(5));
1208        assert!(
1209            result.weights[(0, 30)] < 1.0,
1210            "Weight at outlier should be < 1.0: {}",
1211            result.weights[(0, 30)]
1212        );
1213        assert!(
1214            result.weights[(0, 60)] < 1.0,
1215            "Weight at outlier should be < 1.0: {}",
1216            result.weights[(0, 60)]
1217        );
1218        let non_outlier_weight = result.weights[(0, 15)];
1219        assert!(
1220            non_outlier_weight > result.weights[(0, 30)],
1221            "Non-outlier weight {} should be > outlier weight {}",
1222            non_outlier_weight,
1223            result.weights[(0, 30)]
1224        );
1225    }
1226
1227    #[test]
1228    fn test_stl_decompose_default_params() {
1229        let period = 10;
1230        let m = period * 8;
1231        let data_vec: Vec<f64> = (0..m)
1232            .map(|i| (2.0 * PI * i as f64 / period as f64).sin())
1233            .collect();
1234        let data = FdMatrix::from_column_major(data_vec, 1, m).unwrap();
1235        let result = stl_decompose(&data, period, None, None, None, false, None, None);
1236        assert_eq!(result.trend.ncols(), m);
1237        assert_eq!(result.seasonal.ncols(), m);
1238        assert!(result.s_window >= 3);
1239        assert!(result.t_window >= 3);
1240        assert_eq!(result.inner_iterations, 2);
1241        assert_eq!(result.outer_iterations, 1);
1242    }
1243
1244    #[test]
1245    fn test_stl_decompose_invalid() {
1246        let data = FdMatrix::from_column_major(vec![1.0, 2.0], 1, 2).unwrap();
1247        let result = stl_decompose(&data, 1, None, None, None, false, None, None);
1248        assert_eq!(result.s_window, 0);
1249        let data = FdMatrix::from_column_major(vec![1.0, 2.0, 3.0], 1, 3).unwrap();
1250        let result = stl_decompose(&data, 5, None, None, None, false, None, None);
1251        assert_eq!(result.s_window, 0);
1252        let data = FdMatrix::zeros(0, 0);
1253        let result = stl_decompose(&data, 10, None, None, None, false, None, None);
1254        assert_eq!(result.trend.len(), 0);
1255    }
1256
1257    #[test]
1258    fn test_stl_fdata() {
1259        let n = 3;
1260        let period = 10;
1261        let m = period * 5;
1262        let argvals: Vec<f64> = (0..m).map(|i| i as f64).collect();
1263        let mut data = FdMatrix::zeros(n, m);
1264        for i in 0..n {
1265            let amp = (i + 1) as f64;
1266            for j in 0..m {
1267                data[(i, j)] = amp * (2.0 * PI * argvals[j] / period as f64).sin();
1268            }
1269        }
1270        let result = stl_fdata(&data, &argvals, period, None, None, false);
1271        assert_eq!(result.trend.shape(), (n, m));
1272        assert_eq!(result.seasonal.shape(), (n, m));
1273        assert_eq!(result.remainder.shape(), (n, m));
1274        for i in 0..n {
1275            for j in 0..m {
1276                let reconstructed =
1277                    result.trend[(i, j)] + result.seasonal[(i, j)] + result.remainder[(i, j)];
1278                assert!(
1279                    (reconstructed - data[(i, j)]).abs() < 1e-8,
1280                    "Reconstruction error for sample {} at {}: {} vs {}",
1281                    i,
1282                    j,
1283                    reconstructed,
1284                    data[(i, j)]
1285                );
1286            }
1287        }
1288    }
1289
1290    #[test]
1291    fn test_stl_decompose_multi_sample() {
1292        let n = 5;
1293        let period = 10;
1294        let m = period * 6;
1295        let mut data = FdMatrix::zeros(n, m);
1296        for i in 0..n {
1297            let offset = i as f64 * 0.5;
1298            for j in 0..m {
1299                data[(i, j)] =
1300                    offset + 0.01 * j as f64 + (2.0 * PI * j as f64 / period as f64).sin();
1301            }
1302        }
1303        let result = stl_decompose(&data, period, None, None, None, false, None, None);
1304        assert_eq!(result.trend.shape(), (n, m));
1305        assert_eq!(result.seasonal.shape(), (n, m));
1306        assert_eq!(result.remainder.shape(), (n, m));
1307        assert_eq!(result.weights.shape(), (n, m));
1308    }
1309
1310    #[test]
1311    fn test_detrend_diff_order2() {
1312        let m = 50;
1313        let data_vec: Vec<f64> = (0..m).map(|i| (i as f64).powi(2)).collect();
1314        let data = FdMatrix::from_column_major(data_vec, 1, m).unwrap();
1315        let result = detrend_diff(&data, 2);
1316        for j in 0..m - 2 {
1317            assert!(
1318                (result.detrended[(0, j)] - 2.0).abs() < 1e-10,
1319                "Second diff at {}: expected 2.0, got {}",
1320                j,
1321                result.detrended[(0, j)]
1322            );
1323        }
1324    }
1325
1326    #[test]
1327    fn test_detrend_polynomial_degree3() {
1328        let m = 100;
1329        let argvals: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64 * 5.0).collect();
1330        let data_vec: Vec<f64> = argvals
1331            .iter()
1332            .map(|&t| 1.0 + 2.0 * t - 0.5 * t * t + 0.1 * t * t * t)
1333            .collect();
1334        let data = FdMatrix::from_column_major(data_vec, 1, m).unwrap();
1335        let result = detrend_polynomial(&data, &argvals, 3);
1336        assert_eq!(result.method, "polynomial(3)");
1337        assert!(result.coefficients.is_some());
1338        let max_detrend: f64 = (0..m)
1339            .map(|j| result.detrended[(0, j)].abs())
1340            .fold(0.0, f64::max);
1341        assert!(
1342            max_detrend < 0.1,
1343            "Pure cubic should be nearly zero after degree-3 detrend: {}",
1344            max_detrend
1345        );
1346    }
1347
1348    #[test]
1349    fn test_detrend_loess_invalid() {
1350        let data = FdMatrix::from_column_major(vec![1.0, 2.0, 3.0, 4.0, 5.0], 1, 5).unwrap();
1351        let argvals = vec![0.0, 1.0, 2.0, 3.0, 4.0];
1352        let result = detrend_loess(&data, &argvals, -0.1, 1);
1353        assert_eq!(result.detrended.as_slice(), data.as_slice());
1354        let data2 = FdMatrix::from_column_major(vec![1.0, 2.0], 1, 2).unwrap();
1355        let result = detrend_loess(&data2, &[0.0, 1.0], 0.3, 1);
1356        assert_eq!(result.detrended.as_slice(), &[1.0, 2.0]);
1357    }
1358}