diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py new file mode 100644 index 00000000..84f94008 --- /dev/null +++ b/tests/test_algorithm.py @@ -0,0 +1,57 @@ +from unittest2 import TestCase +from datetime import timedelta + +from zipline.utils.test_utils import setup_logger +import zipline.utils.factory as factory +from zipline.test_algorithms import TestRegisterTransformAlgorithm +from zipline.gens.tradegens import SpecificEquityTrades, DataFrameSource +from zipline.gens.mavg import MovingAverage + +class TestTransformAlgorithm(TestCase): + def setUp(self): + setup_logger(self) + self.trading_environment = factory.create_trading_environment() + setup_logger(self) + + trade_history = factory.create_trade_history( + 133, + [10.0, 10.0, 11.0, 11.0], + [100, 100, 100, 300], + timedelta(days=1), + self.trading_environment + ) + self.source = SpecificEquityTrades(event_list=trade_history) + + self.df_source, self.df = factory.create_test_df_source() + + def test_source_as_input(self): + algo = TestRegisterTransformAlgorithm(sids=[133]) + algo.run(self.source) + self.assertEqual(len(algo.sources), 1) + assert isinstance(algo.sources[0], SpecificEquityTrades) + + def test_multi_source_as_input_no_start_end(self): + algo = TestRegisterTransformAlgorithm(sids=[133]) + with self.assertRaises(AssertionError): + algo.run([self.source, self.df_source]) + + def test_multi_source_as_input(self): + algo = TestRegisterTransformAlgorithm(sids=[0, 1, 133]) + algo.run([self.source, self.df_source], start=self.df.index[0], end=self.df.index[-1]) + self.assertEqual(len(algo.sources), 2) + + def test_df_as_input(self): + algo = TestRegisterTransformAlgorithm(sids=[0, 1]) + algo.run(self.df) + assert isinstance(algo.sources[0], DataFrameSource) + + 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'] == {'days': 2, 'market_aware': True} + assert algo.registered_transforms['mavg']['class'] is MovingAverage + + diff --git a/tests/test_optimize.py b/tests/test_optimize.py index 365a2a23..6d5ccce2 100644 --- a/tests/test_optimize.py +++ b/tests/test_optimize.py @@ -5,12 +5,7 @@ from collections import defaultdict import numpy as np -from zipline.core.devsimulator import AddressAllocator -# TODO: refactor the factory to use generators -# from zipline.optimize.factory import create_predictable_zipline - -DEFAULT_TIMEOUT = 15 # seconds -EXTENDED_TIMEOUT = 90 +from zipline.optimize.factory import create_predictable_zipline from zipline.utils.test_utils import setup_logger, teardown_logger @@ -24,7 +19,7 @@ class TestUpDown(TestCase): def setUp(self): self.zipline_test_config = { - 'sid' : 133, + 'sid' : [0], 'trade_count' : 5, 'amplitude' : 30, 'base_price' : 50 @@ -36,7 +31,6 @@ class TestUpDown(TestCase): teardown_logger(self) @skip - @timed(DEFAULT_TIMEOUT) def test_source_and_orders(self): """verify that UpDownSource is having the correct behavior and that BuySellAlgorithm places the buy/sell @@ -44,17 +38,17 @@ class TestUpDown(TestCase): UpDownSource and BuySellAlgorithm interact correctly." """ - zipline, config = create_predictable_zipline( + + algo, config = create_predictable_zipline( self.zipline_test_config, - offset=0, - simulate=False + offset=0 ) #extract arguments base_price = self.zipline_test_config['base_price'] amplitude = self.zipline_test_config['amplitude'] - prices = np.array([event.price for event in config['trade_source'].event_list]) + prices = config['trade_source'][0].values max_price_idx = np.where(prices==prices.max())[0] min_price_idx = np.where(prices==prices.min())[0] self.assertTrue(np.all(max_price_idx % 2 == 1), @@ -70,9 +64,9 @@ class TestUpDown(TestCase): "Minimum price does not equal expected maximum price." ) - zipline.simulate(blocking=True) + stats = algo.run(config['trade_source']) - algo = config['algorithm'] + self.assertTrue(len(stats) != 0) orders = np.asarray(algo.orders) max_order_idx = np.where(orders==orders.max())[0] @@ -108,12 +102,15 @@ class TestUpDown(TestCase): compound_returns = np.empty(len(test_offsets)) ziplines = [] for i, offset in enumerate(test_offsets): - zipline, config = create_predictable_zipline( + algo, config = create_predictable_zipline( self.zipline_test_config, offset=offset, ) - ziplines.append(zipline) - compound_returns[i] = zipline.get_cumulative_performance()['returns'] + results = algo.run(config['trade_source']) + ziplines.append(algo) + + compound_returns[i] = results.returns.sum() + self.assertTrue(np.all(compound_returns[supposed_max] > compound_returns[np.logical_not(supposed_max)]), "Maximum compound returns are not where they are supposed to be." diff --git a/tests/test_sources.py b/tests/test_sources.py new file mode 100644 index 00000000..548f8754 --- /dev/null +++ b/tests/test_sources.py @@ -0,0 +1,22 @@ +from unittest2 import TestCase + +import zipline.utils.factory as factory +from zipline.gens.tradegens import DataFrameSource + +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() + + 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 d617ce20..ea1024f6 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -1,5 +1,5 @@ import pytz -import numpy +import numpy as np from datetime import timedelta, datetime from unittest2 import TestCase @@ -15,9 +15,10 @@ 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.test_algorithms import BatchTransformAlgorithm + def to_dt(msg): return ndict({'dt': msg}) @@ -42,26 +43,26 @@ class EventWindowTestCase(TestCase): def setUp(self): setup_logger(self) - + self.monday = datetime(2012, 7, 9, 16, tzinfo=pytz.utc) - self.eleven_normal_days = [self.monday + i*timedelta(days=1) + self.eleven_normal_days = [self.monday + i*timedelta(days=1) for i in xrange(11)] # Modify the end of the period slightly to exercise the # incomplete day logic. self.eleven_normal_days[-1] -= timedelta(minutes = 1) self.eleven_normal_days.append(self.monday+timedelta(days=11,seconds=1)) - + # Second set of dates to test holiday handling. self.jul4_monday = datetime(2012, 7, 2, 16, tzinfo=pytz.utc) self.week_of_jul4 = [self.jul4_monday + i*timedelta(days=1) for i in xrange(5)] def test_event_window_with_timedelta(self): - + # Keep all events within a 5 minute window. window = NoopEventWindow( - market_aware = False, + market_aware = False, delta = timedelta(minutes = 5), days = None ) @@ -91,7 +92,7 @@ class EventWindowTestCase(TestCase): def test_market_aware_window_normal_week(self): window = NoopEventWindow( - market_aware = True, + market_aware = True, delta = None, days = 3 ) @@ -102,7 +103,7 @@ class EventWindowTestCase(TestCase): window.update(event) # Record the length of the window after each event. lengths.append(len(window.ticks)) - + # The window stretches out during the weekend because we wait # to drop events until the weekend ends. The last window is # briefly longer because it doesn't complete a full day. The @@ -113,7 +114,7 @@ class EventWindowTestCase(TestCase): def test_market_aware_window_holiday(self): window = NoopEventWindow( - market_aware = True, + market_aware = True, delta = None, days = 2 ) @@ -125,11 +126,11 @@ class EventWindowTestCase(TestCase): window.update(event) # Record the length of the window after each event. lengths.append(len(window.ticks)) - + assert lengths == [1, 2, 3, 3, 2] assert window.added == events assert window.removed == events[:-2] - + def tearDown(self): setup_logger(self) @@ -186,7 +187,7 @@ class FinanceTransformsTestCase(TestCase): expected = [0.0, 0.0, 0.1, 0.0] assert tnfm_vals == expected - + # Two-day returns. An extra kink here is that the # factory will automatically skip a weekend for the # last event. Results shouldn't notice this blip. @@ -222,12 +223,12 @@ class FinanceTransformsTestCase(TestCase): fields = ['price', 'volume'], delta = timedelta(days = 2), ) - + transformed = list(mavg.transform(self.source)) # Output values. tnfm_prices = [message.tnfm_value.price for message in transformed] tnfm_volumes = [message.tnfm_value.volume for message in transformed] - + # "Hand-calculated" values expected_prices = [ ((10.0) / 1.0), @@ -267,16 +268,16 @@ class FinanceTransformsTestCase(TestCase): transformed = list(stddev.transform(self.source)) vals = [message.tnfm_value for message in transformed] - + expected = [ None, - numpy.std([10.0, 15.0], ddof = 1), - numpy.std([10.0, 15.0, 13.0], ddof = 1), - numpy.std([15.0, 13.0, 12.0], ddof = 1), + np.std([10.0, 15.0], ddof = 1), + np.std([10.0, 15.0, 13.0], ddof = 1), + np.std([15.0, 13.0, 12.0], ddof = 1), ] - # numpy has odd rounding behavior, cf. - # http://docs.scipy.org/doc/numpy/reference/generated/numpy.std.html + # np has odd rounding behavior, cf. + # http://docs.scipy.org/doc/np/reference/generated/np.std.html for v1, v2 in zip(vals, expected): if v1 == None: @@ -285,8 +286,23 @@ class FinanceTransformsTestCase(TestCase): assert round(v1, 5) == round(v2, 5) +############################################################ +# Test BatchTransform - - - +class BatchTransformTestCase(TestCase): + def setUp(self): + setup_logger(self) + self.source, self.df = factory.create_test_df_source() + + def test_batch_inherit(self): + algo = BatchTransformAlgorithm(sids=[0, 1]) + algo.run(self.source) + + assert algo.history_class[:2] == algo.history_decorator[:2] == [None, None], "First two iterations should return None" + + # test overloaded class + 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/algorithm.py b/zipline/algorithm.py new file mode 100644 index 00000000..08d388da --- /dev/null +++ b/zipline/algorithm.py @@ -0,0 +1,184 @@ +import pandas as pd +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 + + +class TradingAlgorithm(object): + """Base class for trading algorithms. Inherit and overload + initialize() and handle_data(data). + + A new algorithm could look like this: + ``` + class MyAlgo(TradingAlgorithm): + def initialize(amount): + self.amount = amount + + def handle_data(data): + sid = self.sids[0] + self.order(sid, amount) + ``` + To then to run this algorithm: + + >>> my_algo = MyAlgo(100, sids=[0]) + >>> stats = my_algo.run(data) + + """ + def __init__(self, sids, *args, **kwargs): + """ + Initialize sids and other state variables. + + Calls user-defined initialize and forwarding *args and **kwargs. + """ + self.sids = sids + self.done = False + self.order = None + self.frame_count = 0 + self.portfolio = None + + self.registered_transforms = {} + + # call to user-defined initialize method + self.initialize(*args, **kwargs) + + def _create_simulator(self, start, end): + """ + 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 + + transforms.append(sf) + + # SimulatedTrading is the main class handling data streaming, + # application of transforms and calling of the user algo. + return SimulatedTrading( + self.sources, + transforms, + self, + environment, + FixedSlippage() + ) + + def run(self, source, start=None, end=None): + """Run the algorithm. + + :Arguments: + 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, (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 + self.simulated_trading = self._create_simulator(start=start, end=end) + + # loop through simulated_trading, each iteration returns a + # perf ndict + perfs = list(self.simulated_trading) + + # convert perf ndict to pandas dataframe + daily_stats = self._create_daily_stats(perfs) + + return daily_stats + + + def _create_daily_stats(self, perfs): + # create daily and cumulative stats dataframe + daily_perfs = [] + cum_perfs = [] + for perf in perfs: + if 'daily_perf' in perf: + daily_perfs.append(perf['daily_perf']) + else: + cum_perfs.append(perf) + + daily_dts = [np.datetime64(perf['period_close'], utc=True) for perf in daily_perfs] + daily_stats = pd.DataFrame(daily_perfs, index=daily_dts) + + return daily_stats + + + def add_transform(self, transform_class, tag, *args, **kwargs): + """Add a single-sid, sequential transform to the model. + + :Arguments: + transform_class : class + Which transform to use. E.g. mavg. + tag : str + How to name the transform. Can later be access via: + data[sid].tag() + + Extra args and kwargs will be forwarded to the transform + instantiation. + + """ + self.registered_transforms[tag] = {'class': transform_class, + 'args': args, + 'kwargs': kwargs} + + def set_portfolio(self, portfolio): + self.portfolio = portfolio + + 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): + pass + + def set_slippage_override(self, slippage_callable): + pass + + + diff --git a/zipline/gens/composites.py b/zipline/gens/composites.py index 05839a28..7effcf19 100644 --- a/zipline/gens/composites.py +++ b/zipline/gens/composites.py @@ -1,10 +1,7 @@ -from itertools import tee, chain +from itertools import chain from zipline.gens.utils import roundrobin, done_message from zipline.gens.sort import date_sort -from zipline.gens.merge import merge -from zipline.gens.transform import StatefulTransform - def date_sorted_sources(*sources): """ @@ -14,7 +11,7 @@ def date_sorted_sources(*sources): for source in sources: assert iter(source), "Source %s not iterable" % source - assert 'get_hash' in source.__class__.__dict__, "No get_hash" + assert hasattr(source, 'get_hash'), "No get_hash" # Get name hashes to pass to date_sort. names = [source.get_hash() for source in sources] @@ -29,46 +26,6 @@ def date_sorted_sources(*sources): return date_sort(stream_in, names) - -def merged_transforms(sorted_stream, *transforms): - """ - A generator that takes the expected output of a date_sort, pipes - it through a given set of transforms, and runs the results - through a merge to output a unified stream. tnfms should be a - list of pointers to generator functions. tnfm_args should be a - list of tuples, representing the arguments to be passed to each - transform. tnfm_kwargs should be a list of dictionaries - representing keyword arguments to each transform. - """ - for transform in transforms: - assert isinstance(transform, StatefulTransform) - transform.merged = True - transform.sequential = False - - # Generate expected hashes for each transform - namestrings = [tnfm.get_hash() for tnfm in transforms] - - # Create a copy of the stream for each transform. - split = tee(sorted_stream, len(transforms)) - - # Package a stream copy with each StatefulTransform instance. - bundles = zip(transforms, split) - - # Convert the copies into transform streams. - tnfm_gens = [tnfm.transform(stream) for tnfm, stream in bundles] - - # Roundrobin the outputs of our transforms to create a single flat - # stream. - to_merge = roundrobin(tnfm_gens, namestrings) - - # Pipe the stream into merge. - merged = merge(to_merge, namestrings) - - dt_aliased = alias_dt(merged) - # Return the merged events. - return add_done(dt_aliased) - - def sequential_transforms(stream_in, *transforms): """ Apply each transform in transforms sequentially to each event in stream_in. @@ -87,6 +44,7 @@ def sequential_transforms(stream_in, *transforms): transforms, stream_in) + dt_aliased = alias_dt(stream_out) return add_done(dt_aliased) diff --git a/zipline/gens/tradegens.py b/zipline/gens/tradegens.py index 0842b3bd..f088ae8b 100644 --- a/zipline/gens/tradegens.py +++ b/zipline/gens/tradegens.py @@ -7,7 +7,12 @@ import pytz 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 from zipline.gens.utils import hash_args, create_trade def date_gen(start = datetime(2006, 6, 6, 12, tzinfo=pytz.utc), @@ -73,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) @@ -137,6 +156,7 @@ class SpecificEquityTrades(object): volumes = mock_volumes(self.count) sids = cycle(self.sids) + # Combine the iterators into a single iterator of arguments arg_gen = izip(sids, prices, volumes, dates) @@ -157,33 +177,54 @@ class SpecificEquityTrades(object): return filtered -# !!!!!!! Deprecated for now !!!!!!!!! +class DataFrameSource(SpecificEquityTrades): + """ + Yields all events in event_list that match the given sid_filter. + If no event_list is specified, generates an internal stream of events + to filter. Returns all events if filter is None. -def RandomEquityTrades(object): + Configuration options: - def __init__(self): - # We shouldn't get any positional args. - assert args == () + count : integer representing number of trades + sids : list of values representing simulated internal sids + start : start date + delta : timedelta between internal events + filter : filter to remove the sids + """ - self.count = config.get('count', 500) - self.sids = config.get('sids', [1,2]) - self.filter = config.get('filter') + def __init__(self, data, **kwargs): + assert isinstance(data.index, pd.tseries.index.DatetimeIndex) - dates = fuzzy_dates(count) - prices = mock_prices(count, rand = True) - volumes = mock_volumes(count, rand = True) - sids = cycle(sids) + self.data = data + # Unpack config dictionary with default values. + self.count = kwargs.get('count', len(data)) + self.sids = kwargs.get('sids', data.columns) + self.start = kwargs.get('start', data.index[0]) + self.end = kwargs.get('end', data.index[-1]) + self.delta = kwargs.get('delta', data.index[1]-data.index[0]) - arg_gen = izip(sids, prices, volumes, dates) + # Hash_value for downstream sorting. + self.arg_string = hash_args(data, **kwargs) - unfiltered = (create_trade(*args) for args in arg_gen) + self.generator = self.create_fresh_generator() - if filter: - filtered = ifilter(lambda event: event.sid in filter, unfiltered) - else: - filtered = unfiltered - return filtered + def create_fresh_generator(self): + def _generator(df=self.data): + for dt, series in df.iterrows(): + if (dt < self.start) or (dt > self.end): + continue + event = {'dt': dt, + 'source_id': self.get_hash(), + 'type': DATASOURCE_TYPE.TRADE + } -# if __name__ == "__main__": -# import nose.tools; nose.tools.set_trace() -# trades = SpecificEquityTrades(filter = [1]) + for sid, price in series.iterkv(): + event = copy(event) + event['sid'] = sid + event['price'] = price + + yield ndict(event) + + # Return the filtered event stream. + 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 9c2df532..5efb1505 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -2,6 +2,7 @@ from logbook import Logger, Processor from datetime import datetime from itertools import groupby +from operator import attrgetter from zipline import ndict from zipline.utils.timeout import Heartbeat, Timeout @@ -216,62 +217,61 @@ class AlgorithmSimulator(object): # Capture any output of this generator to stdout and pipe it # to a logbook interface. Also inject the current algo # snapshot time to any log record generated. - with self.processor.threadbound(), self.stdout_capture(Logger('Print'),''): + #with self.processor.threadbound(), self.stdout_capture(Logger('Print'),''): - # Call user's initialize method with a timeout. - with Timeout(INIT_TIMEOUT, message="Call to initialize timed out"): - self.algo.initialize() + # Call user's initialize method with a timeout. + with Timeout(INIT_TIMEOUT, message="Call to initialize timed out"): + self.algo.initialize() - # Group together events with the same dt field. This depends on the - # events already being sorted. - for date, snapshot in groupby(stream_in, lambda e: e.dt): + # Group together events with the same dt field. This depends on the + # events already being sorted. + for date, snapshot in groupby(stream_in, attrgetter('dt')): + # Set the simulation date to be the first event we see. + # This should only occur once, at the start of the test. + if self.simulation_dt == None: + self.simulation_dt = date - # Set the simulation date to be the first event we see. - # This should only occur once, at the start of the test. - if self.simulation_dt == None: - self.simulation_dt = date + # Done message has the risk report, so we yield before exiting. + if date == 'DONE': + for event in snapshot: + yield event.perf_message + raise StopIteration() - # Done message has the risk report, so we yield before exiting. - if date == 'DONE': - for event in snapshot: + # We're still in the warmup period. Use the event to + # update our universe, but don't yield any perf messages, + # and don't send a snapshot to handle_data. + elif date < self.algo_start: + for event in snapshot: + del event['perf_message'] + self.update_universe(event) + + # The algo has taken so long to process events that + # its simulated time is later than the event time. + # Update the universe and yield any perf messages + # encountered, but don't call handle_data. + elif date < self.simulation_dt: + for event in snapshot: + # Only yield if we have something interesting to say. + if event.perf_message != None: yield event.perf_message - raise StopIteration() + # Delete the message before updating so we don't send it + # to the user. + del event['perf_message'] + self.update_universe(event) - # We're still in the warmup period. Use the event to - # update our universe, but don't yield any perf messages, - # and don't send a snapshot to handle_data. - elif date < self.algo_start: - for event in snapshot: - del event['perf_message'] - self.update_universe(event) + # Regular snapshot. Update the universe and send a snapshot + # to handle data. + else: + for event in snapshot: + # Only yield if we have something interesting to say. + if event.perf_message != None: + yield event.perf_message + del event['perf_message'] - # The algo has taken so long to process events that - # its simulated time is later than the event time. - # Update the universe and yield any perf messages - # encountered, but don't call handle_data. - elif date < self.simulation_dt: - for event in snapshot: - # Only yield if we have something interesting to say. - if event.perf_message != None: - yield event.perf_message - # Delete the message before updating so we don't send it - # to the user. - del event['perf_message'] - self.update_universe(event) + self.update_universe(event) - # Regular snapshot. Update the universe and send a snapshot - # to handle data. - else: - for event in snapshot: - # Only yield if we have something interesting to say. - if event.perf_message != None: - yield event.perf_message - del event['perf_message'] - - self.update_universe(event) - - # Send the current state of the universe to the user's algo. - self.simulate_snapshot(date) + # Send the current state of the universe to the user's algo. + self.simulate_snapshot(date) def update_universe(self, event): """ diff --git a/zipline/gens/transform.py b/zipline/gens/transform.py index 03eb41ea..8ec862eb 100644 --- a/zipline/gens/transform.py +++ b/zipline/gens/transform.py @@ -9,6 +9,8 @@ from datetime import datetime from collections import deque from abc import ABCMeta, abstractmethod +import pandas as pd + from zipline import ndict from zipline.utils.tradingcalendar import non_trading_days from zipline.gens.utils import assert_sort_unframe_protocol, hash_args @@ -36,7 +38,7 @@ class TransformMeta(type): still recover an instance of a "raw" Foo by introspecting the resulting StatefulTransform's 'state' field. """ - + def __call__(cls, *args, **kwargs): return StatefulTransform(cls, *args, **kwargs) @@ -53,19 +55,19 @@ class StatefulTransform(object): def __init__(self, tnfm_class, *args, **kwargs): assert isinstance(tnfm_class, (types.ObjectType, types.ClassType)), \ "Stateful transform requires a class." - assert tnfm_class.__dict__.has_key('update'), \ + assert hasattr(tnfm_class, 'update'), \ "Stateful transform requires the class to have an update method" # Flag set inside the Passthrough transform class to signify special # behavior if we are being fed to merged_transforms. - self.passthrough = tnfm_class.__dict__.get('PASSTHROUGH', False) - + self.passthrough = hasattr(tnfm_class, 'PASSTHROUGH') + # Flags specifying how to append the calculated value. # Merged is the default for ease of testing, but we use sequential # in production. self.sequential = False self.merged = True - + # Create an instance of our transform class. if isinstance(tnfm_class, TransformMeta): # Classes derived TransformMeta have their __call__ @@ -104,12 +106,12 @@ class StatefulTransform(object): continue assert_sort_unframe_protocol(message) - + # This flag is set by by merged_transforms to ensure # isolation of messages. if self.merged: message = deepcopy(message) - + tnfm_value = self.state.update(message) # PASSTHROUGH flag means we want to keep all original @@ -133,7 +135,7 @@ class StatefulTransform(object): out_message.tnfm_value = tnfm_value out_message.dt = message.dt yield out_message - + # Sequential flag should be used to add a single new # key-value pair to the event. The new key is this # transform's namestring, and its value is the value @@ -147,9 +149,11 @@ class StatefulTransform(object): out_message = message out_message[self.namestring] = tnfm_value yield out_message - + log.info('Finished StatefulTransform [%s]' % self.get_hash()) + + class EventWindow(object): """ Abstract base class for transform classes that calculate iterative @@ -171,13 +175,13 @@ class EventWindow(object): # Mark this as an abstract base class. __metaclass__ = ABCMeta - def __init__(self, market_aware, days = None, delta = None): + def __init__(self, market_aware, days=None, delta=None): self.market_aware = market_aware self.days = days self.delta = delta - self.ticks = deque() + self.ticks = deque() # Market-aware mode only works with full-day windows. if self.market_aware: @@ -213,12 +217,12 @@ class EventWindow(object): self.assert_well_formed(event) # Add new event and increment totals. - self.ticks.append(event) + self.ticks.append(deepcopy(event)) # Subclasses should override handle_add to define behavior for # adding new ticks. self.handle_add(event) - + if self.market_aware: self.add_new_holidays(event.dt) @@ -229,14 +233,14 @@ class EventWindow(object): # | | # V V while self.drop_condition(self.ticks[0].dt, self.ticks[-1].dt): - + # popleft removes and returns the oldest tick in self.ticks popped = self.ticks.popleft() # Subclasses should override handle_remove to define # behavior for removing ticks. self.handle_remove(popped) - + def add_new_holidays(self, newest): # Add to our tracked window any untracked holidays that are # older than our newest event. (newest should always be @@ -256,12 +260,13 @@ class EventWindow(object): calendar_dates_between = (newest.date() - oldest.date()).days holidays_between = len(self.cur_holidays) trading_days_between = calendar_dates_between - holidays_between - + # "Put back" a day if oldest is earlier in its day than newest, # reflecting the fact that we haven't yet completed the last # day in the window. if oldest.time() > newest.time(): trading_days_between -= 1 + return trading_days_between >= self.days def out_of_delta(self, oldest, newest): @@ -277,3 +282,144 @@ class EventWindow(object): # Something is wrong if new event is older than previous. assert event.dt >= self.ticks[-1].dt, \ "Events arrived out of order in EventWindow: %s -> %s" % (event, self.ticks[0]) + + +class BatchTransform(EventWindow): + """Base class for batch transforms with a trailing window of + variable length. As opposed to pure EventWindows that get a stream + of events and are bound to a single SID, this class creates stream + of pandas DataFrames with each colum representing a sid. + + There are two ways to create a new batch window: + (i) Inherit from BatchTransform and overload get_value(data). + E.g.: + ``` + class MyBatchTransform(BatchTransform): + def get_value(self, data): + # compute difference between the means of sid 0 and sid 1 + return data[0].mean() - data[1].mean() + ``` + + (ii) Use the batch_transform decorator. + E.g.: + ``` + @batch_transform + def my_batch_transform(data): + return data[0].mean() - data[1].mean() + + ``` + + In you algorithm you would then have to instantiate this in the initialize() method: + ``` + self.my_batch_transform = MyBatchTransform() + ``` + + To then use it, inside of the algorithm handle_data(), call the + handle_data() of the BatchTransform and pass it the current event: + ``` + result = self.my_batch_transform(data) + ``` + + """ + + def __init__(self, func=None, refresh_period=None, market_aware=True, delta=None, days=None, sids=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 + + self.full = False + self.last_refresh = None + + self.updated = False + self.data = None + + def handle_data(self, data): + """ + New method to handle a data frame as sent to the algorithm's handle_data + method. + """ + # extract dates + dts = [data[sid].datetime for sid in self.sids] + # 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) + + # append data frame to window. update() will call handle_add() and + # handle_remove() appropriately + self.update(data) + + # return newly computed or cached value + return self.get_transform_value() + + def handle_add(self, event): + if not self.last_refresh: + self.last_refresh = event.dt + return + + age = event.dt - self.last_refresh + if age.days >= self.refresh_period: + # Create a pandas.Panel (i.e. 3d DataFrame) from the + # events in the current window. + # + # The resulting panel looks like this: + # index : field_name (e.g. price) + # major axis/rows : dt + # minor axis/colums : sid + # + # This Panel data structure ultimately gets passed to the + # user-overloaded get_value() method. + fields = {} + for field_name in ['price', 'volume']: + # Skip non-existant fields + if field_name not in self.ticks[0][self.sids[0]]: + continue + + values_per_sid = {} + for sid in self.sids: + values_per_sid[sid] = pd.Series({tick[sid].dt: tick[sid][field_name] for tick in self.ticks}) + + # concatenate different sids into one df + fields[field_name] = pd.DataFrame.from_dict(values_per_sid) + + self.data = pd.Panel.from_dict(fields, orient='items') + + self.updated = True + self.last_refresh = event.dt + else: + self.updated = False + + def handle_remove(self, event): + # since an event is expiring, we know the window is full + self.full = True + + def get_value(self, *args, **kwargs): + raise NotImplementedError("Either overwrite get_value or provide a func argument.") + + def get_transform_value(self, *args, **kwargs): + if self.data is None: + return None + + if self.updated: + self.cached = self.compute_transform_value(self.data, *args, **kwargs) + + return self.cached + + +def batch_transform(func): + """Decorator function to use instead of inheriting from BatchTransform. + For an example on how to use this, see the doc string of BatchTransform. + """ + + def create_window(*args, **kwargs): + # passes the user defined function to BatchTransform which it + # will call instead of self.get_value() + return BatchTransform(*args, func=func, **kwargs) + + return create_window diff --git a/zipline/optimize/algorithms.py b/zipline/optimize/algorithms.py index c4e3fb9d..976979d2 100644 --- a/zipline/optimize/algorithms.py +++ b/zipline/optimize/algorithms.py @@ -1,4 +1,9 @@ -class BuySellAlgorithm(object): +from logbook import Logger +from zipline.algorithm import TradingAlgorithm + +logger = Logger('Algo') + +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 @@ -10,39 +15,19 @@ 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): + def handle_data(self, data): order_size = self.buy_or_sell * (self.amount - (self.offset**2)) - self.order(self.sid, order_size) + self.order(self.sids[0], order_size) + logger.debug("ordering" + str(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] diff --git a/zipline/optimize/example.py b/zipline/optimize/example.py new file mode 100644 index 00000000..c8912558 --- /dev/null +++ b/zipline/optimize/example.py @@ -0,0 +1,159 @@ +# WARNING: This file is still work in progress and contains rather +# random code snippets. + +import pandas as pd + +import numpy as np + +import matplotlib.pyplot as plt +import cProfile +from zipline.gens.mavg import MovingAverage +from zipline.gens.cov import CovTransform, cov +from zipline.algorithm import TradingAlgorithm +from zipline.gens.transform import BatchTransform, batch_transform + +@batch_transform +def cov(data): + return data.price.cov() + +class DMA(TradingAlgorithm): + """Dual Moving Average algorithm. + """ + def initialize(self, amount=100, short_window=20, long_window=40): + self.amount = amount + self.events = 0 + + self.invested = {} + for sid in self.sids: + self.invested[sid] = False + + self.add_transform(MovingAverage, 'short_mavg', ['price'], + market_aware=True, + days=short_window) + + self.add_transform(MovingAverage, 'long_mavg', ['price'], + market_aware=True, + days=long_window) + + self.cov = cov(sids=self.sids, refresh_period=1, days=5) + + def handle_data(self, data): + self.events += 1 + + cov = self.cov.handle_data(data) + print cov + + for sid in self.sids: + # access transforms via their user-defined tag + if (data[sid].short_mavg['price'] > data[sid].long_mavg['price']) and not self.invested[sid]: + self.order(sid, self.amount) + self.invested[sid] = True + elif (data[sid].short_mavg['price'] < data[sid].long_mavg['price']) and self.invested[sid]: + self.order(sid, -self.amount) + self.invested[sid] = False + + +def load_close_px(indexes=None, stocks=None): + from pandas.io.data import DataReader + import pytz + + if indexes is None: + indexes = {'SPX' : '^GSPC'} + if stocks is None: + stocks = ['AAPL'] #, 'GE', 'IBM', 'MSFT', 'XOM', 'AA', 'JNJ', 'PEP'] + + #start = pd.datetime(1990, 1, 1) + start = pd.datetime(1990, 1, 1, 0, 0, 0, 0, pytz.utc) + end = pd.datetime(1992, 1, 1, 0, 0, 0, 0, pytz.utc) #pd.datetime.today() + + data = {} + + for stock in stocks: + print stock + stkd = DataReader(stock, 'yahoo', start, end).sort_index() + data[stock] = stkd + + for name, ticker in indexes.iteritems(): + print name + stkd = DataReader(ticker, 'yahoo', start, end).sort_index() + data[name] = stkd + + #df = pd.DataFrame({key: d['Close'] for key, d in data.iteritems()}) + df = pd.DataFrame({i: d['Close'] for i, d in enumerate(data.itervalues())}) + df.index = df.index.tz_localize(pytz.utc) + + return df + + +def run((short_window, long_window)): + #data = pd.DataFrame.from_csv('SP500.csv') + #data = pd.DataFrame.from_csv('aapl.csv') #load_close_px() + data = load_close_px() + myalgo = DMA([0, 1], amount=100, short_window=short_window, long_window=long_window) + stats = myalgo.run(data) + stats['sw'] = short_window + stats['lw'] = long_window + return stats + +def explore_params(): + sws, lws = np.mgrid[10:20:5, 10:20:5] + + stats_all = map(run, zip(sws.flatten(), lws.flatten())) + stats = pd.concat(stats_all) + returns = stats.groupby(['sw', 'lw']).sum() + + plt.contourf(sws, lws, returns.returns.reshape(sws.shape)) + plt.xlabel('Short window length') + plt.ylabel('Long window length') + plt.savefig('DMA_contour.png') + plt.show() + +#stats = run((10, 50)) + +def get_opt_holdings_qp(univ_rets, track_rets): + from cvxopt import matrix + from cvxopt.solvers import qp + # set up the QP for CVXOPT + # .5 x' P x + q'x + # P = 2 * R'R + # q = - 2 * bmk'R + R = univ_rets.values + b = track_rets.values + P = matrix(2 * np.dot(R.T, R)) + q = matrix(-2 * np.dot(R.T, b)) + result = qp(P, q) + if result['status'] != 'optimal': + raise Exception('optimum not reached by QP') + return pd.Series(np.array(result['x']).ravel(), index=univ_rets.columns) + +def opt_portfolio(cov, budget, min_return): + from cvxopt import matrix + from cvxopt.solvers import qp + n = len(cov) + cov = matrix(2 * cov) + q = matrix(np.zeros(n)) + + h = matrix(budget) # G*x < h + # coneqp + result = qp(cov, q, h=h) + if result['status'] != 'optimal': + raise Exception('optimum not reached by QP') + + return pd.Series(np.array(result['x']).ravel()) + +def calc_te(weights, univ_rets, track_rets): + port_rets = (univ_rets * weights).sum(1) + return (port_rets - track_rets).std() + +def plot_returns(port_returns, bmk_returns): + plt.figure() + cum_port = ((1 + port_returns).cumprod() - 1) + cum_bmk = ((1 + bmk_returns).cumprod() - 1) + # cum_port = port_returns.cumsum() + # cum_bmk = bmk_returns.cumsum() + cum_port.plot(label='Portfolio returns') + cum_bmk.plot(label='Benchmark') + plt.title('Portfolio performance') + plt.legend(loc='best') + +print run((10, 20)) \ No newline at end of file diff --git a/zipline/optimize/factory.py b/zipline/optimize/factory.py index 8967d28a..14fea2ad 100644 --- a/zipline/optimize/factory.py +++ b/zipline/optimize/factory.py @@ -4,14 +4,12 @@ Factory functions to prepare useful data for optimize tests. Author: Thomas V. Wiecki (thomas.wiecki@gmail.com), 2012 """ from datetime import timedelta -import pandas as pd import zipline.protocol as zp from zipline.utils.factory import get_next_trading_dt, create_trading_environment -from zipline.finance.sources import SpecificEquityTrades +from zipline.gens.tradegens import SpecificEquityTrades from zipline.optimize.algorithms import BuySellAlgorithm -from zipline.lines import SimulatedTrading from zipline.finance.slippage import FixedSlippage from copy import copy @@ -122,7 +120,7 @@ def create_predictable_zipline(config, offset=0, simulate=True): amplitude) if 'algorithm' not in config: - config['algorithm'] = BuySellAlgorithm(sid, 100, offset) + algorithm = BuySellAlgorithm(sids=[sid], amount=100, offset=offset) config['order_count'] = trade_count - 1 config['trade_count'] = trade_count @@ -131,9 +129,4 @@ def create_predictable_zipline(config, offset=0, simulate=True): config['slippage'] = FixedSlippage() config['devel'] = True - zipline = SimulatedTrading.create_test_zipline(**config) - - if simulate: - zipline.simulate(blocking=True) - - return zipline, config + return algorithm, config diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index bcff7f81..10795e33 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -69,6 +69,7 @@ class TestAlgorithm(): self.order = None self.frame_count = 0 self.portfolio = None + if sid_filter: self.sid_filter = sid_filter else: @@ -99,7 +100,7 @@ class TestAlgorithm(): def set_slippage_override(self, slippage_callable): pass - # + class HeavyBuyAlgorithm(): """ 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 7146c0de..7517fff5 100644 --- a/zipline/utils/factory.py +++ b/zipline/utils/factory.py @@ -8,14 +8,14 @@ import random from os.path import join, abspath, dirname from operator import attrgetter +import pandas as pd +import numpy as np from datetime import datetime, timedelta -from zipline.utils.date_utils import tuple_to_date -from zipline.utils.protocol_utils import ndict import zipline.finance.risk as risk - -from zipline.gens.tradegens import RandomEquityTrades -from zipline.gens.tradegens import SpecificEquityTrades +from zipline.utils.date_utils import tuple_to_date +from zipline.utils.protocol_utils import ndict +from zipline.gens.tradegens import SpecificEquityTrades, DataFrameSource from zipline.gens.utils import create_trade from zipline.finance.trading import TradingEnvironment @@ -57,12 +57,15 @@ def load_market_data(): return bm_returns, tr_curves -def create_trading_environment(year=2006): +def create_trading_environment(year=2006, start=None, end=None): """Construct a complete environment with reasonable defaults""" benchmark_returns, treasury_curves = load_market_data() - start = datetime(year, 1, 1, tzinfo=pytz.utc) - end = datetime(year, 12, 31, tzinfo=pytz.utc) + if start is None: + start = datetime(year, 1, 1, tzinfo=pytz.utc) + if end is None: + end = datetime(year, 12, 31, tzinfo=pytz.utc) + trading_environment = TradingEnvironment( benchmark_returns, treasury_curves, @@ -88,7 +91,6 @@ def create_trade_history(sid, prices, amounts, interval, trading_calendar): current = trading_calendar.first_open for price, amount in zip(prices, amounts): - trade = create_trade(sid, price, amount, current) trades.append(trade) current = get_next_trading_dt(current, interval, trading_calendar) @@ -233,3 +235,12 @@ def create_trade_source(sids, trade_count, trade_time_increment, trading_environ #trading_environment.period_end = trade_history[-1].dt return source + +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(2., 14.).reshape((6, 2)) + df = pd.DataFrame(x, index=index, columns=[0, 1]) + + return DataFrameSource(df), df