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,