BUG: Fix groupby with string columns.

The previous algorithm assumed that the group labels were integers. It
produced nonsense with LabelArrays (though sadly didn't crash because
numpy promotes None and void to object).
This commit is contained in:
Scott Sanderson
2016-05-10 12:40:36 -04:00
parent 2431aaefb5
commit f7e9281b14
3 changed files with 36 additions and 7 deletions
+13
View File
@@ -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)),
+4
View File
@@ -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
+19 -7
View File
@@ -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,