From 5bae74addac5cbb3768f1e91d1ee148b83e0bbae Mon Sep 17 00:00:00 2001 From: dmichalowicz Date: Tue, 29 Mar 2016 17:16:58 -0400 Subject: [PATCH] ENH: Allow passing a mask when creating a factor --- tests/pipeline/test_engine.py | 65 +++++++++++++++++++++++++++++++++++ zipline/pipeline/engine.py | 7 ++-- zipline/pipeline/graph.py | 19 ++++++++-- zipline/pipeline/mixins.py | 15 +++++--- 4 files changed, 97 insertions(+), 9 deletions(-) diff --git a/tests/pipeline/test_engine.py b/tests/pipeline/test_engine.py index 339f2f9b..9dff54fa 100644 --- a/tests/pipeline/test_engine.py +++ b/tests/pipeline/test_engine.py @@ -62,6 +62,7 @@ from zipline.pipeline.factors import ( MaxDrawdown, SimpleMovingAverage, ) +from zipline.pipeline.term import NotSpecified from zipline.testing import ( make_rotating_equity_info, make_simple_equity_info, @@ -95,6 +96,14 @@ class AssetID(CustomFactor): out[:] = assets +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( @@ -354,6 +363,62 @@ 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:10] + assets = self.assets + asset_ids = self.asset_ids + constants = self.constants + num_dates = len(dates) + open = USEquityPricing.open + engine = SimplePipelineEngine( + lambda column: loader, self.dates, self.asset_finder, + ) + + # These are the expected values for the OpenPrice factor. If we pass + # OpenPrice a mask, any assets that are filtered out should have all + # NaN values. Otherwise, we expect its computed values to be the + # asset's open price. + values = array([constants[open]] * num_dates, dtype=float) + missing_values = array([nan] * num_dates) + + for asset_id in asset_ids: + mask = AssetID() <= asset_id + factor1 = OpenPrice(mask=mask) + + # Test running our pipeline both with and without a second factor. + # We do not explicitly test the resulting values of the second + # factor; we just want to implicitly ensure that the addition of + # another factor to the pipeline term graph does not cause any + # unexpected exceptions when calling `run_pipeline`. + for factor2 in (None, + RollingSumDifference(mask=NotSpecified), + RollingSumDifference(mask=mask)): + if factor2 is None: + columns = {'factor1': factor1} + else: + columns = {'factor1': factor1, 'factor2': factor2} + pipeline = Pipeline(columns=columns) + results = engine.run_pipeline(pipeline, dates[0], dates[-1]) + factor1_results = results['factor1'].unstack() + + expected = { + asset: values if asset.sid <= asset_id else missing_values + for asset in assets + } + + assert_frame_equal( + factor1_results, + DataFrame(expected, index=dates, columns=assets), + ) + def test_rolling_and_nonrolling(self): open_ = USEquityPricing.open close = USEquityPricing.close diff --git a/zipline/pipeline/engine.py b/zipline/pipeline/engine.py index c871530d..39a4776c 100644 --- a/zipline/pipeline/engine.py +++ b/zipline/pipeline/engine.py @@ -242,8 +242,11 @@ class SimplePipelineEngine(object): 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] + dates_offset = ( + graph.extra_rows[self._root_mask_term] - graph.extra_rows[term] + ) + return workspace[mask][mask_offset:], 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..5f815065 100644 --- a/zipline/pipeline/graph.py +++ b/zipline/pipeline/graph.py @@ -9,7 +9,7 @@ from six import itervalues, iteritems from zipline.utils.memoize import lazyval from zipline.pipeline.visualize import display_graph -from .term import LoadableTerm +from .term import ComputableTerm, LoadableTerm class CyclicDependency(Exception): @@ -190,8 +190,23 @@ class TermGraph(DiGraph): # Number of extra rows we need to compute for this term's dependencies. dependency_extra_rows = extra_rows + term.extra_input_rows + if isinstance(term, ComputableTerm): + # For computable terms, we want to manually add the term's mask to + # the graph with zero extra rows. A computable term does not + # directly require its mask to have any extra rows. Only loadable + # terms should dictate how many extra rows a mask should compute. + self._add_to_graph( + term.mask, + parents, + extra_rows=0, + ) + self.add_edge(term.mask, term) + dependencies = term.inputs + else: + dependencies = term.dependencies + # Recursively add dependencies. - for dependency in term.dependencies: + for dependency in dependencies: self._add_to_graph( dependency, parents, 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):