indicators/trend/
atr.rs

1//! Average True Range (ATR).
2//!
3//! Python source: `indicators/trend/volatility/atr.py :: class ATR`
4//!
5//! # Python algorithm (to port)
6//! ```python
7//! high_low        = data["high"] - data["low"]
8//! high_close_prev = abs(data["high"] - data["close"].shift(1))
9//! low_close_prev  = abs(data["low"]  - data["close"].shift(1))
10//! tr  = pd.concat([high_low, high_close_prev, low_close_prev], axis=1).max(axis=1)
11//! atr = tr.rolling(period).mean()           # method=="sma"
12//! # or:
13//! atr = tr.ewm(span=period, adjust=False).mean()  # method=="ema"
14//!
15//! normalized_atr = atr / data["close"] * 100   # percentage
16//! ```
17//!
18//! Output columns: `"ATR_{period}"`, `"ATR_{period}_normalized"`.
19//!
20//! See also: `crate::functions::atr()` and `crate::functions::true_range()`.
21
22use std::collections::HashMap;
23
24use crate::error::IndicatorError;
25use crate::functions::{self};
26use crate::indicator::{Indicator, IndicatorOutput};
27use crate::registry::{param_str, param_usize};
28use crate::types::Candle;
29
30// ── Params ────────────────────────────────────────────────────────────────────
31
32#[derive(Debug, Clone, PartialEq, Eq)]
33pub enum AtrMethod {
34    Sma,
35    Ema,
36}
37
38#[derive(Debug, Clone)]
39pub struct AtrParams {
40    /// Period.  Python default: 14.
41    pub period: usize,
42    /// Smoothing method.  Python default: `"sma"`.
43    pub method: AtrMethod,
44}
45
46impl Default for AtrParams {
47    fn default() -> Self {
48        Self {
49            period: 14,
50            method: AtrMethod::Sma,
51        }
52    }
53}
54
55// ── Indicator struct ──────────────────────────────────────────────────────────
56
57#[derive(Debug, Clone)]
58pub struct Atr {
59    pub params: AtrParams,
60}
61
62impl Atr {
63    pub fn new(params: AtrParams) -> Self {
64        Self { params }
65    }
66    pub fn with_period(period: usize) -> Self {
67        Self::new(AtrParams {
68            period,
69            ..Default::default()
70        })
71    }
72
73    fn output_key(&self) -> String {
74        format!("ATR_{}", self.params.period)
75    }
76    fn norm_key(&self) -> String {
77        format!("ATR_{}_normalized", self.params.period)
78    }
79}
80
81impl Indicator for Atr {
82    fn name(&self) -> &'static str {
83        "ATR"
84    }
85    fn required_len(&self) -> usize {
86        self.params.period + 1
87    } // need prev close
88    fn required_columns(&self) -> &[&'static str] {
89        &["high", "low", "close"]
90    }
91
92    /// Ports the Python ATR calculation.
93    ///
94    /// True range = `max(H−L, |H−prev_C|, |L−prev_C|)`.  For the first bar
95    /// there is no previous close, so `functions::true_range` is expected to
96    /// use `H−L` alone (matching pandas `skipna=True` max behaviour).
97    ///
98    /// SMA path: `tr.rolling(period).mean()` — `NaN` for first `period` bars.
99    /// EMA path: `tr.ewm(span=period, adjust=False).mean()` — value from bar 0.
100    ///
101    /// Normalised ATR = `atr / close * 100` (percentage of price).
102    fn calculate(&self, candles: &[Candle]) -> Result<IndicatorOutput, IndicatorError> {
103        self.check_len(candles)?;
104
105        let high: Vec<f64> = candles.iter().map(|c| c.high).collect();
106        let low: Vec<f64> = candles.iter().map(|c| c.low).collect();
107        let close: Vec<f64> = candles.iter().map(|c| c.close).collect();
108
109        let tr = functions::true_range(&high, &low, &close)?;
110
111        let atr_vals = match self.params.method {
112            // ema_nan_aware seeds from the first TR value (adjust=False),
113            // matching Python's tr.ewm(span=period, adjust=False).mean().
114            AtrMethod::Ema => functions::ema_nan_aware(&tr, self.params.period)?,
115            AtrMethod::Sma => functions::sma(&tr, self.params.period)?,
116        };
117
118        let norm: Vec<f64> = atr_vals
119            .iter()
120            .zip(&close)
121            .map(|(&a, &c)| if c == 0.0 { f64::NAN } else { a / c * 100.0 })
122            .collect();
123
124        Ok(IndicatorOutput::from_pairs([
125            (self.output_key(), atr_vals),
126            (self.norm_key(), norm),
127        ]))
128    }
129}
130
131// ── Registry factory ──────────────────────────────────────────────────────────
132
133pub fn factory<S: ::std::hash::BuildHasher>(
134    params: &HashMap<String, String, S>,
135) -> Result<Box<dyn Indicator>, IndicatorError> {
136    let period = param_usize(params, "period", 14)?;
137    let method = match param_str(params, "method", "sma") {
138        "ema" => AtrMethod::Ema,
139        _ => AtrMethod::Sma,
140    };
141    Ok(Box::new(Atr::new(AtrParams { period, method })))
142}
143
144// ── Tests ─────────────────────────────────────────────────────────────────────
145
146#[cfg(test)]
147mod tests {
148    use super::*;
149
150    fn candles(data: &[(f64, f64, f64)]) -> Vec<Candle> {
151        data.iter()
152            .enumerate()
153            .map(|(i, &(h, l, c))| Candle {
154                time: i64::try_from(i).expect("time index fits i64"),
155                open: c,
156                high: h,
157                low: l,
158                close: c,
159                volume: 1.0,
160            })
161            .collect()
162    }
163
164    #[test]
165    fn atr_output_has_both_columns() {
166        let bars: Vec<(f64, f64, f64)> = (1..=20)
167            .map(|i| (i as f64 + 1.0, i as f64 - 1.0, i as f64))
168            .collect();
169        let atr = Atr::with_period(5);
170        let out = atr.calculate(&candles(&bars)).unwrap();
171        assert!(out.get("ATR_5").is_some());
172        assert!(out.get("ATR_5_normalized").is_some());
173    }
174
175    #[test]
176    fn atr_insufficient_data() {
177        assert!(
178            Atr::with_period(14)
179                .calculate(&candles(&[(10.0, 8.0, 9.0)]))
180                .is_err()
181        );
182    }
183
184    #[test]
185    fn atr_normalized_is_percentage() {
186        let bars: Vec<(f64, f64, f64)> = (1..=20)
187            .map(|i| (i as f64 + 1.0, i as f64 - 1.0, i as f64))
188            .collect();
189        let atr = Atr::with_period(5);
190        let out = atr.calculate(&candles(&bars)).unwrap();
191        let atr_vals = out.get("ATR_5").unwrap();
192        let norm_vals = out.get("ATR_5_normalized").unwrap();
193        let close: Vec<f64> = bars.iter().map(|&(_, _, c)| c).collect();
194        for i in 0..bars.len() {
195            if !atr_vals[i].is_nan() {
196                let expected = atr_vals[i] / close[i] * 100.0;
197                assert!((norm_vals[i] - expected).abs() < 1e-9);
198            }
199        }
200    }
201
202    #[test]
203    fn factory_creates_atr() {
204        let params = [
205            ("period".into(), "14".into()),
206            ("method".into(), "ema".into()),
207        ]
208        .into();
209        let ind = factory(&params).unwrap();
210        assert_eq!(ind.name(), "ATR");
211    }
212}
indicators/trend/atr.rs

indicators/trend/
atr.rs
