rustalib 1.0.9

A library of technical indicators for financial analysis, similar to TA-Lib
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

use polars::prelude::*;

/// Calculates Relative Strength Index (RSI)
///
/// # Arguments
///
/// * `df` - DataFrame containing price data
/// * `window` - RSI calculation period (typically 14)
/// * `column` - Column name to use for calculations (default "close")
///
/// # Returns
///
/// * `PolarsResult<Series>` - RSI values as a Series
pub fn calculate_rsi(df: &DataFrame, window: usize, column: &str) -> PolarsResult<Series> {
    // Check we have enough data
    if df.height() < window + 1 {
        return Err(PolarsError::ComputeError(
            format!(
                "Not enough data points for RSI calculation with window size {}",
                window
            )
            .into(),
        ));
    }

    // Get price data
    let close = df.column(column)?.f64()?.clone().into_series();

    // Calculate price changes
    let prev_close = close.shift(1);
    let price_diff: Vec<f64> = close
        .f64()?
        .iter()
        .zip(prev_close.f64()?.iter())
        .map(|(curr, prev)| match (curr, prev) {
            (Some(c), Some(p)) => c - p,
            _ => f64::NAN,
        })
        .collect();

    // Separate gains and losses
    let mut gains: Vec<f64> = Vec::with_capacity(df.height());
    let mut losses: Vec<f64> = Vec::with_capacity(df.height());

    // First value is NaN (no previous value to compare)
    gains.push(0.0);
    losses.push(0.0);

    for &diff in &price_diff[1..] {
        if diff.is_nan() {
            gains.push(f64::NAN);
            losses.push(f64::NAN);
        } else if diff > 0.0 {
            gains.push(diff);
            losses.push(0.0);
        } else {
            gains.push(0.0);
            losses.push(diff.abs());
        }
    }

    // Calculate RSI using Wilder's smoothing method
    let mut avg_gain = 0.0;
    let mut avg_loss = 0.0;
    let mut rsi: Vec<f64> = Vec::with_capacity(df.height());

    // Fill initial values with NaN
    for _i in 0..window {
        rsi.push(f64::NAN);
    }

    // First average gain/loss is a simple average
    for i in 1..=window {
        avg_gain += gains[i];
        avg_loss += losses[i];
    }
    avg_gain /= window as f64;
    avg_loss /= window as f64;

    // First RSI value
    let rs = if avg_loss == 0.0 {
        100.0 // Prevent division by zero
    } else {
        avg_gain / avg_loss
    };
    let rsi_val = 100.0 - (100.0 / (1.0 + rs));
    rsi[window - 1] = rsi_val;

    // Calculate smoothed RSI for the rest of the series
    for i in window + 1..df.height() {
        // Update using Wilder's smoothing
        avg_gain = ((avg_gain * (window - 1) as f64) + gains[i]) / window as f64;
        avg_loss = ((avg_loss * (window - 1) as f64) + losses[i]) / window as f64;

        // Calculate RSI
        let rs = if avg_loss == 0.0 {
            100.0 // Prevent division by zero
        } else {
            avg_gain / avg_loss
        };
        let rsi_val = 100.0 - (100.0 / (1.0 + rs));
        rsi.push(rsi_val);
    }

    Ok(Series::new(format!("rsi_{}", window).into(), rsi))
}