Skip to main content

kestrel_chartkit/indicator/
smoothing.rs

1use std::collections::VecDeque;
2
3#[cfg(feature = "serde")]
4use serde::{Deserialize, Serialize};
5
6/// Pine's `ta.ema`: seeds with the input value itself on the first sample
7#[derive(Debug, Clone, Copy, Default)]
8pub struct Ema {
9    len: usize,
10    state: Option<f64>,
11}
12
13impl Ema {
14    pub fn new(len: usize) -> Self {
15        Self { len, state: None }
16    }
17
18    pub fn update(&mut self, src: f64) -> f64 {
19        let alpha = 2.0 / (self.len as f64 + 1.0);
20        let next = match self.state {
21            None => src,
22            Some(prev) => alpha * src + (1.0 - alpha) * prev,
23        };
24        self.state = Some(next);
25        next
26    }
27
28    pub fn reset(&mut self) {
29        self.state = None;
30    }
31}
32
33/// Pine's `ta.wma` weighted moving average for scalar streams
34#[derive(Debug, Clone)]
35pub struct Wma {
36    len: usize,
37    window: VecDeque<f64>,
38}
39
40impl Wma {
41    pub fn new(len: usize) -> Self {
42        Self {
43            len: len.max(1),
44            window: VecDeque::with_capacity(len),
45        }
46    }
47
48    pub fn update(&mut self, src: f64) -> Option<f64> {
49        self.window.push_back(src);
50        if self.window.len() > self.len {
51            self.window.pop_front();
52        }
53        if self.window.len() < self.len {
54            return None;
55        }
56
57        let denom = (self.len * (self.len + 1)) as f64 / 2.0;
58        let mut sum = 0.0;
59        for (i, &val) in self.window.iter().enumerate() {
60            sum += val * (i + 1) as f64;
61        }
62        Some(sum / denom)
63    }
64
65    pub fn reset(&mut self) {
66        self.window.clear();
67    }
68}
69
70/// Pine's `ta.rma` (Wilder smoothing): seeds with SMA of the first `len` samples
71#[derive(Debug, Clone)]
72pub struct Rma {
73    len: usize,
74    seed: VecDeque<f64>,
75    state: Option<f64>,
76}
77
78impl Rma {
79    pub fn new(len: usize) -> Self {
80        Self {
81            len,
82            seed: VecDeque::with_capacity(len),
83            state: None,
84        }
85    }
86
87    pub fn update(&mut self, src: f64) -> Option<f64> {
88        if let Some(prev) = self.state {
89            let alpha = 1.0 / self.len as f64;
90            let next = alpha * src + (1.0 - alpha) * prev;
91            self.state = Some(next);
92            return Some(next);
93        }
94        self.seed.push_back(src);
95        if self.seed.len() < self.len {
96            return None;
97        }
98        let sma = self.seed.iter().sum::<f64>() / self.len as f64;
99        self.state = Some(sma);
100        Some(sma)
101    }
102
103    pub fn reset(&mut self) {
104        self.seed.clear();
105        self.state = None;
106    }
107}
108
109/// Pine's `ta.sma`: plain windowed average
110#[derive(Debug, Clone)]
111pub struct Sma {
112    len: usize,
113    window: VecDeque<f64>,
114    sum: f64,
115}
116
117impl Sma {
118    pub fn new(len: usize) -> Self {
119        Self {
120            len,
121            window: VecDeque::with_capacity(len),
122            sum: 0.0,
123        }
124    }
125
126    pub fn update(&mut self, src: f64) -> Option<f64> {
127        self.window.push_back(src);
128        self.sum += src;
129        if self.window.len() > self.len {
130            self.sum -= self.window.pop_front().unwrap();
131        }
132        if self.window.len() < self.len {
133            return None;
134        }
135        Some(self.sum / self.len as f64)
136    }
137
138    pub fn reset(&mut self) {
139        self.window.clear();
140        self.sum = 0.0;
141    }
142}
143
144/// Rolling window for finding highest and lowest values
145#[derive(Debug, Clone)]
146pub struct ExtremeWindow {
147    len: usize,
148    window: VecDeque<f64>,
149}
150
151impl ExtremeWindow {
152    pub fn new(len: usize) -> Self {
153        Self {
154            len,
155            window: VecDeque::with_capacity(len),
156        }
157    }
158
159    pub fn push(&mut self, value: f64) -> Option<(f64, f64)> {
160        if self.window.len() == self.len {
161            self.window.pop_front();
162        }
163        self.window.push_back(value);
164        if self.window.len() < self.len {
165            return None;
166        }
167        let lowest = self.window.iter().cloned().fold(f64::INFINITY, f64::min);
168        let highest = self
169            .window
170            .iter()
171            .cloned()
172            .fold(f64::NEG_INFINITY, f64::max);
173        Some((lowest, highest))
174    }
175
176    pub fn reset(&mut self) {
177        self.window.clear();
178    }
179}
180
181pub fn crossed_over(prev_a: f64, prev_b: f64, a: f64, b: f64) -> bool {
182    prev_a <= prev_b && a > b
183}
184
185pub fn crossed_under(prev_a: f64, prev_b: f64, a: f64, b: f64) -> bool {
186    prev_a >= prev_b && a < b
187}
188
189#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
190#[cfg_attr(
191    feature = "serde",
192    derive(Serialize, Deserialize),
193    serde(rename_all = "snake_case")
194)]
195pub enum SmootherKind {
196    #[default]
197    Ema,
198    Sma,
199    Rma,
200    Alma,
201    Jma,
202}
203
204impl SmootherKind {
205    /// Builds a boxed [`Smoother`] of this kind with `len` bars of lookback/decay, using common
206    /// Pine defaults for `Alma` (`offset = 0.85`, `sigma = 6.0`) and `Jma` (`phase = 0.0`,
207    /// `power = 2.0`). Use the concrete constructors directly to override those.
208    pub fn build(self, len: usize) -> Box<dyn Smoother> {
209        match self {
210            SmootherKind::Ema => Box::new(Ema::new(len)),
211            SmootherKind::Sma => Box::new(Sma::new(len)),
212            SmootherKind::Rma => Box::new(Rma::new(len)),
213            SmootherKind::Alma => Box::new(Alma::new(len, 0.85, 6.0)),
214            SmootherKind::Jma => Box::new(Jma::new(len, 0.0, 2.0)),
215        }
216    }
217}
218
219/// Common contract for streaming smoothers, letting them be chained ([`SmootherChain`]) or
220/// selected dynamically ([`SmootherKind::build`]) regardless of concrete type.
221pub trait Smoother: Send + Sync {
222    /// Feeds one value. Returns `None` while still inside this smoother's own warmup.
223    fn update(&mut self, src: f64) -> Option<f64>;
224    fn reset(&mut self);
225    /// Bars needed before this smoother first returns `Some`. `0` for smoothers that emit from
226    /// the first sample (`Ema`, `Jma`).
227    fn warmup_period(&self) -> usize {
228        0
229    }
230}
231
232impl Smoother for Ema {
233    fn update(&mut self, src: f64) -> Option<f64> {
234        Some(Ema::update(self, src))
235    }
236    fn reset(&mut self) {
237        Ema::reset(self)
238    }
239}
240
241impl Smoother for Sma {
242    fn update(&mut self, src: f64) -> Option<f64> {
243        Sma::update(self, src)
244    }
245    fn reset(&mut self) {
246        Sma::reset(self)
247    }
248    fn warmup_period(&self) -> usize {
249        self.len
250    }
251}
252
253impl Smoother for Rma {
254    fn update(&mut self, src: f64) -> Option<f64> {
255        Rma::update(self, src)
256    }
257    fn reset(&mut self) {
258        Rma::reset(self)
259    }
260    fn warmup_period(&self) -> usize {
261        self.len
262    }
263}
264
265impl Smoother for Wma {
266    fn update(&mut self, src: f64) -> Option<f64> {
267        Wma::update(self, src)
268    }
269    fn reset(&mut self) {
270        Wma::reset(self)
271    }
272    fn warmup_period(&self) -> usize {
273        self.len
274    }
275}
276
277impl Smoother for Alma {
278    fn update(&mut self, src: f64) -> Option<f64> {
279        Alma::update(self, src)
280    }
281    fn reset(&mut self) {
282        Alma::reset(self)
283    }
284    fn warmup_period(&self) -> usize {
285        self.len
286    }
287}
288
289impl Smoother for Jma {
290    fn update(&mut self, src: f64) -> Option<f64> {
291        Some(Jma::update(self, src))
292    }
293    fn reset(&mut self) {
294        Jma::reset(self)
295    }
296}
297
298/// A typed, ordered pipeline of [`Smoother`] stages: each stage's output feeds the next stage's
299/// input. Its own warmup/reset contract composes cleanly from its stages':
300/// [`SmootherChain::warmup_period`] is the sum of every stage's warmup (a downstream stage cannot
301/// start accumulating until its upstream first emits), and [`SmootherChain::reset`] resets every
302/// stage. Works with any mix of existing `Smoother` impls, and with future ones without changes
303/// here — implement [`Smoother`] and it is chainable.
304pub struct SmootherChain {
305    stages: Vec<Box<dyn Smoother>>,
306}
307
308impl SmootherChain {
309    pub fn new(stages: Vec<Box<dyn Smoother>>) -> Self {
310        Self { stages }
311    }
312
313    pub fn warmup_period(&self) -> usize {
314        self.stages.iter().map(|s| s.warmup_period()).sum()
315    }
316
317    pub fn reset(&mut self) {
318        for stage in &mut self.stages {
319            stage.reset();
320        }
321    }
322
323    /// Feeds `src` through every stage in order. Returns `None` if any stage is still inside its
324    /// own warmup this bar.
325    pub fn update(&mut self, src: f64) -> Option<f64> {
326        let mut value = src;
327        for stage in &mut self.stages {
328            value = stage.update(value)?;
329        }
330        Some(value)
331    }
332}
333
334/// A [`SmootherChain`] is itself a [`Smoother`], so chains compose (a chain can be one stage of
335/// another chain) and can be used anywhere a single smoother is expected, e.g. as one leg of
336/// [`super::trend_relationship::AdaptiveTrendRelationship`].
337impl Smoother for SmootherChain {
338    fn update(&mut self, src: f64) -> Option<f64> {
339        SmootherChain::update(self, src)
340    }
341    fn reset(&mut self) {
342        SmootherChain::reset(self)
343    }
344    fn warmup_period(&self) -> usize {
345        SmootherChain::warmup_period(self)
346    }
347}
348
349/// Arnaud Legoux Moving Average (ALMA)
350#[derive(Debug, Clone)]
351pub struct Alma {
352    len: usize,
353    offset: f64,
354    sigma: f64,
355    window: VecDeque<f64>,
356    weights: Vec<f64>,
357    sum_weights: f64,
358}
359
360impl Alma {
361    pub fn new(len: usize, offset: f64, sigma: f64) -> Self {
362        let len = len.max(1);
363        let m = offset * (len - 1) as f64;
364        let s = (len as f64 / sigma).max(1e-6);
365
366        let mut weights = Vec::with_capacity(len);
367        let mut sum_weights = 0.0;
368        for i in 0..len {
369            let w = (-(i as f64 - m).powi(2) / (2.0 * s * s)).exp();
370            weights.push(w);
371            sum_weights += w;
372        }
373
374        Self {
375            len,
376            offset,
377            sigma,
378            window: VecDeque::with_capacity(len),
379            weights,
380            sum_weights,
381        }
382    }
383
384    pub fn offset(&self) -> f64 {
385        self.offset
386    }
387
388    pub fn sigma(&self) -> f64 {
389        self.sigma
390    }
391
392    pub fn update(&mut self, src: f64) -> Option<f64> {
393        self.window.push_back(src);
394        if self.window.len() > self.len {
395            self.window.pop_front();
396        }
397        if self.window.len() < self.len {
398            return None;
399        }
400
401        let mut weighted_sum = 0.0;
402        for (i, &val) in self.window.iter().enumerate() {
403            weighted_sum += val * self.weights[i];
404        }
405        Some(weighted_sum / self.sum_weights)
406    }
407
408    pub fn reset(&mut self) {
409        self.window.clear();
410    }
411}
412
413/// Open Jurik-style moving-average approximation used by the Pine sources.
414#[derive(Debug, Clone)]
415pub struct Jma {
416    len: usize,
417    phase: f64,
418    power: f64,
419    e0: f64,
420    e1: f64,
421    e2: f64,
422    jma: f64,
423    initialized: bool,
424}
425
426impl Jma {
427    pub fn new(len: usize, phase: f64, power: f64) -> Self {
428        Self {
429            len: len.max(1),
430            phase: phase.clamp(-100.0, 100.0),
431            power: power.max(1.0),
432            e0: 0.0,
433            e1: 0.0,
434            e2: 0.0,
435            jma: 0.0,
436            initialized: false,
437        }
438    }
439
440    pub fn phase(&self) -> f64 {
441        self.phase
442    }
443
444    pub fn update(&mut self, src: f64) -> f64 {
445        if !self.initialized {
446            self.e0 = src;
447            self.e1 = 0.0;
448            self.e2 = 0.0;
449            self.jma = src;
450            self.initialized = true;
451            return src;
452        }
453
454        let phase_ratio = self.phase / 100.0 + 1.5;
455        let length_term = 0.45 * (self.len.saturating_sub(1)) as f64;
456        let beta = length_term / (length_term + 2.0);
457        let alpha = beta.powf(self.power);
458        self.e0 = (1.0 - alpha) * src + alpha * self.e0;
459        self.e1 = (src - self.e0) * (1.0 - beta) + beta * self.e1;
460        self.e2 = (self.e0 + phase_ratio * self.e1 - self.jma) * (1.0 - alpha).powi(2)
461            + alpha.powi(2) * self.e2;
462        self.jma += self.e2;
463        self.jma
464    }
465
466    pub fn reset(&mut self) {
467        self.e0 = 0.0;
468        self.e1 = 0.0;
469        self.e2 = 0.0;
470        self.jma = 0.0;
471        self.initialized = false;
472    }
473}
474
475#[cfg(test)]
476mod jma_tests {
477    use super::Jma;
478
479    #[test]
480    fn phase_changes_the_open_jurik_approximation() {
481        let mut leading = Jma::new(7, 100.0, 2.0);
482        let mut lagging = Jma::new(7, -100.0, 2.0);
483        let input = [10.0, 11.0, 13.0, 12.0, 15.0];
484        let leading_value = input.into_iter().map(|v| leading.update(v)).last().unwrap();
485        let lagging_value = input.into_iter().map(|v| lagging.update(v)).last().unwrap();
486        assert!(leading_value > lagging_value);
487    }
488
489    #[test]
490    fn matches_pine_reference_formula_fixture() {
491        let mut jma = Jma::new(3, 0.0, 2.0);
492        let actual: Vec<_> = [1.0, 2.0, 3.0, 4.0]
493            .into_iter()
494            .map(|value| jma.update(value))
495            .collect();
496        let expected = [
497            1.0,
498            1.819_360_773_771_215_6,
499            2.819_354_006_908_239,
500            3.831_322_396_018_062,
501        ];
502        for (actual, expected) in actual.iter().zip(expected) {
503            assert!((actual - expected).abs() < 1e-12, "{actual} != {expected}");
504        }
505    }
506}
507
508#[cfg(test)]
509mod chain_tests {
510    use super::*;
511
512    #[test]
513    fn test_chain_warmup_is_sum_of_stage_warmups() {
514        let chain =
515            SmootherChain::new(vec![SmootherKind::Sma.build(3), SmootherKind::Rma.build(4)]);
516        assert_eq!(chain.warmup_period(), 3 + 4);
517    }
518
519    #[test]
520    fn test_chain_none_until_every_stage_warm() {
521        let mut chain =
522            SmootherChain::new(vec![SmootherKind::Sma.build(2), SmootherKind::Sma.build(2)]);
523        assert_eq!(chain.update(1.0), None); // stage 1 still cold
524        assert_eq!(chain.update(2.0), None); // stage 1 warm (sma=1.5), stage 2 gets its 1st input
525                                             // stage 1 sma(2,3)=2.5; stage 2 now has both its inputs (1.5, 2.5) -> warm.
526        let value = chain.update(3.0).unwrap();
527        assert!((value - 2.0).abs() < 1e-9);
528        let value = chain.update(4.0).unwrap();
529        assert!((value - 3.0).abs() < 1e-9);
530    }
531
532    #[test]
533    fn test_chain_reset_clears_every_stage() {
534        let mut chain = SmootherChain::new(vec![SmootherKind::Sma.build(2)]);
535        chain.update(1.0);
536        assert!(chain.update(2.0).is_some());
537        chain.reset();
538        assert_eq!(chain.update(5.0), None, "reset stage must re-enter warmup");
539    }
540
541    #[test]
542    fn test_smoother_kind_build_matches_direct_construction() {
543        let mut via_kind = SmootherKind::Ema.build(5);
544        let mut direct = Ema::new(5);
545        for v in [10.0, 11.0, 12.0, 9.0] {
546            assert_eq!(via_kind.update(v), Some(Ema::update(&mut direct, v)));
547        }
548    }
549}