⚡️ Optimize single-output indicator paths (#9)

* optimize single-output indicator paths

* hide internal bband helper

* reuse slow stochastic percent k

Author: sjquant

core/src/indicators/bband.rs

@@ -12,6 +12,20 @@ pub fn bband(
     (upper_band, center, lower_band)
 }
 
+pub fn bband_middle(data: &[f64], period: usize) -> Vec<Option<f64>> {
+    sma(data, period)
+}
+
+pub fn bband_upper(data: &[f64], period: usize, sigma: Option<f64>) -> Vec<Option<f64>> {
+    let (upper_band, _) = bband_bands(data, period, sigma);
+    upper_band
+}
+
+pub fn bband_lower(data: &[f64], period: usize, sigma: Option<f64>) -> Vec<Option<f64>> {
+    let (_, lower_band) = bband_bands(data, period, sigma);
+    lower_band
+}
+
 fn bband_bands(
     data: &[f64],
     period: usize,

core/src/indicators/macd.rs

@@ -22,6 +22,39 @@ pub fn macd(
     (macd_line, signal_line, histogram)
 }
 
+pub fn macd_line(data: &[f64], fast_period: usize, slow_period: usize) -> Vec<Option<f64>> {
+    calc_macd_line(data, fast_period, slow_period)
+}
+
+pub fn macd_signal(
+    data: &[f64],
+    fast_period: usize,
+    slow_period: usize,
+    signal_period: usize,
+) -> Vec<Option<f64>> {
+    let macd_line = calc_macd_line(data, fast_period, slow_period);
+    calc_macd_signal(&macd_line, signal_period)
+}
+
+pub fn macd_histogram(
+    data: &[f64],
+    fast_period: usize,
+    slow_period: usize,
+    signal_period: usize,
+) -> Vec<Option<f64>> {
+    let macd_line = calc_macd_line(data, fast_period, slow_period);
+    let signal_line = calc_macd_signal(&macd_line, signal_period);
+
+    macd_line
+        .iter()
+        .zip(signal_line.iter())
+        .map(|(&macd, &signal)| match (macd, signal) {
+            (Some(m), Some(s)) => Some(m - s),
+            _ => None,
+        })
+        .collect()
+}
+
 fn calc_macd_line(data: &[f64], fast_period: usize, slow_period: usize) -> Vec<Option<f64>> {
     let mut macd_line = vec![None; data.len()];
 

core/src/indicators/stochf.rs

@@ -1,46 +1,70 @@
 use crate::utils::{calc_mean, find_max, find_min};
 
-pub fn stochf(
+pub fn stochf_percent_k(
     highs: &[f64],
     lows: &[f64],
     closes: &[f64],
     fastk_period: usize,
-    fastd_period: usize,
-) -> (Vec<Option<f64>>, Vec<Option<f64>>) {
+) -> Vec<Option<f64>> {
     let len = closes.len();
     let mut percent_k = vec![None; len];
-    let mut percent_d = vec![None; len];
 
     if len < fastk_period {
-        return (percent_k, percent_d);
+        return percent_k;
     }
 
     for i in (fastk_period - 1)..len {
         let max_high = find_max(&highs[i + 1 - fastk_period..=i]);
         let min_low = find_min(&lows[i + 1 - fastk_period..=i]);
 
-        let k = if max_high == min_low {
+        percent_k[i] = if max_high == min_low {
             None
         } else {
             Some(((closes[i] - min_low) / (max_high - min_low)) * 100.0)
         };
+    }
+
+    percent_k
+}
 
-        percent_k[i] = k;
+pub fn stochf_percent_d(
+    percent_k: &[Option<f64>],
+    fastk_period: usize,
+    fastd_period: usize,
+) -> Vec<Option<f64>> {
+    let len = percent_k.len();
+    let mut percent_d = vec![None; len];
 
+    if len < fastk_period {
+        return percent_d;
+    }
+
+    for i in (fastk_period - 1)..len {
         if fastd_period == 1 {
-            percent_d[i] = k;
+            percent_d[i] = percent_k[i];
         } else if i >= fastk_period - 1 + (fastd_period - 1) {
             let slice = &percent_k[i + 1 - fastd_period..=i];
             let valid_values: Vec<f64> = slice.iter().filter_map(|&x| x).collect();
-            let d = if valid_values.len() == fastd_period {
+            percent_d[i] = if valid_values.len() == fastd_period {
                 Some(calc_mean(&valid_values))
             } else {
                 None
             };
-            percent_d[i] = d;
         }
     }
 
+    percent_d
+}
+
+pub fn stochf(
+    highs: &[f64],
+    lows: &[f64],
+    closes: &[f64],
+    fastk_period: usize,
+    fastd_period: usize,
+) -> (Vec<Option<f64>>, Vec<Option<f64>>) {
+    let percent_k = stochf_percent_k(highs, lows, closes, fastk_period);
+    let percent_d = stochf_percent_d(&percent_k, fastk_period, fastd_period);
     (percent_k, percent_d)
 }
 

core/src/indicators/stochs.rs

@@ -1,22 +1,18 @@
 use crate::utils::{calc_mean, find_max, find_min};
 
-pub fn stochs(
+fn stochs_raw_k(
     highs: &[f64],
     lows: &[f64],
     closes: &[f64],
     fastk_period: usize,
-    slowk_period: usize,
-    slowd_period: usize,
-) -> (Vec<Option<f64>>, Vec<Option<f64>>) {
+) -> Vec<Option<f64>> {
     let len = closes.len();
-    let mut percent_k = vec![None; len];
-    let mut percent_d = vec![None; len];
+    let mut raw_k = vec![None; len];
 
     if len < fastk_period {
-        return (percent_k, percent_d);
+        return raw_k;
     }
 
-    let mut raw_k = vec![None; len];
     for i in (fastk_period - 1)..len {
         let max_high = find_max(&highs[i + 1 - fastk_period..=i]);
         let min_low = find_min(&lows[i + 1 - fastk_period..=i]);
@@ -27,6 +23,25 @@ pub fn stochs(
             Some(((closes[i] - min_low) / (max_high - min_low)) * 100.0)
         };
     }
+
+    raw_k
+}
+
+pub fn stoch_percent_k(
+    highs: &[f64],
+    lows: &[f64],
+    closes: &[f64],
+    fastk_period: usize,
+    slowk_period: usize,
+) -> Vec<Option<f64>> {
+    let len = closes.len();
+    let mut percent_k = vec![None; len];
+
+    if len < fastk_period {
+        return percent_k;
+    }
+
+    let raw_k = stochs_raw_k(highs, lows, closes, fastk_period);
     for i in (fastk_period + slowk_period - 2)..len {
         let slice = &raw_k[i + 1 - slowk_period..=i];
         let valid_values: Vec<f64> = slice.iter().filter_map(|&x| x).collect();
@@ -37,6 +52,34 @@ pub fn stochs(
         };
     }
 
+    percent_k
+}
+
+pub fn stoch_percent_d(
+    highs: &[f64],
+    lows: &[f64],
+    closes: &[f64],
+    fastk_period: usize,
+    slowk_period: usize,
+    slowd_period: usize,
+) -> Vec<Option<f64>> {
+    let percent_k = stoch_percent_k(highs, lows, closes, fastk_period, slowk_period);
+    stoch_percent_d_from_k(&percent_k, fastk_period, slowk_period, slowd_period)
+}
+
+fn stoch_percent_d_from_k(
+    percent_k: &[Option<f64>],
+    fastk_period: usize,
+    slowk_period: usize,
+    slowd_period: usize,
+) -> Vec<Option<f64>> {
+    let len = percent_k.len();
+    let mut percent_d = vec![None; len];
+
+    if len < fastk_period {
+        return percent_d;
+    }
+
     for i in (fastk_period + slowk_period + slowd_period - 3)..len {
         let slice = &percent_k[i + 1 - slowd_period..=i];
         let valid_values: Vec<f64> = slice.iter().filter_map(|&x| x).collect();
@@ -47,6 +90,19 @@ pub fn stochs(
         };
     }
 
+    percent_d
+}
+
+pub fn stochs(
+    highs: &[f64],
+    lows: &[f64],
+    closes: &[f64],
+    fastk_period: usize,
+    slowk_period: usize,
+    slowd_period: usize,
+) -> (Vec<Option<f64>>, Vec<Option<f64>>) {
+    let percent_k = stoch_percent_k(highs, lows, closes, fastk_period, slowk_period);
+    let percent_d = stoch_percent_d_from_k(&percent_k, fastk_period, slowk_period, slowd_period);
     (percent_k, percent_d)
 }
 

polars/src/expressions.rs

@@ -2,13 +2,17 @@ use polars::prelude::*;
 use pyo3_polars::derive::polars_expr;
 use serde::Deserialize;
 use techr::{
-    bband as techr_bband, disparity as techr_disparity, ema as techr_ema,
+    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,
     ichimoku_base_line as techr_ichimoku_base_line,
     ichimoku_conversion_line as techr_ichimoku_conversion_line,
     ichimoku_lagging_span as techr_ichimoku_lagging_span,
     ichimoku_leading_span_a as techr_ichimoku_leading_span_a,
-    ichimoku_leading_span_b as techr_ichimoku_leading_span_b, macd as techr_macd, sma as techr_sma,
-    stochf as techr_stochf, stochs as techr_stochs, wma as techr_wma,
+    ichimoku_leading_span_b as techr_ichimoku_leading_span_b,
+    macd_histogram as techr_macd_histogram, macd_line as techr_macd_line,
+    macd_signal as techr_macd_signal, sma as techr_sma, stoch_percent_d as techr_stoch_percent_d,
+    stoch_percent_k as techr_stoch_percent_k, stochf_percent_d as techr_stochf_percent_d,
+    stochf_percent_k as techr_stochf_percent_k, wma as techr_wma,
 };
 
 #[derive(Deserialize)]
@@ -116,19 +120,18 @@ fn disparity(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult<Series> {
 #[polars_expr(output_type=Float64)]
 fn macd(inputs: &[Series], kwargs: FastSlowKwargs) -> PolarsResult<Series> {
     let input = series_to_f64_vec(&inputs[0])?;
-    let (macd_line, _, _) = techr_macd(
+    let macd_line = techr_macd_line(
         &input,
         kwargs.fast_period as usize,
         kwargs.slow_period as usize,
-        9,
     );
     Ok(option_vec_to_series(macd_line))
 }
 
 #[polars_expr(output_type=Float64)]
 fn macd_signal(inputs: &[Series], kwargs: FastSlowSignalKwargs) -> PolarsResult<Series> {
     let input = series_to_f64_vec(&inputs[0])?;
-    let (_, signal_line, _) = techr_macd(
+    let signal_line = techr_macd_signal(
         &input,
         kwargs.fast_period as usize,
         kwargs.slow_period as usize,
@@ -140,7 +143,7 @@ fn macd_signal(inputs: &[Series], kwargs: FastSlowSignalKwargs) -> PolarsResult<
 #[polars_expr(output_type=Float64)]
 fn macd_hist(inputs: &[Series], kwargs: FastSlowSignalKwargs) -> PolarsResult<Series> {
     let input = series_to_f64_vec(&inputs[0])?;
-    let (_, _, histogram) = techr_macd(
+    let histogram = techr_macd_histogram(
         &input,
         kwargs.fast_period as usize,
         kwargs.slow_period as usize,
@@ -152,21 +155,21 @@ fn macd_hist(inputs: &[Series], kwargs: FastSlowSignalKwargs) -> PolarsResult<Se
 #[polars_expr(output_type=Float64)]
 fn bband_middle(inputs: &[Series], kwargs: PeriodKwargs) -> PolarsResult<Series> {
     let input = series_to_f64_vec(&inputs[0])?;
-    let (_, middle, _) = techr_bband(&input, kwargs.period as usize, None);
+    let middle = techr_bband_middle(&input, kwargs.period as usize);
     Ok(option_vec_to_series(middle))
 }
 
 #[polars_expr(output_type=Float64)]
 fn bband_lower(inputs: &[Series], kwargs: BBandKwargs) -> PolarsResult<Series> {
     let input = series_to_f64_vec(&inputs[0])?;
-    let (_, _, lower) = techr_bband(&input, kwargs.period as usize, Some(kwargs.sigma));
+    let lower = techr_bband_lower(&input, kwargs.period as usize, Some(kwargs.sigma));
     Ok(option_vec_to_series(lower))
 }
 
 #[polars_expr(output_type=Float64)]
 fn bband_upper(inputs: &[Series], kwargs: BBandKwargs) -> PolarsResult<Series> {
     let input = series_to_f64_vec(&inputs[0])?;
-    let (upper, _, _) = techr_bband(&input, kwargs.period as usize, Some(kwargs.sigma));
+    let upper = techr_bband_upper(&input, kwargs.period as usize, Some(kwargs.sigma));
     Ok(option_vec_to_series(upper))
 }
 
@@ -175,13 +178,7 @@ fn stochf_percent_k(inputs: &[Series], kwargs: StochFKwargs) -> PolarsResult<Ser
     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 (percent_k, _) = techr_stochf(
-        &highs,
-        &lows,
-        &closes,
-        kwargs.fastk_period as usize,
-        kwargs.fastd_period as usize,
-    );
+    let percent_k = techr_stochf_percent_k(&highs, &lows, &closes, kwargs.fastk_period as usize);
     Ok(option_vec_to_series(percent_k))
 }
 
@@ -190,10 +187,9 @@ fn stochf_percent_d(inputs: &[Series], kwargs: StochFKwargs) -> PolarsResult<Ser
     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 (_, percent_d) = techr_stochf(
-        &highs,
-        &lows,
-        &closes,
+    let percent_k = techr_stochf_percent_k(&highs, &lows, &closes, kwargs.fastk_period as usize);
+    let percent_d = techr_stochf_percent_d(
+        &percent_k,
         kwargs.fastk_period as usize,
         kwargs.fastd_period as usize,
     );
@@ -205,13 +201,12 @@ fn stoch_percent_k(inputs: &[Series], kwargs: StochKwargs) -> PolarsResult<Serie
     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 (percent_k, _) = techr_stochs(
+    let percent_k = techr_stoch_percent_k(
         &highs,
         &lows,
         &closes,
         kwargs.fastk_period as usize,
         kwargs.slowk_period as usize,
-        kwargs.slowd_period as usize,
     );
     Ok(option_vec_to_series(percent_k))
 }
@@ -221,7 +216,7 @@ fn stoch_percent_d(inputs: &[Series], kwargs: StochKwargs) -> PolarsResult<Serie
     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 (_, percent_d) = techr_stochs(
+    let percent_d = techr_stoch_percent_d(
         &highs,
         &lows,
         &closes,