diff --git a/polars/README.md b/polars/README.md index 0d6f06c..ddb00e1 100644 --- a/polars/README.md +++ b/polars/README.md @@ -10,13 +10,11 @@ uv add techr ## Supported indicators -- `sma`, `wma`, `ema`, `disparity` -- `macd`, `macd_signal`, `macd_hist` -- `bband_middle`, `bband_lower`, `bband_upper` -- `stochf_percent_k`, `stochf_percent_d` -- `stoch_percent_k`, `stoch_percent_d` -- `ichimoku_base_line`, `ichimoku_conversion_line` -- `ichimoku_leading_span_a`, `ichimoku_leading_span_b`, `ichimoku_lagging_span` +- Price/trend: `sma`, `wma`, `ema`, `disparity`, `mom`, `roc`, `rsi`, `psl` +- Bands/channels: `bband_middle`, `bband_lower`, `bband_upper`, `env_upper`, `env_middle`, `env_lower`, `pchan_upper`, `pchan_middle`, `pchan_lower` +- Momentum/oscillators: `macd`, `macd_line`, `macd_signal`, `macd_hist`, `macd_histogram`, `ppo_line`, `ppo_signal`, `ppo_histogram`, `pvo_line`, `pvo_signal`, `pvo_histogram`, `sonar_line`, `sonar_signal`, `stochf_percent_k`, `stochf_percent_d`, `stoch_percent_k`, `stoch_percent_d`, `stochrsi_percent_k`, `stochrsi_percent_d` +- High/low/volume indicators: `ad`, `adx`, `adxr`, `aroon_up`, `aroon_down`, `aroonosc`, `atr`, `cci`, `cmf`, `co`, `cv`, `dmi_plus`, `dmi_minus`, `efi`, `eom_line`, `eom_signal`, `erbear`, `erbull`, `massi_line`, `massi_signal`, `mfi`, `nvi_line`, `nvi_signal`, `obv_line`, `obv_signal`, `psar`, `pvi_line`, `pvi_signal`, `ultosc`, `vr`, `willr` +- Ichimoku: `ichimoku_base_line`, `ichimoku_conversion_line`, `ichimoku_leading_span_a`, `ichimoku_leading_span_b`, `ichimoku_lagging_span` ## Usage diff --git a/polars/src/expressions.rs b/polars/src/expressions.rs index b5236cd..849f6ca 100644 --- a/polars/src/expressions.rs +++ b/polars/src/expressions.rs @@ -1,6 +1,7 @@ use polars::prelude::*; use pyo3_polars::derive::polars_expr; use serde::Deserialize; +use techr_core as core; use techr_core::{ bband_lower as techr_bband_lower, bband_middle as techr_bband_middle, bband_upper as techr_bband_upper, disparity as techr_disparity, ema as techr_ema, @@ -69,6 +70,96 @@ struct IchimokuLaggingSpanKwargs { base_line_period: u32, } +#[derive(Deserialize)] +struct DmiAdxKwargs { + dmi_period: u32, + adx_period: u32, +} + +#[derive(Deserialize)] +struct AdxrKwargs { + dmi_period: u32, + adx_period: u32, + adxr_period: u32, +} + +#[derive(Deserialize)] +struct EnvKwargs { + period: u32, + shift_percentage: f64, +} + +#[derive(Deserialize)] +struct EomLineKwargs { + period: u32, + scale: f64, +} + +#[derive(Deserialize)] +struct EomSignalKwargs { + period: u32, + signal_period: u32, + scale: f64, +} + +#[derive(Deserialize)] +struct MassiLineKwargs { + period_ema: u32, + period_sum: u32, +} + +#[derive(Deserialize)] +struct MassiSignalKwargs { + period_ema: u32, + period_sum: u32, + period_signal: u32, +} + +#[derive(Deserialize)] +struct SignalKwargs { + signal_period: u32, +} + +#[derive(Deserialize)] +struct PeriodShortLongKwargs { + period_short: u32, + period_long: u32, +} + +#[derive(Deserialize)] +struct PsarKwargs { + increment: f64, + initial_acceleration_factor: f64, + max_acceleration_factor: f64, +} + +#[derive(Deserialize)] +struct SonarLineKwargs { + period: u32, + step: u32, +} + +#[derive(Deserialize)] +struct SonarSignalKwargs { + period: u32, + step: u32, + signal_period: u32, +} + +#[derive(Deserialize)] +struct StochRsiKwargs { + period_rsi: u32, + period_k: u32, + period_d: u32, +} + +#[derive(Deserialize)] +struct UltOscKwargs { + period_short: u32, + period_medium: u32, + period_long: u32, +} + fn series_to_f64_vec(series: &Series) -> PolarsResult>> { let casted = series.cast(&DataType::Float64)?; Ok(casted.f64()?.to_vec()) @@ -121,6 +212,16 @@ fn macd(inputs: &[Series], kwargs: FastSlowKwargs) -> PolarsResult { Ok(option_vec_to_series(macd_line)) } +#[polars_expr(output_type=Float64)] +fn macd_line(inputs: &[Series], kwargs: FastSlowKwargs) -> PolarsResult { + let input = series_to_f64_vec(&inputs[0])?; + Ok(option_vec_to_series(techr_macd_line( + &input, + kwargs.fast_period as usize, + kwargs.slow_period as usize, + ))) +} + #[polars_expr(output_type=Float64)] fn macd_signal(inputs: &[Series], kwargs: FastSlowSignalKwargs) -> PolarsResult { let input = series_to_f64_vec(&inputs[0])?; @@ -145,6 +246,17 @@ fn macd_hist(inputs: &[Series], kwargs: FastSlowSignalKwargs) -> PolarsResult PolarsResult { + let input = series_to_f64_vec(&inputs[0])?; + Ok(option_vec_to_series(techr_macd_histogram( + &input, + kwargs.fast_period as usize, + kwargs.slow_period as usize, + kwargs.signal_period as usize, + ))) +} + #[polars_expr(output_type=Float64)] fn bband_middle(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { let input = series_to_f64_vec(&inputs[0])?; @@ -291,3 +403,556 @@ fn ichimoku_lagging_span( kwargs.base_line_period as usize, ))) } + +#[polars_expr(output_type=Float64)] +fn ad(inputs: &[Series]) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let closes = series_to_f64_vec(&inputs[2])?; + let volumes = series_to_f64_vec(&inputs[3])?; + Ok(option_vec_to_series(core::ad( + &highs, &lows, &closes, &volumes, + ))) +} + +#[polars_expr(output_type=Float64)] +fn adx(inputs: &[Series], kwargs: DmiAdxKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let closes = series_to_f64_vec(&inputs[2])?; + Ok(option_vec_to_series(core::adx( + &highs, + &lows, + &closes, + kwargs.dmi_period as usize, + kwargs.adx_period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn adxr(inputs: &[Series], kwargs: AdxrKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let closes = series_to_f64_vec(&inputs[2])?; + Ok(option_vec_to_series(core::adxr( + &highs, + &lows, + &closes, + kwargs.dmi_period as usize, + kwargs.adx_period as usize, + kwargs.adxr_period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn aroon_up(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let (up, _) = core::aroon(&highs, &lows, kwargs.period as usize); + Ok(option_vec_to_series(up)) +} + +#[polars_expr(output_type=Float64)] +fn aroon_down(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let (_, down) = core::aroon(&highs, &lows, kwargs.period as usize); + Ok(option_vec_to_series(down)) +} + +#[polars_expr(output_type=Float64)] +fn aroonosc(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + Ok(option_vec_to_series(core::aroonosc( + &highs, + &lows, + kwargs.period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn atr(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let closes = series_to_f64_vec(&inputs[2])?; + Ok(option_vec_to_series(core::atr( + &highs, + &lows, + &closes, + kwargs.period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn cci(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let closes = series_to_f64_vec(&inputs[2])?; + Ok(option_vec_to_series(core::cci( + &highs, + &lows, + &closes, + kwargs.period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn cmf(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let closes = series_to_f64_vec(&inputs[2])?; + let volumes = series_to_f64_vec(&inputs[3])?; + Ok(option_vec_to_series(core::cmf( + &highs, + &lows, + &closes, + &volumes, + kwargs.period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn co(inputs: &[Series], kwargs: PeriodShortLongKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let closes = series_to_f64_vec(&inputs[2])?; + let volumes = series_to_f64_vec(&inputs[3])?; + Ok(option_vec_to_series(core::co( + &highs, + &lows, + &closes, + &volumes, + kwargs.period_short as usize, + kwargs.period_long as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn cv(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + Ok(option_vec_to_series(core::cv( + &highs, + &lows, + kwargs.period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn dmi_plus(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let closes = series_to_f64_vec(&inputs[2])?; + let (plus, _) = core::dmi(&highs, &lows, &closes, kwargs.period as usize); + Ok(option_vec_to_series(plus)) +} + +#[polars_expr(output_type=Float64)] +fn dmi_minus(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let closes = series_to_f64_vec(&inputs[2])?; + let (_, minus) = core::dmi(&highs, &lows, &closes, kwargs.period as usize); + Ok(option_vec_to_series(minus)) +} + +#[polars_expr(output_type=Float64)] +fn efi(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let closes = series_to_f64_vec(&inputs[0])?; + let volumes = series_to_f64_vec(&inputs[1])?; + Ok(option_vec_to_series(core::efi( + &closes, + &volumes, + kwargs.period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn env_upper(inputs: &[Series], kwargs: EnvKwargs) -> PolarsResult { + let input = series_to_f64_vec(&inputs[0])?; + let (upper, _, _) = core::env(&input, kwargs.period as usize, kwargs.shift_percentage); + Ok(option_vec_to_series(upper)) +} + +#[polars_expr(output_type=Float64)] +fn env_middle(inputs: &[Series], kwargs: EnvKwargs) -> PolarsResult { + let input = series_to_f64_vec(&inputs[0])?; + let (_, middle, _) = core::env(&input, kwargs.period as usize, kwargs.shift_percentage); + Ok(option_vec_to_series(middle)) +} + +#[polars_expr(output_type=Float64)] +fn env_lower(inputs: &[Series], kwargs: EnvKwargs) -> PolarsResult { + let input = series_to_f64_vec(&inputs[0])?; + let (_, _, lower) = core::env(&input, kwargs.period as usize, kwargs.shift_percentage); + Ok(option_vec_to_series(lower)) +} + +#[polars_expr(output_type=Float64)] +fn eom_line(inputs: &[Series], kwargs: EomLineKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let volumes = series_to_f64_vec(&inputs[2])?; + Ok(option_vec_to_series(core::eom_line( + &highs, + &lows, + &volumes, + kwargs.period as usize, + kwargs.scale, + ))) +} + +#[polars_expr(output_type=Float64)] +fn eom_signal(inputs: &[Series], kwargs: EomSignalKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let volumes = series_to_f64_vec(&inputs[2])?; + Ok(option_vec_to_series(core::eom_signal( + &highs, + &lows, + &volumes, + kwargs.period as usize, + kwargs.signal_period as usize, + kwargs.scale, + ))) +} + +#[polars_expr(output_type=Float64)] +fn erbear(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let lows = series_to_f64_vec(&inputs[0])?; + let closes = series_to_f64_vec(&inputs[1])?; + Ok(option_vec_to_series(core::erbear( + &lows, + &closes, + kwargs.period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn erbull(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let closes = series_to_f64_vec(&inputs[1])?; + Ok(option_vec_to_series(core::erbull( + &highs, + &closes, + kwargs.period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn massi_line(inputs: &[Series], kwargs: MassiLineKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + Ok(option_vec_to_series(core::massi_line( + &highs, + &lows, + kwargs.period_ema as usize, + kwargs.period_sum as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn massi_signal(inputs: &[Series], kwargs: MassiSignalKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + Ok(option_vec_to_series(core::massi_signal( + &highs, + &lows, + kwargs.period_ema as usize, + kwargs.period_sum as usize, + kwargs.period_signal as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn mfi(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let closes = series_to_f64_vec(&inputs[2])?; + let volumes = series_to_f64_vec(&inputs[3])?; + Ok(option_vec_to_series(core::mfi( + &highs, + &lows, + &closes, + &volumes, + kwargs.period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn mom(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let closes = series_to_f64_vec(&inputs[0])?; + Ok(option_vec_to_series(core::mom( + &closes, + kwargs.period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn nvi_line(inputs: &[Series]) -> PolarsResult { + let closes = series_to_f64_vec(&inputs[0])?; + let volumes = series_to_f64_vec(&inputs[1])?; + Ok(option_vec_to_series(core::nvi_line(&closes, &volumes))) +} + +#[polars_expr(output_type=Float64)] +fn nvi_signal(inputs: &[Series], kwargs: SignalKwargs) -> PolarsResult { + let closes = series_to_f64_vec(&inputs[0])?; + let volumes = series_to_f64_vec(&inputs[1])?; + Ok(option_vec_to_series(core::nvi_signal( + &closes, + &volumes, + kwargs.signal_period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn obv_line(inputs: &[Series]) -> PolarsResult { + let closes = series_to_f64_vec(&inputs[0])?; + let volumes = series_to_f64_vec(&inputs[1])?; + Ok(option_vec_to_series(core::obv_line(&closes, &volumes))) +} + +#[polars_expr(output_type=Float64)] +fn obv_signal(inputs: &[Series], kwargs: SignalKwargs) -> PolarsResult { + let closes = series_to_f64_vec(&inputs[0])?; + let volumes = series_to_f64_vec(&inputs[1])?; + Ok(option_vec_to_series(core::obv_signal( + &closes, + &volumes, + kwargs.signal_period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn pchan_upper(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let (upper, _, _) = core::pchan(&highs, &lows, kwargs.period as usize); + Ok(option_vec_to_series(upper)) +} + +#[polars_expr(output_type=Float64)] +fn pchan_middle(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let (_, middle, _) = core::pchan(&highs, &lows, kwargs.period as usize); + Ok(option_vec_to_series(middle)) +} + +#[polars_expr(output_type=Float64)] +fn pchan_lower(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let (_, _, lower) = core::pchan(&highs, &lows, kwargs.period as usize); + Ok(option_vec_to_series(lower)) +} + +#[polars_expr(output_type=Float64)] +fn ppo_line(inputs: &[Series], kwargs: FastSlowKwargs) -> PolarsResult { + let input = series_to_f64_vec(&inputs[0])?; + Ok(option_vec_to_series(core::ppo_line( + &input, + kwargs.fast_period as usize, + kwargs.slow_period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn ppo_signal(inputs: &[Series], kwargs: FastSlowSignalKwargs) -> PolarsResult { + let input = series_to_f64_vec(&inputs[0])?; + Ok(option_vec_to_series(core::ppo_signal( + &input, + kwargs.fast_period as usize, + kwargs.slow_period as usize, + kwargs.signal_period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn ppo_histogram(inputs: &[Series], kwargs: FastSlowSignalKwargs) -> PolarsResult { + let input = series_to_f64_vec(&inputs[0])?; + Ok(option_vec_to_series(core::ppo_histogram( + &input, + kwargs.fast_period as usize, + kwargs.slow_period as usize, + kwargs.signal_period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn psar(inputs: &[Series], kwargs: PsarKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let closes = series_to_f64_vec(&inputs[2])?; + Ok(option_vec_to_series(core::psar( + &highs, + &lows, + &closes, + kwargs.increment, + kwargs.initial_acceleration_factor, + kwargs.max_acceleration_factor, + ))) +} + +#[polars_expr(output_type=Float64)] +fn psl(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let closes = series_to_f64_vec(&inputs[0])?; + Ok(option_vec_to_series(core::psl( + &closes, + kwargs.period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn pvi_line(inputs: &[Series]) -> PolarsResult { + let closes = series_to_f64_vec(&inputs[0])?; + let volumes = series_to_f64_vec(&inputs[1])?; + Ok(option_vec_to_series(core::pvi_line(&closes, &volumes))) +} + +#[polars_expr(output_type=Float64)] +fn pvi_signal(inputs: &[Series], kwargs: SignalKwargs) -> PolarsResult { + let closes = series_to_f64_vec(&inputs[0])?; + let volumes = series_to_f64_vec(&inputs[1])?; + Ok(option_vec_to_series(core::pvi_signal( + &closes, + &volumes, + kwargs.signal_period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn pvo_line(inputs: &[Series], kwargs: FastSlowKwargs) -> PolarsResult { + let input = series_to_f64_vec(&inputs[0])?; + Ok(option_vec_to_series(core::pvo_line( + &input, + kwargs.fast_period as usize, + kwargs.slow_period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn pvo_signal(inputs: &[Series], kwargs: FastSlowSignalKwargs) -> PolarsResult { + let input = series_to_f64_vec(&inputs[0])?; + Ok(option_vec_to_series(core::pvo_signal( + &input, + kwargs.fast_period as usize, + kwargs.slow_period as usize, + kwargs.signal_period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn pvo_histogram(inputs: &[Series], kwargs: FastSlowSignalKwargs) -> PolarsResult { + let input = series_to_f64_vec(&inputs[0])?; + Ok(option_vec_to_series(core::pvo_histogram( + &input, + kwargs.fast_period as usize, + kwargs.slow_period as usize, + kwargs.signal_period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn roc(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let closes = series_to_f64_vec(&inputs[0])?; + Ok(option_vec_to_series(core::roc( + &closes, + kwargs.period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn rsi(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let input = series_to_f64_vec(&inputs[0])?; + Ok(option_vec_to_series(core::rsi( + &input, + kwargs.period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn sonar_line(inputs: &[Series], kwargs: SonarLineKwargs) -> PolarsResult { + let input = series_to_f64_vec(&inputs[0])?; + Ok(option_vec_to_series(core::sonar_line( + &input, + kwargs.period as usize, + kwargs.step as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn sonar_signal(inputs: &[Series], kwargs: SonarSignalKwargs) -> PolarsResult { + let input = series_to_f64_vec(&inputs[0])?; + Ok(option_vec_to_series(core::sonar_signal( + &input, + kwargs.period as usize, + kwargs.step as usize, + kwargs.signal_period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn stochrsi_percent_k(inputs: &[Series], kwargs: StochRsiKwargs) -> PolarsResult { + let closes = series_to_f64_vec(&inputs[0])?; + let (percent_k, _) = core::stochrsi( + &closes, + kwargs.period_rsi as usize, + kwargs.period_k as usize, + kwargs.period_d as usize, + ); + Ok(option_vec_to_series(percent_k)) +} + +#[polars_expr(output_type=Float64)] +fn stochrsi_percent_d(inputs: &[Series], kwargs: StochRsiKwargs) -> PolarsResult { + let closes = series_to_f64_vec(&inputs[0])?; + let (_, percent_d) = core::stochrsi( + &closes, + kwargs.period_rsi as usize, + kwargs.period_k as usize, + kwargs.period_d as usize, + ); + Ok(option_vec_to_series(percent_d)) +} + +#[polars_expr(output_type=Float64)] +fn ultosc(inputs: &[Series], kwargs: UltOscKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let closes = series_to_f64_vec(&inputs[2])?; + Ok(option_vec_to_series(core::ultosc( + &highs, + &lows, + &closes, + kwargs.period_short as usize, + kwargs.period_medium as usize, + kwargs.period_long as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn vr(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let closes = series_to_f64_vec(&inputs[0])?; + let volumes = series_to_f64_vec(&inputs[1])?; + Ok(option_vec_to_series(core::vr( + &closes, + &volumes, + kwargs.period as usize, + ))) +} + +#[polars_expr(output_type=Float64)] +fn willr(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult { + let highs = series_to_f64_vec(&inputs[0])?; + let lows = series_to_f64_vec(&inputs[1])?; + let closes = series_to_f64_vec(&inputs[2])?; + Ok(option_vec_to_series(core::willr( + &highs, + &lows, + &closes, + kwargs.period as usize, + ))) +} diff --git a/polars/techr/__init__.py b/polars/techr/__init__.py index 5113d5d..1153e66 100644 --- a/polars/techr/__init__.py +++ b/polars/techr/__init__.py @@ -9,11 +9,32 @@ LIB = Path(__file__).resolve().parent __all__ = [ + "ad", + "adx", + "adxr", + "aroon_down", + "aroon_up", + "aroonosc", + "atr", "bband_lower", "bband_middle", "bband_upper", + "cci", + "cmf", + "co", + "cv", "disparity", + "dmi_minus", + "dmi_plus", + "efi", "ema", + "env_lower", + "env_middle", + "env_upper", + "eom_line", + "eom_signal", + "erbear", + "erbull", "ichimoku_base_line", "ichimoku_conversion_line", "ichimoku_lagging_span", @@ -21,12 +42,44 @@ "ichimoku_leading_span_b", "macd", "macd_hist", + "macd_histogram", + "macd_line", "macd_signal", + "massi_line", + "massi_signal", + "mfi", + "mom", + "nvi_line", + "nvi_signal", + "obv_line", + "obv_signal", + "pchan_lower", + "pchan_middle", + "pchan_upper", + "ppo_histogram", + "ppo_line", + "ppo_signal", + "psar", + "psl", + "pvi_line", + "pvi_signal", + "pvo_histogram", + "pvo_line", + "pvo_signal", + "roc", + "rsi", "sma", + "sonar_line", + "sonar_signal", "stoch_percent_d", "stoch_percent_k", "stochf_percent_d", "stochf_percent_k", + "stochrsi_percent_d", + "stochrsi_percent_k", + "ultosc", + "vr", + "willr", "wma", ] @@ -40,7 +93,7 @@ def _register( plugin_path=LIB, function_name=function_name, args=args, - kwargs=kwargs, + kwargs=kwargs or None, is_elementwise=False, ) @@ -69,6 +122,14 @@ def macd(expr: IntoExpr, *, fast_period: int, slow_period: int) -> pl.Expr: ) +def macd_line(expr: IntoExpr, *, fast_period: int, slow_period: int) -> pl.Expr: + return _register( + "macd_line", + [expr], + {"fast_period": fast_period, "slow_period": slow_period}, + ) + + def macd_signal( expr: IntoExpr, *, @@ -105,6 +166,24 @@ def macd_hist( ) +def macd_histogram( + expr: IntoExpr, + *, + fast_period: int, + slow_period: int, + signal_period: int, +) -> pl.Expr: + return _register( + "macd_histogram", + [expr], + { + "fast_period": fast_period, + "slow_period": slow_period, + "signal_period": signal_period, + }, + ) + + def bband_middle(expr: IntoExpr, *, period: int) -> pl.Expr: return _register("bband_middle", [expr], {"period": period}) @@ -237,3 +316,463 @@ def ichimoku_lagging_span(close: IntoExpr, *, base_line_period: int) -> pl.Expr: [close], {"base_line_period": base_line_period}, ) + + +def ad(high: IntoExpr, low: IntoExpr, close: IntoExpr, volume: IntoExpr) -> pl.Expr: + return _register("ad", [high, low, close, volume], {}) + + +def adx( + high: IntoExpr, + low: IntoExpr, + close: IntoExpr, + *, + dmi_period: int, + adx_period: int, +) -> pl.Expr: + return _register( + "adx", + [high, low, close], + {"dmi_period": dmi_period, "adx_period": adx_period}, + ) + + +def adxr( + high: IntoExpr, + low: IntoExpr, + close: IntoExpr, + *, + dmi_period: int, + adx_period: int, + adxr_period: int, +) -> pl.Expr: + return _register( + "adxr", + [high, low, close], + { + "dmi_period": dmi_period, + "adx_period": adx_period, + "adxr_period": adxr_period, + }, + ) + + +def aroon_up(high: IntoExpr, low: IntoExpr, *, period: int) -> pl.Expr: + return _register("aroon_up", [high, low], {"period": period}) + + +def aroon_down(high: IntoExpr, low: IntoExpr, *, period: int) -> pl.Expr: + return _register("aroon_down", [high, low], {"period": period}) + + +def aroonosc(high: IntoExpr, low: IntoExpr, *, period: int) -> pl.Expr: + return _register("aroonosc", [high, low], {"period": period}) + + +def atr(high: IntoExpr, low: IntoExpr, close: IntoExpr, *, period: int) -> pl.Expr: + return _register("atr", [high, low, close], {"period": period}) + + +def cci(high: IntoExpr, low: IntoExpr, close: IntoExpr, *, period: int) -> pl.Expr: + return _register("cci", [high, low, close], {"period": period}) + + +def cmf( + high: IntoExpr, + low: IntoExpr, + close: IntoExpr, + volume: IntoExpr, + *, + period: int, +) -> pl.Expr: + return _register("cmf", [high, low, close, volume], {"period": period}) + + +def co( + high: IntoExpr, + low: IntoExpr, + close: IntoExpr, + volume: IntoExpr, + *, + period_short: int, + period_long: int, +) -> pl.Expr: + return _register( + "co", + [high, low, close, volume], + {"period_short": period_short, "period_long": period_long}, + ) + + +def cv(high: IntoExpr, low: IntoExpr, *, period: int) -> pl.Expr: + return _register("cv", [high, low], {"period": period}) + + +def dmi_plus(high: IntoExpr, low: IntoExpr, close: IntoExpr, *, period: int) -> pl.Expr: + return _register("dmi_plus", [high, low, close], {"period": period}) + + +def dmi_minus( + high: IntoExpr, low: IntoExpr, close: IntoExpr, *, period: int +) -> pl.Expr: + return _register("dmi_minus", [high, low, close], {"period": period}) + + +def efi(close: IntoExpr, volume: IntoExpr, *, period: int) -> pl.Expr: + return _register("efi", [close, volume], {"period": period}) + + +def env_upper( + expr: IntoExpr, + *, + period: int, + shift_percentage: float, +) -> pl.Expr: + return _register( + "env_upper", + [expr], + {"period": period, "shift_percentage": shift_percentage}, + ) + + +def env_middle( + expr: IntoExpr, + *, + period: int, + shift_percentage: float, +) -> pl.Expr: + return _register( + "env_middle", + [expr], + {"period": period, "shift_percentage": shift_percentage}, + ) + + +def env_lower( + expr: IntoExpr, + *, + period: int, + shift_percentage: float, +) -> pl.Expr: + return _register( + "env_lower", + [expr], + {"period": period, "shift_percentage": shift_percentage}, + ) + + +def eom_line( + high: IntoExpr, + low: IntoExpr, + volume: IntoExpr, + *, + period: int, + scale: float, +) -> pl.Expr: + return _register( + "eom_line", + [high, low, volume], + {"period": period, "scale": scale}, + ) + + +def eom_signal( + high: IntoExpr, + low: IntoExpr, + volume: IntoExpr, + *, + period: int, + signal_period: int, + scale: float, +) -> pl.Expr: + return _register( + "eom_signal", + [high, low, volume], + {"period": period, "signal_period": signal_period, "scale": scale}, + ) + + +def erbear(low: IntoExpr, close: IntoExpr, *, period: int) -> pl.Expr: + return _register("erbear", [low, close], {"period": period}) + + +def erbull(high: IntoExpr, close: IntoExpr, *, period: int) -> pl.Expr: + return _register("erbull", [high, close], {"period": period}) + + +def massi_line( + high: IntoExpr, + low: IntoExpr, + *, + period_ema: int, + period_sum: int, +) -> pl.Expr: + return _register( + "massi_line", + [high, low], + {"period_ema": period_ema, "period_sum": period_sum}, + ) + + +def massi_signal( + high: IntoExpr, + low: IntoExpr, + *, + period_ema: int, + period_sum: int, + period_signal: int, +) -> pl.Expr: + return _register( + "massi_signal", + [high, low], + { + "period_ema": period_ema, + "period_sum": period_sum, + "period_signal": period_signal, + }, + ) + + +def mfi( + high: IntoExpr, + low: IntoExpr, + close: IntoExpr, + volume: IntoExpr, + *, + period: int, +) -> pl.Expr: + return _register("mfi", [high, low, close, volume], {"period": period}) + + +def mom(close: IntoExpr, *, period: int) -> pl.Expr: + return _register("mom", [close], {"period": period}) + + +def nvi_line(close: IntoExpr, volume: IntoExpr) -> pl.Expr: + return _register("nvi_line", [close, volume], {}) + + +def nvi_signal(close: IntoExpr, volume: IntoExpr, *, signal_period: int) -> pl.Expr: + return _register("nvi_signal", [close, volume], {"signal_period": signal_period}) + + +def obv_line(close: IntoExpr, volume: IntoExpr) -> pl.Expr: + return _register("obv_line", [close, volume], {}) + + +def obv_signal(close: IntoExpr, volume: IntoExpr, *, signal_period: int) -> pl.Expr: + return _register("obv_signal", [close, volume], {"signal_period": signal_period}) + + +def pchan_upper(high: IntoExpr, low: IntoExpr, *, period: int) -> pl.Expr: + return _register("pchan_upper", [high, low], {"period": period}) + + +def pchan_middle(high: IntoExpr, low: IntoExpr, *, period: int) -> pl.Expr: + return _register("pchan_middle", [high, low], {"period": period}) + + +def pchan_lower(high: IntoExpr, low: IntoExpr, *, period: int) -> pl.Expr: + return _register("pchan_lower", [high, low], {"period": period}) + + +def ppo_line(expr: IntoExpr, *, fast_period: int, slow_period: int) -> pl.Expr: + return _register( + "ppo_line", + [expr], + {"fast_period": fast_period, "slow_period": slow_period}, + ) + + +def ppo_signal( + expr: IntoExpr, + *, + fast_period: int, + slow_period: int, + signal_period: int, +) -> pl.Expr: + return _register( + "ppo_signal", + [expr], + { + "fast_period": fast_period, + "slow_period": slow_period, + "signal_period": signal_period, + }, + ) + + +def ppo_histogram( + expr: IntoExpr, + *, + fast_period: int, + slow_period: int, + signal_period: int, +) -> pl.Expr: + return _register( + "ppo_histogram", + [expr], + { + "fast_period": fast_period, + "slow_period": slow_period, + "signal_period": signal_period, + }, + ) + + +def psar( + high: IntoExpr, + low: IntoExpr, + close: IntoExpr, + *, + increment: float, + initial_acceleration_factor: float, + max_acceleration_factor: float, +) -> pl.Expr: + return _register( + "psar", + [high, low, close], + { + "increment": increment, + "initial_acceleration_factor": initial_acceleration_factor, + "max_acceleration_factor": max_acceleration_factor, + }, + ) + + +def psl(close: IntoExpr, *, period: int) -> pl.Expr: + return _register("psl", [close], {"period": period}) + + +def pvi_line(close: IntoExpr, volume: IntoExpr) -> pl.Expr: + return _register("pvi_line", [close, volume], {}) + + +def pvi_signal(close: IntoExpr, volume: IntoExpr, *, signal_period: int) -> pl.Expr: + return _register("pvi_signal", [close, volume], {"signal_period": signal_period}) + + +def pvo_line(expr: IntoExpr, *, fast_period: int, slow_period: int) -> pl.Expr: + return _register( + "pvo_line", + [expr], + {"fast_period": fast_period, "slow_period": slow_period}, + ) + + +def pvo_signal( + expr: IntoExpr, + *, + fast_period: int, + slow_period: int, + signal_period: int, +) -> pl.Expr: + return _register( + "pvo_signal", + [expr], + { + "fast_period": fast_period, + "slow_period": slow_period, + "signal_period": signal_period, + }, + ) + + +def pvo_histogram( + expr: IntoExpr, + *, + fast_period: int, + slow_period: int, + signal_period: int, +) -> pl.Expr: + return _register( + "pvo_histogram", + [expr], + { + "fast_period": fast_period, + "slow_period": slow_period, + "signal_period": signal_period, + }, + ) + + +def roc(close: IntoExpr, *, period: int) -> pl.Expr: + return _register("roc", [close], {"period": period}) + + +def rsi(expr: IntoExpr, *, period: int) -> pl.Expr: + return _register("rsi", [expr], {"period": period}) + + +def sonar_line(expr: IntoExpr, *, period: int, step: int) -> pl.Expr: + return _register("sonar_line", [expr], {"period": period, "step": step}) + + +def sonar_signal( + expr: IntoExpr, + *, + period: int, + step: int, + signal_period: int, +) -> pl.Expr: + return _register( + "sonar_signal", + [expr], + {"period": period, "step": step, "signal_period": signal_period}, + ) + + +def stochrsi_percent_k( + close: IntoExpr, + *, + period_rsi: int, + period_k: int, + period_d: int, +) -> pl.Expr: + return _register( + "stochrsi_percent_k", + [close], + {"period_rsi": period_rsi, "period_k": period_k, "period_d": period_d}, + ) + + +def stochrsi_percent_d( + close: IntoExpr, + *, + period_rsi: int, + period_k: int, + period_d: int, +) -> pl.Expr: + return _register( + "stochrsi_percent_d", + [close], + {"period_rsi": period_rsi, "period_k": period_k, "period_d": period_d}, + ) + + +def ultosc( + high: IntoExpr, + low: IntoExpr, + close: IntoExpr, + *, + period_short: int, + period_medium: int, + period_long: int, +) -> pl.Expr: + return _register( + "ultosc", + [high, low, close], + { + "period_short": period_short, + "period_medium": period_medium, + "period_long": period_long, + }, + ) + + +def vr(close: IntoExpr, volume: IntoExpr, *, period: int) -> pl.Expr: + return _register("vr", [close, volume], {"period": period}) + + +def willr(high: IntoExpr, low: IntoExpr, close: IntoExpr, *, period: int) -> pl.Expr: + return _register("willr", [high, low, close], {"period": period}) diff --git a/polars/tests/test_indicators.py b/polars/tests/test_indicators.py index d71659a..8901f4c 100644 --- a/polars/tests/test_indicators.py +++ b/polars/tests/test_indicators.py @@ -59,18 +59,169 @@ def select_expr(df: pl.DataFrame, expr: pl.Expr, alias: str, lazy: bool) -> pl.S SeriesExprBuilder = Callable[[], pl.Expr] +CASE_TOLERANCES = { + "ad": 1e-4, + "co": 1e-4, + "stochrsi_percent_d": 1e-4, + "stochrsi_percent_k": 1e-4, +} # Indicators whose expected fixture length matches the input row count. CORE_EXPECTED_CASES: list[tuple[str, SeriesExprBuilder, str]] = [ + ( + "ad", + lambda: ta.ad(pl.col("high"), pl.col("low"), pl.col("close"), pl.col("volume")), + "ad", + ), + ( + "adx", + lambda: ta.adx( + pl.col("high"), + pl.col("low"), + pl.col("close"), + dmi_period=14, + adx_period=14, + ), + "adx", + ), + ( + "adxr", + lambda: ta.adxr( + pl.col("high"), + pl.col("low"), + pl.col("close"), + dmi_period=14, + adx_period=14, + adxr_period=14, + ), + "adxr", + ), + ( + "aroon_up", + lambda: ta.aroon_up(pl.col("high"), pl.col("low"), period=25), + "aroon_up", + ), + ( + "aroon_down", + lambda: ta.aroon_down(pl.col("high"), pl.col("low"), period=25), + "aroon_down", + ), + ( + "aroonosc", + lambda: ta.aroonosc(pl.col("high"), pl.col("low"), period=25), + "aroonosc", + ), + ( + "atr", + lambda: ta.atr(pl.col("high"), pl.col("low"), pl.col("close"), period=20), + "atr", + ), ("sma", lambda: ta.sma(pl.col("close"), period=20), "sma"), ("wma", lambda: ta.wma(pl.col("close"), period=20), "wma"), ("ema", lambda: ta.ema(pl.col("close"), period=20), "ema"), ("disparity", lambda: ta.disparity(pl.col("close"), period=20), "disparity"), + ( + "cci", + lambda: ta.cci(pl.col("high"), pl.col("low"), pl.col("close"), period=20), + "cci", + ), + ( + "cmf", + lambda: ta.cmf( + pl.col("high"), + pl.col("low"), + pl.col("close"), + pl.col("volume"), + period=21, + ), + "cmf", + ), + ( + "co", + lambda: ta.co( + pl.col("high"), + pl.col("low"), + pl.col("close"), + pl.col("volume"), + period_short=3, + period_long=10, + ), + "co", + ), + ("cv", lambda: ta.cv(pl.col("high"), pl.col("low"), period=10), "cv"), + ( + "dmi_plus", + lambda: ta.dmi_plus(pl.col("high"), pl.col("low"), pl.col("close"), period=14), + "dmi_plus", + ), + ( + "dmi_minus", + lambda: ta.dmi_minus(pl.col("high"), pl.col("low"), pl.col("close"), period=14), + "dmi_minus", + ), + ( + "efi", + lambda: ta.efi(pl.col("close"), pl.col("volume"), period=14), + "efi", + ), + ( + "env_upper", + lambda: ta.env_upper(pl.col("close"), period=20, shift_percentage=10.0), + "env_upper", + ), + ( + "env_middle", + lambda: ta.env_middle(pl.col("close"), period=20, shift_percentage=10.0), + "sma", + ), + ( + "env_lower", + lambda: ta.env_lower(pl.col("close"), period=20, shift_percentage=10.0), + "env_lower", + ), + ( + "eom_line", + lambda: ta.eom_line( + pl.col("high"), + pl.col("low"), + pl.col("volume"), + period=14, + scale=10000.0, + ), + "eom_line", + ), + ( + "eom_signal", + lambda: ta.eom_signal( + pl.col("high"), + pl.col("low"), + pl.col("volume"), + period=14, + signal_period=3, + scale=10000.0, + ), + "eom_signal", + ), + ( + "erbear", + lambda: ta.erbear(pl.col("low"), pl.col("close"), period=13), + "erbear", + ), + ( + "erbull", + lambda: ta.erbull(pl.col("high"), pl.col("close"), period=13), + "erbull", + ), ( "macd", lambda: ta.macd(pl.col("close"), fast_period=12, slow_period=26), "macd_line", ), + ( + "macd_line", + lambda: ta.macd_line(pl.col("close"), fast_period=12, slow_period=26), + "macd_line", + ), ( "macd_signal", lambda: ta.macd_signal( @@ -91,6 +242,156 @@ def select_expr(df: pl.DataFrame, expr: pl.Expr, alias: str, lazy: bool) -> pl.S ), "macd_histogram", ), + ( + "macd_histogram", + lambda: ta.macd_histogram( + pl.col("close"), + fast_period=12, + slow_period=26, + signal_period=9, + ), + "macd_histogram", + ), + ( + "massi_line", + lambda: ta.massi_line( + pl.col("high"), pl.col("low"), period_ema=9, period_sum=25 + ), + "massi_line", + ), + ( + "massi_signal", + lambda: ta.massi_signal( + pl.col("high"), + pl.col("low"), + period_ema=9, + period_sum=25, + period_signal=9, + ), + "massi_signal", + ), + ( + "mfi", + lambda: ta.mfi( + pl.col("high"), + pl.col("low"), + pl.col("close"), + pl.col("volume"), + period=14, + ), + "mfi", + ), + ("mom", lambda: ta.mom(pl.col("close"), period=10), "mom"), + ( + "nvi_line", + lambda: ta.nvi_line(pl.col("close"), pl.col("volume")), + "nvi_line", + ), + ( + "nvi_signal", + lambda: ta.nvi_signal(pl.col("close"), pl.col("volume"), signal_period=255), + "nvi_signal", + ), + ( + "obv_line", + lambda: ta.obv_line(pl.col("close"), pl.col("volume")), + "obv_line", + ), + ( + "obv_signal", + lambda: ta.obv_signal(pl.col("close"), pl.col("volume"), signal_period=9), + "obv_signal", + ), + ( + "pchan_upper", + lambda: ta.pchan_upper(pl.col("high"), pl.col("low"), period=20), + "pchan_upper", + ), + ( + "pchan_middle", + lambda: ta.pchan_middle(pl.col("high"), pl.col("low"), period=20), + "pchan_middle", + ), + ( + "pchan_lower", + lambda: ta.pchan_lower(pl.col("high"), pl.col("low"), period=20), + "pchan_lower", + ), + ( + "ppo_line", + lambda: ta.ppo_line(pl.col("close"), fast_period=12, slow_period=26), + "ppo_line", + ), + ( + "ppo_signal", + lambda: ta.ppo_signal( + pl.col("close"), + fast_period=12, + slow_period=26, + signal_period=9, + ), + "ppo_signal", + ), + ( + "ppo_histogram", + lambda: ta.ppo_histogram( + pl.col("close"), + fast_period=12, + slow_period=26, + signal_period=9, + ), + "ppo_histogram", + ), + ( + "psar", + lambda: ta.psar( + pl.col("high"), + pl.col("low"), + pl.col("close"), + increment=0.02, + initial_acceleration_factor=0.02, + max_acceleration_factor=0.2, + ), + "psar", + ), + ("psl", lambda: ta.psl(pl.col("close"), period=12), "psl"), + ( + "pvi_line", + lambda: ta.pvi_line(pl.col("close"), pl.col("volume")), + "pvi_line", + ), + ( + "pvi_signal", + lambda: ta.pvi_signal(pl.col("close"), pl.col("volume"), signal_period=255), + "pvi_signal", + ), + ( + "pvo_line", + lambda: ta.pvo_line(pl.col("volume"), fast_period=12, slow_period=26), + "pvo_line", + ), + ( + "pvo_signal", + lambda: ta.pvo_signal( + pl.col("volume"), + fast_period=12, + slow_period=26, + signal_period=9, + ), + "pvo_signal", + ), + ( + "pvo_histogram", + lambda: ta.pvo_histogram( + pl.col("volume"), + fast_period=12, + slow_period=26, + signal_period=9, + ), + "pvo_histogram", + ), + ("roc", lambda: ta.roc(pl.col("close"), period=20), "roc"), + ("rsi", lambda: ta.rsi(pl.col("close"), period=14), "rsi"), ("bband_middle", lambda: ta.bband_middle(pl.col("close"), period=20), "sma"), ( "bband_lower", @@ -163,6 +464,52 @@ def select_expr(df: pl.DataFrame, expr: pl.Expr, alias: str, lazy: bool) -> pl.S lambda: ta.ichimoku_lagging_span(pl.col("close"), base_line_period=26), "ichimoku_lagging_span", ), + ( + "sonar_line", + lambda: ta.sonar_line(pl.col("close"), period=9, step=6), + "sonar_line", + ), + ( + "sonar_signal", + lambda: ta.sonar_signal(pl.col("close"), period=9, step=6, signal_period=5), + "sonar_signal", + ), + ( + "stochrsi_percent_k", + lambda: ta.stochrsi_percent_k( + pl.col("close"), period_rsi=14, period_k=14, period_d=3 + ), + "stochrsi_K", + ), + ( + "stochrsi_percent_d", + lambda: ta.stochrsi_percent_d( + pl.col("close"), period_rsi=14, period_k=14, period_d=3 + ), + "stochrsi_D", + ), + ( + "ultosc", + lambda: ta.ultosc( + pl.col("high"), + pl.col("low"), + pl.col("close"), + period_short=7, + period_medium=14, + period_long=28, + ), + "ultosc", + ), + ( + "vr", + lambda: ta.vr(pl.col("close"), pl.col("volume"), period=20), + "vr", + ), + ( + "willr", + lambda: ta.willr(pl.col("high"), pl.col("low"), pl.col("close"), period=14), + "willr", + ), ] # Leading spans are forward-projected in core fixtures, so Polars compares a @@ -210,7 +557,11 @@ def test_indicator_matches_core_expected( result = select_expr(df, expr_builder(), name, lazy) # then - assert_values_close(result.to_list(), expected) + assert_values_close( + result.to_list(), + expected, + abs_tol=CASE_TOLERANCES.get(name, 1e-8), + ) @pytest.mark.parametrize("symbol", SYMBOLS) @@ -275,13 +626,16 @@ def test_multi_input_integer_columns_are_cast_to_float() -> None: ).get_column("value") # then - assert_values_close(result.to_list(), [ - None, - None, - 58.33333333, - 58.33333333, - 58.33333333, - ]) + assert_values_close( + result.to_list(), + [ + None, + None, + 58.33333333, + 58.33333333, + 58.33333333, + ], + ) @pytest.mark.parametrize("lazy", [False, True])