indexes_rs/v2/atr/
main.rs1use super::types::{ATRError, ATRResult, OHLCData};
2
3pub struct ATR {
11 period: usize,
12 atr_value: Option<f64>,
13 true_ranges: Vec<f64>,
14 prev_close: Option<f64>,
15 initialized: bool,
16}
17
18impl ATR {
19 pub fn new(period: usize) -> Result<Self, ATRError> {
21 if period == 0 {
22 return Err(ATRError::InvalidPeriod);
23 }
24
25 Ok(Self {
26 period,
27 atr_value: None,
28 true_ranges: Vec::with_capacity(period),
29 prev_close: None,
30 initialized: false,
31 })
32 }
33
34 pub fn update(&mut self, data: OHLCData) -> Option<ATRResult> {
36 if data.high.is_nan() || data.low.is_nan() || data.close.is_nan() {
38 return None;
39 }
40 if data.high.is_infinite() || data.low.is_infinite() || data.close.is_infinite() {
41 return None;
42 }
43 if data.high < data.low {
44 return None; }
46
47 let true_range = self.calculate_true_range(data.high, data.low);
49
50 self.prev_close = Some(data.close);
52
53 self.calculate_atr(true_range)
55 }
56
57 pub fn update_hlc(&mut self, high: f64, low: f64, close: f64) -> Option<ATRResult> {
59 self.update(OHLCData::new(high, low, close))
60 }
61
62 pub fn update_close_only(&mut self, close: f64) -> Option<ATRResult> {
64 let estimated_range = close * 0.01; let high = close + estimated_range / 2.0;
67 let low = close - estimated_range / 2.0;
68
69 self.update(OHLCData::new(high, low, close))
70 }
71
72 pub fn value(&self) -> Option<f64> {
74 self.atr_value
75 }
76
77 pub fn result(&self) -> Option<ATRResult> {
79 self.atr_value.map(|atr| ATRResult {
80 atr,
81 true_range: self.true_ranges.last().copied().unwrap_or(0.0),
82 })
83 }
84
85 pub fn reset(&mut self) {
87 self.atr_value = None;
88 self.true_ranges.clear();
89 self.prev_close = None;
90 self.initialized = false;
91 }
92
93 fn calculate_true_range(&self, high: f64, low: f64) -> f64 {
95 if let Some(prev_close) = self.prev_close {
96 let hl = high - low;
101 let hc = (high - prev_close).abs();
102 let lc = (low - prev_close).abs();
103
104 hl.max(hc).max(lc)
105 } else {
106 high - low
108 }
109 }
110
111 fn calculate_atr(&mut self, true_range: f64) -> Option<ATRResult> {
113 if !self.initialized {
114 self.true_ranges.push(true_range);
116
117 if self.true_ranges.len() >= self.period {
118 let initial_atr = self.true_ranges.iter().sum::<f64>() / self.period as f64;
119 self.atr_value = Some(initial_atr);
120 self.initialized = true;
121
122 return Some(ATRResult { atr: initial_atr, true_range });
123 }
124
125 None
126 } else {
127 if let Some(prev_atr) = self.atr_value {
129 let new_atr = (prev_atr * (self.period - 1) as f64 + true_range) / self.period as f64;
130 self.atr_value = Some(new_atr);
131
132 self.true_ranges.clear();
134 self.true_ranges.push(true_range);
135
136 Some(ATRResult { atr: new_atr, true_range })
137 } else {
138 None
139 }
140 }
141 }
142
143 pub fn period(&self) -> usize {
145 self.period
146 }
147
148 pub fn is_ready(&self) -> bool {
150 self.initialized
151 }
152
153 pub fn calculate_batch(period: usize, data: &[OHLCData]) -> Result<Vec<Option<ATRResult>>, ATRError> {
155 let mut atr = Self::new(period)?;
156 let mut results = Vec::with_capacity(data.len());
157
158 for ohlc in data {
159 results.push(atr.update(*ohlc));
160 }
161
162 Ok(results)
163 }
164}