import pandas as pd import numpy as np import matplotlib.pyplot as plt plt.figsize(14,8) from zipline.algorithm import TradingAlgorithm from zipline.transforms import MovingAverage, batch_transform from zipline.utils.factory import load_from_yahoo data = pd.load('talk_px.dat') #data = load_from_yahoo(stocks=['AAPL', 'PEP', 'KO']); data.save('talk_px.dat') data['AAPL'].plot() # Simplest possible algorithm class BuyApple(TradingAlgorithm): # inherit from TradingAlgorithm def handle_data(self, data): # overload handle_data() method self.order('AAPL', 1) # stock (='AAPL') to order and amount (=1 shares) my_algo = BuyApple() # Instantiate our algorithm results_buy_apple = my_algo.run(data) # Backtest algorithm on dataframe. print results_buy_apple results_buy_apple.portfolio_value.plot() class DualMovingAverage(TradingAlgorithm): """Dual Moving Average Crossover algorithm. This algorithm buys apple once its short moving average crosses its long moving average (indicating upwards momentum) and sells its shares once the averages cross again (indicating downwards momentum). """ def initialize(self): # Add 2 mavg transforms, one with a long window, one # with a short window. self.add_transform(MovingAverage, 'short_mavg', ['price'], days=200) self.add_transform(MovingAverage, 'long_mavg', ['price'], days=400) # To keep track of whether we invested in the stock or not self.invested = False self.short_mavg = [] self.long_mavg = [] self.buy_orders = [] self.sell_orders = [] def handle_data(self, data): if (data['AAPL'].short_mavg['price'] > data['AAPL'].long_mavg['price']) and not self.invested: self.order('AAPL', 200) self.invested = True self.buy_orders.append(data['AAPL'].datetime) print "{dt}: Buying 100 AAPL shares.".format(dt=data['AAPL'].datetime) elif (data['AAPL'].short_mavg['price'] < data['AAPL'].long_mavg['price']) and self.invested: self.order('AAPL', -200) self.invested = False self.sell_orders.append(data['AAPL'].datetime) print "{dt}: Selling 100 AAPL shares.".format(dt=data['AAPL'].datetime) # Save mavgs for later analysis. self.short_mavg.append(data['AAPL'].short_mavg['price']) self.long_mavg.append(data['AAPL'].long_mavg['price']) dma = DualMovingAverage() results = dma.run(data) ax1 = plt.subplot(211) data['short'] = dma.short_mavg data['long'] = dma.long_mavg data[['AAPL', 'short', 'long']].plot(ax=ax1) plt.plot(dma.buy_orders, data['short'].ix[dma.buy_orders], '^', c='m', markersize=10, label='buy') plt.plot(dma.sell_orders, data['short'].ix[dma.sell_orders], 'v', c='k', markersize=10, label='sell') plt.legend(loc=0) ax2 = plt.subplot(212) results.portfolio_value.plot(ax=ax2) data[['KO', 'PEP']].plot(); plt.ylabel('price') import statsmodels.api as sm WINDOW_LENGTH = 100 @batch_transform def ols_transform(data): # receives constantly updated dataframe """Computes regression coefficient (slope and intercept) via Ordinary Least Squares between two SIDs. """ p0 = data.price['PEP'] p1 = sm.add_constant(data.price['KO'], prepend=False) slope, intercept = sm.OLS(p0, p1).fit().params return slope, intercept class Pairtrade(TradingAlgorithm): def initialize(self): self.ols_transform = ols_transform(refresh_period=WINDOW_LENGTH, days=WINDOW_LENGTH) self.spreads = [] self.zscores = [] self.invested = False self.buy_orders = [] self.sell_orders = [] def handle_data(self, data): ###################################################### # 1. Compute regression coefficients between PEP and KO params = self.ols_transform.handle_data(data) if params is None: # Still in the warm-up period? return slope, intercept = params ###################################################### # 2. Compute zscore of spread (remove mean and divide by std) zscore = self.compute_zscore(data, slope, intercept) self.zscores.append(zscore) ###################################################### # 3. Place orders self.place_orders(data, zscore) def compute_zscore(self, data, slope, intercept): """1. Compute the spread given slope and intercept. 2. zscore the spread. """ spread = (data['PEP'].price - (slope * data['KO'].price + intercept)) self.spreads.append(spread) zscore = (spread - np.mean(self.spreads[-WINDOW_LENGTH:])) / np.std(self.spreads[-WINDOW_LENGTH:]) return zscore def place_orders(self, data, zscore): """Buy spread if zscore is > 2, sell if zscore < .5. """ if zscore >= 2.0 and not self.invested: self.order('PEP', int(100 / data['PEP'].price)) self.order('KO', -int(100 / data['KO'].price)) self.invested = True self.buy_orders.append(data['PEP'].datetime) elif zscore <= -2.0 and not self.invested: self.order('KO', -int(100 / data['KO'].price)) self.order('PEP', int(100 / data['PEP'].price)) self.invested = True self.buy_orders.append(data['PEP'].datetime) elif abs(zscore) < .5 and self.invested: self.sell_spread() self.invested = False self.sell_orders.append(data['PEP'].datetime) def sell_spread(self): """ decrease exposure, regardless of position long/short. buy for a short position, sell for a long. """ ko_amount = self.portfolio.positions['KO'].amount self.order('KO', -1 * ko_amount) pep_amount = self.portfolio.positions['PEP'].amount self.order('KO', -1 * pep_amount) pairtrade = Pairtrade() stats = pairtrade.run(data) data['zscores'] = np.nan data.zscores[70:] = pairtrade.zscores ax1 = plt.subplot(311) data[['PEP', 'KO']].plot(ax=ax1) plt.ylabel('price') plt.setp(ax1.get_xticklabels(), visible=False) ax2 = plt.subplot(312, sharex=ax1) data.zscores.plot(ax=ax2, color='r') plt.plot(pairtrade.buy_orders, data.zscores.ix[pairtrade.buy_orders], '^', c='m', markersize=10, label='buy') plt.plot(pairtrade.sell_orders, data.zscores.ix[pairtrade.sell_orders], 'v', c='k', markersize=10, label='sell') plt.ylabel('zscore of spread') plt.setp(ax2.get_xticklabels(), visible=False) plt.legend() ax3 = plt.subplot(313, sharex=ax1) stats.portfolio_value.plot()