diff --git a/stockstats.py b/stockstats.py index f45372f..dafec71 100644 --- a/stockstats.py +++ b/stockstats.py @@ -362,6 +362,31 @@ def _get_rsi(cls, df, n_days): closenm_smma_column] cls._drop_columns(df, columns_to_remove) + @classmethod + def _get_stochrsi(cls, df, n_days): + """ Calculate the Stochastic RSI + + calculated based on the formula at: + https://www.investopedia.com/terms/s/stochrsi.asp + + :param df: data + :param n_days: N days + :return: None + """ + n_days = int(n_days) + column_name = 'stochrsi_{}'.format(n_days) + + cls._get_rsi(df, n_days) + + rsi = df['rsi_{}'.format(n_days)] + rsi_min = rsi.rolling( + min_periods=1, window=n_days, center=False).min() + rsi_max = rsi.rolling( + min_periods=1, window=n_days, center=False).max() + + cv = (rsi - rsi_min) / (rsi_max - rsi_min) + df[column_name] = cv.fillna(0).astype('float64') * 100 + @staticmethod def _drop_columns(df, columns): df.drop(columns, inplace=True, axis=1) @@ -372,6 +397,29 @@ def _ensure_type(self, obj): This patch is not the perfect way but could make the lib work. """ return obj + + @classmethod + def _get_wavetrend(cls, df, n1=10, n2=21): + """ Calculate LazyBear's Wavetrend + + Check the algorithm described below: + https://medium.com/@samuel.mcculloch/lets-take-a-look-at-wavetrend-with-crosses-lazybear-s-indicator-2ece1737f72f + :param df: data frame + :param n1: period of EMA on typical price + :param n2: period of EMA + :return: None + """ + df["tp"] = df["middle"] + df["esa"] = df["tp_{}_ema".format(n1)] + df["dd"] = np.abs(df["tp"]-df["esa"]) + df["d"] = df["dd_{}_ema".format(n1)] + df["ci"] = (df["tp"] - df["esa"]) / (0.015*df["d"]) + df["tci"] = df["ci_{}_ema".format(n2)] + df["wt1"] = df["tci"] + df["wt2"] = df["wt1_4_sma"] + cls._drop_columns(df, ["tp", "esa", "dd", "d", "ci", "tci", "middle", + "tp_{}_ema".format(n1), "dd_{}_ema".format(n1), + "ci_{}_ema".format(n2), "wt1_4_sma",]) @classmethod def _get_smma(cls, df, column, windows): @@ -1074,6 +1122,8 @@ def __init_not_exist_column(cls, df, key): cls._get_vr(df) elif key in ['dma']: cls._get_dma(df) + elif key in ['wt1', 'wt2']: + cls._get_wavetrend(df) elif key == 'log-ret': cls._get_log_ret(df) elif key.endswith('_delta'):