ta_lib_in_rust/indicators/stats/
beta.rs

1use crate::util::dataframe_utils::check_window_size;
2use polars::prelude::*;
3
4/// Calculates Beta - regression coefficient between two series
5///
6/// # Arguments
7///
8/// * `df` - DataFrame containing the price data
9/// * `price_column` - Column name for the price series (default "close")
10/// * `market_column` - Column name for the market/benchmark series
11/// * `window` - Window size for the calculation (typically 5)
12///
13/// # Returns
14///
15/// Returns a PolarsResult containing the Beta Series
16pub fn calculate_beta(
17    df: &DataFrame,
18    price_column: &str,
19    market_column: &str,
20    window: usize,
21) -> PolarsResult<Series> {
22    check_window_size(df, window, "Beta")?;
23
24    if !df.schema().contains(price_column) || !df.schema().contains(market_column) {
25        return Err(PolarsError::ComputeError(
26            format!("Beta calculation requires {price_column} and {market_column} columns").into(),
27        ));
28    }
29
30    let price = df.column(price_column)?.f64()?;
31    let market = df.column(market_column)?.f64()?;
32
33    let mut beta_values = Vec::with_capacity(df.height());
34
35    // Fill initial values with NaN
36    for _ in 0..window - 1 {
37        beta_values.push(f64::NAN);
38    }
39
40    // Calculate Beta for each window
41    for i in window - 1..df.height() {
42        let mut sum_xy = 0.0;
43        let mut sum_x = 0.0;
44        let mut sum_y = 0.0;
45        let mut sum_x2 = 0.0;
46        let mut count = 0;
47
48        for j in 0..window {
49            let x_idx = i - j;
50            let x = market.get(x_idx).unwrap_or(f64::NAN);
51            let y = price.get(x_idx).unwrap_or(f64::NAN);
52
53            if !x.is_nan() && !y.is_nan() {
54                sum_xy += x * y;
55                sum_x += x;
56                sum_y += y;
57                sum_x2 += x * x;
58                count += 1;
59            }
60        }
61
62        if count > 1 {
63            // Beta formula: (n*sum_xy - sum_x*sum_y) / (n*sum_x2 - sum_x^2)
64            let numerator = (count as f64 * sum_xy) - (sum_x * sum_y);
65            let denominator = (count as f64 * sum_x2) - (sum_x * sum_x);
66
67            if denominator != 0.0 {
68                beta_values.push(numerator / denominator);
69            } else {
70                beta_values.push(f64::NAN);
71            }
72        } else {
73            beta_values.push(f64::NAN);
74        }
75    }
76
77    Ok(Series::new("beta".into(), beta_values))
78}
ta_lib_in_rust/indicators/stats/beta.rs

ta_lib_in_rust/indicators/stats/
beta.rs
