In [1]:
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
In [3]:
data = pd.load('talk_px.dat') #data = load_from_yahoo(stocks=['AAPL', 'PEP', 'KO']); data.save('talk_px.dat')
In [4]:
data['AAPL'].plot()
Out[4]:
<matplotlib.axes.AxesSubplot at 0xa80784c>
In [5]:
# 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)
In [6]:
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

In [7]:
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)

In [8]:
results_buy_apple.portfolio_value.plot()
Out[8]:
<matplotlib.axes.AxesSubplot at 0xaf2c20c>
In [9]:
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'])
In [10]:
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



In [11]:
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)
Out[11]:
<matplotlib.axes.AxesSubplot at 0xbac262c>
In [12]:
data[['KO', 'PEP']].plot(); plt.ylabel('price')
Out[12]:
<matplotlib.text.Text at 0xcbfc8ec>
In [16]:
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)
In [17]:
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

In [18]:
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()
Out[18]:
<matplotlib.axes.AxesSubplot at 0xe83afec>
In [ ]: