egui_charts/studies/builtin/
trend_strength_index.rs1use crate::model::math::Moments;
2use crate::model::{Bar, PriceSource};
3use crate::studies::palette;
4use crate::studies::stepper::Stepper;
5use crate::studies::{Indicator, IndicatorValue, IndicatorValues};
6use egui::Color32;
7use std::collections::VecDeque;
8
9#[derive(Clone)]
19pub struct TrendStrengthIndex {
20 period: usize,
21 source: PriceSource,
22 values: IndicatorValues,
23 stepper: Stepper<VecDeque<Bar>>,
24 color: Color32,
25 visible: bool,
26}
27
28impl TrendStrengthIndex {
29 pub fn new(period: usize) -> Self {
31 Self {
32 period: period.max(2),
33 source: PriceSource::Close,
34 values: IndicatorValues::new(),
35 stepper: Stepper::default(),
36 color: palette::PURPLE,
37 visible: true,
38 }
39 }
40
41 pub fn with_source(mut self, source: PriceSource) -> Self {
43 self.source = source;
44 self
45 }
46
47 pub fn with_color(mut self, color: Color32) -> Self {
49 self.color = color;
50 self
51 }
52}
53
54impl Default for TrendStrengthIndex {
55 fn default() -> Self {
56 Self::new(14)
57 }
58}
59
60impl Indicator for TrendStrengthIndex {
61 fn source(&self) -> Option<crate::studies::StudySource> {
62 Some(crate::studies::StudySource::Price(self.source))
63 }
64
65 crate::studies::inputs::study_inputs!(self;
66 period: length("Period", 2, usize::MAX),
67 source: choice("Source", PriceSource::all())
68 ; {});
69
70 fn name(&self) -> &str {
71 "Trend Strength Index"
72 }
73
74 fn desc(&self) -> &str {
75 "Trend Strength Index - Rolling correlation of price with bar position (-1 to 1)"
76 }
77
78 fn calculate(&mut self, data: &[Bar]) {
79 self.update(data, 0);
80 }
81
82 fn update(&mut self, data: &[Bar], from: usize) {
83 let (period, source) = (self.period, self.source);
84 self.stepper
85 .run(data, from, &mut self.values, |window, bar| {
86 let price = source.compute(bar.open, bar.high, bar.low, bar.close);
87 window.push_back(Bar::new(bar.time, price, price, price, price, 0.0));
88 if window.len() > period {
89 window.pop_front();
90 }
91 if window.len() < period || window.iter().any(|bar| !bar.close.is_finite()) {
92 return IndicatorValue::None;
93 }
94 Moments::of(window.make_contiguous())
95 .fit()
96 .and_then(|fit| fit.correlation)
97 .map_or(IndicatorValue::None, IndicatorValue::Single)
98 });
99 }
100
101 fn is_incremental(&self) -> bool {
102 true
103 }
104
105 fn values(&self) -> &IndicatorValues {
106 &self.values
107 }
108
109 fn colors(&self) -> Vec<Color32> {
110 vec![self.color]
111 }
112
113 fn set_colors(&mut self, colors: Vec<Color32>) {
114 if !colors.is_empty() {
115 self.color = colors[0];
116 }
117 }
118
119 fn value_format(&self) -> crate::scales::StudyFormat {
120 crate::scales::StudyFormat::Fixed(3)
121 }
122
123 fn pane(&self) -> crate::studies::StudyPane {
124 crate::studies::StudyPane {
125 range: Some((-1.0, 1.0)),
126 levels: vec![(0.0, palette::GRAY.gamma_multiply_u8(128), "Zero")],
127 ..Default::default()
128 }
129 }
130
131 fn is_overlay(&self) -> bool {
132 false
133 }
134
135 fn is_visible(&self) -> bool {
136 self.visible
137 }
138
139 fn set_visible(&mut self, visible: bool) {
140 self.visible = visible;
141 }
142
143 fn clone_box(&self) -> Box<dyn Indicator> {
144 Box::new(self.clone())
145 }
146
147 fn line_names(&self) -> Vec<String> {
148 vec![format!("Trend Strength({})", self.period)]
149 }
150}
151
152#[cfg(test)]
153mod tests {
154 use super::*;
155 use chrono::{TimeZone, Utc};
156
157 fn bars(closes: &[f64]) -> Vec<Bar> {
158 closes
159 .iter()
160 .enumerate()
161 .map(|(i, &c)| {
162 Bar::new(
163 Utc.timestamp_opt(i as i64 * 60, 0).unwrap(),
164 c,
165 c,
166 c,
167 c,
168 1.0,
169 )
170 })
171 .collect()
172 }
173
174 fn rolling(closes: &[f64], period: usize) -> Vec<Option<f64>> {
175 let mut tsi = TrendStrengthIndex::new(period);
176 tsi.calculate(&bars(closes));
177 tsi.values()
178 .to_rows()
179 .into_iter()
180 .map(|v| match v {
181 IndicatorValue::Single(v) => Some(v),
182 _ => None,
183 })
184 .collect()
185 }
186
187 #[test]
191 fn rolling_correlation_matches_hand_values() {
192 let got = rolling(&[1.0, 2.0, 4.0, 3.0, 3.0, 3.0], 3);
193 let want = [
194 None,
195 None,
196 Some(3.0 / (2.0 * 14.0 / 3.0_f64).sqrt()),
197 Some(1.0 / (2.0 * 2.0_f64).sqrt()),
198 Some(-1.0 / (2.0 * 2.0 / 3.0_f64).sqrt()),
199 None,
200 ];
201 for (got, want) in got.iter().zip(want) {
202 match (got, want) {
203 (Some(got), Some(want)) => assert!((got - want).abs() < 1e-12, "{got} {want}"),
204 _ => assert_eq!(*got, want),
205 }
206 }
207 assert_eq!(rolling(&[5.0, 4.0, 3.0, 2.0], 4)[3], Some(-1.0));
208 }
209
210 #[test]
212 fn missing_price_blanks_its_windows() {
213 let got = rolling(&[1.0, 2.0, f64::NAN, 4.0, 5.0, 6.0], 3);
214 assert_eq!(got[..5], [None; 5]);
215 assert!((got[5].unwrap() - 1.0).abs() < 1e-12);
216 }
217}