mirror of
https://github.com/wassname/catalyst.git
synced 2026-08-18 11:50:11 +08:00
ENH: Add support for strings in Pipeline.
- Adds a new class, ``LabelArray``, which is a subclass of np.ndarray. LabelArray is conceptually similar to pandas.Categorical, in that it stores data with many duplicate values as indices into an array of unique values. For string data with many duplicates (e.g. time-series of tickers or or industry classifications), this provides multiple orders of magnitude of improvement when doing string operations, especially string comparison/matching operations. - Adds a new generic object "specialization" for `AdjustedArrayWindow`, and a corresponding ObjectOverwrite adjustment. - Adds a new ``postprocess`` method to ``zipline.pipeline.term.Term``. This method is called on the final result of any pipeline expression after screen filtering has occurred. The default implementation of ``postprocess`` is identity, but Classifier overrides it to coerce string columns into pandas.Categoricals before presenting them to the user.
This commit is contained in:
Scott Sanderson
@@ -9,19 +9,23 @@ from nose_parameterized import parameterized
|
||||
from numpy import (
|
||||
arange,
|
||||
array,
|
||||
asarray,
|
||||
dtype,
|
||||
full,
|
||||
where,
|
||||
)
|
||||
from numpy.testing import assert_array_equal
|
||||
from six.moves import zip_longest
|
||||
from toolz import curry
|
||||
|
||||
from zipline.errors import WindowLengthNotPositive, WindowLengthTooLong
|
||||
from zipline.lib.adjustment import (
|
||||
Datetime64Overwrite,
|
||||
Float64Multiply,
|
||||
Float64Overwrite,
|
||||
ObjectOverwrite,
|
||||
)
|
||||
from zipline.lib.adjusted_array import AdjustedArray, NOMASK
|
||||
from zipline.lib.labelarray import LabelArray
|
||||
from zipline.testing import check_arrays, parameter_space
|
||||
from zipline.utils.numpy_utils import (
|
||||
coerce_to_dtype,
|
||||
@@ -29,6 +33,7 @@ from zipline.utils.numpy_utils import (
|
||||
default_missing_value_for_dtype,
|
||||
float64_dtype,
|
||||
int64_dtype,
|
||||
object_dtype,
|
||||
)
|
||||
|
||||
|
||||
@@ -62,12 +67,41 @@ def valid_window_lengths(underlying_buffer_length):
|
||||
return iter(range(1, underlying_buffer_length + 1))
|
||||
|
||||
|
||||
def _gen_unadjusted_cases(dtype):
|
||||
@curry
|
||||
def as_dtype(dtype, data):
|
||||
"""
|
||||
Curried wrapper around array.astype for when you have the dtype before you
|
||||
have the data.
|
||||
"""
|
||||
return asarray(data).astype(dtype)
|
||||
|
||||
|
||||
@curry
|
||||
def as_labelarray(initial_dtype, missing_value, array):
|
||||
"""
|
||||
Curried wrapper around LabelArray, that round-trips the input data through
|
||||
`initial_dtype` first.
|
||||
"""
|
||||
return LabelArray(
|
||||
array.astype(initial_dtype),
|
||||
missing_value=initial_dtype.type(''),
|
||||
)
|
||||
|
||||
|
||||
bytes_dtype = dtype('S3')
|
||||
unicode_dtype = dtype('U3')
|
||||
|
||||
|
||||
def _gen_unadjusted_cases(name,
|
||||
make_input,
|
||||
make_expected_output,
|
||||
missing_value):
|
||||
nrows = 6
|
||||
ncols = 3
|
||||
data = arange(nrows * ncols).astype(dtype).reshape(nrows, ncols)
|
||||
missing_value = default_missing_value_for_dtype(dtype)
|
||||
|
||||
raw_data = arange(nrows * ncols).reshape(nrows, ncols)
|
||||
input_array = make_input(raw_data)
|
||||
expected_output_array = make_expected_output(raw_data)
|
||||
|
||||
for windowlen in valid_window_lengths(nrows):
|
||||
|
||||
@@ -76,13 +110,13 @@ def _gen_unadjusted_cases(dtype):
|
||||
)
|
||||
|
||||
yield (
|
||||
"dtype_%s_length_%d" % (dtype, windowlen),
|
||||
data,
|
||||
"%s_length_%d" % (name, windowlen),
|
||||
input_array,
|
||||
windowlen,
|
||||
{},
|
||||
missing_value,
|
||||
[
|
||||
data[offset:offset + windowlen]
|
||||
expected_output_array[offset:offset + windowlen]
|
||||
for offset in range(num_legal_windows)
|
||||
],
|
||||
)
|
||||
@@ -156,84 +190,125 @@ def _gen_multiplicative_adjustment_cases(dtype):
|
||||
[1, 6, 1],
|
||||
[1, 1, 1]], dtype=dtype)
|
||||
|
||||
return _gen_expectations(baseline, adjustments, buffer_as_of, nrows)
|
||||
|
||||
|
||||
def _gen_overwrite_adjustment_cases(dtype):
|
||||
"""
|
||||
Generate test cases for overwrite adjustments.
|
||||
|
||||
The algorithm used here is the same as the one used above for
|
||||
multiplicative adjustments. The only difference is the semantics of how
|
||||
the adjustments are expected to modify the arrays.
|
||||
"""
|
||||
adjustment_type = {
|
||||
float64_dtype: Float64Overwrite,
|
||||
datetime64ns_dtype: Datetime64Overwrite,
|
||||
}[dtype]
|
||||
|
||||
nrows, ncols = 6, 3
|
||||
adjustments = {}
|
||||
buffer_as_of = [None] * 6
|
||||
baseline = buffer_as_of[0] = full((nrows, ncols), 2, dtype=dtype)
|
||||
|
||||
# Note that row indices are inclusive!
|
||||
adjustments[1] = [
|
||||
adjustment_type(0, 0, 0, 0, coerce_to_dtype(dtype, 1)),
|
||||
]
|
||||
buffer_as_of[1] = array([[1, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2]], dtype=dtype)
|
||||
|
||||
# No adjustment at index 2.
|
||||
buffer_as_of[2] = buffer_as_of[1]
|
||||
|
||||
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[3] = array([[4, 2, 2],
|
||||
[4, 3, 2],
|
||||
[2, 3, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2]], dtype=dtype)
|
||||
|
||||
adjustments[4] = [
|
||||
adjustment_type(0, 3, 2, 2, coerce_to_dtype(dtype, 5))
|
||||
]
|
||||
buffer_as_of[4] = array([[4, 2, 5],
|
||||
[4, 3, 5],
|
||||
[2, 3, 5],
|
||||
[2, 2, 5],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2]], 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[5] = array([[4, 6, 5],
|
||||
[4, 6, 5],
|
||||
[2, 6, 7],
|
||||
[2, 6, 5],
|
||||
[2, 6, 2],
|
||||
[2, 2, 2]], dtype=dtype)
|
||||
|
||||
return _gen_expectations(
|
||||
baseline,
|
||||
default_missing_value_for_dtype(dtype),
|
||||
adjustments,
|
||||
buffer_as_of,
|
||||
nrows,
|
||||
)
|
||||
|
||||
|
||||
def _gen_expectations(baseline, adjustments, buffer_as_of, nrows):
|
||||
def _gen_overwrite_adjustment_cases(name,
|
||||
make_input,
|
||||
make_expected_output,
|
||||
dtype,
|
||||
missing_value):
|
||||
"""
|
||||
Generate test cases for overwrite adjustments.
|
||||
|
||||
The algorithm used here is the same as the one used above for
|
||||
multiplicative adjustments. The only difference is the semantics of how
|
||||
the adjustments are expected to modify the arrays.
|
||||
|
||||
This is parameterized on `make_input` and make_expected_output functions,
|
||||
which take 2-D lists of values and transform them into desired input/output
|
||||
arrays. We do this so that we can easily test both vanilla numpy ndarrays
|
||||
and our own LabelArray class for strings.
|
||||
"""
|
||||
adjustment_type = {
|
||||
float64_dtype: Float64Overwrite,
|
||||
datetime64ns_dtype: Datetime64Overwrite,
|
||||
bytes_dtype: ObjectOverwrite,
|
||||
unicode_dtype: ObjectOverwrite,
|
||||
object_dtype: ObjectOverwrite,
|
||||
}[dtype]
|
||||
|
||||
if dtype == object_dtype:
|
||||
# When we're testing object dtypes, we expect to have strings, but
|
||||
# coerce_to_dtype(object, 3) just gives 3 as a Python integer.
|
||||
def make_overwrite_value(dtype, value):
|
||||
return str(value)
|
||||
else:
|
||||
make_overwrite_value = coerce_to_dtype
|
||||
|
||||
adjustments = {}
|
||||
buffer_as_of = [None] * 6
|
||||
baseline = make_input([[2, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2]])
|
||||
|
||||
buffer_as_of[0] = make_expected_output([[2, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2]])
|
||||
|
||||
# Note that row indices are inclusive!
|
||||
adjustments[1] = [
|
||||
adjustment_type(0, 0, 0, 0, make_overwrite_value(dtype, 1)),
|
||||
]
|
||||
buffer_as_of[1] = make_expected_output([[1, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2]])
|
||||
|
||||
# No adjustment at index 2.
|
||||
buffer_as_of[2] = buffer_as_of[1]
|
||||
|
||||
adjustments[3] = [
|
||||
adjustment_type(1, 2, 1, 1, make_overwrite_value(dtype, 3)),
|
||||
adjustment_type(0, 1, 0, 0, make_overwrite_value(dtype, 4)),
|
||||
]
|
||||
buffer_as_of[3] = make_expected_output([[4, 2, 2],
|
||||
[4, 3, 2],
|
||||
[2, 3, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2]])
|
||||
|
||||
adjustments[4] = [
|
||||
adjustment_type(0, 3, 2, 2, make_overwrite_value(dtype, 5))
|
||||
]
|
||||
buffer_as_of[4] = make_expected_output([[4, 2, 5],
|
||||
[4, 3, 5],
|
||||
[2, 3, 5],
|
||||
[2, 2, 5],
|
||||
[2, 2, 2],
|
||||
[2, 2, 2]])
|
||||
|
||||
adjustments[5] = [
|
||||
adjustment_type(0, 4, 1, 1, make_overwrite_value(dtype, 6)),
|
||||
adjustment_type(2, 2, 2, 2, make_overwrite_value(dtype, 7)),
|
||||
]
|
||||
buffer_as_of[5] = make_expected_output([[4, 6, 5],
|
||||
[4, 6, 5],
|
||||
[2, 6, 7],
|
||||
[2, 6, 5],
|
||||
[2, 6, 2],
|
||||
[2, 2, 2]])
|
||||
|
||||
return _gen_expectations(
|
||||
baseline,
|
||||
missing_value,
|
||||
adjustments,
|
||||
buffer_as_of,
|
||||
nrows=6,
|
||||
)
|
||||
|
||||
|
||||
def _gen_expectations(baseline,
|
||||
missing_value,
|
||||
adjustments,
|
||||
buffer_as_of,
|
||||
nrows):
|
||||
|
||||
missing_value = default_missing_value_for_dtype(baseline.dtype)
|
||||
for windowlen in valid_window_lengths(nrows):
|
||||
|
||||
num_legal_windows = num_windows_of_length_M_on_buffers_of_length_N(
|
||||
@@ -263,8 +338,60 @@ class AdjustedArrayTestCase(TestCase):
|
||||
|
||||
@parameterized.expand(
|
||||
chain(
|
||||
_gen_unadjusted_cases(float64_dtype),
|
||||
_gen_unadjusted_cases(datetime64ns_dtype),
|
||||
_gen_unadjusted_cases(
|
||||
'float',
|
||||
make_input=as_dtype(float64_dtype),
|
||||
make_expected_output=as_dtype(float64_dtype),
|
||||
missing_value=default_missing_value_for_dtype(float64_dtype),
|
||||
),
|
||||
_gen_unadjusted_cases(
|
||||
'datetime',
|
||||
make_input=as_dtype(datetime64ns_dtype),
|
||||
make_expected_output=as_dtype(datetime64ns_dtype),
|
||||
missing_value=default_missing_value_for_dtype(
|
||||
datetime64ns_dtype
|
||||
),
|
||||
),
|
||||
# Test passing an array of strings to AdjustedArray.
|
||||
_gen_unadjusted_cases(
|
||||
'bytes_ndarray',
|
||||
make_input=as_dtype(bytes_dtype),
|
||||
make_expected_output=as_labelarray(bytes_dtype, b''),
|
||||
missing_value=b'',
|
||||
),
|
||||
_gen_unadjusted_cases(
|
||||
'unicode_ndarray',
|
||||
make_input=as_dtype(unicode_dtype),
|
||||
make_expected_output=as_labelarray(unicode_dtype, u''),
|
||||
missing_value=u'',
|
||||
),
|
||||
_gen_unadjusted_cases(
|
||||
'object_ndarray',
|
||||
make_input=lambda a: a.astype(str).astype(object),
|
||||
make_expected_output=as_labelarray(bytes_dtype, b''),
|
||||
missing_value=b'',
|
||||
),
|
||||
# Test passing a LabelArray directly to AdjustedArray.
|
||||
_gen_unadjusted_cases(
|
||||
'bytes_labelarray',
|
||||
make_input=as_labelarray(bytes_dtype, b''),
|
||||
make_expected_output=as_labelarray(bytes_dtype, b''),
|
||||
missing_value=b'',
|
||||
),
|
||||
_gen_unadjusted_cases(
|
||||
'unicode_labelarray',
|
||||
make_input=as_labelarray(unicode_dtype, u''),
|
||||
make_expected_output=as_labelarray(bytes_dtype, u''),
|
||||
missing_value=u'',
|
||||
),
|
||||
_gen_unadjusted_cases(
|
||||
'object_labelarray',
|
||||
make_input=(
|
||||
lambda a: LabelArray(a.astype(str).astype(object), b'')
|
||||
),
|
||||
make_expected_output=as_labelarray(bytes_dtype, b''),
|
||||
missing_value=b'',
|
||||
),
|
||||
)
|
||||
)
|
||||
def test_no_adjustments(self,
|
||||
@@ -273,14 +400,13 @@ class AdjustedArrayTestCase(TestCase):
|
||||
lookback,
|
||||
adjustments,
|
||||
missing_value,
|
||||
expected):
|
||||
expected_output):
|
||||
|
||||
array = AdjustedArray(data, NOMASK, adjustments, missing_value)
|
||||
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):
|
||||
self.assertEqual(yielded.dtype, data.dtype)
|
||||
assert_array_equal(yielded, expected_yield)
|
||||
in_out = zip(array.traverse(lookback), expected_output)
|
||||
for yielded, expected_yield in in_out:
|
||||
check_arrays(yielded, expected_yield)
|
||||
|
||||
@parameterized.expand(_gen_multiplicative_adjustment_cases(float64_dtype))
|
||||
def test_multiplicative_adjustments(self,
|
||||
@@ -295,12 +421,70 @@ class AdjustedArrayTestCase(TestCase):
|
||||
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):
|
||||
assert_array_equal(yielded, expected_yield)
|
||||
check_arrays(yielded, expected_yield)
|
||||
|
||||
@parameterized.expand(
|
||||
chain(
|
||||
_gen_overwrite_adjustment_cases(float64_dtype),
|
||||
_gen_overwrite_adjustment_cases(datetime64ns_dtype),
|
||||
_gen_overwrite_adjustment_cases(
|
||||
'float',
|
||||
make_input=as_dtype(float64_dtype),
|
||||
make_expected_output=as_dtype(float64_dtype),
|
||||
dtype=float64_dtype,
|
||||
missing_value=default_missing_value_for_dtype(float64_dtype),
|
||||
),
|
||||
_gen_overwrite_adjustment_cases(
|
||||
'datetime',
|
||||
make_input=as_dtype(datetime64ns_dtype),
|
||||
make_expected_output=as_dtype(datetime64ns_dtype),
|
||||
dtype=datetime64ns_dtype,
|
||||
missing_value=default_missing_value_for_dtype(
|
||||
datetime64ns_dtype,
|
||||
),
|
||||
),
|
||||
# There are six cases here:
|
||||
# Using np.bytes/np.unicode/python string arrays as inputs.
|
||||
# Passing np.bytes/np.unicode/python string arrays to LabelArray,
|
||||
# and using those as input.
|
||||
#
|
||||
# The outputs should always be LabelArrays.
|
||||
_gen_unadjusted_cases(
|
||||
'bytes_ndarray',
|
||||
make_input=as_dtype(bytes_dtype),
|
||||
make_expected_output=as_labelarray(bytes_dtype, b''),
|
||||
missing_value=b'',
|
||||
),
|
||||
_gen_unadjusted_cases(
|
||||
'unicode_ndarray',
|
||||
make_input=as_dtype(unicode_dtype),
|
||||
make_expected_output=as_labelarray(unicode_dtype, u''),
|
||||
missing_value=u'',
|
||||
),
|
||||
_gen_unadjusted_cases(
|
||||
'object_ndarray',
|
||||
make_input=lambda a: a.astype(str).astype(object),
|
||||
make_expected_output=as_labelarray(bytes_dtype, b''),
|
||||
missing_value=b'',
|
||||
),
|
||||
_gen_unadjusted_cases(
|
||||
'bytes_labelarray',
|
||||
make_input=as_labelarray(bytes_dtype, b''),
|
||||
make_expected_output=as_labelarray(bytes_dtype, b''),
|
||||
missing_value=b'',
|
||||
),
|
||||
_gen_unadjusted_cases(
|
||||
'unicode_labelarray',
|
||||
make_input=as_labelarray(unicode_dtype, u''),
|
||||
make_expected_output=as_labelarray(bytes_dtype, u''),
|
||||
missing_value=u'',
|
||||
),
|
||||
_gen_unadjusted_cases(
|
||||
'object_labelarray',
|
||||
make_input=(
|
||||
lambda a: LabelArray(a.astype(str).astype(object), b'')
|
||||
),
|
||||
make_expected_output=as_labelarray(bytes_dtype, b''),
|
||||
missing_value=b'',
|
||||
),
|
||||
)
|
||||
)
|
||||
def test_overwrite_adjustment_cases(self,
|
||||
@@ -314,11 +498,15 @@ class AdjustedArrayTestCase(TestCase):
|
||||
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):
|
||||
self.assertEqual(yielded.dtype, data.dtype)
|
||||
assert_array_equal(yielded, expected_yield)
|
||||
check_arrays(yielded, expected_yield)
|
||||
|
||||
@parameter_space(
|
||||
dtype=[float64_dtype, int64_dtype, datetime64ns_dtype],
|
||||
__fail_fast=True,
|
||||
dtype=[
|
||||
float64_dtype,
|
||||
int64_dtype,
|
||||
datetime64ns_dtype,
|
||||
],
|
||||
missing_value=[0, 10000],
|
||||
window_length=[2, 3],
|
||||
)
|
||||
@@ -341,6 +529,37 @@ class AdjustedArrayTestCase(TestCase):
|
||||
for expected, actual in zip(gen_expected, gen_actual):
|
||||
check_arrays(expected, actual)
|
||||
|
||||
@parameter_space(
|
||||
__fail_fast=True,
|
||||
dtype=[bytes_dtype, unicode_dtype, object_dtype],
|
||||
missing_value=["0", "-1", ""],
|
||||
window_length=[2, 3],
|
||||
)
|
||||
def test_masking_with_strings(self, dtype, missing_value, window_length):
|
||||
missing_value = coerce_to_dtype(dtype, missing_value)
|
||||
baseline_ints = arange(15).reshape(5, 3)
|
||||
|
||||
# Coerce to string first so that coercion to object gets us an array of
|
||||
# string objects.
|
||||
baseline = baseline_ints.astype(str).astype(dtype)
|
||||
mask = (baseline_ints % 2).astype(bool)
|
||||
|
||||
masked_baseline = LabelArray(baseline, missing_value=missing_value)
|
||||
masked_baseline[~mask] = missing_value
|
||||
|
||||
array = AdjustedArray(
|
||||
baseline,
|
||||
mask,
|
||||
adjustments={},
|
||||
missing_value=missing_value,
|
||||
)
|
||||
|
||||
gen_expected = moving_window(masked_baseline, window_length)
|
||||
gen_actual = array.traverse(window_length=window_length)
|
||||
|
||||
for expected, actual in zip(gen_expected, gen_actual):
|
||||
check_arrays(expected, actual)
|
||||
|
||||
def test_invalid_lookback(self):
|
||||
|
||||
data = arange(30, dtype=float).reshape(6, 5)
|
||||
|
||||
@@ -1,12 +1,21 @@
|
||||
import numpy as np
|
||||
|
||||
from zipline.lib.labelarray import LabelArray
|
||||
from zipline.pipeline import Classifier
|
||||
from zipline.testing import parameter_space
|
||||
from zipline.utils.numpy_utils import int64_dtype
|
||||
from zipline.utils.numpy_utils import (
|
||||
categorical_dtype,
|
||||
coerce_to_dtype,
|
||||
int64_dtype,
|
||||
)
|
||||
|
||||
from .base import BasePipelineTestCase
|
||||
|
||||
|
||||
bytes_dtype = np.dtype('S3')
|
||||
unicode_dtype = np.dtype('U3')
|
||||
|
||||
|
||||
class ClassifierTestCase(BasePipelineTestCase):
|
||||
|
||||
@parameter_space(mv=[-1, 0, 1, 999])
|
||||
@@ -69,10 +78,56 @@ class ClassifierTestCase(BasePipelineTestCase):
|
||||
mask=self.build_mask(self.ones_mask(shape=data.shape)),
|
||||
)
|
||||
|
||||
@parameter_space(missing=[-1, 0, 1])
|
||||
def test_disallow_comparison_to_missing_value(self, missing):
|
||||
@parameter_space(
|
||||
__fail_fast=True,
|
||||
compval=['a', 'ab', 'not in the array'],
|
||||
labelarray_dtype=(bytes_dtype, categorical_dtype, unicode_dtype),
|
||||
)
|
||||
def test_string_eq(self, compval, labelarray_dtype):
|
||||
|
||||
compval = labelarray_dtype.type(compval)
|
||||
|
||||
class C(Classifier):
|
||||
dtype = int64_dtype
|
||||
dtype = categorical_dtype
|
||||
missing_value = ''
|
||||
inputs = ()
|
||||
window_length = 0
|
||||
|
||||
c = C()
|
||||
|
||||
# There's no significance to the values here other than that they
|
||||
# contain a mix of the comparison value and other values.
|
||||
data = LabelArray(
|
||||
np.asarray(
|
||||
[['', 'a', 'ab', 'ba'],
|
||||
['z', 'ab', 'a', 'ab'],
|
||||
['aa', 'ab', '', 'ab'],
|
||||
['aa', 'a', 'ba', 'ba']],
|
||||
dtype=labelarray_dtype,
|
||||
),
|
||||
missing_value='',
|
||||
)
|
||||
|
||||
self.check_terms(
|
||||
terms={
|
||||
'eq': c.eq(compval),
|
||||
},
|
||||
expected={
|
||||
'eq': (data == compval),
|
||||
},
|
||||
initial_workspace={c: data},
|
||||
mask=self.build_mask(self.ones_mask(shape=data.shape)),
|
||||
)
|
||||
|
||||
@parameter_space(
|
||||
missing=[-1, 0, 1],
|
||||
dtype_=[int64_dtype, categorical_dtype],
|
||||
)
|
||||
def test_disallow_comparison_to_missing_value(self, missing, dtype_):
|
||||
missing = coerce_to_dtype(dtype_, missing)
|
||||
|
||||
class C(Classifier):
|
||||
dtype = dtype_
|
||||
missing_value = missing
|
||||
inputs = ()
|
||||
window_length = 0
|
||||
@@ -82,7 +137,7 @@ class ClassifierTestCase(BasePipelineTestCase):
|
||||
errmsg = str(e.exception)
|
||||
self.assertEqual(
|
||||
errmsg,
|
||||
"Comparison against self.missing_value ({v}) in C.eq().\n"
|
||||
"Comparison against self.missing_value ({v!r}) in C.eq().\n"
|
||||
"Missing values have NaN semantics, so the requested comparison"
|
||||
" would always produce False.\n"
|
||||
"Use the isnull() method to check for missing values.".format(
|
||||
@@ -118,3 +173,100 @@ class ClassifierTestCase(BasePipelineTestCase):
|
||||
initial_workspace={c: data},
|
||||
mask=self.build_mask(self.ones_mask(shape=data.shape)),
|
||||
)
|
||||
|
||||
@parameter_space(
|
||||
__fail_fast=True,
|
||||
compval=['a', 'ab', '', 'not in the array'],
|
||||
missing=['a', 'ab', '', 'not in the array'],
|
||||
labelarray_dtype=(bytes_dtype, unicode_dtype, categorical_dtype),
|
||||
)
|
||||
def test_string_not_equal(self, compval, missing, labelarray_dtype):
|
||||
|
||||
compval = labelarray_dtype.type(compval)
|
||||
|
||||
class C(Classifier):
|
||||
dtype = categorical_dtype
|
||||
missing_value = missing
|
||||
inputs = ()
|
||||
window_length = 0
|
||||
|
||||
c = C()
|
||||
|
||||
# There's no significance to the values here other than that they
|
||||
# contain a mix of the comparison value and other values.
|
||||
data = LabelArray(
|
||||
np.asarray(
|
||||
[['', 'a', 'ab', 'ba'],
|
||||
['z', 'ab', 'a', 'ab'],
|
||||
['aa', 'ab', '', 'ab'],
|
||||
['aa', 'a', 'ba', 'ba']],
|
||||
dtype=labelarray_dtype,
|
||||
),
|
||||
missing_value=missing,
|
||||
)
|
||||
|
||||
expected = (
|
||||
(data.as_int_array() != data.reverse_categories.get(compval, -1)) &
|
||||
(data.as_int_array() != data.reverse_categories[C.missing_value])
|
||||
)
|
||||
|
||||
self.check_terms(
|
||||
terms={
|
||||
'ne': c != compval,
|
||||
},
|
||||
expected={
|
||||
'ne': expected,
|
||||
},
|
||||
initial_workspace={c: data},
|
||||
mask=self.build_mask(self.ones_mask(shape=data.shape)),
|
||||
)
|
||||
|
||||
@parameter_space(
|
||||
__fail_fast=True,
|
||||
compval=['a', 'b', 'ab', 'not in the array'],
|
||||
missing=['a', 'ab', '', 'not in the array'],
|
||||
labelarray_dtype=(categorical_dtype, bytes_dtype, unicode_dtype),
|
||||
)
|
||||
def test_string_elementwise_predicates(self,
|
||||
compval,
|
||||
missing,
|
||||
labelarray_dtype):
|
||||
|
||||
missing = labelarray_dtype.type(missing)
|
||||
compval = labelarray_dtype.type(compval)
|
||||
|
||||
class C(Classifier):
|
||||
dtype = categorical_dtype
|
||||
missing_value = missing
|
||||
inputs = ()
|
||||
window_length = 0
|
||||
|
||||
c = C()
|
||||
|
||||
# There's no significance to the values here other than that they
|
||||
# contain a mix of the comparison value and other values.
|
||||
data = LabelArray(
|
||||
np.asarray(
|
||||
[['', 'a', 'ab', 'ba'],
|
||||
['z', 'ab', 'a', 'ab'],
|
||||
['aa', 'ab', '', 'ab'],
|
||||
['aa', 'a', 'ba', 'ba']],
|
||||
dtype=labelarray_dtype,
|
||||
),
|
||||
missing_value=missing,
|
||||
)
|
||||
|
||||
self.check_terms(
|
||||
terms={
|
||||
'startswith': c.startswith(compval),
|
||||
'endswith': c.endswith(compval),
|
||||
'contains': c.contains(compval),
|
||||
},
|
||||
expected={
|
||||
'startswith': (data.startswith(compval) & (data != missing)),
|
||||
'endswith': (data.endswith(compval) & (data != missing)),
|
||||
'contains': (data.contains(compval) & (data != missing)),
|
||||
},
|
||||
initial_workspace={c: data},
|
||||
mask=self.build_mask(self.ones_mask(shape=data.shape)),
|
||||
)
|
||||
|
||||
@@ -7,6 +7,7 @@ from unittest import TestCase
|
||||
from pandas import date_range, DataFrame
|
||||
from pandas.util.testing import assert_frame_equal
|
||||
|
||||
from zipline.lib.labelarray import LabelArray
|
||||
from zipline.pipeline import Pipeline
|
||||
from zipline.pipeline.data.testing import TestingDataSet as TDS
|
||||
from zipline.testing import chrange, temp_pipeline_engine
|
||||
@@ -35,6 +36,21 @@ class LatestTestCase(TestCase):
|
||||
|
||||
def expected_latest(self, column, slice_):
|
||||
loader = self.engine.get_loader(column)
|
||||
|
||||
index = self.calendar[slice_]
|
||||
columns = self.assets
|
||||
values = loader.values(column.dtype, self.calendar, self.sids)[slice_]
|
||||
|
||||
if column.dtype.kind in ('O', 'S', 'U'):
|
||||
# For string columns, we expect a categorical in the output.
|
||||
return LabelArray(
|
||||
values,
|
||||
missing_value=column.missing_value,
|
||||
).as_categorical_frame(
|
||||
index=index,
|
||||
columns=columns,
|
||||
)
|
||||
|
||||
return DataFrame(
|
||||
loader.values(column.dtype, self.calendar, self.sids)[slice_],
|
||||
index=self.calendar[slice_],
|
||||
@@ -55,6 +71,6 @@ class LatestTestCase(TestCase):
|
||||
dates_to_test[-1],
|
||||
)
|
||||
for column in columns:
|
||||
float_result = result[column.name].unstack()
|
||||
expected_float_result = self.expected_latest(column, cal_slice)
|
||||
assert_frame_equal(float_result, expected_float_result)
|
||||
col_result = result[column.name].unstack()
|
||||
expected_col_result = self.expected_latest(column, cal_slice)
|
||||
assert_frame_equal(col_result, expected_col_result)
|
||||
|
||||
@@ -471,7 +471,8 @@ class ObjectIdentityTestCase(TestCase):
|
||||
for column in TestingDataSet.columns:
|
||||
if column.dtype == bool_dtype:
|
||||
self.assertIsInstance(column.latest, Filter)
|
||||
elif column.dtype == int64_dtype:
|
||||
elif (column.dtype == int64_dtype
|
||||
or column.dtype.kind in ('O', 'S', 'U')):
|
||||
self.assertIsInstance(column.latest, Classifier)
|
||||
elif column.dtype in factor_dtypes:
|
||||
self.assertIsInstance(column.latest, Factor)
|
||||
|
||||
@@ -0,0 +1,151 @@
|
||||
from itertools import product
|
||||
from operator import eq, ne
|
||||
import numpy as np
|
||||
|
||||
from zipline.lib.labelarray import LabelArray
|
||||
from zipline.testing import check_arrays, parameter_space, ZiplineTestCase
|
||||
|
||||
|
||||
def rotN(l, N):
|
||||
"""
|
||||
Rotate a list of elements.
|
||||
|
||||
Pulls N elements off the end of the list and appends them to the front.
|
||||
|
||||
>>> rotN(['a', 'b', 'c', 'd'], 2)
|
||||
['c', 'd', 'a', 'b']
|
||||
>>> rotN(['a', 'b', 'c', 'd'], 3)
|
||||
['d', 'a', 'b', 'c']
|
||||
"""
|
||||
assert len(l) >= N, "Can't rotate list by longer than its length."
|
||||
return l[N:] + l[:N]
|
||||
|
||||
|
||||
class LabelArrayTestCase(ZiplineTestCase):
|
||||
|
||||
@classmethod
|
||||
def init_class_fixtures(cls):
|
||||
super(LabelArrayTestCase, cls).init_class_fixtures()
|
||||
|
||||
cls.rowvalues = row = ['', 'a', 'b', 'ab', 'a', '', 'b', 'ab', 'z']
|
||||
cls.strs = np.array([rotN(row, i) for i in range(3)])
|
||||
|
||||
def test_fail_on_direct_construction(self):
|
||||
# See http://docs.scipy.org/doc/numpy-1.10.0/user/basics.subclassing.html#simple-example-adding-an-extra-attribute-to-ndarray # noqa
|
||||
|
||||
with self.assertRaises(TypeError) as e:
|
||||
np.ndarray.__new__(LabelArray, (5, 5))
|
||||
|
||||
self.assertEqual(
|
||||
str(e.exception),
|
||||
"Direct construction of LabelArrays is not supported."
|
||||
)
|
||||
|
||||
@parameter_space(
|
||||
__fail_fast=True,
|
||||
s=['', 'a', 'z', 'aa', 'not in the array'],
|
||||
shape=[(27,), (9, 3), (3, 9), (3, 3, 3)],
|
||||
array_astype=(bytes, unicode, object),
|
||||
scalar_astype=(bytes, unicode, object),
|
||||
)
|
||||
def test_compare_to_str(self, s, shape, array_astype, scalar_astype):
|
||||
strs = self.strs.reshape(shape).astype(array_astype)
|
||||
arr = LabelArray(strs, missing_value='')
|
||||
check_arrays(strs == s, arr == s)
|
||||
check_arrays(strs != s, arr != s)
|
||||
|
||||
np_startswith = np.vectorize(lambda elem: elem.startswith(s))
|
||||
check_arrays(arr.startswith(s), np_startswith(strs))
|
||||
|
||||
np_endswith = np.vectorize(lambda elem: elem.endswith(s))
|
||||
check_arrays(arr.endswith(s), np_endswith(strs))
|
||||
|
||||
np_contains = np.vectorize(lambda elem: s in elem)
|
||||
check_arrays(arr.contains(s), np_contains(strs))
|
||||
|
||||
def test_compare_to_str_array(self):
|
||||
strs = self.strs
|
||||
shape = strs.shape
|
||||
arr = LabelArray(strs, missing_value='')
|
||||
check_arrays(strs == arr, np.full_like(strs, True, dtype=bool))
|
||||
check_arrays(strs != arr, np.full_like(strs, False, dtype=bool))
|
||||
|
||||
def broadcastable_row(value, dtype):
|
||||
return np.full((shape[0], 1), value, dtype=strs.dtype)
|
||||
|
||||
def broadcastable_col(value, dtype):
|
||||
return np.full((1, shape[1]), value, dtype=strs.dtype)
|
||||
|
||||
for comparator, dtype, value in product((eq, ne),
|
||||
(bytes, unicode, object),
|
||||
set(self.rowvalues)):
|
||||
check_arrays(
|
||||
comparator(arr, np.full_like(strs, value)),
|
||||
comparator(strs, value),
|
||||
)
|
||||
check_arrays(
|
||||
comparator(arr, broadcastable_row(value, dtype=dtype)),
|
||||
comparator(strs, value),
|
||||
)
|
||||
check_arrays(
|
||||
comparator(arr, broadcastable_col(value, dtype=dtype)),
|
||||
comparator(strs, value),
|
||||
)
|
||||
|
||||
@parameter_space(
|
||||
__fail_fast=True,
|
||||
slice_=[
|
||||
0, 1, -1,
|
||||
slice(None),
|
||||
slice(0, 0),
|
||||
slice(0, 3),
|
||||
slice(1, 4),
|
||||
slice(0),
|
||||
slice(None, 1),
|
||||
slice(0, 4, 2),
|
||||
(slice(None), 1),
|
||||
(slice(None), slice(None)),
|
||||
(slice(None), slice(1, 2)),
|
||||
]
|
||||
)
|
||||
def test_slicing_preserves_attributes(self, slice_):
|
||||
arr = LabelArray(self.strs.reshape((9, 3)), missing_value='')
|
||||
sliced = arr[slice_]
|
||||
self.assertIsInstance(sliced, LabelArray)
|
||||
self.assertIs(sliced.categories, arr.categories)
|
||||
self.assertIs(sliced.reverse_categories, arr.reverse_categories)
|
||||
self.assertIs(sliced.missing_value, arr.missing_value)
|
||||
|
||||
def test_infer_categories(self):
|
||||
arr1d = LabelArray(self.strs, missing_value='')
|
||||
codes1d = arr1d.as_int_array()
|
||||
self.assertEqual(arr1d.shape, self.strs.shape)
|
||||
self.assertEqual(arr1d.shape, codes1d.shape)
|
||||
|
||||
categories = arr1d.categories
|
||||
unique_rowvalues = set(self.rowvalues)
|
||||
|
||||
# There should be an entry in categories for each unique row value, and
|
||||
# each integer stored in the data array should be an index into
|
||||
# categories.
|
||||
self.assertEqual(list(categories), sorted(set(self.rowvalues)))
|
||||
self.assertEqual(
|
||||
set(codes1d.ravel()),
|
||||
set(range(len(unique_rowvalues)))
|
||||
)
|
||||
for idx, value in enumerate(arr1d.categories):
|
||||
check_arrays(
|
||||
self.strs == value,
|
||||
arr1d.view(type=np.ndarray) == idx,
|
||||
)
|
||||
|
||||
for shape in (9, 3), (3, 9), (3, 3, 3):
|
||||
strs2d = self.strs.reshape(shape)
|
||||
arr2d = LabelArray(strs2d, missing_value='')
|
||||
codes2d = arr2d.as_int_array()
|
||||
|
||||
self.assertEqual(arr2d.shape, shape)
|
||||
check_arrays(arr2d.categories, categories)
|
||||
|
||||
for idx, value in enumerate(arr2d.categories):
|
||||
check_arrays(strs2d == value, codes2d == idx)
|
||||
Reference in New Issue
Block a user