diff --git a/tests/pipeline/test_downsampling.py b/tests/pipeline/test_downsampling.py index 47f35eee..e57beed6 100644 --- a/tests/pipeline/test_downsampling.py +++ b/tests/pipeline/test_downsampling.py @@ -12,11 +12,13 @@ from zipline.pipeline import ( ) from zipline.pipeline.data.testing import TestingDataSet from zipline.pipeline.factors import SimpleMovingAverage +from zipline.pipeline.filters.smoothing import All from zipline.testing import ZiplineTestCase, parameter_space from zipline.testing.fixtures import ( WithTradingSessions, WithSeededRandomPipelineEngine, ) +from zipline.utils.numpy_utils import int64_dtype class NDaysAgoFactor(CustomFactor): @@ -552,12 +554,9 @@ class DownsampledPipelineTestCase(WithSeededRandomPipelineEngine, # Extend into the first few days of 2015 to test year/quarter boundaries. END_DATE = pd.Timestamp('2015-01-06', tz='UTC') - def test_downsample_windowed_factor(self): + ASSET_FINDER_EQUITY_SIDS = tuple(range(10)) - f = SimpleMovingAverage( - inputs=[TestingDataSet.float_col], - window_length=5, - ) + def check_downsampled_term(self, term): # June 2014 # Mo Tu We Th Fr Sa Su @@ -574,34 +573,34 @@ class DownsampledPipelineTestCase(WithSeededRandomPipelineEngine, start_date, end_date = compute_dates[[0, -1]] pipe = Pipeline({ - 'year': f.downsample(frequency='Y'), - 'quarter': f.downsample(frequency='Q'), - 'month': f.downsample(frequency='M'), - 'week': f.downsample(frequency='W'), + 'year': term.downsample(frequency='Y'), + 'quarter': term.downsample(frequency='Q'), + 'month': term.downsample(frequency='M'), + 'week': term.downsample(frequency='W'), }) - # Raw values for f, computed each day from 2014 to the end of the + # Raw values for term, computed each day from 2014 to the end of the # target period. - raw_f_results = self.run_pipeline( - Pipeline({'f': f}), + raw_term_results = self.run_pipeline( + Pipeline({'term': term}), start_date=pd.Timestamp('2014-01-02', tz='UTC'), end_date=pd.Timestamp('2015-01-06', tz='UTC'), - )['f'].unstack() + )['term'].unstack() expected_results = { - 'year': (raw_f_results + 'year': (raw_term_results .groupby(pd.TimeGrouper('AS')) .first() .reindex(compute_dates, method='ffill')), - 'quarter': (raw_f_results + 'quarter': (raw_term_results .groupby(pd.TimeGrouper('QS')) .first() .reindex(compute_dates, method='ffill')), - 'month': (raw_f_results + 'month': (raw_term_results .groupby(pd.TimeGrouper('MS')) .first() .reindex(compute_dates, method='ffill')), - 'week': (raw_f_results + 'week': (raw_term_results .groupby(pd.TimeGrouper('W', label='left')) .first() .reindex(compute_dates, method='ffill')), @@ -613,3 +612,53 @@ class DownsampledPipelineTestCase(WithSeededRandomPipelineEngine, result = results[frequency].unstack() expected = expected_results[frequency] assert_frame_equal(result, expected) + + def test_downsample_windowed_factor(self): + self.check_downsampled_term( + SimpleMovingAverage( + inputs=[TestingDataSet.float_col], + window_length=5, + ) + ) + + def test_downsample_non_windowed_factor(self): + sma = SimpleMovingAverage( + inputs=[TestingDataSet.float_col], + window_length=5, + ) + + self.check_downsampled_term(((sma + sma) / 2).rank()) + + def test_downsample_windowed_filter(self): + sma = SimpleMovingAverage( + inputs=[TestingDataSet.float_col], + window_length=5, + ) + self.check_downsampled_term(All(inputs=[sma.top(4)], window_length=5)) + + def test_downsample_nonwindowed_filter(self): + sma = SimpleMovingAverage( + inputs=[TestingDataSet.float_col], + window_length=5, + ) + self.check_downsampled_term(sma > 5) + + def test_downsample_windowed_classifier(self): + + class IntSumClassifier(CustomClassifier): + inputs = [TestingDataSet.float_col] + window_length = 8 + dtype = int64_dtype + missing_value = -1 + + def compute(self, today, assets, out, floats): + out[:] = floats.sum(axis=0).astype(int) % 4 + + self.check_downsampled_term(IntSumClassifier()) + + def test_downsample_nonwindowed_classifier(self): + sma = SimpleMovingAverage( + inputs=[TestingDataSet.float_col], + window_length=5, + ) + self.check_downsampled_term(sma.quantiles(5)) diff --git a/zipline/pipeline/mixins.py b/zipline/pipeline/mixins.py index 73b6d111..9fa0ad6a 100644 --- a/zipline/pipeline/mixins.py +++ b/zipline/pipeline/mixins.py @@ -388,7 +388,7 @@ class DownsampledMixin(StandardOutputs): return min_extra_rows + (current_start_pos - new_start_pos) - def _compute(self, windows, dates, assets, mask): + def _compute(self, inputs, dates, assets, mask): """ Compute by delegating to self._wrapped_term._compute on sample dates. @@ -400,6 +400,27 @@ class DownsampledMixin(StandardOutputs): real_compute = self._wrapped_term._compute + if self.windowed: + # If we're windowed, inputs are stateful AdjustedArrays. We don't + # need to do any preparation before forwarding to real_compute, but + # we need to call `next` on them if we want to skip an iteration. + def prepare_inputs(): + return inputs + + def skip_this_input(): + for w in inputs: + next(w) + else: + # If we're not windowed, inputs are just ndarrays. We need to + # slice off one row when forwarding to real_compute, but we don't + # need to do anything to skip an input. + def prepare_inputs(): + # i is the loop iteration variable below. + return [a[[i]] for a in inputs] + + def skip_this_input(): + pass + results = [] samples = iter(to_sample) next_sample = next(samples) @@ -407,7 +428,7 @@ class DownsampledMixin(StandardOutputs): if next_sample == compute_date: results.append( real_compute( - windows, + prepare_inputs(), dates[i:i + 1], assets, mask[i:i + 1], @@ -420,13 +441,10 @@ class DownsampledMixin(StandardOutputs): # compares False with any other datetime. next_sample = pd_NaT else: + skip_this_input() # Copy results from previous sample period. results.append(results[-1]) - # Force adjusted arrays forward one tick. - for w in windows: - next(w) - # We should have exhausted our sample dates. try: next_sample = next(samples)