quantwave_plugins/
hilbert.rs1use 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}