HB : Add StochRSI indicator

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'):