1use polars::prelude::*;
2use pyo3_polars::derive::polars_expr;
3use serde::Deserialize;
4
5use quantwave_core::indicators::incremental::overlap_ta::{TRIMA, MIDPOINT, MIDPRICE, KAMA, T3};
6use quantwave_core::indicators::incremental::dema::DEMA;
7use quantwave_core::indicators::incremental::macd_ext::{MACDEXT, MACDFIX};
8use quantwave_core::traits::Next;
9use talib_rs::MaType;
10
11#[derive(Deserialize)]
12struct SinglePeriodKwargs {
13 timeperiod: usize,
14}
15
16#[polars_expr(output_type=Float64)]
17fn trima(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
18 let s = inputs[0].f64()?;
19 let mut indicator = TRIMA::new(kwargs.timeperiod);
20
21 let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
22 Some(v) if !v.is_nan() => Some(indicator.next(v)),
23 Some(_) => Some(f64::NAN),
24 None => None,
25 }).collect();
26
27 Ok(out.into_series())
28}
29
30#[polars_expr(output_type=Float64)]
31fn midpoint(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
32 let s = inputs[0].f64()?;
33 let mut indicator = MIDPOINT::new(kwargs.timeperiod);
34
35 let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
36 Some(v) if !v.is_nan() => Some(indicator.next(v)),
37 Some(_) => Some(f64::NAN),
38 None => None,
39 }).collect();
40
41 Ok(out.into_series())
42}
43
44#[polars_expr(output_type=Float64)]
45fn midprice(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
46 let high = inputs[0].f64()?;
47 let low = inputs[1].f64()?;
48 let mut indicator = MIDPRICE::new(kwargs.timeperiod);
49
50 let out: Float64Chunked = high.into_iter().zip(low.into_iter()).map(|(h, l)| {
51 match (h, l) {
52 (Some(hv), Some(lv)) if !hv.is_nan() && !lv.is_nan() => Some(indicator.next((hv, lv))),
53 (Some(_), Some(_)) => Some(f64::NAN),
54 _ => None,
55 }
56 }).collect();
57
58 Ok(out.into_series())
59}
60
61#[polars_expr(output_type=Float64)]
62fn kama(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
63 let s = inputs[0].f64()?;
64 let mut indicator = KAMA::new(kwargs.timeperiod);
65
66 let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
67 Some(v) if !v.is_nan() => Some(indicator.next(v)),
68 Some(_) => Some(f64::NAN),
69 None => None,
70 }).collect();
71
72 Ok(out.into_series())
73}
74
75#[derive(Deserialize)]
76struct T3Kwargs {
77 timeperiod: usize,
78 vfactor: f64,
79}
80
81#[polars_expr(output_type=Float64)]
82fn t3(inputs: &[Series], kwargs: T3Kwargs) -> PolarsResult<Series> {
83 let s = inputs[0].f64()?;
84 let mut indicator = T3::new(kwargs.timeperiod, kwargs.vfactor);
85
86 let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
87 Some(v) if !v.is_nan() => Some(indicator.next(v)),
88 Some(_) => Some(f64::NAN),
89 None => None,
90 }).collect();
91
92 Ok(out.into_series())
93}
94
95#[polars_expr(output_type=Float64)]
96fn dema(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
97 let s = inputs[0].f64()?;
98 let mut indicator = DEMA::new(kwargs.timeperiod);
99
100 let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
101 Some(v) if !v.is_nan() => Some(indicator.next(v)),
102 Some(_) => Some(f64::NAN),
103 None => None,
104 }).collect();
105
106 Ok(out.into_series())
107}
108
109#[derive(Deserialize)]
110struct MacdExtKwargs {
111 fastperiod: usize,
112 fastmatype: u8,
113 slowperiod: usize,
114 slowmatype: u8,
115 signalperiod: usize,
116 signalmatype: u8,
117}
118
119pub fn macd_ext_output(_: &[Field]) -> PolarsResult<Field> {
120 Ok(Field::new(
121 "macdext".into(),
122 DataType::Struct(vec![
123 Field::new("macd".into(), DataType::Float64),
124 Field::new("signal".into(), DataType::Float64),
125 Field::new("hist".into(), DataType::Float64),
126 ]),
127 ))
128}
129
130fn u8_to_matype(matype: u8) -> MaType {
131 match matype {
132 0 => MaType::Sma,
133 1 => MaType::Ema,
134 2 => MaType::Wma,
135 3 => MaType::Dema,
136 4 => MaType::Tema,
137 5 => MaType::Trima,
138 6 => MaType::Kama,
139 7 => MaType::Mama,
140 8 => MaType::T3,
141 _ => MaType::Sma,
142 }
143}
144
145#[polars_expr(output_type_func=macd_ext_output)]
146fn macdext(inputs: &[Series], kwargs: MacdExtKwargs) -> PolarsResult<Series> {
147 let s = inputs[0].f64()?;
148
149 let mut indicator = MACDEXT::new(
150 kwargs.fastperiod, u8_to_matype(kwargs.fastmatype),
151 kwargs.slowperiod, u8_to_matype(kwargs.slowmatype),
152 kwargs.signalperiod, u8_to_matype(kwargs.signalmatype),
153 );
154
155 let mut macd_vec = Vec::with_capacity(s.len());
156 let mut signal_vec = Vec::with_capacity(s.len());
157 let mut hist_vec = Vec::with_capacity(s.len());
158
159 for opt_v in s.into_iter() {
160 match opt_v {
161 Some(v) if !v.is_nan() => {
162 let (m, s, h) = indicator.next(v);
163 macd_vec.push(Some(m));
164 signal_vec.push(Some(s));
165 hist_vec.push(Some(h));
166 }
167 Some(_) => {
168 macd_vec.push(Some(f64::NAN));
169 signal_vec.push(Some(f64::NAN));
170 hist_vec.push(Some(f64::NAN));
171 }
172 None => {
173 macd_vec.push(None);
174 signal_vec.push(None);
175 hist_vec.push(None);
176 }
177 }
178 }
179
180 let ca_macd = Float64Chunked::new("macd".into(), macd_vec);
181 let ca_signal = Float64Chunked::new("signal".into(), signal_vec);
182 let ca_hist = Float64Chunked::new("hist".into(), hist_vec);
183
184 let series_vec = vec![ca_macd.into_series(), ca_signal.into_series(), ca_hist.into_series()];
185 let out = StructChunked::from_series("macdext".into(), s.len(), series_vec.iter())?;
186
187 Ok(out.into_series())
188}
189
190#[derive(Deserialize)]
191struct MacdFixKwargs {
192 signalperiod: usize,
193}
194
195#[polars_expr(output_type_func=macd_ext_output)]
196fn macdfix(inputs: &[Series], kwargs: MacdFixKwargs) -> PolarsResult<Series> {
197 let s = inputs[0].f64()?;
198
199 let mut indicator = MACDFIX::new(kwargs.signalperiod);
200
201 let mut macd_vec = Vec::with_capacity(s.len());
202 let mut signal_vec = Vec::with_capacity(s.len());
203 let mut hist_vec = Vec::with_capacity(s.len());
204
205 for opt_v in s.into_iter() {
206 match opt_v {
207 Some(v) if !v.is_nan() => {
208 let (m, s, h) = indicator.next(v);
209 macd_vec.push(Some(m));
210 signal_vec.push(Some(s));
211 hist_vec.push(Some(h));
212 }
213 Some(_) => {
214 macd_vec.push(Some(f64::NAN));
215 signal_vec.push(Some(f64::NAN));
216 hist_vec.push(Some(f64::NAN));
217 }
218 None => {
219 macd_vec.push(None);
220 signal_vec.push(None);
221 hist_vec.push(None);
222 }
223 }
224 }
225
226 let ca_macd = Float64Chunked::new("macd".into(), macd_vec);
227 let ca_signal = Float64Chunked::new("signal".into(), signal_vec);
228 let ca_hist = Float64Chunked::new("hist".into(), hist_vec);
229
230 let series_vec = vec![ca_macd.into_series(), ca_signal.into_series(), ca_hist.into_series()];
231 let out = StructChunked::from_series("macdfix".into(), s.len(), series_vec.iter())?;
232
233 Ok(out.into_series())
234}