Skip to main content

quantwave_plugins/
custom_8.rs

1use polars::prelude::*;
2use pyo3_polars::derive::polars_expr;
3use serde::Deserialize;
4use quantwave_core::*;
5use quantwave_core::traits::Next;
6
7#[derive(Deserialize)]
8struct AutoTuneFilterKwargs {
9    window: usize,
10    bandwidth: f64,
11}
12
13#[polars_expr(output_type=Float64)]
14fn autotune_filter(inputs: &[Series], kwargs: AutoTuneFilterKwargs) -> PolarsResult<Series> {
15    let s = inputs[0].f64()?;
16    let mut indicator = AutoTuneFilter::new(kwargs.window, kwargs.bandwidth);
17    let mut values = Vec::with_capacity(s.len());
18    for i in 0..s.len() {
19        let val = s.get(i).unwrap_or(f64::NAN);
20        values.push(indicator.next(val));
21    }
22    Ok(Series::new("autotune".into(), values))
23}
24
25#[derive(Deserialize)]
26struct TRAdjEMAKwargs {
27    period: usize,
28    pds: usize,
29    mltp: f64,
30}
31
32#[polars_expr(output_type=Float64)]
33fn tradj_ema(inputs: &[Series], kwargs: TRAdjEMAKwargs) -> PolarsResult<Series> {
34    let high = inputs[0].f64()?;
35    let low = inputs[1].f64()?;
36    let close = inputs[2].f64()?;
37
38    let mut indicator = TRAdjEMA::new(kwargs.period, kwargs.pds, kwargs.mltp);
39    let mut values = Vec::with_capacity(high.len());
40    for i in 0..high.len() {
41        let h = high.get(i).unwrap_or(f64::NAN);
42        let l = low.get(i).unwrap_or(f64::NAN);
43        let c = close.get(i).unwrap_or(f64::NAN);
44        values.push(indicator.next((h, l, c)));
45    }
46    Ok(Series::new("tradj_ema".into(), values))
47}
48
49#[derive(Deserialize)]
50struct TEMAKwargs {
51    period: usize,
52}
53
54#[polars_expr(output_type=Float64)]
55fn tema(inputs: &[Series], kwargs: TEMAKwargs) -> PolarsResult<Series> {
56    let s = inputs[0].f64()?;
57    let mut indicator = TEMA::new(kwargs.period);
58    let mut values = Vec::with_capacity(s.len());
59    for i in 0..s.len() {
60        let val = s.get(i).unwrap_or(0.0);
61        values.push(indicator.next(val));
62    }
63    Ok(Series::new("tema".into(), values))
64}
65
66#[derive(Deserialize)]
67struct RSMKKwargs {
68    length: usize,
69    ema_length: usize,
70}
71
72#[polars_expr(output_type=Float64)]
73fn rsmk(inputs: &[Series], kwargs: RSMKKwargs) -> PolarsResult<Series> {
74    let price = inputs[0].f64()?;
75    let benchmark = inputs[1].f64()?;
76
77    let mut indicator = RSMK::new(kwargs.length, kwargs.ema_length);
78    let mut values = Vec::with_capacity(price.len());
79    for i in 0..price.len() {
80        let p = price.get(i).unwrap_or(f64::NAN);
81        let b = benchmark.get(i).unwrap_or(f64::NAN);
82        values.push(indicator.next((p, b)));
83    }
84    Ok(Series::new("rsmk".into(), values))
85}
86
87#[derive(Deserialize)]
88struct ExpDevBandsKwargs {
89    period: usize,
90    multiplier: f64,
91    use_sma: bool,
92}
93
94pub fn exp_dev_bands_output(_: &[Field]) -> PolarsResult<Field> {
95    Ok(Field::new(
96        "exp_dev_bands".into(),
97        DataType::Struct(vec![
98            Field::new("upper".into(), DataType::Float64),
99            Field::new("middle".into(), DataType::Float64),
100            Field::new("lower".into(), DataType::Float64),
101        ]),
102    ))
103}
104
105#[polars_expr(output_type_func=exp_dev_bands_output)]
106fn exp_dev_bands(inputs: &[Series], kwargs: ExpDevBandsKwargs) -> PolarsResult<Series> {
107    let s = inputs[0].f64()?;
108    let mut indicator = ExpDevBands::new(kwargs.period, kwargs.multiplier, kwargs.use_sma);
109    let mut upper_vals = Vec::with_capacity(s.len());
110    let mut mid_vals = Vec::with_capacity(s.len());
111    let mut lower_vals = Vec::with_capacity(s.len());
112
113    for i in 0..s.len() {
114        let val = s.get(i).unwrap_or(f64::NAN);
115        let (upper, mid, lower) = indicator.next(val);
116        upper_vals.push(upper);
117        mid_vals.push(mid);
118        lower_vals.push(lower);
119    }
120
121    let s_upper = Series::new("upper".into(), upper_vals);
122    let s_mid = Series::new("middle".into(), mid_vals);
123    let s_lower = Series::new("lower".into(), lower_vals);
124    
125    let out = StructChunked::from_series(
126        "exp_dev_bands".into(),
127        s.len(),
128        [s_upper, s_mid, s_lower].iter(),
129    )?;
130    Ok(out.into_series())
131}