ta_lib_in_rust/indicators/trend/
vortex.rs

1use polars::prelude::*;
2
3/// Calculate Vortex Indicator (VI+ and VI-)
4///
5/// Returns (vi_plus, vi_minus) as Series
6pub fn calculate_vortex(
7    df: &DataFrame,
8    high_col: &str,
9    low_col: &str,
10    close_col: &str,
11    period: usize,
12) -> PolarsResult<(Series, Series)> {
13    let high = df.column(high_col)?.f64()?;
14    let low = df.column(low_col)?.f64()?;
15    let close = df.column(close_col)?.f64()?;
16    let len = df.height();
17    let mut tr = vec![f64::NAN; len];
18    let mut vm_plus = vec![f64::NAN; len];
19    let mut vm_minus = vec![f64::NAN; len];
20    for i in 1..len {
21        let h = high.get(i).unwrap_or(f64::NAN);
22        let l = low.get(i).unwrap_or(f64::NAN);
23        let h_prev = high.get(i - 1).unwrap_or(f64::NAN);
24        let l_prev = low.get(i - 1).unwrap_or(f64::NAN);
25        let c_prev = close.get(i - 1).unwrap_or(f64::NAN);
26        tr[i] = h.max(c_prev) - l.min(c_prev);
27        vm_plus[i] = (h - l_prev).abs();
28        vm_minus[i] = (l - h_prev).abs();
29    }
30    let mut vi_plus = vec![f64::NAN; len];
31    let mut vi_minus = vec![f64::NAN; len];
32    for i in 0..len {
33        if i + 1 >= period {
34            let sum_tr: f64 = tr[(i + 1 - period)..=i]
35                .iter()
36                .filter(|x| !x.is_nan())
37                .sum();
38            let sum_vm_plus: f64 = vm_plus[(i + 1 - period)..=i]
39                .iter()
40                .filter(|x| !x.is_nan())
41                .sum();
42            let sum_vm_minus: f64 = vm_minus[(i + 1 - period)..=i]
43                .iter()
44                .filter(|x| !x.is_nan())
45                .sum();
46            if sum_tr != 0.0 {
47                vi_plus[i] = sum_vm_plus / sum_tr;
48                vi_minus[i] = sum_vm_minus / sum_tr;
49            }
50        }
51    }
52    Ok((
53        Series::new("vi_plus".into(), vi_plus),
54        Series::new("vi_minus".into(), vi_minus),
55    ))
56}
ta_lib_in_rust/indicators/trend/vortex.rs

ta_lib_in_rust/indicators/trend/
vortex.rs
