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()
<matplotlib.axes.AxesSubplot at 0xa80784c>
# 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.
[2012-10-26 15:19] INFO: Performance: Simulated 2267.0 trading days out of 2268. [2012-10-26 15:19] INFO: Performance: first open: 1993-01-04 14:30:00+00:00
print results_buy_apple
<class 'pandas.core.frame.DataFrame'> DatetimeIndex: 2267 entries, 1993-01-04 21:00:00 to 2001-12-28 21:00:00 Data columns: capital_used 2267 non-null values cumulative_capital_used 2267 non-null values ending_cash 2267 non-null values ending_value 2267 non-null values max_capital_used 2267 non-null values max_leverage 2267 non-null values period_close 2267 non-null values period_open 2267 non-null values pnl 2267 non-null values portfolio_value 2267 non-null values positions 2267 non-null values returns 2267 non-null values starting_cash 2267 non-null values starting_value 2267 non-null values transactions 2267 non-null values dtypes: float64(11), object(4)
results_buy_apple.portfolio_value.plot()
<matplotlib.axes.AxesSubplot at 0xaf2c20c>
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)
[2012-10-26 15:19] INFO: Transform: Running StatefulTransform [short_mavg] [2012-10-26 15:19] INFO: Transform: Running StatefulTransform [long_mavg]
1994-07-01 00:00:00+00:00: Buying 100 AAPL shares. 1996-01-09 00:00:00+00:00: Selling 100 AAPL shares. 1997-12-22 00:00:00+00:00: Buying 100 AAPL shares. 2000-09-13 00:00:00+00:00: Selling 100 AAPL shares. 2001-10-26 00:00:00+00:00: Buying 100 AAPL shares. 2001-11-26 00:00:00+00:00: Selling 100 AAPL shares.
[2012-10-26 15:19] INFO: Transform: Finished StatefulTransform [long_mavg] [2012-10-26 15:19] INFO: Transform: Finished StatefulTransform [short_mavg] [2012-10-26 15:19] INFO: Performance: Simulated 2267.0 trading days out of 2268. [2012-10-26 15:19] INFO: Performance: first open: 1993-01-04 14:30:00+00:00
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)
<matplotlib.axes.AxesSubplot at 0xbac262c>
data[['KO', 'PEP']].plot(); plt.ylabel('price')
<matplotlib.text.Text at 0xcbfc8ec>
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)
[2012-10-26 15:21] INFO: Performance: Simulated 2267.0 trading days out of 2268. [2012-10-26 15:21] INFO: Performance: first open: 1993-01-04 14:30:00+00:00
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()
<matplotlib.axes.AxesSubplot at 0xe83afec>