Skip to main content

wowlab_types/stats/summary/
series.rs

1use super::Streaming;
2
3const MIN_POINTS_FOR_PAIR_STATS: usize = 2;
4const SQUARE_EXP: i32 = 2;
5const HALVING_DIVISOR: f64 = 2.0;
6
7#[derive(Clone, Debug)]
8pub struct LinearRegression {
9    pub slope: f64,
10    pub intercept: f64,
11    pub r_squared: f64,
12}
13
14/// Compute sample covariance between two slices.
15#[must_use]
16#[expect(
17    clippy::cast_precision_loss,
18    reason = "the slice length is converted to the f64 statistics domain"
19)]
20pub fn covariance(x: &[f64], y: &[f64]) -> f64 {
21    if x.len() != y.len() || x.len() < MIN_POINTS_FOR_PAIR_STATS {
22        return f64::NAN;
23    }
24
25    let n = x.len() as f64;
26    let x_mean: f64 = x.iter().sum::<f64>() / n;
27    let y_mean: f64 = y.iter().sum::<f64>() / n;
28
29    let cov: f64 = x
30        .iter()
31        .zip(y.iter())
32        .map(|(xi, yi)| (xi - x_mean) * (yi - y_mean))
33        .sum();
34
35    cov / (n - 1.0)
36}
37
38/// Compute Pearson correlation coefficient.
39#[must_use]
40pub fn correlation(x: &[f64], y: &[f64]) -> f64 {
41    if x.len() != y.len() || x.len() < MIN_POINTS_FOR_PAIR_STATS {
42        return f64::NAN;
43    }
44
45    let x_stats: Streaming = x.iter().copied().collect();
46    let y_stats: Streaming = y.iter().copied().collect();
47
48    let x_std = x_stats.std_dev();
49    let y_std = y_stats.std_dev();
50
51    if x_std.abs() < f64::EPSILON || y_std.abs() < f64::EPSILON {
52        return f64::NAN;
53    }
54
55    covariance(x, y) / (x_std * y_std)
56}
57
58/// Compute simple linear regression using least squares.
59#[must_use]
60pub fn linear_regression(x: &[f64], y: &[f64]) -> Option<LinearRegression> {
61    if x.len() < MIN_POINTS_FOR_PAIR_STATS || x.len() != y.len() {
62        return None;
63    }
64
65    let x_stats: Streaming = x.iter().copied().collect();
66    let y_stats: Streaming = y.iter().copied().collect();
67
68    let x_var = x_stats.variance();
69
70    if x_var.abs() < f64::EPSILON {
71        return None;
72    }
73
74    let cov = covariance(x, y);
75    let slope = cov / x_var;
76    let intercept = y_stats.mean() - slope * x_stats.mean();
77
78    let y_mean = y_stats.mean();
79    let ss_tot: f64 = y.iter().map(|yi| (yi - y_mean).powi(SQUARE_EXP)).sum();
80    let ss_res: f64 = x
81        .iter()
82        .zip(y.iter())
83        .map(|(xi, yi)| {
84            let predicted = slope * xi + intercept;
85
86            (yi - predicted).powi(SQUARE_EXP)
87        })
88        .sum();
89
90    let r_squared = if ss_tot.abs() < f64::EPSILON {
91        1.0
92    } else {
93        1.0 - ss_res / ss_tot
94    };
95
96    Some(LinearRegression {
97        slope,
98        intercept,
99        r_squared,
100    })
101}
102
103/// Exponential moving average with smoothing factor alpha (0-1).
104#[must_use]
105pub fn ema(data: &[f64], alpha: f64) -> Vec<f64> {
106    if data.is_empty() {
107        return Vec::new();
108    }
109
110    let alpha = alpha.clamp(0.0, 1.0);
111    let mut result = Vec::with_capacity(data.len());
112    let mut prev = *data.first().unwrap_or(&0.0);
113
114    result.push(prev);
115
116    // #t(rust_unchecked_indexing) data has at least one element per .get(0) check above
117    for &x in &data[1..] {
118        prev = alpha * x + (1.0 - alpha) * prev;
119        result.push(prev);
120    }
121
122    result
123}
124
125/// EMA with span-based smoothing: alpha = 2 / (span + 1).
126#[must_use]
127#[expect(
128    clippy::cast_precision_loss,
129    reason = "the smoothing span is converted to the f64 calculation domain"
130)]
131pub fn ema_span(data: &[f64], span: usize) -> Vec<f64> {
132    if span == 0 {
133        return data.to_vec();
134    }
135
136    ema(data, HALVING_DIVISOR / (span as f64 + 1.0))
137}
138
139/// Simple moving average with window size.
140#[must_use]
141#[expect(
142    clippy::cast_precision_loss,
143    reason = "bounded moving-average window lengths are converted to f64 divisors"
144)]
145pub fn sma(data: &[f64], window: usize) -> Vec<f64> {
146    if data.is_empty() || window == 0 {
147        return Vec::new();
148    }
149
150    let mut result = Vec::with_capacity(data.len());
151    let mut sum = 0.0;
152
153    for (i, &x) in data.iter().enumerate() {
154        sum += x;
155
156        if i >= window {
157            sum -= data.get(i - window).copied().unwrap_or(0.0);
158            result.push(sum / window as f64);
159        } else {
160            result.push(sum / (i + 1) as f64);
161        }
162    }
163
164    result
165}