diff --git a/tests/modelling/base.py b/tests/modelling/base.py index e0e8490d..b1c1305d 100644 --- a/tests/modelling/base.py +++ b/tests/modelling/base.py @@ -10,7 +10,6 @@ from six import iteritems from zipline.finance.trading import TradingEnvironment from zipline.modelling.engine import SimpleFFCEngine -from zipline.modelling.graph import TermGraph from zipline.modelling.term import AssetExists from zipline.utils.pandas_utils import explode from zipline.utils.test_utils import make_simple_asset_info, ExplodingObject @@ -72,15 +71,15 @@ class BaseFFCTestCase(TestCase): """Default shape for methods that build test data.""" return self.__mask.shape - def run_terms(self, terms, initial_workspace, mask=None): + def run_graph(self, graph, initial_workspace, mask=None): """ - Compute the given terms, seeding the workspace of our FFCEngine with + Compute the given TermGraph, seeding the workspace of our engine with `initial_workspace`. Parameters ---------- - terms : dict - Mapping from termname -> term object. + graph : zipline.pipeline.graph.TermGraph + Graph to run. initial_workspace : dict Initial workspace to forward to SimpleFFCEngine.compute_chunk. mask : DataFrame, optional @@ -104,7 +103,7 @@ class BaseFFCTestCase(TestCase): dates, assets, mask_values = explode(mask) initial_workspace.setdefault(AssetExists(), mask_values) return engine.compute_chunk( - TermGraph(terms), + graph, dates, assets, initial_workspace, diff --git a/tests/modelling/test_engine.py b/tests/modelling/test_engine.py index c1b82b7f..153bf2e2 100644 --- a/tests/modelling/test_engine.py +++ b/tests/modelling/test_engine.py @@ -6,11 +6,12 @@ from unittest import TestCase from itertools import product from numpy import ( + array, full, nan, + tile, zeros, ) -from numpy.testing import assert_array_equal from pandas import ( DataFrame, date_range, @@ -45,6 +46,7 @@ from zipline.modelling.factor.technical import ( MaxDrawdown, SimpleMovingAverage, ) +from zipline.modelling.pipeline import Pipeline from zipline.utils.memoize import lazyval from zipline.utils.test_utils import ( make_rotating_asset_info, @@ -62,6 +64,22 @@ class RollingSumDifference(CustomFactor): out[:] = (open - close).sum(axis=0) +class AssetID(CustomFactor): + """ + CustomFactor that returns the AssetID of each asset. + + Useful for providing a Factor that produces a different value for each + asset. + """ + window_length = 1 + # HACK: We currently decide whether to load or compute a Term based on the + # length of its inputs. This means we have to provide a dummy input. + inputs = [USEquityPricing.close] + + def compute(self, today, assets, out, close): + out[:] = assets + + def assert_multi_index_is_product(testcase, index, *levels): """Assert that a MultiIndex contains the product of `*levels`.""" testcase.assertIsInstance( @@ -102,11 +120,36 @@ class ConstantInputTestCase(TestCase): loader = self.loader engine = SimpleFFCEngine(loader, self.dates, self.asset_finder) + p = Pipeline('test') + msg = "start_date must be before end_date .*" with self.assertRaisesRegexp(ValueError, msg): - engine.factor_matrix({}, self.dates[2], self.dates[1]) + engine.run_pipeline(p, self.dates[2], self.dates[1]) with self.assertRaisesRegexp(ValueError, msg): - engine.factor_matrix({}, self.dates[2], self.dates[2]) + engine.run_pipeline(p, self.dates[2], self.dates[2]) + + def test_screen(self): + loader = self.loader + finder = self.asset_finder + assets = array(self.assets) + engine = SimpleFFCEngine(loader, self.dates, self.asset_finder) + num_dates = 5 + dates = self.dates[10:10 + num_dates] + + factor = AssetID() + for asset in assets: + p = Pipeline('test', columns={'f': factor}, screen=factor <= asset) + result = engine.run_pipeline(p, dates[0], dates[-1]) + + expected_sids = assets[assets <= asset] + expected_assets = finder.retrieve_all(expected_sids) + expected_result = DataFrame( + index=MultiIndex.from_product([dates, expected_assets]), + data=tile(expected_sids.astype(float), [len(dates)]), + columns=['f'], + ) + + assert_frame_equal(result, expected_result) def test_single_factor(self): loader = self.loader @@ -117,17 +160,29 @@ class ConstantInputTestCase(TestCase): dates = self.dates[10:10 + num_dates] factor = RollingSumDifference() + expected_result = -factor.window_length - result = engine.factor_matrix({'f': factor}, dates[0], dates[-1]) - self.assertEqual(set(result.columns), {'f'}) - assert_multi_index_is_product( - self, result.index, dates, finder.retrieve_all(assets) - ) + # Since every asset will pass the screen, these should be equivalent. + pipelines = [ + Pipeline('test', columns={'f': factor}), + Pipeline( + 'test', + columns={'f': factor}, + screen=factor.eq(expected_result), + ), + ] - assert_array_equal( - result['f'].unstack().values, - full(result_shape, -factor.window_length), - ) + for p in pipelines: + result = engine.run_pipeline(p, dates[0], dates[-1]) + self.assertEqual(set(result.columns), {'f'}) + assert_multi_index_is_product( + self, result.index, dates, finder.retrieve_all(assets) + ) + + check_arrays( + result['f'].unstack().values, + full(result_shape, expected_result), + ) def test_multiple_rolling_factors(self): @@ -145,27 +200,32 @@ class ConstantInputTestCase(TestCase): inputs=[USEquityPricing.open, USEquityPricing.high], ) - results = engine.factor_matrix( - {'short': short_factor, 'long': long_factor, 'high': high_factor}, - dates[0], - dates[-1], + pipeline = Pipeline( + 'test', + columns={ + 'short': short_factor, + 'long': long_factor, + 'high': high_factor, + } ) + results = engine.run_pipeline(pipeline, dates[0], dates[-1]) + self.assertEqual(set(results.columns), {'short', 'high', 'long'}) assert_multi_index_is_product( self, results.index, dates, finder.retrieve_all(assets) ) # row-wise sum over an array whose values are all (1 - 2) - assert_array_equal( + check_arrays( results['short'].unstack().values, full(shape, -short_factor.window_length), ) - assert_array_equal( + check_arrays( results['long'].unstack().values, full(shape, -long_factor.window_length), ) # row-wise sum over an array whose values are all (1 - 3) - assert_array_equal( + check_arrays( results['high'].unstack().values, full(shape, -2 * high_factor.window_length), ) @@ -183,12 +243,15 @@ class ConstantInputTestCase(TestCase): open_minus_close = RollingSumDifference(inputs=[open, close]) avg = (high_minus_low + open_minus_close) / 2 - results = engine.factor_matrix( - { - 'high_low': high_minus_low, - 'open_close': open_minus_close, - 'avg': avg, - }, + results = engine.run_pipeline( + Pipeline( + 'test', + columns={ + 'high_low': high_minus_low, + 'open_close': open_minus_close, + 'avg': avg, + }, + ), dates[0], dates[-1], ) @@ -311,8 +374,11 @@ class FrameInputTestCase(TestCase): ) bounds = product_upper_triangle(range(window_length, len(dates))) for start, stop in bounds: - results = engine.factor_matrix( - {'low': low_mavg, 'high': high_mavg}, + results = engine.run_pipeline( + Pipeline( + 'test', + columns={'low': low_mavg, 'high': high_mavg} + ), dates[start], dates[stop], ) @@ -424,8 +490,8 @@ class SyntheticBcolzTestCase(TestCase): window_length=window_length, ) - results = engine.factor_matrix( - {'sma': SMA}, + results = engine.run_pipeline( + Pipeline('test', columns={'sma': SMA}), dates_to_test[0], dates_to_test[-1], ) @@ -476,8 +542,8 @@ class SyntheticBcolzTestCase(TestCase): window_length=window_length, ) - results = engine.factor_matrix( - {'drawdown': drawdown}, + results = engine.run_pipeline( + Pipeline('test', columns={'drawdown': drawdown}), dates_to_test[0], dates_to_test[-1], ) @@ -529,13 +595,16 @@ class MultiColumnLoaderTestCase(TestCase): sumdiff = RollingSumDifference() - result = engine.factor_matrix( - { - 'sumdiff': sumdiff, - 'open': open_.latest, - 'close': close.latest, - 'volume': volume.latest, - }, + result = engine.run_pipeline( + Pipeline( + 'test', + columns={ + 'sumdiff': sumdiff, + 'open': open_.latest, + 'close': close.latest, + 'volume': volume.latest, + }, + ), dates_to_test[0], dates_to_test[-1] ) diff --git a/tests/modelling/test_factor.py b/tests/modelling/test_factor.py index d0c0312c..a7d4c39f 100644 --- a/tests/modelling/test_factor.py +++ b/tests/modelling/test_factor.py @@ -5,6 +5,7 @@ from numpy import array, eye, nan, ones from zipline.errors import UnknownRankMethod from zipline.modelling.factor import Factor from zipline.modelling.filter import Filter +from zipline.modelling.graph import TermGraph from zipline.utils.test_utils import check_arrays from .base import BaseFFCTestCase @@ -68,8 +69,9 @@ class FactorTestCase(BaseFFCTestCase): } def check(terms): - results = self.run_terms( - terms, + graph = TermGraph(terms) + results = self.run_graph( + graph, initial_workspace={self.f: data}, mask=self.build_mask(ones((5, 5))), ) @@ -123,8 +125,9 @@ class FactorTestCase(BaseFFCTestCase): } def check(terms): - results = self.run_terms( - terms, + graph = TermGraph(terms) + results = self.run_graph( + graph, initial_workspace={self.f: data}, mask=self.build_mask(ones((5, 5))), ) @@ -148,12 +151,14 @@ class FactorTestCase(BaseFFCTestCase): mask_data = ~eye(5, dtype=bool) initial_workspace = {self.f: data, Mask(): mask_data} - terms = { - "ascending_nomask": self.f.rank(ascending=True), - "ascending_mask": self.f.rank(ascending=True, mask=Mask()), - "descending_nomask": self.f.rank(ascending=False), - "descending_mask": self.f.rank(ascending=False, mask=Mask()), - } + graph = TermGraph( + { + "ascending_nomask": self.f.rank(ascending=True), + "ascending_mask": self.f.rank(ascending=True, mask=Mask()), + "descending_nomask": self.f.rank(ascending=False), + "descending_mask": self.f.rank(ascending=False, mask=Mask()), + } + ) expected = { "ascending_nomask": array([[1., 3., 4., 5., 2.], @@ -180,8 +185,8 @@ class FactorTestCase(BaseFFCTestCase): [4., 3., 2., 1., nan]]), } - results = self.run_terms( - terms, + results = self.run_graph( + graph, initial_workspace, mask=self.build_mask(ones((5, 5))), ) diff --git a/tests/modelling/test_filter.py b/tests/modelling/test_filter.py index 4349754a..3290a3a5 100644 --- a/tests/modelling/test_filter.py +++ b/tests/modelling/test_filter.py @@ -22,6 +22,7 @@ from numpy.random import randn, seed as random_seed from zipline.errors import BadPercentileBounds from zipline.modelling.filter import Filter from zipline.modelling.factor import Factor +from zipline.modelling.graph import TermGraph from zipline.utils.test_utils import check_arrays from .base import BaseFFCTestCase, with_default_shape @@ -108,7 +109,7 @@ class FilterTestCase(BaseFFCTestCase): nan_data[:, 0] = nan mask = Mask() - initial_workspace = {self.f: data, mask: mask_data} + workspace = {self.f: data, mask: mask_data} methods = ['top', 'bottom'] counts = 2, 3, 10 @@ -127,7 +128,7 @@ class FilterTestCase(BaseFFCTestCase): term = getattr(self.f, method)(**kwargs) terms[termname(method, count, masked)] = term - results = self.run_terms(terms, initial_workspace=initial_workspace) + results = self.run_graph(TermGraph(terms), initial_workspace=workspace) def expected_result(method, count, masked): # Ranking with a mask is equivalent to ranking with nans applied on @@ -155,8 +156,10 @@ class FilterTestCase(BaseFFCTestCase): def test_bottom(self): counts = 2, 3, 10 data = self.randn_data(seed=5) # Arbitrary seed choice. - results = self.run_terms( - terms={'bottom_' + str(c): self.f.bottom(c) for c in counts}, + results = self.run_graph( + TermGraph( + {'bottom_' + str(c): self.f.bottom(c) for c in counts} + ), initial_workspace={self.f: data}, ) for c in counts: @@ -179,15 +182,17 @@ class FilterTestCase(BaseFFCTestCase): filter_names = ['pct_' + str(q) for q in quintiles] iter_quintiles = zip(filter_names, quintiles) - terms = { - name: self.f.percentile_between(q * 20.0, (q + 1) * 20.0) - for name, q in zip(filter_names, quintiles) - } + graph = TermGraph( + { + name: self.f.percentile_between(q * 20.0, (q + 1) * 20.0) + for name, q in zip(filter_names, quintiles) + } + ) # Test with 5 columns and no NaNs. eye5 = eye(5, dtype=float64) - results = self.run_terms( - terms, + results = self.run_graph( + graph, initial_workspace={self.f: eye5}, mask=self.build_mask(ones((5, 5))), ) @@ -211,8 +216,8 @@ class FilterTestCase(BaseFFCTestCase): [1, 1, 1, 0, 1, 1], [1, 1, 1, 1, 0, 1]], dtype=bool) - results = self.run_terms( - terms, + results = self.run_graph( + graph, initial_workspace={self.f: eye6}, mask=self.build_mask(mask) ) @@ -231,8 +236,8 @@ class FilterTestCase(BaseFFCTestCase): # In particular, the NaNs should never pass any filters. eye6_withnans = eye6.copy() putmask(eye6_withnans, ~mask, nan) - results = self.run_terms( - terms, + results = self.run_graph( + graph, initial_workspace={self.f: eye6}, mask=self.build_mask(mask) ) @@ -258,12 +263,14 @@ class FilterTestCase(BaseFFCTestCase): quartiles = range(4) filter_names = ['pct_' + str(q) for q in quartiles] - terms = { - name: self.f.percentile_between(q * 25.0, (q + 1) * 25.0) - for name, q in zip(filter_names, quartiles) - } - results = self.run_terms( - terms, + graph = TermGraph( + { + name: self.f.percentile_between(q * 25.0, (q + 1) * 25.0) + for name, q in zip(filter_names, quartiles) + } + ) + results = self.run_graph( + graph, initial_workspace={self.f: data}, mask=self.build_mask(ones((5, 5))), ) @@ -287,14 +294,16 @@ class FilterTestCase(BaseFFCTestCase): without_mask = self.g.percentile_between(80, 100) with_mask = self.g.percentile_between(80, 100, mask=custom_mask) - terms = { - 'custom_mask': custom_mask, - 'without': without_mask, - 'with': with_mask, - } + graph = TermGraph( + { + 'custom_mask': custom_mask, + 'without': without_mask, + 'with': with_mask, + } + ) - results = self.run_terms( - terms, + results = self.run_graph( + graph, initial_workspace={self.f: f_input, self.g: g_input}, mask=initial_mask, ) diff --git a/tests/modelling/test_modelling_algo.py b/tests/modelling/test_modelling_algo.py index db8750fc..1892ba10 100644 --- a/tests/modelling/test_modelling_algo.py +++ b/tests/modelling/test_modelling_algo.py @@ -29,7 +29,8 @@ from testfixtures import TempDirectory from zipline.algorithm import TradingAlgorithm from zipline.api import ( - add_factor, + attach_pipeline, + drain_pipeline, get_datetime, ) from zipline.data.equities import USEquityPricing @@ -41,9 +42,15 @@ from zipline.data.ffc.loaders.us_equity_pricing import ( SQLiteAdjustmentWriter, USEquityPricingLoader, ) +from zipline.errors import ( + AttachPipelineAfterInitialize, + DrainPipelineDuringInitialize, + NoSuchPipeline, +) from zipline.finance import trading from zipline.modelling.factor.technical import VWAP +from zipline.modelling.pipeline import Pipeline from zipline.utils.test_utils import ( make_simple_asset_info, str_to_seconds, @@ -157,26 +164,127 @@ class ClosesOnly(TestCase): def exists(self, date, asset): return asset.start_date <= date <= asset.end_date + def test_attach_pipeline_after_initialize(self): + """ + Assert that calling attach_pipeline after initialize raises correctly. + """ + def initialize(context): + pass + + def late_attach(context, data): + attach_pipeline(Pipeline('test')) + raise AssertionError("Shouldn't make it past attach_pipeline!") + + algo = TradingAlgorithm( + initialize=initialize, + handle_data=late_attach, + data_frequency='daily', + ffc_loader=self.ffc_loader, + start=self.first_asset_start - trading_day, + end=self.last_asset_end + trading_day, + env=self.env, + ) + + with self.assertRaises(AttachPipelineAfterInitialize): + algo.run(source=self.closes) + + def barf(context, data): + raise AssertionError("Shouldn't make it past before_trading_start") + + algo = TradingAlgorithm( + initialize=initialize, + before_trading_start=late_attach, + handle_data=barf, + data_frequency='daily', + ffc_loader=self.ffc_loader, + start=self.first_asset_start - trading_day, + end=self.last_asset_end + trading_day, + env=self.env, + ) + + with self.assertRaises(AttachPipelineAfterInitialize): + algo.run(source=self.closes) + + def test_drain_pipeline_after_initialize(self): + """ + Assert that calling drain_pipeline after initialize raises correctly. + """ + def initialize(context): + attach_pipeline(Pipeline('test')) + drain_pipeline('test') + raise AssertionError("Shouldn't make it past drain_pipeline()") + + def handle_data(context, data): + raise AssertionError("Shouldn't make it past initialize!") + + def before_trading_start(context, data): + raise AssertionError("Shouldn't make it past initialize!") + + algo = TradingAlgorithm( + initialize=initialize, + handle_data=handle_data, + before_trading_start=before_trading_start, + data_frequency='daily', + ffc_loader=self.ffc_loader, + start=self.first_asset_start - trading_day, + end=self.last_asset_end + trading_day, + env=self.env, + ) + + with self.assertRaises(DrainPipelineDuringInitialize): + algo.run(source=self.closes) + + def test_drain_nonexistent_pipeline(self): + """ + Assert that calling add_pipeline after initialize raises appropriately. + """ + def initialize(context): + attach_pipeline(Pipeline('test')) + + def handle_data(context, data): + raise AssertionError("Shouldn't make it past before_trading_start") + + def before_trading_start(context, data): + drain_pipeline('not_test') + raise AssertionError("Shouldn't make it past drain_pipeline!") + + algo = TradingAlgorithm( + initialize=initialize, + handle_data=handle_data, + before_trading_start=before_trading_start, + data_frequency='daily', + ffc_loader=self.ffc_loader, + start=self.first_asset_start - trading_day, + end=self.last_asset_end + trading_day, + env=self.env, + ) + + with self.assertRaises(NoSuchPipeline): + algo.run(source=self.closes) + def test_assets_appear_on_correct_days(self): """ Assert that assets appear at correct times during a backtest, with correctly-adjusted close price values. """ def initialize(context): - add_factor(USEquityPricing.close.latest, 'close') + p = Pipeline('test') + p.add(USEquityPricing.close.latest, 'close') + + attach_pipeline(p) def handle_data(context, data): - factors = data.factors + results = drain_pipeline('test') date = get_datetime().normalize() for asset in self.assets: # Assets should appear iff they exist today and yesterday. exists_today = self.exists(date, asset) existed_yesterday = self.exists(date - trading_day, asset) if exists_today and existed_yesterday: - latest = factors.loc[asset, 'close'] + latest = results.loc[asset, 'close'] self.assertEqual(latest, self.expected_close(date, asset)) else: - self.assertNotIn(asset, factors.index) + self.assertNotIn(asset, results.index) before_trading_start = handle_data @@ -355,17 +463,20 @@ class FFCAlgorithmTestCase(TestCase): ) def initialize(context): + pipeline = Pipeline('test') context.vwaps = [] for length, key in iteritems(vwap_keys): context.vwaps.append(VWAP(window_length=length)) - add_factor(context.vwaps[-1], name=key) + pipeline.add(context.vwaps[-1], name=key) + + attach_pipeline(pipeline) def handle_data(context, data): today = get_datetime() - factors = data.factors + results = drain_pipeline('test') for length, key in iteritems(vwap_keys): for asset in assets: - computed = factors.loc[asset, key] + computed = results.loc[asset, key] expected = vwaps[length][asset].loc[today] # Only having two places of precision here is a bit diff --git a/tests/modelling/test_pipeline.py b/tests/modelling/test_pipeline.py new file mode 100644 index 00000000..4f6793b1 --- /dev/null +++ b/tests/modelling/test_pipeline.py @@ -0,0 +1,127 @@ +""" +Tests for zipline.modelling.pipeline.Pipeline +""" +from unittest import TestCase + +from zipline.data.equities import USEquityPricing +from zipline.modelling.pipeline import Pipeline +from zipline.modelling.factor import Factor +from zipline.modelling.filter import Filter + + +class SomeFactor(Factor): + window_length = 5 + inputs = [USEquityPricing.close, USEquityPricing.high] + + +class SomeOtherFactor(Factor): + window_length = 5 + inputs = [USEquityPricing.close, USEquityPricing.high] + + +class SomeFilter(Filter): + window_length = 5 + inputs = [USEquityPricing.close, USEquityPricing.high] + + +class SomeOtherFilter(Filter): + window_length = 5 + inputs = [USEquityPricing.close, USEquityPricing.high] + + +class PipelineTestCase(TestCase): + + def test_construction(self): + p0 = Pipeline('arglebargle') + self.assertEqual(p0.name, 'arglebargle') + self.assertEqual(p0.columns, {}) + self.assertIs(p0.screen, None) + + columns = {'f': SomeFactor()} + p1 = Pipeline('test', columns=columns) + self.assertEqual(p1.columns, columns) + + screen = SomeFilter() + p2 = Pipeline('test', screen=screen) + self.assertEqual(p2.columns, {}) + self.assertEqual(p2.screen, screen) + + p3 = Pipeline('test', columns=columns, screen=screen) + self.assertEqual(p3.columns, columns) + self.assertEqual(p3.screen, screen) + + def test_construction_bad_input_types(self): + + with self.assertRaises(TypeError): + Pipeline(1) + + with self.assertRaises(TypeError): + Pipeline('test', 1) + + Pipeline('test', {}) + + with self.assertRaises(TypeError): + Pipeline('test', {}, 1) + + with self.assertRaises(TypeError): + Pipeline('test', {}, SomeFactor()) + + Pipeline('test', {}, SomeFactor() > 5) + + def test_add(self): + p = Pipeline('test') + f = SomeFactor() + + p.add(f, 'f') + self.assertEqual(p.columns, {'f': f}) + + p.add(f > 5, 'g') + self.assertEqual(p.columns, {'f': f, 'g': f > 5}) + + with self.assertRaises(TypeError): + p.add(f, 1) + + def test_overwrite(self): + p = Pipeline('test') + f = SomeFactor() + other_f = SomeOtherFactor() + + p.add(f, 'f') + self.assertEqual(p.columns, {'f': f}) + + with self.assertRaises(KeyError) as e: + p.add(other_f, 'f') + [message] = e.exception.args + self.assertEqual(message, "Column 'f' already exists.") + + p.add(other_f, 'f', overwrite=True) + self.assertEqual(p.columns, {'f': other_f}) + + def test_remove(self): + f = SomeFactor() + p = Pipeline('test', columns={'f': f}) + + with self.assertRaises(KeyError) as e: + p.remove('not_a_real_name') + + self.assertEqual(f, p.remove('f')) + + with self.assertRaises(KeyError) as e: + p.remove('f') + + self.assertEqual(e.exception.args, ('f',)) + + def test_set_screen(self): + f, g = SomeFilter(), SomeOtherFilter() + + p = Pipeline('test') + self.assertEqual(p.screen, None) + + p.set_screen(f) + self.assertEqual(p.screen, f) + + with self.assertRaises(ValueError): + p.set_screen(f) + + p.set_screen(g, overwrite=True) + self.assertEqual(p.screen, g) diff --git a/zipline/algorithm.py b/zipline/algorithm.py index e10db5a0..0c3e37af 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -17,6 +17,7 @@ import warnings import pytz import pandas as pd +from pandas.tseries.tools import normalize_date import numpy as np from datetime import datetime @@ -33,16 +34,18 @@ from operator import attrgetter from zipline.errors import ( - AddTermPostInit, + AttachPipelineAfterInitialize, + NoSuchPipeline, OrderDuringInitialize, OverrideCommissionPostInit, OverrideSlippagePostInit, + DrainPipelineDuringInitialize, RegisterAccountControlPostInit, RegisterTradingControlPostInit, UnsupportedCommissionModel, + UnsupportedDatetimeFormat, UnsupportedOrderParameters, UnsupportedSlippageModel, - UnsupportedDatetimeFormat, ) from zipline.finance.trading import TradingEnvironment from zipline.finance.blotter import Blotter @@ -78,9 +81,11 @@ from zipline.modelling.engine import ( from zipline.sources import DataFrameSource, DataPanelSource from zipline.utils.api_support import ( api_method, + require_initialized, require_not_initialized, ZiplineAPI, ) +from zipline.utils.cache import CachedObject, Expired import zipline.utils.events from zipline.utils.events import ( EventManager, @@ -223,12 +228,13 @@ class TradingAlgorithm(object): # Pull in the environment's new AssetFinder for quick reference self.asset_finder = self.trading_environment.asset_finder - self.init_engine(kwargs.pop('ffc_loader', None)) - # Maps from name to Term - self._filters = {} - self._factors = {} - self._classifiers = {} + # Initialize Modeling API data. + self.init_engine(kwargs.pop('ffc_loader', None)) + self._pipelines = [] + # Create an always-expired cache so that we compute the first time data + # is requested. + self._pipeline_cache = CachedObject(None, pd.Timestamp(0, tz='UTC')) self.blotter = kwargs.pop('blotter', None) if not self.blotter: @@ -1326,41 +1332,96 @@ class TradingAlgorithm(object): """ self.register_trading_control(LongOnly()) - ########### - # FFC API # - ########### + ############## + # Modeling API + ############## @api_method - @require_not_initialized(AddTermPostInit()) - def add_factor(self, factor, name): - if name in self._factors: - raise ValueError("Name %r is already a factor!" % name) - self._factors[name] = factor - - @api_method - @require_not_initialized(AddTermPostInit()) - def add_filter(self, filter): - name = "anon_filter_%d" % len(self._filters) - self._filters[name] = filter - - # Note: add_classifier is not yet implemented since you can't do anything - # useful with classifiers yet. - - def _all_terms(self): - # Merge all three dicts. - return dict( - chain.from_iterable( - iteritems(terms) - for terms in (self._filters, self._factors, self._classifiers) - ) - ) - - def compute_factor_matrix(self, start_date): + @require_not_initialized(AttachPipelineAfterInitialize()) + def attach_pipeline(self, pipeline): """ - Compute a factor matrix containing at least the data necessary to - provide values for `start_date`. + Register a pipeline to be computed at the start of each day. + """ + if self._pipelines: + raise NotImplementedError("Multiple pipelines are not supported.") + self._pipelines.append(pipeline) - Loads a factor matrix with data extending from `start_date` until a - year from `start_date`, or until the end of the simulation. + @api_method + @require_initialized(DrainPipelineDuringInitialize()) + def drain_pipeline(self, name=None): + """ + Get the results of pipeline with name `name`. + + Parameters + ---------- + name : str or None + Name of the pipeline for which results are requested. + + Returns + ------- + results : pd.DataFrame + DataFrame containing the results of the requested pipeline for + the current simulation date. + + Raises + ------ + NoSuchPipeline + Raised when no pipeline with the name `name` has been registered. + + See Also + -------- + :meth:`zipline.modelling.FFCEngine.run_pipeline` + """ + # NOTE: We don't currently support multiple pipelines, but we plan to + # in the future. + for p in self._pipelines: + if p.name == name: + break + # This is a for-else block. Yes, that's a thing in Python. + else: + raise NoSuchPipeline( + name=name, + valid=[p.name for p in self._pipelines], + ) + return self._pipeline_results(p) + + def _pipeline_results(self, pipeline): + """ + Internal implementation of `drain_pipeline`. + """ + today = normalize_date(self.get_datetime()) + try: + data = self._pipeline_cache.unwrap(today) + except Expired: + data, valid_until = self._run_pipeline(pipeline, today) + self._pipeline_cache = CachedObject(data, valid_until) + + # Now that we have a cached result, try to return the data for today. + try: + return data.loc[today] + except KeyError: + # This happens if no assets passed the pipeline screen on a given + # day. + return pd.DataFrame(index=[], columns=data.columns) + + def _run_pipeline(self, pipeline, start_date): + """ + Compute `pipeline`, providing values for at least `start_date`. + + Produces a DataFrame containing data for days between `start_date` and + `end_date`, where `end_date` is defined by: + + `end_date = min(start_date + 252 trading days, simulation_end)` + + 252 is a mostly-arbitrary number based on napkin math. The window + length will likely become dynamic and/or configurable in the future. + + Returns + ------- + (data, valid_until) : tuple (pd.DataFrame, pd.Timestamp) + + See Also + -------- + FFCEngine.run_pipeline """ days = self.trading_environment.trading_days @@ -1369,16 +1430,19 @@ class TradingAlgorithm(object): # ...continuing until either the day before the simulation end, or # until 252 days of data have been loaded. 252 is a totally arbitrary - # choice that seemed reasonable based on napkin math. + # choice that seemed reasonable based on napkin math. In the future, + # this number will likely become dynamic and/or customizable, so don't + # rely on it being 252. sim_end = self.sim_params.last_close.normalize() end_loc = min(start_date_loc + 252, days.get_loc(sim_end)) end_date = days[end_loc] - return self.engine.factor_matrix( - self._all_terms(), - start_date, - end_date, - ), end_date + return \ + self.engine.run_pipeline(pipeline, start_date, end_date), end_date + + ################## + # End Modeling API + ################## def current_universe(self): return self._current_universe diff --git a/zipline/errors.py b/zipline/errors.py index ca4c99e3..3a328442 100644 --- a/zipline/errors.py +++ b/zipline/errors.py @@ -377,15 +377,33 @@ class UnknownRankMethod(ZiplineError): ) -class AddTermPostInit(ZiplineError): +class AttachPipelineAfterInitialize(ZiplineError): """ - Raised when a user tries to call add_{filter,factor,classifier} - outside of initialize. + Raised when a user tries to call add_pipeline outside of initialize. """ msg = ( - "Attempted to add a new filter, factor, or classifier " - "outside of initialize.\n" - "New FFC terms may only be added during initialize." + "Attempted to attach a pipeline after initialize()." + "attach_pipeline() can only be called during initialize." + ) + + +class DrainPipelineDuringInitialize(ZiplineError): + """ + Raised when a user tries to call `drain_pipeline` during initialize. + """ + msg = ( + "Attempted to call drain_pipeline() during initialize. " + "drain_pipeline() can only be called once initialize has completed." + ) + + +class NoSuchPipeline(ZiplineError, KeyError): + """ + Raised when a user tries to access a non-existent pipeline by name. + """ + msg = ( + "No pipeline named '{name}' exists. Valid pipeline names are {valid}. " + "Did you forget to call attach_pipeline()?" ) diff --git a/zipline/modelling/engine.py b/zipline/modelling/engine.py index 6fe4cc2d..cb160cf4 100644 --- a/zipline/modelling/engine.py +++ b/zipline/modelling/engine.py @@ -5,21 +5,15 @@ from abc import ( ABCMeta, abstractmethod, ) -from operator import and_ +from uuid import uuid4 + from six import ( iteritems, - itervalues, with_metaclass, ) -from six.moves import ( - reduce, - zip_longest, -) +from six.moves import zip_longest +from numpy import array -from numpy import ( - add, - empty_like, -) from pandas import ( DataFrame, date_range, @@ -28,32 +22,25 @@ from pandas import ( from zipline.lib.adjusted_array import ensure_ndarray from zipline.errors import NoFurtherDataError +from zipline.utils.numpy_utils import repeat_first_axis, repeat_last_axis from zipline.utils.pandas_utils import explode -from .classifier import Classifier -from .factor import Factor -from .filter import Filter -from .graph import TermGraph from .term import AssetExists class FFCEngine(with_metaclass(ABCMeta)): @abstractmethod - def factor_matrix(self, terms, start_date, end_date): + def run_pipeline(self, pipeline, start_date, end_date): """ - Compute values for `terms` between `start_date` and `end_date`. + Compute values for `pipeline` between `start_date` and `end_date`. - Returns a DataFrame with a MultiIndex of (date, asset) pairs on the - index. On each date, we return a row for each asset that passed all - instances of `Filter` in `terms, and the columns of the returned frame - will be the keys in `terms` whose values are instances of `Factor`. + Returns a DataFrame with a MultiIndex of (date, asset) pairs Parameters ---------- - terms : dict[str -> zipline.modelling.term.Term] - Dict mapping term names to instances. The supplied names are used - as column names in our output frame. + pipeline : zipline.modelling.pipeline.Pipeline + The pipeline to run. start_date : pd.Timestamp Start date of the computed matrix. end_date : pd.Timestamp @@ -61,23 +48,31 @@ class FFCEngine(with_metaclass(ABCMeta)): Returns ------- - matrix : pd.DataFrame - A matrix of computed results. + result : pd.DataFrame + A frame of computed results. + + The columns `result` correspond wil be the computed results of + `pipeline.columns`, which should be a dictionary mapping strings to + instances of `zipline.modelling.term.Term`. + + For each date between `start_date` and `end_date`, `result` will + contain a row for each asset that passed `pipeline.screen`. A + screen of None indicates that a row should be returned for each + asset that existed each day. """ - raise NotImplementedError("factor_matrix") + raise NotImplementedError("run_pipeline") class NoOpFFCEngine(FFCEngine): """ - FFCEngine that doesn't do anything. + An FFCEngine that doesn't do anything. """ - - def factor_matrix(self, terms, start_date, end_date): + def run_pipeline(self, pipeline, start_date, end_date): return DataFrame( index=MultiIndex.from_product( [date_range(start=start_date, end=end_date, freq='D'), ()], ), - columns=sorted(terms.keys()) + columns=sorted(pipeline.columns.keys()), ) @@ -110,15 +105,14 @@ class SimpleFFCEngine(object): self._finder = asset_finder self._root_mask_term = AssetExists() - def factor_matrix(self, terms, start_date, end_date): + def run_pipeline(self, pipeline, start_date, end_date): """ - Compute a factor matrix. + Compute a pipeline. Parameters ---------- - terms : dict[str -> zipline.modelling.term.Term] - Dict mapping term names to instances. The supplied names are used - as column names in our output frame. + pipeline : zipline.modelling.pipeline.Pipeline + The pipeline to run. start_date : pd.Timestamp Start date of the computed matrix. end_date : pd.Timestamp @@ -155,7 +149,7 @@ class SimpleFFCEngine(object): See Also -------- - FFCEngine.factor_matrix + FFCEngine.run_pipeline """ if end_date <= start_date: raise ValueError( @@ -163,36 +157,23 @@ class SimpleFFCEngine(object): "start_date=%s, end_date=%s" % (start_date, end_date) ) - graph = TermGraph(terms) + screen_name = uuid4().hex + graph = pipeline.to_graph(screen_name, self._root_mask_term) extra_rows = graph.extra_rows[self._root_mask_term] - root_mask = self._compute_root_mask(start_date, end_date, extra_rows) dates, assets, root_mask_values = explode(root_mask) - raw_outputs = self.compute_chunk( + + outputs = self.compute_chunk( graph, dates, assets, initial_workspace={self._root_mask_term: root_mask_values}, ) - # Collect the results that we'll actually show to the user. - filters, factors = {}, {} - for name, term in iteritems(terms): - if isinstance(term, Filter): - filters[name] = raw_outputs[name] - elif isinstance(term, Factor): - factors[name] = raw_outputs[name] - elif isinstance(term, Classifier): - continue - else: - raise ValueError("Unknown term type: %s" % term) - - # Add the root mask as an implicit filter, truncating off the extra - # rows that we only needed to compute other terms. - filters['base'] = root_mask_values[extra_rows:] out_dates = dates[extra_rows:] + screen_values = outputs.pop(screen_name) - return self._format_factor_matrix(out_dates, assets, filters, factors) + return self._to_narrow(outputs, screen_values, out_dates, assets) def _compute_root_mask(self, start_date, end_date, extra_rows): """ @@ -360,98 +341,41 @@ class SimpleFFCEngine(object): out[name] = workspace[term][graph_extra_rows[term]:] return out - def _format_factor_matrix(self, dates, assets, filters, factors): + def _to_narrow(self, data, mask, dates, assets): """ - Convert raw computed filters/factors into a DataFrame for public APIs. + Convert raw computed pipeline results into a DataFrame for public APIs. Parameters ---------- - dates : np.array[datetime64] - Row index for arrays in `filters` and `factors.` - assets : np.array[int64] - Column index for arrays in `filters` and `factors.` - filters : dict - Dict mapping filter names -> computed filters. - factors : dict - Dict mapping factor names -> computed factors. + data : dict[str -> ndarray[ndim=2]] + Dict mapping column names to computed results. + mask : ndarray[bool, ndim=2] + Mask array of values to keep. + dates : ndarray[datetime64, ndim=1] + Row index for arrays `data` and `mask` + assets : ndarray[int64, ndim=2] + Column index for arrays `data` and `mask` Returns ------- - factor_matrix : pd.DataFrame - The indices of `factor_matrix` are as follows: + results : pd.DataFrame + The indices of `results` are as follows: index : two-tiered MultiIndex of (date, asset). - For each date, we return a row for each asset that passed all - filters on that date. - columns : keys from `factor_data` + Contains an entry for each (date, asset) pair corresponding to + a `True` value in `mask`. + columns : Index of str + One column per entry in `data`. - Each date/asset/factor triple contains the computed value of the given - factor on the given date for the given asset. + If mask[date, asset] is True, then result.loc[(date, asset), colname] + will contain the value of data[colname][date, asset]. """ - # FUTURE OPTIMIZATION: Cythonize all of this. - - # Boolean mask of values that passed all filters. - unioned = reduce(and_, itervalues(filters)) - - # Parallel arrays of (x,y) coords for (date, asset) pairs that passed - # all filters. Each entry here will correspond to a row in our output - # frame. - nonzero_xs, nonzero_ys = unioned.nonzero() - - # Raw arrays storing (date, asset) pairs. - # These will form the index of our output frame. - raw_dates_index = empty_like(nonzero_xs, dtype='datetime64[ns]') - raw_assets_index = empty_like(nonzero_xs, dtype=int) - - # Mapping from column_name -> array. - # This will be the `data` arg to our output frame. - columns = { - name: empty_like(nonzero_xs, dtype=factor.dtype) - for name, factor in iteritems(factors) - } - # We're going to iterate over `iteritems(columns)` a whole bunch of - # times down below. It's faster to construct iterate over a tuple of - # pairs. - columns_iter = tuple(iteritems(columns)) - - # This is tricky. - - # unioned.sum(axis=1) gives us an array of the same size as `dates` - # containing, for each date, the number of assets that passed our - # filters on that date. - - # Running this through add.accumulate gives us an array containing, for - # each date, the running total of the number of assets that passed our - # filters on or before that date. - - # This means that (bounds[i - 1], bounds[i]) gives us the indices of - # the first and last rows in our output frame for each date in `dates`. - bounds = add.accumulate(unioned.sum(axis=1)) - day_start = 0 - for day_idx, day_end in enumerate(bounds): - - day_bounds = slice(day_start, day_end) - column_indices = nonzero_ys[day_bounds] - - raw_dates_index[day_bounds] = dates[day_idx] - raw_assets_index[day_bounds] = assets[column_indices] - for name, colarray in columns_iter: - colarray[day_bounds] = factors[name][day_idx, column_indices] - - # Upper bound of current row becomes lower bound for next row. - day_start = day_end - + resolved_assets = array(self._finder.retrieve_all(assets)) + dates_kept = repeat_last_axis(dates.values, len(assets))[mask] + assets_kept = repeat_first_axis(resolved_assets, len(dates))[mask] return DataFrame( - data=columns, - index=MultiIndex.from_arrays( - [ - raw_dates_index, - # FUTURE OPTIMIZATION: - # Avoid duplicate lookups by grouping and only looking up - # each unique sid once. - self._finder.retrieve_all(raw_assets_index), - ], - ) + data={name: arr[mask] for name, arr in iteritems(data)}, + index=MultiIndex.from_arrays([dates_kept, assets_kept]), ).tz_localize('UTC', level=0) def _validate_compute_chunk_params(self, dates, assets, initial_workspace): diff --git a/zipline/modelling/pipeline.py b/zipline/modelling/pipeline.py new file mode 100644 index 00000000..5ec6b6ed --- /dev/null +++ b/zipline/modelling/pipeline.py @@ -0,0 +1,157 @@ +from zipline.utils.preprocess import expect_types, optional +from zipline.modelling.term import Term +from zipline.modelling.filter import Filter +from zipline.modelling.graph import TermGraph + + +class Pipeline(object): + """ + Parameters + ---------- + name : str, optional + Name for this pipeline. + columns : dict, optional + Initial columns. + screen : zipline.modelling.term.Filter, optional + Initial screen. + + Methods + ------- + add + remove + apply_screen + + Attributes + ---------- + columns + screen + """ + __slots__ = ('_name', '_columns', '_screen', '__weakref__') + + @expect_types( + name=str, + columns=optional(dict), + screen=optional(Filter), + ) + def __init__(self, name, columns=None, screen=None): + self._name = name + if columns is None: + columns = {} + self._columns = columns + self._screen = screen + + @property + def name(self): + """ + The name of this pipeline. + """ + return self._name + + @property + def columns(self): + """ + The columns currently applied to this pipeline. + """ + return self._columns + + @property + def screen(self): + """ + The screen applied to the rows of this pipeline. + """ + return self._screen + + @expect_types(term=Term, name=str) + def add(self, term, name, overwrite=False): + """ + Add a column. + + The results of computing `term` will show up as a column in the + DataFrame produced by running this pipeline. + + Parameters + ---------- + column : zipline.modelling.Term + A Filter, Factor, or Classifier to add to the pipeline. + name : str + Name of the column to add. + overwrite : bool + Whether to overwrite the existing entry if we already have a column + named `name`. + """ + columns = self.columns + if name in columns: + if overwrite: + self.remove(name) + else: + raise KeyError("Column '{}' already exists.".format(name)) + + self._columns[name] = term + + @expect_types(name=str) + def remove(self, name): + """ + Remove a column. + + Parameters + ---------- + name : str + The name of the column to remove. + + Raises + ------ + KeyError + If `name` is not in self.columns. + + Returns + ------- + removed : zipline.modelling.term.Term + The removed term. + """ + return self.columns.pop(name) + + @expect_types(screen=Filter) + def set_screen(self, screen, overwrite=False): + """ + Apply a screen to this Pipeline. + + If no screen has yet been applied to the pipeline, this method sets + `screen` as the current screen. + + Parameter + --------- + filter : zipline.modelling.filter.Filter + The screen to apply. + overwrite : bool + Whether to overwrite any existing screen. If overwrite is False + and self.screen is not None, we raise an error. + """ + if self._screen is not None and not overwrite: + raise ValueError( + "set_screen() called with overwrite=False and screen already " + "set.\n" + "If you want to apply multiple filters as a screen use " + "set_screen(filter1 & filter2 & ...).\n" + "If you want to replace the previous screen with a new one, " + "use set_screen(new_filter, overwrite=True)." + ) + self._screen = screen + + def to_graph(self, screen_name, default_screen): + """ + Compile into a TermGraph. + + Parameters + ---------- + screen_name : str + Name to supply for self.screen. + default_screen : zipline.modelling.term.Term + Term to use as a screen if self.screen is None. + """ + columns = self.columns.copy() + screen = self.screen + if screen is None: + screen = default_screen + columns[screen_name] = screen + + return TermGraph(columns) diff --git a/zipline/protocol.py b/zipline/protocol.py index c8f80b10..7353a76b 100644 --- a/zipline/protocol.py +++ b/zipline/protocol.py @@ -17,7 +17,6 @@ from copy import copy from six import iteritems, iterkeys import pandas as pd -from pandas.tseries.tools import normalize_date import numpy as np from . utils.protocol_utils import Enum @@ -494,24 +493,6 @@ class BarData(object): def __init__(self, data=None): self._data = data or {} self._contains_override = None - self._factor_matrix = None - self._factor_matrix_expires = pd.Timestamp(0, tz='UTC') - - @property - def factors(self): - algo = get_algo_instance() - today = normalize_date(algo.get_datetime()) - if today > self._factor_matrix_expires: - self._factor_matrix, self._factor_matrix_expires = \ - algo.compute_factor_matrix(today) - try: - return self._factor_matrix.loc[today] - except KeyError: - # This happens if no assets passed our filters on a given day. - return pd.DataFrame( - index=[], - columns=self._factor_matrix.columns, - ) def __contains__(self, name): if self._contains_override: diff --git a/zipline/utils/api_support.py b/zipline/utils/api_support.py index 729991a0..5096fc63 100644 --- a/zipline/utils/api_support.py +++ b/zipline/utils/api_support.py @@ -76,3 +76,25 @@ def require_not_initialized(exception): return method(self, *args, **kwargs) return wrapped_method return decorator + + +def require_initialized(exception): + """ + Decorator for API methods that should only be called after + TradingAlgorithm.initialize. `exception` will be raised if the method is + called before initialize has completed. + + Usage + ----- + @require_initialized(SomeException("Don't do that!")) + def method(self): + # Do stuff that should only be allowed after initialize. + """ + def decorator(method): + @wraps(method) + def wrapped_method(self, *args, **kwargs): + if not self.initialized: + raise exception + return method(self, *args, **kwargs) + return wrapped_method + return decorator