Skip to main content

quantwave_plugins/
volatility.rs

1use polars::prelude::*;
2use pyo3_polars::derive::polars_expr;
3use serde::Deserialize;
4
5use quantwave_core::indicators::incremental::ta_atr::TaATR;
6use quantwave_core::indicators::incremental::trange::{TaTRANGE, TaNATR};
7use quantwave_core::traits::Next;
8
9#[derive(Deserialize)]
10struct SinglePeriodKwargs {
11    timeperiod: usize,
12}
13
14#[polars_expr(output_type=Float64)]
15fn atr(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
16    let high = inputs[0].f64()?;
17    let low = inputs[1].f64()?;
18    let close = inputs[2].f64()?;
19    let mut indicator = TaATR::new(kwargs.timeperiod);
20    
21    let out: Float64Chunked = high.into_iter().zip(low.into_iter()).zip(close.into_iter()).map(|((h, l), c)| {
22        match (h, l, c) {
23            (Some(hv), Some(lv), Some(cv)) if !hv.is_nan() && !lv.is_nan() && !cv.is_nan() => Some(indicator.next((hv, lv, cv))),
24            (Some(_), Some(_), Some(_)) => Some(f64::NAN),
25            _ => None,
26        }
27    }).collect();
28    
29    Ok(out.into_series())
30}
31
32#[polars_expr(output_type=Float64)]
33fn natr(inputs: &[Series], kwargs: SinglePeriodKwargs) -> PolarsResult<Series> {
34    let high = inputs[0].f64()?;
35    let low = inputs[1].f64()?;
36    let close = inputs[2].f64()?;
37    let mut indicator = TaNATR::new(kwargs.timeperiod);
38    
39    let out: Float64Chunked = high.into_iter().zip(low.into_iter()).zip(close.into_iter()).map(|((h, l), c)| {
40        match (h, l, c) {
41            (Some(hv), Some(lv), Some(cv)) if !hv.is_nan() && !lv.is_nan() && !cv.is_nan() => Some(indicator.next((hv, lv, cv))),
42            (Some(_), Some(_), Some(_)) => Some(f64::NAN),
43            _ => None,
44        }
45    }).collect();
46    
47    Ok(out.into_series())
48}
49
50#[polars_expr(output_type=Float64)]
51fn trange(inputs: &[Series]) -> PolarsResult<Series> {
52    let high = inputs[0].f64()?;
53    let low = inputs[1].f64()?;
54    let close = inputs[2].f64()?;
55    let mut indicator = TaTRANGE::new();
56    
57    let out: Float64Chunked = high.into_iter().zip(low.into_iter()).zip(close.into_iter()).map(|((h, l), c)| {
58        match (h, l, c) {
59            (Some(hv), Some(lv), Some(cv)) if !hv.is_nan() && !lv.is_nan() && !cv.is_nan() => Some(indicator.next((hv, lv, cv))),
60            (Some(_), Some(_), Some(_)) => Some(f64::NAN),
61            _ => None,
62        }
63    }).collect();
64    
65    Ok(out.into_series())
66}