mantis_ta/indicators/trend/
wma.rs

1use crate::indicators::Indicator;
2use crate::types::Candle;
3use crate::utils::ringbuf::RingBuf;
4
5/// Weighted Moving Average over closing prices.
6///
7/// Assigns linearly increasing weights to recent prices, giving more importance
8/// to recent data compared to SMA.
9///
10/// # Examples
11/// ```rust
12/// use mantis_ta::indicators::{Indicator, WMA};
13/// use mantis_ta::types::Candle;
14///
15/// let candles: Vec<Candle> = [1.0, 2.0, 3.0, 4.0]
16///     .iter()
17///     .enumerate()
18///     .map(|(i, c)| Candle {
19///         timestamp: i as i64,
20///         open: *c,
21///         high: *c,
22///         low: *c,
23///         close: *c,
24///         volume: 0.0,
25///     })
26///     .collect();
27///
28/// let values = WMA::new(3).calculate(&candles);
29/// assert_eq!(values[0], None);
30/// assert_eq!(values[1], None);
31/// assert!(values[2].is_some());
32/// assert!(values[3].is_some());
33/// ```
34#[derive(Debug, Clone)]
35pub struct WMA {
36    period: usize,
37    window: RingBuf<f64>,
38    divisor: f64,
39}
40
41impl WMA {
42    pub fn new(period: usize) -> Self {
43        assert!(period > 0, "period must be > 0");
44        let divisor = (period * (period + 1) / 2) as f64;
45        Self {
46            period,
47            window: RingBuf::new(period, 0.0),
48            divisor,
49        }
50    }
51
52    #[inline]
53    fn update(&mut self, value: f64) -> Option<f64> {
54        self.window.push(value);
55
56        if self.window.len() < self.period {
57            return None;
58        }
59
60        let weighted_sum: f64 = self
61            .window
62            .iter()
63            .enumerate()
64            .map(|(i, &v)| v * ((i + 1) as f64))
65            .sum();
66
67        Some(weighted_sum / self.divisor)
68    }
69}
70
71impl Indicator for WMA {
72    type Output = f64;
73
74    fn next(&mut self, candle: &Candle) -> Option<Self::Output> {
75        self.update(candle.close)
76    }
77
78    fn reset(&mut self) {
79        self.window = RingBuf::new(self.period, 0.0);
80    }
81
82    fn warmup_period(&self) -> usize {
83        self.period
84    }
85
86    fn clone_boxed(&self) -> Box<dyn Indicator<Output = Self::Output>> {
87        Box::new(self.clone())
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn computes_wma_after_warmup() {
97        let mut wma = WMA::new(3);
98        let candles = [1.0, 2.0, 3.0, 4.0]
99            .iter()
100            .map(|c| Candle {
101                timestamp: 0,
102                open: *c,
103                high: *c,
104                low: *c,
105                close: *c,
106                volume: 0.0,
107            })
108            .collect::<Vec<_>>();
109
110        let mut outputs = Vec::new();
111        for c in &candles {
112            outputs.push(wma.next(c));
113        }
114
115        assert_eq!(outputs[0], None);
116        assert_eq!(outputs[1], None);
117        // WMA(3) at [1,2,3]: (1*1 + 2*2 + 3*3) / 6 = 14/6 = 2.333...
118        assert!(outputs[2].is_some());
119        let wma_val = outputs[2].unwrap();
120        assert!((wma_val - 2.333333).abs() < 0.0001);
121        assert!(outputs[3].is_some());
122    }
123
124    #[test]
125    fn wma_reset_clears_state() {
126        let mut wma = WMA::new(3);
127        let candle = Candle {
128            timestamp: 0,
129            open: 1.0,
130            high: 1.0,
131            low: 1.0,
132            close: 1.0,
133            volume: 0.0,
134        };
135
136        wma.next(&candle);
137        wma.next(&candle);
138        wma.next(&candle);
139        assert!(wma.next(&candle).is_some());
140
141        wma.reset();
142        assert_eq!(wma.next(&candle), None);
143    }
144
145    #[test]
146    fn wma_with_constant_values() {
147        let mut wma = WMA::new(2);
148        let candles = [5.0, 5.0, 5.0]
149            .iter()
150            .map(|c| Candle {
151                timestamp: 0,
152                open: *c,
153                high: *c,
154                low: *c,
155                close: *c,
156                volume: 0.0,
157            })
158            .collect::<Vec<_>>();
159
160        let outputs: Vec<_> = candles.iter().map(|c| wma.next(c)).collect();
161        assert_eq!(outputs[0], None);
162        assert_eq!(outputs[1], Some(5.0));
163        assert_eq!(outputs[2], Some(5.0));
164    }
165
166    #[test]
167    fn wma_warmup_period() {
168        let wma = WMA::new(5);
169        assert_eq!(wma.warmup_period(), 5);
170    }
171}
mantis_ta/indicators/trend/wma.rs

mantis_ta/indicators/trend/
wma.rs
