vector_common/stats/
mod.rs

1#![allow(missing_docs)]
2
3mod ewma_gauge;
4mod time_ewma;
5
6pub use ewma_gauge::{EwmaGauge, TimeEwmaGauge};
7pub use time_ewma::TimeEwma;
8
9use std::sync::atomic::Ordering;
10
11use crate::atomic::AtomicF64;
12
13/// The default alpha parameter used when constructing EWMA-backed gauges.
14pub const DEFAULT_EWMA_ALPHA: f64 = 0.9;
15
16/// Exponentially Weighted Moving Average
17#[derive(Clone, Copy, Debug)]
18pub struct Ewma {
19    average: Option<f64>,
20    alpha: f64,
21}
22
23impl Ewma {
24    #[must_use]
25    pub const fn new(alpha: f64) -> Self {
26        let average = None;
27        Self { average, alpha }
28    }
29
30    #[must_use]
31    pub const fn average(&self) -> Option<f64> {
32        self.average
33    }
34
35    /// Update the current average and return it for convenience
36    pub fn update(&mut self, point: f64) -> f64 {
37        let average = match self.average {
38            None => point,
39            Some(avg) => point.mul_add(self.alpha, avg * (1.0 - self.alpha)),
40        };
41        self.average = Some(average);
42        average
43    }
44}
45
46/// Exponentially Weighted Moving Average that starts with a default average value
47#[derive(Clone, Copy, Debug)]
48pub struct EwmaDefault {
49    average: f64,
50    alpha: f64,
51}
52
53impl EwmaDefault {
54    #[must_use]
55    pub const fn new(alpha: f64, initial_value: f64) -> Self {
56        Self {
57            average: initial_value,
58            alpha,
59        }
60    }
61
62    #[must_use]
63    pub const fn average(&self) -> f64 {
64        self.average
65    }
66
67    /// Update the current average and return it for convenience
68    pub fn update(&mut self, point: f64) -> f64 {
69        self.average = point.mul_add(self.alpha, self.average * (1.0 - self.alpha));
70        self.average
71    }
72}
73
74/// Exponentially Weighted Moving Average with variance calculation
75#[derive(Clone, Copy, Debug)]
76pub struct EwmaVar {
77    state: Option<MeanVariance>,
78    alpha: f64,
79}
80
81#[derive(Clone, Copy, Debug, PartialEq)]
82pub struct MeanVariance {
83    pub mean: f64,
84    pub variance: f64,
85}
86
87impl EwmaVar {
88    #[must_use]
89    pub const fn new(alpha: f64) -> Self {
90        let state = None;
91        Self { state, alpha }
92    }
93
94    #[must_use]
95    pub const fn state(&self) -> Option<MeanVariance> {
96        self.state
97    }
98
99    #[cfg(test)]
100    #[must_use]
101    pub fn average(&self) -> Option<f64> {
102        self.state.map(|state| state.mean)
103    }
104
105    #[cfg(test)]
106    #[must_use]
107    pub fn variance(&self) -> Option<f64> {
108        self.state.map(|state| state.variance)
109    }
110
111    /// Update the current average and variance, and return them for convenience
112    pub fn update(&mut self, point: f64) -> MeanVariance {
113        let (mean, variance) = match self.state {
114            None => (point, 0.0),
115            Some(state) => {
116                let difference = point - state.mean;
117                let increment = self.alpha * difference;
118                (
119                    state.mean + increment,
120                    (1.0 - self.alpha) * difference.mul_add(increment, state.variance),
121                )
122            }
123        };
124        let state = MeanVariance { mean, variance };
125        self.state = Some(state);
126        state
127    }
128}
129
130/// Simple unweighted arithmetic mean
131#[derive(Clone, Copy, Debug, Default)]
132pub struct Mean {
133    mean: f64,
134    count: usize,
135}
136
137impl Mean {
138    /// Update the and return the current average
139    #[expect(
140        clippy::cast_precision_loss,
141        reason = "We have to convert count to f64 for the calculation, it's okay to lose precision for very large counts."
142    )]
143    pub fn update(&mut self, point: f64) {
144        self.count += 1;
145        self.mean += (point - self.mean) / self.count as f64;
146    }
147
148    #[must_use]
149    pub const fn average(&self) -> Option<f64> {
150        match self.count {
151            0 => None,
152            _ => Some(self.mean),
153        }
154    }
155}
156
157/// Atomic EWMA that uses an `AtomicF64` to store the current average.
158#[derive(Debug)]
159pub struct AtomicEwma {
160    average: AtomicF64,
161    alpha: f64,
162}
163
164impl AtomicEwma {
165    #[must_use]
166    pub fn new(alpha: f64) -> Self {
167        Self {
168            average: AtomicF64::new(f64::NAN),
169            alpha,
170        }
171    }
172
173    pub fn update(&self, point: f64) -> f64 {
174        let mut result = f64::NAN;
175        self.average
176            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
177                let average = if current.is_nan() {
178                    point
179                } else {
180                    point.mul_add(self.alpha, current * (1.0 - self.alpha))
181                };
182                result = average;
183                average
184            });
185        result
186    }
187
188    pub fn average(&self) -> Option<f64> {
189        let value = self.average.load(Ordering::Relaxed);
190        if value.is_nan() { None } else { Some(value) }
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197
198    #[test]
199    fn mean_update_works() {
200        let mut mean = Mean::default();
201        assert_eq!(mean.average(), None);
202        mean.update(0.0);
203        assert_eq!(mean.average(), Some(0.0));
204        mean.update(2.0);
205        assert_eq!(mean.average(), Some(1.0));
206        mean.update(4.0);
207        assert_eq!(mean.average(), Some(2.0));
208    }
209
210    #[test]
211    #[expect(clippy::float_cmp, reason = "none of the values will be rounded")]
212    fn ewma_update_works() {
213        let mut mean = Ewma::new(0.5);
214        assert_eq!(mean.average(), None);
215        assert_eq!(mean.update(2.0), 2.0);
216        assert_eq!(mean.average(), Some(2.0));
217        assert_eq!(mean.update(2.0), 2.0);
218        assert_eq!(mean.average(), Some(2.0));
219        assert_eq!(mean.update(1.0), 1.5);
220        assert_eq!(mean.average(), Some(1.5));
221        assert_eq!(mean.update(2.0), 1.75);
222        assert_eq!(mean.average(), Some(1.75));
223
224        assert_eq!(mean.average, Some(1.75));
225    }
226
227    #[test]
228    fn ewma_variance_update_works() {
229        let mut mean = EwmaVar::new(0.5);
230        assert_eq!(mean.average(), None);
231        assert_eq!(mean.variance(), None);
232        mean.update(2.0);
233        assert_eq!(mean.average(), Some(2.0));
234        assert_eq!(mean.variance(), Some(0.0));
235        mean.update(2.0);
236        assert_eq!(mean.average(), Some(2.0));
237        assert_eq!(mean.variance(), Some(0.0));
238        mean.update(1.0);
239        assert_eq!(mean.average(), Some(1.5));
240        assert_eq!(mean.variance(), Some(0.25));
241        mean.update(2.0);
242        assert_eq!(mean.average(), Some(1.75));
243        assert_eq!(mean.variance(), Some(0.1875));
244
245        assert_eq!(
246            mean.state,
247            Some(MeanVariance {
248                mean: 1.75,
249                variance: 0.1875
250            })
251        );
252    }
253
254    #[test]
255    #[expect(clippy::float_cmp, reason = "none of the values will be rounded")]
256    fn atomic_ewma_update_works() {
257        let ewma = AtomicEwma::new(0.5);
258        assert_eq!(ewma.average(), None);
259        assert_eq!(ewma.update(2.0), 2.0);
260        assert_eq!(ewma.average(), Some(2.0));
261        assert_eq!(ewma.update(1.0), 1.5);
262        assert_eq!(ewma.average(), Some(1.5));
263    }
264}
vector_common/stats/mod.rs

vector_common/stats/
mod.rs
