Merge pull request #48 from quantopian/optimize_refactor

Test-framework for optimization.

pavement.py

@@ -39,6 +39,7 @@ version='dev'
 install_requires = parse_requirements('./etc/requirements.txt') + parse_requirements('./etc/requirements_sci.txt')
 tests_require = install_requires + parse_requirements('./etc/requirements_dev.txt')
 
+
 options(
     sphinx=Bunch(
         builddir="_build",
@@ -48,6 +49,8 @@ options(
           version              = version,
           classifiers          = [],
           packages             = find_packages(),
+          package_data         = find_package_data("zipline", package="zipline",
+                                                   only_in_packages=False),
           install_requires     = install_requires,
           tests_require        = tests_require,
           test_suite           = 'nose.collector',

tests/test_optimize.py

@@ -0,0 +1,206 @@
+"""Tests for the zipline.finance package"""
+import unittest
+from unittest2 import TestCase, skip
+from nose.tools import timed
+from collections import defaultdict
+from datetime import datetime, timedelta
+
+import numpy as np
+
+from zipline.optimize.factory import create_updown_trade_source
+import zipline.utils.factory as factory
+import zipline.util as qutil
+
+from zipline.simulator import AddressAllocator, Simulator
+from zipline.optimize.algorithms import BuySellAlgorithm
+from zipline.finance.trading import TradingEnvironment
+from zipline.lines import SimulatedTrading
+from zipline.finance.trading import SIMULATION_STYLE
+
+DEFAULT_TIMEOUT = 15 # seconds
+EXTENDED_TIMEOUT = 90
+
+allocator = AddressAllocator(1000)
+
+class TestUpDown(TestCase):
+    """This unittest verifies that the BuySellAlgorithm in
+    combination with the UpDownSource are suitable for usage in an
+    optimization framework.
+
+    """
+    leased_sockets = defaultdict(list)
+
+    def setUp(self):
+        qutil.configure_logging()
+        self.zipline_test_config = {
+            'allocator':allocator,
+            'sid':133
+        }
+
+    @timed(DEFAULT_TIMEOUT)
+    def test_source_and_orders(self):
+        """verify that UpDownSource is having the correct
+        behavior and that BuySellAlgorithm places the buy/sell
+        orders at the right time. Moreover, establishes that
+        UpDownSource and BuySellAlgorithm interact correctly."
+
+        """
+        #generate events
+        trade_count = 5
+        sid = 133
+        base_price = 50
+        amplitude = 6
+        offset = 0
+        self.zipline_test_config['order_count'] = trade_count - 1
+        self.zipline_test_config['trade_count'] = trade_count
+        self.zipline_test_config['simulation_style'] = \
+        SIMULATION_STYLE.FIXED_SLIPPAGE
+
+        trading_environment = factory.create_trading_environment()
+        source = create_updown_trade_source(sid,
+            trade_count,
+            trading_environment,
+            base_price,
+            amplitude
+        )
+
+        prices = np.array([event.price for event in source.event_list])
+        max_price_idx = np.where(prices==prices.max())[0]
+        min_price_idx = np.where(prices==prices.min())[0]
+        self.assertTrue(np.all(max_price_idx % 2 == 1),
+            "Maximum prices are not periodic."
+        )
+        self.assertTrue(np.all(min_price_idx % 2 == 0),
+            "Minimum prices are not periodic."
+        )
+        self.assertEqual(prices.max(), base_price+amplitude/2.,
+            "Maximum price does not equal expected maximum price."
+        )
+        self.assertEqual(prices.min(), base_price-amplitude/2.,
+            "Minimum price does not equal expected maximum price."
+        )
+
+        algo = BuySellAlgorithm(sid, 100, 0)
+
+        self.zipline_test_config['trade_source'] = source
+        self.zipline_test_config['algorithm'] = algo
+        self.zipline_test_config['environment'] = trading_environment
+
+        zipline = SimulatedTrading.create_test_zipline(**self.zipline_test_config)
+        zipline.simulate(blocking=True)
+
+        orders = np.asarray(algo.orders)
+        max_order_idx = np.where(orders==orders.max())[0]
+        min_order_idx = np.where(orders==orders.min())[0]
+
+        self.assertTrue(np.all(max_order_idx % 2 == 1),
+            "Maximum orders are not periodic."
+        )
+        self.assertTrue(np.all(min_order_idx % 2 == 0),
+            "Minimum orders are not periodic."
+        )
+        self.assertTrue(np.all(max_order_idx == max_price_idx),
+            "Algorithm did not buy when price was going to drop."
+        )
+        self.assertTrue(np.all(min_order_idx == min_price_idx),
+            "Algorithm did not sell when price was going to increase."
+        )
+
+    def test_concavity_of_returns(self):
+        """verify concave relationship between of free parameter and
+        returns in certain region around the max. Moreover,
+        establishes that the max returns is at the correct value
+        (i.e. 0).
+
+        """
+        #generate events
+        trade_count = 6
+        sid = 133
+        amplitude = 30
+        base_price = 50
+        self.zipline_test_config['order_count'] = trade_count - 1
+        self.zipline_test_config['trade_count'] = trade_count
+        self.zipline_test_config['simulation_style'] = \
+        SIMULATION_STYLE.FIXED_SLIPPAGE
+
+        #test whether return-function is concave wrt repeats.
+        test_offsets = np.arange(-9, 9, 1.)
+        supposed_max = np.zeros(len(test_offsets), dtype=bool)
+        supposed_max[len(test_offsets) // 2] = True
+
+        compound_returns = np.empty(len(test_offsets))
+        ziplines = []
+        for i, test_offset in enumerate(test_offsets):
+            trading_environment = factory.create_trading_environment()
+            source = create_updown_trade_source(sid,
+                trade_count,
+                trading_environment,
+                base_price,
+                amplitude
+            )
+
+            algo = BuySellAlgorithm(sid, 100, test_offset)
+            self.zipline_test_config['algorithm'] = algo
+            self.zipline_test_config['trade_source'] = source
+            self.zipline_test_config['environment'] = trading_environment
+            zipline = SimulatedTrading.create_test_zipline(**self.zipline_test_config)
+            zipline.simulate(blocking=True)
+            ziplines.append(zipline)
+            compound_returns[i] = zipline.get_cumulative_performance()['returns']
+
+        self.assertTrue(np.all(compound_returns[supposed_max] > compound_returns[np.logical_not(supposed_max)]),
+            "Maximum compound returns are not where they are supposed to be."
+        )
+
+        # test for concavity
+        max_idx = np.where(supposed_max)[0][0]
+        idx = np.array([max_idx, max_idx])
+        for i in range((len(test_offsets)-1)/2):
+            # going outwards, returns must decrease
+            self.assertTrue(compound_returns[idx[0]-1] < compound_returns[idx[0]],
+                "Compound returns are not convex."
+            )
+            self.assertTrue(compound_returns[idx[1]+1] < compound_returns[idx[1]],
+                "Compound returns are not convex."
+            )
+            idx[0] -= 1
+            idx[1] += 1
+
+    @skip
+    def test_optimize(self):
+        """verify that gradient descent (Powell's method) can find
+        the optimal free parameter under which the BuySellAlgorithm produces
+        maximum returns.
+
+        """
+        def simulate(offset):
+            #generate events
+            trade_count = 3
+            sid = 133
+            amplitude = 10
+            base_price = 50
+            self.zipline_test_config['order_count'] = trade_count - 1
+            self.zipline_test_config['trade_count'] = trade_count
+            self.zipline_test_config['simulation_style'] = \
+            SIMULATION_STYLE.FIXED_SLIPPAGE
+            trading_environment = factory.create_trading_environment()
+            source = create_updown_trade_source(sid,
+                trade_count,
+                trading_environment,
+                base_price,
+                amplitude
+            )
+
+            algo = BuySellAlgorithm(sid, 100, offset)
+            self.zipline_test_config['algorithm'] = algo
+            self.zipline_test_config['trade_source'] = source
+            self.zipline_test_config['environment'] = trading_environment
+            zipline = SimulatedTrading.create_test_zipline(**self.zipline_test_config)
+            zipline.simulate(blocking=True)
+            zipline.shutdown()
+            #function is getting minimized, so have to return negative cum returns.
+            return -zipline.get_cumulative_performance()['returns']
+
+        from scipy import optimize
+        opt = optimize.fmin_powell(simulate, 1.5)
+        np.testing.assert_almost_equal(opt, 0, 5)

zipline/core/component.py

@@ -99,6 +99,7 @@ class Component(object):
         """
         pass
 
+
     # ------------
     # Core Methods
     # ------------

zipline/finance/sources.py

@@ -4,6 +4,7 @@ Provides data handlers that can push messages to a zipline.core.DataFeed
 import datetime
 import random
 import pytz
+from mock import Mock
 
 from zipline.components import DataSource
 from zipline.utils import ndict
@@ -91,9 +92,21 @@ class SpecificEquityTrades(TradeDataSource):
         self.event_list = event_list
         self.count = 0
 
+        # TODO temporary hack
+        self.control_out = Mock()
+
     def get_type(self):
         zp.COMPONENT_TYPE.SOURCE
 
+    @property
+    def get_id(self):
+        """
+        The descriptive name of the component.
+        """
+        # Prevents the bug that Thomas ran into
+        return "Unique ID"
+
+
     def do_work(self):
         if(len(self.event_list) == 0):
             self.signal_done()

zipline/optimize/algorithms.py

@@ -0,0 +1,49 @@
+class BuySellAlgorithm():
+    """Algorithm that buys and sells alternatingly. The amount for
+    each order can be specified. In addition, an offset that will
+    quadratically reduce the amount that will be bought can be
+    specified.
+
+    This algorithm is used to test the parameter optimization
+    framework. If combined with the UpDown trade source, an offset of
+    0 will produce maximum returns.
+
+    """
+
+    def __init__(self, sid, amount, offset):
+        self.sid = sid
+        self.amount = amount
+        self.incr = 0
+        self.done = False
+        self.order = None
+        self.frame_count = 0
+        self.portfolio = None
+        self.buy_or_sell = -1
+        self.offset = offset
+        self.orders = []
+        self.prices = []
+
+    def initialize(self):
+        pass
+
+    def set_order(self, order_callable):
+        self.order = order_callable
+
+    def set_portfolio(self, portfolio):
+        self.portfolio = portfolio
+
+    def handle_data(self, frame):
+        order_size = self.buy_or_sell * (self.amount - (self.offset**2))
+        self.order(self.sid, order_size)
+
+        #sell next time around.
+        self.buy_or_sell *= -1
+
+        self.orders.append(order_size)
+        #self.prices.append(frame['price'])
+
+        self.frame_count += 1
+        self.incr += 1
+
+    def get_sid_filter(self):
+        return [self.sid]

zipline/optimize/factory.py

@@ -0,0 +1,65 @@
+"""
+Factory functions to prepare useful data for optimize tests.
+
+Author: Thomas V. Wiecki (thomas.wiecki@gmail.com), 2012
+"""
+from datetime import datetime, timedelta
+
+import zipline.protocol as zp
+
+from zipline.test.factory import get_next_trading_dt
+from zipline.finance.sources import SpecificEquityTrades
+from zipline.optimize.algorithms import BuySellAlgorithm
+from zipline.lines import SimulatedTrading
+
+def create_updown_trade_source(sid, trade_count, trading_environment, start_price, amplitude):
+    from itertools import cycle
+    volume = 1000
+    events = []
+    price = start_price-amplitude/2.
+
+    cur = trading_environment.first_open
+    one_day = timedelta(days = 1)
+
+    #create iterator to cycle through up and down phases
+    change = cycle([1,-1])
+
+    for i in xrange(trade_count + 2):
+        cur = get_next_trading_dt(cur, one_day, trading_environment)
+
+        event = zp.ndict({
+            "type"      : zp.DATASOURCE_TYPE.TRADE,
+            "sid"       : sid,
+            "price"     : price,
+            "volume"    : volume,
+            "dt"        : cur,
+        })
+
+        events.append(event)
+
+        price += change.next()*amplitude
+
+    trading_environment.period_end = cur
+
+    source = SpecificEquityTrades(sid, events)
+
+    return source
+
+
+def create_predictable_zipline(config, sid=133, amplitude=10, base_price=50, offset=0):
+    config = deepcopy(config)
+    trading_environment = create_trading_environment()
+    source = create_updown_trade_source(sid,
+                                        config['trade_count'],
+                                        trading_environment,
+                                        base_price,
+                                        amplitude)
+
+    algo = RegularIntervalBuySellAlgorithm(sid, 100, offset)
+    config['algorithm'] = algo
+    config['trade_source'] = source
+    config['environment'] = trading_environment
+    zipline = SimulatedTrading.create_test_zipline(**config)
+    zipline.simulate(blocking=True)
+
+    return zipline

zipline/utils/factory.py

@@ -5,6 +5,8 @@ Factory functions to prepare useful data for tests.
 import pytz
 import msgpack
 import random
+from os.path import join
+from operator import attrgetter
 
 from datetime import datetime, timedelta
 import zipline.finance.risk as risk
@@ -19,15 +21,16 @@ def load_market_data():
     for packed_date, returns in bm_list:
         event_dt = zp.tuple_to_date(packed_date)
         #event_dt = event_dt.replace(
-        #    hour=0, 
-        #    minute=0, 
-        #    second=0, 
+        #    hour=0,
+        #    minute=0,
+        #    second=0,
         #    tzinfo=pytz.utc
         #)
 
         daily_return = risk.DailyReturn(date=event_dt, returns=returns)
         bm_returns.append(daily_return)
-    bm_returns = sorted(bm_returns, key=lambda(x): x.date)
+
+    bm_returns = sorted(bm_returns, key=attrgetter('date'))
     fp_tr = open(".//tests/treasury_curves.msgpack", "rb")
     tr_list = msgpack.loads(fp_tr.read())
     tr_curves = {}
@@ -223,7 +226,7 @@ def create_trade_source(sids, trade_count, trade_time_increment, trading_environ
 
         trade_history.extend(generated_trades)
 
-    trade_history = sorted(trade_history, key=lambda(x): x.dt)
+    trade_history = sorted(trade_history, key=attrgetter('dt'))
 
     #set the trading environment's end to same dt as the last trade in the
     #history.