Skip to main content

quantwave_plugins/
hilbert.rs

1use polars::prelude::*;
2use pyo3_polars::derive::polars_expr;
3
4use quantwave_core::indicators::incremental::hilbert_ta::{
5    HT_DCPERIOD, HT_DCPHASE, HT_PHASOR, HT_SINE, HT_TRENDMODE, HT_TRENDLINE
6};
7use quantwave_core::traits::Next;
8
9#[polars_expr(output_type=Float64)]
10fn ht_dcperiod(inputs: &[Series]) -> PolarsResult<Series> {
11    let s = inputs[0].f64()?;
12    let mut indicator = HT_DCPERIOD::new();
13    
14    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
15        Some(v) if !v.is_nan() => Some(indicator.next(v)),
16        Some(_) => Some(f64::NAN),
17        None => None,
18    }).collect();
19    
20    Ok(out.into_series())
21}
22
23#[polars_expr(output_type=Float64)]
24fn ht_dcphase(inputs: &[Series]) -> PolarsResult<Series> {
25    let s = inputs[0].f64()?;
26    let mut indicator = HT_DCPHASE::new();
27    
28    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
29        Some(v) if !v.is_nan() => Some(indicator.next(v)),
30        Some(_) => Some(f64::NAN),
31        None => None,
32    }).collect();
33    
34    Ok(out.into_series())
35}
36
37pub fn ht_phasor_output(_: &[Field]) -> PolarsResult<Field> {
38    Ok(Field::new(
39        "ht_phasor".into(),
40        DataType::Struct(vec![
41            Field::new("inphase".into(), DataType::Float64),
42            Field::new("quadrature".into(), DataType::Float64),
43        ]),
44    ))
45}
46
47#[polars_expr(output_type_func=ht_phasor_output)]
48fn ht_phasor(inputs: &[Series]) -> PolarsResult<Series> {
49    let s = inputs[0].f64()?;
50    let mut indicator = HT_PHASOR::new();
51    
52    let mut inphase_vec = Vec::with_capacity(s.len());
53    let mut quadrature_vec = Vec::with_capacity(s.len());
54    
55    for opt_v in s.into_iter() {
56        match opt_v {
57            Some(v) if !v.is_nan() => {
58                let (i, q) = indicator.next(v);
59                inphase_vec.push(Some(i));
60                quadrature_vec.push(Some(q));
61            }
62            Some(_) => {
63                inphase_vec.push(Some(f64::NAN));
64                quadrature_vec.push(Some(f64::NAN));
65            }
66            None => {
67                inphase_vec.push(None);
68                quadrature_vec.push(None);
69            }
70        }
71    }
72    
73    let ca_inphase = Float64Chunked::new("inphase".into(), inphase_vec);
74    let ca_quadrature = Float64Chunked::new("quadrature".into(), quadrature_vec);
75    
76    let series_vec = vec![ca_inphase.into_series(), ca_quadrature.into_series()];
77    let out = StructChunked::from_series("ht_phasor".into(), s.len(), series_vec.iter())?;
78    
79    Ok(out.into_series())
80}
81
82pub fn ht_sine_output(_: &[Field]) -> PolarsResult<Field> {
83    Ok(Field::new(
84        "ht_sine".into(),
85        DataType::Struct(vec![
86            Field::new("sine".into(), DataType::Float64),
87            Field::new("leadsine".into(), DataType::Float64),
88        ]),
89    ))
90}
91
92#[polars_expr(output_type_func=ht_sine_output)]
93fn ht_sine(inputs: &[Series]) -> PolarsResult<Series> {
94    let s = inputs[0].f64()?;
95    let mut indicator = HT_SINE::new();
96    
97    let mut sine_vec = Vec::with_capacity(s.len());
98    let mut leadsine_vec = Vec::with_capacity(s.len());
99    
100    for opt_v in s.into_iter() {
101        match opt_v {
102            Some(v) if !v.is_nan() => {
103                let (si, ls) = indicator.next(v);
104                sine_vec.push(Some(si));
105                leadsine_vec.push(Some(ls));
106            }
107            Some(_) => {
108                sine_vec.push(Some(f64::NAN));
109                leadsine_vec.push(Some(f64::NAN));
110            }
111            None => {
112                sine_vec.push(None);
113                leadsine_vec.push(None);
114            }
115        }
116    }
117    
118    let ca_sine = Float64Chunked::new("sine".into(), sine_vec);
119    let ca_leadsine = Float64Chunked::new("leadsine".into(), leadsine_vec);
120    
121    let series_vec = vec![ca_sine.into_series(), ca_leadsine.into_series()];
122    let out = StructChunked::from_series("ht_sine".into(), s.len(), series_vec.iter())?;
123    
124    Ok(out.into_series())
125}
126
127#[polars_expr(output_type=Int32)]
128fn ht_trendmode(inputs: &[Series]) -> PolarsResult<Series> {
129    let s = inputs[0].f64()?;
130    let mut indicator = HT_TRENDMODE::new();
131    
132    let out: Int32Chunked = s.into_iter().map(|opt_v| match opt_v {
133        Some(v) if !v.is_nan() => Some(indicator.next(v) as i32),
134        _ => None,
135    }).collect();
136    
137    Ok(out.into_series())
138}
139
140#[polars_expr(output_type=Float64)]
141fn ht_trendline(inputs: &[Series]) -> PolarsResult<Series> {
142    let s = inputs[0].f64()?;
143    let mut indicator = HT_TRENDLINE::new();
144    
145    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
146        Some(v) if !v.is_nan() => Some(indicator.next(v)),
147        Some(_) => Some(f64::NAN),
148        None => None,
149    }).collect();
150    
151    Ok(out.into_series())
152}