diff --git a/tests/test_sources.py b/tests/test_sources.py index 1257ac55..837ee459 100644 --- a/tests/test_sources.py +++ b/tests/test_sources.py @@ -25,11 +25,9 @@ class TestDataFrameSource(TestCase): for expected_dt, expected_price in df.iterrows(): sid0 = source.next() - sid1 = source.next() - assert expected_dt == sid0.dt == sid1.dt + assert expected_dt == sid0.dt assert expected_price[0] == sid0.price - assert expected_price[1] == sid1.price def test_sid_filtering(self): _, df = factory.create_test_df_source() diff --git a/tests/test_transforms.py b/tests/test_transforms.py index ae340ca8..2f1b2d0d 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -123,11 +123,7 @@ class TestEventWindow(TestCase): # 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 - # window then shrinks once the day completes - assert lengths == [1, 2, 3, 3, 3, 4, 5, 5, 5, 3, 4, 3] + assert lengths == [1, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3] assert window.added == events assert window.removed == events[:-3] @@ -146,7 +142,7 @@ class TestEventWindow(TestCase): # Record the length of the window after each event. lengths.append(len(window.ticks)) - assert lengths == [1, 2, 3, 3, 2] + assert lengths == [1, 2, 2, 2, 2] assert window.added == events assert window.removed == events[:-2] @@ -317,23 +313,16 @@ class TestBatchTransform(TestCase): def test_event_window(self): algo = BatchTransformAlgorithm() algo.run(self.source) - - self.assertEqual(algo.history_return_price_class[:2], - [None, None], + wl = algo.window_length + self.assertEqual(algo.history_return_price_class[:wl], + [None] * wl, "First two iterations should return None") - self.assertEqual(algo.history_return_price_decorator[:2], - [None, None], + self.assertEqual(algo.history_return_price_decorator[:wl], + [None] * wl, "First two iterations should return None") - self.assertEqual(algo.history_return_price_market_aware[:2], - [None, None], - "First two iterations should return None") - self.assertEqual(algo.history_return_more_days_than_refresh[:3], - [None, None, None], - "First five iterations should return None") self.assertTrue(isinstance( - algo.history_return_more_days_than_refresh[4], - pd.DataFrame), - "Sixth iteration should not be None" + algo.history_return_price_class[wl + 1], + pd.DataFrame) ) # Test whether arbitrary fields can be added to datapanel @@ -344,27 +333,21 @@ class TestBatchTransform(TestCase): ) self.assertTrue(all( - field['arbitrary'].values.flatten() == ['test'] * 8), + field['arbitrary'].values.flatten() == + [123] * algo.window_length), 'arbitrary dataframe should contain only "test"' ) # test overloaded class for test_history in [algo.history_return_price_class, algo.history_return_price_decorator]: - np.testing.assert_array_equal( - range(2, 8), - test_history[2].values.flatten() - ) - - np.testing.assert_array_equal( - range(2, 8), - test_history[3].values.flatten() - ) - - np.testing.assert_array_equal( - range(4, 12), - test_history[4].values.flatten() - ) + # starting at window length, the window should contain + # consecutive (of window length) numbers up till the end. + for i in range(algo.window_length, len(test_history)): + np.testing.assert_array_equal( + range(i - algo.window_length + 1, i + 1), + test_history[i].values.flatten() + ) def test_passing_of_args(self): algo = BatchTransformAlgorithm(1, kwarg='str') @@ -375,29 +358,5 @@ class TestBatchTransform(TestCase): expected_item = ((1, ), {'kwarg': 'str'}) self.assertEqual( algo.history_return_args, - [None, None, expected_item, expected_item, - expected_item, expected_item]) - - -class TestBatchTransformMarketAware(TestCase): - def setUp(self): - setup_logger(self) - start = pd.datetime(1993, 1, 1, 0, 0, 0, 0, pytz.utc) - end = pd.datetime(1994, 1, 1, 0, 0, 0, 0, pytz.utc) - - self.data = factory.load_from_yahoo(stocks=['AAPL'], - indexes={}, - start=start, end=end) - - def test_event_window(self): - days = 50 - algo = BatchTransformAlgorithm(days=days, refresh_period=days) - algo.run(self.data) - - self.assertEqual(algo.history_return_price_market_aware[:days], - [None] * days, - "First {days} iterations should return None" - .format(days=days)) - self.assertFalse(algo.history_return_price_market_aware[days + 1] - is None, - "Window is contains too many Nones.") + [None, None, None, expected_item, expected_item, + expected_item]) diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index 8fa80b22..99f92a02 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -214,7 +214,6 @@ class TimeoutAlgorithm(TradingAlgorithm): time.sleep(100) pass -from datetime import timedelta from zipline.algorithm import TradingAlgorithm from zipline.transforms import BatchTransform, batch_transform from zipline.transforms import MovingAverage @@ -237,6 +236,7 @@ class TestRegisterTransformAlgorithm(TradingAlgorithm): class ReturnPriceBatchTransform(BatchTransform): def get_value(self, data): + assert data.shape[1] == self.window_length return data.price @@ -257,7 +257,7 @@ def return_data(data, *args, **kwargs): class BatchTransformAlgorithm(TradingAlgorithm): def initialize(self, *args, **kwargs): - self.refresh_period = kwargs.pop('refresh_period', 2) + self.refresh_period = kwargs.pop('refresh_period', 1) self.window_length = kwargs.pop('window_length', 3) self.args = args @@ -266,46 +266,47 @@ class BatchTransformAlgorithm(TradingAlgorithm): self.history_return_price_class = [] self.history_return_price_decorator = [] self.history_return_args = [] - self.history_return_price_market_aware = [] - self.history_return_more_days_than_refresh = [] self.history_return_arbitrary_fields = [] + self.history_return_nan = [] self.return_price_class = ReturnPriceBatchTransform( - market_aware=False, refresh_period=self.refresh_period, - delta=timedelta(days=self.window_length) + window_length=self.window_length, + fillna=False ) self.return_price_decorator = return_price_batch_decorator( - market_aware=False, refresh_period=self.refresh_period, - delta=timedelta(days=self.window_length) + window_length=self.window_length, + fillna=False ) self.return_args_batch = return_args_batch_decorator( - market_aware=False, refresh_period=self.refresh_period, - delta=timedelta(days=self.window_length) + window_length=self.window_length, + fillna=False ) self.return_price_market_aware = ReturnPriceBatchTransform( - market_aware=True, refresh_period=self.refresh_period, - window_length=self.window_length - ) - - self.return_price_more_days_than_refresh = ReturnPriceBatchTransform( - market_aware=True, - refresh_period=1, - window_length=3 + window_length=self.window_length, + fillna=False ) self.return_arbitrary_fields = return_data( - market_aware=True, - refresh_period=1, - window_length=3 + refresh_period=self.refresh_period, + window_length=self.window_length, + fillna=False ) + self.return_nan = return_price_batch_decorator( + refresh_period=self.refresh_period, + window_length=self.window_length, + fillna=True + ) + + self.iter = 0 + self.set_slippage(FixedSlippage()) def handle_data(self, data): @@ -316,18 +317,28 @@ class BatchTransformAlgorithm(TradingAlgorithm): self.history_return_args.append( self.return_args_batch.handle_data( data, *self.args, **self.kwargs)) - self.history_return_price_market_aware.append( - self.return_price_market_aware.handle_data(data)) - self.history_return_more_days_than_refresh.append( - self.return_price_more_days_than_refresh.handle_data(data)) new_data = deepcopy(data) for sid in new_data: - new_data[sid]['arbitrary'] = 'test' + new_data[sid]['arbitrary'] = 123 self.history_return_arbitrary_fields.append( self.return_arbitrary_fields.handle_data(new_data)) + # nan every second event price + if self.iter % 2 == 0: + self.history_return_nan.append( + self.return_nan.handle_data(data)) + else: + nan_data = deepcopy(data) + import numpy as np + for sid in nan_data.iterkeys(): + nan_data[sid].price = np.nan + self.history_return_nan.append( + self.return_nan.handle_data(nan_data)) + + self.iter += 1 + class SetPortfolioAlgorithm(TradingAlgorithm): """ diff --git a/zipline/transforms/utils.py b/zipline/transforms/utils.py index f1e7b551..c65644ca 100644 --- a/zipline/transforms/utils.py +++ b/zipline/transforms/utils.py @@ -239,17 +239,9 @@ class EventWindow(object): # adding new ticks. self.handle_add(event) - if self.market_aware: - self.add_new_holidays(event.dt) - # Clear out any expired events. drop_condition changes depending # on whether or not we are running in market_aware mode. - # - # oldest newest - # | | - # V V - while self.drop_condition(self.ticks[0].dt, self.ticks[-1].dt): - + while self.drop_condition(): # popleft removes and returns the oldest tick in self.ticks popped = self.ticks.popleft() @@ -257,36 +249,17 @@ class EventWindow(object): # 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 - # self.ticks[-1]) - while len(self.all_holidays) > 0 and self.all_holidays[0] <= newest: - self.cur_holidays.append(self.all_holidays.popleft()) + def out_of_market_window(self): + # Find number of unique days in window + # Note that this assumes that each day we received an + # event is a trading day. + unique_dts = set([event.dt.date() for event in self.ticks]) - def drop_old_holidays(self, oldest): - # Drop from our tracked window any holidays that are older - # than our oldest tracked event. (oldest should always - # be self.ticks[0]) - while len(self.cur_holidays) > 0 and self.cur_holidays[0] < oldest: - self.cur_holidays.popleft() + return len(unique_dts) > self.window_length - def out_of_market_window(self, oldest, newest): - self.drop_old_holidays(oldest) - 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.window_length - - def out_of_delta(self, oldest, newest): - return (newest - oldest) >= self.delta + def out_of_delta(self): + # newest - oldest + return (self.ticks[-1].dt - self.ticks[0].dt) >= self.delta # All event windows expect to receive events with datetime fields # that arrive in sorted order. @@ -344,19 +317,19 @@ class BatchTransform(EventWindow): def __init__(self, func=None, refresh_period=None, - market_aware=True, - delta=None, - window_length=None): + window_length=None, + fillna=True): - super(BatchTransform, self).__init__(market_aware, - window_length=window_length, - delta=delta) + super(BatchTransform, self).__init__(True, + window_length=window_length) if func is not None: self.compute_transform_value = func else: self.compute_transform_value = self.get_value + self.fillna = fillna + self.refresh_period = refresh_period self.window_length = window_length self.trading_days_since_update = 0 @@ -366,7 +339,7 @@ class BatchTransform(EventWindow): self.last_dt = None self.updated = False - self.data = None + self.cached = None self.field_names = None @@ -394,20 +367,22 @@ class BatchTransform(EventWindow): # return newly computed or cached value return self.get_transform_value(*args, **kwargs) - def handle_add(self, event): - if not self.last_dt: - self.last_dt = event.dt - return - + def _extract_field_names(self, event): # extract field names from sids (price, volume etc), make sure # every sid has the same fields. sid_keys = [set(sid.keys()) for sid in event.data.itervalues()] assert sid_keys[0] == set.intersection(*sid_keys),\ "Each sid must have the same keys." - if self.field_names is None: - unwanted_fields = set(['portfolio', 'sid', 'dt', 'type', - 'datetime']) - self.field_names = sid_keys[0] - unwanted_fields + + unwanted_fields = set(['portfolio', 'sid', 'dt', 'type', + 'datetime', 'source_id']) + return sid_keys[0] - unwanted_fields + + def handle_add(self, event): + if not self.last_dt: + self.field_names = self._extract_field_names(event) + self.last_dt = event.dt + return # update trading day counters if self.last_dt.day != event.dt.day: @@ -419,13 +394,11 @@ class BatchTransform(EventWindow): self.trading_days_total >= self.window_length and self.trading_days_since_update >= self.refresh_period ): - - # Create datapanel of running event window. - self.data = self.get_data() # Setting updated to True will cause get_transform_value() # to call the user-defined batch-transform with the most # recent datapanel self.updated = True + self.full = True self.trading_days_since_update = 0 else: self.updated = False @@ -442,36 +415,28 @@ 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 = {} + sids = set.union(*[set(tick.data.keys()) for tick in self.ticks]) + dts = [tick.dt for tick in self.ticks] - for field_name in self.field_names: - sids = self.ticks[0].data.keys() + data = pd.Panel(items=self.field_names, major_axis=dts, + minor_axis=sids) - values_per_sid = {} + # Fill data panel + for tick in self.ticks: + dt = tick.dt + for sid, fields in tick.data.iteritems(): + for field_name in self.field_names: + data[field_name][sid].ix[dt] = fields[field_name] - for sid in sids: - values_per_sid[sid] = pd.Series( - {tick.data[sid].dt: tick.data[sid][field_name] - for tick in self.ticks} - ) + if self.fillna: + # Fills in gaps of missing data during transform + # of multiple stocks. E.g. we may be missing + # minute data because of illiquidity of one stock + data = data.fillna(method='ffill') - # concatenate different sids into one df - df = pd.DataFrame.from_dict(values_per_sid) - # Fills in gaps of missing data during transform of multiple - # stocks. - # e.g. we may be missing minute data because of illiquidity - # of one stock - df = df.fillna(method='ffill') - # Drop any empty rows after the fill. - # This will drop a leading row of N/A - df = df.dropna() - - fields[field_name] = df - - data = pd.Panel.from_dict(fields, orient='items') + # Drop any empty rows after the fill. + # This will drop a leading row of N/A + data = data.dropna(axis=1) return data @@ -491,11 +456,11 @@ class BatchTransform(EventWindow): has actually been updated. Otherwise, the previously, cached value will be returned. """ - if self.data is None: + if not self.full: return None if self.updated: - self.cached = self.compute_transform_value(self.data, + self.cached = self.compute_transform_value(self.get_data(), *args, **kwargs) return self.cached diff --git a/zipline/utils/factory.py b/zipline/utils/factory.py index 6d4ade15..4f2d2ce5 100644 --- a/zipline/utils/factory.py +++ b/zipline/utils/factory.py @@ -272,9 +272,9 @@ 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., len(index) * 2 + 2).reshape((-1, 2)) + x = np.arange(0, len(index)) - df = pd.DataFrame(x, index=index, columns=[0, 1]) + df = pd.DataFrame(x, index=index, columns=[0]) return DataFrameSource(df), df