From 5bae74addac5cbb3768f1e91d1ee148b83e0bbae Mon Sep 17 00:00:00 2001 From: dmichalowicz Date: Tue, 29 Mar 2016 17:16:58 -0400 Subject: [PATCH 1/4] 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): From 88eeb3689d59c072686df3c17ad0e0884980249a Mon Sep 17 00:00:00 2001 From: Scott Sanderson Date: Wed, 6 Apr 2016 13:25:56 -0400 Subject: [PATCH 2/4] MAINT: Make dependencies a dict. --- zipline/pipeline/data/dataset.py | 1 - zipline/pipeline/graph.py | 35 ++++++++------------------------ zipline/pipeline/term.py | 25 +++++++++++++++-------- 3 files changed, 24 insertions(+), 37 deletions(-) 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/graph.py b/zipline/pipeline/graph.py index 5f815065..6290c937 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,30 +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 - - 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 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/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 ( From 8db59b387bbf3ce72c20fe4dd92e6a866c5fc3fa Mon Sep 17 00:00:00 2001 From: dmichalowicz Date: Wed, 6 Apr 2016 17:20:52 -0400 Subject: [PATCH 3/4] TST: Overhaul test case --- tests/pipeline/test_engine.py | 102 +++++++++++++++++++++------------- zipline/pipeline/engine.py | 8 ++- zipline/pipeline/graph.py | 2 +- 3 files changed, 70 insertions(+), 42 deletions(-) diff --git a/tests/pipeline/test_engine.py b/tests/pipeline/test_engine.py index 9dff54fa..871c7a80 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 ( @@ -62,7 +63,6 @@ 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, @@ -96,6 +96,14 @@ 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] @@ -166,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, @@ -372,52 +380,68 @@ class ConstantInputTestCase(TestCase): values. """ loader = self.loader - dates = self.dates[5:10] + dates = self.dates[5:8] assets = self.assets asset_ids = self.asset_ids constants = self.constants - num_dates = len(dates) open = USEquityPricing.open + close = USEquityPricing.close 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) + factor1_value = constants[open] + factor2_value = 3.0 * (constants[open] - constants[close]) - for asset_id in asset_ids: - mask = AssetID() <= asset_id - factor1 = OpenPrice(mask=mask) + def create_expected_results(expected_value, mask): + expected_values = where(mask, expected_value, nan) + return DataFrame(expected_values, index=dates, columns=assets) - # 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() + # Produce a mask that looks like: + # + # Equity(0 [A]) Equity(1 [B]) Equity(2 [C]) Equity(3 [D]) + # Day 1 True True True False + # Day 2 True True False False + # Day 3 True False False False + # + cascading_mask = AssetIDPlusDay() < (asset_ids[-1] + dates[0].day) - expected = { - asset: values if asset.sid <= asset_id else missing_values - for asset in assets - } + # And another one that looks like: + # + # Equity(0 [A]) Equity(1 [B]) Equity(2 [C]) Equity(3 [D]) + # Day 1 False True False True + # Day 2 True False True False + # Day 3 False True False True + # + alternating_mask = (AssetIDPlusDay() % 2).eq(0) - assert_frame_equal( - factor1_results, - DataFrame(expected, index=dates, columns=assets), - ) + for mask in (cascading_mask, alternating_mask): + # 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() + 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() + 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 diff --git a/zipline/pipeline/engine.py b/zipline/pipeline/engine.py index 39a4776c..8a751bc5 100644 --- a/zipline/pipeline/engine.py +++ b/zipline/pipeline/engine.py @@ -237,16 +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 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:], dates[dates_offset:] + + 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 6290c937..529e2459 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 ComputableTerm, LoadableTerm +from .term import LoadableTerm class CyclicDependency(Exception): From 4449f289c25129d698ef73f03ca53457440253c9 Mon Sep 17 00:00:00 2001 From: Scott Sanderson Date: Thu, 7 Apr 2016 17:04:14 -0400 Subject: [PATCH 4/4] TEST: Test that the mask is what we expect. --- tests/pipeline/test_engine.py | 39 +++++++++++++++++++++-------------- 1 file changed, 24 insertions(+), 15 deletions(-) diff --git a/tests/pipeline/test_engine.py b/tests/pipeline/test_engine.py index 871c7a80..0f3be397 100644 --- a/tests/pipeline/test_engine.py +++ b/tests/pipeline/test_engine.py @@ -397,30 +397,36 @@ class ConstantInputTestCase(TestCase): expected_values = where(mask, expected_value, nan) return DataFrame(expected_values, index=dates, columns=assets) - # Produce a mask that looks like: - # - # Equity(0 [A]) Equity(1 [B]) Equity(2 [C]) Equity(3 [D]) - # Day 1 True True True False - # Day 2 True True False False - # Day 3 True False False False - # 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, + ) - # And another one that looks like: - # - # Equity(0 [A]) Equity(1 [B]) Equity(2 [C]) Equity(3 [D]) - # Day 1 False True False True - # Day 2 True False True False - # Day 3 False True False True - # 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, + ) - for mask in (cascading_mask, alternating_mask): + 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) @@ -433,7 +439,10 @@ class ConstantInputTestCase(TestCase): 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,