diff --git a/tests/pipeline/test_engine.py b/tests/pipeline/test_engine.py index 6a7037eb..23c843b6 100644 --- a/tests/pipeline/test_engine.py +++ b/tests/pipeline/test_engine.py @@ -33,7 +33,7 @@ from pandas import ( from pandas.compat.chainmap import ChainMap from pandas.util.testing import assert_frame_equal from six import iteritems, itervalues -from toolz import merge +from toolz import merge, assoc from zipline.assets.synthetic import make_rotating_equity_info from zipline.errors import NoFurtherDataError @@ -163,14 +163,14 @@ class RollingSumSum(CustomFactor): out[:] = sum(inputs).sum(axis=0) -class ConstantInputTestCase(WithTradingEnvironment, ZiplineTestCase): +class WithConstantInputs(WithTradingEnvironment): asset_ids = ASSET_FINDER_EQUITY_SIDS = 1, 2, 3, 4 START_DATE = Timestamp('2014-01-01', tz='utc') END_DATE = Timestamp('2014-03-01', tz='utc') @classmethod def init_class_fixtures(cls): - super(ConstantInputTestCase, cls).init_class_fixtures() + super(WithConstantInputs, cls).init_class_fixtures() cls.constants = { # Every day, assume every stock starts at 2, goes down to 1, # goes up to 4, and finishes at 3. @@ -192,6 +192,8 @@ class ConstantInputTestCase(WithTradingEnvironment, ZiplineTestCase): ) cls.assets = cls.asset_finder.retrieve_all(cls.asset_ids) + +class ConstantInputTestCase(WithConstantInputs, ZiplineTestCase): def test_bad_dates(self): loader = self.loader engine = SimplePipelineEngine( @@ -1315,3 +1317,51 @@ class StringColumnTestCase(WithSeededRandomPipelineEngine, columns=self.asset_finder.retrieve_all(self.asset_finder.sids), ) assert_frame_equal(result.c.unstack(), expected_final_result) + + +class PopulateInitialWorkspaceTestCase(WithConstantInputs, ZiplineTestCase): + def make_engine(self, populate_initial_workspace): + return SimplePipelineEngine( + lambda column: self.loader, + self.dates, + self.asset_finder, + populate_initial_workspace=populate_initial_workspace, + ) + + def test_populate_default_workspace(self): + column = USEquityPricing.low + base_term = column.latest + term = base_term + 1 + column_value = self.constants[column] + precomputed_value = -column_value + + def populate_initial_workspace(initial_workspace, + root_mask_term, + execution_plan, + dates, + assets): + return assoc( + initial_workspace, + term, + full((len(dates), len(assets)), precomputed_value), + ) + + # I resisted the urge to use ``make_engine`` as a decorator here + # because Scott would have yelled at me. + engine = self.make_engine(populate_initial_workspace) + + results = engine.run_pipeline( + Pipeline({ + 'term-in-initial-workspace': term, + 'term-not-in-initial-workspace': base_term, + }), + self.dates[0], + self.dates[-1], + ) + + self.assertTrue( + (results['term-in-initial-workspace'] == precomputed_value).all(), + ) + self.assertTrue( + (results['term-not-in-initial-workspace'] == column_value).all(), + ) diff --git a/zipline/pipeline/engine.py b/zipline/pipeline/engine.py index 4bb6f6b0..e2a098e0 100644 --- a/zipline/pipeline/engine.py +++ b/zipline/pipeline/engine.py @@ -81,6 +81,14 @@ class ExplodingPipelineEngine(PipelineEngine): ) +def _default_populate_initial_workspace(initial_workspace, + root_mask_term, + execution_plan, + dates, + assets): + return initial_workspace + + class SimplePipelineEngine(object): """ PipelineEngine class that computes each term independently. @@ -96,6 +104,12 @@ class SimplePipelineEngine(object): asset_finder : zipline.assets.AssetFinder An AssetFinder instance. We depend on the AssetFinder to determine which assets are in the top-level universe at any point in time. + populate_initial_workspace : callable, optional + A function which will be used to populate the initial workspace when + computing a pipeline. This function will be passed the + initial_workspace, the root mask term, the execution_plan, the dates + being computed for, and the assets requested and should return a new + dictionary which will be used as the initial_workspace. """ __slots__ = ( '_get_loader', @@ -103,10 +117,15 @@ class SimplePipelineEngine(object): '_finder', '_root_mask_term', '_root_mask_dates_term', + '_populate_initial_workspace', '__weakref__', ) - def __init__(self, get_loader, calendar, asset_finder): + def __init__(self, + get_loader, + calendar, + asset_finder, + populate_initial_workspace=None): self._get_loader = get_loader self._calendar = calendar self._finder = asset_finder @@ -114,6 +133,10 @@ class SimplePipelineEngine(object): self._root_mask_term = AssetExists() self._root_mask_dates_term = InputDates() + self._populate_initial_workspace = ( + populate_initial_workspace or _default_populate_initial_workspace + ) + def run_pipeline(self, pipeline, start_date, end_date): """ Compute a pipeline. @@ -179,14 +202,22 @@ class SimplePipelineEngine(object): root_mask = self._compute_root_mask(start_date, end_date, extra_rows) dates, assets, root_mask_values = explode(root_mask) + initial_workspace = self._populate_initial_workspace( + { + self._root_mask_term: root_mask_values, + self._root_mask_dates_term: as_column(dates.values) + }, + self._root_mask_term, + graph, + dates, + assets, + ) + results = self.compute_chunk( graph, dates, assets, - initial_workspace={ - self._root_mask_term: root_mask_values, - self._root_mask_dates_term: as_column(dates.values) - }, + initial_workspace, ) return self._to_narrow( @@ -255,21 +286,6 @@ 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, 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:], all_dates[dates_offset:] - @staticmethod def _inputs_for_term(term, workspace, graph): """ @@ -356,8 +372,11 @@ class SimplePipelineEngine(object): # Asset labels are always the same, but date labels vary by how # many extra rows are needed. - mask, mask_dates = self._mask_and_dates_for_term( - term, workspace, graph, dates + mask, mask_dates = graph.mask_and_dates_for_term( + term, + self._root_mask_term, + workspace, + dates, ) if isinstance(term, LoadableTerm): diff --git a/zipline/pipeline/graph.py b/zipline/pipeline/graph.py index 667c3be9..cef8e7a0 100644 --- a/zipline/pipeline/graph.py +++ b/zipline/pipeline/graph.py @@ -375,3 +375,40 @@ class ExecutionPlan(TermGraph): """ attrs = self.node[term] attrs['extra_rows'] = max(N, attrs.get('extra_rows', 0)) + + def mask_and_dates_for_term(self, + term, + root_mask_term, + workspace, + all_dates): + """ + Load mask and mask row labels for term. + + Parameters + ---------- + term : Term + The term to load the mask and labels for. + root_mask_term : Term + The term that represents the root asset exists mask. + workspace : dict[Term, any] + The values that have been computed for each term. + all_dates : pd.DatetimeIndex + All of the dates that are being computed for in the pipeline. + + Returns + ------- + mask : np.ndarray + The correct mask for this term. + dates : np.ndarray + The slice of dates for this term. + """ + mask = term.mask + mask_offset = self.extra_rows[mask] - self.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 = ( + self.extra_rows[root_mask_term] - self.extra_rows[term] + ) + + return workspace[mask][mask_offset:], all_dates[dates_offset:]