vector/
stats.rs

1#![allow(missing_docs)]
2/// Exponentially Weighted Moving Average
3#[derive(Clone, Copy, Debug)]
4pub struct Ewma {
5    average: Option<f64>,
6    alpha: f64,
7}
8
9impl Ewma {
10    pub const fn new(alpha: f64) -> Self {
11        let average = None;
12        Self { average, alpha }
13    }
14
15    pub const fn average(&self) -> Option<f64> {
16        self.average
17    }
18
19    /// Update the current average and return it for convenience
20    pub fn update(&mut self, point: f64) -> f64 {
21        let average = match self.average {
22            None => point,
23            Some(avg) => point.mul_add(self.alpha, avg * (1.0 - self.alpha)),
24        };
25        self.average = Some(average);
26        average
27    }
28}
29
30/// Exponentially Weighted Moving Average that starts with a default average value
31#[derive(Clone, Copy, Debug)]
32pub struct EwmaDefault {
33    average: f64,
34    alpha: f64,
35}
36
37impl EwmaDefault {
38    pub const fn new(alpha: f64, initial_value: f64) -> Self {
39        Self {
40            average: initial_value,
41            alpha,
42        }
43    }
44
45    pub const fn average(&self) -> f64 {
46        self.average
47    }
48
49    /// Update the current average and return it for convenience
50    pub fn update(&mut self, point: f64) -> f64 {
51        self.average = point.mul_add(self.alpha, self.average * (1.0 - self.alpha));
52        self.average
53    }
54}
55
56/// Exponentially Weighted Moving Average with variance calculation
57#[derive(Clone, Copy, Debug)]
58pub struct EwmaVar {
59    state: Option<MeanVariance>,
60    alpha: f64,
61}
62
63#[derive(Clone, Copy, Debug, PartialEq)]
64pub struct MeanVariance {
65    pub mean: f64,
66    pub variance: f64,
67}
68
69impl EwmaVar {
70    pub const fn new(alpha: f64) -> Self {
71        let state = None;
72        Self { state, alpha }
73    }
74
75    pub const fn state(&self) -> Option<MeanVariance> {
76        self.state
77    }
78
79    #[cfg(test)]
80    pub fn average(&self) -> Option<f64> {
81        self.state.map(|state| state.mean)
82    }
83
84    #[cfg(test)]
85    pub fn variance(&self) -> Option<f64> {
86        self.state.map(|state| state.variance)
87    }
88
89    /// Update the current average and variance, and return them for convenience
90    pub fn update(&mut self, point: f64) -> MeanVariance {
91        let (mean, variance) = match self.state {
92            None => (point, 0.0),
93            Some(state) => {
94                let difference = point - state.mean;
95                let increment = self.alpha * difference;
96                (
97                    state.mean + increment,
98                    (1.0 - self.alpha) * difference.mul_add(increment, state.variance),
99                )
100            }
101        };
102        let state = MeanVariance { mean, variance };
103        self.state = Some(state);
104        state
105    }
106}
107
108/// Simple unweighted arithmetic mean
109#[derive(Clone, Copy, Debug, Default)]
110pub struct Mean {
111    mean: f64,
112    count: usize,
113}
114
115impl Mean {
116    /// Update the and return the current average
117    pub fn update(&mut self, point: f64) {
118        self.count += 1;
119        self.mean += (point - self.mean) / self.count as f64;
120    }
121
122    pub const fn average(&self) -> Option<f64> {
123        match self.count {
124            0 => None,
125            _ => Some(self.mean),
126        }
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133
134    #[test]
135    fn mean_update_works() {
136        let mut mean = Mean::default();
137        assert_eq!(mean.average(), None);
138        mean.update(0.0);
139        assert_eq!(mean.average(), Some(0.0));
140        mean.update(2.0);
141        assert_eq!(mean.average(), Some(1.0));
142        mean.update(4.0);
143        assert_eq!(mean.average(), Some(2.0));
144    }
145
146    #[test]
147    fn ewma_update_works() {
148        let mut mean = Ewma::new(0.5);
149        assert_eq!(mean.average(), None);
150        mean.update(2.0);
151        assert_eq!(mean.average(), Some(2.0));
152        mean.update(2.0);
153        assert_eq!(mean.average(), Some(2.0));
154        mean.update(1.0);
155        assert_eq!(mean.average(), Some(1.5));
156        mean.update(2.0);
157        assert_eq!(mean.average(), Some(1.75));
158
159        assert_eq!(mean.average, Some(1.75));
160    }
161
162    #[test]
163    fn ewma_variance_update_works() {
164        let mut mean = EwmaVar::new(0.5);
165        assert_eq!(mean.average(), None);
166        assert_eq!(mean.variance(), None);
167        mean.update(2.0);
168        assert_eq!(mean.average(), Some(2.0));
169        assert_eq!(mean.variance(), Some(0.0));
170        mean.update(2.0);
171        assert_eq!(mean.average(), Some(2.0));
172        assert_eq!(mean.variance(), Some(0.0));
173        mean.update(1.0);
174        assert_eq!(mean.average(), Some(1.5));
175        assert_eq!(mean.variance(), Some(0.25));
176        mean.update(2.0);
177        assert_eq!(mean.average(), Some(1.75));
178        assert_eq!(mean.variance(), Some(0.1875));
179
180        assert_eq!(
181            mean.state,
182            Some(MeanVariance {
183                mean: 1.75,
184                variance: 0.1875
185            })
186        );
187    }
188}
vector/stats.rs

vector/
stats.rs
