Add strategy RSI

msaltnet · Jul 10, 2022 · 59ef24c · 59ef24c
1 parent a12b565
commit 59ef24c
Show file tree

Hide file tree

Showing 10 changed files with 614 additions and 29 deletions.
diff --git a/smtm/__init__.py b/smtm/__init__.py
@@ -10,6 +10,7 @@
 from .simulation_operator import SimulationOperator
 from .strategy_bnh import StrategyBuyAndHold
 from .strategy_sma_0 import StrategySma0
+from .strategy_rsi import StrategyRsi
 from .virtual_market import VirtualMarket
 from .worker import Worker
 from .simulator import Simulator
@@ -33,6 +34,7 @@
     "SimulationOperator",
     "StrategyBuyAndHold",
     "StrategySma0",
+    "StrategyRsi",
     "VirtualMarket",
     "Worker",
     "Simulator",

diff --git a/smtm/__main__.py b/smtm/__main__.py
@@ -7,13 +7,14 @@
     4: mass simulation with config file
     5: make config file for mass simulation
 
-Example) python -m smtm --mode 0
-Example) python -m smtm --mode 1
-Example) python -m smtm --mode 1 --budget 500 --from_dash_to 201220.170000-201221 --term 1 --strategy 0 --currency BTC
-Example) python -m smtm --mode 2 --budget 50000 --term 60 --strategy 0 --currency ETH
-Example) python -m smtm --mode 3
-Example) python -m smtm --mode 4 --config /data/sma0_simulation.json
-Example) python -m smtm --mode 5 --budget 50000 --title SMA_2H_week --strategy 1 --currency ETH --from_dash_to 210804.000000-210811.000000 --offset 120 --file generated_config.json
+Example)
+python -m smtm --mode 0
+python -m smtm --mode 1
+python -m smtm --mode 1 --budget 500 --from_dash_to 201220.170000-201221 --term 1 --strategy 0 --currency BTC
+python -m smtm --mode 2 --budget 50000 --term 60 --strategy 0 --currency ETH
+python -m smtm --mode 3
+python -m smtm --mode 4 --config /data/sma0_simulation.json
+python -m smtm --mode 5 --budget 50000 --title SMA_2H_week --strategy 1 --currency ETH --from_dash_to 210804.000000-210811.000000 --offset 120 --file generated_config.json
 """
 import argparse
 from argparse import RawTextHelpFormatter
@@ -33,18 +34,19 @@
     4: mass simulation with config file
     5: make config file for mass simulation
 
-Example) python -m smtm --mode 0
-Example) python -m smtm --mode 1 --budget 50000 --from_dash_to 201220.170000-201221 --term 0.1 --strategy 0 --currency BTC
-Example) python -m smtm --mode 2 --budget 50000 --term 60 --strategy 0 --currency ETH
-Example) python -m smtm --mode 3
-Example) python -m smtm --mode 4 --config /data/sma0_simulation.json
-Example) python -m smtm --mode 5 --budget 50000 --title SMA_6H_week --strategy 1 --currency ETH --from_dash_to 210804.000000-210811.000000 --offset 360 --file generated_config.json
+Example)
+python -m smtm --mode 0
+python -m smtm --mode 1 --budget 50000 --from_dash_to 201220.170000-201221 --term 0.1 --strategy 0 --currency BTC
+python -m smtm --mode 2 --budget 50000 --term 60 --strategy 0 --currency ETH
+python -m smtm --mode 3
+python -m smtm --mode 4 --config /data/sma0_simulation.json
+python -m smtm --mode 5 --budget 50000 --title SMA_6H_week --strategy 1 --currency ETH --from_dash_to 210804.000000-210811.000000 --offset 360 --file generated_config.json
 """,
         formatter_class=RawTextHelpFormatter,
     )
     parser.add_argument("--budget", help="budget", type=int, default=10000)
     parser.add_argument("--term", help="trading tick interval (seconds)", type=float, default="60")
-    parser.add_argument("--strategy", help="strategy 0: buy and hold, 1: sma0", default="0")
+    parser.add_argument("--strategy", help="strategy 0: buy and hold, 1: sma0, 2: rsi", default="0")
     parser.add_argument("--trader", help="trader 0: Upbit, 1: Bithumb", default="0")
     parser.add_argument("--currency", help="trading currency e.g.BTC", default="BTC")
     parser.add_argument("--config", help="mass simulation config file", default="")

diff --git a/smtm/analyzer.py b/smtm/analyzer.py
@@ -312,6 +312,9 @@ def make_rsi(prices, count=14):
         http://stockcharts.com/school/doku.php?id=chart_school:glossary_r#relativestrengthindex
         http://www.investopedia.com/terms/r/rsi.asp
         """
+        if len(prices) <= count:
+            return None
+
         deltas = np.diff(prices)
         seed = deltas[:count]
         up_avg = seed[seed >= 0].sum() / count
@@ -654,11 +657,18 @@ def __draw_graph(
 
         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 rsi is not None:
+                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), secondary_y=False)
+                )
+                apds.append(
+                    mpf.make_addplot(rsi_low, panel=2, color="red", width=0.5, secondary_y=False)
+                )
+                apds.append(
+                    mpf.make_addplot(rsi_high, panel=2, color="red", width=0.5, secondary_y=False)
+                )
 
         if "buy" in total.columns:
             apds.append(mpf.make_addplot(total["buy"], type="scatter", markersize=100, marker="^"))

diff --git a/smtm/controller.py b/smtm/controller.py
@@ -10,6 +10,7 @@
     BithumbDataProvider,
     StrategyBuyAndHold,
     StrategySma0,
+    StrategyRsi,
     Operator,
 )
 
@@ -36,12 +37,19 @@ def __init__(
         self.command_list = []
         self.create_command()
         self.is_bithumb = is_bithumb
-        self.strategy = StrategySma0()
+        self.strategy = None
         self.currency = currency
         LogManager.set_stream_level(30)
 
-        if int(strategy) == 0:
+        strategy_num = int(strategy)
+        if strategy_num == 0:
             self.strategy = StrategyBuyAndHold()
+        elif strategy_num == 1:
+            self.strategy = StrategySma0()
+        elif strategy_num == 2:
+            self.strategy = StrategyRsi()
+        else:
+            raise UserWarning(f"Invalid Strategy! {self.strategy}")
 
     def create_command(self):
         """명령어 정보를 생성한다"""

diff --git a/smtm/mass_simulator.py b/smtm/mass_simulator.py
@@ -19,6 +19,7 @@
     SimulationDataProvider,
     StrategyBuyAndHold,
     StrategySma0,
+    StrategyRsi,
     SimulationOperator,
     SimulationTrader,
     Analyzer,
@@ -127,7 +128,16 @@ def get_initialized_operator(budget, strategy_num, interval, currency, start, en
         data_provider = SimulationDataProvider(currency=currency)
         data_provider.initialize_simulation(end=end, count=count)
 
-        strategy = StrategyBuyAndHold() if strategy_num == 0 else StrategySma0()
+        strategy_number = int(strategy_num)
+        if strategy_number == 0:
+            strategy = StrategyBuyAndHold()
+        elif strategy_number == 1:
+            strategy = StrategySma0()
+        elif strategy_number == 2:
+            strategy = StrategyRsi()
+        else:
+            raise UserWarning(f"Invalid Strategy! {strategy_number}")
+
         strategy.is_simulation = True
 
         trader = SimulationTrader(currency=currency)
@@ -250,7 +260,16 @@ def analyze_result(self, result_list, config):
         """수익률 비교 결과를 파일로 저장"""
         title = config["title"]
         period_list = config["period_list"]
-        strategy = StrategyBuyAndHold.NAME if config["strategy"] == 0 else StrategySma0.NAME
+
+        strategy_number = int(config["strategy"])
+        if strategy_number == 0:
+            strategy_name = StrategyBuyAndHold.NAME
+        elif strategy_number == 1:
+            strategy_name = StrategySma0.NAME
+        elif strategy_number == 2:
+            strategy_name = StrategyRsi.NAME
+        else:
+            raise UserWarning(f"Invalid Strategy! {strategy_number}")
 
         final_return_list = []
         min_return_list = []
@@ -287,7 +306,7 @@ def analyze_result(self, result_list, config):
             result_file.write(f"Title: {title}\n")
             result_file.write(f"Description: {config['description']}\n")
             result_file.write(
-                f"Strategy: {strategy}, Budget: {config['budget']}, Currency: {config['currency']}\n"
+                f"Strategy: {strategy_name}, Budget: {config['budget']}, Currency: {config['currency']}\n"
             )
             result_file.write(
                 f"{period_list[0]['start']} ~ {period_list[-1]['end']} ({len(final_return_list)})\n"
@@ -383,7 +402,7 @@ def make_config_json(
             "title": title,
             "description": "write descriptionf here",
             "budget": budget,
-            "strategy": strategy_num,
+            "strategy": int(strategy_num),
             "interval": interval,
             "currency": currency,
             "period_list": [],

diff --git a/smtm/simulator.py b/smtm/simulator.py
@@ -8,6 +8,7 @@
     SimulationDataProvider,
     StrategyBuyAndHold,
     StrategySma0,
+    StrategyRsi,
     SimulationOperator,
     DateConverter,
 )
@@ -123,7 +124,7 @@ def __init__(
                 "action": self._set_budget,
             },
             {
-                "guide": "전략 번호 입력. 0: Buy and Hold, 1: SMA-0",
+                "guide": "전략 번호 입력. 0: Buy and Hold, 1: SMA-0, 2: RSI",
                 "value": self.strategy,
                 "action": self._set_strategy,
             },
@@ -143,8 +144,12 @@ def initialize(self):
 
         if self.strategy == 0:
             strategy = StrategyBuyAndHold()
-        else:
+        elif self.strategy == 1:
             strategy = StrategySma0()
+        elif self.strategy == 2:
+            strategy = StrategyRsi()
+        else:
+            raise UserWarning(f"Invalid Strategy! {self.strategy}")
 
         strategy.is_simulation = True
         self.operator = SimulationOperator()