ENH: Stream benchmark returns as events.

Instead of creating a list of benchmarks in the risk module,
stream benchmarks through the system as events, starting from the
algorithm generator.

Works towards more easily setting arbritrary pricing data as
a a benchmark, as well as working towards live minutely benchmarks.

tests/test_finance.py

@@ -28,10 +28,14 @@ import numpy as np
 
 from nose.tools import timed
 
+import zipline.protocol
+from zipline.protocol import Event, DATASOURCE_TYPE
+
 import zipline.utils.factory as factory
 import zipline.utils.simfactory as simfactory
 
 from zipline.gens.tradesimulation import Order, Blotter
+from zipline.gens.composites import date_sorted_sources
 
 import zipline.finance.trading as trading
 from zipline.finance.trading import SimulationParameters
@@ -165,8 +169,9 @@ class FinanceTestCase(TestCase):
         # No transactions can be filled on the first trade, so
         # we have one extra trade to ensure all orders are filled.
         self.zipline_test_config['trade_count'] = 101
-        zipline = simfactory.create_test_zipline(**self.zipline_test_config)
-        assert_single_position(self, zipline)
+        full_zipline = simfactory.create_test_zipline(
+            **self.zipline_test_config)
+        assert_single_position(self, full_zipline)
 
     # TODO: write tests for short sales
     # TODO: write a test to do massive buying or shorting.
@@ -340,18 +345,34 @@ class FinanceTestCase(TestCase):
 
         tracker = PerformanceTracker(sim_params)
 
+        benchmark_returns = [
+            Event({'dt': ret.date,
+                   'returns': ret.returns,
+                   'type':
+                   zipline.protocol.DATASOURCE_TYPE.BENCHMARK,
+                   'source_id': 'benchmarks'})
+            for ret in trading.environment.benchmark_returns
+            if ret.date.date() >= sim_params.period_start.date()
+            and ret.date.date() <= sim_params.period_end.date()
+        ]
+
+        generated_events = date_sorted_sources(generated_trades,
+                                               benchmark_returns)
+
         # this approximates the loop inside TradingSimulationClient
         transactions = []
-        for dt, trades in itertools.groupby(generated_trades,
+        for dt, events in itertools.groupby(generated_events,
                                             operator.attrgetter('dt')):
-            for trade in trades:
+            for event in events:
+                if event.type == DATASOURCE_TYPE.TRADE:
 
-                txns = blotter.process_trade(trade)
+                    txns = blotter.process_trade(event)
 
-                for txn in txns:
-                    transactions.append(txn)
-                    tracker.process_event(txn)
-                tracker.process_event(trade)
+                    for txn in txns:
+                        transactions.append(txn)
+                        tracker.process_event(txn)
+
+                tracker.process_event(event)
 
         if complete_fill:
             self.assertEqual(len(transactions), len(order_list))

tests/test_perf_tracking.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import collections
+import heapq
 
 import unittest
 from nose_parameterized import parameterized
@@ -32,6 +33,8 @@ from zipline.gens.tradesimulation import Order
 import zipline.finance.trading as trading
 from zipline.protocol import DATASOURCE_TYPE
 from zipline.utils.factory import create_random_simulation_parameters
+import zipline.protocol
+from zipline.protocol import Event
 
 onesec = datetime.timedelta(seconds=1)
 oneday = datetime.timedelta(days=1)
@@ -42,6 +45,19 @@ def create_txn(sid, price, amount, dt):
     return create_transaction(sid, amount, price, dt, "fakeuid")
 
 
+def benchmark_events_in_range(sim_params):
+    return [
+        Event({'dt': ret.date,
+               'returns': ret.returns,
+               'type':
+               zipline.protocol.DATASOURCE_TYPE.BENCHMARK,
+               'source_id': 'benchmarks'})
+        for ret in trading.environment.benchmark_returns
+        if ret.date.date() >= sim_params.period_start.date()
+        and ret.date.date() <= sim_params.period_end.date()
+    ]
+
+
 class TestDividendPerformance(unittest.TestCase):
 
     def setUp(self):
@@ -51,6 +67,8 @@ class TestDividendPerformance(unittest.TestCase):
 
         self.sim_params.capital_base = 10e3
 
+        self.benchmark_events = benchmark_events_in_range(self.sim_params)
+
     def test_market_hours_calculations(self):
         with trading.TradingEnvironment():
             # DST in US/Eastern began on Sunday March 14, 2010
@@ -87,10 +105,15 @@ class TestDividendPerformance(unittest.TestCase):
         txn = create_txn(1, 10.0, 100, events[0].dt)
         events.insert(0, txn)
         events.insert(1, dividend)
+
         perf_tracker = perf.PerformanceTracker(self.sim_params)
+
+        all_events = (msg[1] for msg in heapq.merge(
+            ((event.dt, event) for event in events),
+            ((event.dt, event) for event in self.benchmark_events)))
+
         transformed_events = list(perf_tracker.transform(
-            ((event.dt, [event]) for event in events))
-        )
+            itertools.groupby(all_events, attrgetter('dt'))))
 
         #flatten the list of events
         results = []
@@ -145,9 +168,13 @@ class TestDividendPerformance(unittest.TestCase):
         txn = create_txn(1, 10.0, 100, events[3].dt)
         events.insert(4, txn)
         perf_tracker = perf.PerformanceTracker(self.sim_params)
+
+        all_events = (msg[1] for msg in heapq.merge(
+            ((event.dt, event) for event in events),
+            ((event.dt, event) for event in self.benchmark_events)))
+
         transformed_events = list(perf_tracker.transform(
-            ((event.dt, [event]) for event in events))
-        )
+            itertools.groupby(all_events, attrgetter('dt'))))
 
         #flatten the list of events
         results = []
@@ -195,9 +222,13 @@ class TestDividendPerformance(unittest.TestCase):
         events.insert(4, sell_txn)
         events.insert(0, dividend)
         perf_tracker = perf.PerformanceTracker(self.sim_params)
+
+        all_events = (msg[1] for msg in heapq.merge(
+            ((event.dt, event) for event in events),
+            ((event.dt, event) for event in self.benchmark_events)))
+
         transformed_events = list(perf_tracker.transform(
-            ((event.dt, [event]) for event in events))
-        )
+            itertools.groupby(all_events, attrgetter('dt'))))
 
         #flatten the list of events
         results = []
@@ -245,9 +276,13 @@ class TestDividendPerformance(unittest.TestCase):
         events.insert(4, sell_txn)
         events.insert(1, dividend)
         perf_tracker = perf.PerformanceTracker(self.sim_params)
+
+        all_events = heapq.merge(
+            ((event.dt, event) for event in events),
+            ((event.dt, event) for event in self.benchmark_events))
+
         transformed_events = list(perf_tracker.transform(
-            ((event.dt, [event]) for event in events))
-        )
+            (event[0], [event[1]]) for event in all_events))
 
         #flatten the list of events
         results = []
@@ -293,9 +328,13 @@ class TestDividendPerformance(unittest.TestCase):
         events.insert(2, buy_txn)
         events.insert(1, dividend)
         perf_tracker = perf.PerformanceTracker(self.sim_params)
+
+        all_events = (msg[1] for msg in heapq.merge(
+            ((event.dt, event) for event in events),
+            ((event.dt, event) for event in self.benchmark_events)))
+
         transformed_events = list(perf_tracker.transform(
-            ((event.dt, [event]) for event in events))
-        )
+            itertools.groupby(all_events, attrgetter('dt'))))
 
         #flatten the list of events
         results = []
@@ -344,9 +383,13 @@ class TestDividendPerformance(unittest.TestCase):
         events.insert(1, txn)
         events.insert(0, dividend)
         perf_tracker = perf.PerformanceTracker(self.sim_params)
+
+        all_events = (msg[1] for msg in heapq.merge(
+            ((event.dt, event) for event in events),
+            ((event.dt, event) for event in self.benchmark_events)))
+
         transformed_events = list(perf_tracker.transform(
-            ((event.dt, [event]) for event in events))
-        )
+            itertools.groupby(all_events, attrgetter('dt'))))
 
         #flatten the list of events
         results = []
@@ -390,9 +433,13 @@ class TestDividendPerformance(unittest.TestCase):
 
         events.insert(1, dividend)
         perf_tracker = perf.PerformanceTracker(self.sim_params)
+
+        all_events = (msg[1] for msg in heapq.merge(
+            ((event.dt, event) for event in events),
+            ((event.dt, event) for event in self.benchmark_events)))
+
         transformed_events = list(perf_tracker.transform(
-            ((event.dt, [event]) for event in events))
-        )
+            itertools.groupby(all_events, attrgetter('dt'))))
 
         #flatten the list of events
         results = []
@@ -423,6 +470,8 @@ class TestPositionPerformance(unittest.TestCase):
         self.sim_params, self.dt, self.end_dt = \
             create_random_simulation_parameters()
 
+        self.benchmark_events = benchmark_events_in_range(self.sim_params)
+
     def test_long_position(self):
         """
             verify that the performance period calculates properly for a
@@ -935,6 +984,8 @@ class TestPerformanceTracker(unittest.TestCase):
             period_end=end_dt
         )
 
+        benchmark_events = benchmark_events_in_range(sim_params)
+
         trade_history = factory.create_trade_history(
             sid,
             price_list,
@@ -1000,12 +1051,17 @@ class TestPerformanceTracker(unittest.TestCase):
         orders = [event for event in
                   events if event.type == DATASOURCE_TYPE.ORDER]
 
+        all_events = (msg[1] for msg in heapq.merge(
+            ((event.dt, event) for event in events),
+            ((event.dt, event) for event in benchmark_events)))
+
+        # Extract events with transactions to use for verification.
         perf_messages = \
-            [msg for date, snapshot in
+            [m for date, snapshot in
              perf_tracker.transform(
-                 itertools.groupby(events, attrgetter('dt')))
-             for event in snapshot
-             for msg in event.perf_messages]
+                 itertools.groupby(all_events, attrgetter('dt')))
+             for e in snapshot
+             for m in e.perf_messages]
 
         end_perf_messages, risk_message = perf_tracker.handle_simulation_end()
 

zipline/algorithm.py

@@ -39,6 +39,9 @@ from zipline.finance.slippage import (
 )
 from zipline.finance.commission import PerShare, PerTrade
 from zipline.finance.constants import ANNUALIZER
+import zipline.finance.trading as trading
+import zipline.protocol
+from zipline.protocol import Event
 
 from zipline.gens.composites import (
     date_sorted_sources,
@@ -129,17 +132,31 @@ class TradingAlgorithm(object):
         processed by the zipline, and False for those that should be
         skipped.
         """
+        benchmark_return_source = [
+            Event({'dt': ret.date,
+                   'returns': ret.returns,
+                   'type': zipline.protocol.DATASOURCE_TYPE.BENCHMARK,
+                   'source_id': 'benchmarks'})
+            for ret in trading.environment.benchmark_returns
+            if ret.date.date() >= self.sim_params.period_start.date()
+            and ret.date.date() <= self.sim_params.period_end.date()
+        ]
 
         date_sorted = date_sorted_sources(*self.sources)
+
         if source_filter:
             date_sorted = ifilter(source_filter, date_sorted)
+
         with_tnfms = sequential_transforms(date_sorted,
                                            *self.transforms)
         with_alias_dt = alias_dt(with_tnfms)
 
+        with_benchmarks = date_sorted_sources(benchmark_return_source,
+                                              with_alias_dt)
+
         # Group together events with the same dt field. This depends on the
         # events already being sorted.
-        return groupby(with_alias_dt, attrgetter('dt'))
+        return groupby(with_benchmarks, attrgetter('dt'))
 
     def _create_generator(self, sim_params, source_filter=None):
         """

zipline/finance/performance.py

@@ -165,11 +165,8 @@ class PerformanceTracker(object):
         self.emission_rate = sim_params.emission_rate
 
         # Temporarily hold these here as we work on streaming benchmarks.
-        self.all_benchmark_returns = pd.Series({
-            x.date: x.returns
-            for x in trading.environment.benchmark_returns
-            if x.date >= self.period_start
-        })
+        self.all_benchmark_returns = pd.Series(
+            index=trading.environment.trading_days)
 
         # this performance period will span the entire simulation.
         self.cumulative_performance = PerformancePeriod(
@@ -324,6 +321,8 @@ class PerformanceTracker(object):
             # we just want to relay this event unchanged.
             messages = []
             return messages
+        elif event.type == zp.DATASOURCE_TYPE.BENCHMARK:
+            self.all_benchmark_returns[event.dt] = event.returns
 
         #calculate performance as of last trade
         self.cumulative_performance.calculate_performance()

zipline/protocol.py

@@ -30,7 +30,8 @@ DATASOURCE_TYPE = Enum(
     'ORDER',
     'EMPTY',
     'DONE',
-    'CUSTOM'
+    'CUSTOM',
+    'BENCHMARK'
 )
 
 

zipline/transforms/utils.py

@@ -142,10 +142,7 @@ class StatefulTransform(object):
                     and message.type not in (
                         DATASOURCE_TYPE.TRADE,
                         DATASOURCE_TYPE.CUSTOM)):
-                # TODO: this should be yielding the original message
-                # instead of swallowing it. Will be an issue when we
-                # have a transaction source from brokers etc.
-                continue
+                yield message
             # allow upstream generators to yield None to avoid
             # blocking.
             if message is None: