kestrel_chartkit/indicator/
smoothing.rs1use std::collections::VecDeque;
2
3#[cfg(feature = "serde")]
4use serde::{Deserialize, Serialize};
5
6#[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#[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#[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#[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#[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 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
219pub trait Smoother: Send + Sync {
222 fn update(&mut self, src: f64) -> Option<f64>;
224 fn reset(&mut self);
225 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
298pub 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 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
334impl 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#[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#[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); assert_eq!(chain.update(2.0), None); 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}