egui_charts/studies/builtin/
linear_regression_channel.rs1use 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#[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 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), palette::RED.gamma_multiply_u8(13), ],
47 visible: true,
48 source: PriceSource::Close,
49 extend: false,
50 fit: None,
51 }
52 }
53
54 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 pub fn with_source(mut self, source: PriceSource) -> Self {
64 self.source = source;
65 self
66 }
67
68 pub fn with_extend_right(mut self, extend: bool) -> Self {
70 self.extend = extend;
71 self
72 }
73
74 pub fn correlation(&self) -> Option<f64> {
76 self.fit.and_then(|fit| fit.correlation)
77 }
78
79 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}