diff --git a/tests/pipeline/test_adjusted_array.py b/tests/pipeline/test_adjusted_array.py index 0327bdd0..5f4b4def 100644 --- a/tests/pipeline/test_adjusted_array.py +++ b/tests/pipeline/test_adjusted_array.py @@ -1,7 +1,8 @@ """ Tests for chunked adjustments. """ -from itertools import chain +from collections import namedtuple +from itertools import chain, product from textwrap import dedent from unittest import TestCase @@ -95,6 +96,20 @@ bytes_dtype = dtype('S3') unicode_dtype = dtype('U3') +AdjustmentCase = namedtuple( + 'AdjustmentCase', + [ + 'name', + 'baseline', + 'window_length', + 'adjustments', + 'missing_value', + 'perspective_offset', + 'expected_result', + ] +) + + def _gen_unadjusted_cases(name, make_input, make_expected_output, @@ -112,13 +127,14 @@ def _gen_unadjusted_cases(name, windowlen, nrows ) - yield ( - "%s_length_%d" % (name, windowlen), - input_array, - windowlen, - {}, - missing_value, - [ + yield AdjustmentCase( + name="%s_length_%d" % (name, windowlen), + baseline=input_array, + window_length=windowlen, + adjustments={}, + missing_value=missing_value, + perspective_offset=0, + expected_result=[ expected_output_array[offset:offset + windowlen] for offset in range(num_legal_windows) ], @@ -199,85 +215,7 @@ def _gen_multiplicative_adjustment_cases(dtype): adjustments, buffer_as_of, nrows, - ) - - -def _gen_multiplicative_adjustment_cases_with_perpsective_offset(dtype): - """ - Generate expected moving windows on a buffer with adjustments. - - We proceed by constructing, at each row, the view of the array we expect in - in all windows anchored on that row. - - In general, if we have an adjustment to be applied once we process the row - at index N, should see that adjustment applied to the underlying buffer for - any window containing the row at index N - 1. - - We then build all legal windows over these buffers. - """ - adjustment_type = { - float64_dtype: Float64Multiply, - }[dtype] - - nrows, ncols = 6, 3 - adjustments = {} - buffer_as_of = [None] * 6 - baseline = full((nrows, ncols), 1, dtype=dtype) - - # Note that row indices are inclusive! - adjustments[1] = [ - adjustment_type(0, 0, 0, 0, coerce_to_dtype(dtype, 2)), - ] - buffer_as_of[0] = array([[2, 1, 1], - [1, 1, 1], - [1, 1, 1], - [1, 1, 1], - [1, 1, 1], - [1, 1, 1]], dtype=dtype) - - # No adjustment at index 2. - buffer_as_of[1] = buffer_as_of[0] - - adjustments[3] = [ - adjustment_type(1, 2, 1, 1, coerce_to_dtype(dtype, 3)), - adjustment_type(0, 1, 0, 0, coerce_to_dtype(dtype, 4)), - ] - buffer_as_of[2] = array([[8, 1, 1], - [4, 3, 1], - [1, 3, 1], - [1, 1, 1], - [1, 1, 1], - [1, 1, 1]], dtype=dtype) - - adjustments[4] = [ - adjustment_type(0, 3, 2, 2, coerce_to_dtype(dtype, 5)) - ] - buffer_as_of[3] = array([[8, 1, 5], - [4, 3, 5], - [1, 3, 5], - [1, 1, 5], - [1, 1, 1], - [1, 1, 1]], dtype=dtype) - - adjustments[5] = [ - adjustment_type(0, 4, 1, 1, coerce_to_dtype(dtype, 6)), - adjustment_type(2, 2, 2, 2, coerce_to_dtype(dtype, 7)), - ] - buffer_as_of[4] = array([[8, 6, 5], - [4, 18, 5], - [1, 18, 35], - [1, 6, 5], - [1, 6, 1], - [1, 1, 1]], dtype=dtype) - - buffer_as_of[5] = buffer_as_of[4] - - return _gen_expectations( - baseline, - default_missing_value_for_dtype(dtype), - adjustments, - buffer_as_of, - nrows, + perspective_offsets=(0, 1, 2, 1000), ) @@ -380,6 +318,7 @@ def _gen_overwrite_adjustment_cases(dtype): adjustments, buffer_as_of, nrows=6, + perspective_offsets=(0, 1, 2, 1000), ) @@ -476,13 +415,13 @@ def _gen_overwrite_1d_array_adjustment_case(dtype): [2, 4, 5], [2, 5, 2], [2, 2, 2]]) - return _gen_expectations( baseline, missing_value, adjustments, buffer_as_of, nrows=6, + perspective_offsets=(0, 1, 2, 1000), ) @@ -490,30 +429,63 @@ def _gen_expectations(baseline, missing_value, adjustments, buffer_as_of, - nrows): - - for windowlen in valid_window_lengths(nrows): + nrows, + perspective_offsets): + for windowlen, perspective_offset in product(valid_window_lengths(nrows), + perspective_offsets): + # How long an an iterator of length-N windows on this buffer? + # For example, for a window of length 3 on a buffer of length 6, there + # are four valid windows. num_legal_windows = num_windows_of_length_M_on_buffers_of_length_N( windowlen, nrows ) - yield ( - "dtype_%s_length_%d" % (baseline.dtype, windowlen), - baseline, + # Build the sequence of regions in the underlying buffer we expect to + # see. For example, with a window length of 3 on a buffer of length 6, + # we expect to see: + # (buffer[0:3], buffer[1:4], buffer[2:5], buffer[3:6]) + # + slices = [slice(i, i + windowlen) for i in range(num_legal_windows)] + + # The sequence of perspectives we expect to take on the underlying + # data. For example, with a window length of 3 and a perspective offset + # of 1, we expect to see: + # (buffer_as_of[3], buffer_as_of[4], buffer_as_of[5], buffer_as_of[5]) + # + initial_perspective = windowlen + perspective_offset - 1 + perspectives = range( + initial_perspective, + initial_perspective + num_legal_windows + ) + + def as_of(p): + # perspective_offset can push us past the end of the underlying + # buffer/adjustments. When it does, we should always see the latest + # version of the buffer. + if p >= len(buffer_as_of): + return buffer_as_of[-1] + return buffer_as_of[p] + + expected_iterator_results = [ + as_of(perspective)[slice_] + for slice_, perspective in zip(slices, perspectives) + ] + + test_name = "dtype_{}_length_{}_perpective_offset_{}".format( + baseline.dtype, windowlen, - adjustments, - missing_value, - [ - # This is a nasty expression... - # - # Reading from right to left: we want a slice of length - # 'windowlen', starting at 'offset', from the buffer on which - # we've applied all adjustments corresponding to the last row - # of the data, which will be (offset + windowlen - 1). - buffer_as_of[offset + windowlen - 1][offset:offset + windowlen] - for offset in range(num_legal_windows) - ], + perspective_offset, + ) + + yield AdjustmentCase( + name=test_name, + baseline=baseline, + window_length=windowlen, + adjustments=adjustments, + missing_value=missing_value, + perspective_offset=perspective_offset, + expected_result=expected_iterator_results ) @@ -583,6 +555,7 @@ class AdjustedArrayTestCase(TestCase): lookback, adjustments, missing_value, + perspective_offset, expected_output): array = AdjustedArray(data, NOMASK, adjustments, missing_value) @@ -598,31 +571,18 @@ class AdjustedArrayTestCase(TestCase): lookback, adjustments, missing_value, + perspective_offset, expected): array = AdjustedArray(data, NOMASK, adjustments, missing_value) for _ in range(2): # Iterate 2x ensure adjusted_arrays are re-usable. - window_iter = array.traverse(lookback) + window_iter = array.traverse( + lookback, + perspective_offset=perspective_offset, + ) for yielded, expected_yield in zip_longest(window_iter, expected): check_arrays(yielded, expected_yield) - @parameterized.expand( - _gen_multiplicative_adjustment_cases_with_perpsective_offset( - float64_dtype)) - def test_multiplicative_adjustments_with_perspective_offset(self, - name, - data, - lookback, - adjustments, - missing_value, - expected): - array = AdjustedArray(data, NOMASK, adjustments, missing_value, 1) - for _ in range(2): # Iterate 2x ensure adjusted_arrays are re-usable. - window_iter = array.traverse(lookback) - for yielded, expected_yield in zip_longest(window_iter, expected): - print yielded - check_arrays(yielded, expected_yield) - @parameterized.expand( chain( _gen_overwrite_adjustment_cases(float64_dtype), @@ -680,14 +640,19 @@ class AdjustedArrayTestCase(TestCase): ) def test_overwrite_adjustment_cases(self, name, - data, + baseline, lookback, adjustments, missing_value, + perspective_offset, expected): - array = AdjustedArray(data, NOMASK, adjustments, missing_value) + array = AdjustedArray(baseline, NOMASK, adjustments, missing_value) + for _ in range(2): # Iterate 2x ensure adjusted_arrays are re-usable. - window_iter = array.traverse(lookback) + window_iter = array.traverse( + lookback, + perspective_offset=perspective_offset, + ) for yielded, expected_yield in zip_longest(window_iter, expected): check_arrays(yielded, expected_yield) diff --git a/zipline/lib/adjusted_array.py b/zipline/lib/adjusted_array.py index cc4d6407..b88b1cbb 100644 --- a/zipline/lib/adjusted_array.py +++ b/zipline/lib/adjusted_array.py @@ -152,22 +152,16 @@ class AdjustedArray(object): missing_value : object A value to use to fill missing data in yielded windows. Should be a value coercible to `data.dtype`. - perspective_offset : int - The number of rows after the current end of the window, from which the - data is being viewed. This value is used so that adjustments that occur - between the end of the window and the vantage point are applied. """ __slots__ = ( '_data', '_view_kwargs', 'adjustments', 'missing_value', - 'perspective_offset', '__weakref__', ) - def __init__(self, data, mask, adjustments, missing_value, - perspective_offset=0): + def __init__(self, data, mask, adjustments, missing_value): self._data, self._view_kwargs = _normalize_array(data, missing_value) self.adjustments = adjustments @@ -183,8 +177,6 @@ class AdjustedArray(object): ) self._data[~mask] = self.missing_value - self.perspective_offset = perspective_offset - @lazyval def data(self): """ @@ -208,7 +200,10 @@ class AdjustedArray(object): return LabelWindow return CONCRETE_WINDOW_TYPES[self._data.dtype] - def traverse(self, window_length, offset=0): + def traverse(self, + window_length, + offset=0, + perspective_offset=0): """ Produce an iterator rolling windows rows over our data. Each emitted window will have `window_length` rows. @@ -218,7 +213,10 @@ class AdjustedArray(object): window_length : int The number of rows in each emitted window. offset : int, optional - Number of rows to skip before the first window. + Number of rows to skip before the first window. Default is 0. + perspective_offset : int, optional + Number of rows past the end of the current window from which to + "view" the underlying data. """ data = self._data.copy() _check_window_params(data, window_length) @@ -228,7 +226,7 @@ class AdjustedArray(object): self.adjustments, offset, window_length, - self.perspective_offset + perspective_offset, ) def inspect(self):