ENH: Add isnull and notnull methods to Factor.

This commit is contained in:
Scott Sanderson
2016-03-07 16:19:08 -05:00
parent e810f26097
commit f635a14289
7 changed files with 215 additions and 47 deletions
+14 -26
View File
@@ -23,11 +23,11 @@ from zipline.lib.adjustment import (
)
from zipline.lib.adjusted_array import AdjustedArray, NOMASK
from zipline.utils.numpy_utils import (
coerce_to_dtype,
datetime64ns_dtype,
default_missing_value_for_dtype,
float64_dtype,
int64_dtype,
make_datetime64ns,
)
from zipline.utils.test_utils import check_arrays, parameter_space
@@ -62,18 +62,6 @@ def valid_window_lengths(underlying_buffer_length):
return iter(range(1, underlying_buffer_length + 1))
def value_with_dtype(dtype, value):
"""
Make a value with the specified numpy dtype.
"""
name = dtype.name
if name.startswith('datetime64'):
if name != 'datetime64[ns]':
raise TypeError("Expected datetime64[ns], but got %s." % name)
return make_datetime64ns(value)
return dtype.type(value)
def _gen_unadjusted_cases(dtype):
nrows = 6
@@ -124,7 +112,7 @@ def _gen_multiplicative_adjustment_cases(dtype):
# Note that row indices are inclusive!
adjustments[1] = [
adjustment_type(0, 0, 0, 0, value_with_dtype(dtype, 2)),
adjustment_type(0, 0, 0, 0, coerce_to_dtype(dtype, 2)),
]
buffer_as_of[1] = array([[2, 1, 1],
[1, 1, 1],
@@ -137,8 +125,8 @@ def _gen_multiplicative_adjustment_cases(dtype):
buffer_as_of[2] = buffer_as_of[1]
adjustments[3] = [
adjustment_type(1, 2, 1, 1, value_with_dtype(dtype, 3)),
adjustment_type(0, 1, 0, 0, value_with_dtype(dtype, 4)),
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([[8, 1, 1],
[4, 3, 1],
@@ -148,7 +136,7 @@ def _gen_multiplicative_adjustment_cases(dtype):
[1, 1, 1]], dtype=dtype)
adjustments[4] = [
adjustment_type(0, 3, 2, 2, value_with_dtype(dtype, 5))
adjustment_type(0, 3, 2, 2, coerce_to_dtype(dtype, 5))
]
buffer_as_of[4] = array([[8, 1, 5],
[4, 3, 5],
@@ -158,8 +146,8 @@ def _gen_multiplicative_adjustment_cases(dtype):
[1, 1, 1]], dtype=dtype)
adjustments[5] = [
adjustment_type(0, 4, 1, 1, value_with_dtype(dtype, 6)),
adjustment_type(2, 2, 2, 2, value_with_dtype(dtype, 7)),
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([[8, 6, 5],
[4, 18, 5],
@@ -191,7 +179,7 @@ def _gen_overwrite_adjustment_cases(dtype):
# Note that row indices are inclusive!
adjustments[1] = [
adjustment_type(0, 0, 0, 0, value_with_dtype(dtype, 1)),
adjustment_type(0, 0, 0, 0, coerce_to_dtype(dtype, 1)),
]
buffer_as_of[1] = array([[1, 2, 2],
[2, 2, 2],
@@ -204,8 +192,8 @@ def _gen_overwrite_adjustment_cases(dtype):
buffer_as_of[2] = buffer_as_of[1]
adjustments[3] = [
adjustment_type(1, 2, 1, 1, value_with_dtype(dtype, 3)),
adjustment_type(0, 1, 0, 0, value_with_dtype(dtype, 4)),
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],
@@ -215,7 +203,7 @@ def _gen_overwrite_adjustment_cases(dtype):
[2, 2, 2]], dtype=dtype)
adjustments[4] = [
adjustment_type(0, 3, 2, 2, value_with_dtype(dtype, 5))
adjustment_type(0, 3, 2, 2, coerce_to_dtype(dtype, 5))
]
buffer_as_of[4] = array([[4, 2, 5],
[4, 3, 5],
@@ -225,8 +213,8 @@ def _gen_overwrite_adjustment_cases(dtype):
[2, 2, 2]], dtype=dtype)
adjustments[5] = [
adjustment_type(0, 4, 1, 1, value_with_dtype(dtype, 6)),
adjustment_type(2, 2, 2, 2, value_with_dtype(dtype, 7)),
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],
@@ -335,7 +323,7 @@ class AdjustedArrayTestCase(TestCase):
window_length=[2, 3],
)
def test_masking(self, dtype, missing_value, window_length):
missing_value = value_with_dtype(dtype, missing_value)
missing_value = coerce_to_dtype(dtype, missing_value)
baseline_ints = arange(15).reshape(5, 3)
baseline = baseline_ints.astype(dtype)
mask = (baseline_ints % 2).astype(bool)
+73 -1
View File
@@ -22,10 +22,15 @@ from zipline.pipeline.factors import (
Returns,
RSI,
)
from zipline.utils.test_utils import check_allclose, check_arrays
from zipline.utils.test_utils import (
check_allclose,
check_arrays,
parameter_space,
)
from zipline.utils.numpy_utils import (
datetime64ns_dtype,
float64_dtype,
int64_dtype,
NaTns,
)
@@ -59,6 +64,73 @@ class FactorTestCase(BasePipelineTestCase):
with self.assertRaises(UnknownRankMethod):
self.f.rank("not a real rank method")
@parameter_space(method_name=['isnan', 'notnan', 'isfinite'])
def test_float64_only_ops(self, method_name):
class NotFloat(Factor):
dtype = datetime64ns_dtype
inputs = ()
window_length = 0
nf = NotFloat()
meth = getattr(nf, method_name)
with self.assertRaises(TypeError):
meth()
@parameter_space(custom_missing_value=[-1, 0])
def test_isnull_int_dtype(self, custom_missing_value):
class CustomMissingValue(Factor):
dtype = int64_dtype
window_length = 0
missing_value = custom_missing_value
inputs = ()
factor = CustomMissingValue()
data = arange(25).reshape(5, 5)
data[eye(5, dtype=bool)] = custom_missing_value
graph = TermGraph(
{
'isnull': factor.isnull(),
'notnull': factor.notnull(),
}
)
results = self.run_graph(
graph,
initial_workspace={factor: data},
mask=self.build_mask(ones((5, 5))),
)
check_arrays(results['isnull'], eye(5, dtype=bool))
check_arrays(results['notnull'], ~eye(5, dtype=bool))
def test_isnull_datetime_dtype(self):
class DatetimeFactor(Factor):
dtype = datetime64ns_dtype
window_length = 0
inputs = ()
factor = DatetimeFactor()
data = arange(25).reshape(5, 5).astype('datetime64[ns]')
data[eye(5, dtype=bool)] = NaTns
graph = TermGraph(
{
'isnull': factor.isnull(),
'notnull': factor.notnull(),
}
)
results = self.run_graph(
graph,
initial_workspace={factor: data},
mask=self.build_mask(ones((5, 5))),
)
check_arrays(results['isnull'], eye(5, dtype=bool))
check_arrays(results['notnull'], ~eye(5, dtype=bool))
@for_each_factor_dtype
def test_rank_ascending(self, name, factor_dtype):
+10 -2
View File
@@ -345,10 +345,14 @@ class FilterTestCase(BasePipelineTestCase):
data[diag] = nan
results = self.run_graph(
TermGraph({'isnan': self.f.isnan()}),
TermGraph({
'isnan': self.f.isnan(),
'isnull': self.f.isnull(),
}),
initial_workspace={self.f: data},
)
check_arrays(results['isnan'], diag)
check_arrays(results['isnull'], diag)
def test_notnan(self):
data = self.randn_data(seed=10)
@@ -356,10 +360,14 @@ class FilterTestCase(BasePipelineTestCase):
data[diag] = nan
results = self.run_graph(
TermGraph({'notnan': self.f.notnan()}),
TermGraph({
'notnan': self.f.notnan(),
'notnull': self.f.notnull(),
}),
initial_workspace={self.f: data},
)
check_arrays(results['notnan'], ~diag)
check_arrays(results['notnull'], ~diag)
def test_isfinite(self):
data = self.randn_data(seed=10)