finance_solution/stocks/ta/
atr.rs1use 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#[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#[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#[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
267pub 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]; 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()); }
334
335 #[test]
336 fn gap_increases_atr_vs_no_gap() {
337 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 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}