From 16b0d7150641df0f220408da84489a825328abd1 Mon Sep 17 00:00:00 2001 From: fawce Date: Fri, 5 Oct 2012 16:32:27 -0400 Subject: [PATCH] refactoring of algorithm to make it work for both batch style run method, and generator style consumption. removed the portfolio property from the data parameter. added set_slippage and set_commission methods to algorithm. removed timeout tracking. --- tests/test_algorithm.py | 1 - tests/test_exception_handling.py | 60 +--------- tests/test_finance.py | 2 +- tests/test_perf_tracking.py | 3 +- tests/test_transforms.py | 5 +- zipline/algorithm.py | 121 ++++++++++++------- zipline/finance/performance.py | 13 +- zipline/finance/slippage.py | 1 - zipline/finance/trading.py | 12 +- zipline/gens/tradesimulation.py | 66 ++--------- zipline/gens/transform.py | 37 ++++-- zipline/lines.py | 13 +- zipline/test_algorithms.py | 196 +++++-------------------------- zipline/utils/protocol_utils.py | 10 +- zipline/utils/simfactory.py | 99 ++++++++-------- 15 files changed, 228 insertions(+), 411 deletions(-) diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index f908ed7d..52666921 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -65,7 +65,6 @@ class TestTransformAlgorithm(TestCase): def test_transform_registered(self): algo = TestRegisterTransformAlgorithm(sids=[133]) algo.run(self.source) - assert algo.get_sid_filter() == algo.sids == [133] assert 'mavg' in algo.registered_transforms assert algo.registered_transforms['mavg']['args'] == (['price'],) assert algo.registered_transforms['mavg']['kwargs'] == \ diff --git a/tests/test_exception_handling.py b/tests/test_exception_handling.py index 269522b8..fefd2bfe 100644 --- a/tests/test_exception_handling.py +++ b/tests/test_exception_handling.py @@ -17,11 +17,12 @@ from unittest2 import TestCase from collections import defaultdict import zipline.utils.simfactory as simfactory -from zipline.test_algorithms import ExceptionAlgorithm, DivByZeroAlgorithm, \ - InitializeTimeoutAlgorithm, TooMuchProcessingAlgorithm +from zipline.test_algorithms import ( + ExceptionAlgorithm, + DivByZeroAlgorithm, +) from zipline.finance.slippage import FixedSlippage from zipline.gens.transform import StatefulTransform -from zipline.utils.timeout import TimeoutException from zipline.utils.test_utils import ( @@ -78,25 +79,6 @@ class ExceptionTestCase(TestCase): self.assertEqual(ctx.exception.message, 'An assertion message') - def test_exception_in_init(self): - # Simulation - # ---------- - self.zipline_test_config['algorithm'] = \ - ExceptionAlgorithm( - 'initialize', - self.zipline_test_config['sid'] - ) - - zipline = simfactory.create_test_zipline( - **self.zipline_test_config - ) - - with self.assertRaises(Exception) as ctx: - output, _ = drain_zipline(self, zipline) - - self.assertEqual(ctx.exception.message, - 'Algo exception in initialize') - def test_exception_in_handle_data(self): # Simulation # ---------- @@ -134,37 +116,3 @@ class ExceptionTestCase(TestCase): self.assertEqual(ctx.exception.message, 'integer division or modulo by zero') - - def test_initialize_timeout(self): - - self.zipline_test_config['algorithm'] = \ - InitializeTimeoutAlgorithm( - self.zipline_test_config['sid'] - ) - - zipline = simfactory.create_test_zipline( - **self.zipline_test_config - ) - - with self.assertRaises(TimeoutException) as ctx: - output, _ = drain_zipline(self, zipline) - - self.assertEqual(ctx.exception.message, 'Call to initialize timed out') - - def test_heartbeat(self): - - self.zipline_test_config['algorithm'] = \ - TooMuchProcessingAlgorithm( - self.zipline_test_config['sid'] - ) - zipline = simfactory.create_test_zipline( - **self.zipline_test_config - ) - - with self.assertRaises(TimeoutException) as ctx: - output, _ = drain_zipline(self, zipline) - - self.assertEqual( - ctx.exception.message, - 'Too much time spent in handle_data call' - ) diff --git a/tests/test_finance.py b/tests/test_finance.py index 224b5db2..647e076d 100644 --- a/tests/test_finance.py +++ b/tests/test_finance.py @@ -320,7 +320,7 @@ class FinanceTestCase(TestCase): self.assertEqual(order.sid, sid) self.assertEqual(order.amount, order_amount * alternator ** i) - tracker = PerformanceTracker(trading_environment, [sid]) + tracker = PerformanceTracker(trading_environment) # this approximates the loop inside TradingSimulationClient transactions = [] diff --git a/tests/test_perf_tracking.py b/tests/test_perf_tracking.py index 1fcf8baa..57f62d78 100644 --- a/tests/test_perf_tracking.py +++ b/tests/test_perf_tracking.py @@ -567,8 +567,7 @@ shares in position" 'price', 'changed'] perf_tracker = perf.PerformanceTracker( - self.trading_environment, - [sid, sid2] + self.trading_environment ) for event in trade_history: diff --git a/tests/test_transforms.py b/tests/test_transforms.py index e87551c8..85daa29c 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -32,6 +32,7 @@ from zipline.gens.stddev import MovingStandardDev from zipline.gens.returns import Returns import zipline.utils.factory as factory + from zipline.test_algorithms import BatchTransformAlgorithm @@ -315,7 +316,7 @@ class BatchTransformTestCase(TestCase): self.source, self.df = factory.create_test_df_source() def test_event_window(self): - algo = BatchTransformAlgorithm(sids=[0, 1]) + algo = BatchTransformAlgorithm() algo.run(self.source) self.assertEqual(algo.history_return_price_class[:2], @@ -344,7 +345,7 @@ class BatchTransformTestCase(TestCase): ) def test_passing_of_args(self): - algo = BatchTransformAlgorithm([0, 1], 1, kwarg='str') + algo = BatchTransformAlgorithm(1, kwarg='str') self.assertEqual(algo.args, (1,)) self.assertEqual(algo.kwargs, {'kwarg': 'str'}) diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 4379c5b6..f4f26c02 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -19,9 +19,20 @@ import numpy as np from zipline.gens.tradegens import DataFrameSource from zipline.utils.factory import create_trading_environment from zipline.gens.transform import StatefulTransform -from zipline.lines import SimulatedTrading -from zipline.finance.slippage import FixedSlippage, transact_partial -from zipline.finance.commission import PerShare +from zipline.finance.slippage import ( + VolumeShareSlippage, + FixedSlippage, + transact_partial +) +from zipline.finance.commission import PerShare, PerTrade + + +from zipline.gens.composites import ( + date_sorted_sources, + sequential_transforms +) +from zipline.gens.tradesimulation import TradeSimulationClient as tsc +from zipline import MESSAGES class TradingAlgorithm(object): @@ -44,54 +55,46 @@ class TradingAlgorithm(object): >>> stats = my_algo.run(data) """ - def __init__(self, sids, *args, **kwargs): + def __init__(self, *args, **kwargs): """ Initialize sids and other state variables. - - Calls user-defined initialize() forwarding *args and **kwargs. """ - self.sids = sids self.done = False self.order = None self.frame_count = 0 self.portfolio = None self.registered_transforms = {} + self.transforms = [] + self.sources = [] - # call to user-defined initialize method + # default components for transact + self.slippage = VolumeShareSlippage() + self.commission = PerShare() + + # an algorithm subclass needs to set initialized to True + # when it is fully initialized. + self.initialized = False + + # call to user-defined constructor method self.initialize(*args, **kwargs) - self.initialized = True - - def _create_simulator(self, start, end): + def _create_generator(self, environment): """ Create trading environment, transforms and SimulatedTrading object. Gets called by self.run(). """ - environment = create_trading_environment(start=start, end=end) - # Create transforms by wrapping them into StatefulTransforms - transforms = [] - for namestring, trans_descr in self.registered_transforms.iteritems(): - sf = StatefulTransform( - trans_descr['class'], - *trans_descr['args'], - **trans_descr['kwargs'] - ) - sf.namestring = namestring + self.date_sorted = date_sorted_sources(*self.sources) + self.with_tnfms = sequential_transforms(self.date_sorted, + *self.transforms) + self.trading_client = tsc(self, environment) - transforms.append(sf) + transact_method = transact_partial(self.slippage, self.commission) + self.set_transact(transact_method) - # SimulatedTrading is the main class handling data streaming, - # application of transforms and calling of the user algo. - return SimulatedTrading( - self.sources, - transforms, - self, - environment, - transact_partial(FixedSlippage(), PerShare(0.0)) - ) + return self.trading_client.simulate(self.with_tnfms) def run(self, source, start=None, end=None): """Run the algorithm. @@ -121,7 +124,7 @@ start and end date have to be specified.""" elif isinstance(source, pd.DataFrame): assert isinstance(source.index, pd.tseries.index.DatetimeIndex) # if DataFrame provided, wrap in DataFrameSource - source = DataFrameSource(source, sids=self.sids) + source = DataFrameSource(source) # If values not set, try to extract from source. if start is None: @@ -134,12 +137,25 @@ start and end date have to be specified.""" else: self.sources = source + # Create transforms by wrapping them into StatefulTransforms + for namestring, trans_descr in self.registered_transforms.iteritems(): + sf = StatefulTransform( + trans_descr['class'], + *trans_descr['args'], + **trans_descr['kwargs'] + ) + sf.namestring = namestring + + self.transforms.append(sf) + + environment = create_trading_environment(start=start, end=end) + # create transforms and zipline - self.simulated_trading = self._create_simulator(start=start, end=end) + self.gen = self._create_generator(environment) # loop through simulated_trading, each iteration returns a # perf ndict - perfs = list(self.simulated_trading) + perfs = list(self.gen) # convert perf ndict to pandas dataframe daily_stats = self._create_daily_stats(perfs) @@ -186,14 +202,39 @@ start and end date have to be specified.""" def set_order(self, order_callable): self.order = order_callable - def get_sid_filter(self): - return self.sids - def set_logger(self, logger): self.logger = logger - def initialize(self, *args, **kwargs): + def init(self, *args, **kwargs): + """Called from constructor.""" pass - def set_transact_setter(self, transact_setter): - pass + def set_transact(self, transact): + """ + Set the method that will be called to create a + transaction from open orders and trade events. + """ + self.trading_client.ordering_client.transact = transact + + def set_slippage(self, slippage): + assert isinstance(slippage, (VolumeShareSlippage, FixedSlippage)), \ + MESSAGES.ERRORS.UNSUPPORTED_SLIPPAGE_MODEL + if self.initialized: + raise Exception(MESSAGES.ERRORS.OVERRIDE_SLIPPAGE_POST_INIT) + self.slippage = slippage + + def set_commission(self, commission): + assert isinstance(commission, (PerShare, PerTrade)), \ + MESSAGES.ERRORS.UNSUPPORTED_COMMISSION_MODEL + + if self.initialized: + raise Exception(MESSAGES.ERRORS.OVERRIDE_COMMISSION_POST_INIT) + self.commission = commission + + def set_sources(self, sources): + assert isinstance(sources, list) + self.sources = sources + + def set_transforms(self, transforms): + assert isinstance(transforms, list) + self.transforms = transforms diff --git a/zipline/finance/performance.py b/zipline/finance/performance.py index 228a98b8..ff01a397 100644 --- a/zipline/finance/performance.py +++ b/zipline/finance/performance.py @@ -154,7 +154,7 @@ class PerformanceTracker(object): """ - def __init__(self, trading_environment, sid_list): + def __init__(self, trading_environment): self.trading_environment = trading_environment self.trading_day = datetime.timedelta(hours=6, minutes=30) @@ -203,9 +203,8 @@ class PerformanceTracker(object): keep_transactions=True ) - for sid in sid_list: - self.cumulative_performance.positions[sid] = Position(sid) - self.todays_performance.positions[sid] = Position(sid) + self.cumulative_performance.positions = positiondict() + self.todays_performance.positions = positiondict() def transform(self, stream_in): """ @@ -571,3 +570,9 @@ class PerformancePeriod(object): cur = pos.to_dict() positions.append(cur) return positions + + +class positiondict(dict): + + def __missing__(self, key): + return Position(key) diff --git a/zipline/finance/slippage.py b/zipline/finance/slippage.py index fc53fa1d..016911dd 100644 --- a/zipline/finance/slippage.py +++ b/zipline/finance/slippage.py @@ -26,7 +26,6 @@ def transact_stub(slippage, commission, open_orders, events): This is intended to be wrapped in a partial, so that the slippage and commission models can be enclosed. """ - transaction = slippage.simulate(open_orders, events) if transaction and transaction.amount != 0: direction = abs(transaction.amount) / transaction.amount diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py index 133b45ba..7a4ebd38 100644 --- a/zipline/finance/trading.py +++ b/zipline/finance/trading.py @@ -31,16 +31,8 @@ log = logbook.Logger('Transaction Simulator') class TransactionSimulator(object): - def __init__(self, transact=None): - - if transact is not None: - self.transact = transact - else: - self.transact = transact_partial( - VolumeShareSlippage(), - PerShare() - ) - + def __init__(self): + self.transact = transact_partial(VolumeShareSlippage(), PerShare()) self.open_orders = defaultdict(list) def place_order(self, order): diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index 819b21d6..b94f3c86 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -21,7 +21,6 @@ from itertools import groupby from operator import attrgetter from zipline import ndict -from zipline.utils.timeout import Heartbeat, Timeout from zipline.finance.trading import TransactionSimulator from zipline.finance.performance import PerformanceTracker @@ -29,11 +28,6 @@ from zipline.gens.utils import hash_args log = Logger('Trade Simulation') -# TODO: make these arguments rather than global constants -INIT_TIMEOUT = 5 -HEARTBEAT_INTERVAL = 1 # seconds -MAX_HEARTBEAT_INTERVALS = 15 # count - class TradeSimulationClient(object): """ @@ -69,15 +63,13 @@ class TradeSimulationClient(object): is sent to the algo. """ - def __init__(self, algo, environment, transact): + def __init__(self, algo, environment): self.algo = algo - self.sids = algo.get_sid_filter() self.environment = environment - self.transact = transact - self.ordering_client = TransactionSimulator(self.transact) - self.perf_tracker = PerformanceTracker(self.environment, self.sids) + self.ordering_client = TransactionSimulator() + self.perf_tracker = PerformanceTracker(self.environment) self.algo_start = self.environment.first_open self.algo_sim = AlgorithmSimulator( @@ -139,7 +131,6 @@ class AlgorithmSimulator(object): self.order_book = order_book self.algo = algo - self.sids = algo.get_sid_filter() self.algo_start = algo_start # Monkey patch the user algorithm to place orders in the @@ -148,35 +139,18 @@ class AlgorithmSimulator(object): self.algolog = Logger("AlgoLog") self.algo.set_logger(self.algolog) - # Provide user algorithm with a setter for the transact - # method (method that constructs transactions based on - # open orders and trade events). - self.algo.set_transact_setter(self.set_transact) - - # Handler for heartbeats during calls to handle_data. - def log_heartbeats(beat_count, stackframe): - t = beat_count * HEARTBEAT_INTERVAL - warning = "handle_data has been processing for %i seconds" % t - self.algolog.warn(warning) - - # Context manager that calls log_heartbeats every HEARTBEAT_INTERVAL - # seconds, raising an exception after MAX_HEARTBEATS - self.heartbeat_monitor = Heartbeat( - HEARTBEAT_INTERVAL, - MAX_HEARTBEAT_INTERVALS, - frame_handler=log_heartbeats, - timeout_message="Too much time spent in handle_data call" - ) - # ============== # Snapshot Setup # ============== # The algorithm's universe as of our most recent event. - self.universe = ndict() - for sid in self.sids: - self.universe[sid] = ndict() - self.universe.portfolio = None + # We want an ndict that will have empty ndicts as default + # values on missing keys. + self.universe = ndict(default=ndict) + # TODO: these keys are being inserted because universe + # has a default dictionary backing __internal. + # del self.universe['__members__'] + # del self.universe['__methods__'] # We don't have a datetime for the current snapshot until we # receive a message. @@ -193,19 +167,11 @@ class AlgorithmSimulator(object): record.extra['algo_dt'] = self.snapshot_dt self.processor = Processor(inject_algo_dt) - def set_transact(self, transact): - """ - Set the method that will be called to create a - transaction from open orders and trade events. - """ - self.order_book.transact = transact - def order(self, sid, amount): """ Closure to pass into the user's algo to allow placing orders into the transaction simulator's dict of open orders. """ - assert sid in self.sids, "Order on invalid sid: %i" % sid order = ndict({ 'dt': self.simulation_dt, 'sid': sid, @@ -233,12 +199,6 @@ class AlgorithmSimulator(object): """ Main generator work loop. """ - # Call user's initialize method with a timeout (only if - # initialize wasn't called already). - if not getattr(self.algo, 'initialized', False): - with Timeout(INIT_TIMEOUT, message="Call to initialize timed out"): - self.algo.initialize() - # inject the current algo # snapshot time to any log record generated. with self.processor.threadbound(): @@ -298,7 +258,7 @@ class AlgorithmSimulator(object): Update the universe with new event information. """ # Update our portfolio. - self.universe.portfolio = event.portfolio + self.algo.set_portfolio(event.portfolio) # Update our knowledge of this event's sid for field in event.keys(): @@ -312,10 +272,8 @@ class AlgorithmSimulator(object): # Needs to be set so that we inject the proper date into algo # log/print lines. self.snapshot_dt = date - start_tic = datetime.now() - with self.heartbeat_monitor: - self.algo.handle_data(self.universe) + self.algo.handle_data(self.universe) stop_tic = datetime.now() # How long did you take? diff --git a/zipline/gens/transform.py b/zipline/gens/transform.py index c9e455a5..f1c99293 100644 --- a/zipline/gens/transform.py +++ b/zipline/gens/transform.py @@ -348,16 +348,16 @@ class BatchTransform(EventWindow): refresh_period=None, market_aware=True, delta=None, - days=None, - sids=None): - super(BatchTransform, self).__init__( - market_aware, days=days, delta=delta) + days=None): + + super(BatchTransform, self).__init__(market_aware, + days=days, delta=delta) + if func is not None: self.compute_transform_value = func else: self.compute_transform_value = self.get_value - self.sids = sids self.refresh_period = refresh_period self.days = days @@ -373,15 +373,20 @@ class BatchTransform(EventWindow): handle_data method. """ # extract dates - dts = [data[sid].datetime for sid in self.sids] + #dts = [data[sid].datetime for sid in self.sids] + dts = [event.datetime for event in data.itervalues()] # we have to provide the event with a dt. This is only for # checking if the event is outside the window or not so a - # couple of seconds shouldn't matter - data.dt = max(dts) + # couple of seconds shouldn't matter. We don't add it to + # the data parameter, because it would mix dt with the + # sid keys. + event = ndict() + event.dt = max(dts) + event.data = data # append data frame to window. update() will call handle_add() and # handle_remove() appropriately - self.update(data) + self.update(event) # return newly computed or cached value return self.get_transform_value(*args, **kwargs) @@ -403,17 +408,23 @@ class BatchTransform(EventWindow): # # This Panel data structure ultimately gets passed to the # user-overloaded get_value() method. + # + # self.ticks contains ndicts with data, dt keys. + # event parameter is an ndict with data, dt keys. fields = {} for field_name in ['price', 'volume']: + sids = self.ticks[0].data.keys() # Skip non-existant fields - if field_name not in self.ticks[0][self.sids[0]]: + if field_name not in self.ticks[0].data[sids[0]]: continue values_per_sid = {} - for sid in self.sids: + + for sid in sids: values_per_sid[sid] = pd.Series( - {tick[sid].dt: tick[sid][field_name] - for tick in self.ticks}) + {tick.data[sid].dt: tick.data[sid][field_name] + for tick in self.ticks} + ) # concatenate different sids into one df fields[field_name] = pd.DataFrame.from_dict(values_per_sid) diff --git a/zipline/lines.py b/zipline/lines.py index f7b2baf4..3ad0ae83 100644 --- a/zipline/lines.py +++ b/zipline/lines.py @@ -59,6 +59,7 @@ before invoking simulate. | __init__. | +---------------------------------+ """ + from zipline.gens.composites import ( date_sorted_sources, sequential_transforms @@ -76,8 +77,8 @@ class SimulatedTrading(object): sources, transforms, algorithm, - environment, - sim_method=None): + environment + ): """ @sources - an iterable of iterables These iterables must yield ndicts that contain: @@ -104,7 +105,13 @@ class SimulatedTrading(object): # Formerly merged_transforms. self.with_tnfms = sequential_transforms(self.date_sorted, *self.transforms) - self.trading_client = tsc(algorithm, environment, sim_method) + self.trading_client = tsc(algorithm, environment) + + # give the algorithm access to the simulator to control + # state such as universe, commissions, and slippage. With + # great power comes great responsibility. + algorithm.simulator = self + self.gen = self.trading_client.simulate(self.with_tnfms) def __iter__(self): diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index 5d79499b..7255df9c 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -67,42 +67,28 @@ The algorithm must expose methods: and trade events. """ +from zipline.algorithm import TradingAlgorithm +from zipline.finance.slippage import FixedSlippage -class TestAlgorithm(): +class TestAlgorithm(TradingAlgorithm): """ 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, sid_filter=None): + def initialize(self, sid, amount, order_count, sid_filter=None): self.count = order_count self.sid = sid self.amount = amount self.incr = 0 - self.done = False - self.order = None - self.frame_count = 0 - self.portfolio = None if sid_filter: self.sid_filter = sid_filter else: self.sid_filter = [self.sid] - def initialize(self): - pass - - def set_order(self, order_callable): - self.order = order_callable - - def set_logger(self, logger): - pass - - def set_portfolio(self, portfolio): - self.portfolio = portfolio - def handle_data(self, data): self.frame_count += 1 #place an order for 100 shares of sid @@ -110,40 +96,18 @@ class TestAlgorithm(): self.order(self.sid, self.amount) self.incr += 1 - def get_sid_filter(self): - return self.sid_filter - def set_transact_setter(self, txn_sim_callable): - pass - - -class HeavyBuyAlgorithm(): +class HeavyBuyAlgorithm(TradingAlgorithm): """ 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): + def initialize(self, sid, amount): self.sid = sid self.amount = amount self.incr = 0 - self.done = False - self.order = None - self.frame_count = 0 - self.portfolio = None - - def initialize(self): - pass - - def set_order(self, order_callable): - self.order = order_callable - - def set_logger(self, logger): - pass - - def set_portfolio(self, portfolio): - self.portfolio = portfolio def handle_data(self, data): self.frame_count += 1 @@ -158,26 +122,10 @@ class HeavyBuyAlgorithm(): pass -class NoopAlgorithm(object): +class NoopAlgorithm(TradingAlgorithm): """ Dolce fa niente. """ - - def initialize(self): - pass - - def set_order(self, order_callable): - pass - - def set_logger(self, logger): - pass - - def set_portfolio(self, portfolio): - pass - - def handle_data(self, data): - pass - def get_sid_filter(self): return [] @@ -185,17 +133,17 @@ class NoopAlgorithm(object): pass -class ExceptionAlgorithm(object): +class ExceptionAlgorithm(TradingAlgorithm): """ Throw an exception from the method name specified in the constructor. """ - def __init__(self, throw_from, sid): + def initialize(self, throw_from, sid): + self.throw_from = throw_from self.sid = sid - def initialize(self): if self.throw_from == "initialize": raise Exception("Algo exception in initialize") else: @@ -207,9 +155,6 @@ class ExceptionAlgorithm(object): else: pass - def set_logger(self, logger): - pass - def set_portfolio(self, portfolio): if self.throw_from == "set_portfolio": raise Exception("Algo exception in set_portfolio") @@ -232,43 +177,24 @@ class ExceptionAlgorithm(object): pass -class DivByZeroAlgorithm(): +class DivByZeroAlgorithm(TradingAlgorithm): - def __init__(self, sid): + def initialize(self, sid): self.sid = sid self.incr = 0 - def initialize(self): - pass - - def set_order(self, order_callable): - pass - - def set_logger(self, logger): - pass - - def set_portfolio(self, portfolio): - pass - def handle_data(self, data): self.incr += 1 if self.incr > 4: 5 / 0 pass - def get_sid_filter(self): - return [self.sid] - def set_transact_setter(self, txn_sim_callable): - pass +class InitializeTimeoutAlgorithm(TradingAlgorithm): - -class InitializeTimeoutAlgorithm(): - def __init__(self, sid): + def initialize(self, sid): self.sid = sid self.incr = 0 - - def initialize(self): import time from zipline.gens.tradesimulation import INIT_TIMEOUT time.sleep(INIT_TIMEOUT + 1000) @@ -292,121 +218,54 @@ class InitializeTimeoutAlgorithm(): pass -class TooMuchProcessingAlgorithm(): - def __init__(self, sid): +class TooMuchProcessingAlgorithm(TradingAlgorithm): + + def initialize(self, sid): self.sid = sid - def initialize(self): - pass - - def set_order(self, order_callable): - pass - - def set_logger(self, logger): - pass - - def set_portfolio(self, portfolio): - pass - def handle_data(self, data): # Unless we're running on some sort of # supercomputer this will hit timeout. for i in xrange(1000000000): self.foo = i - def get_sid_filter(self): - return [self.sid] - def set_transact_setter(self, txn_sim_callable): - pass +class TimeoutAlgorithm(TradingAlgorithm): - -class TimeoutAlgorithm(): - - def __init__(self, sid): + def initialize(self, sid): self.sid = sid self.incr = 0 - def initialize(self): - pass - - def set_order(self, order_callable): - pass - - def set_logger(self, logger): - pass - - def set_portfolio(self, portfolio): - pass - def handle_data(self, data): if self.incr > 4: import time time.sleep(100) pass - def get_sid_filter(self): - return [self.sid] - def set_transact_setter(self, txn_sim_callable): - pass +class TestPrintAlgorithm(TradingAlgorithm): - -class TestPrintAlgorithm(): - - def __init__(self, sid): + def initialize(self, sid): self.sid = sid - - def initialize(self): print "Initializing..." - def set_order(self, order_callable): - pass - - def set_logger(self, logger): - pass - - def set_portfolio(self, portfolio): - pass - def handle_data(self, data): print "Handling Data..." pass - def get_sid_filter(self): - return [self.sid] - def set_transact_setter(self, txn_sim_callable): - pass +class TestLoggingAlgorithm(TradingAlgorithm): - -class TestLoggingAlgorithm(): - - def __init__(self, sid): + def initialize(self, sid): self.log = None self.sid = sid - def initialize(self): - self.log.info("Initializing...") - - def set_order(self, order_callable): - pass - - def set_logger(self, logger): - self.log = logger - def set_portfolio(self, portfolio): pass def handle_data(self, data): self.log.info("Handling Data...") - def get_sid_filter(self): - return [self.sid] - - def set_transact_setter(self, txn_sim_callable): - pass - from datetime import timedelta from zipline.algorithm import TradingAlgorithm @@ -415,11 +274,13 @@ from zipline.gens.mavg import MovingAverage class TestRegisterTransformAlgorithm(TradingAlgorithm): - def initialize(self): + def initialize(self, *args, **kwargs): self.add_transform(MovingAverage, 'mavg', ['price'], market_aware=True, days=2) + self.set_slippage(FixedSlippage()) + def handle_data(self, data): pass @@ -454,26 +315,25 @@ class BatchTransformAlgorithm(TradingAlgorithm): self.kwargs = kwargs self.return_price_class = ReturnPriceBatchTransform( - sids=self.sids, market_aware=False, refresh_period=2, delta=timedelta(days=self.days) ) self.return_price_decorator = return_price_batch_decorator( - sids=self.sids, market_aware=False, refresh_period=2, delta=timedelta(days=self.days) ) self.return_args_batch = return_args_batch_decorator( - sids=self.sids, market_aware=False, refresh_period=2, delta=timedelta(days=self.days) ) + self.set_slippage(FixedSlippage()) + def handle_data(self, data): self.history_return_price_class.append( self.return_price_class.handle_data(data)) diff --git a/zipline/utils/protocol_utils.py b/zipline/utils/protocol_utils.py index f1c77102..a37895fe 100644 --- a/zipline/utils/protocol_utils.py +++ b/zipline/utils/protocol_utils.py @@ -1,7 +1,7 @@ import copy import pandas from ctypes import Structure, c_ubyte -from collections import MutableMapping +from collections import MutableMapping, defaultdict def Enum(*options): @@ -45,9 +45,11 @@ class ndict(MutableMapping): cls = None __slots__ = ['cls', '__internal'] - def __init__(self, dct=None): - self.__internal = dict() - + def __init__(self, dct=None, default=None): + if default is not None: + self.__internal = dict() + else: + self.__internal = defaultdict(default) if not ndict.cls: ndict.cls = frozenset(dir(self)) diff --git a/zipline/utils/simfactory.py b/zipline/utils/simfactory.py index 7cf18130..b9dd78c4 100644 --- a/zipline/utils/simfactory.py +++ b/zipline/utils/simfactory.py @@ -1,36 +1,33 @@ import zipline.utils.factory as factory from zipline.test_algorithms import TestAlgorithm -from zipline.lines import SimulatedTrading -from zipline.finance.slippage import FixedSlippage, transact_partial -from zipline.finance.commission import PerShare def create_test_zipline(**config): """ - :param config: A configuration object that is a dict with: + :param config: A configuration object that is a dict with: - - environment - a \ - :py:class:`zipline.finance.trading.TradingEnvironment` - - sid - an integer, which will be used as the security ID. - - order_count - the number of orders the test algo will place, - defaults to 100 - - order_amount - the number of shares per order, defaults to 100 - - trade_count - the number of trades to simulate, defaults to 101 - to ensure all orders are processed. - - 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:`zipline.sources.SpecificEquityTrades` - is the source, with daily frequency in trades. - - slippage: optional parameter that configures the - :py:class:`zipline.gens.tradingsimulation.TransactionSimulator`. Expects - an object with a simulate mehod, such as - :py:class:`zipline.gens.tradingsimulation.FixedSlippage`. - :py:mod:`zipline.finance.trading` - - transforms: optional parameter that provides a list - of StatefulTransform objects. - """ + - environment - a \ + :py:class:`zipline.finance.trading.TradingEnvironment` + - sid - an integer, which will be used as the security ID. + - order_count - the number of orders the test algo will place, + defaults to 100 + - order_amount - the number of shares per order, defaults to 100 + - trade_count - the number of trades to simulate, defaults to 101 + to ensure all orders are processed. + - 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:`zipline.sources.SpecificEquityTrades` + is the source, with daily frequency in trades. + - slippage: optional parameter that configures the + :py:class:`zipline.gens.tradingsimulation.TransactionSimulator`. + Expects an object with a simulate mehod, such as + :py:class:`zipline.gens.tradingsimulation.FixedSlippage`. + :py:mod:`zipline.finance.trading` + - transforms: optional parameter that provides a list + of StatefulTransform objects. + """ assert isinstance(config, dict) sid_list = config.get('sid_list') if not sid_list: @@ -64,12 +61,20 @@ def create_test_zipline(**config): # trade than order trade_count = 101 - slippage = config.get('slippage', FixedSlippage()) - commission = PerShare() - transact_method = transact_partial(slippage, commission) + #------------------- + # Create the Algo + #------------------- + if 'algorithm' in config: + test_algo = config['algorithm'] + else: + test_algo = TestAlgorithm( + sid, + order_amount, + order_count + ) #------------------- - # Trade Source + # Trade Source #------------------- if 'trade_source' in config: trade_source = config['trade_source'] @@ -84,31 +89,21 @@ def create_test_zipline(**config): #------------------- # Transforms #------------------- - transforms = config.get('transforms', []) + test_algo.set_sources([trade_source]) + + transforms = config.get('transforms', None) + if transforms is not None: + test_algo.set_transforms(transforms) #------------------- - # Create the Algo - #------------------- - if 'algorithm' in config: - test_algo = config['algorithm'] - else: - test_algo = TestAlgorithm( - sid, - order_amount, - order_count - ) + # Slippage + # ------------------ + slippage = config.get('slippage', None) + if slippage is not None: + test_algo.set_slippage(slippage) - #------------------- - # Simulation - #------------------- - - sim = SimulatedTrading( - [trade_source], - transforms, - test_algo, - trading_environment, - transact_method - ) - #------------------- + # ------------------ + # generator/simulator + sim = test_algo._create_generator(trading_environment) return sim