indexes_rs/v2/atr/
main.rs

1use super::types::{ATRError, ATRResult, OHLCData};
2
3/// Average True Range (ATR) Calculator
4///
5/// ATR measures market volatility by calculating the average of true ranges over a period.
6/// True Range is the greatest of:
7/// - Current High - Current Low
8/// - |Current High - Previous Close|
9/// - |Current Low - Previous Close|
10pub 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    /// Creates a new ATR calculator with the specified period
20    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    /// Updates ATR with OHLC data
35    pub fn update(&mut self, data: OHLCData) -> Option<ATRResult> {
36        // Validate input
37        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; // Invalid OHLC data
45        }
46
47        // Calculate True Range
48        let true_range = self.calculate_true_range(data.high, data.low);
49
50        // Store current close as previous close for next calculation
51        self.prev_close = Some(data.close);
52
53        // Calculate ATR
54        self.calculate_atr(true_range)
55    }
56
57    /// Updates ATR with separate high, low, close values
58    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    /// Calculate with just closing price (simplified - less accurate)
63    pub fn update_close_only(&mut self, close: f64) -> Option<ATRResult> {
64        // Estimate high/low using close with a volatility factor
65        let estimated_range = close * 0.01; // 1% estimated daily range
66        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    /// Returns the current ATR value
73    pub fn value(&self) -> Option<f64> {
74        self.atr_value
75    }
76
77    /// Returns the current ATR result with true range
78    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    /// Resets the calculator
86    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    /// Calculate True Range
94    fn calculate_true_range(&self, high: f64, low: f64) -> f64 {
95        if let Some(prev_close) = self.prev_close {
96            // True Range = max of:
97            // 1. Current High - Current Low
98            // 2. |Current High - Previous Close|
99            // 3. |Current Low - Previous Close|
100            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            // First candle: use High - Low
107            high - low
108        }
109    }
110
111    /// Calculate ATR using Wilder's smoothing method
112    fn calculate_atr(&mut self, true_range: f64) -> Option<ATRResult> {
113        if !self.initialized {
114            // Initial ATR calculation: Simple average of first N true ranges
115            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            // Wilder's smoothing: ATR = ((ATR_prev * (n-1)) + TR) / n
128            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                // Keep only the last true range for reference
133                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    /// Get the period
144    pub fn period(&self) -> usize {
145        self.period
146    }
147
148    /// Check if ATR is initialized
149    pub fn is_ready(&self) -> bool {
150        self.initialized
151    }
152
153    /// Batch calculation for historical data
154    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}