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()