wassname/catalyst
1
0
Fork 0
mirror of https://github.com/wassname/catalyst.git synced 2026-07-13 17:42:42 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1095 from quantopian/factor-mask

Browse Source 
Pass a mask (filter) to custom factors
This commit is contained in:
Scott Sanderson
2016-04-07 20:34:10 -04:00
parent c5e15f2e8e 4449f289c2
commit 675eff4f6d
6 changed files with 146 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/pipeline/test_engine.py
+102 -4
View File
@@ -10,13 +10,14 @@ from nose_parameterized import parameterized
from numpy import (
arange,
array,
concatenate,
float32,
full,
log,
nan,
tile,
where,
zeros,
float32,
concatenate,
log,
)
from numpy.testing import assert_almost_equal
from pandas import (
@@ -95,6 +96,22 @@ class AssetID(CustomFactor):
out[:] = assets
class AssetIDPlusDay(CustomFactor):
window_length = 1
inputs = [USEquityPricing.close]
def compute(self, today, assets, out, close):
out[:] = assets + today.day
class OpenPrice(CustomFactor):
window_length = 1
inputs = [USEquityPricing.open]
def compute(self, today, assets, out, open):
out[:] = open
def assert_multi_index_is_product(testcase, index, *levels):
"""Assert that a MultiIndex contains the product of `*levels`."""
testcase.assertIsInstance(
@@ -157,7 +174,7 @@ class ConstantInputTestCase(TestCase):
USEquityPricing.close: 3,
USEquityPricing.high: 4,
}
self.asset_ids = [1, 2, 3]
self.asset_ids = [1, 2, 3, 4]
self.dates = date_range('2014-01', '2014-03', freq='D', tz='UTC')
self.loader = PrecomputedLoader(
constants=self.constants,
@@ -354,6 +371,87 @@ class ConstantInputTestCase(TestCase):
DataFrame(expected_avg, index=dates, columns=self.assets),
)
def test_masked_factor(self):
"""
Test that a Custom Factor computes the correct values when passed a
mask. The mask/filter should be applied prior to computing any values,
as opposed to computing the factor across the entire universe of
assets. Any assets that are filtered out should be filled with missing
values.
"""
loader = self.loader
dates = self.dates[5:8]
assets = self.assets
asset_ids = self.asset_ids
constants = self.constants
open = USEquityPricing.open
close = USEquityPricing.close
engine = SimplePipelineEngine(
lambda column: loader, self.dates, self.asset_finder,
)
factor1_value = constants[open]
factor2_value = 3.0 * (constants[open] - constants[close])
def create_expected_results(expected_value, mask):
expected_values = where(mask, expected_value, nan)
return DataFrame(expected_values, index=dates, columns=assets)
cascading_mask = AssetIDPlusDay() < (asset_ids[-1] + dates[0].day)
expected_cascading_mask_result = array(
[[True, True, True, False],
[True, True, False, False],
[True, False, False, False]],
dtype=bool,
)
alternating_mask = (AssetIDPlusDay() % 2).eq(0)
expected_alternating_mask_result = array(
[[False, True, False, True],
[True, False, True, False],
[False, True, False, True]],
dtype=bool,
)
masks = cascading_mask, alternating_mask
expected_mask_results = (
expected_cascading_mask_result,
expected_alternating_mask_result,
)
for mask, expected_mask in zip(masks, expected_mask_results):
# Test running a pipeline with a single masked factor.
columns = {'factor1': OpenPrice(mask=mask), 'mask': mask}
pipeline = Pipeline(columns=columns)
results = engine.run_pipeline(pipeline, dates[0], dates[-1])
mask_results = results['mask'].unstack()
check_arrays(mask_results.values, expected_mask)
factor1_results = results['factor1'].unstack()
factor1_expected = create_expected_results(factor1_value,
mask_results)
assert_frame_equal(factor1_results, factor1_expected)
# Test running a pipeline with a second factor. This ensures that
# adding another factor to the pipeline with a different window
# length does not cause any unexpected behavior, especially when
# both factors share the same mask.
columns['factor2'] = RollingSumDifference(mask=mask)
pipeline = Pipeline(columns=columns)
results = engine.run_pipeline(pipeline, dates[0], dates[-1])
mask_results = results['mask'].unstack()
check_arrays(mask_results.values, expected_mask)
factor1_results = results['factor1'].unstack()
factor2_results = results['factor2'].unstack()
factor1_expected = create_expected_results(factor1_value,
mask_results)
factor2_expected = create_expected_results(factor2_value,
mask_results)
assert_frame_equal(factor1_results, factor1_expected)
assert_frame_equal(factor2_results, factor2_expected)
def test_rolling_and_nonrolling(self):
open_ = USEquityPricing.open
close = USEquityPricing.close
zipline/pipeline/data/dataset.py
-1
View File
@@ -114,7 +114,6 @@ class BoundColumn(LoadableTerm):
The name of this column.
"""
mask = AssetExists()
extra_input_rows = 0
inputs = ()
def __new__(cls, dtype, missing_value, dataset, name):
zipline/pipeline/engine.py
+10 -3
View File
@@ -237,13 +237,20 @@ class SimplePipelineEngine(object):
assert shape[0] * shape[1] != 0, 'root mask cannot be empty'
return ret
def _mask_and_dates_for_term(self, term, workspace, graph, dates):
def _mask_and_dates_for_term(self, term, workspace, graph, all_dates):
"""
Load mask and mask row labels for term.
"""
mask = term.mask
offset = graph.extra_rows[mask] - graph.extra_rows[term]
return workspace[mask][offset:], dates[offset:]
mask_offset = graph.extra_rows[mask] - graph.extra_rows[term]
# This offset is computed against _root_mask_term because that is what
# determines the shape of the top-level dates array.
dates_offset = (
graph.extra_rows[self._root_mask_term] - graph.extra_rows[term]
)
return workspace[mask][mask_offset:], all_dates[dates_offset:]
@staticmethod
def _inputs_for_term(term, workspace, graph):
zipline/pipeline/graph.py
+8 -12
View File
@@ -104,9 +104,9 @@ class TermGraph(DiGraph):
zipline.pipeline.engine.SimplePipelineEngine._inputs_for_term
zipline.pipeline.engine.SimplePipelineEngine._mask_and_dates_for_term
"""
return {(term, dep): self.extra_rows[dep] - term.extra_input_rows
return {(term, dep): self.extra_rows[dep] - additional_extra_rows
for term in self
for dep in term.dependencies}
for dep, additional_extra_rows in term.dependencies.items()}
@lazyval
def extra_rows(self):
@@ -119,10 +119,9 @@ class TermGraph(DiGraph):
Notes
----
This value depends on the other terms in the graph that require `term`
**as an input**. This is not to be confused with
`term.extra_input_rows`, which is how many extra rows of `term`'s
inputs we need to load, and which is determined entirely by `Term`
itself.
**as an input**. This is not to be confused with `term.dependencies`,
which describes how many additional rows of `term`'s inputs we need to
load, and which is determined entirely by `Term` itself.
Example
-------
@@ -144,7 +143,7 @@ class TermGraph(DiGraph):
See Also
--------
zipline.pipeline.graph.TermGraph.offset
zipline.pipeline.term.Term.extra_input_rows
zipline.pipeline.term.Term.dependencies
"""
return {
term: attrs['extra_rows']
@@ -187,15 +186,12 @@ class TermGraph(DiGraph):
# Make sure we're going to compute at least `extra_rows` of `term`.
self._ensure_extra_rows(term, extra_rows)
# Number of extra rows we need to compute for this term's dependencies.
dependency_extra_rows = extra_rows + term.extra_input_rows
# Recursively add dependencies.
for dependency in term.dependencies:
for dependency, additional_extra_rows in term.dependencies.items():
self._add_to_graph(
dependency,
parents,
extra_rows=dependency_extra_rows,
extra_rows=extra_rows + additional_extra_rows,
)
self.add_edge(dependency, term)
zipline/pipeline/mixins.py
+10 -5
View File
@@ -70,6 +70,7 @@ class CustomTermMixin(object):
def __new__(cls,
inputs=NotSpecified,
window_length=NotSpecified,
mask=NotSpecified,
dtype=NotSpecified,
missing_value=NotSpecified,
**kwargs):
@@ -88,6 +89,7 @@ class CustomTermMixin(object):
cls,
inputs=inputs,
window_length=window_length,
mask=mask,
dtype=dtype,
missing_value=missing_value,
**kwargs
@@ -104,7 +106,6 @@ class CustomTermMixin(object):
Call the user's `compute` function on each window with a pre-built
output array.
"""
# TODO: Make mask available to user's `compute`.
compute = self.compute
missing_value = self.missing_value
params = self.params
@@ -113,14 +114,18 @@ class CustomTermMixin(object):
# TODO: Consider pre-filtering columns that are all-nan at each
# time-step?
for idx, date in enumerate(dates):
col_mask = mask[idx]
masked_out = out[idx][col_mask]
masked_assets = assets[col_mask]
compute(
date,
assets,
out[idx],
*(next(w) for w in windows),
masked_assets,
masked_out,
*(next(w)[:, col_mask] for w in windows),
**params
)
out[~mask] = missing_value
out[idx][col_mask] = masked_out
return out
def short_repr(self):
zipline/pipeline/term.py
+16 -9
View File
@@ -294,13 +294,13 @@ class Term(with_metaclass(ABCMeta, object)):
"""
raise NotImplementedError('mask')
@lazyval
@abstractproperty
def dependencies(self):
"""
A tuple containing all terms that must be computed before this term can
be loaded or computed.
A dictionary mapping terms that must be computed before `self` to the
number of extra rows needed for those terms.
"""
return self.inputs + (self.mask,)
raise NotImplementedError('dependencies')
class AssetExists(Term):
@@ -319,9 +319,8 @@ class AssetExists(Term):
"""
dtype = bool_dtype
dataset = None
extra_input_rows = 0
inputs = ()
dependencies = ()
dependencies = {}
mask = None
windowed = False
@@ -335,9 +334,12 @@ class LoadableTerm(Term):
This is the base class for :class:`zipline.pipeline.data.BoundColumn`.
"""
inputs = ()
windowed = False
@lazyval
def dependencies(self):
return {self.mask: 0}
class ComputableTerm(Term):
"""
@@ -442,12 +444,17 @@ class ComputableTerm(Term):
)
@lazyval
def extra_input_rows(self):
def dependencies(self):
"""
The number of extra rows needed for each of our inputs to compute this
term.
"""
return max(0, self.window_length - 1)
extra_input_rows = max(0, self.window_length - 1)
out = {}
for term in self.inputs:
out[term] = extra_input_rows
out[self.mask] = 0
return out
def __repr__(self):
return (