Merge pull request #31 from quantopian/slippage_model_bug

Slippage model bug

tests/test_algorithm_gen.py

@@ -0,0 +1,103 @@
+#!/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.
+
+from unittest import TestCase
+from nose.tools import timed
+
+from datetime import datetime
+import pytz
+
+from zipline.algorithm import TradingAlgorithm
+from zipline.finance import slippage
+from zipline.utils import factory
+from zipline.utils.test_utils import (
+    setup_logger,
+    teardown_logger
+)
+
+DEFAULT_TIMEOUT = 15  # seconds
+EXTENDED_TIMEOUT = 90
+
+
+class RecordDateSlippage(slippage.FixedSlippage):
+    def __init__(self, spread):
+        super(RecordDateSlippage, self).__init__(spread=spread)
+        self.latest_date = None
+
+    def simulate(self, event, open_orders):
+        self.latest_date = event['datetime']
+        result = super(RecordDateSlippage, self).simulate(event, open_orders)
+        return result
+
+
+class TestAlgo(TradingAlgorithm):
+
+    def __init__(self, asserter, *args, **kwargs):
+        super(TestAlgo, self).__init__(*args, **kwargs)
+        self.asserter = asserter
+
+    def initialize(self, window_length=100):
+        self.latest_date = None
+
+        self.set_slippage(RecordDateSlippage(spread=0.05))
+        self.stocks = [8229]
+        self.ordered = False
+
+    def handle_data(self, data):
+        self.latest_date = self.get_datetime()
+
+        if not self.ordered:
+            for stock in self.stocks:
+                self.order(stock, 100)
+
+            self.ordered = True
+
+        self.asserter.assertGreaterEqual(
+            self.latest_date,
+            self.slippage.latest_date
+        )
+
+
+class AlgorithmGeneratorTestCase(TestCase):
+    def setUp(self):
+        setup_logger(self)
+
+    def tearDown(self):
+        teardown_logger(self)
+
+    @timed(DEFAULT_TIMEOUT)
+    def test_generator_dates(self):
+        """
+        Ensure the pipeline of generators are in sync, at least as far as
+        their current dates.
+        """
+        algo = TestAlgo(self)
+        trading_environment = factory.create_trading_environment(
+            start=datetime(2012, 1, 3, tzinfo=pytz.utc),
+            end=datetime(2012, 7, 30, tzinfo=pytz.utc)
+        )
+        trade_source = factory.create_daily_trade_source(
+            [8229],
+            200,
+            trading_environment
+        )
+        algo.set_sources([trade_source])
+
+        gen = algo.get_generator(trading_environment)
+        self.assertTrue(list(gen))
+
+        self.assertTrue(algo.slippage.latest_date)
+        self.assertTrue(algo.latest_date)

zipline/algorithm.py

@@ -20,6 +20,9 @@ import numpy as np
 
 from datetime import datetime
 
+from itertools import groupby
+from operator import attrgetter
+
 from zipline.sources import DataFrameSource
 from zipline.utils.factory import create_trading_environment
 from zipline.transforms.utils import StatefulTransform
@@ -95,12 +98,15 @@ class TradingAlgorithm(object):
         self.date_sorted = date_sorted_sources(*self.sources)
         self.with_tnfms = sequential_transforms(self.date_sorted,
                                                 *self.transforms)
+        # Group together events with the same dt field. This depends on the
+        # events already being sorted.
+        self.grouped_by_date = groupby(self.with_tnfms, attrgetter('dt'))
         self.trading_client = tsc(self, environment)
 
         transact_method = transact_partial(self.slippage, self.commission)
         self.set_transact(transact_method)
 
-        return self.trading_client.simulate(self.with_tnfms)
+        return self.trading_client.simulate(self.grouped_by_date)
 
     def get_generator(self, environment):
         """

zipline/finance/performance.py

@@ -210,16 +210,18 @@ class PerformanceTracker(object):
         """
         Main generator work loop.
         """
-        for event in stream_in:
-            if event.dt == "DONE":
-                event.perf_message = self.handle_simulation_end()
-                del event['TRANSACTION']
-                yield event
-            else:
-                event.perf_message = self.process_event(event)
-                event.portfolio = self.get_portfolio()
-                del event['TRANSACTION']
-                yield event
+        for date, snapshot in stream_in:
+            yield date, [self._transform_event(event) for event in snapshot]
+
+    def _transform_event(self, event):
+        if event.dt == "DONE":
+            event.perf_message = self.handle_simulation_end()
+        else:
+            event.perf_message = self.process_event(event)
+            event.portfolio = self.get_portfolio()
+
+        del event['TRANSACTION']
+        return event
 
     def get_portfolio(self):
         return self.cumulative_performance.as_portfolio()

zipline/finance/slippage.py

@@ -21,19 +21,27 @@ from functools import partial
 
 from zipline.utils.protocol_utils import ndict
 
+from logbook import Processor
+
 
 def transact_stub(slippage, commission, event, open_orders):
     """
     This is intended to be wrapped in a partial, so that the
     slippage and commission models can be enclosed.
     """
-    transaction = slippage.simulate(event, open_orders)
-    if transaction and transaction.amount != 0:
-        direction = abs(transaction.amount) / transaction.amount
-        per_share, total_commission = commission.calculate(transaction)
-        transaction.price = transaction.price + (per_share * direction)
-        transaction.commission = total_commission
-    return transaction
+    def inject_algo_dt(record):
+        if not 'algo_dt' in record.extra:
+            record.extra['algo_dt'] = event['dt']
+
+    with Processor(inject_algo_dt).threadbound():
+
+        transaction = slippage.simulate(event, open_orders)
+        if transaction and transaction.amount != 0:
+            direction = abs(transaction.amount) / transaction.amount
+            per_share, total_commission = commission.calculate(transaction)
+            transaction.price = transaction.price + (per_share * direction)
+            transaction.commission = total_commission
+        return transaction
 
 
 def transact_partial(slippage, commission):

zipline/finance/trading.py

@@ -44,8 +44,8 @@ class TransactionSimulator(object):
         """
         Main generator work loop.
         """
-        for event in stream_in:
-            yield self.update(event)
+        for date, snapshot in stream_in:
+            yield date, [self.update(event) for event in snapshot]
 
     def update(self, event):
         event.TRANSACTION = None

zipline/gens/tradesimulation.py

@@ -18,8 +18,6 @@ from logbook import Logger, Processor
 from collections import defaultdict
 
 from datetime import datetime
-from itertools import groupby
-from operator import attrgetter
 
 from zipline import ndict
 
@@ -161,7 +159,8 @@ class AlgorithmSimulator(object):
         # Processor function for injecting the algo_dt into
         # user prints/logs.
         def inject_algo_dt(record):
-            record.extra['algo_dt'] = self.snapshot_dt
+            if not 'algo_dt' in record.extra:
+                record.extra['algo_dt'] = self.snapshot_dt
         self.processor = Processor(inject_algo_dt)
 
     def order(self, sid, amount):
@@ -199,9 +198,8 @@ class AlgorithmSimulator(object):
         # inject the current algo
         # snapshot time to any log record generated.
         with self.processor.threadbound():
-            # Group together events with the same dt field. This depends on the
-            # events already being sorted.
-            for date, snapshot in groupby(stream_in, attrgetter('dt')):
+
+            for date, snapshot in stream_in:
                 # Set the simulation date to be the first event we see.
                 # This should only occur once, at the start of the test.
                 if self.simulation_dt is None: