Merge pull request #1530 from quantopian/add-specific-assets

ENH: Add `SpecificAssets` filter.

tests/pipeline/test_filter.py

@@ -26,11 +26,13 @@ from numpy import (
 from numpy.random import randn, seed as random_seed
 
 from zipline.errors import BadPercentileBounds
-from zipline.pipeline import Filter, Factor
+from zipline.pipeline import Filter, Factor, Pipeline
 from zipline.pipeline.classifiers import Classifier
 from zipline.pipeline.factors import CustomFactor
-from zipline.pipeline.filters import All, Any, AtLeastN
-from zipline.testing import parameter_space, permute_rows
+from zipline.pipeline.filters import All, Any, AtLeastN, SpecificAssets
+from zipline.testing import parameter_space, permute_rows, ZiplineTestCase
+from zipline.testing.fixtures import WithSeededRandomPipelineEngine
+from zipline.testing.predicates import assert_equal
 from zipline.utils.numpy_utils import float64_dtype, int64_dtype
 from .base import BasePipelineTestCase, with_default_shape
 
@@ -820,3 +822,40 @@ class FilterTestCase(BasePipelineTestCase):
             },
             mask=self.build_mask(permute(rot90(self.eye_mask(shape=shape)))),
         )
+
+
+class SpecificAssetsTestCase(WithSeededRandomPipelineEngine,
+                             ZiplineTestCase):
+
+    ASSET_FINDER_EQUITY_SIDS = tuple(range(10))
+
+    def test_specific_assets(self):
+        assets = self.asset_finder.retrieve_all(self.ASSET_FINDER_EQUITY_SIDS)
+
+        class SidFactor(CustomFactor):
+            """A factor that just returns each asset's sid."""
+            inputs = ()
+            window_length = 1
+
+            def compute(self, today, sids, out):
+                out[:] = sids
+
+        pipe = Pipeline(
+            columns={
+                'sid': SidFactor(),
+                'evens': SpecificAssets(assets[::2]),
+                'odds': SpecificAssets(assets[1::2]),
+                'first_five': SpecificAssets(assets[:5]),
+                'last_three': SpecificAssets(assets[-3:]),
+            },
+        )
+
+        start, end = self.trading_days[[-10, -1]]
+        results = self.run_pipeline(pipe, start, end).unstack()
+
+        sids = results.sid.astype(int64_dtype)
+
+        assert_equal(results.evens, ~(sids % 2).astype(bool))
+        assert_equal(results.odds, (sids % 2).astype(bool))
+        assert_equal(results.first_five, sids < 5)
+        assert_equal(results.last_three, sids >= 7)

zipline/pipeline/filters/__init__.py

@@ -8,6 +8,7 @@ from .filter import (
     NumExprFilter,
     PercentileFilter,
     SingleAsset,
+    SpecificAssets,
 )
 from .smoothing import All, Any, AtLeastN
 
@@ -24,4 +25,5 @@ __all__ = [
     'NumExprFilter',
     'PercentileFilter',
     'SingleAsset',
+    'SpecificAssets',
 ]

zipline/pipeline/filters/filter.py

@@ -494,3 +494,29 @@ class SingleAsset(Filter):
                 asset=self._asset, start_date=dates[0], end_date=dates[-1],
             )
         return out
+
+
+class SpecificAssets(Filter):
+    """
+    A Filter that computes True for a specific set of predetermined assets.
+
+    ``SpecificAssets`` is mostly useful for debugging or for interactively
+    computing pipeline terms for a fixed set of assets that are known ahead of
+    time.
+
+    Parameters
+    ----------
+    assets : iterable[Asset]
+        An iterable of assets for which to filter.
+    """
+    inputs = ()
+    window_length = 0
+    params = ('sids',)
+
+    def __new__(cls, assets):
+        sids = frozenset(asset.sid for asset in assets)
+        return super(SpecificAssets, cls).__new__(cls, sids=sids)
+
+    def _compute(self, arrays, dates, sids, mask):
+        my_columns = sids.isin(self.params['sids'])
+        return repeat_first_axis(my_columns, len(mask)) & mask