can draw rsi index graph to 2nd panel

smtm/analyzer.py

@@ -9,6 +9,7 @@
 import pandas as pd
 import mplfinance as mpf
 import psutil
+import numpy as np
 from .log_manager import LogManager
 
 matplotlib.use("Agg")
@@ -30,6 +31,7 @@ class Analyzer:
     RECORD_INTERVAL = 60
     GRAPH_MAX_COUNT = 1440
     DEBUG_MODE = True
+    RSI = None  # set (low, high, count) tuple to draw e.g. (30, 70, 14)
 
     def __init__(self, sma_info=(10, 40, 60)):
         self.request_list = []
@@ -303,6 +305,39 @@ def get_return_report(self, graph_filename=None, index_info=None):
             spot_list=spot_list,
         )
 
+    @staticmethod
+    def make_rsi(prices, count=14):
+        """
+        compute the n period relative strength indicator
+        http://stockcharts.com/school/doku.php?id=chart_school:glossary_r#relativestrengthindex
+        http://www.investopedia.com/terms/r/rsi.asp
+        """
+        deltas = np.diff(prices)
+        seed = deltas[:count]
+        up_avg = seed[seed >= 0].sum() / count
+        down_avg = -seed[seed < 0].sum() / count
+        r_strength = up_avg / down_avg
+        rsi = np.zeros_like(prices)
+        rsi[: count + 1] = 100.0 - 100.0 / (1.0 + r_strength)
+
+        for i in range(count + 1, len(prices)):
+            delta = deltas[i - 1]  # cause the diff is 1 shorter
+
+            if delta > 0:
+                upval = delta
+                downval = 0.0
+            else:
+                upval = 0.0
+                downval = -delta
+
+            up_avg = (up_avg * (count - 1) + upval) / count
+            down_avg = (down_avg * (count - 1) + downval) / count
+
+            r_strength = up_avg / down_avg
+            rsi[i] = 100.0 - 100.0 / (1.0 + r_strength)
+
+        return rsi
+
     def __make_interval_data(self, index_info):
         period = index_info[0]
         index = index_info[1]
@@ -617,6 +652,14 @@ def __draw_graph(
         total.index = pd.to_datetime(total.index)
         apds = []
 
+        if self.RSI is not None:
+            rsi = self.make_rsi(total["Close"], count=self.RSI[2])
+            rsi_low = np.full(len(rsi), self.RSI[0])
+            rsi_high = np.full(len(rsi), self.RSI[1])
+            apds.append(mpf.make_addplot(rsi, panel=2, color="lime", ylim=(10, 90)))
+            apds.append(mpf.make_addplot(rsi_low, panel=2, color="red", width=0.5))
+            apds.append(mpf.make_addplot(rsi_high, panel=2, color="red", width=0.5))
+
         if "buy" in total.columns:
             apds.append(mpf.make_addplot(total["buy"], type="scatter", markersize=100, marker="^"))
         if "sell" in total.columns:

tests/analyzer_test.py

@@ -957,7 +957,7 @@ def test_create_report_draw_correct_graph(self, mock_file, mock_plot, mock_make_
         analyzer.update_asset_info = MagicMock()
         filename = "apple"
 
-        dummy_data = self.fill_test_data_for_report(analyzer)
+        self.fill_test_data_for_report(analyzer)
         analyzer.create_report(filename)
 
         self.assertEqual(
@@ -990,6 +990,91 @@ def test_create_report_draw_correct_graph(self, mock_file, mock_plot, mock_make_
         analyzer.update_asset_info.assert_called_once()
         analyzer._get_rss_memory.assert_called_once()
 
+    @patch("mplfinance.make_addplot")
+    @patch("mplfinance.plot")
+    @patch("builtins.open", new_callable=mock_open)
+    def test_create_report_draw_correct_graph_when_rsi_enabled(
+        self, mock_file, mock_plot, mock_make_addplot
+    ):
+        analyzer = Analyzer()
+        analyzer.RSI = (30, 70, 7)
+        analyzer._get_rss_memory = MagicMock(return_value=123.45678)
+        analyzer.initialize("mango")
+        analyzer.update_asset_info = MagicMock()
+        filename = "apple"
+
+        self.fill_test_data_for_report(analyzer)
+        analyzer.create_report(filename)
+
+        self.assertEqual(
+            mock_make_addplot.call_args_list[0][1],
+            {"panel": 2, "color": "lime", "ylim": (10, 90)},
+        )
+
+        self.assertEqual(
+            mock_make_addplot.call_args_list[1][1],
+            {"panel": 2, "color": "red", "width": 0.5},
+        )
+
+        self.assertEqual(
+            mock_make_addplot.call_args_list[2][1],
+            {"panel": 2, "color": "red", "width": 0.5},
+        )
+
+        mock_plot.assert_called_once_with(
+            ANY,
+            type="candle",
+            volume=True,
+            addplot=ANY,
+            mav=analyzer.sma_info,
+            style="starsandstripes",
+            savefig=dict(
+                fname=analyzer.OUTPUT_FOLDER + filename + ".jpg", dpi=300, pad_inches=0.25
+            ),
+            figscale=1.25,
+        )
+        analyzer.update_asset_info.assert_called_once()
+        analyzer._get_rss_memory.assert_called_once()
+
+    def test_make_rsi_return_correct_rsi(self):
+        prices = [
+            26026000.0,
+            26075000.0,
+            26051000.0,
+            26039000.0,
+            26007000.0,
+            25981000.0,
+            26004000.0,
+            26020000.0,
+            25997000.0,
+            25981000.0,
+            26002000.0,
+            26002000.0,
+            25970000.0,
+            25940000.0,
+            25924000.0,
+        ]
+        expected = [
+            34.27,
+            34.27,
+            34.27,
+            34.27,
+            34.27,
+            34.27,
+            45.27,
+            52.22,
+            42.52,
+            36.6,
+            48.38,
+            48.38,
+            33.54,
+            24.67,
+            20.97,
+        ]
+        rsi = Analyzer.make_rsi(prices, count=5)
+        for i in range(len(rsi)):
+            self.assertAlmostEqual(rsi[i], expected[i], 2)
+
     def test_get_trading_results_return_result(self):
         analyzer = Analyzer()
         analyzer.result_list = "mango"