quantwave_plugins/
volatility.rs1use 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}