Implement SMA, price/SMA, and bollinger band

manual_strategy/indicators.py

@@ -11,7 +11,7 @@ def author():
     return "felixm"
 
 
-def test_policy():
+def test_optimal_strategy():
     symbol = "JPM"
     start_value = 100000
     sd = dt.datetime(2008, 1, 1)
@@ -55,7 +55,7 @@ def test_policy():
     portvals["Bench"] = bench["Portval"]
     portvals.drop(columns=["Portval"], inplace=True)
 
-    ax = portvals.plot(title="Optimal strategy versus 1000 shares of JPM")
+    portvals.plot(title="Optimal strategy versus 1000 shares of JPM")
     plt.savefig('figure_5.png')
 
 
@@ -63,23 +63,46 @@ def normalize(timeseries):
     return timeseries / timeseries.iloc[0]
 
 
-def bollinger_band(prices):
-    pass
+def bollinger_band(df, symbol, period=20, m=2):
+    boll_sma = df[symbol].rolling(period).mean()
+    std = df[symbol].rolling(period).std()
+    boll_up = boll_sma + m * std
+    boll_lo = boll_sma - m * std
+    df[f"Boll({symbol}, {period})-sma"] = boll_sma
+    df[f"Boll({symbol}, {period})-up"] = boll_up
+    df[f"Boll({symbol}, {period})-lo"] = boll_lo
+
+
+def sma(df, symbol, period):
+    """Adds a new column to the dataframe SMA(period)"""
+    df[f"{symbol}-sma({period})"] = df[symbol].rolling(period).mean()
+
+
+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()
+    df[f"{symbol}-price/sma({period})"] = df[symbol] / sma
 
 
 def main():
-    test_policy()
+    # test_optimal_strategy()
+    symbol = "JPM"
 
-    # plot_data(prices)
-    # sd = dt.datetime(2008, 1, 1)
-    # ed = dt.datetime(2009, 12, 31)
-    # prices = get_data(['JPM'], pd.date_range(sd, ed))
-    # prices['JPM'] = normalize(prices['JPM'])
-    # print(prices)
+    sd = dt.datetime(2008, 1, 1)
+    ed = dt.datetime(2009, 12, 31)
+    df = get_data([symbol], pd.date_range(sd, ed))
+    df.drop(columns=["SPY"], inplace=True)
+    df = normalize(df)
 
-    # plot_data(prices)
-    # prices_appl = get_data(['AAPL'], pd.date_range(sd, ed), 'High')
-    # prices['AAPL'] = prices_appl['AAPL']
+    sma(df, symbol, 8)
+    # price_sma(df, symbol, 8)
+    # bollinger_band(df, symbol)
+
+    # TODO
+    # rsi(df, symbol)
+    # macd(df, symbol)
+
+    plot_data(df)
 
 
 if __name__ == "__main__":