From 13a2b1c6379dd370e9611f893208f32e86a8b8e6 Mon Sep 17 00:00:00 2001 From: Eddie Hebert Date: Tue, 1 Jan 2013 20:56:25 -0500 Subject: [PATCH 1/4] Changes date_sort to use heapq module. Uses heapq.merge to sort input from mulitple sources instead of our own sort module. From profiling heapq.merge is more efficient than our own efforts. --- zipline/gens/composites.py | 30 ++++++++++++------------------ 1 file changed, 12 insertions(+), 18 deletions(-) diff --git a/zipline/gens/composites.py b/zipline/gens/composites.py index 02b40f48..04824c3a 100644 --- a/zipline/gens/composites.py +++ b/zipline/gens/composites.py @@ -13,10 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +import heapq + from itertools import chain -from zipline.gens.utils import roundrobin, done_message -from zipline.gens.sort import date_sort +from zipline.gens.utils import done_message + + +def decorate_source(source): + for message in source: + yield ((message.dt, message.source_id), message) def date_sorted_sources(*sources): @@ -24,23 +30,11 @@ def date_sorted_sources(*sources): Takes an iterable of sources, generating namestrings and piping their output into date_sort. """ + sorted_stream = heapq.merge(*(decorate_source(s) for s in sources)) - for source in sources: - assert iter(source), "Source %s not iterable" % source - assert hasattr(source, 'get_hash'), "No get_hash" - - # Get name hashes to pass to date_sort. - names = [source.get_hash() for source in sources] - - # Convert the list of generators into a flat stream by pulling - # one element at a time from each. - stream_in = roundrobin(sources, names) - - # Guarantee the flat stream will be sorted by date, using - # source_id as tie-breaker, which is fully deterministic (given - # deterministic string representation for all args/kwargs) - - return date_sort(stream_in, names) + # Strip out key decoration + for _, message in sorted_stream: + yield message def sequential_transforms(stream_in, *transforms): From e7a31c0661e17daa6950dca3621a253598fdd80d Mon Sep 17 00:00:00 2001 From: Eddie Hebert Date: Tue, 1 Jan 2013 23:01:25 -0500 Subject: [PATCH 2/4] Removes sort module. Sort module is now unneeded with use of heapq.merge Also adapts test_perf_tracking so that it uses date_sorted_sources. --- tests/test_perf_tracking.py | 38 +++--- tests/test_sorting.py | 264 ------------------------------------ zipline/gens/sort.py | 116 ---------------- zipline/gens/utils.py | 22 --- 4 files changed, 16 insertions(+), 424 deletions(-) delete mode 100644 tests/test_sorting.py delete mode 100644 zipline/gens/sort.py diff --git a/tests/test_perf_tracking.py b/tests/test_perf_tracking.py index 76698f28..7cab2518 100644 --- a/tests/test_perf_tracking.py +++ b/tests/test_perf_tracking.py @@ -27,8 +27,9 @@ from operator import attrgetter import zipline.utils.factory as factory import zipline.finance.performance as perf from zipline.utils.protocol_utils import ndict -from zipline.gens.sort import date_sort -from zipline.protocol import DATASOURCE_TYPE +from zipline.protocol import Event + +from zipline.gens.composites import date_sorted_sources from zipline.finance.trading import TradingEnvironment @@ -642,8 +643,6 @@ class TestPerformanceTracker(unittest.TestCase): del trade_history[-days_to_delete.end:] del trade_history2[-days_to_delete.end:] - trade_history.extend(trade_history2) - trading_environment.first_open = \ trading_environment.calculate_first_open() trading_environment.last_close = \ @@ -659,25 +658,21 @@ class TestPerformanceTracker(unittest.TestCase): trading_environment ) - # date_sort requires 'DONE' messages from each source - events = itertools.chain(trade_history, - [ndict({ - 'source_id': 'factory1', - 'dt': 'DONE', - 'type': DATASOURCE_TYPE.TRADE - }), - ndict({ - 'source_id': 'factory2', - 'dt': 'DONE', - 'type': DATASOURCE_TYPE.TRADE - })]) - events = date_sort(events, ('factory1', 'factory2')) - events = itertools.chain(events, - [ndict({'dt': 'DONE'})]) + events = date_sorted_sources(trade_history, trade_history2) events = [self.event_with_txn(event, trade_history[0].dt) for event in events] + # Extract events with transactions to use for verification. + events_with_txns = [event for event in events if event.TRANSACTION] + + done_message = Event({ + 'dt': 'DONE', + 'TRANSACTION': None + }) + + events = itertools.chain(events, [done_message]) + perf_messages = \ [msg for date, snapshot in perf_tracker.transform( @@ -686,11 +681,10 @@ class TestPerformanceTracker(unittest.TestCase): for msg in event.perf_messages] #we skip two trades, to test case of None transaction - txn_count = len(trade_history) - 2 - self.assertEqual(perf_tracker.txn_count, txn_count) + self.assertEqual(perf_tracker.txn_count, len(events_with_txns)) cumulative_pos = perf_tracker.cumulative_performance.positions[sid] - expected_size = txn_count / 2 * -25 + expected_size = len(events_with_txns) / 2 * -25 self.assertEqual(cumulative_pos.amount, expected_size) self.assertEqual(perf_tracker.last_close, diff --git a/tests/test_sorting.py b/tests/test_sorting.py deleted file mode 100644 index a1034bf1..00000000 --- a/tests/test_sorting.py +++ /dev/null @@ -1,264 +0,0 @@ -# -# 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 pytz - -from unittest import TestCase -from itertools import chain, izip_longest -from datetime import datetime, timedelta -from collections import deque - -from zipline import ndict -from zipline.gens.sort import ( - date_sort, - done, - queue_is_done -) -from zipline.gens.utils import alternate, done_message -from zipline.sources import SpecificEquityTrades -from zipline.gens.composites import date_sorted_sources - - -class HelperTestCase(TestCase): - - def setUp(self): - pass - - def tearDown(self): - pass - - def test_individual_queue_logic(self): - queue = deque() - # Empty queues are neither done nor ready. - assert not queue_is_done(queue) - - queue.append(to_dt('foo')) - assert not queue_is_done(queue) - - queue.appendleft(to_dt('DONE')) - - # Checking done when we have a message after done will trip an assert. - self.assertRaises(AssertionError, queue_is_done, queue) - - queue.pop() - assert queue_is_done(queue) - - def test_pop_logic(self): - sources = {} - ids = ['a', 'b', 'c'] - for id in ids: - sources[id] = deque() - - assert not done(sources) - - # All sources must have a message to be ready/done - sources['a'].append(to_dt("datetime")) - assert not done(sources) - sources['a'].pop() - - for id in ids: - sources[id].append(to_dt("datetime")) - - assert not done(sources) - - for id in ids: - sources[id].appendleft(to_dt("DONE")) - - # ["DONE", message] will trip an assert in queue_is_done. - self.assertRaises(AssertionError, done, sources) - - for id in ids: - sources[id].pop() - - assert done(sources) - - -class DateSortTestCase(TestCase): - - def setUp(self): - pass - - def tearDown(self): - pass - - def run_date_sort(self, event_stream, expected, source_ids): - """ - Take a list of events, their source_ids, and an expected sorting. - Assert that date_sort's output agrees with expected. - """ - sort_out = date_sort(event_stream, source_ids) - for m1, m2 in izip_longest(sort_out, expected): - assert m1 == m2 - - def test_single_source(self): - - # Just using the built-in defaults. See - # zipline.sources.py - source = SpecificEquityTrades() - expected = list(source) - source.rewind() - # The raw source doesn't handle done messaging, so we need to - # append a done message for sort to work properly. - with_done = chain(source, [done_message(source.get_hash())]) - self.run_date_sort(with_done, expected, [source.get_hash()]) - - def test_multi_source(self): - - filter = [2, 3] - args_a = tuple() - kwargs_a = { - 'count': 100, - 'sids': [1, 2, 3], - 'start': datetime(2012, 1, 3, 15, tzinfo=pytz.utc), - 'delta': timedelta(minutes=6), - 'filter': filter - } - source_a = SpecificEquityTrades(*args_a, **kwargs_a) - - args_b = tuple() - kwargs_b = { - 'count': 100, - 'sids': [2, 3, 4], - 'start': datetime(2012, 1, 3, 15, tzinfo=pytz.utc), - 'delta': timedelta(minutes=5), - 'filter': filter - } - source_b = SpecificEquityTrades(*args_b, **kwargs_b) - - all_events = list(chain(source_a, source_b)) - - # The expected output is all events, sorted by dt with - # source_id as a tiebreaker. - expected = sorted(all_events, comp) - source_ids = [source_a.get_hash(), source_b.get_hash()] - - # Generating the events list consumes the sources. Rewind them - # for testing. - source_a.rewind() - source_b.rewind() - - # Append a done message to each source. - with_done_a = chain(source_a, [done_message(source_a.get_hash())]) - with_done_b = chain(source_b, [done_message(source_b.get_hash())]) - - interleaved = alternate(with_done_a, with_done_b) - - # Test sort with alternating messages from source_a and - # source_b. - self.run_date_sort(interleaved, expected, source_ids) - - source_a.rewind() - source_b.rewind() - with_done_a = chain(source_a, [done_message(source_a.get_hash())]) - with_done_b = chain(source_b, [done_message(source_b.get_hash())]) - - sequential = chain(with_done_a, with_done_b) - - # Test sort with all messages from a, followed by all messages - # from b. - - self.run_date_sort(sequential, expected, source_ids) - - def test_sort_composite(self): - - filter = [1, 2] - - #Set up source a. One hour between events. - args_a = tuple() - kwargs_a = { - 'count': 100, - 'sids': [1], - 'start': datetime(2012, 6, 6, 0), - 'delta': timedelta(hours=1), - 'filter': filter - } - source_a = SpecificEquityTrades(*args_a, **kwargs_a) - - #Set up source b. One day between events. - args_b = tuple() - kwargs_b = { - 'count': 50, - 'sids': [2], - 'start': datetime(2012, 6, 6, 0), - 'delta': timedelta(days=1), - 'filter': filter - } - source_b = SpecificEquityTrades(*args_b, **kwargs_b) - - #Set up source c. One minute between events. - args_c = tuple() - kwargs_c = { - 'count': 150, - 'sids': [1, 2], - 'start': datetime(2012, 6, 6, 0), - 'delta': timedelta(minutes=1), - 'filter': filter - } - source_c = SpecificEquityTrades(*args_c, **kwargs_c) - # Set up source d. This should produce no events because the - # internal sids don't match the filter. - args_d = tuple() - kwargs_d = { - 'count': 50, - 'sids': [3], - 'start': datetime(2012, 6, 6, 0), - 'delta': timedelta(minutes=1), - 'filter': filter - } - source_d = SpecificEquityTrades(*args_d, **kwargs_d) - sources = [source_a, source_b, source_c, source_d] - hashes = [source.get_hash() for source in sources] - - sort_out = date_sorted_sources(*sources) - - # Read all the values from sort and assert that they arrive in - # the correct sorting with the expected hash values. - to_list = list(sort_out) - copy = to_list[:] - - # We should have 300 events (100 from a, 150 from b, 50 from c) - assert len(to_list) == 300 - - for e in to_list: - # All events should match one of our expected source_ids. - assert e.source_id in hashes - # But none of them should match source_d. - assert e.source_id != source_d.get_hash() - - # The events should be sorted by dt, with source_id as tiebreaker. - expected = sorted(copy, comp) - - assert to_list == expected - - -def compare_by_dt_source_id(x, y): - if x.dt < y.dt: - return -1 - elif x.dt > y.dt: - return 1 - - elif x.source_id < y.source_id: - return -1 - elif x.source_id > y.source_id: - return 1 - else: - return 0 - -#Alias for ease of use -comp = compare_by_dt_source_id - - -def to_dt(msg): - return ndict({'dt': msg}) diff --git a/zipline/gens/sort.py b/zipline/gens/sort.py deleted file mode 100644 index 19e33cb5..00000000 --- a/zipline/gens/sort.py +++ /dev/null @@ -1,116 +0,0 @@ -# -# 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. - -""" -Sorting generator. -""" -import logbook - -from collections import deque - -log = logbook.Logger('Sorting') - - -def date_sort(stream_in, source_ids): - """ - A generator that takes a generator and a list of source_ids. We - maintain an internal queue for each id in source_ids. While we - have messages pending from all sources, we pull the earliest - message and yield it. - """ - assert isinstance(source_ids, (list, tuple)) - - # Set up an internal queue for each expected source. - sources = {} - for id in source_ids: - assert isinstance(id, basestring), "Bad source_id %s" % id - sources[id] = deque() - - # Process incoming streams. - for message in stream_in: - # Only allow messages from sources we expect. - assert message.source_id in sources, "Unexpected source: %s" % message - - sources[message.source_id].append(message) - - # Only pop messages when we have a pending message from - # all datasources. Stop if all sources have signalled done. - - while all(sources.values()) and not done(sources): - message = pop_oldest(sources) - yield message - - # We should have only a done message left in each queue. - for queue in sources.itervalues(): - assert len(queue) == 1, "Bad queue in date_sort on exit: %s" % queue - assert queue[0].dt == "DONE", \ - "Bad last message in date_sort on exit: %s" % queue - - -def done(sources): - """Feed is done when all internal queues have only a "DONE" message.""" - assert isinstance(sources, dict) - return all((queue_is_done(source) for source in sources.itervalues())) - - -def queue_is_done(queue): - assert isinstance(queue, deque) - if len(queue) == 0: - return False - if queue[0].dt == "DONE": - assert len(queue) == 1, "Message after DONE in date_sort: %s" % queue - return True - else: - return False - - -def pop_oldest(sources): - - oldest_event = None - - # Iterate over the dict, checking internal queues for the oldest - # pending event. - - for queue in sources.itervalues(): - current_event = queue[0] - # Skip queues that are done. - if current_event.dt == "DONE": - continue - # Any event is older than nothing. - elif oldest_event is None: - oldest_event = current_event - # Keep the older event. Break ties by source_id. This will - # trip an assert if we have duplicate sources. - else: - oldest_event = older(oldest_event, current_event) - - # Pop the oldest event we found from its queue and return it. - return sources[oldest_event.source_id].popleft() - - -# Return the event with the older timestamp. Break ties by source_id. -def older(oldest, current): - # Try to compare by dt. - if oldest.dt < current.dt: - return oldest - elif oldest.dt > current.dt: - return current - # Break ties by source_id. - elif oldest.source_id < current.source_id: - return oldest - elif oldest.source_id > current.source_id: - return current - else: - assert False, "Duplicate event" diff --git a/zipline/gens/utils.py b/zipline/gens/utils.py index c3b08cb4..2cdc2581 100644 --- a/zipline/gens/utils.py +++ b/zipline/gens/utils.py @@ -17,7 +17,6 @@ import pytz import numbers -from collections import OrderedDict from hashlib import md5 from datetime import datetime from itertools import izip_longest @@ -59,27 +58,6 @@ def alternate(g1, g2): yield e2 -def roundrobin(sources, namestrings): - """ - Takes N generators, pulling one element off each until all inputs - are empty. - """ - assert len(sources) == len(namestrings) - mapping = OrderedDict(zip(namestrings, sources)) - - # While our generators have not been exhausted, pull elements - while mapping.keys() != []: - for namestring, source in mapping.iteritems(): - try: - message = source.next() - # allow sources to yield None to avoid blocking. - if message: - yield message - except StopIteration: - yield done_message(namestring) - del mapping[namestring] - - def hash_args(*args, **kwargs): """Define a unique string for any set of representable args.""" arg_string = '_'.join([str(arg) for arg in args]) From fc03e80cdf72184f644b8e47f9cf25a8ec351242 Mon Sep 17 00:00:00 2001 From: Eddie Hebert Date: Wed, 2 Jan 2013 03:35:59 -0500 Subject: [PATCH 3/4] Removes done message. Instead of checking for 'DONE' on each call uses generators builtin StopIteration for signalling the end of input. --- tests/test_perf_tracking.py | 14 +++++--------- zipline/finance/performance.py | 8 ++------ zipline/gens/composites.py | 12 +----------- zipline/gens/tradesimulation.py | 21 ++++++++++++--------- zipline/gens/utils.py | 13 ------------- 5 files changed, 20 insertions(+), 48 deletions(-) diff --git a/tests/test_perf_tracking.py b/tests/test_perf_tracking.py index 7cab2518..b7619441 100644 --- a/tests/test_perf_tracking.py +++ b/tests/test_perf_tracking.py @@ -27,7 +27,6 @@ from operator import attrgetter import zipline.utils.factory as factory import zipline.finance.performance as perf from zipline.utils.protocol_utils import ndict -from zipline.protocol import Event from zipline.gens.composites import date_sorted_sources @@ -666,13 +665,6 @@ class TestPerformanceTracker(unittest.TestCase): # Extract events with transactions to use for verification. events_with_txns = [event for event in events if event.TRANSACTION] - done_message = Event({ - 'dt': 'DONE', - 'TRANSACTION': None - }) - - events = itertools.chain(events, [done_message]) - perf_messages = \ [msg for date, snapshot in perf_tracker.transform( @@ -680,6 +672,10 @@ class TestPerformanceTracker(unittest.TestCase): for event in snapshot for msg in event.perf_messages] + end_perf_messages, risk_message = perf_tracker.handle_simulation_end() + + perf_messages.extend(end_perf_messages) + #we skip two trades, to test case of None transaction self.assertEqual(perf_tracker.txn_count, len(events_with_txns)) @@ -696,7 +692,7 @@ class TestPerformanceTracker(unittest.TestCase): def event_with_txn(self, event, no_txn_dt): #create a transaction for all but #first trade in each sid, to simulate None transaction - if event.dt != no_txn_dt and event.dt != 'DONE': + if event.dt != no_txn_dt: txn = ndict({ 'sid': event.sid, 'amount': -25, diff --git a/zipline/finance/performance.py b/zipline/finance/performance.py index 67428ede..d8776220 100644 --- a/zipline/finance/performance.py +++ b/zipline/finance/performance.py @@ -219,12 +219,8 @@ class PerformanceTracker(object): new_snapshot = [] for event in snapshot: - if date != "DONE": - event.perf_messages = self.process_event(event) - event.portfolio = self.get_portfolio() - else: - event.perf_messages, event.risk_message = \ - self.handle_simulation_end() + event.perf_messages = self.process_event(event) + event.portfolio = self.get_portfolio() del event['TRANSACTION'] new_snapshot.append(event) diff --git a/zipline/gens/composites.py b/zipline/gens/composites.py index 04824c3a..08543593 100644 --- a/zipline/gens/composites.py +++ b/zipline/gens/composites.py @@ -15,10 +15,6 @@ import heapq -from itertools import chain - -from zipline.gens.utils import done_message - def decorate_source(source): for message in source: @@ -55,8 +51,7 @@ def sequential_transforms(stream_in, *transforms): transforms, stream_in) - dt_aliased = alias_dt(stream_out) - return add_done(dt_aliased) + return alias_dt(stream_out) def alias_dt(stream_in): @@ -66,8 +61,3 @@ def alias_dt(stream_in): for message in stream_in: message['datetime'] = message['dt'] yield message - - -# Add a done message to a stream. -def add_done(stream_in): - return chain(stream_in, [done_message('Composite')]) diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index 7616ab48..4cf41196 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -71,6 +71,7 @@ class TradeSimulationClient(object): self.algo_start = self.environment.first_open self.algo_sim = AlgorithmSimulator( self.ordering_client, + self.perf_tracker, self.algo, self.algo_start ) @@ -116,6 +117,7 @@ class AlgorithmSimulator(object): def __init__(self, order_book, + perf_tracker, algo, algo_start): @@ -126,6 +128,7 @@ class AlgorithmSimulator(object): # We extract the order book from the txn client so that # the algo can place new orders. self.order_book = order_book + self.perf_tracker = perf_tracker self.algo = algo self.algo_start = algo_start.replace(hour=0, minute=0, @@ -203,18 +206,10 @@ class AlgorithmSimulator(object): if self.simulation_dt is None: self.simulation_dt = date - # Done message has the risk report, so we yield before exiting. - if date == 'DONE': - for event in snapshot: - for perf_message in event.perf_messages: - yield perf_message - yield event.risk_message - raise StopIteration - # We're still in the warmup period. Use the event to # update our universe, but don't yield any perf messages, # and don't send a snapshot to handle_data. - elif date < self.algo_start: + if date < self.algo_start: for event in snapshot: del event['perf_messages'] self.update_universe(event) @@ -233,6 +228,14 @@ class AlgorithmSimulator(object): # to the user's algo. self.simulate_snapshot(date) + perf_messages, risk_message = \ + self.perf_tracker.handle_simulation_end() + + for message in perf_messages: + yield message + + yield risk_message + def update_universe(self, event): """ Update the universe with new event information. diff --git a/zipline/gens/utils.py b/zipline/gens/utils.py index 2cdc2581..dd72cb41 100644 --- a/zipline/gens/utils.py +++ b/zipline/gens/utils.py @@ -20,7 +20,6 @@ import numbers from hashlib import md5 from datetime import datetime from itertools import izip_longest -from zipline import ndict from zipline.protocol import ( DATASOURCE_TYPE, Event @@ -37,18 +36,6 @@ def mock_raw_event(sid, dt): return event -def mock_done(id): - return ndict({ - 'dt': "DONE", - "source_id": id, - 'tnfm_id': id, - 'tnfm_value': None, - 'type': DATASOURCE_TYPE.DONE - }) - -done_message = mock_done - - def alternate(g1, g2): """Specialized version of roundrobin for just 2 generators.""" for e1, e2 in izip_longest(g1, g2): From ee02ff644573d02105eacb128ec06f9ed99951d5 Mon Sep 17 00:00:00 2001 From: Eddie Hebert Date: Mon, 7 Jan 2013 13:03:12 -0500 Subject: [PATCH 4/4] Adds _ prefix to decorate_source to imply internal usage. --- zipline/gens/composites.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zipline/gens/composites.py b/zipline/gens/composites.py index 08543593..e58390ed 100644 --- a/zipline/gens/composites.py +++ b/zipline/gens/composites.py @@ -16,7 +16,7 @@ import heapq -def decorate_source(source): +def _decorate_source(source): for message in source: yield ((message.dt, message.source_id), message) @@ -26,7 +26,7 @@ def date_sorted_sources(*sources): Takes an iterable of sources, generating namestrings and piping their output into date_sort. """ - sorted_stream = heapq.merge(*(decorate_source(s) for s in sources)) + sorted_stream = heapq.merge(*(_decorate_source(s) for s in sources)) # Strip out key decoration for _, message in sorted_stream: