From 26da4316d1c4a05c07e7ad3fec7e362aac75dcb5 Mon Sep 17 00:00:00 2001 From: fawce Date: Tue, 20 Mar 2012 15:51:10 -0400 Subject: [PATCH] documented algorithm protocol, split example algorithms into their own module. added filter methods to algo and to datasources. The top level zipline is responsible for piping algo's filter to the datasource. --- zipline/finance/trading.py | 25 ++++++----- zipline/lines.py | 53 +++++++++++------------ zipline/messaging.py | 10 +++++ zipline/test/algorithms.py | 82 ++++++++++++++++++++++++++++++++++++ zipline/test/client.py | 33 --------------- zipline/test/test_finance.py | 4 +- 6 files changed, 130 insertions(+), 77 deletions(-) create mode 100644 zipline/test/algorithms.py diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py index 83d1e763..5cc33f30 100644 --- a/zipline/finance/trading.py +++ b/zipline/finance/trading.py @@ -17,11 +17,11 @@ class TradeSimulationClient(qmsg.Component): self.received_count = 0 self.prev_dt = None self.event_queue = None - self.event_callbacks = [] self.txn_count = 0 self.trading_environment = trading_environment self.current_dt = trading_environment.period_start self.last_iteration_dur = datetime.timedelta(seconds=0) + self.algorithm = None assert self.trading_environment.frame_index != None self.event_frame = pandas.DataFrame( @@ -41,15 +41,15 @@ class TradeSimulationClient(qmsg.Component): @property def get_id(self): return str(zp.FINANCE_COMPONENT.TRADING_CLIENT) - - def add_event_callback(self, callback): + + def set_algorithm(self, algorithm): """ - :param callable callback: must be a function with the signature - f(event), where event is a namedict whose properties depend on the - upstream configuration of the zipline. It will include datasource and - transformations. + :param algorithm: must implement the algorithm protocol. See + algorithm_protocol.rst. """ - self.event_callbacks.append(callback) + self.algorithm = algorithm + #register the trading_client's order method with the algorithm + self.algorithm.set_order(self.order) def open(self): self.result_feed = self.connect_result() @@ -66,7 +66,7 @@ class TradeSimulationClient(qmsg.Component): if msg == str(zp.CONTROL_PROTOCOL.DONE): qutil.LOGGER.info("Client is DONE!") - self.run_callbacks() + self.run_algorithm() self.signal_order_done() self.signal_done() return @@ -83,7 +83,7 @@ class TradeSimulationClient(qmsg.Component): self.queue_event(event) if event.dt >= self.current_dt: - self.run_callbacks() + self.run_algorithm() #update time based on receipt of the order self.last_iteration_dur = datetime.datetime.utcnow() - event_start @@ -93,10 +93,9 @@ class TradeSimulationClient(qmsg.Component): #signal done to order source. self.order_socket.send(str(zp.ORDER_PROTOCOL.BREAK)) - def run_callbacks(self): + def run_algorithm(self): frame = self.get_frame() - for cb in self.event_callbacks: - cb(frame) + self.algorithm.handle_frame(frame) def connect_order(self): return self.connect_push_socket(self.addresses['order_address']) diff --git a/zipline/lines.py b/zipline/lines.py index f9773fd7..16502fe3 100644 --- a/zipline/lines.py +++ b/zipline/lines.py @@ -84,7 +84,7 @@ import zipline.protocol as zp import zipline.finance.performance as perf import zipline.messaging as zmsg -from zipline.test.client import TestAlgorithm +from zipline.test.algorithms import TestAlgorithm from zipline.sources import SpecificEquityTrades from zipline.finance.trading import TransactionSimulator, OrderDataSource, \ TradeSimulationClient @@ -114,11 +114,9 @@ class SimulatedTrading(object): def __init__(self, **config): """ :param config: a dict with the following required properties:: - - algorithm: a class that follows the algorithm protocol. Must - have a handle_frame method that accepts a pandas.Dataframe of the - current state of the simulation universe. Must have an order - property which can be set equal to the order method of - trading_client. (TODO: where should this protocol be documented?) + - algorithm: a class that follows the algorithm protocol. See + :py:meth:`zipline.finance.trading.TradingSimulationClient.add_algorithm` + for details. - trading_environment: an instance of :py:class:`zipline.trading.TradingEnvironment` - allocator: an instance of @@ -149,43 +147,30 @@ class SimulatedTrading(object): sockets[7], logging = qutil.LOGGER ) - + + self.started = False self.sim = config['simulator_class'](addresses) self.clients = {} self.trading_client = TradeSimulationClient(self.trading_environment) - self.clients[self.trading_client.get_id] = self.trading_client + self.add_client(self.trading_client) # setup all sources self.sources = {} self.order_source = OrderDataSource() - self.sources[self.order_source.get_id] = self.order_source + self.add_source(self.order_source) #setup transforms self.transaction_sim = TransactionSimulator() self.transforms = {} - self.transforms[self.transaction_sim.get_id] = self.transaction_sim + self.add_transform(self.transaction_sim) - #register all components - self.sim.register_components([ - self.trading_client, - self.order_source, - self.transaction_sim - ]) - self.sim.register_controller( self.con ) self.sim.on_done = self.shutdown() - self.started = False - ################################################################## - #TODO: the next two lines of code need refactoring from RealDiehl - ################################################################## - #wire up a callback inside the algorithm to receive frames from the - #trading client - self.trading_client.add_event_callback(self.algorithm.handle_frame) - #register the trading_client's order method with the algorithm - self.algorithm.set_order(self.trading_client.order) + + self.trading_client.set_algorithm(self.algorithm) @staticmethod def create_test_zipline(**config): @@ -202,7 +187,7 @@ class SimulatedTrading(object): subclass of ComponentHost to hold the whole zipline. Defaults to :py:class:`zipline.simulator.Simulator` - algorithm - optional parameter providing an algorithm. defaults - to :py:class:`zipline.test.client.TestAlgorithm` + to :py:class:`zipline.test.algorithms.TestAlgorithm` """ assert isinstance(config, dict) @@ -270,17 +255,27 @@ class SimulatedTrading(object): return zipline def add_source(self, source): + """ + Adds the source to the zipline, sets the sid filter of the + source to the algorithm's sid filter. + """ assert isinstance(source, zmsg.DataSource) self.check_started() + source.set_filter('SID', self.algorithm.get_sid_filter) self.sim.register_components([source]) self.sources[source.get_id] = source - def add_transform(self, transform): assert isinstance(transform, zmsg.BaseTransform) self.check_started() self.sim.register_components([transform]) - self.sources[transform.get_id] = transform + self.transforms[transform.get_id] = transform + + def add_client(self, client): + assert isinstance(client, TradeSimulationClient) + self.check_started() + self.sim.register_components([client]) + self.clients[client.get_id] = client def check_started(self): if self.started: diff --git a/zipline/messaging.py b/zipline/messaging.py index b87ca2ec..8f1a394a 100644 --- a/zipline/messaging.py +++ b/zipline/messaging.py @@ -545,16 +545,26 @@ 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). + + Every datasource has a dict property to hold filters:: + - key -- name of the filter, e.g. SID + - value -- a primitive representing the filter. e.g. a list of ints. + + Modify the datasource's filters via the set_filter(name, value) """ def __init__(self, source_id): Component.__init__(self) self.id = source_id self.init() + self.filter = {} def init(self): self.cur_event = None + def set_filter(self, name, value): + self.filter[name] = value + @property def get_id(self): return self.id diff --git a/zipline/test/algorithms.py b/zipline/test/algorithms.py new file mode 100644 index 00000000..8d8366bb --- /dev/null +++ b/zipline/test/algorithms.py @@ -0,0 +1,82 @@ +""" +Algorithm Protocol +=================== + +For a class to be passed as a trading algorithm to the +:py:class:`zipline.lines.SimulatedTrading` zipline +it must follow an implementation protocol. Examples of this algorithm protocol +are provided below. + +The algorithm must expose methods:: + - get_sid_filter: method that takes no args, and returns a list + of valid sids. List must have a length between 1 and 10. If None is returned + the filter will block all events. + + - handle_frame: method that accepts a :py:class:`pandas.Dataframe` of the + current state of the simulation universe. An example frame: ++-----------------+--------------+----------------+--------------------+ +| | SID(133) | SID(134) | SID(135) | ++=================+==============+=====================================+ +| price | $10.10 | $22.50 | $13.37 | ++-----------------+--------------+----------------+--------------------+ +| volume | 10,000 | 5,000 | 50,000 | ++-----------------+--------------+----------------+--------------------+ +| mvg_avg_30 | $9.97 | $22.61 | $13.37 | ++-----------------+--------------+----------------+--------------------+ +| dt | 6/30/2012 | 6/30/2012 | 6/29/2012 | ++-----------------+--------------+----------------+--------------------+ + +The algorithm must also expose settable properties: + - order: property which can be set equal to the order method of + trading_client. An algorithm can then place orders with a valid + SID and a number of shares:: + self.order(SID(133), share_count) + +""" + +import zipline.protocol as zp + +class TestAlgorithm(): + """ + This algorithm will send a specified number of orders, to allow unit tests + to verify the orders sent/received, transactions created, and positions + at the close of a simulation. + """ + + def __init__(self, sid, amount, order_count): + self.count = order_count + self.sid = sid + self.amount = amount + self.incr = 0 + self.done = False + self.order = None + + def set_order(self, order_callable): + self.order = order_callable + + def handle_frame(self, frame): + for dt, s in frame.iteritems(): + data = {} + data.update(s) + event = zp.namedict(data) + #place an order for 100 shares of sid:133 + if self.incr < self.count: + self.order(self.sid, self.amount) + self.incr += 1 + + def get_sid_filter(self): + return [self.sid] + +class NoopAlgorithm(object): + """ + Dolce fa niente. + """ + + def set_order(self, order_callable): + pass + + def handle_frame(self, frame): + pass + + def get_sid_filter(): + return None \ No newline at end of file diff --git a/zipline/test/client.py b/zipline/test/client.py index 74c556f7..324f52a7 100644 --- a/zipline/test/client.py +++ b/zipline/test/client.py @@ -83,36 +83,3 @@ class TestClient(qmsg.Component): def unframe(self, msg): return zp.MERGE_UNFRAME(msg) - - -class TestAlgorithm(): - - def __init__(self, sid, amount, order_count): - self.count = order_count - self.sid = sid - self.amount = amount - self.incr = 0 - self.done = False - self.order = None - - def set_order(self, order_callable): - self.order = order_callable - - def handle_frame(self, frame): - for dt, s in frame.iteritems(): - data = {} - data.update(s) - event = zp.namedict(data) - #place an order for 100 shares of sid:133 - if self.incr < self.count: - self.order(self.sid, self.amount) - self.incr += 1 - -class NoopAlgorithm(object): - - def set_order(self, order_callable): - pass - - def handle_frame(self, frame): - pass - \ No newline at end of file diff --git a/zipline/test/test_finance.py b/zipline/test/test_finance.py index 7371cc78..5edc0b0f 100644 --- a/zipline/test/test_finance.py +++ b/zipline/test/test_finance.py @@ -14,7 +14,7 @@ import zipline.finance.risk as risk import zipline.protocol as zp import zipline.finance.performance as perf -from zipline.test.client import TestAlgorithm +from zipline.test.algorithms import TestAlgorithm from zipline.sources import SpecificEquityTrades from zipline.finance.trading import TransactionSimulator, OrderDataSource, \ TradeSimulationClient, TradingEnvironment @@ -86,7 +86,7 @@ class FinanceTestCase(TestCase): zipline.algorithm.count, "The order source should have sent as many orders as the algo." ) - + transaction_sim = zipline.transforms[zp.TRANSFORM_TYPE.TRANSACTION] self.assertEqual( transaction_sim.txn_count,