hopr_primitive_types/
sma.rs

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

hopr_primitive_types/
sma.rs
