quantwave_plugins/
statistics.rs

1use polars::prelude::*;
2use pyo3_polars::derive::polars_expr;
3use serde::Deserialize;
4
5use quantwave_core::indicators::incremental::statistics_ta::{
6    TaSTDDEV, TaVAR, TaLINEARREG, TaLINEARREG_SLOPE, TaLINEARREG_INTERCEPT,
7    TaLINEARREG_ANGLE, TaTSF, TaCORREL, TaBETA
8};
9use quantwave_core::traits::Next;
10
11#[derive(Deserialize)]
12struct SinglePeriodKwargs {
13    timeperiod: usize,
14}
15
16#[derive(Deserialize)]
17struct DevKwargs {
18    timeperiod: usize,
19    nbdev: f64,
20}
21
22#[polars_expr(output_type=Float64)]
23fn stddev(inputs: &[Series], kwargs: DevKwargs) -> PolarsResult<Series> {
24    let s = inputs[0].f64()?;
25    let mut indicator = TaSTDDEV::new(kwargs.timeperiod, kwargs.nbdev);
26    
27    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
28        Some(v) if !v.is_nan() => Some(indicator.next(v)),
29        Some(_) => Some(f64::NAN),
30        None => None,
31    }).collect();
32    
33    Ok(out.into_series())
34}
35
36#[polars_expr(output_type=Float64)]
37fn var(inputs: &[Series], kwargs: DevKwargs) -> PolarsResult<Series> {
38    let s = inputs[0].f64()?;
39    let mut indicator = TaVAR::new(kwargs.timeperiod, kwargs.nbdev);
40    
41    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
42        Some(v) if !v.is_nan() => Some(indicator.next(v)),
43        Some(_) => Some(f64::NAN),
44        None => None,
45    }).collect();
46    
47    Ok(out.into_series())
48}
49
50#[polars_expr(output_type=Float64)]
51fn linearreg(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
52    let s = inputs[0].f64()?;
53    let mut indicator = TaLINEARREG::new(kwargs.timeperiod);
54    
55    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
56        Some(v) if !v.is_nan() => Some(indicator.next(v)),
57        Some(_) => Some(f64::NAN),
58        None => None,
59    }).collect();
60    
61    Ok(out.into_series())
62}
63
64#[polars_expr(output_type=Float64)]
65fn linearreg_slope(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
66    let s = inputs[0].f64()?;
67    let mut indicator = TaLINEARREG_SLOPE::new(kwargs.timeperiod);
68    
69    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
70        Some(v) if !v.is_nan() => Some(indicator.next(v)),
71        Some(_) => Some(f64::NAN),
72        None => None,
73    }).collect();
74    
75    Ok(out.into_series())
76}
77
78#[polars_expr(output_type=Float64)]
79fn linearreg_intercept(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
80    let s = inputs[0].f64()?;
81    let mut indicator = TaLINEARREG_INTERCEPT::new(kwargs.timeperiod);
82    
83    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
84        Some(v) if !v.is_nan() => Some(indicator.next(v)),
85        Some(_) => Some(f64::NAN),
86        None => None,
87    }).collect();
88    
89    Ok(out.into_series())
90}
91
92#[polars_expr(output_type=Float64)]
93fn linearreg_angle(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
94    let s = inputs[0].f64()?;
95    let mut indicator = TaLINEARREG_ANGLE::new(kwargs.timeperiod);
96    
97    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
98        Some(v) if !v.is_nan() => Some(indicator.next(v)),
99        Some(_) => Some(f64::NAN),
100        None => None,
101    }).collect();
102    
103    Ok(out.into_series())
104}
105
106#[polars_expr(output_type=Float64)]
107fn tsf(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
108    let s = inputs[0].f64()?;
109    let mut indicator = TaTSF::new(kwargs.timeperiod);
110    
111    let out: Float64Chunked = s.into_iter().map(|opt_v| match opt_v {
112        Some(v) if !v.is_nan() => Some(indicator.next(v)),
113        Some(_) => Some(f64::NAN),
114        None => None,
115    }).collect();
116    
117    Ok(out.into_series())
118}
119
120#[polars_expr(output_type=Float64)]
121fn correl(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
122    let s0 = inputs[0].f64()?;
123    let s1 = inputs[1].f64()?;
124    let mut indicator = TaCORREL::new(kwargs.timeperiod);
125    
126    let out: Float64Chunked = s0.into_iter().zip(s1.into_iter()).map(|(v0, v1)| {
127        match (v0, v1) {
128            (Some(x), Some(y)) if !x.is_nan() && !y.is_nan() => Some(indicator.next((x, y))),
129            (Some(_), Some(_)) => Some(f64::NAN),
130            _ => None,
131        }
132    }).collect();
133    
134    Ok(out.into_series())
135}
136
137#[polars_expr(output_type=Float64)]
138fn beta(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
139    let s0 = inputs[0].f64()?;
140    let s1 = inputs[1].f64()?;
141    let mut indicator = TaBETA::new(kwargs.timeperiod);
142    
143    let out: Float64Chunked = s0.into_iter().zip(s1.into_iter()).map(|(v0, v1)| {
144        match (v0, v1) {
145            (Some(x), Some(y)) if !x.is_nan() && !y.is_nan() => Some(indicator.next((x, y))),
146            (Some(_), Some(_)) => Some(f64::NAN),
147            _ => None,
148        }
149    }).collect();
150    
151    Ok(out.into_series())
152}
quantwave_plugins/statistics.rs

quantwave_plugins/
statistics.rs
