quantwave_plugins/
custom.rs

1use polars::prelude::*;
2use pyo3_polars::derive::polars_expr;
3use serde::Deserialize;
4use quantwave_core::SuperTrend;
5use quantwave_core::traits::Next;
6
7#[derive(Deserialize)]
8struct SupertrendKwargs {
9    period: usize,
10    multiplier: f64,
11}
12
13pub fn supertrend_output(_: &[Field]) -> PolarsResult<Field> {
14    Ok(Field::new(
15        "supertrend".into(),
16        DataType::Struct(vec![
17            Field::new("supertrend".into(), DataType::Float64),
18            Field::new("direction".into(), DataType::Float64),
19        ]),
20    ))
21}
22
23#[polars_expr(output_type_func=supertrend_output)]
24fn supertrend(inputs: &[Series], kwargs: SupertrendKwargs) -> PolarsResult<Series> {
25    let high = inputs[0].f64()?;
26    let low = inputs[1].f64()?;
27    let close = inputs[2].f64()?;
28    
29    let mut indicator = SuperTrend::new(kwargs.period, kwargs.multiplier);
30    
31    let mut st_vec = Vec::with_capacity(high.len());
32    let mut dir_vec = Vec::with_capacity(high.len());
33    
34    for i in 0..high.len() {
35        let h = high.get(i).unwrap_or(f64::NAN);
36        let l = low.get(i).unwrap_or(f64::NAN);
37        let c = close.get(i).unwrap_or(f64::NAN);
38        
39        if h.is_nan() || l.is_nan() || c.is_nan() {
40            st_vec.push(Some(f64::NAN));
41            dir_vec.push(Some(f64::NAN));
42        } else {
43            let (val, dir) = indicator.next((h, l, c));
44            st_vec.push(Some(val));
45            dir_vec.push(Some(dir as f64));
46        }
47    }
48    
49    let ca_st = Float64Chunked::new("supertrend".into(), st_vec);
50    let ca_dir = Float64Chunked::new("direction".into(), dir_vec);
51    
52    let series_vec = vec![ca_st.into_series(), ca_dir.into_series()];
53    let out = StructChunked::from_series("supertrend".into(), high.len(), series_vec.iter())?;
54    
55    Ok(out.into_series())
56}
quantwave_plugins/custom.rs

quantwave_plugins/
custom.rs
