Start working on strategy evaluation

strategy_evaluation/experiment1.py

@@ -0,0 +1,141 @@
+import pandas as pd
+import datetime as dt
+import marketsim.marketsim as marketsim
+import indicators
+import util
+import matplotlib.pyplot as plt
+from matplotlib.widgets import MultiCursor
+from BenchmarkStrategy import BenchmarkStrategy
+from ManualStrategy import ManualStrategy
+
+
+def macd_strat(macd):
+    def strat(x):
+        print(x)
+
+    macd['macd_trigger'] = macd.rolling(2).apply(strat)
+
+    # for i, row in macd.iterrows():
+        # if i == 0:
+            # continue
+        # print(row)
+        # prev_macd, prev_signal, _ = row
+        # cur_macd, cur_signal, _ = row
+        # if cur_macd < -.5 and (prev_macd < prev_signal) \
+                          # and (cur_macd > cur_signal):
+            # macd.iloc[i]['macd_buy_sell'] = 1
+
+def experiment1():
+    symbol = "JPM"
+    start_value = 10000
+    sd = dt.datetime(2008, 1, 1)
+    ed = dt.datetime(2009, 12, 31)
+    df = util.get_data([symbol], pd.date_range(sd, ed))
+    df.drop(columns=["SPY"], inplace=True)
+    # df = pd.DataFrame(index=df.index)
+
+    bs = BenchmarkStrategy()
+    orders = bs.testPolicy(symbol, sd, ed, start_value)
+    df["Benchmark"] = marketsim.compute_portvals(orders, start_value)
+
+    ms = ManualStrategy()
+    orders = ms.testPolicy(symbol, sd, ed, start_value)
+    df["Manual"] = marketsim.compute_portvals(orders, start_value)
+
+    # indicators.price_sma(df, symbol, 21)
+    # sma = indicators.sma(df, symbol, [9, 21])
+    # rsi = indicators.rsi(df, symbol)
+    macd = indicators.macd(df, symbol).copy()
+    # macd_strat(macd)
+
+    fig, ax = plt.subplots(2, sharex=True)
+    df[symbol].plot(ax=ax[0])
+    # sma.plot(ax=ax[0])
+    macd.plot(ax=ax[1])
+    # macd.iloc[:,0].plot(ax=ax[1])
+    # rsi.plot(ax=ax[2])
+    # df[["Benchmark", "Manual"]].plot(ax=ax[3])
+
+    # XXX: Plot where we buy and sell.
+
+    for a in ax: a.grid()
+    multi = MultiCursor(fig.canvas, ax, color='r', lw=0.5)
+    plt.show()
+
+    # df.plot(title="results", subplots=True)
+    #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_orig = df.copy()
+    #df = indicators.normalize(df)
+    #indicators.price_sma(df, symbol, 21)
+    #df.plot(title="21 SMA and EMA")
+    #plt.show()
+
+    # You may use data from other symbols (such as SPY) to inform both your
+    # Manual Learner and Strategy Learner. The in-sample/development period is
+    # January 1, 2008 to December 31 2009. The out-of-sample/testing period is
+    # January 1, 2010 to December 31 2011.
+
+
+class BlittedCursor:
+    """
+    A cross hair cursor using blitting for faster redraw.
+    """
+    def __init__(self, ax):
+        self.ax = ax
+        self.background = None
+        self.horizontal_line = ax.axhline(color='k', lw=0.8, ls='--')
+        self.vertical_line = ax.axvline(color='k', lw=0.8, ls='--')
+        # text location in axes coordinates
+        self.text = ax.text(0.72, 0.9, '', transform=ax.transAxes)
+        self._creating_background = False
+        ax.figure.canvas.mpl_connect('draw_event', self.on_draw)
+
+    def on_draw(self, event):
+        self.create_new_background()
+
+    def set_cross_hair_visible(self, visible):
+        need_redraw = self.horizontal_line.get_visible() != visible
+        self.horizontal_line.set_visible(visible)
+        self.vertical_line.set_visible(visible)
+        self.text.set_visible(visible)
+        return need_redraw
+
+    def create_new_background(self):
+        if self._creating_background:
+            # discard calls triggered from within this function
+            return
+        self._creating_background = True
+        self.set_cross_hair_visible(False)
+        self.ax.figure.canvas.draw()
+        self.background = self.ax.figure.canvas.copy_from_bbox(self.ax.bbox)
+        self.set_cross_hair_visible(True)
+        self._creating_background = False
+
+    def on_mouse_move(self, event):
+        if self.background is None:
+            self.create_new_background()
+        if not event.inaxes:
+            need_redraw = self.set_cross_hair_visible(False)
+            if need_redraw:
+                self.ax.figure.canvas.restore_region(self.background)
+                self.ax.figure.canvas.blit(self.ax.bbox)
+        else:
+            self.set_cross_hair_visible(True)
+            # update the line positions
+            x, y = event.xdata, event.ydata
+            self.horizontal_line.set_ydata(y)
+            self.vertical_line.set_xdata(x)
+            self.text.set_text('x=%1.2f, y=%1.2f' % (x, y))
+
+            self.ax.figure.canvas.restore_region(self.background)
+            self.ax.draw_artist(self.horizontal_line)
+            self.ax.draw_artist(self.vertical_line)
+            self.ax.draw_artist(self.text)
+            self.ax.figure.canvas.blit(self.ax.bbox)
+
+
+if __name__ == "__main__":
+    experiment1()