1use 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#[derive(Debug, Clone)]
22pub struct TrendResult {
23 pub trend: FdMatrix,
25 pub detrended: FdMatrix,
27 pub method: Cow<'static, str>,
29 pub coefficients: Option<FdMatrix>,
32 pub rss: Vec<f64>,
34 pub n_params: usize,
36}
37
38impl TrendResult {
39 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#[derive(Debug, Clone)]
55pub struct DecomposeResult {
56 pub trend: FdMatrix,
58 pub seasonal: FdMatrix,
60 pub remainder: FdMatrix,
62 pub period: f64,
64 pub method: Cow<'static, str>,
66}
67
68pub 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 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 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 * β
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
225fn 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
244pub 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
284pub 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
308pub 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
348pub 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
418pub 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
470pub 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#[derive(Debug, Clone)]
531pub struct StlResult {
532 pub trend: FdMatrix,
534 pub seasonal: FdMatrix,
536 pub remainder: FdMatrix,
538 pub weights: FdMatrix,
540 pub period: usize,
542 pub s_window: usize,
544 pub t_window: usize,
546 pub inner_iterations: usize,
548 pub outer_iterations: usize,
550}
551
552pub 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
804pub 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}