diff --git a/tests/pipeline/test_adjusted_array.py b/tests/pipeline/test_adjusted_array.py index 4f94f27d..858db648 100644 --- a/tests/pipeline/test_adjusted_array.py +++ b/tests/pipeline/test_adjusted_array.py @@ -95,7 +95,7 @@ def _gen_multiplicative_adjustment_cases(dtype): # Note that row indices are inclusive! adjustments[1] = [ - adjustment_type(0, 0, 0, dtype(2)), + adjustment_type(0, 0, 0, 0, dtype(2)), ] buffer_as_of[1] = array([[2, 1, 1], [1, 1, 1], @@ -108,8 +108,8 @@ def _gen_multiplicative_adjustment_cases(dtype): buffer_as_of[2] = buffer_as_of[1] adjustments[3] = [ - adjustment_type(1, 2, 1, dtype(3)), - adjustment_type(0, 1, 0, dtype(4)), + adjustment_type(1, 2, 1, 1, dtype(3)), + adjustment_type(0, 1, 0, 0, dtype(4)), ] buffer_as_of[3] = array([[8, 1, 1], [4, 3, 1], @@ -119,7 +119,7 @@ def _gen_multiplicative_adjustment_cases(dtype): [1, 1, 1]], dtype=dtype) adjustments[4] = [ - adjustment_type(0, 3, 2, dtype(5)) + adjustment_type(0, 3, 2, 2, dtype(5)) ] buffer_as_of[4] = array([[8, 1, 5], [4, 3, 5], @@ -129,8 +129,8 @@ def _gen_multiplicative_adjustment_cases(dtype): [1, 1, 1]], dtype=dtype) adjustments[5] = [ - adjustment_type(0, 4, 1, dtype(6)), - adjustment_type(2, 2, 2, dtype(7)), + adjustment_type(0, 4, 1, 1, dtype(6)), + adjustment_type(2, 2, 2, 2, dtype(7)), ] buffer_as_of[5] = array([[8, 6, 5], [4, 18, 5], @@ -162,7 +162,7 @@ def _gen_overwrite_adjustment_cases(dtype): # Note that row indices are inclusive! adjustments[1] = [ - adjustment_type(0, 0, 0, dtype(1)), + adjustment_type(0, 0, 0, 0, dtype(1)), ] buffer_as_of[1] = array([[1, 2, 2], [2, 2, 2], @@ -175,8 +175,8 @@ def _gen_overwrite_adjustment_cases(dtype): buffer_as_of[2] = buffer_as_of[1] adjustments[3] = [ - adjustment_type(1, 2, 1, dtype(3)), - adjustment_type(0, 1, 0, dtype(4)), + adjustment_type(1, 2, 1, 1, dtype(3)), + adjustment_type(0, 1, 0, 0, dtype(4)), ] buffer_as_of[3] = array([[4, 2, 2], [4, 3, 2], @@ -186,7 +186,7 @@ def _gen_overwrite_adjustment_cases(dtype): [2, 2, 2]], dtype=dtype) adjustments[4] = [ - adjustment_type(0, 3, 2, dtype(5)) + adjustment_type(0, 3, 2, 2, dtype(5)) ] buffer_as_of[4] = array([[4, 2, 5], [4, 3, 5], @@ -196,8 +196,8 @@ def _gen_overwrite_adjustment_cases(dtype): [2, 2, 2]], dtype=dtype) adjustments[5] = [ - adjustment_type(0, 4, 1, dtype(6)), - adjustment_type(2, 2, 2, dtype(7)), + adjustment_type(0, 4, 1, 1, dtype(6)), + adjustment_type(2, 2, 2, 2, dtype(7)), ] buffer_as_of[5] = array([[4, 6, 5], [4, 6, 5], diff --git a/tests/pipeline/test_frameload.py b/tests/pipeline/test_frameload.py index 622e6a0c..5b5b20c9 100644 --- a/tests/pipeline/test_frameload.py +++ b/tests/pipeline/test_frameload.py @@ -189,14 +189,38 @@ class DataFrameLoaderTestCase(TestCase): ) expected_formatted_adjustments = { 6: [ - Float64Multiply(first_row=0, last_row=5, col=0, value=0.5), - Float64Add(first_row=0, last_row=5, col=1, value=1.0), + Float64Multiply( + first_row=0, + last_row=5, + first_col=0, + last_col=0, + value=0.5, + ), + Float64Add( + first_row=0, + last_row=5, + first_col=1, + last_col=1, + value=1.0, + ), ], 7: [ - Float64Add(first_row=5, last_row=6, col=1, value=1.0), + Float64Add( + first_row=5, + last_row=6, + first_col=1, + last_col=1, + value=1.0, + ), ], 8: [ - Float64Overwrite(first_row=6, last_row=7, col=2, value=99.0) + Float64Overwrite( + first_row=6, + last_row=7, + first_col=2, + last_col=2, + value=99.0, + ) ], } self.assertEqual(formatted_adjustments, expected_formatted_adjustments) diff --git a/tests/pipeline/test_us_equity_pricing_loader.py b/tests/pipeline/test_us_equity_pricing_loader.py index b75f4799..0c687553 100644 --- a/tests/pipeline/test_us_equity_pricing_loader.py +++ b/tests/pipeline/test_us_equity_pricing_loader.py @@ -337,7 +337,8 @@ class USEquityPricingLoaderTestCase(TestCase): Float64Multiply( first_row=0, last_row=delta, - col=sid - 1, + first_col=sid - 1, + last_col=sid - 1, value=ratio, ) ) @@ -347,7 +348,8 @@ class USEquityPricingLoaderTestCase(TestCase): Float64Multiply( first_row=0, last_row=delta, - col=sid - 1, + first_col=sid - 1, + last_col=sid - 1, value=1.0 / ratio, ) ) diff --git a/zipline/data/_adjustments.pyx b/zipline/data/_adjustments.pyx index 8f923ee0..b53474f2 100644 --- a/zipline/data/_adjustments.pyx +++ b/zipline/data/_adjustments.pyx @@ -225,7 +225,7 @@ cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection asset_ixs[sid] = assets.get_loc(sid) asset_ix = asset_ixs[sid] - price_adj = Float64Multiply(0, date_loc, asset_ix, ratio) + price_adj = Float64Multiply(0, date_loc, asset_ix, asset_ix, ratio) for i, column in enumerate(columns): col_adjustments = results[i] if column != 'volume': @@ -235,7 +235,7 @@ cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection col_adjustments[date_loc] = [price_adj] else: volume_adj = Float64Multiply( - 0, date_loc, asset_ix, 1.0 / ratio + 0, date_loc, asset_ix, asset_ix, 1.0 / ratio ) try: col_adjustments[date_loc].append(volume_adj) @@ -253,7 +253,7 @@ cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection asset_ixs[sid] = assets.get_loc(sid) asset_ix = asset_ixs[sid] - adj = Float64Multiply(0, date_loc, asset_ix, ratio) + adj = Float64Multiply(0, date_loc, asset_ix, asset_ix, ratio) for i, column in enumerate(columns): col_adjustments = results[i] if column != 'volume': @@ -273,7 +273,7 @@ cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection asset_ixs[sid] = assets.get_loc(sid) asset_ix = asset_ixs[sid] - adj = Float64Multiply(0, date_loc, asset_ix, ratio) + adj = Float64Multiply(0, date_loc, asset_ix, asset_ix, ratio) for i, column in enumerate(columns): col_adjustments = results[i] if column != 'volume': diff --git a/zipline/lib/adjustment.pyx b/zipline/lib/adjustment.pyx index 245e1616..2b64642c 100644 --- a/zipline/lib/adjustment.pyx +++ b/zipline/lib/adjustment.pyx @@ -87,7 +87,7 @@ cpdef _from_assets_and_dates(cls, end_date, asset_id, ) - return cls(first_row, last_row, col, value) + return cls(first_row, last_row, col, col, value) cdef class Float64Adjustment: @@ -95,30 +95,37 @@ cdef class Float64Adjustment: Base class for adjustments that operate on Float64 buffers. """ cdef: - readonly Py_ssize_t col, first_row, last_row + readonly Py_ssize_t first_col, last_col, first_row, last_row readonly float64_t value def __cinit__(self, Py_ssize_t first_row, Py_ssize_t last_row, - Py_ssize_t col, + Py_ssize_t first_col, + Py_ssize_t last_col, object value): assert 0 <= first_row <= last_row + assert 0 <= first_col <= last_col self.first_row = first_row self.last_row = last_row - self.col = col + self.first_col = first_col + self.last_col = last_col self.value = float(value) from_assets_and_dates = classmethod(_from_assets_and_dates) def __repr__(self): - return "%s(first_row=%d, last_row=%d, col=%d, value=%f)" % ( - type(self).__name__, - self.first_row, - self.last_row, - self.col, - self.value, + return ( + "%s(first_row=%d, last_row=%d," + " first_col=%d, last_col=%d, value=%f)" % ( + type(self).__name__, + self.first_row, + self.last_row, + self.first_col, + self.last_col, + self.value, + ) ) def __richcmp__(self, object other, int op): @@ -128,9 +135,15 @@ cdef class Float64Adjustment: if op != Py_EQ or type(self) != type(other): return NotImplemented + return self._key() == other._key() + + cpdef _key(self): return ( - (self.first_row, self.last_row, self.col, self.value) == \ - (other.first_row, other.last_row, other.col, other.value) + self.first_row, + self.last_row, + self.first_col, + self.last_col, + self.value, ) @@ -148,21 +161,28 @@ cdef class Float64Multiply(Float64Adjustment): [ 3., 4., 5.], [ 6., 7., 8.]]) - >>> adj = Float64Multiply(first_row=1, last_row=2, col=1, value=4.0) + >>> adj = Float64Multiply( + ... first_row=1, + ... last_row=2, + ... first_col=1, + ... last_col=2, + ... value=4.0, + ... ) >>> adj.mutate(arr) >>> arr array([[ 0., 1., 2.], - [ 3., 16., 5.], - [ 6., 28., 8.]]) + [ 3., 16., 20.], + [ 6., 28., 32.]]) """ cpdef mutate(self, float64_t[:, :] data): cdef Py_ssize_t row, col - col = self.col - # last_row + 1 because last_row should also be affected. - for row in range(self.first_row, self.last_row + 1): - data[row, col] *= self.value + # last_col + 1 because last_col should also be affected. + for col in range(self.first_col, self.last_col + 1): + # last_row + 1 because last_row should also be affected. + for row in range(self.first_row, self.last_row + 1): + data[row, col] *= self.value cdef class Float64Overwrite(Float64Adjustment): @@ -179,21 +199,28 @@ cdef class Float64Overwrite(Float64Adjustment): [ 3., 4., 5.], [ 6., 7., 8.]]) - >>> adj = Float64Overwrite(first_row=1, last_row=2, col=1, value=0.0) + >>> adj = Float64Overwrite( + ... first_row=1, + ... last_row=2, + ... first_col=1, + ... last_col=2, + ... value=0.0, + ... ) >>> adj.mutate(arr) >>> arr array([[ 0., 1., 2.], - [ 3., 0., 5.], - [ 6., 0., 8.]]) + [ 3., 0., 0.], + [ 6., 0., 0.]]) """ cpdef mutate(self, float64_t[:, :] data): cdef Py_ssize_t row, col - col = self.col - # last_row + 1 because last_row should also be affected. - for row in range(self.first_row, self.last_row + 1): - data[row, col] = self.value + # last_col + 1 because last_col should also be affected. + for col in range(self.first_col, self.last_col + 1): + # last_row + 1 because last_row should also be affected. + for row in range(self.first_row, self.last_row + 1): + data[row, col] = self.value cdef class Float64Add(Float64Adjustment): @@ -210,18 +237,25 @@ cdef class Float64Add(Float64Adjustment): [ 3., 4., 5.], [ 6., 7., 8.]]) - >>> adj = Float64Add(first_row=1, last_row=2, col=1, value=1.0) + >>> adj = Float64Add( + ... first_row=1, + ... last_row=2, + ... first_col=1, + ... last_col=2, + ... value=1.0, + ... ) >>> adj.mutate(arr) >>> arr array([[ 0., 1., 2.], - [ 3., 5., 5.], - [ 6., 8., 8.]]) + [ 3., 5., 6.], + [ 6., 8., 9.]]) """ cpdef mutate(self, float64_t[:, :] data): cdef Py_ssize_t row, col - col = self.col - # last_row + 1 because last_row should also be affected. - for row in range(self.first_row, self.last_row + 1): - data[row, col] += self.value + # last_col + 1 because last_col should also be affected. + for col in range(self.first_col, self.last_col + 1): + # last_row + 1 because last_row should also be affected. + for row in range(self.first_row, self.last_row + 1): + data[row, col] += self.value