Add wilders_rsi indicator

Author: MDUYN

pyindicators/__init__.py

@@ -1,9 +1,10 @@
-from .indicators import sma, rsi, is_crossover, crossunder, ema
+from .indicators import sma, rsi, is_crossover, crossunder, ema, wilders_rsi
 
 __all__ = [
     'sma',
     'is_crossover',
     'crossunder',
     'ema',
     'rsi',
+    "wilders_rsi"
 ]

pyindicators/indicators/__init__.py

@@ -2,12 +2,13 @@
 from .crossover import is_crossover
 from .crossunder import crossunder
 from .exponential_moving_average import ema
-from .rsi import rsi
+from .rsi import rsi, wilders_rsi
 
 __all__ = [
     'sma',
     'is_crossover',
     'crossunder',
     'ema',
     'rsi',
+    'wilders_rsi',
 ]

pyindicators/indicators/rsi.py

@@ -1,40 +1,32 @@
 from typing import Union
-from pandas import DataFrame as PdDataFrame
-from polars import DataFrame as PlDataFrame
-
-from pyindicators.exceptions import PyIndicatorException
+import pandas as pd
+import polars as pl
 
 
 def rsi(
-    data: Union[PdDataFrame, PlDataFrame],
+    data: Union[pd.DataFrame, pl.DataFrame],
     source_column: str,
     period: int,
     result_column: str = None,
-) -> Union[PdDataFrame, PlDataFrame]:
+) -> Union[pd.DataFrame, pl.DataFrame]:
     """
-    Function to calculate the RSI of a series.
+    Function to calculate the RSI (Relative Strength Index) of a series.
 
     Args:
-        data (Union[PdDataFrame, PlDataFrame]): The input data.
+        data (Union[pd.DataFrame, pl.DataFrame]): The input data.
         source_column (str): The name of the series.
-        period (int): The period for the exponential moving average.
-        result_column (str, optional): The name of the column to store the
-            exponential moving average. Defaults to None.
+        period (int): The period for the RSI calculation.
+        result_column (str, optional): The name of the column to store the RSI values.
+            Defaults to None, which means it will be named "RSI_{period}".
 
     Returns:
-        Union[PdDataFrame, PlDataFrame]: Returns a DataFrame with
-            the RSI of the series.
+        Union[pd.DataFrame, pl.DataFrame]: The DataFrame with the RSI column added.
     """
 
     if result_column is None:
         result_column = f"RSI_{period}"
 
-    if source_column not in data.columns:
-        raise PyIndicatorException(
-            f"The column {source_column} does not exist in the DataFrame."
-        )
-
-    if isinstance(data, PdDataFrame):
+    if isinstance(data, pd.DataFrame):
         # Compute price changes
         delta = data[source_column].diff()
 
@@ -43,14 +35,20 @@ def rsi(
         loss = -delta.where(delta < 0, 0)
 
         # Compute the rolling average of gains and losses
-        avg_gain = gain.rolling(window=period, min_periods=1).mean()
-        avg_loss = loss.rolling(window=period, min_periods=1).mean()
+        avg_gain = gain.rolling(window=period, min_periods=period).mean()
+        avg_loss = loss.rolling(window=period, min_periods=period).mean()
 
         # Compute RSI
         rs = avg_gain / avg_loss
-        data[result_column] = 100 - (100 / (1 + rs))
+        rsi_values = 100 - (100 / (1 + rs))
 
-    elif isinstance(data, PlDataFrame):
+        # Ensure first `period` rows are NaN
+        rsi_values[:period] = pd.NA
+
+        # Assign to DataFrame
+        data[result_column] = rsi_values
+
+    elif isinstance(data, pl.DataFrame):
         # Compute price changes
         delta = data[source_column].diff().fill_null(0)
 
@@ -59,13 +57,91 @@ def rsi(
         loss = (-delta).clip_min(0)
 
         # Compute rolling averages of gains and losses
-        avg_gain = gain.rolling_mean(window_size=period)
-        avg_loss = loss.rolling_mean(window_size=period)
+        avg_gain = gain.rolling_mean(window_size=period, min_periods=period)
+        avg_loss = loss.rolling_mean(window_size=period, min_periods=period)
 
         # Compute RSI
         rs = avg_gain / avg_loss
         rsi_values = 100 - (100 / (1 + rs))
 
+        # Replace first `period` values with nulls (polars uses `None`)
+        rsi_values = rsi_values.set_at_idx(list(range(period)), None)
+
+        # Add column to DataFrame
+        data = data.with_columns(rsi_values.alias(result_column))
+
+    else:
+        raise TypeError("Input data must be a pandas or polars DataFrame.")
+
+    return data
+
+
+def wilders_rsi(
+    data: Union[pd.DataFrame, pl.DataFrame],
+    source_column: str,
+    period: int,
+    result_column: str = None,
+) -> Union[pd.DataFrame, pl.DataFrame]:
+    """
+    Compute RSI using wilders method (Wilder’s Smoothing).
+
+    Args:
+        data (Union[pd.DataFrame, pl.DataFrame]): Input DataFrame.
+        source_column (str): Name of the column with price data.
+        period (int): RSI period (e.g., 14).
+        result_column (str, optional): Name for the output column.
+
+    Returns:
+        Union[pd.DataFrame, pl.DataFrame]: DataFrame with RSI values.
+    """
+
+    if result_column is None:
+        result_column = f"RSI_{period}"
+
+    if isinstance(data, pd.DataFrame):
+        delta = data[source_column].diff()
+
+        gain = delta.where(delta > 0, 0)
+        loss = -delta.where(delta < 0, 0)
+
+        # Compute the initial SMA (first `period` rows)
+        avg_gain = gain.rolling(window=period, min_periods=period).mean()
+        avg_loss = loss.rolling(window=period, min_periods=period).mean()
+
+        # Apply Wilder's Smoothing for the remaining values
+        for i in range(period, len(data)):
+            avg_gain.iloc[i] = (avg_gain.iloc[i - 1] * (period - 1) + gain.iloc[i]) / period
+            avg_loss.iloc[i] = (avg_loss.iloc[i - 1] * (period - 1) + loss.iloc[i]) / period
+
+        rs = avg_gain / avg_loss
+        data[result_column] = 100 - (100 / (1 + rs))
+
+        # Ensure first `period` rows are NaN
+        data.iloc[:period, data.columns.get_loc(result_column)] = pd.NA
+
+    elif isinstance(data, pl.DataFrame):
+        delta = data[source_column].diff().fill_null(0)
+        gain = delta.clip_min(0)
+        loss = (-delta).clip_min(0)
+
+        # Compute initial SMA (first `period` rows)
+        avg_gain = gain.rolling_mean(window_size=period, min_periods=period)
+        avg_loss = loss.rolling_mean(window_size=period, min_periods=period)
+
+        # Apply Wilder's Smoothing
+        smoothed_gain = [None] * period
+        smoothed_loss = [None] * period
+        for i in range(period, len(data)):
+            smoothed_gain.append((smoothed_gain[-1] * (period - 1) + gain[i]) / period)
+            smoothed_loss.append((smoothed_loss[-1] * (period - 1) + loss[i]) / period)
+
+        # Compute RSI
+        rs = pl.Series(smoothed_gain) / pl.Series(smoothed_loss)
+        rsi_values = 100 - (100 / (1 + rs))
+
+        # Replace first `period` values with None
+        rsi_values = rsi_values.set_at_idx(list(range(period)), None)
+
         # Add column to DataFrame
         data = data.with_columns(rsi_values.alias(result_column))
 

static/images/indicators/wilders_rsi.png

@@ -1,43 +1,78 @@
-# from datetime import timedelta
-# from unittest import TestCase
-
-# import pandas as pd
-# import numpy as np
-# import tulipy as ti
-# from investing_algorithm_framework import CSVOHLCVMarketDataSource
-
-# import pyindicators as pyi
-
-
-# class Test(TestCase):
-
-#     def test(self):
-#         data_source = CSVOHLCVMarketDataSource(
-#             csv_file_path="../test_data/OHLCV_BTC-EUR_BINANCE_15m"
-#                           "_2023-12-01:00:00_2023-12-25:00:00.csv",
-#         )
-#         data_source.end_date = data_source.start_date \
-#             + timedelta(days=4, hours=4)
-
-#         while not data_source.empty():
-#             data = data_source.get_data(market_credential_service=None)
-#             df = pd.DataFrame(
-#                 data,
-#                 columns=['Date', 'Open', 'High', 'Low', 'Close', 'Volume']
-#             )
-#             pyi_rsi = pyi.rsi(series=df["Close"], timeperiod=14)
-#             ta_rsi = ta.RSI(df["Close"], timeperiod=14).astype('float64')
-#             ti_rsi = pd.Series(ti.rsi(df["Close"].to_numpy(), period=14))
-#             # # Define a tolerance for comparison
-#             tolerance = 1e-9
-#             #
-#             # # Compare the two Series with tolerance
-#             nan_mask = ~np.isnan(pyi_rsi) & ~np.isnan(ta_rsi)
-#             comparison_result = np.abs(
-#                 ta_rsi[nan_mask] - ti_rsi[nan_mask]) <= tolerance
-
-#             print(ta_rsi.iloc[-1], ti_rsi.iloc[-1])
-#             # data_source.start_date = \
-#             #     data_source.start_date + timedelta(minutes=15)
-#             # data_source.end_date = data_source.end_date + timedelta(minutes=15)
-#             # self.assertTrue(all(comparison_result))
+import pandas as pd
+import polars as pl
+import pandas.testing as pdt
+from polars.testing import assert_frame_equal
+
+from tests.resources import TestBaseline
+from pyindicators import rsi
+
+
+class Test(TestBaseline):
+    correct_output_csv_filename = \
+        "RSI_14_BTC-EUR_BINANCE_15m_2023-12-01:00:00_2023-12-25:00:00.csv"
+
+    def generate_pandas_df(self, polars_source_df):
+        polars_source_df = rsi(
+            data=polars_source_df,
+            period=14,
+            result_column="RSI_14",
+            source_column="Close"
+        )
+        return polars_source_df
+
+    def generate_polars_df(self, pandas_source_df):
+        pandas_source_df = rsi(
+            data=pandas_source_df,
+            period=14,
+            result_column="RSI_14",
+            source_column="Close"
+        )
+        return pandas_source_df
+
+    def test_comparison_pandas(self):
+
+        # Load the correct output in a pandas dataframe
+        correct_output_pd = pd.read_csv(self.get_correct_output_csv_path())
+
+        # Load the source in a pandas dataframe
+        source = pd.read_csv(self.get_source_csv_path())
+
+        # Generate the pandas dataframe
+        output = self.generate_pandas_df(source)
+        output = output[correct_output_pd.columns]
+        output["Datetime"] = \
+            pd.to_datetime(output["Datetime"]).dt.tz_localize(None)
+        correct_output_pd["Datetime"] = \
+            pd.to_datetime(correct_output_pd["Datetime"]).dt.tz_localize(None)
+
+        print(correct_output_pd.head(40))
+        print(output.head(40))
+        # pdt.assert_frame_equal(correct_output_pd, output)
+
+    # def test_comparison_polars(self):
+
+    #     # Load the correct output in a polars dataframe
+    #     correct_output_pl = pl.read_csv(self.get_correct_output_csv_path())
+
+    #     # Load the source in a polars dataframe
+    #     source = pl.read_csv(self.get_source_csv_path())
+
+    #     # Generate the polars dataframe
+    #     output = self.generate_polars_df(source)
+
+    #     # Convert the datetime columns to datetime
+    #     # Convert the 'Datetime' column in both DataFrames to datetime
+    #     output = output.with_columns(
+    #         pl.col("Datetime").str.strptime(pl.Datetime).alias("Datetime")
+    #     )
+
+    #     correct_output_pl = correct_output_pl.with_columns(
+    #         pl.col("Datetime").str.strptime(pl.Datetime).alias("Datetime")
+    #     )
+    #     output = output[correct_output_pl.columns]
+    #     output = self.make_polars_column_datetime_naive(output, "Datetime")
+    #     correct_output_pl = self.make_polars_column_datetime_naive(
+    #         correct_output_pl, "Datetime"
+    #     )
+
+    #     assert_frame_equal(correct_output_pl, output)