quantwave_plugins/
lib.rs

1#![allow(clippy::unused_unit)]
2pub mod momentum;
3pub mod volatility;
4pub mod volume;
5pub mod price_transform;
6pub mod overlap;
7pub mod statistics;
8pub mod hilbert;
9
10use polars::prelude::*;
11use pyo3_polars::derive::polars_expr;
12use pyo3::prelude::*;
13use serde::Deserialize;
14
15use quantwave_core::indicators::smoothing::{SMA, EMA};
16use quantwave_core::indicators::incremental::rsi::RSI;
17use quantwave_core::indicators::incremental::macd::MACD;
18use quantwave_core::indicators::incremental::bbands::BBANDS;
19use quantwave_core::traits::Next;
20use talib_rs::MaType;
21
22#[derive(Deserialize)]
23struct SmaKwargs {
24    period: usize,
25}
26
27#[polars_expr(output_type=Float64)]
28fn sma(inputs: &[Series], kwargs: SmaKwargs) -> PolarsResult<Series> {
29    let s = &inputs[0];
30    let s_f64 = s.f64()?;
31    
32    let mut indicator = SMA::new(kwargs.period);
33    
34    let out: Float64Chunked = s_f64.into_iter().map(|opt_v| {
35        match opt_v {
36            Some(v) if !v.is_nan() => Some(indicator.next(v)),
37            Some(_) => Some(f64::NAN),
38            None => None,
39        }
40    }).collect();
41    
42    Ok(out.into_series())
43}
44
45#[derive(Deserialize)]
46struct EmaKwargs {
47    period: usize,
48}
49
50#[polars_expr(output_type=Float64)]
51fn ema(inputs: &[Series], kwargs: EmaKwargs) -> PolarsResult<Series> {
52    let s = &inputs[0];
53    let s_f64 = s.f64()?;
54    
55    let mut indicator = EMA::new(kwargs.period);
56    
57    let out: Float64Chunked = s_f64.into_iter().map(|opt_v| {
58        match opt_v {
59            Some(v) if !v.is_nan() => Some(indicator.next(v)),
60            Some(_) => Some(f64::NAN),
61            None => None,
62        }
63    }).collect();
64    
65    Ok(out.into_series())
66}
67
68#[derive(Deserialize)]
69struct RsiKwargs {
70    timeperiod: usize,
71}
72
73#[polars_expr(output_type=Float64)]
74fn rsi(inputs: &[Series], kwargs: RsiKwargs) -> PolarsResult<Series> {
75    let s = &inputs[0];
76    let s_f64 = s.f64()?;
77    
78    let mut indicator = RSI::new(kwargs.timeperiod);
79    
80    let out: Float64Chunked = s_f64.into_iter().map(|opt_v| {
81        match opt_v {
82            Some(v) if !v.is_nan() => Some(indicator.next(v)),
83            Some(_) => Some(f64::NAN),
84            None => None,
85        }
86    }).collect();
87    
88    Ok(out.into_series())
89}
90
91#[derive(Deserialize)]
92struct MacdKwargs {
93    fast: usize,
94    slow: usize,
95    signal: usize,
96}
97
98pub fn macd_output(_: &[Field]) -> PolarsResult<Field> {
99    Ok(Field::new(
100        "macd".into(),
101        DataType::Struct(vec![
102            Field::new("macd".into(), DataType::Float64),
103            Field::new("signal".into(), DataType::Float64),
104            Field::new("hist".into(), DataType::Float64),
105        ]),
106    ))
107}
108
109#[polars_expr(output_type_func=macd_output)]
110fn macd(inputs: &[Series], kwargs: MacdKwargs) -> PolarsResult<Series> {
111    let s = &inputs[0];
112    let s_f64 = s.f64()?;
113    
114    let mut indicator = MACD::new(kwargs.fast, kwargs.slow, kwargs.signal);
115    
116    let mut macd_vec = Vec::with_capacity(s_f64.len());
117    let mut signal_vec = Vec::with_capacity(s_f64.len());
118    let mut hist_vec = Vec::with_capacity(s_f64.len());
119    
120    for opt_v in s_f64.into_iter() {
121        match opt_v {
122            Some(v) if !v.is_nan() => {
123                let (m, s, h) = indicator.next(v);
124                macd_vec.push(Some(m));
125                signal_vec.push(Some(s));
126                hist_vec.push(Some(h));
127            }
128            Some(_) => {
129                macd_vec.push(Some(f64::NAN));
130                signal_vec.push(Some(f64::NAN));
131                hist_vec.push(Some(f64::NAN));
132            }
133            None => {
134                macd_vec.push(None);
135                signal_vec.push(None);
136                hist_vec.push(None);
137            }
138        }
139    }
140    
141    let ca_macd = Float64Chunked::new("macd".into(), macd_vec);
142    let ca_signal = Float64Chunked::new("signal".into(), signal_vec);
143    let ca_hist = Float64Chunked::new("hist".into(), hist_vec);
144    
145    let series_vec = vec![ca_macd.into_series(), ca_signal.into_series(), ca_hist.into_series()];
146    let out = StructChunked::from_series("macd".into(), s_f64.len(), series_vec.iter())?;
147    
148    Ok(out.into_series())
149}
150
151#[derive(Deserialize)]
152struct BbandsKwargs {
153    timeperiod: usize,
154    nbdevup: f64,
155    nbdevdn: f64,
156    matype: u8,
157}
158
159pub fn bbands_output(_: &[Field]) -> PolarsResult<Field> {
160    Ok(Field::new(
161        "bbands".into(),
162        DataType::Struct(vec![
163            Field::new("upper".into(), DataType::Float64),
164            Field::new("middle".into(), DataType::Float64),
165            Field::new("lower".into(), DataType::Float64),
166        ]),
167    ))
168}
169
170#[polars_expr(output_type_func=bbands_output)]
171fn bbands(inputs: &[Series], kwargs: BbandsKwargs) -> PolarsResult<Series> {
172    let s = &inputs[0];
173    let s_f64 = s.f64()?;
174    
175    let ma_type = match kwargs.matype {
176        0 => MaType::Sma,
177        1 => MaType::Ema,
178        2 => MaType::Wma,
179        3 => MaType::Dema,
180        4 => MaType::Tema,
181        5 => MaType::Trima,
182        6 => MaType::Kama,
183        7 => MaType::Mama,
184        8 => MaType::T3,
185        _ => MaType::Sma,
186    };
187    
188    let mut indicator = BBANDS::new(kwargs.timeperiod, kwargs.nbdevup, kwargs.nbdevdn, ma_type);
189    
190    let mut upper_vec = Vec::with_capacity(s_f64.len());
191    let mut middle_vec = Vec::with_capacity(s_f64.len());
192    let mut lower_vec = Vec::with_capacity(s_f64.len());
193    
194    for opt_v in s_f64.into_iter() {
195        match opt_v {
196            Some(v) if !v.is_nan() => {
197                let (u, m, l) = indicator.next(v);
198                upper_vec.push(Some(u));
199                middle_vec.push(Some(m));
200                lower_vec.push(Some(l));
201            }
202            Some(_) => {
203                upper_vec.push(Some(f64::NAN));
204                middle_vec.push(Some(f64::NAN));
205                lower_vec.push(Some(f64::NAN));
206            }
207            None => {
208                upper_vec.push(None);
209                middle_vec.push(None);
210                lower_vec.push(None);
211            }
212        }
213    }
214    
215    let ca_upper = Float64Chunked::new("upper".into(), upper_vec);
216    let ca_middle = Float64Chunked::new("middle".into(), middle_vec);
217    let ca_lower = Float64Chunked::new("lower".into(), lower_vec);
218    
219    let series_vec = vec![ca_upper.into_series(), ca_middle.into_series(), ca_lower.into_series()];
220    let out = StructChunked::from_series("bbands".into(), s_f64.len(), series_vec.iter())?;
221    
222    Ok(out.into_series())
223}
224
225#[pymodule]
226#[pyo3(name = "quantwave_plugins")]
227fn quantwave_plugins(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
228    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
229    Ok(())
230}
231pub mod generated;
232pub mod custom;
233pub mod custom_0;
234pub mod custom_1;
235pub mod custom_2;
236pub mod custom_3;
237pub mod custom_4;
238pub mod custom_5;
239pub mod custom_6;
240pub mod custom_7;
241pub mod custom_8;
242pub mod custom_9;
243pub mod custom_10;
quantwave_plugins/lib.rs

quantwave_plugins/
lib.rs
