diff --git a/tests/pipeline/test_engine.py b/tests/pipeline/test_engine.py index 339f2f9b..0f3be397 100644 --- a/tests/pipeline/test_engine.py +++ b/tests/pipeline/test_engine.py @@ -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 diff --git a/zipline/pipeline/data/dataset.py b/zipline/pipeline/data/dataset.py index da236541..b5980d83 100644 --- a/zipline/pipeline/data/dataset.py +++ b/zipline/pipeline/data/dataset.py @@ -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): diff --git a/zipline/pipeline/engine.py b/zipline/pipeline/engine.py index c871530d..8a751bc5 100644 --- a/zipline/pipeline/engine.py +++ b/zipline/pipeline/engine.py @@ -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): diff --git a/zipline/pipeline/graph.py b/zipline/pipeline/graph.py index 6bcbe8c0..529e2459 100644 --- a/zipline/pipeline/graph.py +++ b/zipline/pipeline/graph.py @@ -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) diff --git a/zipline/pipeline/mixins.py b/zipline/pipeline/mixins.py index 4e8a27b5..da9297b8 100644 --- a/zipline/pipeline/mixins.py +++ b/zipline/pipeline/mixins.py @@ -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): diff --git a/zipline/pipeline/term.py b/zipline/pipeline/term.py index 07b0ef8f..16faf227 100644 --- a/zipline/pipeline/term.py +++ b/zipline/pipeline/term.py @@ -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 (