ta_lib_in_rust/indicators/moving_averages/
ema.rs

1use crate::util::dataframe_utils::check_window_size;
2use polars::prelude::*;
3
4/// Calculates Exponential Moving Average (EMA)
5///
6/// # Arguments
7///
8/// * `df` - DataFrame containing the input data
9/// * `column` - Column name to calculate EMA on
10/// * `window` - Window size for the EMA
11///
12/// # Returns
13///
14/// Returns a PolarsResult containing the EMA Series
15pub fn calculate_ema(df: &DataFrame, column: &str, window: usize) -> PolarsResult<Series> {
16    // Check we have enough data
17    check_window_size(df, window, "EMA")?;
18
19    let series = df.column(column)?.f64()?.clone().into_series();
20    let series_ca = series.f64()?;
21    let alpha = 2.0 / (window as f64 + 1.0);
22
23    let mut ema_values = Vec::with_capacity(series.len());
24
25    // Initialize with SMA for first window points
26    let mut sma_sum = 0.0;
27    for i in 0..window {
28        let val = series_ca.get(i).unwrap_or(0.0);
29        sma_sum += val;
30
31        // Fill with nulls until we have enough data
32        if i < window - 1 {
33            ema_values.push(f64::NAN);
34        }
35    }
36
37    // Add the initial SMA value
38    let initial_ema = sma_sum / window as f64;
39    ema_values.push(initial_ema);
40
41    // Calculate EMA using the recursive formula
42    let mut prev_ema = initial_ema;
43    for i in window..series.len() {
44        let price = series_ca.get(i).unwrap_or(0.0);
45        let ema = alpha * price + (1.0 - alpha) * prev_ema;
46        ema_values.push(ema);
47        prev_ema = ema;
48    }
49
50    Ok(Series::new("ema".into(), ema_values))
51}
ta_lib_in_rust/indicators/moving_averages/ema.rs

ta_lib_in_rust/indicators/moving_averages/
ema.rs
