From b69ea6b790a8c207eb7444957d5cc20e966ecdb6 Mon Sep 17 00:00:00 2001 From: fawce Date: Thu, 8 Mar 2012 19:21:28 -0500 Subject: [PATCH] refactoring performance to be a plain class, not a component. --- zipline/finance/performance.py | 55 +++++++++++++++----------------- zipline/finance/trading.py | 26 +++++++++------ zipline/test/client.py | 13 ++++---- zipline/test/test_finance.py | 58 +++++++++++++++++++++++++++------- 4 files changed, 95 insertions(+), 57 deletions(-) diff --git a/zipline/finance/performance.py b/zipline/finance/performance.py index 8f3d9aa1..67b05be6 100644 --- a/zipline/finance/performance.py +++ b/zipline/finance/performance.py @@ -9,10 +9,9 @@ import zipline.util as qutil import zipline.protocol as zp import zipline.finance.risk as risk -class PortfolioClient(qmsg.Component): +class PerformanceTracker(): def __init__(self, period_start, period_end, capital_base, trading_environment): - qmsg.Component.__init__(self) self.trading_day = datetime.timedelta(hours=6, minutes=30) self.calendar_day = datetime.timedelta(hours=24) self.period_start = period_start @@ -27,35 +26,33 @@ class PortfolioClient(qmsg.Component): self.capital_base = capital_base self.trading_environment = trading_environment self.returns = [] - self.cumulative_performance = PerformancePeriod(self.period_start, self.period_end, {}, 0, capital_base = capital_base) - self.todays_performance = PerformancePeriod(self.market_open, self.market_close, {}, 0, capital_base = capital_base) - - @property - def get_id(self): - return str(zp.FINANCE_COMPONENT.PORTFOLIO_CLIENT) - - def open(self): - self.result_feed = self.connect_result() - - def do_work(self): - #next feed event - socks = dict(self.poll.poll(self.heartbeat_timeout)) - - if self.result_feed in socks and socks[self.result_feed] == self.zmq.POLLIN: - msg = self.result_feed.recv() - - if msg == str(zp.CONTROL_PROTOCOL.DONE): - self.handle_simulation_end() - qutil.LOGGER.info("Portfolio Client is DONE!") - self.signal_done() - return - - event = zp.MERGE_UNFRAME(msg) + self.txn_count = 0 + self.event_count = 0 + self.cumulative_performance = PerformancePeriod( + self.period_start, + self.period_end, + {}, + capital_base, + capital_base = capital_base + ) + self.todays_performance = PerformancePeriod( + self.market_open, + self.market_close, + {}, + capital_base, + capital_base = capital_base + ) + + + + def update(self, event): + self.event_count += 1 if(event.dt >= self.market_close): self.handle_market_close() - if event.TRANSACTION: + if event.TRANSACTION != None: + self.txn_count += 1 self.cumulative_performance.execute_transaction(event.TRANSACTION) self.todays_performance.execute_transaction(event.TRANSACTION) @@ -73,9 +70,7 @@ class PortfolioClient(qmsg.Component): #calculate performance as of last trade self.cumulative_performance.calculate_performance() self.todays_performance.calculate_performance() - - - + def handle_market_close(self): self.market_open = self.market_open + self.calendar_day while not self.trading_environment.is_trading_day(self.market_open): diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py index 431323c9..34c19b62 100644 --- a/zipline/finance/trading.py +++ b/zipline/finance/trading.py @@ -15,11 +15,19 @@ class TradeSimulationClient(qmsg.Component): self.received_count = 0 self.prev_dt = None self.event_queue = [] + self.event_callbacks = [] @property def get_id(self): return str(zp.FINANCE_COMPONENT.TRADING_CLIENT) + def add_event_callback(self, callback): + """ + :param callable callback: must be a function with the signature + f(frame). + """ + self.event_callbacks.append(callback) + def open(self): self.result_feed = self.connect_result() self.order_socket = self.connect_order() @@ -39,19 +47,15 @@ class TradeSimulationClient(qmsg.Component): return event = zp.MERGE_UNFRAME(msg) - self._handle_event(event) + for cb in self.event_callbacks: + cb(event) + + #signal done to order source. + self.order_socket.send(str(zp.ORDER_PROTOCOL.BREAK)) def connect_order(self): return self.connect_push_socket(self.addresses['order_address']) - def _handle_event(self, event): - self.handle_event(event) - #signal done to order source. - self.order_socket.send(str(zp.ORDER_PROTOCOL.BREAK)) - - def handle_event(self, event): - raise NotImplementedError - def order(self, sid, amount): self.order_socket.send(zp.ORDER_FRAME(sid, amount)) @@ -151,6 +155,7 @@ class TransactionSimulator(qmsg.BaseTransform): qmsg.BaseTransform.__init__(self, zp.TRANSFORM_TYPE.TRANSACTION) self.open_orders = {} self.order_count = 0 + self.txn_count = 0 self.trade_windwo = datetime.timedelta(seconds=30) self.orderTTL = datetime.timedelta(days=1) self.volume_share = 0.05 @@ -231,7 +236,8 @@ class TransactionSimulator(qmsg.BaseTransform): return self.create_transaction(event.sid, amount, event.price + impact, dt.replace(tzinfo = pytz.utc), direction) - def create_transaction(self, sid, amount, price, dt, direction): + def create_transaction(self, sid, amount, price, dt, direction): + self.txn_count += 1 txn = {'sid' : sid, 'amount' : int(amount), 'dt' : dt, diff --git a/zipline/test/client.py b/zipline/test/client.py index 205cabff..ccc6e397 100644 --- a/zipline/test/client.py +++ b/zipline/test/client.py @@ -66,10 +66,11 @@ class TestClient(qmsg.Component): return zp.MERGE_UNFRAME(msg) -class TestTradingClient(TradeSimulationClient): +class TestAlgorithm(): - def __init__(self, sid, amount, order_count): - TradeSimulationClient.__init__(self) + def __init__(self, sid, amount, order_count, trading_client): + self.trading_client = trading_client + self.trading_client.add_event_callback(self.handle_event) self.count = order_count self.sid = sid self.amount = amount @@ -78,8 +79,8 @@ class TestTradingClient(TradeSimulationClient): def handle_event(self, event): #place an order for 100 shares of sid:133 if(self.incr < self.count): - self.order(self.sid, self.amount) + self.trading_client.order(self.sid, self.amount) self.incr += 1 else: - self.signal_order_done() - self.signal_done() + self.trading_client.signal_order_done() + self.trading_client.signal_done() diff --git a/zipline/test/test_finance.py b/zipline/test/test_finance.py index 829f5fa2..8d34938d 100644 --- a/zipline/test/test_finance.py +++ b/zipline/test/test_finance.py @@ -10,9 +10,10 @@ import zipline.finance.risk as risk import zipline.protocol as zp import zipline.finance.performance as perf -from zipline.test.client import TestTradingClient +from zipline.test.client import TestAlgorithm from zipline.sources import SpecificEquityTrades -from zipline.finance.trading import TransactionSimulator, OrderDataSource +from zipline.finance.trading import TransactionSimulator, OrderDataSource, \ +TradeSimulationClient from zipline.simulator import AddressAllocator, Simulator from zipline.monitor import Controller @@ -172,15 +173,21 @@ class FinanceTestCase(TestCase): ) set1 = SpecificEquityTrades("flat-133", trade_history) - - #client sill send 10 orders for 100 shares of 133 - client = TestTradingClient(133, 100, 10) + + trading_client = TradeSimulationClient() + #client will send 10 orders for 100 shares of 133 + test_algo = TestAlgorithm(133, 100, 10, trading_client) ts = datetime.strptime("02/1/2012","%m/%d/%Y").replace(tzinfo=pytz.utc) order_source = OrderDataSource(ts) transaction_sim = TransactionSimulator() - sim.register_components([client, order_source, transaction_sim, set1]) + sim.register_components([ + trading_client, + order_source, + transaction_sim, + set1 + ]) sim.register_controller( con ) # Simulation @@ -242,24 +249,27 @@ class FinanceTestCase(TestCase): set1 = SpecificEquityTrades("flat-133", trade_history) #client sill send 10 orders for 100 shares of 133 - client = TestTradingClient(133, 100, 10) + trading_client = TradeSimulationClient() + test_algo = TestAlgorithm(133, 100, 10, trading_client) ts = datetime.strptime("02/1/2012","%m/%d/%Y") ts = ts.replace(tzinfo=pytz.utc) order_source = OrderDataSource(ts) transaction_sim = TransactionSimulator() - portfolio_client = perf.PortfolioClient( + perf_tracker = perf.PerformanceTracker( trade_history[0]['dt'], trade_history[-1]['dt'], 1000000.0, self.trading_environment) + + #register perf_tracker to receive callbacks from the client. + trading_client.add_event_callback(perf_tracker.update) sim.register_components([ - client, + trading_client, order_source, transaction_sim, set1, - portfolio_client, ]) sim.register_controller( con ) @@ -268,8 +278,34 @@ class FinanceTestCase(TestCase): sim_context = sim.simulate() sim_context.join() + # TODO: Make more assertions about the final state of the components. self.assertEqual(sim.feed.pending_messages(), 0, \ "The feed should be drained of all messages, found {n} remaining." \ - .format(n=sim.feed.pending_messages())) \ No newline at end of file + .format(n=sim.feed.pending_messages())) + + self.assertEqual( + order_source.sent_count, + test_algo.count, + "The order source should have sent as many orders as the algo." + ) + + self.assertEqual( + transaction_sim.txn_count, + perf_tracker.txn_count, + "The perf tracker should handle the same number of transactions as\ + as the simulator emits." + ) + + self.assertEqual( + len(perf_tracker.cumulative_performance.positions), + 1, + "Portfolio should have one position." + ) + + self.assertEqual( + perf_tracker.cumulative_performance.positions[133].sid, + 133, + "Portfolio should have one position in 133." + ) \ No newline at end of file