Skip to main content

kestrel_chartkit/indicator/
vortex.rs

1use super::{Indicator, IndicatorAlert, IndicatorOutput};
2use crate::model::Bar;
3use std::collections::{HashMap, VecDeque};
4
5/// Vortex Indicator (+VI and -VI) Engine.
6#[derive(Debug, Clone)]
7pub struct VortexEngine {
8    period: usize,
9    prev_bar: Option<Bar>,
10    vm_plus_window: VecDeque<f64>,
11    vm_minus_window: VecDeque<f64>,
12    tr_window: VecDeque<f64>,
13}
14
15impl VortexEngine {
16    pub fn new(period: usize) -> Self {
17        Self {
18            period: period.max(1),
19            prev_bar: None,
20            vm_plus_window: VecDeque::with_capacity(period),
21            vm_minus_window: VecDeque::with_capacity(period),
22            tr_window: VecDeque::with_capacity(period),
23        }
24    }
25}
26
27impl Indicator for VortexEngine {
28    fn name(&self) -> &str {
29        "vortex"
30    }
31
32    fn warmup_period(&self) -> usize {
33        self.period + 1
34    }
35
36    fn reset(&mut self) {
37        self.prev_bar = None;
38        self.vm_plus_window.clear();
39        self.vm_minus_window.clear();
40        self.tr_window.clear();
41    }
42
43    fn on_bar(&mut self, bar: &Bar) -> Option<IndicatorOutput> {
44        let prev = match &self.prev_bar {
45            Some(p) => p.clone(),
46            None => {
47                self.prev_bar = Some(bar.clone());
48                return None;
49            }
50        };
51        self.prev_bar = Some(bar.clone());
52
53        let vm_plus = (bar.high - prev.low).abs();
54        let vm_minus = (bar.low - prev.high).abs();
55        let tr = (bar.high - bar.low)
56            .max((bar.high - prev.close).abs())
57            .max((bar.low - prev.close).abs());
58
59        self.vm_plus_window.push_back(vm_plus);
60        self.vm_minus_window.push_back(vm_minus);
61        self.tr_window.push_back(tr);
62
63        if self.vm_plus_window.len() > self.period {
64            self.vm_plus_window.pop_front();
65            self.vm_minus_window.pop_front();
66            self.tr_window.pop_front();
67        }
68
69        if self.vm_plus_window.len() < self.period {
70            return None;
71        }
72
73        let sum_vm_plus: f64 = self.vm_plus_window.iter().sum();
74        let sum_vm_minus: f64 = self.vm_minus_window.iter().sum();
75        let sum_tr: f64 = self.tr_window.iter().sum::<f64>().max(1e-8);
76
77        let vi_plus = sum_vm_plus / sum_tr;
78        let vi_minus = sum_vm_minus / sum_tr;
79
80        let mut extra = HashMap::new();
81        extra.insert("vi_plus".to_string(), vi_plus);
82        extra.insert("vi_minus".to_string(), vi_minus);
83
84        Some(IndicatorOutput::with_extra(vi_plus, extra))
85    }
86
87    fn alerts(&self) -> Vec<IndicatorAlert> {
88        Vec::new()
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    #[test]
97    fn test_vortex_indicator() {
98        let mut vi = VortexEngine::new(14);
99        let mut out = None;
100        for i in 0..20 {
101            let b = Bar::new(i, 100.0, 105.0, 95.0, 100.0 + i as f64, 1000.0);
102            out = vi.on_bar(&b);
103        }
104        assert!(out.is_some());
105        let o = out.unwrap();
106        assert!(o.extra.contains_key("vi_plus"));
107        assert!(o.extra.contains_key("vi_minus"));
108    }
109}