ENH: Add SpecificAssets filter.

Adds a filter that matches a set of assets.  Mainly useful for testing
and debugging.
This commit is contained in:
Scott Sanderson
2016-10-07 18:11:01 -04:00
parent f92b6afb31
commit c9ebf9f45d
3 changed files with 66 additions and 4 deletions
+42 -3
View File
@@ -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)