egui_charts/studies/builtin/
squeeze_momentum.rs

1use super::{clamp_param, over_finite_bars};
2use crate::model::Bar;
3use crate::studies::helpers;
4use crate::studies::palette;
5use crate::studies::stepper::{update_window, usable_behind};
6use crate::studies::{Indicator, IndicatorValue, IndicatorValues};
7use egui::Color32;
8
9/// Squeeze Momentum Indicator
10/// Combines Bollinger Bands and Keltner Channels to identify periods of low volatility
11/// When BB is inside KC = "squeeze" (low volatility period)
12#[derive(Clone)]
13pub struct SqueezeMomentum {
14    /// Bollinger Bands period (typically 20)
15    bb_period: usize,
16    /// Bollinger Bands standard deviation multiplier (typically 2.0)
17    bb_mult: f64,
18    /// Keltner Channel period (typically 20)
19    kc_period: usize,
20    /// Keltner Channel ATR multiplier (typically 1.5)
21    kc_mult: f64,
22    /// Momentum period (typically 12)
23    momentum_period: usize,
24    values: IndicatorValues,
25    /// Colors: [Momentum positive, Squeeze on (the squeeze line), Momentum negative, Squeeze off]
26    colors: Vec<Color32>,
27    visible: bool,
28}
29
30impl SqueezeMomentum {
31    /// Create with the default Bollinger and Keltner channel settings.
32    pub fn new() -> Self {
33        Self {
34            bb_period: 20,
35            bb_mult: 2.0,
36            kc_period: 20,
37            kc_mult: 1.5,
38            momentum_period: 12,
39            values: IndicatorValues::new(),
40            colors: vec![
41                palette::GREEN,    // Green - Momentum positive
42                palette::DISABLED, // Gray - Squeeze on
43                palette::RED,      // Red - Momentum negative
44                palette::BLUE,     // Blue - Squeeze off
45            ],
46            visible: true,
47        }
48    }
49
50    /// Set the Bollinger Bands period and multiplier; `mult` is clamped into 0
51    /// to `f64::MAX`, NaN to 0.
52    pub fn with_bb_params(mut self, period: usize, mult: f64) -> Self {
53        self.bb_period = period.max(1);
54        self.bb_mult = clamp_param(mult, 0.0, f64::MAX);
55        self
56    }
57
58    /// Set the Keltner Channel period and multiplier; `mult` is clamped into 0
59    /// to `f64::MAX`, NaN to 0.
60    pub fn with_kc_params(mut self, period: usize, mult: f64) -> Self {
61        self.kc_period = period.max(1);
62        self.kc_mult = clamp_param(mult, 0.0, f64::MAX);
63        self
64    }
65
66    /// Set plot colors in the order returned by [`Indicator::line_names`].
67    pub fn with_colors(mut self, colors: Vec<Color32>) -> Self {
68        self.colors = colors;
69        self
70    }
71
72    fn calculate_atr(data: &[Bar], period: usize) -> Vec<f64> {
73        let mut tr = Vec::with_capacity(data.len());
74        tr.push(data[0].high - data[0].low);
75
76        for i in 1..data.len() {
77            let high_low = data[i].high - data[i].low;
78            let high_close = (data[i].high - data[i - 1].close).abs();
79            let low_close = (data[i].low - data[i - 1].close).abs();
80            tr.push(high_low.max(high_close).max(low_close));
81        }
82
83        helpers::sma(&tr, period)
84    }
85
86    fn linear_regression(data: &[f64], period: usize) -> Vec<f64> {
87        let mut result = Vec::with_capacity(data.len());
88
89        for i in 0..data.len() {
90            if i < period - 1 {
91                result.push(f64::NAN);
92            } else {
93                let window = &data[i + 1 - period..=i];
94                let n = period as f64;
95
96                let sum_x: f64 = (0..period).map(|x| x as f64).sum();
97                let sum_y: f64 = window.iter().sum();
98                let sum_xy: f64 = window.iter().enumerate().map(|(x, &y)| x as f64 * y).sum();
99                let sum_xx: f64 = (0..period).map(|x| (x * x) as f64).sum();
100
101                let slope = (n * sum_xy - sum_x * sum_y) / (n * sum_xx - sum_x * sum_x);
102                let intercept = (sum_y - slope * sum_x) / n;
103
104                // Value at the end of the regression line
105                let value = intercept + slope * (period - 1) as f64;
106                result.push(value);
107            }
108        }
109
110        result
111    }
112}
113
114impl Default for SqueezeMomentum {
115    fn default() -> Self {
116        Self::new()
117    }
118}
119
120impl SqueezeMomentum {
121    /// The values over `data`, every bar of which has finite inputs.
122    fn compute(&self, data: &[Bar]) -> IndicatorValues {
123        let mut values = IndicatorValues::new();
124        let required = self.bb_period.max(self.kc_period).max(self.momentum_period);
125        if data.len() < required {
126            for _ in 0..data.len() {
127                values.push(IndicatorValue::None);
128            }
129            return values;
130        }
131
132        // Get closes
133        let closes: Vec<f64> = data.iter().map(|b| b.close).collect();
134
135        // Calculate Bollinger Bands
136        let bb_sma = helpers::sma(&closes, self.bb_period);
137        let bb_std = helpers::stdev(&closes, self.bb_period);
138
139        // Calculate Keltner Channels
140        let kc_sma = helpers::sma(&closes, self.kc_period);
141        let atr = Self::calculate_atr(data, self.kc_period);
142
143        // Calculate midline and deviations for squeeze calculation
144        let hl2: Vec<f64> = data.iter().map(|b| (b.high + b.low) / 2.0).collect();
145        let hl2_sma = helpers::sma(&hl2, self.momentum_period);
146
147        // Calculate momentum using linear regression
148        let mut deviation = Vec::with_capacity(data.len());
149        for i in 0..data.len() {
150            if hl2_sma[i].is_nan() {
151                deviation.push(f64::NAN);
152            } else {
153                deviation.push(closes[i] - hl2_sma[i]);
154            }
155        }
156        let momentum = Self::linear_regression(&deviation, self.momentum_period);
157
158        // Store results: [momentum, squeeze_state (1.0 = squeeze on, 0.0 = off)]
159        for i in 0..data.len() {
160            if bb_sma[i].is_nan()
161                || bb_std[i].is_nan()
162                || kc_sma[i].is_nan()
163                || atr[i].is_nan()
164                || momentum[i].is_nan()
165            {
166                values.push(IndicatorValue::None);
167            } else {
168                let bb_upper = bb_sma[i] + self.bb_mult * bb_std[i];
169                let bb_lower = bb_sma[i] - self.bb_mult * bb_std[i];
170                let kc_upper = kc_sma[i] + self.kc_mult * atr[i];
171                let kc_lower = kc_sma[i] - self.kc_mult * atr[i];
172
173                // Squeeze is ON when BB is inside KC
174                let squeeze_on = bb_lower > kc_lower && bb_upper < kc_upper;
175
176                values.push_lines(&[momentum[i], if squeeze_on { 1.0 } else { 0.0 }]);
177            }
178        }
179
180        values
181    }
182}
183
184impl Indicator for SqueezeMomentum {
185    crate::studies::inputs::study_inputs!(self;
186        bb_period: length("Bb period", 1, usize::MAX),
187        bb_mult: number("Bb mult", 0.0, f64::MAX),
188        kc_period: length("Kc period", 1, usize::MAX),
189        kc_mult: number("Kc mult", 0.0, f64::MAX),
190        momentum_period: length("Momentum period", 1, usize::MAX)
191    ; {});
192
193    fn value_format(&self) -> crate::scales::StudyFormat {
194        crate::scales::StudyFormat::Price
195    }
196
197    fn name(&self) -> &str {
198        "Squeeze"
199    }
200
201    fn desc(&self) -> &str {
202        "Squeeze Momentum - Identifies volatility breakouts"
203    }
204
205    fn calculate(&mut self, data: &[Bar]) {
206        self.update(data, 0);
207    }
208
209    fn update(&mut self, data: &[Bar], from: usize) {
210        let mut values = std::mem::take(&mut self.values);
211        update_window(
212            &mut values,
213            data,
214            from,
215            usable_behind(
216                (self.bb_period - 1)
217                    .max(self.kc_period)
218                    .max((self.momentum_period - 1).saturating_mul(2)),
219                |b| b.high.is_finite() && b.low.is_finite() && b.close.is_finite(),
220            ),
221            |data| {
222                over_finite_bars(
223                    data,
224                    |b| b.high.is_finite() && b.low.is_finite() && b.close.is_finite(),
225                    |data| self.compute(data),
226                )
227            },
228        );
229        self.values = values;
230    }
231
232    fn is_incremental(&self) -> bool {
233        true
234    }
235
236    fn values(&self) -> &IndicatorValues {
237        &self.values
238    }
239
240    fn colors(&self) -> Vec<Color32> {
241        self.colors.clone()
242    }
243
244    fn set_colors(&mut self, colors: Vec<Color32>) {
245        if colors.len() >= 4 {
246            self.colors = colors;
247        }
248    }
249
250    fn is_overlay(&self) -> bool {
251        false
252    }
253
254    fn is_visible(&self) -> bool {
255        self.visible
256    }
257
258    fn set_visible(&mut self, visible: bool) {
259        self.visible = visible;
260    }
261
262    fn clone_box(&self) -> Box<dyn Indicator> {
263        Box::new(self.clone())
264    }
265
266    fn line_names(&self) -> Vec<String> {
267        vec!["Momentum".to_string(), "Squeeze".to_string()]
268    }
269}
270
271#[cfg(test)]
272mod tests {
273    use super::*;
274    use crate::studies::test_support::make_range_bar as make_bar;
275
276    #[test]
277    fn test_squeeze_calculation() {
278        let mut squeeze = SqueezeMomentum::new();
279
280        let data: Vec<Bar> = (0..50)
281            .map(|i| make_bar(102.0 + i as f64, 98.0 + i as f64, 100.0 + i as f64))
282            .collect();
283
284        squeeze.calculate(&data);
285
286        assert_eq!(squeeze.values.len(), 50);
287
288        // Should have momentum and squeeze state
289        let rows = squeeze.values.to_rows();
290        let IndicatorValue::Multiple(v) = rows.last().unwrap() else {
291            panic!("no value: {:?}", squeeze.values.to_rows().last().unwrap());
292        };
293        assert_eq!(v.len(), 2);
294    }
295}
egui_charts/studies/builtin/squeeze_momentum.rs

egui_charts/studies/builtin/
squeeze_momentum.rs
