quantwave_plugins/
custom_1.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 GapMomentumKwargs {
9    period: usize,
10    signal_period: usize,
11}
12
13pub fn gap_momentum_output(_: &[Field]) -> PolarsResult<Field> {
14    Ok(Field::new(
15        "gap_momentum_result".into(),
16        DataType::Struct(vec![
17            Field::new("gap_ratio".into(), DataType::Float64),
18            Field::new("gap_signal".into(), DataType::Float64),
19        ]),
20    ))
21}
22
23#[polars_expr(output_type_func=gap_momentum_output)]
24fn gap_momentum(inputs: &[Series], kwargs: GapMomentumKwargs) -> PolarsResult<Series> {
25    let s_o = &inputs[0];
26    let s_c = &inputs[1];
27    
28    let open = s_o.f64()?;
29    let close = s_c.f64()?;
30    
31    let mut indicator = quantwave_core::GapMomentum::new(kwargs.period, kwargs.signal_period);
32    
33    let mut ratio_vals = Vec::with_capacity(s_o.len());
34    let mut signal_vals = Vec::with_capacity(s_o.len());
35    
36    for i in 0..s_o.len() {
37        let o = open.get(i).unwrap_or(f64::NAN);
38        let c = close.get(i).unwrap_or(f64::NAN);
39        let (ratio, signal) = indicator.next((o, c));
40        ratio_vals.push(ratio);
41        signal_vals.push(signal);
42    }
43    
44    let s_ratio = Series::new("gap_ratio".into(), ratio_vals);
45    let s_signal = Series::new("gap_signal".into(), signal_vals);
46    
47    let series_vec = vec![s_ratio, s_signal];
48    let struct_series = StructChunked::from_series(
49        "gap_momentum_result".into(),
50        s_o.len(),
51        series_vec.iter(),
52    )?;
53    Ok(struct_series.into_series())
54}
55
56#[derive(Deserialize)]
57struct PeltKwargs {
58    penalty: f64,
59    min_dist: usize,
60}
61
62#[polars_expr(output_type=UInt32)]
63fn pelt(inputs: &[Series], kwargs: PeltKwargs) -> PolarsResult<Series> {
64    let s = &inputs[0];
65    let ca = s.f64()?;
66    
67    let data: Vec<f64> = ca.into_iter().map(|v| v.unwrap_or(f64::NAN)).collect();
68    let pelt = quantwave_core::regimes::pelt::PELT::new(kwargs.penalty, kwargs.min_dist);
69    let cps = pelt.detect(&data);
70    
71    let mut values = vec![0u32; s.len()];
72    for cp in cps {
73        if cp < values.len() {
74            values[cp] = 1;
75        }
76    }
77    
78    Ok(Series::new("changepoints".into(), values))
79}
80
81#[derive(Deserialize)]
82struct ZlemaKwargs {
83    period: usize,
84}
85
86#[polars_expr(output_type=Float64)]
87fn zlema(inputs: &[Series], kwargs: ZlemaKwargs) -> PolarsResult<Series> {
88    let s = &inputs[0];
89    let ca = s.f64()?;
90    
91    let mut zlema = quantwave_core::ZLEMA::new(kwargs.period);
92    let mut values = Vec::with_capacity(s.len());
93    
94    for i in 0..s.len() {
95        let val = ca.get(i).unwrap_or(0.0);
96        values.push(zlema.next(val));
97    }
98    
99    Ok(Series::new("zlema".into(), values))
100}
101
102#[derive(Deserialize)]
103struct RodcKwargs {
104    window_size: usize,
105    threshold: f64,
106    smooth_period: usize,
107}
108
109#[polars_expr(output_type=Float64)]
110fn rodc(inputs: &[Series], kwargs: RodcKwargs) -> PolarsResult<Series> {
111    let s = &inputs[0];
112    let ca = s.f64()?;
113    
114    let mut indicator = quantwave_core::RODC::new(kwargs.window_size, kwargs.threshold, kwargs.smooth_period);
115    let mut values = Vec::with_capacity(s.len());
116    
117    for i in 0..s.len() {
118        let val = ca.get(i).unwrap_or(f64::NAN);
119        values.push(indicator.next(val));
120    }
121    
122    Ok(Series::new("rodc".into(), values))
123}
124
125pub fn heikin_ashi_output(_: &[Field]) -> PolarsResult<Field> {
126    Ok(Field::new(
127        "heikin_ashi_output".into(),
128        DataType::Struct(vec![
129            Field::new("ha_open".into(), DataType::Float64),
130            Field::new("ha_high".into(), DataType::Float64),
131            Field::new("ha_low".into(), DataType::Float64),
132            Field::new("ha_close".into(), DataType::Float64),
133        ]),
134    ))
135}
136
137#[polars_expr(output_type_func=heikin_ashi_output)]
138fn heikin_ashi(inputs: &[Series]) -> PolarsResult<Series> {
139    let s_o = &inputs[0];
140    let s_h = &inputs[1];
141    let s_l = &inputs[2];
142    let s_c = &inputs[3];
143    
144    let open = s_o.f64()?;
145    let high = s_h.f64()?;
146    let low = s_l.f64()?;
147    let close = s_c.f64()?;
148    
149    let mut ha = quantwave_core::HeikinAshi::new();
150    let mut ha_opens = Vec::with_capacity(s_o.len());
151    let mut ha_highs = Vec::with_capacity(s_o.len());
152    let mut ha_lows = Vec::with_capacity(s_o.len());
153    let mut ha_closes = Vec::with_capacity(s_o.len());
154    
155    for i in 0..s_o.len() {
156        let o = open.get(i).unwrap_or(0.0);
157        let h = high.get(i).unwrap_or(0.0);
158        let l = low.get(i).unwrap_or(0.0);
159        let c = close.get(i).unwrap_or(0.0);
160        let (ha_o, ha_h, ha_l, ha_c) = ha.next((o, h, l, c));
161        ha_opens.push(ha_o);
162        ha_highs.push(ha_h);
163        ha_lows.push(ha_l);
164        ha_closes.push(ha_c);
165    }
166    
167    let o_series = Series::new("ha_open".into(), ha_opens);
168    let h_series = Series::new("ha_high".into(), ha_highs);
169    let l_series = Series::new("ha_low".into(), ha_lows);
170    let c_series = Series::new("ha_close".into(), ha_closes);
171    
172    let series_vec = vec![o_series, h_series, l_series, c_series];
173    let out = StructChunked::from_series(
174        "heikin_ashi_output".into(),
175        s_o.len(),
176        series_vec.iter(),
177    )?;
178    Ok(out.into_series())
179}
quantwave_plugins/custom_1.rs

quantwave_plugins/
custom_1.rs
