Skip to main content

finance_solution/stocks/ta/
atr.rs

1//! # Average True Range (ATR)
2//!
3//! Wilder ATR of high / low / close:
4//!
5//! ```text
6//! TR[i]  = max( high−low, |high−close_prev|, |low−close_prev| )
7//! ATR[i] = Wilder smooth of TR over `period`
8//!          (seed = SMA of first `period` true ranges)
9//! ```
10//!
11//! Default pack: **period 14** ([`AtrParams::period_14`]).  
12//! Keltner channels reuse this definition internally for their ATR leg.
13//!
14//! ---
15//!
16//! ## Trading perspective
17//!
18//! | Use | Habit |
19//! |-----|-------|
20//! | Volatility level | Wide ATR → large stops / size down |
21//! | Breakout filters | Move in ATR units |
22//! | Keltner width | `mult * ATR` around EMA |
23//!
24//! ---
25//!
26//! ## Engineering perspective
27//!
28//! [`AtrParams`] → [`ValidatedAtr`] / [`atr`] → [`AtrState`] → [`atr_solution`].  
29//! First ATR at index `period − 1` when bar 0 has TR = high−low only (no prior close).
30//!
31//! ## Word problem
32//!
33//! > Constant 2-point range bars, no gaps. What is ATR(14) after warm-up?
34//!
35//! ≈ **2.0**.
36//!
37//! ```
38//! use finance_solution::stocks::ta::{atr, AtrParams};
39//! let n = 30usize;
40//! let high: Vec<_> = (0..n).map(|_| 102.0).collect();
41//! let low: Vec<_> = (0..n).map(|_| 100.0).collect();
42//! let close: Vec<_> = (0..n).map(|_| 101.0).collect();
43//! let s = atr(&high, &low, &close, AtrParams::period_14()).unwrap();
44//! assert!((s.last().unwrap() - 2.0).abs() < 1e-9);
45//! ```
46
47use crate::stocks::ta::common::{opt_cell, require_hlc, true_range};
48use crate::util::error::FinanceResult;
49use crate::util::primitives::PeriodLength;
50use crate::{columns_with_strings, print_table_locale_opt};
51
52/// ATR lookback pack (Wilder).
53#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
54pub struct AtrParams {
55    pub period: usize,
56}
57
58impl AtrParams {
59    pub const fn new(period: usize) -> Self {
60        Self { period }
61    }
62
63    pub const fn period_14() -> Self {
64        Self { period: 14 }
65    }
66
67    pub const fn period_10() -> Self {
68        Self { period: 10 }
69    }
70}
71
72/// Validated ATR config.
73#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
74pub struct ValidatedAtr {
75    params: AtrParams,
76}
77
78impl ValidatedAtr {
79    pub fn new(params: AtrParams) -> FinanceResult<Self> {
80        PeriodLength::new(params.period)?;
81        Ok(Self { params })
82    }
83
84    pub fn params(self) -> AtrParams {
85        self.params
86    }
87
88    pub fn compute(self, high: &[f64], low: &[f64], close: &[f64]) -> FinanceResult<AtrSeries> {
89        atr_validated(high, low, close, self)
90    }
91}
92
93#[derive(Clone, Debug, PartialEq)]
94pub struct AtrSeries {
95    pub atr: Vec<Option<f64>>,
96    pub params: AtrParams,
97}
98
99impl AtrSeries {
100    pub fn last(&self) -> Option<f64> {
101        self.atr.iter().rev().find_map(|x| *x)
102    }
103}
104
105#[derive(Clone, Debug)]
106pub struct AtrSolution {
107    series: AtrSeries,
108    close: Vec<f64>,
109    formula: String,
110    symbolic_formula: String,
111}
112
113impl AtrSolution {
114    pub fn series(&self) -> &AtrSeries {
115        &self.series
116    }
117    pub fn formula(&self) -> &str {
118        &self.formula
119    }
120    pub fn symbolic_formula(&self) -> &str {
121        &self.symbolic_formula
122    }
123
124    pub fn print_table(&self) {
125        self.print_table_locale_opt(None, None);
126    }
127
128    pub fn print_table_locale(&self, locale: &num_format::Locale, precision: usize) {
129        self.print_table_locale_opt(Some(locale), Some(precision));
130    }
131
132    fn print_table_locale_opt(
133        &self,
134        locale: Option<&num_format::Locale>,
135        precision: Option<usize>,
136    ) {
137        let columns = columns_with_strings(&[
138            ("period", "i", true),
139            ("close", "f", true),
140            ("atr", "f", true),
141        ]);
142        let data = self
143            .close
144            .iter()
145            .enumerate()
146            .map(|(i, c)| vec![i.to_string(), c.to_string(), opt_cell(self.series.atr[i])])
147            .collect();
148        print_table_locale_opt(&columns, data, locale, precision);
149    }
150}
151
152/// Incremental Wilder ATR.
153#[derive(Clone, Debug, PartialEq)]
154pub struct AtrState {
155    params: AtrParams,
156    prev_close: Option<f64>,
157    atr: Option<f64>,
158    seed_tr: Vec<f64>,
159    last: Option<f64>,
160}
161
162impl AtrState {
163    pub fn new(params: AtrParams) -> FinanceResult<Self> {
164        PeriodLength::new(params.period)?;
165        Ok(Self {
166            params,
167            prev_close: None,
168            atr: None,
169            seed_tr: Vec::with_capacity(params.period),
170            last: None,
171        })
172    }
173
174    pub fn from_history(
175        params: AtrParams,
176        high: &[f64],
177        low: &[f64],
178        close: &[f64],
179    ) -> FinanceResult<Self> {
180        let mut s = Self::new(params)?;
181        let _ = s.push_bars(high, low, close)?;
182        Ok(s)
183    }
184
185    pub fn push(&mut self, high: f64, low: f64, close: f64) -> FinanceResult<Option<f64>> {
186        crate::util::error::require_finite("high", high)?;
187        crate::util::error::require_finite("low", low)?;
188        crate::util::error::require_finite("close", close)?;
189        if high < low {
190            return Err(crate::util::error::FinanceError::InvalidCashflow {
191                message: "high must be >= low",
192            });
193        }
194        let period = self.params.period;
195        let tr = true_range(high, low, self.prev_close);
196        let out = if self.atr.is_none() {
197            self.seed_tr.push(tr);
198            if self.seed_tr.len() == period {
199                let a = self.seed_tr.iter().sum::<f64>() / period as f64;
200                self.atr = Some(a);
201                self.last = Some(a);
202                self.last
203            } else {
204                None
205            }
206        } else {
207            let a = self.atr.unwrap();
208            let a = (a * (period as f64 - 1.0) + tr) / period as f64;
209            self.atr = Some(a);
210            self.last = Some(a);
211            self.last
212        };
213        self.prev_close = Some(close);
214        Ok(out)
215    }
216
217    pub fn push_bars(
218        &mut self,
219        high: &[f64],
220        low: &[f64],
221        close: &[f64],
222    ) -> FinanceResult<Vec<Option<f64>>> {
223        require_hlc(high, low, close)?;
224        let mut out = Vec::with_capacity(close.len());
225        for i in 0..close.len() {
226            out.push(self.push(high[i], low[i], close[i])?);
227        }
228        Ok(out)
229    }
230
231    pub fn last(&self) -> Option<f64> {
232        self.last
233    }
234
235    pub fn reset(&mut self) {
236        self.prev_close = None;
237        self.atr = None;
238        self.seed_tr.clear();
239        self.last = None;
240    }
241}
242
243pub fn atr(
244    high: &[f64],
245    low: &[f64],
246    close: &[f64],
247    params: AtrParams,
248) -> FinanceResult<AtrSeries> {
249    ValidatedAtr::new(params)?.compute(high, low, close)
250}
251
252fn atr_validated(
253    high: &[f64],
254    low: &[f64],
255    close: &[f64],
256    eng: ValidatedAtr,
257) -> FinanceResult<AtrSeries> {
258    require_hlc(high, low, close)?;
259    let mut state = AtrState::new(eng.params)?;
260    let atr = state.push_bars(high, low, close)?;
261    Ok(AtrSeries {
262        atr,
263        params: eng.params,
264    })
265}
266
267/// # Examples
268/// ```
269/// use finance_solution::stocks::ta::{atr_solution, AtrParams};
270/// let high = vec![10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0];
271/// let low  = vec![ 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0];
272/// let close= vec![ 9.5, 10.5, 11.5, 12.5, 13.5, 14.5, 15.5, 16.5, 17.5, 18.5, 19.5, 20.5, 21.5, 22.5, 23.5];
273/// let sol = atr_solution(&high, &low, &close, AtrParams::period_14()).unwrap();
274/// assert!(sol.series().last().is_some());
275/// ```
276pub fn atr_solution(
277    high: &[f64],
278    low: &[f64],
279    close: &[f64],
280    params: AtrParams,
281) -> FinanceResult<AtrSolution> {
282    let series = atr(high, low, close, params)?;
283    Ok(AtrSolution {
284        series,
285        close: close.to_vec(),
286        formula: format!("ATR({}) = Wilder smooth of true range", params.period),
287        symbolic_formula: "ATR = Wilder(TR); TR = max(H-L, |H-Cprev|, |L-Cprev|)".to_string(),
288    })
289}
290
291#[cfg(test)]
292mod tests {
293    use super::*;
294
295    #[test]
296    fn constant_range() {
297        let n = 30usize;
298        let high: Vec<_> = (0..n).map(|_| 102.0).collect();
299        let low: Vec<_> = (0..n).map(|_| 100.0).collect();
300        let close: Vec<_> = (0..n).map(|_| 101.0).collect();
301        let s = atr(&high, &low, &close, AtrParams::period_14()).unwrap();
302        assert!((s.last().unwrap() - 2.0).abs() < 1e-9);
303    }
304
305    #[test]
306    fn state_parity() {
307        let n = 40usize;
308        let high: Vec<_> = (0..n).map(|i| 101.0 + i as f64 * 0.1).collect();
309        let low: Vec<_> = (0..n).map(|i| 99.0 + i as f64 * 0.1).collect();
310        let close: Vec<_> = (0..n).map(|i| 100.0 + i as f64 * 0.1).collect();
311        let batch = atr(&high, &low, &close, AtrParams::period_10()).unwrap();
312        let st = AtrState::from_history(AtrParams::period_10(), &high, &low, &close).unwrap();
313        assert!((batch.last().unwrap() - st.last().unwrap()).abs() < 1e-9);
314    }
315
316    #[test]
317    fn high_lt_low_err() {
318        let high = vec![10.0, 9.0];
319        let low = vec![9.0, 10.0]; // bar 1 inverted
320        let close = vec![9.5, 9.5];
321        assert!(atr(&high, &low, &close, AtrParams::period_14()).is_err());
322    }
323
324    #[test]
325    fn first_atr_at_period_minus_one() {
326        let n = 20usize;
327        let high: Vec<_> = (0..n).map(|_| 102.0).collect();
328        let low: Vec<_> = (0..n).map(|_| 100.0).collect();
329        let close: Vec<_> = (0..n).map(|_| 101.0).collect();
330        let s = atr(&high, &low, &close, AtrParams::period_14()).unwrap();
331        assert!(s.atr[12].is_none());
332        assert!(s.atr[13].is_some()); // period=14 → index 13
333    }
334
335    #[test]
336    fn gap_increases_atr_vs_no_gap() {
337        // Same ranges but with a large gap mid-path
338        let n = 30usize;
339        let mut high: Vec<f64> = (0..n).map(|_| 102.0).collect();
340        let mut low: Vec<f64> = (0..n).map(|_| 100.0).collect();
341        let mut close: Vec<f64> = (0..n).map(|_| 101.0).collect();
342        let base = atr(&high, &low, &close, AtrParams::period_14())
343            .unwrap()
344            .last()
345            .unwrap();
346        // Introduce gap open after bar 15
347        high[16] = 110.0;
348        low[16] = 108.0;
349        close[16] = 109.0;
350        let gapped = atr(&high, &low, &close, AtrParams::period_14())
351            .unwrap()
352            .last()
353            .unwrap();
354        assert!(gapped > base);
355    }
356}