Skip to main content

quantwave_core/indicators/incremental/
ta_atr.rs

1//! Native O(1) ATR — TA-Lib Wilder smoothing on true range.
2
3use crate::traits::Next;
4
5/// Average True Range — matches `talib_rs::volatility::atr`.
6#[derive(Debug, Clone)]
7#[allow(non_camel_case_types)]
8pub struct TaATR {
9    pub timeperiod: usize,
10    period_f: f64,
11    prev_close: Option<f64>,
12    bars_seen: usize,
13    warmup_tr_count: usize,
14    warmup_sum: f64,
15    atr: f64,
16}
17
18impl TaATR {
19    pub fn new(timeperiod: usize) -> Self {
20        Self {
21            timeperiod,
22            period_f: timeperiod as f64,
23            prev_close: None,
24            bars_seen: 0,
25            warmup_tr_count: 0,
26            warmup_sum: 0.0,
27            atr: 0.0,
28        }
29    }
30
31    #[inline]
32    fn true_range(&self, high: f64, low: f64, prev_close: f64) -> f64 {
33        let hl = high - low;
34        let hc = (high - prev_close).abs();
35        let lc = (low - prev_close).abs();
36        hl.max(hc).max(lc)
37    }
38}
39
40impl Next<(f64, f64, f64)> for TaATR {
41    type Output = f64;
42
43    fn next(&mut self, (high, low, close): (f64, f64, f64)) -> Self::Output {
44        let period = self.timeperiod;
45        if period < 1 {
46            return f64::NAN;
47        }
48        if high.is_nan() || low.is_nan() || close.is_nan() {
49            return f64::NAN;
50        }
51
52        if self.bars_seen == 0 {
53            self.prev_close = Some(close);
54            self.bars_seen = 1;
55            return f64::NAN;
56        }
57
58        let pc = self.prev_close.unwrap_or(close);
59        let tr = self.true_range(high, low, pc);
60        self.prev_close = Some(close);
61        self.bars_seen += 1;
62
63        if self.warmup_tr_count < period {
64            self.warmup_tr_count += 1;
65            self.warmup_sum += tr;
66            if self.warmup_tr_count < period {
67                return f64::NAN;
68            }
69            self.atr = self.warmup_sum / self.period_f;
70            return self.atr;
71        }
72
73        self.atr = (self.atr * (self.period_f - 1.0) + tr) / self.period_f;
74        self.atr
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81    use proptest::prelude::*;
82
83    proptest! {
84        #[test]
85        fn test_ta_atr_parity(
86            h in prop::collection::vec(1.0..100.0, 1..100),
87            l in prop::collection::vec(1.0..100.0, 1..100),
88            c in prop::collection::vec(1.0..100.0, 1..100)
89        ) {
90            let len = h.len().min(l.len()).min(c.len());
91            if len == 0 { return Ok(()); }
92            let mut high = Vec::with_capacity(len);
93            let mut low = Vec::with_capacity(len);
94            let mut close = Vec::with_capacity(len);
95            for i in 0..len {
96                let v_h: f64 = h[i];
97                let v_l: f64 = l[i];
98                let v_c: f64 = c[i];
99                high.push(v_h.max(v_l).max(v_c));
100                low.push(v_h.min(v_l).min(v_c));
101                close.push(v_c);
102            }
103
104            let period = 14;
105            let mut ta_atr = TaATR::new(period);
106            let streaming_results: Vec<f64> =
107                (0..len).map(|i| ta_atr.next((high[i], low[i], close[i]))).collect();
108            let batch_results = talib_rs::volatility::atr(&high, &low, &close, period)
109                .unwrap_or_else(|_| vec![f64::NAN; len]);
110
111            for (s, b) in streaming_results.iter().zip(batch_results.iter()) {
112                if s.is_nan() {
113                    assert!(b.is_nan());
114                } else {
115                    approx::assert_relative_eq!(s, b, epsilon = 1e-6);
116                }
117            }
118        }
119    }
120}