Skip to main content

egui_charts/studies/builtin/
linear_regression_channel.rs

1use super::clamp_param;
2use crate::model::math::{LinearFit, Moments};
3use crate::model::{Bar, PriceSource};
4use crate::studies::palette;
5use crate::studies::{Indicator, IndicatorValue, IndicatorValues};
6use egui::Color32;
7use std::borrow::Cow;
8
9/// Linear regression channel fitted to the closes of the last `period` bars.
10///
11/// Outputs upper, middle and lower lines across that window, with no value on
12/// earlier bars or until a full window is available. The upper band is
13/// `upper_multiplier` and the lower band `lower_multiplier` residual standard deviations
14/// from the line, using an n - 1 denominator, and the band above and the band below the
15/// line are filled.
16/// Every calculation or update refits the whole channel. A window containing
17/// a non-finite close has no values.
18#[derive(Clone)]
19pub struct LinearRegressionChannel {
20    period: usize,
21    upper_multiplier: f64,
22    lower_multiplier: f64,
23    values: IndicatorValues,
24    colors: Vec<Color32>,
25    visible: bool,
26    source: PriceSource,
27    extend: bool,
28    fit: Option<LinearFit>,
29}
30
31impl LinearRegressionChannel {
32    /// Create a channel with a period of at least 2. `multiplier` is clamped
33    /// into 0 to `f64::MAX`, NaN to 0.
34    pub fn new(period: usize, multiplier: f64) -> Self {
35        Self {
36            period: period.max(2),
37            upper_multiplier: clamp_param(multiplier, 0.0, f64::MAX),
38            lower_multiplier: clamp_param(multiplier, 0.0, f64::MAX),
39            values: IndicatorValues::new(),
40            colors: vec![
41                palette::ACCENT,
42                palette::RED,
43                palette::ACCENT,
44                palette::ACCENT.gamma_multiply_u8(13), // Upper band fill
45                palette::RED.gamma_multiply_u8(13),    // Lower band fill
46            ],
47            visible: true,
48            source: PriceSource::Close,
49            extend: false,
50            fit: None,
51        }
52    }
53
54    /// Place the upper band `upper` and the lower band `lower` residual standard deviations from
55    /// the line. Each is clamped into 0 to `f64::MAX`, NaN to 0.
56    pub fn with_multipliers(mut self, upper: f64, lower: f64) -> Self {
57        self.upper_multiplier = clamp_param(upper, 0.0, f64::MAX);
58        self.lower_multiplier = clamp_param(lower, 0.0, f64::MAX);
59        self
60    }
61
62    /// Select the input price for the fit. Defaults to close.
63    pub fn with_source(mut self, source: PriceSource) -> Self {
64        self.source = source;
65        self
66    }
67
68    /// Extend the three fitted lines to the right edge of the plot.
69    pub fn with_extend_right(mut self, extend: bool) -> Self {
70        self.extend = extend;
71        self
72    }
73
74    /// Pearson correlation of the source with bar position, unavailable before a fit or on flat input.
75    pub fn correlation(&self) -> Option<f64> {
76        self.fit.and_then(|fit| fit.correlation)
77    }
78
79    /// Fraction of source variance explained by the line; a flat source is fitted exactly.
80    pub fn r_squared(&self) -> Option<f64> {
81        self.fit.map(|fit| fit.r_squared)
82    }
83}
84
85impl Default for LinearRegressionChannel {
86    fn default() -> Self {
87        Self::new(100, 2.0)
88    }
89}
90
91impl Indicator for LinearRegressionChannel {
92    fn source(&self) -> Option<crate::studies::StudySource> {
93        Some(crate::studies::StudySource::Price(self.source))
94    }
95
96    crate::studies::inputs::study_inputs!(self;
97        period: length("Period", 2, usize::MAX),
98        upper_multiplier: number("Upper deviation", 0.0, f64::MAX),
99        lower_multiplier: number("Lower deviation", 0.0, f64::MAX),
100        source: choice("Source", crate::model::PriceSource::all()),
101        extend: toggle("Extend")
102    ; {});
103
104    fn fills(&self) -> Vec<crate::studies::StudyFill> {
105        use crate::studies::{FillBoundary, StudyFill};
106        vec![
107            StudyFill {
108                first: FillBoundary::Plot(0),
109                second: FillBoundary::Plot(1),
110                color: self.colors[3],
111                inverted_color: None,
112            },
113            StudyFill {
114                first: FillBoundary::Plot(1),
115                second: FillBoundary::Plot(2),
116                color: self.colors[4],
117                inverted_color: None,
118            },
119        ]
120    }
121
122    fn name(&self) -> &str {
123        "LinReg Channel"
124    }
125
126    fn desc(&self) -> &str {
127        "Linear Regression Channel - Best fit line and residual deviation bands over the last bars"
128    }
129
130    fn calculate(&mut self, data: &[Bar]) {
131        self.update(data, 0);
132    }
133
134    fn update(&mut self, data: &[Bar], _from: usize) {
135        self.fit = None;
136        let previous = self.values.len().saturating_sub(self.period);
137        self.values.truncate(previous);
138        self.values.resize(data.len(), IndicatorValue::None);
139        if data.len() < self.period {
140            return;
141        }
142        let start = data.len() - self.period;
143        let window: Cow<'_, [Bar]> = if self.source == PriceSource::Close {
144            Cow::Borrowed(&data[start..])
145        } else {
146            Cow::Owned(
147                data[start..]
148                    .iter()
149                    .map(|bar| {
150                        let mut sourced = bar.clone();
151                        sourced.close = self.source.compute(bar.open, bar.high, bar.low, bar.close);
152                        sourced
153                    })
154                    .collect(),
155            )
156        };
157        if window.iter().any(|bar| !bar.close.is_finite()) {
158            return;
159        }
160        let Some(fit) = Moments::of(&window).fit() else {
161            return;
162        };
163        self.fit = Some(fit);
164        let upper = self.upper_multiplier * fit.sigma;
165        let lower = self.lower_multiplier * fit.sigma;
166        for i in 0..self.period {
167            let middle = fit.intercept + fit.slope * i as f64;
168            self.values.set(
169                start + i,
170                IndicatorValue::Multiple(vec![middle + upper, middle, middle - lower]),
171            );
172        }
173    }
174
175    fn is_incremental(&self) -> bool {
176        true
177    }
178
179    fn extend_right(&self) -> bool {
180        self.extend
181    }
182
183    fn plot_label(&self) -> Option<(usize, String)> {
184        self.fit.map(|fit| {
185            (
186                1,
187                fit.correlation
188                    .map_or_else(|| "R: undefined".to_owned(), |r| format!("R: {r:.3}")),
189            )
190        })
191    }
192
193    fn values(&self) -> &IndicatorValues {
194        &self.values
195    }
196
197    fn colors(&self) -> Vec<Color32> {
198        self.colors.clone()
199    }
200
201    fn set_colors(&mut self, colors: Vec<Color32>) {
202        if colors.len() >= 3 {
203            let n = colors.len().min(self.colors.len());
204            self.colors[..n].copy_from_slice(&colors[..n]);
205        } else if !colors.is_empty() {
206            self.colors[..3].fill(colors[0]);
207        }
208    }
209
210    fn is_overlay(&self) -> bool {
211        true
212    }
213
214    fn line_count(&self) -> usize {
215        3
216    }
217
218    fn is_visible(&self) -> bool {
219        self.visible
220    }
221
222    fn set_visible(&mut self, visible: bool) {
223        self.visible = visible;
224    }
225
226    fn clone_box(&self) -> Box<dyn Indicator> {
227        Box::new(self.clone())
228    }
229
230    fn line_names(&self) -> Vec<String> {
231        vec![
232            format!("LinReg Upper({}, {})", self.period, self.upper_multiplier),
233            format!("LinReg Middle({})", self.period),
234            format!("LinReg Lower({}, {})", self.period, self.lower_multiplier),
235        ]
236    }
237}