Skip to main content

quantwave_core/indicators/incremental/
talib_wma.rs

1//! TA-Lib compatible streaming WMA — O(1) per bar, bounded memory.
2//!
3//! Mirrors the `talib_rs::overlap::wma` recurrence exactly:
4//!   `ws[i] = ws[i-1] + period * x[i] - ps[i-1]`
5//!   `ps[i] = ps[i-1] + x[i] - x[i-period]`
6//! so the streaming and batch paths agree bit-for-bit on the same inputs.
7
8use crate::traits::Next;
9use crate::utils::RingBuffer;
10
11/// Weighted moving average matching `talib_rs::overlap::wma` (lookback `period - 1`).
12#[derive(Debug, Clone)]
13pub struct TalibWma {
14    period: usize,
15    divider: f64,
16    /// Weighted sum over the current window.
17    ws: f64,
18    /// Plain sum over the current window.
19    ps: f64,
20    /// Last `period` inputs — needed only for the value leaving the window.
21    window: RingBuffer<f64>,
22    bars_seen: usize,
23}
24
25impl TalibWma {
26    pub fn new(period: usize) -> Self {
27        let p = period as f64;
28        Self {
29            period,
30            divider: p * (p + 1.0) / 2.0,
31            ws: 0.0,
32            ps: 0.0,
33            window: RingBuffer::with_capacity(period.max(1)),
34            bars_seen: 0,
35        }
36    }
37}
38
39impl Next<f64> for TalibWma {
40    type Output = f64;
41
42    fn next(&mut self, input: f64) -> Self::Output {
43        let p = self.period;
44        if p == 0 {
45            return f64::NAN;
46        }
47        let i = self.bars_seen;
48        self.bars_seen += 1;
49
50        if i < p {
51            // Seeding window: replicate the batch initialisation loop.
52            self.window.push_back(input);
53            self.ws += input * (i + 1) as f64;
54            self.ps += input;
55            if i + 1 < p {
56                return f64::NAN;
57            }
58            return self.ws / self.divider;
59        }
60
61        let old = self.window.pop_front().unwrap_or(0.0);
62        self.window.push_back(input);
63        self.ws = self.ws + p as f64 * input - self.ps;
64        self.ps = self.ps + input - old;
65        self.ws / self.divider
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use super::*;
72    use proptest::prelude::*;
73
74    proptest! {
75        #[test]
76        fn test_talib_wma_parity(input in prop::collection::vec(0.1..100.0, 1..100)) {
77            let period = 10;
78            let mut wma = TalibWma::new(period);
79            let streaming: Vec<f64> = input.iter().map(|&x| wma.next(x)).collect();
80            let batch = talib_rs::overlap::wma(&input, period)
81                .unwrap_or_else(|_| vec![f64::NAN; input.len()]);
82            for (s, b) in streaming.iter().zip(batch.iter()) {
83                if s.is_nan() { assert!(b.is_nan()); }
84                else { approx::assert_relative_eq!(s, b, epsilon = 1e-9); }
85            }
86        }
87    }
88}