Adds example algorithm scripts.

README.md

@@ -104,6 +104,8 @@ dma = DualMovingAverage()
 results = dma.run(data)
 ```
 
+You can find other examples in the zipline/examples directory.
+
 Style Guide
 ===========
 

zipline/examples/buyapple.py

@@ -0,0 +1,36 @@
+#!/usr/bin/python
+#
+# Copyright 2012 Quantopian, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import matplotlib.pyplot as plt
+
+from zipline.algorithm import TradingAlgorithm
+from zipline.utils.factory import load_from_yahoo
+
+
+class BuyApple(TradingAlgorithm):  # inherit from TradingAlgorithm
+    """This is the simplest possible algorithm that does nothing but
+    buy 1 apple share on each event.
+    """
+    def handle_data(self, data):  # overload handle_data() method
+        self.order('AAPL', 1)  # order SID (=0) and amount (=1 shares)
+
+
+if __name__ == '__main__':
+    data = load_from_yahoo(stocks=['AAPL'], indexes={})
+    simple_algo = BuyApple()
+    results = simple_algo.run(data)
+    results.portfolio_value.plot()
+    plt.show()

zipline/examples/dual_moving_average.py

@@ -0,0 +1,74 @@
+#!/usr/bin/python
+#
+# Copyright 2012 Quantopian, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import matplotlib.pyplot as plt
+
+from zipline.algorithm import TradingAlgorithm
+from zipline.transforms import MovingAverage
+from zipline.utils.factory import load_from_yahoo
+
+
+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, short_window=200, long_window=400):
+        # Add 2 mavg transforms, one with a long window, one
+        # with a short window.
+        self.add_transform(MovingAverage, 'short_mavg', ['price'],
+                           days=short_window)
+
+        self.add_transform(MovingAverage, 'long_mavg', ['price'],
+                           days=long_window)
+
+        # To keep track of whether we invested in the stock or not
+        self.invested = False
+
+        self.short_mavgs = []
+        self.long_mavgs = []
+
+    def handle_data(self, data):
+        short_mavg = data['AAPL'].short_mavg['price']
+        long_mavg = data['AAPL'].long_mavg['price']
+        if short_mavg > long_mavg and not self.invested:
+            self.order('AAPL', 100)
+            self.invested = True
+        elif short_mavg < long_mavg and self.invested:
+            self.order('AAPL', -100)
+            self.invested = False
+
+        # Save mavgs for later analysis.
+        self.short_mavgs.append(short_mavg)
+        self.long_mavgs.append(long_mavg)
+
+
+if __name__ == '__main__':
+    data = load_from_yahoo(stocks=['AAPL'], indexes={})
+    dma = DualMovingAverage()
+    results = dma.run(data)
+
+    results.portfolio_value.plot()
+
+    data['short'] = dma.short_mavgs
+    data['long'] = dma.long_mavgs
+    data[['AAPL', 'short', 'long']].plot()
+    plt.legend(loc=0)
+    plt.show()

zipline/examples/pairtrade.py

@@ -0,0 +1,131 @@
+#!/usr/bin/python
+#
+# Copyright 2012 Quantopian, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import matplotlib.pyplot as plt
+import numpy as np
+import statsmodels.api as sm
+
+from zipline.algorithm import TradingAlgorithm
+from zipline.transforms import batch_transform
+from zipline.utils.factory import load_from_yahoo
+
+
+@batch_transform
+def ols_transform(data, sid1, sid2):
+    """Computes regression coefficient (slope and intercept)
+    via Ordinary Least Squares between two SIDs.
+    """
+    p0 = data.price[sid1]
+    p1 = sm.add_constant(data.price[sid2])
+    slope, intercept = sm.OLS(p0, p1).fit().params
+
+    return slope, intercept
+
+
+class Pairtrade(TradingAlgorithm):
+    """Pairtrading relies on cointegration of two stocks.
+
+    The expectation is that once the two stocks drifted apart
+    (i.e. there is spread), they will eventually revert again. Thus,
+    if we short the upward drifting stock and long the downward
+    drifting stock (in short, we buy the spread) once the spread
+    widened we can sell the spread with profit once they converged
+    again. A nice property of this algorithm is that we enter the
+    market in a neutral position.
+
+    This specific algorithm tries to exploit the cointegration of
+    Pepsi and Coca Cola by estimating the correlation between the
+    two. Divergence of the spread is evaluated by z-scoring.
+    """
+
+    def initialize(self, window_length=100):
+        self.spreads = []
+        self.zscores = []
+        self.invested = 0
+        self.window_length = window_length
+        self.ols_transform = ols_transform(refresh_period=self.window_length,
+                                           days=self.window_length)
+
+    def handle_data(self, data):
+        ######################################################
+        # 1. Compute regression coefficients between PEP and KO
+        params = self.ols_transform.handle_data(data, 'PEP', 'KO')
+        if params is None:
+            return
+        slope, intercept = params
+
+        ######################################################
+        # 2. Compute spread and zscore
+        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)
+        spread_wind = self.spreads[-self.window_length:]
+        zscore = (spread - np.mean(spread_wind)) / np.std(spread_wind)
+        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
+        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
+        elif abs(zscore) < .5 and self.invested:
+            self.sell_spread()
+            self.invested = False
+
+    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)
+
+if __name__ == '__main__':
+    data = load_from_yahoo(stocks=['PEP', 'KO'], indexes={})
+
+    pairtrade = Pairtrade()
+    results = pairtrade.run(data)
+    data['spreads'] = np.nan
+    data.spreads[70:] = pairtrade.spreads
+
+    ax1 = plt.subplot(211)
+    data[['PEP', 'KO']].plot(ax=ax1)
+    plt.ylabel('price')
+    plt.setp(ax1.get_xticklabels(), visible=False)
+
+    ax2 = plt.subplot(212, sharex=ax1)
+    data.spreads.plot(ax=ax2, color='r')
+    plt.ylabel('spread')
+
+    plt.show()