quantwave_core/indicators/incremental/
ta_atr.rs

1//! Native O(1) ATR — TA-Lib and classic Wilder smoothing on true range.
2
3use crate::traits::Next;
4use serde::{Deserialize, Serialize};
5
6/// Seeding and warm-up convention for Average True Range.
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
8pub enum AtrSeed {
9    /// TA-Lib convention (default): True Range is undefined on bar 0.
10    /// Emits `NaN` for bars 0..period-1 (the first `period` bars).
11    /// Seeds at bar `period` with the simple average of `TR[1..period]`.
12    #[default]
13    TaLib,
14    /// Classic J. Welles Wilder Jr. (1978) specification:
15    /// Bar 0 True Range is `High[0] - Low[0]`.
16    /// Emits `NaN` for bars 0..period-2 (the first `period - 1` bars).
17    /// Seeds at bar `period - 1` with the simple average of `TR[0..period-1]`.
18    Wilder1978,
19}
20
21impl std::str::FromStr for AtrSeed {
22    type Err = String;
23    fn from_str(s: &str) -> Result<Self, Self::Err> {
24        match s.to_ascii_lowercase().as_str() {
25            "talib" | "standard" | "ta_lib" | "ta-lib" => Ok(AtrSeed::TaLib),
26            "wilder1978" | "wilder_1978" | "wilder-1978" | "classic" | "wilder" | "book" => {
27                Ok(AtrSeed::Wilder1978)
28            }
29            other => Err(format!(
30                "unknown ATR seed policy: '{other}' (expected 'talib' or 'wilder1978'/'classic')"
31            )),
32        }
33    }
34}
35
36/// Average True Range — matches `talib_rs::volatility::atr` by default, or Wilder 1978 specification.
37#[derive(Debug, Clone, Serialize, Deserialize)]
38#[allow(non_camel_case_types)]
39pub struct TaATR {
40    pub timeperiod: usize,
41    period_f: f64,
42    prev_close: Option<f64>,
43    bars_seen: usize,
44    warmup_tr_count: usize,
45    warmup_sum: f64,
46    atr: f64,
47    seed: AtrSeed,
48}
49
50impl TaATR {
51    pub fn new(timeperiod: usize) -> Self {
52        Self::with_seed(timeperiod, AtrSeed::TaLib)
53    }
54
55    pub fn with_seed(timeperiod: usize, seed: AtrSeed) -> Self {
56        Self {
57            timeperiod,
58            period_f: timeperiod as f64,
59            prev_close: None,
60            bars_seen: 0,
61            warmup_tr_count: 0,
62            warmup_sum: 0.0,
63            atr: 0.0,
64            seed,
65        }
66    }
67
68    pub fn classic(timeperiod: usize) -> Self {
69        Self::with_seed(timeperiod, AtrSeed::Wilder1978)
70    }
71
72    #[inline]
73    pub fn seed(&self) -> AtrSeed {
74        self.seed
75    }
76
77    #[inline]
78    fn true_range(&self, high: f64, low: f64, prev_close: f64) -> f64 {
79        let hl = high - low;
80        let hc = (high - prev_close).abs();
81        let lc = (low - prev_close).abs();
82        hl.max(hc).max(lc)
83    }
84}
85
86impl Next<(f64, f64, f64)> for TaATR {
87    type Output = f64;
88
89    fn next(&mut self, (high, low, close): (f64, f64, f64)) -> Self::Output {
90        let period = self.timeperiod;
91        if period < 1 {
92            return f64::NAN;
93        }
94        if high.is_nan() || low.is_nan() || close.is_nan() {
95            return f64::NAN;
96        }
97
98        match self.seed {
99            AtrSeed::TaLib => {
100                if self.bars_seen == 0 {
101                    self.prev_close = Some(close);
102                    self.bars_seen = 1;
103                    return f64::NAN;
104                }
105
106                let pc = self.prev_close.unwrap_or(close);
107                let tr = self.true_range(high, low, pc);
108                self.prev_close = Some(close);
109                self.bars_seen += 1;
110
111                if self.warmup_tr_count < period {
112                    self.warmup_tr_count += 1;
113                    self.warmup_sum += tr;
114                    if self.warmup_tr_count < period {
115                        return f64::NAN;
116                    }
117                    self.atr = self.warmup_sum / self.period_f;
118                    return self.atr;
119                }
120
121                self.atr = (self.atr * (self.period_f - 1.0) + tr) / self.period_f;
122                self.atr
123            }
124            AtrSeed::Wilder1978 => {
125                let tr = match self.prev_close {
126                    Some(pc) => self.true_range(high, low, pc),
127                    None => high - low,
128                };
129                self.prev_close = Some(close);
130                self.bars_seen += 1;
131
132                if self.warmup_tr_count < period {
133                    self.warmup_tr_count += 1;
134                    self.warmup_sum += tr;
135                    if self.warmup_tr_count < period {
136                        return f64::NAN;
137                    }
138                    self.atr = self.warmup_sum / self.period_f;
139                    return self.atr;
140                }
141
142                self.atr = (self.atr * (self.period_f - 1.0) + tr) / self.period_f;
143                self.atr
144            }
145        }
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152    use proptest::prelude::*;
153
154    proptest! {
155        #[test]
156        fn test_ta_atr_parity(
157            h in prop::collection::vec(1.0..100.0, 1..100),
158            l in prop::collection::vec(1.0..100.0, 1..100),
159            c in prop::collection::vec(1.0..100.0, 1..100)
160        ) {
161            let len = h.len().min(l.len()).min(c.len());
162            if len == 0 { return Ok(()); }
163            let mut high = Vec::with_capacity(len);
164            let mut low = Vec::with_capacity(len);
165            let mut close = Vec::with_capacity(len);
166            for i in 0..len {
167                let v_h: f64 = h[i];
168                let v_l: f64 = l[i];
169                let v_c: f64 = c[i];
170                high.push(v_h.max(v_l).max(v_c));
171                low.push(v_h.min(v_l).min(v_c));
172                close.push(v_c);
173            }
174
175            let period = 14;
176            let mut ta_atr = TaATR::new(period);
177            let streaming_results: Vec<f64> =
178                (0..len).map(|i| ta_atr.next((high[i], low[i], close[i]))).collect();
179            let batch_results = talib_rs::volatility::atr(&high, &low, &close, period)
180                .unwrap_or_else(|_| vec![f64::NAN; len]);
181
182            for (s, b) in streaming_results.iter().zip(batch_results.iter()) {
183                if s.is_nan() {
184                    assert!(b.is_nan());
185                } else {
186                    approx::assert_relative_eq!(s, b, epsilon = 1e-6);
187                }
188            }
189        }
190    }
191
192    #[test]
193    fn test_wilder1978_classic_formula() {
194        // Synthetic 20-bar dataset
195        let high: Vec<f64> = vec![
196            10.0, 11.0, 12.0, 11.5, 12.5, 13.0, 12.8, 13.5, 14.0, 13.8, 14.2, 14.5, 14.1, 14.8,
197            15.0, 15.2, 15.1, 15.5, 15.8, 16.0,
198        ];
199        let low: Vec<f64> = vec![
200            8.0, 9.0, 10.0, 9.5, 10.5, 11.0, 10.8, 11.5, 12.0, 11.8, 12.2, 12.5, 12.1, 12.8, 13.0,
201            13.2, 13.1, 13.5, 13.8, 14.0,
202        ];
203        let close: Vec<f64> = vec![
204            9.0, 10.5, 11.0, 10.0, 12.0, 11.5, 12.5, 13.0, 12.5, 13.5, 13.0, 14.0, 13.5, 14.5,
205            14.0, 14.8, 14.2, 15.0, 14.5, 15.5,
206        ];
207
208        // Wilder 1978 hand calculation for period = 14:
209        // TR[0] = H[0] - L[0] = 10.0 - 8.0 = 2.0
210        let mut tr = [0.0; 20];
211        tr[0] = high[0] - low[0];
212        for i in 1..20 {
213            let hl = high[i] - low[i];
214            let hc = (high[i] - close[i - 1]).abs();
215            let lc = (low[i] - close[i - 1]).abs();
216            tr[i] = hl.max(hc).max(lc);
217        }
218
219        let mut expected_atr = [f64::NAN; 20];
220        // Bars 0..12 are NaN.
221        // Bar 13 (14th bar): mean of TR[0..14]
222        let sum_14: f64 = tr[0..14].iter().sum();
223        expected_atr[13] = sum_14 / 14.0;
224        for i in 14..20 {
225            expected_atr[i] = (expected_atr[i - 1] * 13.0 + tr[i]) / 14.0;
226        }
227
228        let mut atr = TaATR::classic(14);
229        let mut actual = Vec::with_capacity(20);
230        for i in 0..20 {
231            actual.push(atr.next((high[i], low[i], close[i])));
232        }
233
234        for i in 0..13 {
235            assert!(actual[i].is_nan(), "Bar {i} must be NaN under Wilder 1978");
236        }
237        for i in 13..20 {
238            approx::assert_relative_eq!(actual[i], expected_atr[i], epsilon = 1e-12);
239        }
240    }
241
242    #[test]
243    fn test_wilder1978_talib_convergence() {
244        // 500-bar series where differences decay below 1e-9 by bar 400
245        let n = 500;
246        let mut high = Vec::with_capacity(n);
247        let mut low = Vec::with_capacity(n);
248        let mut close = Vec::with_capacity(n);
249        let mut p = 100.0;
250        for i in 0..n {
251            let step = (i as f64 * 0.1).sin() * 2.0;
252            p += step;
253            high.push(p + 1.5);
254            low.push(p - 1.5);
255            close.push(p);
256        }
257
258        let mut atr_talib = TaATR::new(14);
259        let mut atr_classic = TaATR::classic(14);
260
261        let mut diff_400 = 0.0;
262        for i in 0..n {
263            let v_t = atr_talib.next((high[i], low[i], close[i]));
264            let v_c = atr_classic.next((high[i], low[i], close[i]));
265            if i == 400 {
266                diff_400 = (v_t - v_c).abs() / v_t;
267            }
268        }
269
270        assert!(
271            diff_400 < 1e-9,
272            "Relative difference at bar 400 must be < 1e-9, got {diff_400}"
273        );
274    }
275}
quantwave_core/indicators/incremental/ta_atr.rs

quantwave_core/indicators/incremental/
ta_atr.rs
