ENH: provide a hook for prepopulating the initial workspace

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(),
+        )

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):

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:]