ta_lib_in_rust/indicators/moving_averages/
wma.rs

1use crate::util::dataframe_utils::check_window_size;
2use polars::prelude::*;
3
4/// Calculates Weighted Moving Average (WMA)
5///
6/// # Arguments
7///
8/// * `df` - DataFrame containing the input data
9/// * `column` - Column name to calculate WMA on
10/// * `window` - Window size for the WMA
11///
12/// # Returns
13///
14/// Returns a PolarsResult containing the WMA Series
15pub fn calculate_wma(df: &DataFrame, column: &str, window: usize) -> PolarsResult<Series> {
16    // Check we have enough data
17    check_window_size(df, window, "WMA")?;
18
19    let series = df.column(column)?.f64()?.clone().into_series();
20
21    // Create linear weights [1, 2, 3, ..., window]
22    let weights: Vec<f64> = (1..=window).map(|i| i as f64).collect();
23
24    // Calculate WMA using rolling_mean with weights
25    series.rolling_mean(RollingOptionsFixedWindow {
26        window_size: window,
27        min_periods: window,
28        center: false,
29        weights: Some(weights),
30        fn_params: None,
31    })
32}