quantwave_plugins/
custom_3.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 HmaKwargs {
9    period: usize,
10}
11
12pub fn hma_output(_: &[Field]) -> PolarsResult<Field> {
13    Ok(Field::new("hma".into(), DataType::Float64))
14}
15
16#[polars_expr(output_type_func=hma_output)]
17fn hma(inputs: &[Series], kwargs: HmaKwargs) -> PolarsResult<Series> {
18    let s = inputs[0].f64()?;
19    let mut hma = HMA::new(kwargs.period);
20    let mut values = Vec::with_capacity(s.len());
21
22    for i in 0..s.len() {
23        let val = s.get(i).unwrap_or(0.0);
24        values.push(hma.next(val));
25    }
26
27    Ok(Series::new("hma".into(), values))
28}
29
30#[derive(Deserialize)]
31struct ObvmKwargs {
32    obvm_period: usize,
33    signal_period: usize,
34}
35
36pub fn obvm_output(_: &[Field]) -> PolarsResult<Field> {
37    Ok(Field::new(
38        "obvm_data".into(),
39        DataType::Struct(vec![
40            Field::new("obvm".into(), DataType::Float64),
41            Field::new("signal".into(), DataType::Float64),
42        ]),
43    ))
44}
45
46#[polars_expr(output_type_func=obvm_output)]
47fn obvm(inputs: &[Series], kwargs: ObvmKwargs) -> PolarsResult<Series> {
48    let high = inputs[0].f64()?;
49    let low = inputs[1].f64()?;
50    let close = inputs[2].f64()?;
51    let volume = inputs[3].f64()?;
52
53    let mut indicator = Obvm::new(kwargs.obvm_period, kwargs.signal_period);
54    let mut obvm_vals = Vec::with_capacity(high.len());
55    let mut signal_vals = Vec::with_capacity(high.len());
56
57    for i in 0..high.len() {
58        let h = high.get(i).unwrap_or(f64::NAN);
59        let l = low.get(i).unwrap_or(f64::NAN);
60        let c = close.get(i).unwrap_or(f64::NAN);
61        let v = volume.get(i).unwrap_or(f64::NAN);
62        let (o, sig) = indicator.next((h, l, c, v));
63        obvm_vals.push(o);
64        signal_vals.push(sig);
65    }
66
67    let s_obvm = Series::new("obvm".into(), obvm_vals);
68    let s_signal = Series::new("signal".into(), signal_vals);
69    let out = StructChunked::from_series(
70        "obvm_data".into(),
71        high.len(),
72        [s_obvm, s_signal].iter(),
73    )?;
74    Ok(out.into_series())
75}
76
77pub fn fractals_output(_: &[Field]) -> PolarsResult<Field> {
78    Ok(Field::new(
79        "fractals_data".into(),
80        DataType::Struct(vec![
81            Field::new("bearish".into(), DataType::Boolean),
82            Field::new("bullish".into(), DataType::Boolean),
83        ]),
84    ))
85}
86
87#[polars_expr(output_type_func=fractals_output)]
88fn bill_williams_fractals(inputs: &[Series]) -> PolarsResult<Series> {
89    let high = inputs[0].f64()?;
90    let low = inputs[1].f64()?;
91
92    let mut fractals = BillWilliamsFractals::new();
93    let mut bearish_vals = Vec::with_capacity(high.len());
94    let mut bullish_vals = Vec::with_capacity(high.len());
95
96    for i in 0..high.len() {
97        let h = high.get(i).unwrap_or(0.0);
98        let l = low.get(i).unwrap_or(0.0);
99        let (bear, bull) = fractals.next((h, l));
100        bearish_vals.push(bear);
101        bullish_vals.push(bull);
102    }
103
104    let bearish_series = Series::new("bearish".into(), bearish_vals);
105    let bullish_series = Series::new("bullish".into(), bullish_vals);
106
107    let out = StructChunked::from_series(
108        "fractals_data".into(),
109        high.len(),
110        [bearish_series, bullish_series].iter(),
111    )?;
112    Ok(out.into_series())
113}
114
115#[derive(Deserialize)]
116struct PeriodKwargs {
117    period: usize,
118}
119
120pub fn donchian_output(_: &[Field]) -> PolarsResult<Field> {
121    Ok(Field::new(
122        "donchian_data".into(),
123        DataType::Struct(vec![
124            Field::new("upper".into(), DataType::Float64),
125            Field::new("middle".into(), DataType::Float64),
126            Field::new("lower".into(), DataType::Float64),
127        ]),
128    ))
129}
130
131#[polars_expr(output_type_func=donchian_output)]
132fn donchian_channels(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult<Series> {
133    let high = inputs[0].f64()?;
134    let low = inputs[1].f64()?;
135
136    let mut dc = DonchianChannels::new(kwargs.period);
137    let mut uppers = Vec::with_capacity(high.len());
138    let mut middles = Vec::with_capacity(high.len());
139    let mut lowers = Vec::with_capacity(high.len());
140
141    for i in 0..high.len() {
142        let h = high.get(i).unwrap_or(0.0);
143        let l = low.get(i).unwrap_or(0.0);
144        let (upper, middle, lower) = dc.next((h, l));
145        uppers.push(upper);
146        middles.push(middle);
147        lowers.push(lower);
148    }
149
150    let upper_series = Series::new("upper".into(), uppers);
151    let middle_series = Series::new("middle".into(), middles);
152    let lower_series = Series::new("lower".into(), lowers);
153
154    let out = StructChunked::from_series(
155        "donchian_data".into(),
156        high.len(),
157        [upper_series, middle_series, lower_series].iter(),
158    )?;
159    Ok(out.into_series())
160}
161
162#[derive(Deserialize)]
163struct GmmKwargs {
164    k: usize,
165}
166
167pub fn gmm_output(_: &[Field]) -> PolarsResult<Field> {
168    Ok(Field::new("gmm_regime".into(), DataType::UInt32))
169}
170
171#[polars_expr(output_type_func=gmm_output)]
172fn gmm(inputs: &[Series], kwargs: GmmKwargs) -> PolarsResult<Series> {
173    let n_rows = inputs[0].len();
174
175    let mut values = Vec::with_capacity(n_rows);
176    for _ in 0..n_rows {
177        values.push(0u32);
178    }
179
180    Ok(Series::new("gmm_regime".into(), values))
181}
quantwave_plugins/custom_3.rs

quantwave_plugins/
custom_3.rs
