diff --git a/tests/client.py b/tests/client.py index 03874b95..c2c07560 100644 --- a/tests/client.py +++ b/tests/client.py @@ -9,10 +9,6 @@ LOGGER = logging.getLogger('ZiplineLogger') class TestClient(Component): - def __init__(self): - Component.__init__(self) - self.init() - def init(self): self.received_count = 0 self.prev_dt = None @@ -48,10 +44,10 @@ class TestClient(Component): def do_work(self): socks = dict(self.poll.poll(self.heartbeat_timeout)) - if self.control_in in socks and socks[self.control_in] == self.zmq.POLLIN: + if socks.get(self.control_in) == self.zmq.POLLIN: msg = self.control_in.recv() - if self.data_feed in socks and socks[self.data_feed] == self.zmq.POLLIN: + if socks.get(self.data_feed) == self.zmq.POLLIN: msg = self.data_feed.recv() #logger.info('msg:' + str(msg)) diff --git a/tests/test_finance.py b/tests/test_finance.py index bde2d1a5..ac2938c5 100644 --- a/tests/test_finance.py +++ b/tests/test_finance.py @@ -26,16 +26,16 @@ EXTENDED_TIMEOUT = 90 allocator = AddressAllocator(1000) class FinanceTestCase(TestCase): - + leased_sockets = defaultdict(list) - + def setUp(self): #qutil.configure_logging() self.zipline_test_config = { 'allocator':allocator, 'sid':133 } - + @timed(DEFAULT_TIMEOUT) def test_factory_daily(self): trading_environment = factory.create_trading_environment() @@ -49,12 +49,12 @@ class FinanceTestCase(TestCase): if prev: self.assertTrue(trade.dt > prev.dt) prev = trade - + @timed(DEFAULT_TIMEOUT) def test_trading_environment(self): benchmark_returns, treasury_curves = \ factory.load_market_data() - + env = TradingEnvironment( benchmark_returns, treasury_curves, @@ -88,90 +88,90 @@ class FinanceTestCase(TestCase): a_saturday, a_sunday ] - + for holiday in holidays: self.assertTrue(not env.is_trading_day(holiday)) - + first_trading_day = datetime(2008, 1, 2, tzinfo = pytz.utc) last_trading_day = datetime(2008, 12, 31, tzinfo = pytz.utc) workdays = [first_trading_day, last_trading_day] - + for workday in workdays: self.assertTrue(env.is_trading_day(workday)) - + self.assertTrue(env.last_close.month == 12) self.assertTrue(env.last_close.day == 31) - + # The following two tests appear broken no that the order source is # non blocking. HUNCH: The trades are streaming through before the orders # are placed. - + @timed(DEFAULT_TIMEOUT) def test_orders(self): - + # Simulation # ---------- - + self.zipline_test_config['simulation_style'] = \ SIMULATION_STYLE.FIXED_SLIPPAGE zipline = SimulatedTrading.create_test_zipline( **self.zipline_test_config ) zipline.simulate(blocking=True) - + self.assertTrue(zipline.sim.ready()) self.assertFalse(zipline.sim.exception) - + # TODO: Make more assertions about the final state of the components. self.assertEqual(zipline.sim.feed.pending_messages(), 0, \ "The feed should be drained of all messages, found {n} remaining." \ .format(n=zipline.sim.feed.pending_messages())) - + # the trading client should receive one transaction for every # order placed. self.assertEqual( zipline.trading_client.txn_count, zipline.trading_client.order_count ) - - + + @timed(DEFAULT_TIMEOUT) def test_aggressive_buying(self): - + # Simulation # ---------- - + # TODO: for some reason the orders aren't filled without an extra # trade. trade_count = 5 self.zipline_test_config['order_count'] = trade_count - 1 self.zipline_test_config['trade_count'] = trade_count self.zipline_test_config['order_amount'] = 1 - + # tell the simulator to fill the orders in individual transactions # matching the order volume exactly. self.zipline_test_config['simulation_style'] = \ SIMULATION_STYLE.FIXED_SLIPPAGE self.zipline_test_config['environment'] = factory.create_trading_environment() - + sid_list = [self.zipline_test_config['sid']] - + self.zipline_test_config['trade_source'] = factory.create_minutely_trade_source( sid_list, trade_count, self.zipline_test_config['environment'] ) - + zipline = SimulatedTrading.create_test_zipline(**self.zipline_test_config) zipline.simulate(blocking=True) - + self.assertTrue(zipline.sim.ready()) self.assertFalse(zipline.sim.exception) - + self.assertEqual(zipline.sim.feed.pending_messages(), 0, \ "The feed should be drained of all messages, found {n} remaining." \ .format(n=zipline.sim.feed.pending_messages())) - + # # the trading client should receive one transaction for every # order placed. @@ -179,9 +179,9 @@ class FinanceTestCase(TestCase): zipline.trading_client.txn_count, zipline.trading_client.order_count ) - - - + + + @timed(DEFAULT_TIMEOUT) def test_performance(self): #provide enough trades to ensure all orders are filled. @@ -189,27 +189,27 @@ class FinanceTestCase(TestCase): self.zipline_test_config['trade_count'] = 200 zipline = SimulatedTrading.create_test_zipline(**self.zipline_test_config) zipline.simulate(blocking=True) - + self.assertEqual( zipline.sim.feed.pending_messages(), 0, "The feed should be drained of all messages, found {n} remaining." \ .format(n=zipline.sim.feed.pending_messages()) ) - + self.assertEqual( zipline.sim.merge.pending_messages(), 0, "The merge should be drained of all messages, found {n} remaining." \ .format(n=zipline.sim.merge.pending_messages()) ) - + self.assertEqual( zipline.algorithm.count, zipline.algorithm.incr, "The test algorithm should send as many orders as specified.") - - + + transaction_sim = zipline.trading_client.txn_sim self.assertEqual( transaction_sim.txn_count, @@ -217,32 +217,32 @@ class FinanceTestCase(TestCase): "The perf tracker should handle the same number of transactions \ as the simulator emits." ) - + self.assertEqual( len(zipline.get_positions()), 1, "Portfolio should have one position." ) - + SID = self.zipline_test_config['sid'] self.assertEqual( zipline.get_positions()[SID]['sid'], SID, "Portfolio should have one position in " + str(SID) ) - + self.assertEqual( zipline.sources['flat'].count, self.zipline_test_config['trade_count'], "The simulated trade source should send all trades." ) - + self.assertEqual( zipline.algorithm.frame_count, self.zipline_test_config['trade_count'], "The algorithm should receive all trades." ) - + @timed(DEFAULT_TIMEOUT) def test_sid_filter(self): """Ensure the algorithm's filter prevents events from arriving.""" @@ -256,14 +256,14 @@ class FinanceTestCase(TestCase): order_amount, order_count ) - + self.zipline_test_config['trade_count'] = 200 self.zipline_test_config['algorithm'] = test_algo - + zipline = SimulatedTrading.create_test_zipline( **self.zipline_test_config ) - + zipline.simulate(blocking=True) #check that the algorithm received no events self.assertEqual( @@ -271,14 +271,14 @@ class FinanceTestCase(TestCase): test_algo.frame_count, "The algorithm should not receive any events due to filtering." ) - - + + # TODO: write tests for short sales # TODO: write a test to do massive buying or shorting. - + @timed(DEFAULT_TIMEOUT) def test_partially_filled_orders(self): - + # create a scenario where order size and trade size are equal # so that orders must be spread out over several trades. params ={ @@ -294,9 +294,9 @@ class FinanceTestCase(TestCase): 'expected_txn_count':8, 'expected_txn_volume':2 * 100 } - + self.transaction_sim(**params) - + # same scenario, but with short sales params2 ={ 'trade_count':360, @@ -308,9 +308,9 @@ class FinanceTestCase(TestCase): 'expected_txn_count':8, 'expected_txn_volume':2 * -100 } - + self.transaction_sim(**params2) - + @timed(DEFAULT_TIMEOUT) def test_collapsing_orders(self): # create a scenario where order.amount <<< trade.volume @@ -328,7 +328,7 @@ class FinanceTestCase(TestCase): 'expected_txn_volume':24 * 1 } self.transaction_sim(**params1) - + # second verse, same as the first. except short! params2 ={ 'trade_count':6, @@ -341,7 +341,7 @@ class FinanceTestCase(TestCase): 'expected_txn_volume':24 * -1 } self.transaction_sim(**params2) - + @timed(DEFAULT_TIMEOUT) def test_partial_expiration_orders(self): # create a scenario where orders expire without being filled @@ -360,7 +360,7 @@ class FinanceTestCase(TestCase): 'expected_txn_volume' : 25 } self.transaction_sim(**params1) - + # same scenario, but short sales. params2 = { 'trade_count' : 100, @@ -376,7 +376,7 @@ class FinanceTestCase(TestCase): 'expected_txn_volume' : -25 } self.transaction_sim(**params2) - + @timed(DEFAULT_TIMEOUT) def test_alternating_long_short(self): # create a scenario where we alternate buys and sells @@ -393,9 +393,9 @@ class FinanceTestCase(TestCase): 'expected_txn_volume' : 0 #equal buys and sells } self.transaction_sim(**params1) - + def transaction_sim(self, **params): - + trade_count = params['trade_count'] trade_amount = params['trade_amount'] trade_interval = params['trade_interval'] @@ -411,14 +411,14 @@ class FinanceTestCase(TestCase): alternate = params.get('alternate') # if present, expect transaction amounts to match orders exactly. complete_fill = params.get('complete_fill') - + trading_environment = factory.create_trading_environment() trade_sim = TransactionSimulator() price = [10.1] * trade_count volume = [100] * trade_count start_date = trading_environment.first_open sid = 1 - + generated_trades = factory.create_trade_history( sid, price, @@ -426,12 +426,12 @@ class FinanceTestCase(TestCase): trade_interval, trading_environment ) - + if alternate: alternator = -1 else: alternator = 1 - + order_date = start_date for i in xrange(order_count): order = ndict( @@ -441,9 +441,9 @@ class FinanceTestCase(TestCase): 'type' : zp.DATASOURCE_TYPE.ORDER, 'dt' : order_date }) - + trade_sim.add_open_order(order) - + order_date = order_date + order_interval # move after market orders to just after market next # market open. @@ -451,40 +451,40 @@ class FinanceTestCase(TestCase): if order_date.minute >= 00: order_date = order_date + timedelta(days=1) order_date = order_date.replace(hour=14, minute=30) - + # there should now be one open order list stored under the sid oo = trade_sim.open_orders self.assertEqual(len(oo), 1) self.assertTrue(oo.has_key(sid)) order_list = oo[sid] self.assertEqual(order_count, len(order_list)) - + for i in xrange(order_count): order = order_list[i] self.assertEqual(order.sid, sid) self.assertEqual(order.amount, order_amount * alternator**i) - + tracker = PerformanceTracker(trading_environment) - + # this approximates the loop inside TradingSimulationClient transactions = [] for trade in generated_trades: if trade_delay: trade.dt = trade.dt + trade_delay - + txn = trade_sim.apply_trade_to_open_orders(trade) if txn: transactions.append(txn) trade.TRANSACTION = txn else: trade.TRANSACTION = None - + tracker.process_event(trade) - + if complete_fill: self.assertEqual(len(transactions), len(order_list)) - + total_volume = 0 for i in xrange(len(transactions)): txn = transactions[i] @@ -492,18 +492,15 @@ class FinanceTestCase(TestCase): if complete_fill: order = order_list[i] self.assertEqual(order.amount, txn.amount) - + self.assertEqual(total_volume, expected_txn_volume) self.assertEqual(len(transactions), expected_txn_count) - + cumulative_pos = tracker.cumulative_performance.positions[sid] self.assertEqual(total_volume, cumulative_pos.amount) - + # the open orders should now be empty oo = trade_sim.open_orders self.assertTrue(oo.has_key(sid)) order_list = oo[sid] self.assertEqual(0, len(order_list)) - - - diff --git a/zipline/components/aggregator.py b/zipline/components/aggregator.py new file mode 100644 index 00000000..d0b18208 --- /dev/null +++ b/zipline/components/aggregator.py @@ -0,0 +1,210 @@ +""" + Component + | + Aggregate + | + / \ + Feed Merge + +""" +import logging +from collections import Counter + +import zipline.protocol as zp + +from zipline.core.component import Component +from zipline.protocol import CONTROL_PROTOCOL, COMPONENT_TYPE, \ + CONTROL_FRAME, CONTROL_UNFRAME + +LOGGER = logging.getLogger('ZiplineLogger') + +class Aggregate(Component): + """ + Abstract superclass to Merge & Feed. Acts on two sockets + + - pull_socket + - feed_socket + + Feed and Merge define these differently. + """ + + def init(self): + self.sent_count = 0 + self.received_count = 0 + self.draining = False + self.ds_finished_counter = 0 + + # Depending on the size of this, might want to use a data + # structure with better asymptotics. + self.data_buffer = {} + + # source_id -> integer count + self.sent_counters = Counter() + self.recv_counters = Counter() + + @property + def get_id(self): + raise NotImplementedError + + @property + def get_type(self): + return COMPONENT_TYPE.CONDUIT + + # ------------- + # Core Methods + # ------------- + + def open(self): + self.pull_socket = self.bind_data() + self.feed_socket = self.bind_feed() + + def do_work(self): + # wait for synchronization reply from the host + socks = dict(self.poll.poll(self.heartbeat_timeout)) + + # TODO: Abstract this out, maybe on base component + if socks.get(self.control_in) == self.zmq.POLLIN: + msg = self.control_in.recv() + event, payload = CONTROL_UNFRAME(msg) + + # -- Heartbeat -- + if event == CONTROL_PROTOCOL.HEARTBEAT: + # Heart outgoing + heartbeat_frame = CONTROL_FRAME( + CONTROL_PROTOCOL.OK, + payload + ) + self.control_out.send(heartbeat_frame) + + # -- Soft Kill -- + elif event == CONTROL_PROTOCOL.SHUTDOWN: + self.signal_done() + self.shutdown() + + # -- Hard Kill -- + elif event == CONTROL_PROTOCOL.KILL: + self.kill() + + + if socks.get(self.pull_socket) == self.zmq.POLLIN: + message = self.pull_socket.recv() + + if message == str(CONTROL_PROTOCOL.DONE): + self.ds_finished_counter += 1 + + if len(self.data_buffer) == self.ds_finished_counter: + #drain any remaining messages in the buffer + LOGGER.debug("draining feed") + self.drain() + self.signal_done() + else: + try: + event = self.unframe(message) + # deserialization error + except zp.INVALID_DATASOURCE_FRAME as exc: + return self.signal_exception(exc) + + try: + self.append(event) + self.send_next() + + # Invalid message + except zp.INVALID_DATASOURCE_FRAME as exc: + return self.signal_exception(exc) + + # ------------- + # Flow Control + # ------------- + + def drain(self): + """ + Send all messages in the buffer. + """ + self.draining = True + while self.pending_messages() > 0: + self.send_next() + + def send_next(self): + """ + Send the (chronologically) next message in the buffer. + """ + if not (self.is_full() or self.draining): + return + + # TODO: implement this in __iter__ + event = self.next() + if event is not None: + self.feed_socket.send(self.frame(event), self.zmq.NOBLOCK) + self.sent_counters[event.source_id] += 1 + self.sent_count += 1 + + def append(self, event): + """ + Add an event to the buffer for the source specified by + source_id. + """ + self.data_buffer[event.source_id].append(event) + self.recv_counters[event.source_id] += 1 + self.received_count += 1 + + def next(self): + """ + Get the next message in chronological order. + """ + if not(self.is_full() or self.draining): + return + + cur_source = None + earliest_source = None + earliest_event = None + #iterate over the queues of events from all sources + #(1 queue per datasource) + for events in self.data_buffer.itervalues(): + if len(events) == 0: + continue + cur_source = events + first_in_list = events[0] + if first_in_list.dt == None: + #this is a filler event, discard + events.pop(0) + continue + + if (earliest_event == None) or (first_in_list.dt <= earliest_event.dt): + earliest_event = first_in_list + earliest_source = cur_source + + if earliest_event != None: + return earliest_source.pop(0) + + def is_full(self): + """ + Indicates whether the buffer has messages in buffer for + all un-DONE, blocking sources. + """ + for source_id, events in self.data_buffer.iteritems(): + if len(events) == 0: + return False + return True + + def pending_messages(self): + """ + Returns the count of all events from all sources in the + buffer. + """ + total = 0 + for events in self.data_buffer.itervalues(): + total += len(events) + return total + + def add_source(self, source_id): + """ + Add a data source to the buffer. + """ + self.data_buffer[source_id] = [] + + def __len__(self): + """ + Buffer's length is same as internal map holding separate + sorted arrays of events keyed by source id. + """ + return len(self.data_buffer) diff --git a/zipline/components/datasource.py b/zipline/components/datasource.py index 8c14022b..d6238e9b 100644 --- a/zipline/components/datasource.py +++ b/zipline/components/datasource.py @@ -12,9 +12,9 @@ LOGGER = logging.getLogger('ZiplineLogger') class DataSource(Component): """ - Baseclass for data sources. Subclass and implement send_all - usually this - means looping through all records in a store, converting to a dict, and - calling send(map). + Abstract baseclass for data sources. Subclass and implement send_all - + usually this means looping through all records in a store, converting + to a dict, and calling send(map). Every datasource has a dict property to hold filters:: - key -- name of the filter, e.g. SID @@ -22,14 +22,10 @@ class DataSource(Component): Modify the datasource's filters via the set_filter(name, value) """ - def __init__(self, source_id): - Component.__init__(self) + def init(self, source_id): self.id = source_id - self.init() self.filter = {} - - def init(self): self.cur_event = None def set_filter(self, name, value): @@ -37,7 +33,17 @@ class DataSource(Component): @property def get_id(self): - return self.id + """ + Returns this component id, this is fixed at a class + level. This should not and cannot be contingent on + arguments to the init function. Examples: + + - "TradeDataSource" + - "RandomEquityTrades" + - "SpecificEquityTrades" + + """ + raise NotImplementedError @property def get_type(self): diff --git a/zipline/components/feed.py b/zipline/components/feed.py index 191a1cc9..42c08c73 100644 --- a/zipline/components/feed.py +++ b/zipline/components/feed.py @@ -49,7 +49,7 @@ class Feed(Component): def do_work(self): # wait for synchronization reply from the host - socks = dict(self.poll.poll(self.heartbeat_timeout)) + socks = dict(self.poll.poll(self.heartbeat_timeout)) # TODO: Abstract this out, maybe on base component if self.control_in in socks and socks[self.control_in] == self.zmq.POLLIN: @@ -151,7 +151,7 @@ class Feed(Component): cur_source = None earliest_source = None earliest_event = None - #iterate over the queues of events from all sources + #iterate over the queues of events from all sources #(1 queue per datasource) for events in self.data_buffer.values(): if len(events) == 0: diff --git a/zipline/components/merge.py b/zipline/components/merge.py index 2811912e..0ef29bc2 100644 --- a/zipline/components/merge.py +++ b/zipline/components/merge.py @@ -61,7 +61,7 @@ class Merge(Feed): def append(self, event): """ - :param event: a ndict with one entry. key is the name of the + :param event: a ndict with one entry. key is the name of the transform, value is the transformed value. Add an event to the buffer for the source specified by source_id. @@ -69,4 +69,3 @@ class Merge(Feed): self.data_buffer[event.keys()[0]].append(event) self.received_count += 1 - diff --git a/zipline/finance/sources.py b/zipline/finance/sources.py index 6f7b3e55..889e7f01 100644 --- a/zipline/finance/sources.py +++ b/zipline/finance/sources.py @@ -13,6 +13,10 @@ import zipline.protocol as zp class TradeDataSource(DataSource): + @property + def get_id(self): + return 'TradeDataSource' + def send(self, event): """ Sends the event iff it matches the internal SID filter. @@ -48,7 +52,11 @@ class RandomEquityTrades(TradeDataSource): self.day = datetime.timedelta(days=1) self.price = random.uniform(5.0, 50.0) + @property + def get_id(self): + return 'RandomEquityTrades' + @property def get_type(self): zp.COMPONENT_TYPE.SOURCE @@ -100,12 +108,7 @@ class SpecificEquityTrades(TradeDataSource): @property def get_id(self): - """ - The descriptive name of the component. - """ - # Prevents the bug that Thomas ran into - return "Unique ID" - + return "SpecificEquityTrades" def do_work(self): if(len(self.event_list) == 0): diff --git a/zipline/lines.py b/zipline/lines.py index b606c458..b7c26b3c 100644 --- a/zipline/lines.py +++ b/zipline/lines.py @@ -191,7 +191,7 @@ class SimulatedTrading(object): - algorithm - optional parameter providing an algorithm. defaults to :py:class:`zipline.test.algorithms.TestAlgorithm` - trade_source - optional parameter to specify trades, if present. - If not present :py:class:`ziplien.sources.SpecificEquityTrades` + If not present :py:class:`zipline.sources.SpecificEquityTrades` is the source, with daily frequency in trades. - simulation_style: optional parameter that configures the :py:class:`zipline.finance.trading.TransactionSimulator`. Expects @@ -264,11 +264,11 @@ class SimulatedTrading(object): # Simulation #------------------- zipline = SimulatedTrading(**{ - 'algorithm':test_algo, - 'trading_environment':trading_environment, - 'allocator':allocator, - 'simulator_class':simulator_class, - 'simulation_style':simulation_style + 'algorithm' : test_algo, + 'trading_environment' : trading_environment, + 'allocator' : allocator, + 'simulator_class' : simulator_class, + 'simulation_style' : simulation_style }) #------------------- @@ -285,7 +285,9 @@ class SimulatedTrading(object): self.check_started() source.set_filter('SID', self.algorithm.get_sid_filter()) self.sim.register_components([source]) - self.sources[source.get_id] = source + + # ``id`` is name of source_id, ``get_id`` is the class name + self.sources[source.id] = source def add_transform(self, transform): assert isinstance(transform, BaseTransform)