Skip to main content

egui_charts/studies/builtin/
trend_strength_index.rs

1use 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/// Trend Strength Index: the Pearson correlation of the source price with bar position over
10/// the last `period` bars, from -1 for a straight fall to 1 for a straight rise.
11///
12/// A bar has no value until `period` bars exist, when the window holds a non-finite price, or
13/// when the prices of the window are all equal, as the correlation is then undefined. A bar
14/// with a non-finite price still takes its place in the window, so the correlation is always
15/// taken against true bar positions. The latest value equals
16/// [`LinearRegressionChannel::correlation`](super::LinearRegressionChannel::correlation) of
17/// the same period.
18#[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    /// Create with a window of `period` bars, clamped to at least 2.
30    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    /// Select the input price. Defaults to close.
42    pub fn with_source(mut self, source: PriceSource) -> Self {
43        self.source = source;
44        self
45    }
46
47    /// Set the plot color.
48    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    /// Period 3 over 1, 2, 4, 3, 3, 3 against positions 0, 1, 2, from the centred sums
188    /// r = sxy / sqrt(sxx * syy), with sxx = 2: 1, 2, 4 has sxy 3 and syy 14/3; 2, 4, 3 has
189    /// sxy 1 and syy 2; 4, 3, 3 has sxy -1 and syy 2/3; 3, 3, 3 is flat. A straight fall is -1.
190    #[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    /// A missing price blanks every window that holds it and keeps later bar positions true.
211    #[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}