Change indicators to return their results and work on three indicator strat

strategy_evaluation/ManualStrategy.py

@@ -57,6 +57,18 @@ class ManualStrategy:
 
         orders['Shares'] = orders['Shares'].rolling(2).apply(strat)
 
+    def three_indicator_strat(self, macd, rsi, price_sma, orders):
+        # XXX: continue here
+
+        def strat(row):
+            m = macd.loc[row.index]
+            print(m)
+            # mac, signal = m.iloc[0]
+            # print(mac, signal)
+            return 0
+
+        orders['Shares'] = orders.apply(strat, axis=1)
+        raise Exception()
 
     # this method should use the existing policy and test it against new data
     def testPolicy(self, symbol="IBM",
@@ -67,22 +79,18 @@ class ManualStrategy:
         self.holding = 0
         df = util.get_data([symbol], pd.date_range(sd, ed))
         df.drop(columns=["SPY"], inplace=True)
-        macd = indicators.macd(df, symbol)
 
         orders = pd.DataFrame(index=df.index)
         orders["Symbol"] = symbol
         orders["Order"] = ""
         orders["Shares"] = 0
 
-        self.macd_strat(macd, orders)
+        macd = indicators.macd(df, symbol)
+        rsi = indicators.rsi(df, symbol)
+        price_sma = indicators.price_sma(df, symbol, [21])
+        # self.macd_strat(macd, orders)
+        self.three_indicator_strat(macd, rsi, price_sma, orders)
 
-        # here we build a fake set of trades
+        # heers = orders[orders["Shares"] != 0]
-        # orders.iloc[0] = [symbol, "BUY", 1000]
-        # orders.iloc[40] = [symbol, "SELL", 1000]
-        # orders.iloc[41] = [symbol, "BUY", 1000]
-        # orders.iloc[60] = [symbol, "SELL", 2000]
-        # orders.iloc[61] = [symbol, "BUY", 2000]
-        # orders.iloc[-1] = [symbol, "SELL", 1000]
-        # orders = orders[orders["Shares"] != 0]
         return orders
 

strategy_evaluation/experiment1.py

@@ -27,16 +27,20 @@ def experiment1():
     df["Manual"] = marketsim.compute_portvals(orders, start_value)
     df["Orders Manual"] = orders["Shares"]
 
-    # indicators.price_sma(df, symbol, 21)
+    price_sma = indicators.price_sma(df, symbol, [21])
-    # sma = indicators.sma(df, symbol, [9, 21])
+    bb = indicators.bollinger_band(df, symbol)
-    # rsi = indicators.rsi(df, symbol)
+    sma = indicators.sma(df, symbol, [9, 21])
+    rsi = indicators.rsi(df, symbol)
     macd = indicators.macd(df, symbol).copy()
 
     fig, ax = plt.subplots(4, sharex=True)
-    df[symbol].plot(ax=ax[0])
+    df[[symbol]].plot(ax=ax[0])
-    macd.plot(ax=ax[3])
+    # bb.plot(ax=ax[0])
-    df[["Benchmark", "Manual"]].plot(ax=ax[1])
+    price_sma.plot(ax=ax[1])
-    df[["Orders Benchmark", "Orders Manual"]].plot(ax=ax[2])
+    macd.plot(ax=ax[2])
+    rsi.plot(ax=ax[3])
+    # df[["Benchmark", "Manual"]].plot(ax=ax[1])
+    # df[["Orders Benchmark", "Orders Manual"]].plot(ax=ax[2])
 
     for a in ax:
         a.grid()

strategy_evaluation/indicators.py

@@ -17,9 +17,13 @@ def bollinger_band(df, symbol, period=20, m=2):
     std = df[symbol].rolling(period).std()
     boll_up = boll_sma + m * std
     boll_lo = boll_sma - m * std
-    df[f"{symbol}-Boll({period})-sma"] = boll_sma
+    key_sma = f"{symbol}-Boll({period})-sma"
-    df[f"{symbol}-Boll({period})-up"] = boll_up
+    key_up = f"{symbol}-Boll({period})-up"
-    df[f"{symbol}-Boll({period})-lo"] = boll_lo
+    key_lo = f"{symbol}-Boll({period})-lo"
+    df[key_sma] = boll_sma
+    df[key_up] = boll_up
+    df[key_lo] = boll_lo
+    return df[[key_sma, key_up, key_lo]]
 
 
 def sma(df, symbol, period):
@@ -35,14 +39,28 @@ def sma(df, symbol, period):
 
 
 def ema(df, symbol, period):
-    """Adds a new column to the dataframe EMA(period)"""
+    """Adds EMA for one or multiple periods to df and returns EMAs"""
-    df[f"{symbol}-ema({period})"] = df[symbol].ewm(span=period).mean()
+    if type(period) is int:
+        period = [period]
+    keys = []
+    for p in period:
+        key = f"{symbol}-ema({p})"
+        df[key] = df[symbol].ewm(span=period).mean()
+        keys.append(key)
+    return df[keys]
 
 
 def price_sma(df, symbol, period):
     """Calculates SMA and adds new column price divided by SMA to the df."""
-    sma = df[symbol].rolling(period).mean()
+    if type(period) is int:
-    df[f"{symbol}-price/sma({period})"] = df[symbol] / sma
+        period = [period]
+    keys = []
+    for p in period:
+        key = f"{symbol}-price/sma({p})"
+        sma = df[symbol].rolling(p).mean()
+        df[key] = df[symbol] / sma
+        keys.append(key)
+    return df[keys]
 
 
 def rsi(df, symbol, period=14):
@@ -61,7 +79,7 @@ def rsi(df, symbol, period=14):
     # Add one to get 'period' price changes (first change is nan).
     period += 1
     df[key] = df[symbol].rolling(period).apply(rsi)
-    return df[key]
+    return df[[key]]
 
 
 def macd(df, symbol):