diff --git a/tests/pipeline/test_factor.py b/tests/pipeline/test_factor.py index c5873c1f..4eee9f2e 100644 --- a/tests/pipeline/test_factor.py +++ b/tests/pipeline/test_factor.py @@ -23,6 +23,7 @@ from numpy import ( from numpy.random import randn, seed from zipline.errors import UnknownRankMethod +from zipline.lib.labelarray import LabelArray from zipline.lib.rank import masked_rankdata_2d from zipline.lib.normalize import naive_grouped_rowwise_apply as grouped_apply from zipline.pipeline import Classifier, Factor, Filter, TermGraph @@ -38,6 +39,7 @@ from zipline.testing import ( ) from zipline.utils.functional import dzip_exact from zipline.utils.numpy_utils import ( + categorical_dtype, datetime64ns_dtype, float64_dtype, int64_dtype, @@ -442,6 +444,7 @@ class FactorTestCase(BasePipelineTestCase): f = self.f m = Mask() c = C() + str_c = C(dtype=categorical_dtype) factor_data = array( [[1.0, 2.0, 3.0, 4.0], @@ -463,12 +466,18 @@ class FactorTestCase(BasePipelineTestCase): [1, 1, 2, 2]], dtype=int64_dtype, ) + string_classifier_data = LabelArray( + classifier_data.astype(str).astype(object), + missing_value=None, + ) terms = { 'vanilla': f.demean(), 'masked': f.demean(mask=m), 'grouped': f.demean(groupby=c), + 'grouped_str': f.demean(groupby=str_c), 'grouped_masked': f.demean(mask=m, groupby=c), + 'grouped_masked_str': f.demean(mask=m, groupby=str_c), } expected = { 'vanilla': array( @@ -496,6 +505,9 @@ class FactorTestCase(BasePipelineTestCase): [-0.500, 0.500, 0.000, nan]] ) } + # Changing the classifier dtype shouldn't affect anything. + expected['grouped_str'] = expected['grouped'] + expected['grouped_masked_str'] = expected['grouped_masked'] graph = TermGraph(terms) results = self.run_graph( @@ -503,6 +515,7 @@ class FactorTestCase(BasePipelineTestCase): initial_workspace={ f: factor_data, c: classifier_data, + str_c: string_classifier_data, m: filter_data, }, mask=self.build_mask(self.ones_mask(shape=factor_data.shape)), diff --git a/zipline/lib/labelarray.py b/zipline/lib/labelarray.py index b707f0cd..a3ef5027 100644 --- a/zipline/lib/labelarray.py +++ b/zipline/lib/labelarray.py @@ -211,6 +211,10 @@ class LabelArray(ndarray): # This is a property because it should be immutable. return self._missing_value + @property + def missing_value_code(self): + return self.reverse_categories[self.missing_value] + def has_label(self, value): return value in self.reverse_categories diff --git a/zipline/pipeline/factors/factor.py b/zipline/pipeline/factors/factor.py index a76002a9..6a0b0b08 100644 --- a/zipline/pipeline/factors/factor.py +++ b/zipline/pipeline/factors/factor.py @@ -47,6 +47,7 @@ from zipline.utils.input_validation import expect_types from zipline.utils.math_utils import nanmean, nanstd from zipline.utils.numpy_utils import ( bool_dtype, + categorical_dtype, coerce_to_dtype, datetime64ns_dtype, float64_dtype, @@ -939,15 +940,26 @@ class GroupedRowTransform(Factor): def _compute(self, arrays, dates, assets, mask): data = arrays[0] - null_group_value = self.inputs[1].missing_value - group_labels = where( - mask, - arrays[1], - null_group_value, - ) + groupby_expr = self.inputs[1] + if groupby_expr.dtype == int64_dtype: + group_labels = arrays[1] + null_label = self.inputs[1].missing_value + elif groupby_expr.dtype == categorical_dtype: + # Coerce our LabelArray into an isomorphic array of ints. This is + # necessary because np.where doesn't know about LabelArrays or the + # void dtype. + group_labels = arrays[1].as_int_array() + null_label = arrays[1].missing_value_code + else: + raise TypeError( + "Unexpected groupby dtype: %s." % groupby_expr.dtype + ) + + # Make a copy with the null code written to masked locations. + group_labels = where(mask, group_labels, null_label) return where( - group_labels != null_group_value, + group_labels != null_label, naive_grouped_rowwise_apply( data=data, group_labels=group_labels,