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

hopr_primitive_types/
sma.rs
