Skip to main content

rill_ml/stats/
ew_mean.rs

1//! Exponentially weighted mean.
2//!
3//! Time complexity per update: `O(1)`. Space complexity: `O(1)`.
4//!
5//! The update rule is `mean = alpha * x + (1 - alpha) * mean`. The first
6//! observation seeds the mean directly.
7
8use crate::error::{RillError, checked_increment, ensure_finite};
9use crate::traits::OnlineStatistic;
10
11/// Exponentially weighted moving average.
12///
13/// `alpha` must satisfy `0 < alpha <= 1`. Smaller values give more weight to
14/// the past; `alpha = 1` reduces to a [`LastValue`](crate::stats::extrema)-like
15/// tracker.
16///
17/// # Examples
18///
19/// ```
20/// use rill_ml::stats::ExponentiallyWeightedMean;
21/// use rill_ml::OnlineStatistic;
22///
23/// let mut ew = ExponentiallyWeightedMean::new(0.5).unwrap();
24/// ew.update(10.0).unwrap();
25/// ew.update(20.0).unwrap();
26/// // 10.0, then 0.5*20 + 0.5*10 = 15.0
27/// assert!((ew.value() - 15.0).abs() < 1e-12);
28/// ```
29#[derive(Debug, Clone)]
30#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
31pub struct ExponentiallyWeightedMean {
32    alpha: f64,
33    mean: f64,
34    count: u64,
35}
36
37impl ExponentiallyWeightedMean {
38    /// Create a new exponentially weighted mean accumulator.
39    ///
40    /// Returns an error if `alpha` is not in `(0, 1]`.
41    pub fn new(alpha: f64) -> Result<Self, RillError> {
42        ensure_finite("alpha", alpha)?;
43        if alpha <= 0.0 || alpha > 1.0 {
44            return Err(RillError::InvalidParameter {
45                name: "alpha",
46                value: alpha,
47            });
48        }
49        Ok(Self {
50            alpha,
51            mean: 0.0,
52            count: 0,
53        })
54    }
55
56    /// The configured alpha.
57    pub const fn alpha(&self) -> f64 {
58        self.alpha
59    }
60
61    /// Current weighted mean, or `0.0` if no observations have been seen.
62    pub const fn value(&self) -> f64 {
63        self.mean
64    }
65
66    /// Number of observations seen so far.
67    pub const fn count(&self) -> u64 {
68        self.count
69    }
70}
71
72impl OnlineStatistic for ExponentiallyWeightedMean {
73    fn update(&mut self, value: f64) -> Result<(), RillError> {
74        ensure_finite("value", value)?;
75        let next_count = checked_increment(self.count, "EW mean sample")?;
76        let next_mean = if self.count == 0 {
77            value
78        } else {
79            self.alpha * value + (1.0 - self.alpha) * self.mean
80        };
81        ensure_finite("EW mean", next_mean)?;
82        self.mean = next_mean;
83        self.count = next_count;
84        Ok(())
85    }
86
87    fn samples_seen(&self) -> u64 {
88        self.count
89    }
90
91    fn reset(&mut self) {
92        self.mean = 0.0;
93        self.count = 0;
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn first_sample_seeds_mean() {
103        let mut ew = ExponentiallyWeightedMean::new(0.3).unwrap();
104        ew.update(10.0).unwrap();
105        assert!((ew.value() - 10.0).abs() < 1e-12);
106    }
107
108    #[test]
109    fn weighted_update_matches_formula() {
110        let mut ew = ExponentiallyWeightedMean::new(0.5).unwrap();
111        ew.update(10.0).unwrap();
112        ew.update(20.0).unwrap();
113        assert!((ew.value() - 15.0).abs() < 1e-12);
114    }
115
116    #[test]
117    fn alpha_one_tracks_last_value() {
118        let mut ew = ExponentiallyWeightedMean::new(1.0).unwrap();
119        ew.update(3.0).unwrap();
120        ew.update(7.0).unwrap();
121        assert!((ew.value() - 7.0).abs() < 1e-12);
122    }
123
124    #[test]
125    fn invalid_alpha_rejected() {
126        assert!(ExponentiallyWeightedMean::new(0.0).is_err());
127        assert!(ExponentiallyWeightedMean::new(-0.1).is_err());
128        assert!(ExponentiallyWeightedMean::new(1.5).is_err());
129        assert!(ExponentiallyWeightedMean::new(f64::NAN).is_err());
130    }
131
132    #[test]
133    fn reset_clears_state() {
134        let mut ew = ExponentiallyWeightedMean::new(0.5).unwrap();
135        ew.update(10.0).unwrap();
136        ew.reset();
137        assert_eq!(ew.count(), 0);
138        assert_eq!(ew.value(), 0.0);
139    }
140}