quantwave_plugins/
momentum.rs

1use polars::prelude::*;
2use pyo3_polars::derive::polars_expr;
3use serde::Deserialize;
4use talib_rs::MaType;
5
6use quantwave_core::indicators::incremental::cci::CCI;
7use quantwave_core::indicators::incremental::cmo::CMO;
8use quantwave_core::indicators::incremental::mom::{MOM, ROC, ROCP, ROCR, ROCR100};
9use quantwave_core::indicators::incremental::ultosc::ULTOSC;
10use quantwave_core::indicators::incremental::willr::WILLR;
11use quantwave_core::indicators::incremental::trix::TRIX;
12use quantwave_core::indicators::incremental::apo::{APO, PPO};
13use quantwave_core::indicators::incremental::sar::SAR;
14use quantwave_core::indicators::incremental::aroon::{AROON, AROONOSC};
15use quantwave_core::indicators::incremental::stoch::{STOCH, STOCHF, STOCHRSI};
16use quantwave_core::indicators::incremental::dmi::{DX, ADX, ADXR, MINUS_DI, PLUS_DI};
17use quantwave_core::indicators::incremental::dm::{MINUS_DM, PLUS_DM};
18use quantwave_core::traits::Next;
19
20#[derive(Deserialize)]
21struct SinglePeriodKwargs {
22    timeperiod: usize,
23}
24
25#[polars_expr(output_type=Float64)]
26fn cci(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
27    let high = inputs[0].f64()?;
28    let low = inputs[1].f64()?;
29    let close = inputs[2].f64()?;
30    let mut indicator = CCI::new(kwargs.timeperiod);
31    let out: Float64Chunked = high.into_iter().zip(low.into_iter()).zip(close.into_iter()).map(|((h, l), c)| {
32        match (h, l, c) {
33            (Some(hv), Some(lv), Some(cv)) if !hv.is_nan() && !lv.is_nan() && !cv.is_nan() => Some(indicator.next((hv, lv, cv))),
34            (Some(_), Some(_), Some(_)) => Some(f64::NAN),
35            _ => None,
36        }
37    }).collect();
38    Ok(out.into_series())
39}
40
41#[polars_expr(output_type=Float64)]
42fn cmo(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
43    let s = inputs[0].f64()?;
44    let mut indicator = CMO::new(kwargs.timeperiod);
45    let out: Float64Chunked = s.into_iter().map(|opt_v| {
46        match opt_v {
47            Some(v) if !v.is_nan() => Some(indicator.next(v)),
48            Some(_) => Some(f64::NAN),
49            None => None,
50        }
51    }).collect();
52    Ok(out.into_series())
53}
54
55#[polars_expr(output_type=Float64)]
56fn mom(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
57    let s = inputs[0].f64()?;
58    let mut indicator = MOM::new(kwargs.timeperiod);
59    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
60        Some(v) if !v.is_nan() => Some(indicator.next(v)),
61        Some(_) => Some(f64::NAN),
62        None => None,
63    }).collect();
64    Ok(out.into_series())
65}
66
67#[polars_expr(output_type=Float64)]
68fn roc(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
69    let s = inputs[0].f64()?;
70    let mut indicator = ROC::new(kwargs.timeperiod);
71    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
72        Some(v) if !v.is_nan() => Some(indicator.next(v)),
73        Some(_) => Some(f64::NAN),
74        None => None,
75    }).collect();
76    Ok(out.into_series())
77}
78
79#[polars_expr(output_type=Float64)]
80fn rocp(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
81    let s = inputs[0].f64()?;
82    let mut indicator = ROCP::new(kwargs.timeperiod);
83    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
84        Some(v) if !v.is_nan() => Some(indicator.next(v)),
85        Some(_) => Some(f64::NAN),
86        None => None,
87    }).collect();
88    Ok(out.into_series())
89}
90
91#[polars_expr(output_type=Float64)]
92fn rocr(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
93    let s = inputs[0].f64()?;
94    let mut indicator = ROCR::new(kwargs.timeperiod);
95    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
96        Some(v) if !v.is_nan() => Some(indicator.next(v)),
97        Some(_) => Some(f64::NAN),
98        None => None,
99    }).collect();
100    Ok(out.into_series())
101}
102
103#[polars_expr(output_type=Float64)]
104fn rocr100(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
105    let s = inputs[0].f64()?;
106    let mut indicator = ROCR100::new(kwargs.timeperiod);
107    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
108        Some(v) if !v.is_nan() => Some(indicator.next(v)),
109        Some(_) => Some(f64::NAN),
110        None => None,
111    }).collect();
112    Ok(out.into_series())
113}
114
115#[polars_expr(output_type=Float64)]
116fn trix(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
117    let s = inputs[0].f64()?;
118    let mut indicator = TRIX::new(kwargs.timeperiod);
119    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
120        Some(v) if !v.is_nan() => Some(indicator.next(v)),
121        Some(_) => Some(f64::NAN),
122        None => None,
123    }).collect();
124    Ok(out.into_series())
125}
126
127#[polars_expr(output_type=Float64)]
128fn willr(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
129    let high = inputs[0].f64()?;
130    let low = inputs[1].f64()?;
131    let close = inputs[2].f64()?;
132    let mut indicator = WILLR::new(kwargs.timeperiod);
133    let out: Float64Chunked = high.into_iter().zip(low.into_iter()).zip(close.into_iter()).map(|((h, l), c)| {
134        match (h, l, c) {
135            (Some(hv), Some(lv), Some(cv)) if !hv.is_nan() && !lv.is_nan() && !cv.is_nan() => Some(indicator.next((hv, lv, cv))),
136            (Some(_), Some(_), Some(_)) => Some(f64::NAN),
137            _ => None,
138        }
139    }).collect();
140    Ok(out.into_series())
141}
142
143#[polars_expr(output_type=Float64)]
144fn adx(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
145    let high = inputs[0].f64()?;
146    let low = inputs[1].f64()?;
147    let close = inputs[2].f64()?;
148    let mut indicator = ADX::new(kwargs.timeperiod);
149    let out: Float64Chunked = high.into_iter().zip(low.into_iter()).zip(close.into_iter()).map(|((h, l), c)| {
150        match (h, l, c) {
151            (Some(hv), Some(lv), Some(cv)) if !hv.is_nan() && !lv.is_nan() && !cv.is_nan() => Some(indicator.next((hv, lv, cv))),
152            (Some(_), Some(_), Some(_)) => Some(f64::NAN),
153            _ => None,
154        }
155    }).collect();
156    Ok(out.into_series())
157}
158
159#[polars_expr(output_type=Float64)]
160fn adxr(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
161    let high = inputs[0].f64()?;
162    let low = inputs[1].f64()?;
163    let close = inputs[2].f64()?;
164    let mut indicator = ADXR::new(kwargs.timeperiod);
165    let out: Float64Chunked = high.into_iter().zip(low.into_iter()).zip(close.into_iter()).map(|((h, l), c)| {
166        match (h, l, c) {
167            (Some(hv), Some(lv), Some(cv)) if !hv.is_nan() && !lv.is_nan() && !cv.is_nan() => Some(indicator.next((hv, lv, cv))),
168            (Some(_), Some(_), Some(_)) => Some(f64::NAN),
169            _ => None,
170        }
171    }).collect();
172    Ok(out.into_series())
173}
174
175#[polars_expr(output_type=Float64)]
176fn dx(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
177    let high = inputs[0].f64()?;
178    let low = inputs[1].f64()?;
179    let close = inputs[2].f64()?;
180    let mut indicator = DX::new(kwargs.timeperiod);
181    let out: Float64Chunked = high.into_iter().zip(low.into_iter()).zip(close.into_iter()).map(|((h, l), c)| {
182        match (h, l, c) {
183            (Some(hv), Some(lv), Some(cv)) if !hv.is_nan() && !lv.is_nan() && !cv.is_nan() => Some(indicator.next((hv, lv, cv))),
184            (Some(_), Some(_), Some(_)) => Some(f64::NAN),
185            _ => None,
186        }
187    }).collect();
188    Ok(out.into_series())
189}
190
191#[polars_expr(output_type=Float64)]
192fn plus_di(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
193    let high = inputs[0].f64()?;
194    let low = inputs[1].f64()?;
195    let close = inputs[2].f64()?;
196    let mut indicator = PLUS_DI::new(kwargs.timeperiod);
197    let out: Float64Chunked = high.into_iter().zip(low.into_iter()).zip(close.into_iter()).map(|((h, l), c)| {
198        match (h, l, c) {
199            (Some(hv), Some(lv), Some(cv)) if !hv.is_nan() && !lv.is_nan() && !cv.is_nan() => Some(indicator.next((hv, lv, cv))),
200            (Some(_), Some(_), Some(_)) => Some(f64::NAN),
201            _ => None,
202        }
203    }).collect();
204    Ok(out.into_series())
205}
206
207#[polars_expr(output_type=Float64)]
208fn minus_di(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
209    let high = inputs[0].f64()?;
210    let low = inputs[1].f64()?;
211    let close = inputs[2].f64()?;
212    let mut indicator = MINUS_DI::new(kwargs.timeperiod);
213    let out: Float64Chunked = high.into_iter().zip(low.into_iter()).zip(close.into_iter()).map(|((h, l), c)| {
214        match (h, l, c) {
215            (Some(hv), Some(lv), Some(cv)) if !hv.is_nan() && !lv.is_nan() && !cv.is_nan() => Some(indicator.next((hv, lv, cv))),
216            (Some(_), Some(_), Some(_)) => Some(f64::NAN),
217            _ => None,
218        }
219    }).collect();
220    Ok(out.into_series())
221}
222
223#[polars_expr(output_type=Float64)]
224fn plus_dm(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
225    let high = inputs[0].f64()?;
226    let low = inputs[1].f64()?;
227    let mut indicator = PLUS_DM::new(kwargs.timeperiod);
228    let out: Float64Chunked = high.into_iter().zip(low.into_iter()).map(|(h, l)| {
229        match (h, l) {
230            (Some(hv), Some(lv)) if !hv.is_nan() && !lv.is_nan() => Some(indicator.next((hv, lv))),
231            (Some(_), Some(_)) => Some(f64::NAN),
232            _ => None,
233        }
234    }).collect();
235    Ok(out.into_series())
236}
237
238#[polars_expr(output_type=Float64)]
239fn minus_dm(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
240    let high = inputs[0].f64()?;
241    let low = inputs[1].f64()?;
242    let mut indicator = MINUS_DM::new(kwargs.timeperiod);
243    let out: Float64Chunked = high.into_iter().zip(low.into_iter()).map(|(h, l)| {
244        match (h, l) {
245            (Some(hv), Some(lv)) if !hv.is_nan() && !lv.is_nan() => Some(indicator.next((hv, lv))),
246            (Some(_), Some(_)) => Some(f64::NAN),
247            _ => None,
248        }
249    }).collect();
250    Ok(out.into_series())
251}
252
253#[polars_expr(output_type=Float64)]
254fn aroonosc(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
255    let high = inputs[0].f64()?;
256    let low = inputs[1].f64()?;
257    let mut indicator = AROONOSC::new(kwargs.timeperiod);
258    let out: Float64Chunked = high.into_iter().zip(low.into_iter()).map(|(h, l)| {
259        match (h, l) {
260            (Some(hv), Some(lv)) if !hv.is_nan() && !lv.is_nan() => Some(indicator.next((hv, lv))),
261            (Some(_), Some(_)) => Some(f64::NAN),
262            _ => None,
263        }
264    }).collect();
265    Ok(out.into_series())
266}
267
268#[derive(Deserialize)]
269struct UltoscKwargs {
270    timeperiod1: usize,
271    timeperiod2: usize,
272    timeperiod3: usize,
273}
274
275#[polars_expr(output_type=Float64)]
276fn ultosc(inputs: &[Series], kwargs: UltoscKwargs) -> PolarsResult<Series> {
277    let high = inputs[0].f64()?;
278    let low = inputs[1].f64()?;
279    let close = inputs[2].f64()?;
280    let mut indicator = ULTOSC::new(kwargs.timeperiod1, kwargs.timeperiod2, kwargs.timeperiod3);
281    let out: Float64Chunked = high.into_iter().zip(low.into_iter()).zip(close.into_iter()).map(|((h, l), c)| {
282        match (h, l, c) {
283            (Some(hv), Some(lv), Some(cv)) if !hv.is_nan() && !lv.is_nan() && !cv.is_nan() => Some(indicator.next((hv, lv, cv))),
284            (Some(_), Some(_), Some(_)) => Some(f64::NAN),
285            _ => None,
286        }
287    }).collect();
288    Ok(out.into_series())
289}
290
291#[derive(Deserialize)]
292struct ApoKwargs {
293    fastperiod: usize,
294    slowperiod: usize,
295    matype: u8,
296}
297
298#[polars_expr(output_type=Float64)]
299fn apo(inputs: &[Series], kwargs: ApoKwargs) -> PolarsResult<Series> {
300    let s = inputs[0].f64()?;
301    let ma_type = match kwargs.matype {
302        0 => MaType::Sma,
303        1 => MaType::Ema,
304        2 => MaType::Wma,
305        3 => MaType::Dema,
306        4 => MaType::Tema,
307        5 => MaType::Trima,
308        6 => MaType::Kama,
309        7 => MaType::Mama,
310        8 => MaType::T3,
311        _ => MaType::Sma,
312    };
313    let mut indicator = APO::new(kwargs.fastperiod, kwargs.slowperiod, ma_type);
314    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
315        Some(v) if !v.is_nan() => Some(indicator.next(v)),
316        Some(_) => Some(f64::NAN),
317        None => None,
318    }).collect();
319    Ok(out.into_series())
320}
321
322#[polars_expr(output_type=Float64)]
323fn ppo(inputs: &[Series], kwargs: ApoKwargs) -> PolarsResult<Series> {
324    let s = inputs[0].f64()?;
325    let ma_type = match kwargs.matype {
326        0 => MaType::Sma,
327        1 => MaType::Ema,
328        2 => MaType::Wma,
329        3 => MaType::Dema,
330        4 => MaType::Tema,
331        5 => MaType::Trima,
332        6 => MaType::Kama,
333        7 => MaType::Mama,
334        8 => MaType::T3,
335        _ => MaType::Sma,
336    };
337    let mut indicator = PPO::new(kwargs.fastperiod, kwargs.slowperiod, ma_type);
338    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
339        Some(v) if !v.is_nan() => Some(indicator.next(v)),
340        Some(_) => Some(f64::NAN),
341        None => None,
342    }).collect();
343    Ok(out.into_series())
344}
345
346#[derive(Deserialize)]
347struct SarKwargs {
348    optinacceleration: f64,
349    optinmaximum: f64,
350}
351
352#[polars_expr(output_type=Float64)]
353fn sar(inputs: &[Series], kwargs: SarKwargs) -> PolarsResult<Series> {
354    let high = inputs[0].f64()?;
355    let low = inputs[1].f64()?;
356    let mut indicator = SAR::new(kwargs.optinacceleration, kwargs.optinmaximum);
357    let out: Float64Chunked = high.into_iter().zip(low.into_iter()).map(|(h, l)| {
358        match (h, l) {
359            (Some(hv), Some(lv)) if !hv.is_nan() && !lv.is_nan() => Some(indicator.next((hv, lv))),
360            (Some(_), Some(_)) => Some(f64::NAN),
361            _ => None,
362        }
363    }).collect();
364    Ok(out.into_series())
365}
366
367pub fn aroon_output(_: &[Field]) -> PolarsResult<Field> {
368    Ok(Field::new(
369        "aroon".into(),
370        DataType::Struct(vec![
371            Field::new("down".into(), DataType::Float64),
372            Field::new("up".into(), DataType::Float64),
373        ]),
374    ))
375}
376
377#[polars_expr(output_type_func=aroon_output)]
378fn aroon(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
379    let high = inputs[0].f64()?;
380    let low = inputs[1].f64()?;
381    let mut indicator = AROON::new(kwargs.timeperiod);
382    
383    let mut down_vec = Vec::with_capacity(high.len());
384    let mut up_vec = Vec::with_capacity(high.len());
385    
386    for (h, l) in high.into_iter().zip(low.into_iter()) {
387        match (h, l) {
388            (Some(hv), Some(lv)) if !hv.is_nan() && !lv.is_nan() => {
389                let (d, u) = indicator.next((hv, lv));
390                down_vec.push(Some(d));
391                up_vec.push(Some(u));
392            }
393            (Some(_), Some(_)) => {
394                down_vec.push(Some(f64::NAN));
395                up_vec.push(Some(f64::NAN));
396            }
397            _ => {
398                down_vec.push(None);
399                up_vec.push(None);
400            }
401        }
402    }
403    
404    let ca_down = Float64Chunked::new("down".into(), down_vec);
405    let ca_up = Float64Chunked::new("up".into(), up_vec);
406    
407    let series_vec = vec![ca_down.into_series(), ca_up.into_series()];
408    let out = StructChunked::from_series("aroon".into(), high.len(), series_vec.iter())?;
409    Ok(out.into_series())
410}
411
412#[derive(Deserialize)]
413struct StochKwargs {
414    fastk_period: usize,
415    slowk_period: usize,
416    slowk_matype: u8,
417    slowd_period: usize,
418    slowd_matype: u8,
419}
420
421pub fn stoch_output(_: &[Field]) -> PolarsResult<Field> {
422    Ok(Field::new(
423        "stoch".into(),
424        DataType::Struct(vec![
425            Field::new("slowk".into(), DataType::Float64),
426            Field::new("slowd".into(), DataType::Float64),
427        ]),
428    ))
429}
430
431#[polars_expr(output_type_func=stoch_output)]
432fn stoch(inputs: &[Series], kwargs: StochKwargs) -> PolarsResult<Series> {
433    let high = inputs[0].f64()?;
434    let low = inputs[1].f64()?;
435    let close = inputs[2].f64()?;
436    
437    let slowk_matype = match kwargs.slowk_matype {
438        0 => MaType::Sma, 1 => MaType::Ema, 2 => MaType::Wma, 3 => MaType::Dema, 4 => MaType::Tema, 5 => MaType::Trima, 6 => MaType::Kama, 7 => MaType::Mama, 8 => MaType::T3, _ => MaType::Sma,
439    };
440    let slowd_matype = match kwargs.slowd_matype {
441        0 => MaType::Sma, 1 => MaType::Ema, 2 => MaType::Wma, 3 => MaType::Dema, 4 => MaType::Tema, 5 => MaType::Trima, 6 => MaType::Kama, 7 => MaType::Mama, 8 => MaType::T3, _ => MaType::Sma,
442    };
443    
444    let mut indicator = STOCH::new(kwargs.fastk_period, kwargs.slowk_period, slowk_matype, kwargs.slowd_period, slowd_matype);
445    
446    let mut slowk_vec = Vec::with_capacity(high.len());
447    let mut slowd_vec = Vec::with_capacity(high.len());
448    
449    for ((h, l), c) in high.into_iter().zip(low.into_iter()).zip(close.into_iter()) {
450        match (h, l, c) {
451            (Some(hv), Some(lv), Some(cv)) if !hv.is_nan() && !lv.is_nan() && !cv.is_nan() => {
452                let (k, d) = indicator.next((hv, lv, cv));
453                slowk_vec.push(Some(k));
454                slowd_vec.push(Some(d));
455            }
456            (Some(_), Some(_), Some(_)) => {
457                slowk_vec.push(Some(f64::NAN));
458                slowd_vec.push(Some(f64::NAN));
459            }
460            _ => {
461                slowk_vec.push(None);
462                slowd_vec.push(None);
463            }
464        }
465    }
466    
467    let ca_k = Float64Chunked::new("slowk".into(), slowk_vec);
468    let ca_d = Float64Chunked::new("slowd".into(), slowd_vec);
469    
470    let series_vec = vec![ca_k.into_series(), ca_d.into_series()];
471    let out = StructChunked::from_series("stoch".into(), high.len(), series_vec.iter())?;
472    Ok(out.into_series())
473}
474
475#[derive(Deserialize)]
476struct StochfKwargs {
477    fastk_period: usize,
478    fastd_period: usize,
479    fastd_matype: u8,
480}
481
482pub fn stochf_output(_: &[Field]) -> PolarsResult<Field> {
483    Ok(Field::new(
484        "stochf".into(),
485        DataType::Struct(vec![
486            Field::new("fastk".into(), DataType::Float64),
487            Field::new("fastd".into(), DataType::Float64),
488        ]),
489    ))
490}
491
492#[polars_expr(output_type_func=stochf_output)]
493fn stochf(inputs: &[Series], kwargs: StochfKwargs) -> PolarsResult<Series> {
494    let high = inputs[0].f64()?;
495    let low = inputs[1].f64()?;
496    let close = inputs[2].f64()?;
497    
498    let fastd_matype = match kwargs.fastd_matype {
499        0 => MaType::Sma, 1 => MaType::Ema, 2 => MaType::Wma, 3 => MaType::Dema, 4 => MaType::Tema, 5 => MaType::Trima, 6 => MaType::Kama, 7 => MaType::Mama, 8 => MaType::T3, _ => MaType::Sma,
500    };
501    
502    let mut indicator = STOCHF::new(kwargs.fastk_period, kwargs.fastd_period, fastd_matype);
503    
504    let mut fastk_vec = Vec::with_capacity(high.len());
505    let mut fastd_vec = Vec::with_capacity(high.len());
506    
507    for ((h, l), c) in high.into_iter().zip(low.into_iter()).zip(close.into_iter()) {
508        match (h, l, c) {
509            (Some(hv), Some(lv), Some(cv)) if !hv.is_nan() && !lv.is_nan() && !cv.is_nan() => {
510                let (k, d) = indicator.next((hv, lv, cv));
511                fastk_vec.push(Some(k));
512                fastd_vec.push(Some(d));
513            }
514            (Some(_), Some(_), Some(_)) => {
515                fastk_vec.push(Some(f64::NAN));
516                fastd_vec.push(Some(f64::NAN));
517            }
518            _ => {
519                fastk_vec.push(None);
520                fastd_vec.push(None);
521            }
522        }
523    }
524    
525    let ca_k = Float64Chunked::new("fastk".into(), fastk_vec);
526    let ca_d = Float64Chunked::new("fastd".into(), fastd_vec);
527    
528    let series_vec = vec![ca_k.into_series(), ca_d.into_series()];
529    let out = StructChunked::from_series("stochf".into(), high.len(), series_vec.iter())?;
530    Ok(out.into_series())
531}
532
533#[derive(Deserialize)]
534struct StochrsiKwargs {
535    timeperiod: usize,
536    fastk_period: usize,
537    fastd_period: usize,
538    fastd_matype: u8,
539}
540
541pub fn stochrsi_output(_: &[Field]) -> PolarsResult<Field> {
542    Ok(Field::new(
543        "stochrsi".into(),
544        DataType::Struct(vec![
545            Field::new("fastk".into(), DataType::Float64),
546            Field::new("fastd".into(), DataType::Float64),
547        ]),
548    ))
549}
550
551#[polars_expr(output_type_func=stochrsi_output)]
552fn stochrsi(inputs: &[Series], kwargs: StochrsiKwargs) -> PolarsResult<Series> {
553    let s = inputs[0].f64()?;
554    
555    let fastd_matype = match kwargs.fastd_matype {
556        0 => MaType::Sma, 1 => MaType::Ema, 2 => MaType::Wma, 3 => MaType::Dema, 4 => MaType::Tema, 5 => MaType::Trima, 6 => MaType::Kama, 7 => MaType::Mama, 8 => MaType::T3, _ => MaType::Sma,
557    };
558    
559    let mut indicator = STOCHRSI::new(kwargs.timeperiod, kwargs.fastk_period, kwargs.fastd_period, fastd_matype);
560    
561    let mut fastk_vec = Vec::with_capacity(s.len());
562    let mut fastd_vec = Vec::with_capacity(s.len());
563    
564    for opt_v in s.into_iter() {
565        match opt_v {
566            Some(v) if !v.is_nan() => {
567                let (k, d) = indicator.next(v);
568                fastk_vec.push(Some(k));
569                fastd_vec.push(Some(d));
570            }
571            Some(_) => {
572                fastk_vec.push(Some(f64::NAN));
573                fastd_vec.push(Some(f64::NAN));
574            }
575            None => {
576                fastk_vec.push(None);
577                fastd_vec.push(None);
578            }
579        }
580    }
581    
582    let ca_k = Float64Chunked::new("fastk".into(), fastk_vec);
583    let ca_d = Float64Chunked::new("fastd".into(), fastd_vec);
584    
585    let series_vec = vec![ca_k.into_series(), ca_d.into_series()];
586    let out = StructChunked::from_series("stochrsi".into(), s.len(), series_vec.iter())?;
587    Ok(out.into_series())
588}
quantwave_plugins/momentum.rs

quantwave_plugins/
momentum.rs
