hopr_primitive_types/
sma.rs

1use std::{
2    collections::VecDeque,
3    fmt::{Display, Formatter},
4    iter::Sum,
5    marker::PhantomData,
6    ops::{AddAssign, Div, SubAssign},
7};
8
9/// Simple Moving Average trait.
10///
11/// The second-most useful filter type, bested only by coffee filters.
12pub trait SMA<T> {
13    /// Pushes a sample.
14    fn push(&mut self, sample: T);
15
16    /// Calculates the moving average value.
17    /// Returns `None` if no samples were added.
18    fn average(&self) -> Option<T>;
19
20    /// Returns the window size.
21    fn window_size(&self) -> usize;
22
23    /// Returns the number of elements in the window.
24    /// This value is always between 0 and `window_size()`.
25    fn len(&self) -> usize;
26
27    /// Indicates whether the window is fully occupied
28    /// with samples.
29    fn is_window_full(&self) -> bool {
30        self.len() == self.window_size()
31    }
32
33    /// Indicates whether there are no samples.
34    fn is_empty(&self) -> bool;
35
36    /// Clears all the samples.
37    fn clear(&mut self);
38}
39
40/// Basic implementation of Simple Moving Average (SMA).
41///
42/// The maximum window size is bound by 2^32 - 1.
43/// Useful mainly for floating-point types, as it does not accumulate floating point error with each sample.
44/// Requires `O(N)` of memory and `O(N)` for average computation, `N` being window size.
45/// The divisor argument `D` is used only for such types `T` that do not implement `From<u32>` (such as `Duration`,...).
46#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
47pub struct NoSumSMA<T, D = T> {
48    window: VecDeque<T>,
49    window_size: usize,
50    _div: PhantomData<D>,
51}
52
53impl<T, D> SMA<T> for NoSumSMA<T, D>
54where
55    T: for<'a> Sum<&'a T> + Div<D, Output = T>,
56    D: From<u32>,
57{
58    fn push(&mut self, sample: T) {
59        if self.is_window_full() {
60            self.window.pop_front();
61        }
62        self.window.push_back(sample);
63    }
64
65    fn average(&self) -> Option<T> {
66        if !self.is_empty() {
67            Some(self.window.iter().sum::<T>() / D::from(self.window.len() as u32))
68        } else {
69            None
70        }
71    }
72
73    fn window_size(&self) -> usize {
74        self.window_size
75    }
76
77    fn len(&self) -> usize {
78        self.window.len()
79    }
80
81    fn is_empty(&self) -> bool {
82        self.window.is_empty()
83    }
84
85    fn clear(&mut self) {
86        self.window.clear()
87    }
88}
89
90impl<T, D> NoSumSMA<T, D>
91where
92    T: for<'a> Sum<&'a T> + Div<D, Output = T>,
93    D: From<u32>,
94{
95    /// Creates an empty SMA instance with the given window size.
96    /// The maximum window size is u32::MAX and must be greater than 1.
97    pub fn new(window_size: usize) -> Self {
98        assert!(window_size > 1, "window size must be greater than 1");
99        Self {
100            window: VecDeque::with_capacity(window_size),
101            window_size,
102            _div: PhantomData,
103        }
104    }
105
106    /// Creates SMA instance given window size and some initial samples.
107    pub fn new_with_samples<I>(window_size: usize, initial_samples: I) -> Self
108    where
109        I: IntoIterator<Item = T>,
110    {
111        let mut ret = Self::new(window_size);
112        initial_samples.into_iter().for_each(|s| ret.push(s));
113        ret
114    }
115}
116
117impl<T, D> Display for NoSumSMA<T, D>
118where
119    T: for<'a> Sum<&'a T> + Div<D, Output = T> + Default + Display,
120    D: From<u32>,
121{
122    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
123        write!(f, "{}", self.average().unwrap_or_default())
124    }
125}
126
127/// Basic implementation of Simple Moving Average (SMA).
128///
129/// The maximum window size is bound by 2^32 - 1.
130/// Useful mainly for integer types, as it does accumulate floating point error with each sample.
131/// Requires `O(N)` of memory and `O(1)` for average computation, `N` being window size.
132/// The divisor argument `D` is used only for such types `T` that do not implement `From<u32>` (such as `Duration`,...).
133#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)]
134pub struct SingleSumSMA<T, D = T> {
135    window: VecDeque<T>,
136    window_size: usize,
137    sum: T,
138    _div: PhantomData<D>,
139}
140
141impl<T, D> SMA<T> for SingleSumSMA<T, D>
142where
143    T: AddAssign + SubAssign + Div<D, Output = T> + Copy,
144    D: From<u32>,
145{
146    fn push(&mut self, sample: T) {
147        self.sum += sample;
148
149        if self.is_window_full() {
150            if let Some(shifted_sample) = self.window.pop_front() {
151                self.sum -= shifted_sample;
152            }
153        }
154
155        self.window.push_back(sample);
156    }
157
158    fn average(&self) -> Option<T> {
159        if !self.is_empty() {
160            Some(self.sum / D::from(self.window.len() as u32))
161        } else {
162            None
163        }
164    }
165
166    fn window_size(&self) -> usize {
167        self.window_size
168    }
169
170    fn len(&self) -> usize {
171        self.window.len()
172    }
173
174    fn is_empty(&self) -> bool {
175        self.window.is_empty()
176    }
177
178    fn clear(&mut self) {
179        self.sum -= self.sum;
180        self.window.clear();
181    }
182}
183
184impl<T, D> SingleSumSMA<T, D>
185where
186    T: AddAssign + SubAssign + Div<D, Output = T> + Copy + Default,
187    D: From<u32>,
188{
189    /// Creates an empty SMA instance with the given window size.
190    /// The maximum window size is u32::MAX and must be greater than 1.
191    pub fn new(window_size: usize) -> Self {
192        assert!(window_size > 1, "window size must be greater than 1");
193        Self {
194            window: VecDeque::with_capacity(window_size),
195            window_size,
196            sum: T::default(),
197            _div: PhantomData,
198        }
199    }
200
201    /// Creates SMA instance given window size and some initial samples.
202    pub fn new_with_samples<I>(window_size: usize, initial_samples: I) -> Self
203    where
204        I: IntoIterator<Item = T>,
205    {
206        let mut ret = Self::new(window_size);
207        initial_samples.into_iter().for_each(|s| ret.push(s));
208        ret
209    }
210}
211
212impl<T, D> Display for SingleSumSMA<T, D>
213where
214    T: AddAssign + SubAssign + Div<D, Output = T> + Copy + Display + Default,
215    D: From<u32>,
216{
217    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
218        write!(f, "{}", self.average().unwrap_or_default())
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    use crate::sma::{NoSumSMA, SMA, SingleSumSMA};
225
226    fn test_sma<T: SMA<u32>>(mut sma: T, window_size: usize) {
227        assert_eq!(window_size, sma.window_size(), "invalid windows size");
228        assert!(sma.average().is_none(), "invalid empty average");
229
230        assert!(sma.is_empty(), "should be empty");
231        assert_eq!(0, sma.len(), "len is invalid");
232
233        sma.push(0);
234        sma.push(1);
235        sma.push(2);
236        sma.push(3);
237
238        assert_eq!(Some(2), sma.average(), "invalid average");
239        assert_eq!(3, sma.window_size(), "window size is invalid");
240
241        assert!(!sma.is_empty(), "should not be empty");
242        assert_eq!(3, sma.len(), "len is invalid");
243    }
244
245    #[test]
246    fn test_nosum_sma_should_calculate_avg_correctly() {
247        test_sma(NoSumSMA::<u32, u32>::new(3), 3);
248    }
249
250    #[test]
251    fn test_single_sum_sma_should_calculate_avg_correctly() {
252        test_sma(SingleSumSMA::<u32, u32>::new(3), 3);
253    }
254}
hopr_primitive_types/sma.rs

hopr_primitive_types/
sma.rs
