diff --git a/.github/workflows/polars-ci.yml b/.github/workflows/polars-ci.yml index 00e4573..fb9e376 100644 --- a/.github/workflows/polars-ci.yml +++ b/.github/workflows/polars-ci.yml @@ -2,12 +2,11 @@ name: Polars CI on: pull_request: - types: [review_requested, ready_for_review] + types: [review_requested, ready_for_review, synchronize] paths: - ".github/workflows/polars-*.yml" - "Cargo.toml" - "Makefile" - - "core/**" - "polars/**" concurrency: @@ -18,42 +17,82 @@ permissions: contents: read jobs: + preflight: + runs-on: ubuntu-latest + permissions: + contents: read + pull-requests: read + outputs: + validate: ${{ steps.filter.outputs.validate }} + steps: + - uses: dorny/paths-filter@d1c1ffe0248fe513906c8e24db8ea791d46f8590 + id: filter + with: + filters: | + validate: + - "Cargo.toml" + - "Makefile" + - "polars/**" + validate: + needs: preflight + if: ${{ always() }} runs-on: ubuntu-latest defaults: run: shell: bash steps: - - uses: actions/checkout@v6 - - uses: actions/setup-python@v6 + - name: Fail if preflight did not succeed + if: ${{ needs.preflight.result != 'success' }} + run: | + echo "preflight job did not succeed: ${{ needs.preflight.result }}" + exit 1 + + - name: Skip validation for non-polars changes + if: ${{ needs.preflight.result == 'success' && needs.preflight.outputs.validate != 'true' }} + run: echo "Skipping Polars CI validate job because this PR does not touch Cargo.toml, Makefile, or polars/**." + + - if: ${{ needs.preflight.outputs.validate == 'true' }} + uses: actions/checkout@v6 + - if: ${{ needs.preflight.outputs.validate == 'true' }} + uses: actions/setup-python@v6 with: python-version: "3.10" - - uses: astral-sh/setup-uv@v8.0.0 - - uses: dtolnay/rust-toolchain@stable + - if: ${{ needs.preflight.outputs.validate == 'true' }} + uses: astral-sh/setup-uv@v8.0.0 + - if: ${{ needs.preflight.outputs.validate == 'true' }} + uses: dtolnay/rust-toolchain@stable - name: Sync polars dependencies + if: ${{ needs.preflight.outputs.validate == 'true' }} working-directory: polars run: uv sync --group dev --no-install-project - name: Test techr-core + if: ${{ needs.preflight.outputs.validate == 'true' }} run: cargo test -p techr-core - name: Build local extension for tests + if: ${{ needs.preflight.outputs.validate == 'true' }} working-directory: polars run: uv run maturin develop --uv - name: Test polars package + if: ${{ needs.preflight.outputs.validate == 'true' }} working-directory: polars run: uv run pytest - name: Build wheel and sdist + if: ${{ needs.preflight.outputs.validate == 'true' }} working-directory: polars run: uv run maturin build --release --sdist --out dist - name: Check artifact contents + if: ${{ needs.preflight.outputs.validate == 'true' }} run: python polars/scripts/check_artifacts.py polars/dist - name: Smoke test built wheel + if: ${{ needs.preflight.outputs.validate == 'true' }} run: | wheel="$(python - <<'PY' from pathlib import Path diff --git a/core/src/indicators/bband.rs b/core/src/indicators/bband.rs index 00f9941..ede803c 100644 --- a/core/src/indicators/bband.rs +++ b/core/src/indicators/bband.rs @@ -1,9 +1,8 @@ use crate::indicators::sma::sma; - -use crate::utils::round_scalar; +use crate::utils::{rolling_mean_stddev_strict, round_scalar}; pub fn bband( - data: &[f64], + data: &[Option], period: usize, sigma: Option, ) -> (Vec>, Vec>, Vec>) { @@ -12,55 +11,42 @@ pub fn bband( (upper_band, center, lower_band) } -pub fn bband_middle(data: &[f64], period: usize) -> Vec> { +pub fn bband_middle(data: &[Option], period: usize) -> Vec> { sma(data, period) } -pub fn bband_upper(data: &[f64], period: usize, sigma: Option) -> Vec> { +pub fn bband_upper(data: &[Option], period: usize, sigma: Option) -> Vec> { let (upper_band, _) = bband_bands(data, period, sigma); upper_band } -pub fn bband_lower(data: &[f64], period: usize, sigma: Option) -> Vec> { +pub fn bband_lower(data: &[Option], period: usize, sigma: Option) -> Vec> { let (_, lower_band) = bband_bands(data, period, sigma); lower_band } fn bband_bands( - data: &[f64], + data: &[Option], period: usize, sigma: Option, ) -> (Vec>, Vec>) { - let mut upper_band = vec![None; data.len()]; - let mut lower_band = vec![None; data.len()]; - let mut sum = 0.0; - let mut sum_sq = 0.0; let sigma = sigma.unwrap_or(2.0); + let (means, stddevs) = rolling_mean_stddev_strict(data, period); - if data.len() < period { - return (upper_band, lower_band); - } - - for i in 0..data.len() { - sum += data[i]; - sum_sq += data[i] * data[i]; - - if i >= period { - sum -= data[i - period]; - sum_sq -= data[i - period] * data[i - period]; - } - - if i >= period - 1 { - let mean = sum / period as f64; - let variance = (sum_sq / period as f64) - (mean * mean); - let stddev = variance.sqrt(); - let deviation = sigma * stddev; - upper_band[i] = Some(round_scalar(mean + deviation, 8)); - lower_band[i] = Some(round_scalar(mean - deviation, 8)); - } - } - - (upper_band, lower_band) + means + .into_iter() + .zip(stddevs) + .map(|(mean, stddev)| match (mean, stddev) { + (Some(mean), Some(stddev)) => { + let deviation = sigma * stddev; + ( + Some(round_scalar(mean + deviation, 8)), + Some(round_scalar(mean - deviation, 8)), + ) + } + _ => (None, None), + }) + .unzip() } #[cfg(test)] @@ -73,7 +59,7 @@ mod tests { fn test_bband() { let test_cases = vec!["005930", "TSLA"]; for symbol in test_cases { - let input = testutils::load_data(&format!("../data/{}.json", symbol), "c"); + let input = testutils::load_data_nullable(&format!("../data/{}.json", symbol), "c"); let (upper, middle, lower) = bband(&input, 20, None); let expected_upper = testutils::load_expected::>(&format!( @@ -109,4 +95,26 @@ mod tests { ); } } + + #[test] + fn test_bband_with_interior_gap_invalidates_full_window() { + let input = vec![Some(1.0), Some(2.0), None, Some(4.0), Some(5.0)]; + + let (upper, middle, lower) = bband(&input, 2, Some(2.0)); + + assert_eq!(middle, vec![None, Some(1.5), None, None, Some(4.5)]); + assert_eq!(upper, vec![None, Some(2.5), None, None, Some(5.5)]); + assert_eq!(lower, vec![None, Some(0.5), None, None, Some(3.5)]); + } + + #[test] + fn test_bband_full_window_invalidation() { + let input = vec![None, Some(2.0), None, Some(4.0)]; + + let (upper, middle, lower) = bband(&input, 2, None); + + assert_eq!(upper, vec![None, None, None, None]); + assert_eq!(middle, vec![None, None, None, None]); + assert_eq!(lower, vec![None, None, None, None]); + } } diff --git a/core/src/indicators/cv.rs b/core/src/indicators/cv.rs index 96478c8..a06fd25 100644 --- a/core/src/indicators/cv.rs +++ b/core/src/indicators/cv.rs @@ -1,4 +1,4 @@ -use crate::indicators::ema::ema; +use crate::indicators::ema::ema_dense; pub fn cv(highs: &[f64], lows: &[f64], period: usize) -> Vec> { let mut cv = vec![None; highs.len()]; @@ -9,7 +9,7 @@ pub fn cv(highs: &[f64], lows: &[f64], period: usize) -> Vec> { } let high_low_diffs: Vec = highs.iter().zip(lows.iter()).map(|(h, l)| h - l).collect(); - let ema_high_low_diffs = ema(&high_low_diffs, period); + let ema_high_low_diffs = ema_dense(&high_low_diffs, period); for i in period * 2 - 1..len { if let (Some(current_ema), Some(previous_ema)) = diff --git a/core/src/indicators/disparity.rs b/core/src/indicators/disparity.rs index 9e8a32f..8d6c13a 100644 --- a/core/src/indicators/disparity.rs +++ b/core/src/indicators/disparity.rs @@ -1,19 +1,23 @@ use crate::indicators::sma::sma; -pub fn disparity(data: &[f64], period: usize) -> Vec> { +pub fn disparity(data: &[Option], period: usize) -> Vec> { let len = data.len(); let mut result = vec![None; len]; - if len < period { + if len < period || period == 0 { return result; } let sma = sma(data, period); for i in period - 1..len { + let Some(value) = data[i] else { + continue; + }; + if let Some(sma_value) = sma[i] { if sma_value != 0.0 { - result[i] = Some((data[i] / sma_value) * 100.0); + result[i] = Some((value / sma_value) * 100.0); } } } @@ -31,7 +35,7 @@ mod tests { fn test_disparity() { let test_cases = vec!["005930", "TSLA"]; for symbol in test_cases { - let close = testutils::load_data(&format!("../data/{}.json", symbol), "c"); + let close = testutils::load_data_nullable(&format!("../data/{}.json", symbol), "c"); let result = disparity(&close, 20); let expected = testutils::load_expected::>(&format!( "../data/expected/disparity_{}.json", @@ -46,4 +50,29 @@ mod tests { ); } } + + #[test] + fn test_disparity_with_interior_gap_invalidates_window() { + let input = vec![Some(1.0), Some(2.0), None, Some(4.0), Some(5.0)]; + let expected = vec![ + None, + Some(133.33333333333331), + None, + None, + Some(111.11111111111111), + ]; + + let result = disparity(&input, 2); + + assert_eq!(result, expected); + } + + #[test] + fn test_disparity_full_window_invalidation() { + let input = vec![None, Some(2.0), None, Some(4.0)]; + + let result = disparity(&input, 2); + + assert_eq!(result, vec![None, None, None, None]); + } } diff --git a/core/src/indicators/efi.rs b/core/src/indicators/efi.rs index 7fce517..bf54d59 100644 --- a/core/src/indicators/efi.rs +++ b/core/src/indicators/efi.rs @@ -1,4 +1,4 @@ -use crate::indicators::ema::ema; +use crate::indicators::ema::ema_dense; pub fn efi(closes: &[f64], volumes: &[f64], period: usize) -> Vec> { let len = closes.len(); @@ -22,7 +22,7 @@ pub fn efi(closes: &[f64], volumes: &[f64], period: usize) -> Vec> { *efi_val = Some(force_val); }); } else { - let ema_result = ema(&force, period); + let ema_result = ema_dense(&force, period); efi.iter_mut() .skip(1) .zip(ema_result.into_iter()) diff --git a/core/src/indicators/ema.rs b/core/src/indicators/ema.rs index 06d9fae..5fdb38d 100644 --- a/core/src/indicators/ema.rs +++ b/core/src/indicators/ema.rs @@ -1,35 +1,17 @@ -/// Computes an exponential moving average over a dense `f64` series. -/// -/// The returned vector keeps the same length as the input and emits `None` -/// until the first full `period` window has been observed. -pub fn ema(data: &[f64], period: usize) -> Vec> { - let mut result = vec![None; data.len()]; - - if data.len() < period { - return result; - } - - let alpha = 2.0 / (period as f64 + 1.0); - let mut ema = data[..period].iter().sum::() / period as f64; - - result[period - 1] = Some(ema); - - for i in period..data.len() { - ema = alpha * data[i] + (1.0 - alpha) * ema; - result[i] = Some(ema); - } - - result +pub(crate) fn ema_dense(data: &[f64], period: usize) -> Vec> { + let nullable = data.iter().copied().map(Some).collect::>(); + ema(&nullable, period) } -/// Computes an EMA over an optional series while preserving original alignment. +/// Computes an exponential moving average over an aligned nullable series. /// -/// The EMA state advances only on `Some(f64)` values, so interior `None` holes -/// are skipped safely. This lets callers avoid compacting a sparse aligned -/// series into a temporary dense vector before applying `ema`. -pub(crate) fn ema_aligned(data: &[Option], period: usize) -> Vec> { +/// The returned vector keeps the same length as the input and emits `None` +/// until the first contiguous run of `period` valid observations has been +/// observed. Once seeded, gaps emit `None` without resetting the EMA state. +pub fn ema(data: &[Option], period: usize) -> Vec> { let mut result = vec![None; data.len()]; - if period == 0 { + + if data.len() < period || period == 0 { return result; } @@ -38,8 +20,12 @@ pub(crate) fn ema_aligned(data: &[Option], period: usize) -> Vec], period: usize) -> Vec>(&format!( "../data/expected/ema_{}.json", @@ -93,31 +81,41 @@ mod tests { } } - /// Verifies that aligned EMA preserves offsets while matching dense EMA values. #[test] - fn test_ema_aligned() { - // Given + fn test_ema_with_prefix_gap() { let aligned = vec![None, None, Some(1.0), Some(2.0), Some(3.0), Some(4.0)]; let expected = vec![None, None, None, Some(1.5), Some(2.5), Some(3.5)]; - // When - let result = ema_aligned(&aligned, 2); + let result = ema(&aligned, 2); - // Then assert_eq!(result, expected); } - /// Verifies that aligned EMA skips interior gaps without panicking. #[test] - fn test_ema_aligned_with_interior_gaps() { - // Given + fn test_ema_with_interior_gaps_resumes_from_prior_state() { let aligned = vec![None, None, Some(1.0), Some(2.0), None, Some(3.0), Some(4.0)]; let expected = vec![None, None, None, Some(1.5), None, Some(2.5), Some(3.5)]; - // When - let result = ema_aligned(&aligned, 2); + let result = ema(&aligned, 2); + + assert_eq!(result, expected); + } + + #[test] + fn test_ema_requires_contiguous_values_before_seed() { + let aligned = vec![ + Some(1.0), + Some(2.0), + None, + Some(3.0), + Some(4.0), + Some(5.0), + Some(6.0), + ]; + let expected = vec![None, None, None, None, None, Some(4.0), Some(5.0)]; + + let result = ema(&aligned, 3); - // Then assert_eq!(result, expected); } } diff --git a/core/src/indicators/env.rs b/core/src/indicators/env.rs index a25eb8f..2cbd546 100644 --- a/core/src/indicators/env.rs +++ b/core/src/indicators/env.rs @@ -1,7 +1,7 @@ use crate::indicators::sma::sma; pub fn env( - data: &[f64], + data: &[Option], period: usize, shift_percentage: f64, ) -> (Vec>, Vec>, Vec>) { @@ -33,7 +33,7 @@ mod tests { fn test_env() { let test_cases = vec!["005930", "TSLA"]; for symbol in test_cases { - let input = testutils::load_data(&format!("../data/{}.json", symbol), "c"); + let input = testutils::load_data_nullable(&format!("../data/{}.json", symbol), "c"); let result = env(&input, 20, 10.0); let (env_upper, sma_values, env_lower) = result; @@ -67,4 +67,32 @@ mod tests { ); } } + + #[test] + fn test_env_with_interior_gap_invalidates_window() { + let input = vec![Some(1.0), Some(2.0), None, Some(4.0), Some(5.0)]; + + let (upper, middle, lower) = env(&input, 2, 10.0); + + assert_eq!(middle, vec![None, Some(1.5), None, None, Some(4.5)]); + assert_eq!( + round_vec(upper, 8), + vec![None, Some(1.65), None, None, Some(4.95)] + ); + assert_eq!( + round_vec(lower, 8), + vec![None, Some(1.35), None, None, Some(4.05)] + ); + } + + #[test] + fn test_env_full_window_invalidation() { + let input = vec![None, Some(2.0), None, Some(4.0)]; + + let (upper, middle, lower) = env(&input, 2, 10.0); + + assert_eq!(upper, vec![None, None, None, None]); + assert_eq!(middle, vec![None, None, None, None]); + assert_eq!(lower, vec![None, None, None, None]); + } } diff --git a/core/src/indicators/eom.rs b/core/src/indicators/eom.rs index 8ef4461..9fed832 100644 --- a/core/src/indicators/eom.rs +++ b/core/src/indicators/eom.rs @@ -1,4 +1,4 @@ -use crate::indicators::sma::{sma, sma_aligned}; +use crate::indicators::sma::{sma_aligned, sma_dense}; pub fn eom( highs: &[f64], @@ -55,7 +55,7 @@ pub fn eom_line( } let mut eom_line = vec![None; len]; - let eom_sma = sma(&eom_values, period); + let eom_sma = sma_dense(&eom_values, period); for (i, &value) in eom_sma.iter().enumerate() { eom_line[i + 1] = value; } diff --git a/core/src/indicators/erbear.rs b/core/src/indicators/erbear.rs index b998700..f5ef92a 100644 --- a/core/src/indicators/erbear.rs +++ b/core/src/indicators/erbear.rs @@ -1,4 +1,4 @@ -use crate::indicators::ema::ema; +use crate::indicators::ema::ema_dense; pub fn erbear(lows: &[f64], closes: &[f64], period: usize) -> Vec> { let mut erbear = vec![None; lows.len()]; @@ -7,7 +7,7 @@ pub fn erbear(lows: &[f64], closes: &[f64], period: usize) -> Vec> { return erbear; } - let ema_values = ema(closes, period); + let ema_values = ema_dense(closes, period); for i in (period - 1)..lows.len() { if let Some(ema_value) = ema_values[i] { diff --git a/core/src/indicators/erbull.rs b/core/src/indicators/erbull.rs index 01fd17d..cda1125 100644 --- a/core/src/indicators/erbull.rs +++ b/core/src/indicators/erbull.rs @@ -1,4 +1,4 @@ -use crate::indicators::ema::ema; +use crate::indicators::ema::ema_dense; pub fn erbull(highs: &[f64], closes: &[f64], period: usize) -> Vec> { let mut erbull = vec![None; highs.len()]; @@ -7,7 +7,7 @@ pub fn erbull(highs: &[f64], closes: &[f64], period: usize) -> Vec> return erbull; } - let ema_values = ema(closes, period); + let ema_values = ema_dense(closes, period); for i in (period - 1)..highs.len() { if let Some(ema_value) = ema_values[i] { diff --git a/core/src/indicators/macd.rs b/core/src/indicators/macd.rs index 0cb3d9f..1a605e2 100644 --- a/core/src/indicators/macd.rs +++ b/core/src/indicators/macd.rs @@ -1,4 +1,4 @@ -use crate::indicators::ema::{ema, ema_aligned}; +use crate::indicators::ema::{ema_aligned, ema_dense}; pub fn macd( data: &[f64], @@ -56,8 +56,8 @@ pub fn macd_line(data: &[f64], fast_period: usize, slow_period: usize) -> Vec = highs.iter().zip(lows.iter()).map(|(h, l)| h - l).collect(); - let s_ema = ema(&high_low_diffs, period_ema); + let s_ema = ema_dense(&high_low_diffs, period_ema); let offset: usize = period_ema - 1; let d_ema = ema_aligned(&s_ema, period_ema); diff --git a/core/src/indicators/ppo.rs b/core/src/indicators/ppo.rs index 100f39d..63ea40d 100644 --- a/core/src/indicators/ppo.rs +++ b/core/src/indicators/ppo.rs @@ -1,4 +1,4 @@ -use crate::indicators::ema::{ema, ema_aligned}; +use crate::indicators::ema::{ema_aligned, ema_dense}; pub fn ppo( data: &[f64], @@ -55,8 +55,8 @@ pub fn ppo_line(data: &[f64], fast_period: usize, slow_period: usize) -> Vec Vec Vec> { - let mut sma = vec![None; data.len()]; - let mut sum = 0.0; - - if data.len() < period { - return sma; - } - - for i in 0..data.len() { - sum += data[i]; - if i >= period { - sum -= data[i - period]; - } - if i >= period - 1 { - sma[i] = Some(sum / period as f64); - } - } +use crate::utils::rolling_mean_strict; - sma +pub(crate) fn sma_dense(data: &[f64], period: usize) -> Vec> { + let nullable = data.iter().copied().map(Some).collect::>(); + rolling_mean_strict(&nullable, period) } -/// Computes an SMA over an aligned optional series. +/// Computes a simple moving average over an aligned nullable series. /// -/// The input is expected to contain an optional prefix of `None` values followed -/// by a contiguous run of `Some(f64)` values. The returned vector preserves the -/// original alignment while avoiding the extra compaction/remapping pass that -/// would otherwise be needed before applying `sma`. -pub(crate) fn sma_aligned(data: &[Option], period: usize) -> Vec> { - let mut result = vec![None; data.len()]; - let Some(first_valid_idx) = data.iter().position(|value| value.is_some()) else { - return result; - }; - - let valid_len = data.len() - first_valid_idx; - if period == 0 || valid_len < period { - return result; - } - - let mut sum = 0.0; - for value in data.iter().skip(first_valid_idx).take(period) { - sum += value.expect("initial SMA window must be fully populated"); - } - - let first_signal_idx = first_valid_idx + period - 1; - result[first_signal_idx] = Some(sum / period as f64); - - for idx in (first_signal_idx + 1)..data.len() { - let entering = - data[idx].expect("aligned SMA input must be contiguous after the first value"); - let leaving = - data[idx - period].expect("aligned SMA input must be contiguous after the first value"); - sum += entering - leaving; - result[idx] = Some(sum / period as f64); - } - - result +/// The returned vector keeps the same length as the input and emits `None` +/// until the first full `period` window has been observed. +pub fn sma(data: &[Option], period: usize) -> Vec> { + rolling_mean_strict(data, period) } +pub(crate) use sma as sma_aligned; + #[cfg(test)] mod tests { use super::*; @@ -74,7 +29,7 @@ mod tests { // When for symbol in test_cases { - let input = testutils::load_data(&format!("../data/{}.json", symbol), "c"); + let input = testutils::load_data_nullable(&format!("../data/{}.json", symbol), "c"); let result = sma(&input, 20); let expected = testutils::load_expected::>(&format!( "../data/expected/sma_{}.json", @@ -91,17 +46,33 @@ mod tests { } } - /// Verifies that aligned SMA preserves offsets while matching dense SMA values. #[test] - fn test_sma_aligned() { - // Given + fn test_sma_with_prefix_gap() { let aligned = vec![None, None, Some(1.0), Some(2.0), Some(3.0), Some(4.0)]; let expected = vec![None, None, None, Some(1.5), Some(2.5), Some(3.5)]; - // When - let result = sma_aligned(&aligned, 2); + let result = sma(&aligned, 2); + + assert_eq!(result, expected); + } + + #[test] + fn test_sma_with_interior_gap_invalidates_window() { + let aligned = vec![Some(1.0), Some(2.0), None, Some(4.0), Some(5.0)]; + let expected = vec![None, Some(1.5), None, None, Some(4.5)]; + + let result = sma(&aligned, 2); + + assert_eq!(result, expected); + } + + #[test] + fn test_sma_full_window_invalidation() { + let aligned = vec![Some(1.0), None, None, Some(4.0)]; + let expected = vec![None, None, None, None]; + + let result = sma(&aligned, 2); - // Then assert_eq!(result, expected); } } diff --git a/core/src/indicators/sonar.rs b/core/src/indicators/sonar.rs index 688aae8..e462d44 100644 --- a/core/src/indicators/sonar.rs +++ b/core/src/indicators/sonar.rs @@ -1,4 +1,4 @@ -use crate::indicators::ema::{ema, ema_aligned}; +use crate::indicators::ema::{ema_aligned, ema_dense}; pub fn sonar( data: &[f64], @@ -29,7 +29,7 @@ pub fn sonar_line(data: &[f64], period: usize, step: usize) -> Vec> return sonar_line; } - let ema_values = ema(data, period); + let ema_values = ema_dense(data, period); for i in (period + step - 1)..data.len() { if let (Some(current_ema), Some(previous_ema)) = (ema_values[i], ema_values[i - step]) { diff --git a/core/src/indicators/wma.rs b/core/src/indicators/wma.rs index 5b6a00a..ea0f44f 100644 --- a/core/src/indicators/wma.rs +++ b/core/src/indicators/wma.rs @@ -1,27 +1,7 @@ -pub fn wma(data: &[f64], period: usize) -> Vec> { - let mut result = vec![None; data.len()]; +use crate::utils::rolling_weighted_mean_strict; - if data.len() < period { - return result; - } - - let weight_sum = (period * (period + 1)) / 2; - let mut weighted_sum = 0.0; - - // Initialize the first period - for i in 0..period { - weighted_sum += data[i] * (i + 1) as f64; - } - - for i in period - 1..data.len() { - result[i] = Some(weighted_sum / weight_sum as f64); - if i + 1 < data.len() { - weighted_sum = weighted_sum + data[i + 1] * period as f64 - - data[i + 1 - period..=i].iter().sum::(); - } - } - - result +pub fn wma(data: &[Option], period: usize) -> Vec> { + rolling_weighted_mean_strict(data, period) } #[cfg(test)] @@ -34,7 +14,7 @@ mod tests { fn test_wma() { let test_cases = vec!["005930", "TSLA"]; for symbol in test_cases { - let input = testutils::load_data(&format!("../data/{}.json", symbol), "c"); + let input = testutils::load_data_nullable(&format!("../data/{}.json", symbol), "c"); let result = wma(&input, 20); let expected = testutils::load_expected::>(&format!( "../data/expected/wma_{}.json", @@ -49,4 +29,34 @@ mod tests { ); } } + + #[test] + fn test_wma_with_prefix_gap() { + let aligned = vec![None, Some(1.0), Some(2.0), Some(3.0)]; + let expected = vec![None, None, Some(5.0 / 3.0), Some(8.0 / 3.0)]; + + let result = wma(&aligned, 2); + + assert_eq!(result, expected); + } + + #[test] + fn test_wma_with_interior_gap_invalidates_window() { + let aligned = vec![Some(1.0), Some(2.0), None, Some(4.0), Some(5.0)]; + let expected = vec![None, Some(5.0 / 3.0), None, None, Some(14.0 / 3.0)]; + + let result = wma(&aligned, 2); + + assert_eq!(result, expected); + } + + #[test] + fn test_wma_full_window_invalidation() { + let aligned = vec![Some(1.0), None, None, Some(4.0)]; + let expected = vec![None, None, None, None]; + + let result = wma(&aligned, 2); + + assert_eq!(result, expected); + } } diff --git a/core/src/testutils.rs b/core/src/testutils.rs index 90e293d..14b1b0b 100644 --- a/core/src/testutils.rs +++ b/core/src/testutils.rs @@ -17,6 +17,11 @@ pub fn load_data(path: &str, field: &str) -> Vec { res.iter().map(|x| x[field_index]).collect() } +#[cfg(test)] +pub fn load_data_nullable(path: &str, field: &str) -> Vec> { + load_data(path, field).into_iter().map(Some).collect() +} + #[cfg(test)] pub fn load_expected(path: &str) -> Vec { use std::fs; diff --git a/core/src/utils.rs b/core/src/utils.rs index 7e264ba..dff8f89 100644 --- a/core/src/utils.rs +++ b/core/src/utils.rs @@ -121,7 +121,7 @@ pub fn rolling_midpoint(highs: &[f64], lows: &[f64], period: usize) -> Vec], period: usize) -> Vec> { let len = data.len(); let mut means = vec![None; len]; @@ -154,6 +154,97 @@ pub fn rolling_mean_strict(data: &[Option], period: usize) -> Vec], period: usize) -> Vec> { + let len = data.len(); + let mut result = vec![None; len]; + + if len < period || period == 0 { + return result; + } + + let weight_sum = (period * (period + 1)) as f64 / 2.0; + let mut sum = 0.0; + let mut weighted_sum = 0.0; + let mut valid_count = 0usize; + + for i in 0..period { + if let Some(value) = data[i] { + sum += value; + weighted_sum += value * (i + 1) as f64; + valid_count += 1; + } + } + + if valid_count == period { + result[period - 1] = Some(weighted_sum / weight_sum); + } + + for end in period..len { + let entering = data[end].unwrap_or(0.0); + let leaving = data[end - period].unwrap_or(0.0); + + weighted_sum = weighted_sum - sum + entering * period as f64; + sum += entering - leaving; + + if data[end].is_some() { + valid_count += 1; + } + if data[end - period].is_some() { + valid_count -= 1; + } + + if valid_count == period { + result[end] = Some(weighted_sum / weight_sum); + } + } + + result +} + +/// Computes rolling mean and standard deviation for fully valid windows only. +pub fn rolling_mean_stddev_strict( + data: &[Option], + period: usize, +) -> (Vec>, Vec>) { + let len = data.len(); + let mut means = vec![None; len]; + let mut stddevs = vec![None; len]; + + if len < period || period == 0 { + return (means, stddevs); + } + + let mut sum = 0.0; + let mut sum_sq = 0.0; + let mut valid_count = 0usize; + + for i in 0..len { + if let Some(value) = data[i] { + sum += value; + sum_sq += value * value; + valid_count += 1; + } + + if i >= period { + if let Some(value) = data[i - period] { + sum -= value; + sum_sq -= value * value; + valid_count -= 1; + } + } + + if i >= period - 1 && valid_count == period { + let mean = sum / period as f64; + let variance = (sum_sq / period as f64) - (mean * mean); + means[i] = Some(mean); + stddevs[i] = Some(variance.max(0.0).sqrt()); + } + } + + (means, stddevs) +} + fn push_max_index(deque: &mut VecDeque, data: &[f64], idx: usize) { if data[idx].is_nan() { return; @@ -393,6 +484,25 @@ mod tests { assert_eq!(means, vec![None, Some(2.0), None, None, Some(6.0)]); } + #[test] + fn test_rolling_weighted_mean_strict() { + let data = vec![None, Some(1.0), Some(2.0), None, Some(4.0)]; + + let means = rolling_weighted_mean_strict(&data, 2); + + assert_eq!(means, vec![None, None, Some(5.0 / 3.0), None, None]); + } + + #[test] + fn test_rolling_mean_stddev_strict() { + let data = vec![Some(1.0), Some(2.0), None, Some(4.0), Some(5.0)]; + + let (means, stddevs) = rolling_mean_stddev_strict(&data, 2); + + assert_eq!(means, vec![None, Some(1.5), None, None, Some(4.5)]); + assert_eq!(stddevs, vec![None, Some(0.5), None, None, Some(0.5)]); + } + #[test] fn test_forward_shift() { let values = vec![Some(1.0), None, Some(3.0)];