From b87d454938388919a034d3b378cad070ab39c828 Mon Sep 17 00:00:00 2001 From: Thomas Wiecki Date: Sat, 11 May 2013 15:26:14 -0400 Subject: [PATCH] BUG: Add bar kwarg to batch_transform. Before the change to the RollingPanel, window_length specified the number of days that should be in a window. The previous commit broke this if data was minute resolution. By passing bar='minute' to the batch_transform we internally multiply the window_length by 60*6.5 to have a full day. Also adds a (still rudamentary) test for batch_transform with minute data. --- tests/test_batchtransform.py | 23 +++++++++++++++++++++++ zipline/test_algorithms.py | 21 +++++++++++++++++++++ zipline/transforms/utils.py | 15 +++++++++++++-- zipline/utils/factory.py | 20 ++++++++++++++++++-- 4 files changed, 75 insertions(+), 4 deletions(-) diff --git a/tests/test_batchtransform.py b/tests/test_batchtransform.py index ad73239d..490e35f7 100644 --- a/tests/test_batchtransform.py +++ b/tests/test_batchtransform.py @@ -28,6 +28,7 @@ from zipline.sources.data_source import DataSource import zipline.utils.factory as factory from zipline.test_algorithms import (BatchTransformAlgorithm, + BatchTransformAlgorithmMinute, batch_transform, ReturnPriceBatchTransform) @@ -125,6 +126,28 @@ class TestChangeOfSids(TestCase): self.assertEqual(df[last_elem][last_elem], last_elem) +class TestBatchTransformMinutely(TestCase): + def setUp(self): + start = pd.datetime(1990, 1, 3, 0, 0, 0, 0, pytz.utc) + end = pd.datetime(1990, 1, 8, 0, 0, 0, 0, pytz.utc) + self.sim_params = factory.create_simulation_parameters( + start=start, + end=end, + ) + self.sim_params.emission_rate = 'daily' + self.sim_params.data_frequency = 'minute' + setup_logger(self) + self.source, self.df = \ + factory.create_test_df_source(bars='minute') + + def test_core(self): + algo = BatchTransformAlgorithmMinute(sim_params=self.sim_params) + algo.run(self.source) + wl = int(algo.window_length * 6.5 * 60) + for bt in algo.history[wl:]: + self.assertEqual(len(bt), wl) + + class TestBatchTransform(TestCase): def setUp(self): self.sim_params = factory.create_simulation_parameters( diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index 76668d86..cd0916ce 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -434,6 +434,27 @@ class BatchTransformAlgorithm(TradingAlgorithm): ) +class BatchTransformAlgorithmMinute(TradingAlgorithm): + def initialize(self, *args, **kwargs): + self.refresh_period = kwargs.pop('refresh_period', 1) + self.window_length = kwargs.pop('window_length', 3) + + self.args = args + self.kwargs = kwargs + + self.history = [] + + self.batch_transform = return_price_batch_decorator( + refresh_period=self.refresh_period, + window_length=self.window_length, + clean_nans=False, + bars='minute' + ) + + def handle_data(self, data): + self.history.append(self.batch_transform.handle_data(data)) + + class SetPortfolioAlgorithm(TradingAlgorithm): """ An algorithm that tries to set the portfolio directly. diff --git a/zipline/transforms/utils.py b/zipline/transforms/utils.py index b28fea8c..cd287388 100644 --- a/zipline/transforms/utils.py +++ b/zipline/transforms/utils.py @@ -315,7 +315,8 @@ class BatchTransform(object): clean_nans=True, sids=None, fields=None, - compute_only_full=True): + compute_only_full=True, + bars='daily'): """Instantiate new batch_transform object. @@ -350,6 +351,15 @@ class BatchTransform(object): self.clean_nans = clean_nans self.compute_only_full = compute_only_full + # How many bars are in a day + self.bars = bars + if self.bars == 'daily': + self.bars_in_day = 1 + elif self.bars == 'minute': + self.bars_in_day = int(6.5 * 60) + else: + raise ValueError('%s bars not understood.' % self.bars) + # The following logic is to allow pre-specified sid filters # to operate on the data, but to also allow new symbols to # enter the batch transform's window IFF a sid filter is not @@ -434,7 +444,8 @@ class BatchTransform(object): # Create rolling panel if not existant if self.rolling_panel is None: - self.rolling_panel = RollingPanel(self.window_length, + self.rolling_panel = RollingPanel(self.window_length * + self.bars_in_day, self.field_names, sids) # Store event in rolling frame diff --git a/zipline/utils/factory.py b/zipline/utils/factory.py index 898376b4..c9591012 100644 --- a/zipline/utils/factory.py +++ b/zipline/utils/factory.py @@ -288,7 +288,13 @@ def create_trade_source(sids, trade_count, return source -def create_test_df_source(sim_params=None): +def create_test_df_source(sim_params=None, bars='daily'): + if bars == 'daily': + freq = pd.datetools.BDay() + elif bars == 'minute': + freq = pd.datetools.Minute() + else: + raise ValueError('%s bars not understood.' % freq) if sim_params: index = sim_params.trading_days @@ -298,9 +304,19 @@ def create_test_df_source(sim_params=None): index = pd.DatetimeIndex( start=start, end=end, - freq=pd.datetools.BDay() + freq=freq ) + if bars == 'minute': + new_index = [] + for i in index: + market_open = i.replace(hour=14, + minute=31) + market_close = i.replace(hour=21, + minute=0) + if i >= market_open and i <= market_close: + new_index.append(i) + index = new_index x = np.arange(0, len(index)) df = pd.DataFrame(x, index=index, columns=[0])