diff --git a/tests/test_sources.py b/tests/test_sources.py index 0711f72f..548f8754 100644 --- a/tests/test_sources.py +++ b/tests/test_sources.py @@ -1,13 +1,22 @@ +from unittest2 import TestCase + import zipline.utils.factory as factory +from zipline.gens.tradegens import DataFrameSource -def test_dataframe_source(): - source, df = factory.create_test_df_source() +class TestDataFrameSource(TestCase): + def test_streaming_of_df(self): + source, df = factory.create_test_df_source() - for expected_dt, expected_price in df.iterrows(): - sid0 = source.next() - sid1 = source.next() + for expected_dt, expected_price in df.iterrows(): + sid0 = source.next() + sid1 = source.next() - assert expected_dt == sid0.dt == sid1.dt - assert expected_price[0] == sid0.price - assert expected_price[1] == sid1.price + assert expected_dt == sid0.dt == sid1.dt + assert expected_price[0] == sid0.price + assert expected_price[1] == sid1.price + def test_sid_filtering(self): + _, df = factory.create_test_df_source() + source = DataFrameSource(df, sids=[0]) + assert 1 not in [event.sid for event in source], \ + "DataFrameSource should only stream selected sid 0, not sid 1." \ No newline at end of file diff --git a/tests/test_transforms.py b/tests/test_transforms.py index d33951b2..ea1024f6 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -10,13 +10,14 @@ from zipline.utils.test_utils import setup_logger from zipline.utils.date_utils import utcnow from zipline.gens.tradegens import SpecificEquityTrades -from zipline.gens.transform import StatefulTransform, EventWindow, BatchTransform, batch_transform +from zipline.gens.transform import StatefulTransform, EventWindow from zipline.gens.vwap import VWAP from zipline.gens.mavg import MovingAverage from zipline.gens.stddev import MovingStandardDev from zipline.gens.returns import Returns import zipline.utils.factory as factory -from zipline import TradingAlgorithm + +from zipline.test_algorithms import BatchTransformAlgorithm def to_dt(msg): return ndict({'dt': msg}) @@ -288,36 +289,7 @@ class FinanceTransformsTestCase(TestCase): ############################################################ # Test BatchTransform -class NoopBatchTransform(BatchTransform): - def get_value(self, data): - return data.price - -@batch_transform -def noop_batch_decorator(data): - return data.price - -class BatchTransformAlgorithm(TradingAlgorithm): - def initialize(self, *args, **kwargs): - self.history_class = [] - self.history_decorator = [] - self.days = 3 - self.noop_class = NoopBatchTransform(sids=[0, 1], - market_aware=False, - refresh_period=2, - delta=timedelta(days=self.days)) - - self.noop_decorator = noop_batch_decorator(sids=[0, 1], - market_aware=False, - refresh_period=2, - delta=timedelta(days=self.days)) - - def handle_data(self, data): - window_class = self.noop_class.handle_data(data) - window_decorator = self.noop_decorator.handle_data(data) - self.history_class.append(window_class) - self.history_decorator.append(window_decorator) - -class BatchTransformTestCase(): +class BatchTransformTestCase(TestCase): def setUp(self): setup_logger(self) self.source, self.df = factory.create_test_df_source() @@ -329,21 +301,8 @@ class BatchTransformTestCase(): assert algo.history_class[:2] == algo.history_decorator[:2] == [None, None], "First two iterations should return None" # test overloaded class - # every 2nd event should be identical because of refresh_period=2 - # not sure why actual length gets up to 4, bug in EventWindow? - assert np.all(algo.history_class[2][0].values == [2, 4, 6]) - assert np.all(algo.history_class[2][1].values == [3, 5, 7]) - assert np.all(algo.history_class[3][0].values == [2, 4, 6]) - assert np.all(algo.history_class[3][1].values == [3, 5, 7]) - assert np.all(algo.history_class[4][0].values == [4, 6, 8, 10]) - assert np.all(algo.history_class[4][1].values == [5, 7, 9, 11]) - - # test decorator - assert np.all(algo.history_decorator[2][0].values == [2, 4, 6]) - assert np.all(algo.history_decorator[2][1].values == [3, 5, 7]) - assert np.all(algo.history_decorator[3][0].values == [2, 4, 6]) - assert np.all(algo.history_decorator[3][1].values == [3, 5, 7]) - assert np.all(algo.history_decorator[4][0].values == [4, 6, 8, 10]) - assert np.all(algo.history_decorator[4][1].values == [5, 7, 9, 11]) - + for test_history in [algo.history_class, algo.history_decorator]: + self.assertTrue(np.all(test_history[2].values.flatten() == range(4, 10))) + self.assertTrue(np.all(test_history[3].values.flatten() == range(4, 10))) + self.assertTrue(np.all(test_history[4].values.flatten() == range(6, 14))) diff --git a/zipline/__init__.py b/zipline/__init__.py index 3916f8cf..4151cd02 100644 --- a/zipline/__init__.py +++ b/zipline/__init__.py @@ -6,9 +6,7 @@ Zipline # it is a place to expose the public interfaces. from utils.protocol_utils import ndict -from algorithm import TradingAlgorithm __all__ = [ - ndict, - TradingAlgorithm + ndict ] diff --git a/zipline/algorithm.py b/zipline/algorithm.py index b102ea7e..08d388da 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -9,8 +9,8 @@ from zipline.finance.slippage import FixedSlippage class TradingAlgorithm(object): - """ - Base class for trading algorithms. Inherit and overload handle_data(data). + """Base class for trading algorithms. Inherit and overload + initialize() and handle_data(data). A new algorithm could look like this: ``` @@ -22,7 +22,7 @@ class TradingAlgorithm(object): sid = self.sids[0] self.order(sid, amount) ``` - To then run this algorithm: + To then to run this algorithm: >>> my_algo = MyAlgo(100, sids=[0]) >>> stats = my_algo.run(data) @@ -45,13 +45,13 @@ class TradingAlgorithm(object): # call to user-defined initialize method self.initialize(*args, **kwargs) - def _create_simulator(self, source): + def _create_simulator(self, start, end): """ Create trading environment, transforms and SimulatedTrading object. Gets called by self.run(). """ - environment = create_trading_environment(start=source.data.index[0], end=source.data.index[-1]) + environment = create_trading_environment(start=start, end=end) # Create transforms by wrapping them into StatefulTransforms transforms = [] @@ -68,39 +68,59 @@ class TradingAlgorithm(object): # SimulatedTrading is the main class handling data streaming, # application of transforms and calling of the user algo. return SimulatedTrading( - [source], + self.sources, transforms, self, environment, FixedSlippage() ) - def run(self, source): - """ - Run the algorithm. + def run(self, source, start=None, end=None): + """Run the algorithm. :Arguments: - data : zipline source or pandas.DataFrame - pandas.DataFrame must have the following structure: - * column names must consist of ints representing the different sids - * index must be TimeStamps - * array contents should be price + source : can be either: + - pandas.DataFrame + - zipline source + - list of zipline sources + + If pandas.DataFrame is provided, it must have the + following structure: + * column names must consist of ints representing the + different sids + * index must be DatetimeIndex + * array contents should be price info. :Returns: daily_stats : pandas.DataFrame Daily performance metrics such as returns, alpha etc. """ - if isinstance(source, pd.DataFrame): + if isinstance(source, (list, tuple)): + assert start is not None and end is not None, \ + "When providing a list of sources, 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) + # If values not set, try to extract from source. + if start is None: + start = source.start + if end is None: + end = source.end + + if not isinstance(source, (list, tuple)): + self.sources = [source] + else: + self.sources = source + # create transforms and zipline - simulated_trading = self._create_simulator(source) + self.simulated_trading = self._create_simulator(start=start, end=end) # loop through simulated_trading, each iteration returns a # perf ndict - perfs = list(simulated_trading) + perfs = list(self.simulated_trading) # convert perf ndict to pandas dataframe daily_stats = self._create_daily_stats(perfs) @@ -123,6 +143,7 @@ class TradingAlgorithm(object): return daily_stats + def add_transform(self, transform_class, tag, *args, **kwargs): """Add a single-sid, sequential transform to the model. diff --git a/zipline/gens/tradegens.py b/zipline/gens/tradegens.py index c1b00c65..f088ae8b 100644 --- a/zipline/gens/tradegens.py +++ b/zipline/gens/tradegens.py @@ -9,6 +9,7 @@ from itertools import chain, cycle, ifilter, izip, repeat from datetime import datetime, timedelta import pandas as pd from copy import copy +import numpy as np from zipline.protocol import DATASOURCE_TYPE from zipline.utils import ndict @@ -77,17 +78,31 @@ class SpecificEquityTrades(object): # We shouldn't get any positional arguments. assert len(args) == 0 - # Unpack config dictionary with default values. - self.count = kwargs.get('count', 500) - self.sids = kwargs.get('sids', [1, 2]) - self.start = kwargs.get('start', datetime(2008, 6, 6, 15, tzinfo = pytz.utc)) - self.delta = kwargs.get('delta', timedelta(minutes = 1)) - self.concurrent = kwargs.get('concurrent', False) - # Default to None for event_list and filter. self.event_list = kwargs.get('event_list') self.filter = kwargs.get('filter') + if self.event_list is not None: + # If event_list is provided, extract parameters from there + # This isn't really clean and ultimately I think this + # class should serve a single purpose (either take an + # event_list or autocreate events). + self.count = kwargs.get('count', len(self.event_list)) + self.sids = kwargs.get('sids', np.unique([event.sid for event in self.event_list]).tolist()) + self.start = kwargs.get('start', self.event_list[0].dt) + self.end = kwargs.get('start', self.event_list[-1].dt) + self.delta = kwargs.get('delta', self.event_list[1].dt - self.event_list[0].dt) + self.concurrent = kwargs.get('concurrent', False) + + else: + # Unpack config dictionary with default values. + self.count = kwargs.get('count', 500) + self.sids = kwargs.get('sids', [1, 2]) + self.start = kwargs.get('start', datetime(2008, 6, 6, 15, tzinfo = pytz.utc)) + self.delta = kwargs.get('delta', timedelta(minutes = 1)) + self.concurrent = kwargs.get('concurrent', False) + + # Hash_value for downstream sorting. self.arg_string = hash_args(*args, **kwargs) @@ -188,9 +203,6 @@ class DataFrameSource(SpecificEquityTrades): self.end = kwargs.get('end', data.index[-1]) self.delta = kwargs.get('delta', data.index[1]-data.index[0]) - # Default to None for event_list and filter. - self.filter = kwargs.get('filter') - # Hash_value for downstream sorting. self.arg_string = hash_args(data, **kwargs) @@ -214,4 +226,5 @@ class DataFrameSource(SpecificEquityTrades): yield ndict(event) # Return the filtered event stream. - return _generator() \ No newline at end of file + drop_sids = lambda x: x.sid in self.sids + return ifilter(drop_sids, _generator()) diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index 1c7f0efe..5efb1505 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -194,8 +194,6 @@ class AlgorithmSimulator(object): 'filled' : 0 }) - log.debug(order) - # Tell the user if they try to buy 0 shares of something. if order.amount == 0: zero_message = "Requested to trade zero shares of {sid}".format( @@ -296,8 +294,8 @@ class AlgorithmSimulator(object): self.snapshot_dt = date start_tic = datetime.now() - #with self.heartbeat_monitor: - self.algo.handle_data(self.universe) + with self.heartbeat_monitor: + self.algo.handle_data(self.universe) stop_tic = datetime.now() # How long did you take? diff --git a/zipline/optimize/algorithms.py b/zipline/optimize/algorithms.py index 6625012a..976979d2 100644 --- a/zipline/optimize/algorithms.py +++ b/zipline/optimize/algorithms.py @@ -1,9 +1,9 @@ from logbook import Logger -from zipline import TradingAlgorithm +from zipline.algorithm import TradingAlgorithm logger = Logger('Algo') -class BuySellAlgorithm(object): +class BuySellAlgorithm(TradingAlgorithm): """Algorithm that buys and sells alternatingly. The amount for each order can be specified. In addition, an offset that will quadratically reduce the amount that will be bought can be @@ -15,69 +15,11 @@ class BuySellAlgorithm(object): """ - def __init__(self, sid, amount, offset): - self.sid = sid + def initialize(self, amount=100, offset=0): self.amount = amount - self.incr = 0 - self.done = False - self.order = None - self.frame_count = 0 - self.portfolio = None self.buy_or_sell = -1 self.offset = offset self.orders = [] - self.prices = [] - - def initialize(self): - pass - - def set_order(self, order_callable): - self.order = order_callable - - def set_portfolio(self, portfolio): - self.portfolio = portfolio - - def handle_data(self, frame): - print frame.sid - order_size = self.buy_or_sell * (self.amount - (self.offset**2)) - self.order(self.sid, order_size) - - #sell next time around. - self.buy_or_sell *= -1 - - self.orders.append(order_size) - - self.frame_count += 1 - self.incr += 1 - - def get_sid_filter(self): - return [self.sid] - - -class BuySellAlgorithmNew(TradingAlgorithm): - """Algorithm that buys and sells alternatingly. The amount for - each order can be specified. In addition, an offset that will - quadratically reduce the amount that will be bought can be - specified. - - This algorithm is used to test the parameter optimization - framework. If combined with the UpDown trade source, an offset of - 0 will produce maximum returns. - - """ - - def __init__(self, sids, amount, offset): - self.sids = sids - self.amount = amount - self.incr = 0 - self.done = False - self.order = None - self.frame_count = 0 - self.portfolio = None - self.buy_or_sell = -1 - self.offset = offset - self.orders = [] - self.prices = [] def handle_data(self, data): order_size = self.buy_or_sell * (self.amount - (self.offset**2)) @@ -89,6 +31,3 @@ class BuySellAlgorithmNew(TradingAlgorithm): self.orders.append(order_size) - self.frame_count += 1 - self.incr += 1 - diff --git a/zipline/optimize/factory.py b/zipline/optimize/factory.py index 901798b3..14fea2ad 100644 --- a/zipline/optimize/factory.py +++ b/zipline/optimize/factory.py @@ -9,7 +9,7 @@ import zipline.protocol as zp from zipline.utils.factory import get_next_trading_dt, create_trading_environment from zipline.gens.tradegens import SpecificEquityTrades -from zipline.optimize.algorithms import BuySellAlgorithmNew +from zipline.optimize.algorithms import BuySellAlgorithm from zipline.finance.slippage import FixedSlippage from copy import copy @@ -120,7 +120,7 @@ def create_predictable_zipline(config, offset=0, simulate=True): amplitude) if 'algorithm' not in config: - algorithm = BuySellAlgorithmNew(sid, 100, offset) + algorithm = BuySellAlgorithm(sids=[sid], amount=100, offset=offset) config['order_count'] = trade_count - 1 config['trade_count'] = trade_count diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index fbb9d5a9..10795e33 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -52,6 +52,7 @@ The algorithm must expose methods: """ + class TestAlgorithm(): """ This algorithm will send a specified number of orders, to allow unit tests @@ -382,3 +383,49 @@ class TestLoggingAlgorithm(): def set_slippage_override(self, slippage_callable): pass + + +from datetime import timedelta +from zipline.algorithm import TradingAlgorithm +from zipline.gens.transform import BatchTransform, batch_transform +from zipline.gens.mavg import MovingAverage + +class TestRegisterTransformAlgorithm(TradingAlgorithm): + def initialize(self): + self.add_transform(MovingAverage, 'mavg', ['price'], + market_aware=True, + days=2) + + def handle_data(self, data): + pass + +class NoopBatchTransform(BatchTransform): + def get_value(self, data): + return data.price + +@batch_transform +def noop_batch_decorator(data): + return data.price + +class BatchTransformAlgorithm(TradingAlgorithm): + def initialize(self, *args, **kwargs): + self.history_class = [] + self.history_decorator = [] + self.days = 3 + self.noop_class = NoopBatchTransform(sids=[0, 1], + market_aware=False, + refresh_period=2, + delta=timedelta(days=self.days)) + + self.noop_decorator = noop_batch_decorator(sids=[0, 1], + market_aware=False, + refresh_period=2, + delta=timedelta(days=self.days)) + + def handle_data(self, data): + window_class = self.noop_class.handle_data(data) + window_decorator = self.noop_decorator.handle_data(data) + self.history_class.append(window_class) + self.history_decorator.append(window_decorator) + + diff --git a/zipline/utils/factory.py b/zipline/utils/factory.py index 1f0026f5..7517fff5 100644 --- a/zipline/utils/factory.py +++ b/zipline/utils/factory.py @@ -240,7 +240,7 @@ def create_test_df_source(): start = pd.datetime(1990, 1, 3, 0, 0, 0, 0, pytz.utc) end = pd.datetime(1990, 1, 8, 0, 0, 0, 0, pytz.utc) index = pd.DatetimeIndex(start=start, end=end, freq=pd.datetools.day) - x = np.arange(0, 12).reshape((6, 2)) + x = np.arange(2., 14.).reshape((6, 2)) df = pd.DataFrame(x, index=index, columns=[0, 1]) return DataFrameSource(df), df