Merge pull request #2420 from quantopian/adjusted-array-map

ENH: Add AdjustedArray.map_labels

Author: Andrew Daniels

tests/pipeline/test_adjusted_array.py

@@ -707,3 +707,41 @@ def test_inspect(self):
         )
         got = adj_array.inspect()
         self.assertEqual(expected, got)
+
+    def test_update_labels(self):
+        data = array([
+            ['aaa', 'bbb', 'ccc'],
+            ['ddd', 'eee', 'fff'],
+            ['ggg', 'hhh', 'iii'],
+            ['jjj', 'kkk', 'lll'],
+            ['mmm', 'nnn', 'ooo'],
+        ])
+        label_array = LabelArray(data, missing_value='')
+
+        adj_array = AdjustedArray(
+            data=label_array,
+            adjustments={4: [ObjectOverwrite(2, 3, 0, 0, 'ppp')]},
+            missing_value='',
+        )
+
+        expected_data = array([
+            ['aaa-foo', 'bbb-foo', 'ccc-foo'],
+            ['ddd-foo', 'eee-foo', 'fff-foo'],
+            ['ggg-foo', 'hhh-foo', 'iii-foo'],
+            ['jjj-foo', 'kkk-foo', 'lll-foo'],
+            ['mmm-foo', 'nnn-foo', 'ooo-foo'],
+        ])
+        expected_label_array = LabelArray(expected_data, missing_value='')
+
+        expected_adj_array = AdjustedArray(
+            data=expected_label_array,
+            adjustments={4: [ObjectOverwrite(2, 3, 0, 0, 'ppp-foo')]},
+            missing_value='',
+        )
+
+        adj_array.update_labels(lambda x: x + '-foo')
+
+        # Check that the mapped AdjustedArray has the expected baseline
+        # values and adjustment values.
+        check_arrays(adj_array.data, expected_adj_array.data)
+        self.assertEqual(adj_array.adjustments, expected_adj_array.adjustments)

zipline/lib/adjusted_array.py

@@ -13,6 +13,7 @@
     uint32,
     uint8,
 )
+from six import iteritems
 from zipline.errors import (
     WindowLengthNotPositive,
     WindowLengthTooLong,
@@ -164,7 +165,7 @@ def __init__(self, data, adjustments, missing_value):
         self.adjustments = adjustments
         self.missing_value = missing_value
 
-    @lazyval
+    @property
     def data(self):
         """
         The data stored in this array.
@@ -237,6 +238,25 @@ def inspect(self):
             adjustments=self.adjustments,
         )
 
+    def update_labels(self, func):
+        """
+        Map a function over baseline and adjustment values in place.
+
+        Note that the baseline data values must be a LabelArray.
+        """
+        if not isinstance(self.data, LabelArray):
+            raise TypeError(
+                'update_labels only supported if data is of type LabelArray.'
+            )
+
+        # Map the baseline values.
+        self._data = self._data.map(func)
+
+        # Map each of the adjustments.
+        for _, row_adjustments in iteritems(self.adjustments):
+            for adjustment in row_adjustments:
+                adjustment.value = func(adjustment.value)
+
 
 def ensure_adjusted_array(ndarray_or_adjusted_array, missing_value):
     if isinstance(ndarray_or_adjusted_array, AdjustedArray):

zipline/lib/adjustment.pxd

@@ -42,7 +42,7 @@ cdef class Float64Adjustment(Adjustment):
     """
     Base class for adjustments that operate on Float64 data.
     """
-    cdef readonly np.float64_t value
+    cdef public np.float64_t value
 
 
 cdef class Float64Multiply(Float64Adjustment):
@@ -146,7 +146,7 @@ cdef class Float641DArrayOverwrite(ArrayAdjustment):
            [ 4.,  16.,  17.,  18.,  19.],
            [ 20.,  21.,  22.,  23.,  24.]])
     """
-    cdef readonly np.float64_t[:] values
+    cdef public np.float64_t[:] values
     cpdef mutate(self, np.float64_t[:, :] data)
 
 
@@ -184,7 +184,7 @@ cdef class Datetime641DArrayOverwrite(ArrayAdjustment):
        [False, False, False],
        [False,  True,  True]], dtype=bool)
     """
-    cdef readonly np.int64_t[:] values
+    cdef public np.int64_t[:] values
     cpdef mutate(self, np.int64_t[:, :] data)
 
 
@@ -226,7 +226,7 @@ cdef class _Int64Adjustment(Adjustment):
     This is private because we never actually operate on integers as data, but
     we use integer arrays to represent datetime and timedelta data.
     """
-    cdef readonly np.int64_t value
+    cdef public np.int64_t value
 
 
 cdef class Int64Overwrite(_Int64Adjustment):
@@ -312,7 +312,7 @@ cdef class _ObjectAdjustment(Adjustment):
     We use only this for categorical data, where our data buffer is an array of
     indices into an array of unique Python string objects.
     """
-    cdef readonly object value
+    cdef public object value
 
 
 cdef class ObjectOverwrite(_ObjectAdjustment):
@@ -330,7 +330,7 @@ cdef class BooleanAdjustment(Adjustment):
     instead we work with uint8 values everywhere, and we do validation/coercion
     at API boundaries.
     """
-    cdef readonly np.uint8_t value
+    cdef public np.uint8_t value
 
 
 cdef class BooleanOverwrite(BooleanAdjustment):