From 26fd6fda8b095d052374fb27812dca3a6cad4869 Mon Sep 17 00:00:00 2001 From: Scott Sanderson Date: Thu, 10 Sep 2015 16:49:15 -0400 Subject: [PATCH] ENH/BUG: Modeling API enhancements. - Fixes an error where Modeling API data known as of the close of `day N` would be shown to algorithms during `before_trading_start` as of the close of the same day. Algorithms should now only receive data during `before_trading_start/handle_data` that was known as of the simulation time at which the function would be called. - All Term instances now have a `mask` attribute that must be a `Filter` or an instance of `AssetExists()`. `mask` can be used to specify that a Factor should be computed in a manner that ignores the values that were not `True` in the mask. - Changed the interface for `FFCLoader.load_adjusted_array` and `Term._compute` from `(columns, mask)`, with mask as a DataFrame, to `(columns, dates, assets, mask)`, where mask is a numpy array. This is primarily to avoid having to reconstruct extra DataFrames when using masks produced by non `AssetExists` filters. - Adds `BoundColumn.latest`, which gives the most-recently-known value of a column. --- setup.py | 8 +- tests/modelling/base.py | 44 ++- tests/modelling/test_engine.py | 188 +++++++----- tests/modelling/test_factor.py | 52 +++- tests/modelling/test_filter.py | 118 +++++--- tests/modelling/test_frameload.py | 27 +- tests/modelling/test_modelling_algo.py | 211 +++++++++++-- tests/modelling/test_numerical_expression.py | 4 +- tests/modelling/test_term.py | 30 +- .../test_us_equity_pricing_loader.py | 88 ++++-- tests/test_assets.py | 29 +- zipline/algorithm.py | 13 +- zipline/assets/assets.py | 22 +- zipline/data/dataset.py | 5 + zipline/data/ffc/frame.py | 12 +- zipline/data/ffc/loaders/_adjustments.pyx | 282 ++++++++++++++++++ zipline/data/ffc/loaders/_equities.pyx | 196 ++++++++++++ .../data/ffc/loaders/_us_equity_pricing.pyx | 189 +----------- zipline/data/ffc/loaders/us_equity_pricing.py | 140 +++++---- zipline/data/ffc/synthetic.py | 10 +- zipline/errors.py | 10 + zipline/lib/adjustment.pyx | 3 +- zipline/modelling/engine.py | 188 ++++++++---- zipline/modelling/expression.py | 6 +- zipline/modelling/factor/factor.py | 56 ++-- zipline/modelling/factor/latest.py | 15 + zipline/modelling/filter.py | 102 +------ zipline/modelling/graph.py | 163 +++++++--- zipline/modelling/term.py | 125 ++++++-- zipline/modelling/visualize.py | 19 +- zipline/protocol.py | 9 +- zipline/utils/pandas_utils.py | 12 + 32 files changed, 1643 insertions(+), 733 deletions(-) create mode 100644 zipline/data/ffc/loaders/_adjustments.pyx create mode 100644 zipline/data/ffc/loaders/_equities.pyx create mode 100644 zipline/modelling/factor/latest.py create mode 100644 zipline/utils/pandas_utils.py diff --git a/setup.py b/setup.py index ad3966f6..52d92132 100644 --- a/setup.py +++ b/setup.py @@ -64,8 +64,12 @@ ext_modules = LazyCythonizingList([ ('zipline.lib.adjustment', ['zipline/lib/adjustment.pyx']), ('zipline.lib.rank', ['zipline/lib/rank.pyx']), ( - 'zipline.data.ffc.loaders._us_equity_pricing', - ['zipline/data/ffc/loaders/_us_equity_pricing.pyx'] + 'zipline.data.ffc.loaders._equities', + ['zipline/data/ffc/loaders/_equities.pyx'], + ), + ( + 'zipline.data.ffc.loaders._adjustments', + ['zipline/data/ffc/loaders/_adjustments.pyx'], ), ]) diff --git a/tests/modelling/base.py b/tests/modelling/base.py index 031d0e03..e0e8490d 100644 --- a/tests/modelling/base.py +++ b/tests/modelling/base.py @@ -5,13 +5,14 @@ from functools import wraps from unittest import TestCase from numpy import arange, prod -from numpy.random import randn, seed as random_seed from pandas import date_range, Int64Index, DataFrame 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 from zipline.utils.tradingcalendar import trading_day @@ -56,9 +57,15 @@ class BaseFFCTestCase(TestCase): assets, self.__calendar[0], self.__calendar[-1], - )) + ), + ) self.__finder = env.asset_finder - self.__mask = self.__finder.lifetimes(self.__calendar[-10:]) + + # Use a 30-day period at the end of the year by default. + self.__mask = self.__finder.lifetimes( + self.__calendar[-30:], + include_start_date=False, + ) @property def default_shape(self): @@ -74,6 +81,12 @@ class BaseFFCTestCase(TestCase): ---------- terms : dict Mapping from termname -> term object. + initial_workspace : dict + Initial workspace to forward to SimpleFFCEngine.compute_chunk. + mask : DataFrame, optional + This is a value to pass to `initial_workspace` as the mask from + `AssetExists()`. Defaults to a frame of shape `self.default_shape` + containing all True values. Returns ------- @@ -85,10 +98,23 @@ class BaseFFCTestCase(TestCase): self.__calendar, self.__finder, ) - mask = mask if mask is not None else self.__mask - return engine.compute_chunk(TermGraph(terms), mask, initial_workspace) + if mask is None: + mask = self.__mask + + dates, assets, mask_values = explode(mask) + initial_workspace.setdefault(AssetExists(), mask_values) + return engine.compute_chunk( + TermGraph(terms), + dates, + assets, + initial_workspace, + ) def build_mask(self, array): + """ + Helper for constructing an AssetExists mask from a boolean-coercible + array. + """ ndates, nassets = array.shape return DataFrame( array, @@ -105,11 +131,3 @@ class BaseFFCTestCase(TestCase): Build a block of testing data from numpy.arange. """ return arange(prod(shape), dtype=dtype).reshape(shape) - - @with_default_shape - def randn_data(self, seed, shape): - """ - Build a block of testing data from numpy.random.randn. - """ - random_seed(seed) - return randn(*shape) diff --git a/tests/modelling/test_engine.py b/tests/modelling/test_engine.py index f2aa6e08..76b4707a 100644 --- a/tests/modelling/test_engine.py +++ b/tests/modelling/test_engine.py @@ -7,8 +7,8 @@ from itertools import product from numpy import ( full, - isnan, nan, + zeros, ) from numpy.testing import assert_array_equal from pandas import ( @@ -191,33 +191,21 @@ class ConstantInputTestCase(TestCase): expected_high_low = 3.0 * (constants[high] - constants[low]) assert_frame_equal( high_low_result, - DataFrame( - expected_high_low, - index=dates, - columns=self.assets, - ) + DataFrame(expected_high_low, index=dates, columns=self.assets), ) open_close_result = results['open_close'].unstack() expected_open_close = 3.0 * (constants[open] - constants[close]) assert_frame_equal( open_close_result, - DataFrame( - expected_open_close, - index=dates, - columns=self.assets, - ) + DataFrame(expected_open_close, index=dates, columns=self.assets), ) avg_result = results['avg'].unstack() expected_avg = (expected_high_low + expected_open_close) / 2.0 assert_frame_equal( avg_result, - DataFrame( - expected_avg, - index=dates, - columns=self.assets, - ) + DataFrame(expected_avg, index=dates, columns=self.assets), ) @@ -338,20 +326,18 @@ class SyntheticBcolzTestCase(TestCase): def setUpClass(cls): cls.first_asset_start = Timestamp('2015-04-01', tz='UTC') cls.env = TradingEnvironment() - cls.trading_day = cls.env.trading_day + cls.trading_day = day = cls.env.trading_day + cls.calendar = date_range('2015', '2015-08', tz='UTC', freq=day) + cls.asset_info = make_rotating_asset_info( num_assets=6, first_start=cls.first_asset_start, - frequency=cls.trading_day, + frequency=day, periods_between_starts=4, asset_lifetime=8, ) + cls.last_asset_end = cls.asset_info['end_date'].max() cls.all_assets = cls.asset_info.index - cls.all_dates = date_range( - start=cls.first_asset_start, - end=cls.asset_info['end_date'].max(), - freq=cls.trading_day, - ) cls.env.write_data(equities_df=cls.asset_info) cls.finder = cls.env.asset_finder @@ -359,34 +345,74 @@ class SyntheticBcolzTestCase(TestCase): cls.temp_dir = TempDirectory() cls.temp_dir.create() - cls.writer = SyntheticDailyBarWriter( - asset_info=cls.asset_info[['start_date', 'end_date']], - calendar=cls.all_dates, - ) - table = cls.writer.write( - cls.temp_dir.getpath('testdata.bcolz'), - cls.all_dates, - cls.all_assets, - ) + try: + cls.writer = SyntheticDailyBarWriter( + asset_info=cls.asset_info[['start_date', 'end_date']], + calendar=cls.calendar, + ) + table = cls.writer.write( + cls.temp_dir.getpath('testdata.bcolz'), + cls.calendar, + cls.all_assets, + ) - cls.ffc_loader = USEquityPricingLoader( - BcolzDailyBarReader(table), - NullAdjustmentReader(), - ) + cls.ffc_loader = USEquityPricingLoader( + BcolzDailyBarReader(table), + NullAdjustmentReader(), + ) + except: + cls.temp_dir.cleanup() + raise @classmethod def tearDownClass(cls): del cls.env cls.temp_dir.cleanup() + def write_nans(self, df): + """ + Write nans to the locations in data corresponding to the (date, asset) + pairs for which we wouldn't have data for `asset` on `date` in a + backtest. + + Parameters + ---------- + df : pd.DataFrame + A DataFrame with a DatetimeIndex as index and an object index of + Assets as columns. + + This means that we write nans for dates after an asset's end_date and + **on or before** an asset's start_date. The assymetry here is because + of the fact that, on the morning of an asset's first date, we haven't + yet seen any trades for that asset, so we wouldn't be able to show any + useful data to the user. + """ + # Mask out with nans all the dates on which each asset didn't exist + index = df.index + min_, max_ = index[[0, -1]] + for asset in df.columns: + if asset.start_date >= min_: + start = index.get_loc(asset.start_date, method='bfill') + df.loc[:start + 1, asset] = nan # +1 to overwrite start_date + if asset.end_date <= max_: + end = index.get_loc(asset.end_date) + df.ix[end + 1:, asset] = nan # +1 to *not* overwrite end_date + def test_SMA(self): engine = SimpleFFCEngine( self.ffc_loader, self.env.trading_days, self.finder, ) - dates, assets = self.all_dates, self.all_assets window_length = 5 + assets = self.all_assets + dates = date_range( + self.first_asset_start + self.trading_day, + self.last_asset_end, + freq=self.trading_day, + ) + dates_to_test = dates[window_length:] + SMA = SimpleMovingAverage( inputs=(USEquityPricing.close,), window_length=window_length, @@ -394,27 +420,30 @@ class SyntheticBcolzTestCase(TestCase): results = engine.factor_matrix( {'sma': SMA}, - dates[window_length], - dates[-1], + dates_to_test[0], + dates_to_test[-1], ) - raw_closes = self.writer.expected_values_2d(dates, assets, 'close') - expected_sma_result = rolling_mean( - raw_closes, + + # Shift back the raw inputs by a trading day because we expect our + # computed results to be computed using values anchored on the + # **previous** day's data. + expected_raw = rolling_mean( + self.writer.expected_values_2d( + dates - self.trading_day, assets, 'close', + ), window_length, min_periods=1, ) - expected_sma_result[isnan(raw_closes)] = nan - expected_sma_result = expected_sma_result[window_length:] - sma_result = results['sma'].unstack() - assert_frame_equal( - sma_result, - DataFrame( - expected_sma_result, - index=dates[window_length:], - columns=assets, - ), + expected = DataFrame( + # Truncate off the extra rows needed to compute the SMAs. + expected_raw[window_length:], + index=dates_to_test, # dates_to_test is dates[window_length:] + columns=self.finder.retrieve_all(assets), ) + self.write_nans(expected) + result = results['sma'].unstack() + assert_frame_equal(result, expected) def test_drawdown(self): # The monotonically-increasing data produced by SyntheticDailyBarWriter @@ -427,8 +456,15 @@ class SyntheticBcolzTestCase(TestCase): self.env.trading_days, self.finder, ) - dates, assets = self.all_dates, self.all_assets window_length = 5 + assets = self.all_assets + dates = date_range( + self.first_asset_start + self.trading_day, + self.last_asset_end, + freq=self.trading_day, + ) + dates_to_test = dates[window_length:] + drawdown = MaxDrawdown( inputs=(USEquityPricing.close,), window_length=window_length, @@ -436,32 +472,27 @@ class SyntheticBcolzTestCase(TestCase): results = engine.factor_matrix( {'drawdown': drawdown}, - dates[window_length], - dates[-1], + dates_to_test[0], + dates_to_test[-1], ) - dd_result = results['drawdown'] - # We expect NaNs when the asset was undefined, otherwise 0 everywhere, # since the input is always increasing. - expected = self.writer.expected_values_2d(dates, assets, 'close') - expected[~isnan(expected)] = 0 - expected = expected[window_length:] - - assert_frame_equal( - dd_result.unstack(), - DataFrame( - expected, - index=dates[window_length:], - columns=assets, - ), + expected = DataFrame( + data=zeros((len(dates_to_test), len(assets)), dtype=float), + index=dates_to_test, + columns=self.finder.retrieve_all(assets), ) + self.write_nans(expected) + result = results['drawdown'].unstack() + + assert_frame_equal(expected, result) class MultiColumnLoaderTestCase(TestCase): def setUp(self): self.assets = [1, 2, 3] - self.dates = date_range('2014-01-01', '2014-02-01', freq='D', tz='UTC') + self.dates = date_range('2014-01', '2014-03', freq='D', tz='UTC') asset_info = make_simple_asset_info( self.assets, @@ -475,6 +506,11 @@ class MultiColumnLoaderTestCase(TestCase): def test_engine_with_multicolumn_loader(self): open_, close = USEquityPricing.open, USEquityPricing.close + # Test for thirty days up to the second to last day that we think all + # the assets existed. If we test the last day of our calendar, no + # assets will be in our output, because their end dates are all + dates_to_test = self.dates[-32:-2] + loader = MultiColumnLoader({ open_: ConstantLoader(dates=self.dates, assets=self.assets, @@ -489,12 +525,14 @@ class MultiColumnLoaderTestCase(TestCase): factor = RollingSumDifference() result = engine.factor_matrix({'f': factor}, - self.dates[2], - self.dates[-1]) + dates_to_test[0], + dates_to_test[-1]) self.assertIsNotNone(result) self.assertEqual({'f'}, set(result.columns)) - # (close - open) * window = (1 - 2) * 3 = -3 - # skipped 2 from the start, so that the window is full - check_arrays(result['f'], - Series([-3] * len(self.assets) * (len(self.dates) - 2))) + result_index = self.assets * len(dates_to_test) + result_shape = (len(result_index),) + check_arrays( + result['f'], + Series(index=result_index, data=full(result_shape, -3)), + ) diff --git a/tests/modelling/test_factor.py b/tests/modelling/test_factor.py index a35ca584..71d63c86 100644 --- a/tests/modelling/test_factor.py +++ b/tests/modelling/test_factor.py @@ -1,13 +1,10 @@ """ Tests for Factor terms. """ -from numpy import ( - array, - ones, -) - +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.utils.test_utils import check_arrays from .base import BaseFFCTestCase @@ -18,6 +15,11 @@ class F(Factor): window_length = 0 +class Mask(Filter): + inputs = () + window_length = 0 + + class FactorTestCase(BaseFFCTestCase): def setUp(self): @@ -135,3 +137,43 @@ class FactorTestCase(BaseFFCTestCase): }) # Not passing a method should default to ordinal. check({'ordinal': self.f.rank(ascending=False)}) + + def test_rank_after_mask(self): + # data = arange(25).reshape(5, 5).transpose() % 4 + data = array([[0, 1, 2, 3, 0], + [1, 2, 3, 0, 1], + [2, 3, 0, 1, 2], + [3, 0, 1, 2, 3], + [0, 1, 2, 3, 0]], dtype=float) + 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()), + } + + expected = { + "ascending_nomask": array([[1., 3., 4., 5., 2.], + [2., 4., 5., 1., 3.], + [3., 5., 1., 2., 4.], + [4., 1., 2., 3., 5.], + [1., 3., 4., 5., 2.]]), + # Diagonal should be all nans, and anything whose rank was less + # than the diagonal in the unmasked calc should go down by 1. + "ascending_mask": array([[nan, 2., 3., 4., 1.], + [2., nan, 4., 1., 3.], + [2., 4., nan, 1., 3.], + [3., 1., 2., nan, 4.], + [1., 2., 3., 4., nan]]), + } + + results = self.run_terms( + terms, + initial_workspace, + mask=self.build_mask(ones((5, 5))), + ) + for method in results: + check_arrays(expected[method], results[method]) diff --git a/tests/modelling/test_filter.py b/tests/modelling/test_filter.py index 74a1b088..4349754a 100644 --- a/tests/modelling/test_filter.py +++ b/tests/modelling/test_filter.py @@ -1,6 +1,7 @@ """ Tests for filter terms. """ +from itertools import product from operator import and_ from numpy import ( @@ -13,17 +14,20 @@ from numpy import ( nan, nanpercentile, ones, + ones_like, putmask, ) +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.utils.test_utils import check_arrays -from .base import BaseFFCTestCase +from .base import BaseFFCTestCase, with_default_shape -def rowwise_rank(array): +def rowwise_rank(array, mask=None): """ Take a 2D array and return the 0-indexed sorted position of each element in the array for each row. @@ -61,6 +65,11 @@ class SomeOtherFactor(Factor): window_length = 0 +class Mask(Filter): + inputs = () + window_length = 0 + + class FilterTestCase(BaseFFCTestCase): def setUp(self): @@ -68,6 +77,14 @@ class FilterTestCase(BaseFFCTestCase): self.f = SomeFactor() self.g = SomeOtherFactor() + @with_default_shape + def randn_data(self, seed, shape): + """ + Build a block of testing data from numpy.random.randn. + """ + random_seed(seed) + return randn(*shape) + def test_bad_percentiles(self): f = self.f @@ -81,25 +98,58 @@ class FilterTestCase(BaseFFCTestCase): with self.assertRaises(BadPercentileBounds): f.percentile_between(min_, max_) - def test_top(self): + def test_top_and_bottom(self): + data = self.randn_data(seed=5) # Fix a seed for determinism. + + mask_data = ones_like(data, dtype=bool) + mask_data[:, 0] = False + + nan_data = data.copy() + nan_data[:, 0] = nan + + mask = Mask() + initial_workspace = {self.f: data, mask: mask_data} + + methods = ['top', 'bottom'] counts = 2, 3, 10 - data = self.randn_data(seed=5) # Arbitrary seed choice. - results = self.run_terms( - terms={'top_' + str(c): self.f.top(c) for c in counts}, - initial_workspace={self.f: data}, - ) - for c in counts: - result = results['top_' + str(c)] + term_combos = list(product(methods, counts, [True, False])) + + def termname(method, count, masked): + return '_'.join([method, str(count), 'mask' if masked else '']) + + # Add a term for each permutation of top/bottom, count, and + # mask/no_mask. + terms = {} + for method, count, masked in term_combos: + kwargs = {'N': count} + if masked: + kwargs['mask'] = mask + term = getattr(self.f, method)(**kwargs) + terms[termname(method, count, masked)] = term + + results = self.run_terms(terms, initial_workspace=initial_workspace) + + def expected_result(method, count, masked): + # Ranking with a mask is equivalent to ranking with nans applied on + # the masked values. + to_rank = nan_data if masked else data + + if method == 'top': + return rowwise_rank(-to_rank) < count + elif method == 'bottom': + return rowwise_rank(to_rank) < count + + for method, count, masked in term_combos: + result = results[termname(method, count, masked)] # Check that `min(c, num_assets)` assets passed each day. passed_per_day = result.sum(axis=1) check_arrays( passed_per_day, - full_like(passed_per_day, min(c, data.shape[1])), + full_like(passed_per_day, min(count, data.shape[1])), ) - # Check that the top `c` assets passed. - expected = rowwise_rank(-data) < c + expected = expected_result(method, count, masked) check_arrays(result, expected) def test_bottom(self): @@ -228,35 +278,19 @@ class FilterTestCase(BaseFFCTestCase): ) check_arrays(result, expected) - def test_sequenced_filter_order_independent(self): - data = self.arange_data() % 5 - results = self.run_terms( - { - # Sequencing is equivalent to &ing for commutative filters. - 'sequenced': (1.5 < self.f).then(self.f < 3.5), - 'anded': (1.5 < self.f) & (self.f < 3.5), - }, - initial_workspace={self.f: data}, - ) - expected = (1.5 < data) & (data < 3.5) - - check_arrays(results['sequenced'], expected) - check_arrays(results['anded'], expected) - - def test_sequenced_filter_order_dependent(self): - - first = self.f < 1 + def test_percentile_after_mask(self): f_input = eye(5) - - second = self.g.percentile_between(80, 100) g_input = arange(25, dtype=float).reshape(5, 5) - initial_mask = self.build_mask(ones((5, 5))) + custom_mask = self.f < 1 + without_mask = self.g.percentile_between(80, 100) + with_mask = self.g.percentile_between(80, 100, mask=custom_mask) + terms = { - 'first': first, - 'second': second, - 'sequenced': first.then(second), + 'custom_mask': custom_mask, + 'without': without_mask, + 'with': with_mask, } results = self.run_terms( @@ -266,11 +300,11 @@ class FilterTestCase(BaseFFCTestCase): ) # First should pass everything but the diagonal. - check_arrays(results['first'], ~eye(5, dtype=bool)) + check_arrays(results['custom_mask'], ~eye(5, dtype=bool)) # Second should pass the largest value each day. Each row is strictly # increasing, so we always select the last value. - expected_second = array( + expected_without = array( [[0, 0, 0, 0, 1], [0, 0, 0, 0, 1], [0, 0, 0, 0, 1], @@ -278,12 +312,12 @@ class FilterTestCase(BaseFFCTestCase): [0, 0, 0, 0, 1]], dtype=bool, ) - check_arrays(results['second'], expected_second) + check_arrays(results['without'], expected_without) # When sequencing, we should remove the diagonal as an option before # computing percentiles. On the last day, we should get the # second-largest value, rather than the largest. - expected_sequenced = array( + expected_with = array( [[0, 0, 0, 0, 1], [0, 0, 0, 0, 1], [0, 0, 0, 0, 1], @@ -291,4 +325,4 @@ class FilterTestCase(BaseFFCTestCase): [0, 0, 0, 1, 0]], # Different from previous! dtype=bool, ) - check_arrays(results['sequenced'], expected_sequenced) + check_arrays(results['with'], expected_with) diff --git a/tests/modelling/test_frameload.py b/tests/modelling/test_frameload.py index 47512edf..708a00af 100644 --- a/tests/modelling/test_frameload.py +++ b/tests/modelling/test_frameload.py @@ -4,7 +4,7 @@ Tests for zipline.data.ffc.frame.DataFrameFFCLoader from unittest import TestCase from mock import patch -from numpy import arange +from numpy import arange, ones from numpy.testing import assert_array_equal from pandas import ( DataFrame, @@ -40,12 +40,7 @@ class DataFrameFFCLoaderTestCase(TestCase): periods=self.ndates, ) - self.mask = DataFrame( - True, - index=self.dates, - columns=self.sids, - dtype=bool, - ) + self.mask = ones((len(self.dates), len(self.sids)), dtype=bool) def tearDown(self): pass @@ -60,13 +55,17 @@ class DataFrameFFCLoaderTestCase(TestCase): with self.assertRaises(ValueError): # Wrong column. - loader.load_adjusted_array([USEquityPricing.open], self.mask) + loader.load_adjusted_array( + [USEquityPricing.open], self.dates, self.sids, self.mask + ) with self.assertRaises(ValueError): # Too many columns. loader.load_adjusted_array( [USEquityPricing.open, USEquityPricing.close], - self.mask + self.dates, + self.sids, + self.mask, ) def test_baseline(self): @@ -81,7 +80,9 @@ class DataFrameFFCLoaderTestCase(TestCase): sids_slice = slice(1, 3, None) [adj_array] = loader.load_adjusted_array( [USEquityPricing.close], - self.mask.iloc[dates_slice, sids_slice] + self.dates[dates_slice], + self.sids[sids_slice], + self.mask[dates_slice, sids_slice], ) for idx, window in enumerate(adj_array.traverse(window_length=3)): @@ -203,10 +204,12 @@ class DataFrameFFCLoaderTestCase(TestCase): } self.assertEqual(formatted_adjustments, expected_formatted_adjustments) - mask = self.mask.iloc[dates_slice, sids_slice] + mask = self.mask[dates_slice, sids_slice] with patch('zipline.data.ffc.frame.adjusted_array') as m: loader.load_adjusted_array( columns=[USEquityPricing.close], + dates=self.dates[dates_slice], + assets=self.sids[sids_slice], mask=mask, ) @@ -214,5 +217,5 @@ class DataFrameFFCLoaderTestCase(TestCase): args, kwargs = m.call_args assert_array_equal(kwargs['data'], expected_baseline.values) - assert_array_equal(kwargs['mask'], mask.values) + assert_array_equal(kwargs['mask'], mask) self.assertEqual(kwargs['adjustments'], expected_formatted_adjustments) diff --git a/tests/modelling/test_modelling_algo.py b/tests/modelling/test_modelling_algo.py index a1060501..db8750fc 100644 --- a/tests/modelling/test_modelling_algo.py +++ b/tests/modelling/test_modelling_algo.py @@ -7,9 +7,9 @@ from os.path import ( join, realpath, ) - from numpy import ( array, + arange, full_like, nan, ) @@ -17,22 +17,23 @@ from numpy.testing import assert_almost_equal from pandas import ( concat, DataFrame, + date_range, DatetimeIndex, Panel, read_csv, Series, Timestamp, ) -from six import iteritems +from six import iteritems, itervalues from testfixtures import TempDirectory from zipline.algorithm import TradingAlgorithm from zipline.api import ( - # add_filter, add_factor, get_datetime, ) -# from zipline.data.equities import USEquityPricing +from zipline.data.equities import USEquityPricing +from zipline.data.ffc.frame import DataFrameFFCLoader, MULTIPLY from zipline.data.ffc.loaders.us_equity_pricing import ( BcolzDailyBarReader, DailyBarWriterFromCSVs, @@ -40,14 +41,17 @@ from zipline.data.ffc.loaders.us_equity_pricing import ( SQLiteAdjustmentWriter, USEquityPricingLoader, ) -# from zipline.modelling.factor import CustomFactor from zipline.finance import trading + from zipline.modelling.factor.technical import VWAP from zipline.utils.test_utils import ( make_simple_asset_info, str_to_seconds, ) -from zipline.utils.tradingcalendar import trading_days +from zipline.utils.tradingcalendar import ( + trading_day, + trading_days, +) TEST_RESOURCE_PATH = join( @@ -70,6 +74,127 @@ def rolling_vwap(df, length): return Series(out, index=df.index) +class ClosesOnly(TestCase): + + def setUp(self): + self.env = env = trading.TradingEnvironment() + self.dates = date_range( + '2014-01-01', '2014-02-01', freq=trading_day, tz='UTC' + ) + asset_info = DataFrame.from_records([ + { + 'sid': 1, + 'symbol': 'A', + 'asset_type': 'equity', + 'start_date': self.dates[10], + 'end_date': self.dates[13], + 'exchange': 'TEST', + }, + { + 'sid': 2, + 'symbol': 'B', + 'asset_type': 'equity', + 'start_date': self.dates[11], + 'end_date': self.dates[14], + 'exchange': 'TEST', + }, + { + 'sid': 3, + 'symbol': 'C', + 'asset_type': 'equity', + 'start_date': self.dates[12], + 'end_date': self.dates[15], + 'exchange': 'TEST', + }, + ]) + self.first_asset_start = min(asset_info.start_date) + self.last_asset_end = max(asset_info.end_date) + env.write_data(equities_df=asset_info) + self.asset_finder = finder = env.asset_finder + + sids = (1, 2, 3) + self.assets = finder.retrieve_all(sids) + + # View of the baseline data. + self.closes = DataFrame( + {sid: arange(1, len(self.dates) + 1) * sid for sid in sids}, + index=self.dates, + dtype=float, + ) + + # Add a split for 'A' on its second date. + self.split_asset = self.assets[0] + self.split_date = self.split_asset.start_date + trading_day + self.split_ratio = 0.5 + self.adjustments = DataFrame.from_records([ + { + 'sid': self.split_asset.sid, + 'value': self.split_ratio, + 'kind': MULTIPLY, + 'start_date': Timestamp('NaT'), + 'end_date': self.split_date, + 'apply_date': self.split_date, + } + ]) + + # View of the data on/after the split. + self.adj_closes = adj_closes = self.closes.copy() + adj_closes.ix[:self.split_date, self.split_asset] *= self.split_ratio + + self.ffc_loader = DataFrameFFCLoader( + column=USEquityPricing.close, + baseline=self.closes, + adjustments=self.adjustments, + ) + + def expected_close(self, date, asset): + if date < self.split_date: + lookup = self.closes + else: + lookup = self.adj_closes + return lookup.loc[date, asset] + + def exists(self, date, asset): + return asset.start_date <= date <= asset.end_date + + 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') + + def handle_data(context, data): + factors = data.factors + 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'] + self.assertEqual(latest, self.expected_close(date, asset)) + else: + self.assertNotIn(asset, factors.index) + + before_trading_start = handle_data + + 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, + ) + + # Run for a week in the middle of our data. + algo.run(source=self.closes.iloc[10:17]) + + class FFCAlgorithmTestCase(TestCase): @classmethod @@ -155,16 +280,21 @@ class FFCAlgorithmTestCase(TestCase): def test_handle_adjustment(self): AAPL, MSFT, BRK_A = assets = self.AAPL, self.MSFT, self.BRK_A - raw_data = self.raw_data - adjusted_data = {k: v.copy() for k, v in iteritems(raw_data)} - AAPL_split_date = Timestamp("2014-06-09", tz='UTC') - split_loc = raw_data[AAPL].index.get_loc(AAPL_split_date) + # Our view of the data before AAPL's split on June 9, 2014. + raw = {k: v.copy() for k, v in iteritems(self.raw_data)} - # Our view of AAPL's history changes after the split. - ohlc = ['open', 'high', 'low', 'close'] - adjusted_data[AAPL].ix[:split_loc, ohlc] /= 7.0 - adjusted_data[AAPL].ix[:split_loc, ['volume']] *= 7.0 + split_date = Timestamp("2014-06-09", tz='UTC') + split_loc = self.dates.get_loc(split_date) + split_ratio = 7.0 + + # Our view of the data after AAPL's split. All prices from before June + # 9 get divided by the split ratio, and volumes get multiplied by the + # split ratio. + adj = {k: v.copy() for k, v in iteritems(self.raw_data)} + for column in 'open', 'high', 'low', 'close': + adj[AAPL].ix[:split_loc, column] /= split_ratio + adj[AAPL].ix[:split_loc, 'volume'] *= split_ratio window_lengths = [1, 2, 5, 10] # length -> asset -> expected vwap @@ -173,14 +303,56 @@ class FFCAlgorithmTestCase(TestCase): for length in window_lengths: vwap_keys[length] = "vwap_%d" % length for asset in AAPL, MSFT, BRK_A: - raw = rolling_vwap(raw_data[asset], length) - adj = rolling_vwap(adjusted_data[asset], length) + raw_vwap = rolling_vwap(raw[asset], length) + adj_vwap = rolling_vwap(adj[asset], length) + # Shift computed results one day forward so that they're + # labelled by the date on which they'll be seen in the + # algorithm. (We can't show the close price for day N until day + # N + 1.) vwaps[length][asset] = concat( [ - raw[:split_loc], - adj[split_loc:] + raw_vwap[:split_loc - 1], + adj_vwap[split_loc - 1:] ] - ) + ).shift(1, trading_day) + + # Make sure all the expected vwaps have the same dates. + vwap_dates = vwaps[1][self.AAPL].index + for dict_ in itervalues(vwaps): + # Each value is a dict mapping sid -> expected series. + for series in itervalues(dict_): + self.assertTrue((vwap_dates == series.index).all()) + + # Spot check expectations near the AAPL split. + # length 1 vwap for the morning before the split should be the close + # price of the previous day. + before_split = vwaps[1][AAPL].loc[split_date - trading_day] + assert_almost_equal(before_split, 647.3499, decimal=2) + assert_almost_equal( + before_split, + raw[AAPL].loc[split_date - (2 * trading_day), 'close'], + decimal=2, + ) + + # length 1 vwap for the morning of the split should be the close price + # of the previous day, **ADJUSTED FOR THE SPLIT**. + on_split = vwaps[1][AAPL].loc[split_date] + assert_almost_equal(on_split, 645.5700 / split_ratio, decimal=2) + assert_almost_equal( + on_split, + raw[AAPL].loc[split_date - trading_day, 'close'] / split_ratio, + decimal=2, + ) + + # length 1 vwap on the day after the split should be the as-traded + # close on the split day. + after_split = vwaps[1][AAPL].loc[split_date + trading_day] + assert_almost_equal(after_split, 93.69999, decimal=2) + assert_almost_equal( + after_split, + raw[AAPL].loc[split_date, 'close'], + decimal=2, + ) def initialize(context): context.vwaps = [] @@ -195,6 +367,7 @@ class FFCAlgorithmTestCase(TestCase): for asset in assets: computed = factors.loc[asset, key] expected = vwaps[length][asset].loc[today] + # Only having two places of precision here is a bit # unfortunate. assert_almost_equal(computed, expected, decimal=2) diff --git a/tests/modelling/test_numerical_expression.py b/tests/modelling/test_numerical_expression.py index 4c3ba24a..3616c70f 100644 --- a/tests/modelling/test_numerical_expression.py +++ b/tests/modelling/test_numerical_expression.py @@ -65,7 +65,9 @@ class NumericalExpressionTestCase(TestCase): def check_output(self, expr, expected): result = expr._compute( [self.fake_raw_data[input_] for input_ in expr.inputs], - self.mask, + self.mask.index, + self.mask.columns, + self.mask.values, ) check_arrays(result, expected) diff --git a/tests/modelling/test_term.py b/tests/modelling/test_term.py index 5a2ceb34..e79f744f 100644 --- a/tests/modelling/test_term.py +++ b/tests/modelling/test_term.py @@ -19,9 +19,10 @@ from zipline.errors import ( TermInputsNotSpecified, WindowLengthNotSpecified, ) +from zipline.modelling.expression import NUMEXPR_MATH_FUNCS from zipline.modelling.factor import Factor from zipline.modelling.graph import TermGraph -from zipline.modelling.expression import NUMEXPR_MATH_FUNCS +from zipline.modelling.term import AssetExists, NotSpecified class SomeDataSet(DataSet): @@ -54,13 +55,16 @@ def gen_equivalent_factors(): object. """ yield SomeFactor() - yield SomeFactor(inputs=None) + yield SomeFactor(inputs=NotSpecified) yield SomeFactor(SomeFactor.inputs) yield SomeFactor(inputs=SomeFactor.inputs) yield SomeFactor([SomeDataSet.foo, SomeDataSet.bar]) yield SomeFactor(window_length=SomeFactor.window_length) - yield SomeFactor(window_length=None) - yield SomeFactor([SomeDataSet.foo, SomeDataSet.bar], window_length=None) + yield SomeFactor(window_length=NotSpecified) + yield SomeFactor( + [SomeDataSet.foo, SomeDataSet.bar], + window_length=NotSpecified, + ) yield SomeFactor( [SomeDataSet.foo, SomeDataSet.bar], window_length=SomeFactor.window_length, @@ -96,9 +100,10 @@ class DependencyResolutionTestCase(TestCase): resolution_order = list(graph.ordered()) - self.assertEqual(len(resolution_order), 3) + self.assertEqual(len(resolution_order), 4) + self.assertIs(resolution_order[0], AssetExists()) self.assertEqual( - set([resolution_order[0], resolution_order[1]]), + set([resolution_order[1], resolution_order[2]]), set([SomeDataSet.foo, SomeDataSet.bar]), ) self.assertEqual(resolution_order[-1], SomeFactor()) @@ -118,9 +123,14 @@ class DependencyResolutionTestCase(TestCase): resolution_order = list(graph.ordered()) - self.assertEqual(len(resolution_order), 3) + # SomeFactor, its inputs, and AssetExists() + self.assertEqual(len(resolution_order), 4) + + self.assertIs(resolution_order[0], AssetExists()) + self.assertEqual(graph.extra_rows[AssetExists()], 4) + self.assertEqual( - set([resolution_order[0], resolution_order[1]]), + set([resolution_order[1], resolution_order[2]]), set([bar, buzz]), ) self.assertEqual( @@ -141,12 +151,14 @@ class DependencyResolutionTestCase(TestCase): resolution_order = list(graph.ordered()) # bar should only appear once. - self.assertEqual(len(resolution_order), 5) + self.assertEqual(len(resolution_order), 6) indices = { term: resolution_order.index(term) for term in resolution_order } + self.assertEqual(indices[AssetExists()], 0) + # Verify that f1's dependencies will be computed before f1. self.assertLess(indices[SomeDataSet.foo], indices[f1]) self.assertLess(indices[SomeDataSet.bar], indices[f1]) diff --git a/tests/modelling/test_us_equity_pricing_loader.py b/tests/modelling/test_us_equity_pricing_loader.py index f9f0de00..54d87d90 100644 --- a/tests/modelling/test_us_equity_pricing_loader.py +++ b/tests/modelling/test_us_equity_pricing_loader.py @@ -21,6 +21,8 @@ from nose_parameterized import parameterized from numpy import ( arange, datetime64, + float64, + ones, uint32, ) from numpy.testing import ( @@ -31,6 +33,7 @@ from pandas import ( concat, DataFrame, DatetimeIndex, + Int64Index, Timestamp, ) from pandas.util.testing import assert_index_equal @@ -206,8 +209,8 @@ class BcolzDailyBarTestCase(TestCase): def _check_read_results(self, columns, assets, start_date, end_date): table = self.writer.write(self.dest, self.trading_days, self.assets) reader = BcolzDailyBarReader(table) + results = reader.load_raw_arrays(columns, start_date, end_date, assets) dates = self.trading_days_between(start_date, end_date) - results = reader.load_raw_arrays(columns, dates, assets) for column, result in zip(columns, results): assert_array_equal( result, @@ -433,10 +436,14 @@ class USEquityPricingLoaderTestCase(TestCase): self.assertGreaterEqual(eff_date, asset_start) self.assertLessEqual(eff_date, asset_end) - def calendar_days_between(self, start_date, end_date): - return self.calendar_days[ - self.calendar_days.slice_indexer(start_date, end_date) - ] + def calendar_days_between(self, start_date, end_date, shift=0): + slice_ = self.calendar_days.slice_indexer(start_date, end_date) + start = slice_.start + shift + stop = slice_.stop + shift + if start < 0: + raise KeyError(start_date, shift) + + return self.calendar_days[start:stop] def expected_adjustments(self, start_date, end_date): price_adjustments = {} @@ -448,23 +455,20 @@ class USEquityPricingLoaderTestCase(TestCase): for eff_date_secs, ratio, sid in table.itertuples(index=False): eff_date = Timestamp(eff_date_secs, unit='s', tz='UTC') - # The boundary conditions here are subtle. An adjustment with - # an effective date equal to the query start can't have an - # effect because adjustments only the array for dates strictly - # less than the adjustment effective date. - if not (start_date < eff_date <= end_date): + # Ignore adjustments outside the query bounds. + if not (start_date <= eff_date <= end_date): continue eff_date_loc = query_days.get_loc(eff_date) delta = eff_date_loc - start_loc - # Pricing adjusments should be applied on the date + # Pricing adjustments should be applied on the date # corresponding to the effective date of the input data. They # should affect all rows **before** the effective date. price_adjustments.setdefault(delta, []).append( Float64Multiply( first_row=0, - last_row=delta - 1, + last_row=delta, col=sid - 1, value=ratio, ) @@ -474,7 +478,7 @@ class USEquityPricingLoaderTestCase(TestCase): volume_adjustments.setdefault(delta, []).append( Float64Multiply( first_row=0, - last_row=delta - 1, + last_row=delta, col=sid - 1, value=1.0 / ratio, ) @@ -486,7 +490,7 @@ class USEquityPricingLoaderTestCase(TestCase): columns = [USEquityPricing.close, USEquityPricing.volume] query_days = self.calendar_days_between( TEST_QUERY_START, - TEST_QUERY_STOP + TEST_QUERY_STOP, ) adjustments = reader.load_adjustments( @@ -510,6 +514,13 @@ class USEquityPricingLoaderTestCase(TestCase): TEST_QUERY_START, TEST_QUERY_STOP ) + # Our expected results for each day are based on values from the + # previous day. + shifted_query_days = self.calendar_days_between( + TEST_QUERY_START, + TEST_QUERY_STOP, + shift=-1, + ) adjustments = adjustment_reader.load_adjustments( columns, @@ -526,16 +537,18 @@ class USEquityPricingLoaderTestCase(TestCase): closes, volumes = pricing_loader.load_adjusted_array( columns, - DataFrame(True, index=query_days, columns=self.assets), + dates=query_days, + assets=self.assets, + mask=ones((len(query_days), len(self.assets)), dtype=bool), ) expected_baseline_closes = self.bcolz_writer.expected_values_2d( - query_days, + shifted_query_days, self.assets, 'close', ) expected_baseline_volumes = self.bcolz_writer.expected_values_2d( - query_days, + shifted_query_days, self.assets, 'volume', ) @@ -562,19 +575,22 @@ class USEquityPricingLoaderTestCase(TestCase): def apply_adjustments(self, dates, assets, baseline_values, adjustments): min_date, max_date = dates[[0, -1]] - values = baseline_values.copy() + # HACK: Simulate the coercion to float64 we do in adjusted_array. This + # should be removed when AdjustedArray properly supports + # non-floating-point types. + orig_dtype = baseline_values.dtype + values = baseline_values.astype(float64).copy() for eff_date_secs, ratio, sid in adjustments.itertuples(index=False): eff_date = seconds_to_timestamp(eff_date_secs) - if eff_date < min_date or eff_date > max_date: + # Don't apply adjustments that aren't in the current date range. + if eff_date not in dates: continue eff_date_loc = dates.get_loc(eff_date) asset_col = assets.get_loc(sid) - # Apply ratio multiplicatively to the asset column on all rows - # **strictly less** than the adjustment effective date. Note that - # this will be a no-op in the case that the effective date is the - # first entry in dates. - values[:eff_date_loc, asset_col] *= ratio - return values + # Apply ratio multiplicatively to the asset column on all rows less + # than or equal adjustment effective date. + values[:eff_date_loc + 1, asset_col] *= ratio + return values.astype(orig_dtype) def test_read_with_adjustments(self): columns = [USEquityPricing.high, USEquityPricing.volume] @@ -582,6 +598,13 @@ class USEquityPricingLoaderTestCase(TestCase): TEST_QUERY_START, TEST_QUERY_STOP ) + # Our expected results for each day are based on values from the + # previous day. + shifted_query_days = self.calendar_days_between( + TEST_QUERY_START, + TEST_QUERY_STOP, + shift=-1, + ) baseline_reader = BcolzDailyBarReader(self.bcolz_path) adjustment_reader = SQLiteAdjustmentReader(self.db_path) @@ -590,18 +613,20 @@ class USEquityPricingLoaderTestCase(TestCase): adjustment_reader, ) - closes, volumes = pricing_loader.load_adjusted_array( + highs, volumes = pricing_loader.load_adjusted_array( columns, - DataFrame(True, index=query_days, columns=arange(1, 7)), + dates=query_days, + assets=Int64Index(arange(1, 7)), + mask=ones((len(query_days), 6), dtype=bool), ) expected_baseline_highs = self.bcolz_writer.expected_values_2d( - query_days, + shifted_query_days, self.assets, 'high', ) expected_baseline_volumes = self.bcolz_writer.expected_values_2d( - query_days, + shifted_query_days, self.assets, 'volume', ) @@ -609,7 +634,7 @@ class USEquityPricingLoaderTestCase(TestCase): # At each point in time, the AdjustedArrays should yield the baseline # with all adjustments up to that date applied. for windowlen in range(1, len(query_days) + 1): - for offset, window in enumerate(closes.traverse(windowlen)): + for offset, window in enumerate(highs.traverse(windowlen)): baseline = expected_baseline_highs[offset:offset + windowlen] baseline_dates = query_days[offset:offset + windowlen] expected_adjusted_highs = self.apply_adjustments( @@ -627,6 +652,7 @@ class USEquityPricingLoaderTestCase(TestCase): # Apply only splits and invert the ratio. adjustments = SPLITS.copy() adjustments.ratio = 1 / adjustments.ratio + expected_adjusted_volumes = self.apply_adjustments( baseline_dates, self.assets, @@ -641,6 +667,6 @@ class USEquityPricingLoaderTestCase(TestCase): # Verify that we checked up to the longest possible window. with self.assertRaises(WindowLengthTooLong): - closes.traverse(windowlen + 1) + highs.traverse(windowlen + 1) with self.assertRaises(WindowLengthTooLong): volumes.traverse(windowlen + 1) diff --git a/tests/test_assets.py b/tests/test_assets.py index d2efde6d..ebe11e37 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -653,7 +653,12 @@ class AssetFinderTestCase(TestCase): ) for dates in all_subindices(all_dates): - expected_mask = full( + expected_with_start_raw = full( + shape=(len(dates), num_assets), + fill_value=False, + dtype=bool, + ) + expected_no_start_raw = full( shape=(len(dates), num_assets), fill_value=False, dtype=bool, @@ -662,18 +667,28 @@ class AssetFinderTestCase(TestCase): for i, date in enumerate(dates): it = frame[['start_date', 'end_date']].itertuples() for j, start, end in it: + # This way of doing the checks is redundant, but very + # clear. if start <= date <= end: - expected_mask[i, j] = True + expected_with_start_raw[i, j] = True + if start < date: + expected_no_start_raw[i, j] = True - # Filter out columns with all-empty columns. - expected_result = pd.DataFrame( - data=expected_mask, + expected_with_start = pd.DataFrame( + data=expected_with_start_raw, index=dates, columns=frame.index.values, ) + result = finder.lifetimes(dates, include_start_date=True) + assert_frame_equal(result, expected_with_start) - actual_result = finder.lifetimes(dates) - assert_frame_equal(actual_result, expected_result) + expected_no_start = pd.DataFrame( + data=expected_no_start_raw, + index=dates, + columns=frame.index.values, + ) + result = finder.lifetimes(dates, include_start_date=False) + assert_frame_equal(result, expected_no_start) def test_sids(self): # Ensure that the sids property of the AssetFinder is functioning diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 97be99b7..4fce2d7b 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -1363,13 +1363,24 @@ class TradingAlgorithm(object): def compute_factor_matrix(self, start_date): """ - Compute a factor matrix starting at start_date. + Compute a factor matrix containing at least the data necessary to + provide values for `start_date`. + + Loads a factor matrix with data extending from `start_date` until a + year from `start_date`, or until the end of the simulation. """ days = self.trading_environment.trading_days + + # Load data starting from the previous trading day... start_date_loc = days.get_loc(start_date) + + # ...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. 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, diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index 94d8077e..8e9ab400 100644 --- a/zipline/assets/assets.py +++ b/zipline/assets/assets.py @@ -625,7 +625,7 @@ class AssetFinder(object): ('end', '= ? AND effective_date <= ? +""" +cdef dict SID_QUERIES = { + tablename: _SID_QUERY_TEMPLATE.format(tablename) + for tablename in ('splits', 'dividends', 'mergers') +} + +ADJ_QUERY_TEMPLATE = """ +SELECT sid, ratio, effective_date +FROM {0} +WHERE sid IN ({1}) AND effective_date >= {2} AND effective_date <= {3} +""" + +cdef int SQLITE_MAX_IN_STATEMENT = 999 +EPOCH = Timestamp(0, tz='UTC') + +cdef set _get_sids_from_table(object db, + str tablename, + int start_date, + int end_date): + """ + Get the unique sids for all adjustments between start_date and end_date + from table `tablename`. + + Parameters + ---------- + db : sqlite3.connection + tablename : str + start_date : int (seconds since epoch) + end_date : int (seconds since epoch) + + Returns + ------- + sids : set + Set of sets + """ + + cdef object cursor = db.execute( + SID_QUERIES[tablename], + (start_date, end_date), + ) + cdef set out = set() + cdef tuple result + for result in cursor.fetchall(): + PySet_Add(out, result[0]) + return out + + +cdef set _get_split_sids(object db, int start_date, int end_date): + return _get_sids_from_table(db, 'splits', start_date, end_date) + + +cdef set _get_merger_sids(object db, int start_date, int end_date): + return _get_sids_from_table(db, 'mergers', start_date, end_date) + + +cdef set _get_dividend_sids(object db, int start_date, int end_date): + return _get_sids_from_table(db, 'dividends', start_date, end_date) + + +cdef _adjustments(object adjustments_db, + set split_sids, + set merger_sids, + set dividends_sids, + int start_date, + int end_date, + Int64Index_t assets): + + c = adjustments_db.cursor() + + splits_to_query = [str(a) for a in assets if a in split_sids] + splits_results = [] + while splits_to_query: + query_len = min(len(splits_to_query), SQLITE_MAX_IN_STATEMENT) + query_assets = splits_to_query[:query_len] + t= [str(a) for a in query_assets] + statement = ADJ_QUERY_TEMPLATE.format('splits', + ",".join(['?' for _ in query_assets]), start_date, end_date) + c.execute(statement, t) + splits_to_query = splits_to_query[query_len:] + splits_results.extend(c.fetchall()) + + mergers_to_query = [str(a) for a in assets if a in merger_sids] + mergers_results = [] + while mergers_to_query: + query_len = min(len(mergers_to_query), SQLITE_MAX_IN_STATEMENT) + query_assets = mergers_to_query[:query_len] + t= [str(a) for a in query_assets] + statement = ADJ_QUERY_TEMPLATE.format('mergers', + ",".join(['?' for _ in query_assets]), start_date, end_date) + c.execute(statement, t) + mergers_to_query = mergers_to_query[query_len:] + mergers_results.extend(c.fetchall()) + + dividends_to_query = [str(a) for a in assets if a in dividends_sids] + dividends_results = [] + while dividends_to_query: + query_len = min(len(dividends_to_query), SQLITE_MAX_IN_STATEMENT) + query_assets = dividends_to_query[:query_len] + t= [str(a) for a in query_assets] + statement = ADJ_QUERY_TEMPLATE.format('dividends', + ",".join(['?' for _ in query_assets]), start_date, end_date) + c.execute(statement, t) + dividends_to_query = dividends_to_query[query_len:] + dividends_results.extend(c.fetchall()) + + return splits_results, mergers_results, dividends_results + + +cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection + list columns, + DatetimeIndex_t dates, + Int64Index_t assets): + """ + Load a dictionary of Adjustment objects from adjustments_db + + Parameters + ---------- + adjustments_db : sqlite3.Connection + Connection to a sqlite3 table in the format written by + SQLiteAdjustmentWriter. + columns : list[str] + List of column names for which adjustments are needed. + dates : pd.DatetimeIndex + Dates for which adjustments are needed + assets : pd.Int64Index + Assets for which adjustments are needed. + """ + + cdef int start_date = int((dates[0] - EPOCH).total_seconds()) + cdef int end_date = int((dates[-1] - EPOCH).total_seconds()) + + cdef set split_sids = _get_split_sids( + adjustments_db, + start_date, + end_date, + ) + cdef set merger_sids = _get_merger_sids( + adjustments_db, + start_date, + end_date, + ) + cdef set dividend_sids = _get_dividend_sids( + adjustments_db, + start_date, + end_date, + ) + + cdef: + list splits, mergers, dividends + splits, mergers, dividends = _adjustments( + adjustments_db, + split_sids, + merger_sids, + dividend_sids, + start_date, + end_date, + assets, + ) + + cdef list results = [{} for column in columns] + cdef dict asset_ixs = {} # Cache sid lookups here. + cdef: + int sid + double ratio + int eff_date + int date_loc + Py_ssize_t asset_ix + int i + dict col_adjustments + + # splits affect prices and volumes, volumes is the inverse + for sid, ratio, eff_date in splits: + if eff_date < start_date: + continue + date_loc = dates.get_loc( + Timestamp(eff_date, unit='s', tz='UTC'), + # Get the first date **on or after** the effective date. + method='bfill', + ) + + if not PyDict_Contains(asset_ixs, sid): + asset_ixs[sid] = assets.get_loc(sid) + asset_ix = asset_ixs[sid] + + price_adj = Float64Multiply(0, date_loc, asset_ix, ratio) + for i, column in enumerate(columns): + col_adjustments = results[i] + if column != 'volume': + try: + col_adjustments[date_loc].append(price_adj) + except KeyError: + col_adjustments[date_loc] = [price_adj] + else: + volume_adj = Float64Multiply( + 0, date_loc, asset_ix, 1.0 / ratio + ) + try: + col_adjustments[date_loc].append(volume_adj) + except KeyError: + col_adjustments[date_loc] = [volume_adj] + + # mergers affect prices only + for sid, ratio, eff_date in mergers: + if eff_date < start_date: + continue + date_loc = dates.get_loc( + Timestamp(eff_date, unit='s', tz='UTC'), + # Get the first date **on or after** the effective date. + method='bfill', + ) + + if not PyDict_Contains(asset_ixs, sid): + asset_ixs[sid] = assets.get_loc(sid) + asset_ix = asset_ixs[sid] + + adj = Float64Multiply(0, date_loc, asset_ix, ratio) + for i, column in enumerate(columns): + col_adjustments = results[i] + if column != 'volume': + try: + col_adjustments[date_loc].append(adj) + except KeyError: + col_adjustments[date_loc] = [adj] + + # dividends affect prices only + for sid, ratio, eff_date in dividends: + if eff_date < start_date: + continue + date_loc = dates.get_loc( + Timestamp(eff_date, unit='s', tz='UTC'), + # Get the first date **on or after** the effective date. + method='bfill', + ) + + if not PyDict_Contains(asset_ixs, sid): + asset_ixs[sid] = assets.get_loc(sid) + asset_ix = asset_ixs[sid] + + adj = Float64Multiply(0, date_loc, asset_ix, ratio) + for i, column in enumerate(columns): + col_adjustments = results[i] + if column != 'volume': + try: + col_adjustments[date_loc].append(adj) + except KeyError: + col_adjustments[date_loc] = [adj] + + return results diff --git a/zipline/data/ffc/loaders/_equities.pyx b/zipline/data/ffc/loaders/_equities.pyx new file mode 100644 index 00000000..36cce95d --- /dev/null +++ b/zipline/data/ffc/loaders/_equities.pyx @@ -0,0 +1,196 @@ +# +# Copyright 2015 Quantopian, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import bcolz +cimport cython + +from numpy import ( + array, + float64, + intp, + uint32, + zeros, +) +from numpy cimport ( + float64_t, + intp_t, + ndarray, + uint32_t, + uint8_t, +) +from numpy.math cimport NAN + +ctypedef object ctable_t +ctypedef object Timestamp_t +ctypedef object DatetimeIndex_t +ctypedef object Int64Index_t + + +@cython.boundscheck(False) +@cython.wraparound(False) +cpdef _compute_row_slices(dict asset_starts_absolute, + dict asset_ends_absolute, + dict asset_starts_calendar, + intp_t query_start, + intp_t query_end, + Int64Index_t requested_assets): + """ + Core indexing functionality for loading raw data from bcolz. + + Parameters + ---------- + asset_starts_absolute : dict + Dictionary containing the index of the first row of each asset in the + bcolz file from which we will query. + + asset_ends_absolute : dict + Dictionary containing the index of the last row of each asset in the + bcolz file from which we will query. + + asset_starts_calendar : dict + Dictionary containing the index of in our calendar corresponding to the + start date of each asset + + query_start : intp + query_end : intp + Start and end indices in our calendar of the dates for which we're + querying. + + requested_assets : pandas.Int64Index + The assets for which we want to load data. + + For each asset in requested assets, computes three values: + 1.) The index in the raw bcolz data of first row to load. + 2.) The index in the raw bcolz data of the last row to load. + 3.) The index in the dates of our query corresponding to the first row for + each asset. This is non-zero iff the asset's lifetime begins partway + through the requested query dates. + + Returns + ------- + first_rows, last_rows, offsets : 3-tuple of ndarrays + """ + cdef: + intp_t nassets = len(requested_assets) + + # For each sid, we need to compute the following: + ndarray[dtype=intp_t, ndim=1] first_row_a = zeros(nassets, dtype=intp) + ndarray[dtype=intp_t, ndim=1] last_row_a = zeros(nassets, dtype=intp) + ndarray[dtype=intp_t, ndim=1] offset_a = zeros(nassets, dtype=intp) + + # Loop variables. + intp_t i + intp_t asset + intp_t asset_start_data + intp_t asset_end_data + intp_t asset_start_calendar + intp_t asset_end_calendar + + for i, asset in enumerate(requested_assets): + asset_start_data = asset_starts_absolute[asset] + asset_end_data = asset_ends_absolute[asset] + asset_start_calendar = asset_starts_calendar[asset] + asset_end_calendar = ( + asset_start_calendar + (asset_end_data - asset_start_data) + ) + + # If the asset started during the query, then start with the asset's + # first row. + # Otherwise start with the asset's first row + the number of rows + # before the query on which the asset existed. + first_row_a[i] = ( + asset_start_data + max(0, (query_start - asset_start_calendar)) + ) + # If the asset ended during the query, the end with the asset's last + # row. + # Otherwise, end with the asset's last row minus the number of rows + # after the query for which the asset + last_row_a[i] = ( + asset_end_data - max(0, asset_end_calendar - query_end) + ) + # If the asset existed on or before the query, no offset. + # Otherwise, offset by the number of rows in the query in which the + # asset did not yet exist. + offset_a[i] = max(0, asset_start_calendar - query_start) + + return first_row_a, last_row_a, offset_a + + +@cython.boundscheck(False) +@cython.wraparound(False) +cpdef _read_bcolz_data(ctable_t table, + tuple shape, + list columns, + intp_t[:] first_rows, + intp_t[:] last_rows, + intp_t[:] offsets): + """ + Load raw bcolz data for the given columns and indices. + + Parameters + ---------- + table : bcolz.ctable + The table from which to read. + shape : tuple (length 2) + The shape of the expected output arrays. + columns : list[str] + List of column names to read. + + first_rows : ndarray[intp] + last_rows : ndarray[intp] + offsets : ndarray[intp + Arrays in the format returned by _compute_row_slices. + + Returns + ------- + results : list of ndarray + A 2D array of shape `shape` for each column in `columns`. + """ + cdef: + int nassets + str column_name + ndarray[dtype=uint32_t, ndim=1] raw_data + ndarray[dtype=uint32_t, ndim=2] outbuf + ndarray[dtype=uint8_t, ndim=2, cast=True] where_nan + ndarray[dtype=float64_t, ndim=2] outbuf_as_float + intp_t asset + intp_t out_idx + intp_t raw_idx + intp_t first_row + intp_t last_row + intp_t offset + list results = [] + + nassets = shape[1] + if not nassets== len(first_rows) == len(last_rows) == len(offsets): + raise ValueError("Incompatible index arrays.") + + for column_name in columns: + raw_data = table[column_name][:] + outbuf = zeros(shape=shape, dtype=uint32) + for asset in range(nassets): + first_row = first_rows[asset] + last_row = last_rows[asset] + offset = offsets[asset] + for out_idx, raw_idx in enumerate(range(first_row, last_row + 1)): + outbuf[out_idx + offset, asset] = raw_data[raw_idx] + + if column_name in {'open', 'high', 'low', 'close'}: + where_nan = (outbuf == 0) + outbuf_as_float = outbuf.astype(float64) * .001 + outbuf_as_float[where_nan] = NAN + results.append(outbuf_as_float) + else: + results.append(outbuf) + return results diff --git a/zipline/data/ffc/loaders/_us_equity_pricing.pyx b/zipline/data/ffc/loaders/_us_equity_pricing.pyx index d71cdf39..daf8204f 100644 --- a/zipline/data/ffc/loaders/_us_equity_pricing.pyx +++ b/zipline/data/ffc/loaders/_us_equity_pricing.pyx @@ -17,26 +17,12 @@ from cpython cimport ( PySet_Add, ) -import bcolz -cimport cython from numpy import ( - array, - float64, - intp, uint32, zeros, ) -from numpy cimport ( - float64_t, - intp_t, - ndarray, - uint32_t, - uint8_t, -) -from numpy.math cimport NAN from pandas import Timestamp -ctypedef object ctable_t ctypedef object Timestamp_t ctypedef object DatetimeIndex_t ctypedef object Int64Index_t @@ -212,7 +198,6 @@ cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection double ratio int eff_date int date_loc - int last_row Py_ssize_t asset_ix int i dict col_adjustments @@ -224,15 +209,12 @@ cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection # Get the first date **on or after** the effective date. method='bfill', ) - last_row = date_loc - 1 - if last_row < 0: - continue if not PyDict_Contains(asset_ixs, sid): asset_ixs[sid] = assets.get_loc(sid) asset_ix = asset_ixs[sid] - price_adj = Float64Multiply(0, last_row, asset_ix, ratio) + price_adj = Float64Multiply(0, date_loc, asset_ix, ratio) for i, column in enumerate(columns): col_adjustments = results[i] if column != 'volume': @@ -242,7 +224,7 @@ cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection col_adjustments[date_loc] = [price_adj] else: volume_adj = Float64Multiply( - 0, last_row, asset_ix, 1.0 / ratio + 0, date_loc, asset_ix, 1.0 / ratio ) try: col_adjustments[date_loc].append(volume_adj) @@ -256,15 +238,12 @@ cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection # Get the first date **on or after** the effective date. method='bfill', ) - last_row = date_loc - 1 - if last_row < 0: - continue if not PyDict_Contains(asset_ixs, sid): asset_ixs[sid] = assets.get_loc(sid) asset_ix = asset_ixs[sid] - adj = Float64Multiply(0, last_row, asset_ix, ratio) + adj = Float64Multiply(0, date_loc, asset_ix, ratio) for i, column in enumerate(columns): col_adjustments = results[i] if column != 'volume': @@ -280,9 +259,6 @@ cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection # Get the first date **on or after** the effective date. method='bfill', ) - last_row = date_loc - 1 - if last_row <= 0: - continue if not PyDict_Contains(asset_ixs, sid): asset_ixs[sid] = assets.get_loc(sid) @@ -298,162 +274,3 @@ cpdef load_adjustments_from_sqlite(object adjustments_db, # sqlite3.Connection col_adjustments[date_loc] = [adj] return results - - -@cython.boundscheck(False) -@cython.wraparound(False) -cpdef _compute_row_slices(dict asset_starts_absolute, - dict asset_ends_absolute, - dict asset_starts_calendar, - intp_t query_start, - intp_t query_end, - Int64Index_t requested_assets): - """ - Core indexing functionality for loading raw data from bcolz. - - Parameters - ---------- - asset_starts_absolute : dict - Dictionary containing the index of the first row of each asset in the - bcolz file from which we will query. - - asset_ends_absolute : dict - Dictionary containing the index of the last row of each asset in the - bcolz file from which we will query. - - asset_starts_calendar : dict - Dictionary containing the index of in our calendar corresponding to the - start date of each asset - - query_start : intp - query_end : intp - Start and end indices in our calendar of the dates for which we're - querying. - - requested_assets : pandas.Int64Index - The assets for which we want to load data. - - For each asset in requested assets, computes three values: - 1.) The index in the raw bcolz data of first row to load. - 2.) The index in the raw bcolz data of the last row to load. - 3.) The index in the dates of our query corresponding to the first row for - each asset. This is non-zero iff the asset's lifetime begins partway - through the requested query dates. - - Returns - ------- - first_rows, last_rows, offsets : 3-tuple of ndarrays - """ - cdef: - intp_t nassets = len(requested_assets) - - # For each sid, we need to compute the following: - ndarray[dtype=intp_t, ndim=1] first_row_a = zeros(nassets, dtype=intp) - ndarray[dtype=intp_t, ndim=1] last_row_a = zeros(nassets, dtype=intp) - ndarray[dtype=intp_t, ndim=1] offset_a = zeros(nassets, dtype=intp) - - # Loop variables. - intp_t i - intp_t asset - intp_t asset_start_data - intp_t asset_end_data - intp_t asset_start_calendar - intp_t asset_end_calendar - - for i, asset in enumerate(requested_assets): - asset_start_data = asset_starts_absolute[asset] - asset_end_data = asset_ends_absolute[asset] - asset_start_calendar = asset_starts_calendar[asset] - asset_end_calendar = ( - asset_start_calendar + (asset_end_data - asset_start_data) - ) - - # If the asset started during the query, then start with the asset's - # first row. - # Otherwise start with the asset's first row + the number of rows - # before the query on which the asset existed. - first_row_a[i] = ( - asset_start_data + max(0, (query_start - asset_start_calendar)) - ) - # If the asset ended during the query, the end with the asset's last - # row. - # Otherwise, end with the asset's last row minus the number of rows - # after the query for which the asset - last_row_a[i] = ( - asset_end_data - max(0, asset_end_calendar - query_end) - ) - # If the asset existed on or before the query, no offset. - # Otherwise, offset by the number of rows in the query in which the - # asset did not yet exist. - offset_a[i] = max(0, asset_start_calendar - query_start) - - return first_row_a, last_row_a, offset_a - - -@cython.boundscheck(False) -@cython.wraparound(False) -cpdef _read_bcolz_data(ctable_t table, - tuple shape, - list columns, - intp_t[:] first_rows, - intp_t[:] last_rows, - intp_t[:] offsets): - """ - Load raw bcolz data for the given columns and indices. - - Parameters - ---------- - table : bcolz.ctable - The table from which to read. - shape : tuple (length 2) - The shape of the expected output arrays. - columns : list[str] - List of column names to read. - - first_rows : ndarray[intp] - last_rows : ndarray[intp] - offsets : ndarray[intp - Arrays in the format returned by _compute_row_slices. - - Returns - ------- - results : list of ndarray - A 2D array of shape `shape` for each column in `columns`. - """ - cdef: - int nassets - str column_name - ndarray[dtype=uint32_t, ndim=1] raw_data - ndarray[dtype=uint32_t, ndim=2] outbuf - ndarray[dtype=uint8_t, ndim=2, cast=True] where_nan - ndarray[dtype=float64_t, ndim=2] outbuf_as_float - intp_t asset - intp_t out_idx - intp_t raw_idx - intp_t first_row - intp_t last_row - intp_t offset - list results = [] - - nassets = shape[1] - if not nassets== len(first_rows) == len(last_rows) == len(offsets): - raise ValueError("Incompatible index arrays.") - - for column_name in columns: - raw_data = table[column_name][:] - outbuf = zeros(shape=shape, dtype=uint32) - for asset in range(nassets): - first_row = first_rows[asset] - last_row = last_rows[asset] - offset = offsets[asset] - for out_idx, raw_idx in enumerate(range(first_row, last_row + 1)): - outbuf[out_idx + offset, asset] = raw_data[raw_idx] - - if column_name in {'open', 'high', 'low', 'close'}: - where_nan = (outbuf == 0) - outbuf_as_float = outbuf.astype(float64) * .001 - outbuf_as_float[where_nan] = NAN - results.append(outbuf_as_float) - else: - results.append(outbuf) - return results diff --git a/zipline/data/ffc/loaders/us_equity_pricing.py b/zipline/data/ffc/loaders/us_equity_pricing.py index 90b6e17c..95c8facc 100644 --- a/zipline/data/ffc/loaders/us_equity_pricing.py +++ b/zipline/data/ffc/loaders/us_equity_pricing.py @@ -27,7 +27,6 @@ from bcolz import ( from click import progressbar from numpy import ( array, - array_equal, float64, floating, full, @@ -50,11 +49,11 @@ import sqlite3 from zipline.data.ffc.base import FFCLoader -from zipline.data.ffc.loaders._us_equity_pricing import ( +from zipline.data.ffc.loaders._equities import ( _compute_row_slices, _read_bcolz_data, - load_adjustments_from_sqlite, ) +from zipline.data.ffc.loaders._adjustments import load_adjustments_from_sqlite from zipline.lib.adjusted_array import ( adjusted_array, ) @@ -351,48 +350,17 @@ class BcolzDailyBarReader(object): for id_, offset in iteritems(table.attrs['calendar_offset']) } - def _slice_locs(self, start_date, end_date): - try: - start = self._calendar.get_loc(start_date) - except KeyError: - if start_date < self._calendar[0]: - raise NoFurtherDataError( - msg=( - "FFC Query requesting data starting on {query_start}, " - "but first known date is {calendar_start}" - ).format( - query_start=str(start_date), - calendar_start=str(self._calendar[0]), - ) - ) - else: - raise ValueError("Query start %s not in calendar" % start_date) - try: - stop = self._calendar.get_loc(end_date) - except: - if end_date > self._calendar[-1]: - raise NoFurtherDataError( - msg=( - "FFC Query requesting data up to {query_end}, " - "but last known date is {calendar_end}" - ).format( - query_end=end_date, - calendar_end=self._calendar[-1], - ) - ) - else: - raise ValueError("Query end %s not in calendar" % end_date) - return start, stop - - def _compute_slices(self, dates, assets): + def _compute_slices(self, start_idx, end_idx, assets): """ Compute the raw row indices to load for each asset on a query for the - given dates. + given dates after applying a shift. Parameters ---------- - dates : pandas.DatetimeIndex - Dates of the query on which we want to compute row indices. + start_idx : int + Index of first date for which we want data. + end_idx : int + Index of last date for which we want data. assets : pandas.Int64Index Assets for which we want to compute row indices @@ -413,31 +381,29 @@ class BcolzDailyBarReader(object): of a query. Otherwise, offset[i] will be equal to the number of entries in `dates` for which the asset did not yet exist. """ - start, stop = self._slice_locs(dates[0], dates[-1]) - - # Sanity check that the requested date range matches our calendar. - # This could be removed in the future if it's materially affecting - # performance. - query_dates = self._calendar[start:stop + 1] - if not array_equal(query_dates.values, dates.values): - raise ValueError("Incompatible calendars!") - # The core implementation of the logic here is implemented in Cython # for efficiency. return _compute_row_slices( self._first_rows, self._last_rows, self._calendar_offsets, - start, - stop, + start_idx, + end_idx, assets, ) - def load_raw_arrays(self, columns, dates, assets): - first_rows, last_rows, offsets = self._compute_slices(dates, assets) + def load_raw_arrays(self, columns, start_date, end_date, assets): + # Assumes that the given dates are actually in calendar. + start_idx = self._calendar.get_loc(start_date) + end_idx = self._calendar.get_loc(end_date) + first_rows, last_rows, offsets = self._compute_slices( + start_idx, + end_idx, + assets, + ) return _read_bcolz_data( self._table, - (len(dates), len(assets)), + (end_idx - start_idx + 1, len(assets)), [column.name for column in columns], first_rows, last_rows, @@ -617,15 +583,28 @@ class USEquityPricingLoader(FFCLoader): def __init__(self, raw_price_loader, adjustments_loader): self.raw_price_loader = raw_price_loader + # HACK: Pull the calendar off our raw_price_loader so that we can + # backshift dates. + self._calendar = self.raw_price_loader._calendar self.adjustments_loader = adjustments_loader - def load_adjusted_array(self, columns, mask): - dates, assets = mask.index, mask.columns + def load_adjusted_array(self, columns, dates, assets, mask): + # load_adjusted_array is called with dates on which the user's algo + # will be shown data, which means we need to return the data that would + # be known at the start of each date. We assume that the latest data + # known on day N is the data from day (N - 1), so we shift all query + # dates back by a day. + start_date, end_date = _shift_dates( + self._calendar, dates[0], dates[-1], shift=1, + ) + raw_arrays = self.raw_price_loader.load_raw_arrays( columns, - dates, + start_date, + end_date, assets, ) + adjustments = self.adjustments_loader.load_adjustments( columns, dates, @@ -633,6 +612,51 @@ class USEquityPricingLoader(FFCLoader): ) return [ - adjusted_array(raw_array, mask.values, col_adjustments) + adjusted_array(raw_array, mask, col_adjustments) for raw_array, col_adjustments in zip(raw_arrays, adjustments) ] + + +def _shift_dates(dates, start_date, end_date, shift): + try: + start = dates.get_loc(start_date) + except KeyError: + if start_date < dates[0]: + raise NoFurtherDataError( + msg=( + "Modeling Query requested data starting on {query_start}, " + "but first known date is {calendar_start}" + ).format( + query_start=str(start_date), + calendar_start=str(dates[0]), + ) + ) + else: + raise ValueError("Query start %s not in calendar" % start_date) + + # Make sure that shifting doesn't push us out of the calendar. + if start < shift: + raise NoFurtherDataError( + msg=( + "Modeling Query requested data from {shift}" + " days before {query_start}, but first known date is only " + "{start} days earlier." + ).format(shift=shift, query_start=start_date, start=start), + ) + + try: + end = dates.get_loc(end_date) + except KeyError: + if end_date > dates[-1]: + raise NoFurtherDataError( + msg=( + "Modeling Query requesting data up to {query_end}, " + "but last known date is {calendar_end}" + ).format( + query_end=end_date, + calendar_end=dates[-1], + ) + ) + else: + raise ValueError("Query end %s not in calendar" % end_date) + return dates[start - shift], dates[end - shift] diff --git a/zipline/data/ffc/synthetic.py b/zipline/data/ffc/synthetic.py index ef3ba13f..8f008fc1 100644 --- a/zipline/data/ffc/synthetic.py +++ b/zipline/data/ffc/synthetic.py @@ -48,17 +48,17 @@ class MultiColumnLoader(FFCLoader): def __init__(self, loaders): self._loaders = loaders - def load_adjusted_array(self, columns, mask): + def load_adjusted_array(self, columns, dates, assets, mask): """ Load by delegating to sub-loaders. """ out = [] - for column in columns: + for col in columns: try: - loader = self._loaders[column] + loader = self._loaders[col] except KeyError: - raise ValueError("Couldn't find loader for %s" % column) - out.extend(loader.load_adjusted_array([column], mask)) + raise ValueError("Couldn't find loader for %s" % col) + out.extend(loader.load_adjusted_array([col], dates, assets, mask)) return out diff --git a/zipline/errors.py b/zipline/errors.py index 8fd97ba5..ca4c99e3 100644 --- a/zipline/errors.py +++ b/zipline/errors.py @@ -344,6 +344,16 @@ class WindowLengthNotSpecified(ZiplineError): ) +class DTypeNotSpecified(ZiplineError): + """ + Raised if a user attempts to construct a term without specifying dtype and + that term does not have class-level default dtype. + """ + msg = ( + "{termname} requires a dtype, but no dtype was passed." + ) + + class BadPercentileBounds(ZiplineError): """ Raised by API functions accepting percentile bounds when the passed bounds diff --git a/zipline/lib/adjustment.pyx b/zipline/lib/adjustment.pyx index 776bd6a1..245e1616 100644 --- a/zipline/lib/adjustment.pyx +++ b/zipline/lib/adjustment.pyx @@ -43,8 +43,6 @@ cpdef tuple get_adjustment_locs(DatetimeIndex_t dates_index, start_date_loc = dates_index.get_loc(start_date, method='bfill') return ( - # start_date is allowed to be None, indicating "everything - # before the end_date" start_date_loc, # Location of latest date on or before start_date. dates_index.get_loc(end_date, method='ffill'), @@ -135,6 +133,7 @@ cdef class Float64Adjustment: (other.first_row, other.last_row, other.col, other.value) ) + cdef class Float64Multiply(Float64Adjustment): """ An adjustment that multiplies by a scalar. diff --git a/zipline/modelling/engine.py b/zipline/modelling/engine.py index 5a1de7ae..6fe4cc2d 100644 --- a/zipline/modelling/engine.py +++ b/zipline/modelling/engine.py @@ -28,11 +28,13 @@ from pandas import ( from zipline.lib.adjusted_array import ensure_ndarray from zipline.errors import NoFurtherDataError +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)): @@ -49,17 +51,18 @@ class FFCEngine(with_metaclass(ABCMeta)): Parameters ---------- - terms : dict - Map from str -> zipline.modelling.term.Term. - start_date : datetime - The first date of the matrix. - end_date : datetime - The last date of the matrix. + 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. + start_date : pd.Timestamp + Start date of the computed matrix. + end_date : pd.Timestamp + End date of the computed matrix. Returns ------- matrix : pd.DataFrame - A matrix of factors + A matrix of computed results. """ raise NotImplementedError("factor_matrix") @@ -97,6 +100,7 @@ class SimpleFFCEngine(object): '_loader', '_calendar', '_finder', + '_root_mask_term', '__weakref__', ] @@ -104,6 +108,7 @@ class SimpleFFCEngine(object): self._loader = loader self._calendar = calendar self._finder = asset_finder + self._root_mask_term = AssetExists() def factor_matrix(self, terms, start_date, end_date): """ @@ -144,7 +149,7 @@ class SimpleFFCEngine(object): 5. Stick the values computed in (4) into a DataFrame and return it. Step 0 is performed by `zipline.modelling.graph.TermGraph`. - Step 1 is performed in `self.build_lifetimes_matrix`. + Step 1 is performed in `self._compute_root_mask`. Step 2 is performed in `self.compute_chunk`. Steps 3, 4, and 5 are performed in self._format_factor_matrix. @@ -159,20 +164,18 @@ class SimpleFFCEngine(object): ) graph = TermGraph(terms) - max_extra_rows = graph.max_extra_rows + extra_rows = graph.extra_rows[self._root_mask_term] - lifetimes = self.build_lifetimes_matrix( - start_date, - end_date, - max_extra_rows, + 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( + graph, + dates, + assets, + initial_workspace={self._root_mask_term: root_mask_values}, ) - raw_outputs = self.compute_chunk(graph, lifetimes, {}) - lifetimes_between_dates = lifetimes[max_extra_rows:] - dates = lifetimes_between_dates.index.values - assets = lifetimes_between_dates.columns.values - - # We only need filters and factors to compute the final output matrix. + # Collect the results that we'll actually show to the user. filters, factors = {}, {} for name, term in iteritems(terms): if isinstance(term, Filter): @@ -184,12 +187,14 @@ class SimpleFFCEngine(object): else: raise ValueError("Unknown term type: %s" % term) - # Treat base_mask as an implicit filter. - # TODO: Is there a clean way to make this actually just be a filter? - filters['base'] = lifetimes_between_dates.values - return self._format_factor_matrix(dates, assets, filters, factors) + # 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:] - def build_lifetimes_matrix(self, start_date, end_date, extra_rows): + return self._format_factor_matrix(out_dates, assets, filters, factors) + + def _compute_root_mask(self, start_date, end_date, extra_rows): """ Compute a lifetimes matrix from our AssetFinder, then drop columns that didn't exist at all during the query dates. @@ -201,9 +206,9 @@ class SimpleFFCEngine(object): end_date : pd.Timestamp End date for the matrix. extra_rows : int - Number of rows prior to `start_date` to include. + Number of extra rows to compute before `start_date`. Extra rows are needed by terms like moving averages that require a - trailing window of data to compute. + trailing window of data. Returns ------- @@ -230,8 +235,10 @@ class SimpleFFCEngine(object): # Build lifetimes matrix reaching back to `extra_rows` days before # `start_date.` lifetimes = finder.lifetimes( - calendar[start_idx - extra_rows:end_idx] + calendar[start_idx - extra_rows:end_idx], + include_start_date=False ) + assert lifetimes.index[extra_rows] == start_date assert lifetimes.index[-1] == end_date if not lifetimes.columns.unique: @@ -244,36 +251,48 @@ class SimpleFFCEngine(object): existed = lifetimes.iloc[extra_rows:].any() return lifetimes.loc[:, existed] - def _inputs_for_term(self, term, workspace, extra_rows): + def _mask_and_dates_for_term(self, term, workspace, graph, dates): + """ + Load mask and mask row labels for term. + """ + mask = term.mask + offset = graph.extra_rows[mask] - graph.extra_rows[term] + return workspace[mask][offset:], dates[offset:] + + def _inputs_for_term(self, term, workspace, graph): """ Compute inputs for the given term. - This is mostly complicated by the fact that for each input we store - as many rows as will be necessary to serve any term requiring that - input. Thus if Factor A needs 5 extra rows of price, and Factor B - needs 3 extra rows of price, we need to remove 2 leading rows from our - stored prices before passing them to Factor B. + This is mostly complicated by the fact that for each input we store as + many rows as will be necessary to serve **any** computation requiring + that input. """ - term_extra_rows = term.extra_input_rows + offsets = graph.offset if term.windowed: + # If term is windowed, then all input data should be instances of + # AdjustedArray. return [ workspace[input_].traverse( - term.window_length, - offset=extra_rows[input_] - term_extra_rows - ) - for input_ in term.inputs - ] - else: - return [ - ensure_ndarray( - workspace[input_][ - extra_rows[input_] - term_extra_rows: - ], + window_length=term.window_length, + offset=offsets[term, input_] ) for input_ in term.inputs ] - def compute_chunk(self, graph, base_mask, initial_workspace): + # If term is not windowed, input_data may be an AdjustedArray or + # np.ndarray. Coerce the former to the latter. + out = [] + for input_ in term.inputs: + input_data = ensure_ndarray(workspace[input_]) + offset = offsets[term, input_] + # OPTIMIZATION: Don't make a copy by doing input_data[0:] if + # offset is zero. + if offset: + input_data = input_data[offset:] + out.append(input_data) + return out + + def compute_chunk(self, graph, dates, assets, initial_workspace): """ Compute the FFC terms in the graph for the requested start and end dates. @@ -281,53 +300,64 @@ class SimpleFFCEngine(object): Parameters ---------- graph : zipline.modelling.graph.TermGraph + dates : pd.DatetimeIndex + Row labels for our root mask. + assets : pd.Int64Index + Column labels for our root mask. + initial_workspace : dict + Map from term -> output. + Must contain at least entry for `self._root_mask_term` whose shape + is `(len(dates), len(assets))`, but may contain additional + pre-computed terms for testing or optimization purposes. Returns ------- results : dict Dictionary mapping requested results to outputs. """ + self._validate_compute_chunk_params(dates, assets, initial_workspace) loader = self._loader - extra_rows = graph.extra_rows - max_extra_rows = graph.max_extra_rows - - workspace = {} - if initial_workspace is not None: - workspace.update(initial_workspace) + # Copy the supplied initial workspace so we don't mutate it in place. + workspace = initial_workspace.copy() for term in graph.ordered(): - # Subclasses are allowed to pre-populate computed values for terms, - # and in the future we may pre-compute atomic terms coming from the - # same dataset. In both cases, it's possible that we already have - # an entry for this term. + # `term` may have been supplied in `initial_workspace`, and in the + # future we may pre-compute atomic terms coming from the same + # dataset. In either case, we will already have an entry for this + # term, which we shouldn't re-compute. if term in workspace: continue - base_mask_for_term = base_mask.iloc[ - max_extra_rows - extra_rows[term]: - ] + # 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 + ) if term.atomic: # FUTURE OPTIMIZATION: Scan the resolution order for terms in # the same dataset and load them here as well. to_load = [term] loaded = loader.load_adjusted_array( - to_load, - base_mask_for_term, + to_load, mask_dates, assets, mask, ) + assert len(to_load) == len(loaded) for loaded_term, adj_array in zip_longest(to_load, loaded): workspace[loaded_term] = adj_array else: workspace[term] = term._compute( - self._inputs_for_term(term, workspace, extra_rows), - base_mask_for_term, + self._inputs_for_term(term, workspace, graph), + mask_dates, + assets, + mask, ) - assert(workspace[term].shape == base_mask_for_term.shape) + assert(workspace[term].shape == mask.shape) out = {} + graph_extra_rows = graph.extra_rows for name, term in iteritems(graph.outputs): # Truncate off extra rows from outputs. - out[name] = workspace[term][extra_rows[term]:] + out[name] = workspace[term][graph_extra_rows[term]:] return out def _format_factor_matrix(self, dates, assets, filters, factors): @@ -423,3 +453,31 @@ class SimpleFFCEngine(object): ], ) ).tz_localize('UTC', level=0) + + def _validate_compute_chunk_params(self, dates, assets, initial_workspace): + """ + Verify that the values passed to compute_chunk are well-formed. + """ + root = self._root_mask_term + clsname = type(self).__name__ + # Writing this out explicitly so this errors in testing if we change + # the name without updating this line. + compute_chunk_name = self.compute_chunk.__name__ + if root not in initial_workspace: + raise AssertionError( + "root_mask values not supplied to {cls}.{method}".format( + cls=clsname, + method=compute_chunk_name, + ) + ) + + shape = initial_workspace[root].shape + implied_shape = len(dates), len(assets) + if shape != implied_shape: + raise AssertionError( + "root_mask shape is {shape}, but received dates/assets " + "imply that shape should be {implied}".format( + shape=shape, + implied=implied_shape, + ) + ) diff --git a/zipline/modelling/expression.py b/zipline/modelling/expression.py index cdbcf982..63568135 100644 --- a/zipline/modelling/expression.py +++ b/zipline/modelling/expression.py @@ -12,7 +12,7 @@ from numpy import ( ) from six import integer_types -from zipline.modelling.term import Term +from zipline.modelling.term import Term, NotSpecified _VARIABLE_NAME_RE = re.compile("^(x_)([0-9]+)$") @@ -180,7 +180,7 @@ class NumericalExpression(Term): # factor(int64) + factor(int64) + 2.5. # The real fix for this is probably for the calling context to specify # dtypes. - if cls.dtype is not None: + if cls.dtype is not NotSpecified: dtype = cls.dtype else: dtype = find_common_type( @@ -228,7 +228,7 @@ class NumericalExpression(Term): ) return super(NumericalExpression, self)._validate() - def _compute(self, arrays, mask): + def _compute(self, arrays, dates, assets, mask): """ Compute our stored expression string with numexpr. """ diff --git a/zipline/modelling/factor/factor.py b/zipline/modelling/factor/factor.py index 4b5c3b0c..c09ac19d 100644 --- a/zipline/modelling/factor/factor.py +++ b/zipline/modelling/factor/factor.py @@ -16,6 +16,7 @@ from zipline.errors import ( from zipline.lib.rank import rankdata_2d_ordinal from zipline.modelling.term import ( CustomTermMixin, + NotSpecified, RequiredWindowLengthMixin, SingleInputMixin, Term, @@ -184,6 +185,8 @@ class Factor(Term): A transformation yielding a timeseries of scalar values associated with an Asset. """ + dtype = float64 + # Dynamically add functions for creating NumExprFactor/NumExprFilter # instances. clsdict = locals() @@ -219,7 +222,7 @@ class Factor(Term): eq = binary_operator('==') - def rank(self, method='ordinal', ascending=True): + def rank(self, method='ordinal', ascending=True, mask=NotSpecified): """ Construct a new Factor representing the sorted rank of each column within each row. @@ -256,9 +259,9 @@ class Factor(Term): zipline.lib.rank zipline.modelling.factor.Rank """ - return Rank(self if ascending else -self, method=method) + return Rank(self if ascending else -self, method=method, mask=mask) - def top(self, N): + def top(self, N, mask=NotSpecified): """ Construct a Filter matching the top N asset values of self each day. @@ -271,9 +274,9 @@ class Factor(Term): ------- filter : zipline.modelling.filter.Filter """ - return self.rank(ascending=False) <= N + return self.rank(ascending=False, mask=mask) <= N - def bottom(self, N): + def bottom(self, N, mask=NotSpecified): """ Construct a Filter matching the bottom N asset values of self each day. @@ -286,9 +289,12 @@ class Factor(Term): ------- filter : zipline.modelling.filter.Filter """ - return self.rank(ascending=True) <= N + return self.rank(ascending=True, mask=mask) <= N - def percentile_between(self, min_percentile, max_percentile): + def percentile_between(self, + min_percentile, + max_percentile, + mask=NotSpecified): """ Construct a new Filter representing entries from the output of this Factor that fall within the percentile range defined by min_percentile @@ -312,6 +318,7 @@ class Factor(Term): self, min_percentile=min_percentile, max_percentile=max_percentile, + mask=mask, ) @@ -359,15 +366,15 @@ class Rank(SingleInputMixin, Factor): Most users should call Factor.rank rather than directly construct an instance of this class. """ - dtype = float64 window_length = 0 - domain = None + dtype = float64 - def __new__(cls, factor, method): + def __new__(cls, factor, method, mask): return super(Rank, cls).__new__( cls, inputs=(factor,), method=method, + mask=mask, ) def _init(self, method, *args, **kwargs): @@ -392,35 +399,35 @@ class Rank(SingleInputMixin, Factor): ) return super(Rank, self)._validate() - def _compute(self, arrays, mask): + def _compute(self, arrays, dates, assets, mask): """ For each row in the input, compute a like-shaped array of per-row ranks. """ - # OPTIMIZATION: Fast path the default value with our own specialized - # implementation. + inv_mask = ~mask + data = arrays[0].copy() + data[inv_mask] = nan + # OPTIMIZATION: Fast path the default case with our own specialized + # Cython implementation. if self._method == 'ordinal': - result = rankdata_2d_ordinal(arrays[0]) + result = rankdata_2d_ordinal(data) else: # FUTURE OPTIMIZATION: # Write a less general "apply to rows" method that doesn't do all # the extra work that apply_along_axis does. - result = apply_along_axis( - rankdata, - 1, - arrays[0], - method=self._method, - ) - # rankdata will sort nan values into last place, but we want our - # nans to propagate, so explicitly re-apply - result[~mask.values] = nan + result = apply_along_axis(rankdata, 1, data, method=self._method) + + # rankdata will sort nan values into last place, but we want our + # nans to propagate, so explicitly re-apply. + result[inv_mask] = nan return result def __repr__(self): - return "{type}({input_}, method='{method}')".format( + return "{type}({input_}, method='{method}', mask={mask})".format( type=type(self).__name__, input_=self.inputs[0], method=self._method, + mask=self.mask, ) @@ -433,7 +440,6 @@ class CustomFactor(RequiredWindowLengthMixin, CustomTermMixin, Factor): We currently only support CustomFactors of type float64. """ - dtype = float64 ctx = nullctx() def _validate(self): diff --git a/zipline/modelling/factor/latest.py b/zipline/modelling/factor/latest.py new file mode 100644 index 00000000..199074ec --- /dev/null +++ b/zipline/modelling/factor/latest.py @@ -0,0 +1,15 @@ +""" +Factor that produces the most most recently-known value of Column. +""" +from .factor import CustomFactor +from ..term import SingleInputMixin + + +class Latest(SingleInputMixin, CustomFactor): + """ + Factor producing the most recently-known value of `inputs[0]` on each day. + """ + window_length = 1 + + def compute(self, today, assets, out, data): + out[:] = data[-1] diff --git a/zipline/modelling/filter.py b/zipline/modelling/filter.py index 04d0fef4..ac79d1c6 100644 --- a/zipline/modelling/filter.py +++ b/zipline/modelling/filter.py @@ -85,7 +85,6 @@ class Filter(Term): """ A boolean predicate on a universe of Assets. """ - domain = None dtype = bool_ clsdict = locals() @@ -96,41 +95,22 @@ class Filter(Term): } ) - def then(self, other): - """ - Create a new filter by computing `self`, then computing `other` on the - data that survived the first filter. - - Parameters - ---------- - other : zipline.modelling.filter.Filter - The Filter to apply next. - - Returns - ------- - filter : zipline.modelling.filter.SequencedFilter - A filter which will compute `self` and then `other`. - - See Also - -------- - zipline.modelling.filter.SequencedFilter - """ - return SequencedFilter(self, other) - class NumExprFilter(NumericalExpression, Filter): """ A Filter computed from a numexpr expression. """ - def _compute(self, arrays, mask): + def _compute(self, arrays, dates, assets, mask): """ - Compute our result with numexpr, then apply `mask`. + Compute our result with numexpr, then re-apply `mask`. """ return super(NumExprFilter, self)._compute( arrays, + dates, + assets, mask, - ) & mask.values + ) & mask class PercentileFilter(SingleInputMixin, Filter): @@ -148,10 +128,11 @@ class PercentileFilter(SingleInputMixin, Filter): """ window_length = 0 - def __new__(cls, factor, min_percentile, max_percentile): + def __new__(cls, factor, min_percentile, max_percentile, mask): return super(PercentileFilter, cls).__new__( cls, inputs=(factor,), + mask=mask, min_percentile=min_percentile, max_percentile=max_percentile, ) @@ -180,15 +161,15 @@ class PercentileFilter(SingleInputMixin, Filter): ) return super(PercentileFilter, self)._validate() - def _compute(self, arrays, mask): + def _compute(self, arrays, dates, assets, mask): """ For each row in the input, compute a mask of all values falling between the given percentiles. """ # TODO: Review whether there's a better way of handling small numbers # of columns. - data = arrays[0].astype(float64) - data[~mask.values] = nan + data = arrays[0].copy().astype(float64) + data[~mask] = nan # FIXME: np.nanpercentile **should** support computing multiple bounds # at once, but there's a bug in the logic for multiple bounds in numpy @@ -207,66 +188,3 @@ class PercentileFilter(SingleInputMixin, Filter): keepdims=True, ) return (lower_bounds <= data) & (data <= upper_bounds) - - -class SequencedFilter(Filter): - """ - Term representing sequenced computation of two Filters. - - Parameters - ---------- - first : zipline.modelling.filter.Filter - The first filter to compute. - second : zipline.modelling.filter.Filter - The second filter to compute. - - Notes - ----- - In general, users should rarely have to construct SequencedFilter instances - directly. Instead, prefer construction via `Filter.then`. - - See Also - -------- - Filter.then - """ - window_length = 0 - - def __new__(cls, first, then): - return super(SequencedFilter, cls).__new__( - cls, - inputs=concat_tuples((first,), then.inputs), - then=then, - ) - - def _init(self, then, *args, **kwargs): - self._then = then - return super(SequencedFilter, self)._init(*args, **kwargs) - - def _validate(self): - """ - Ensure that we're actually sequencing filters. - """ - first, then = self.inputs[0], self._then - if not isinstance(first, Filter): - raise TypeError("Expected Filter, got %s" % type(first).__name__) - if not isinstance(then, Filter): - raise TypeError("Expected Filter, got %s" % type(then).__name__) - return super(SequencedFilter, self)._validate() - - @classmethod - def static_identity(cls, then, *args, **kwargs): - return ( - super(SequencedFilter, cls).static_identity(*args, **kwargs), - then, - ) - - def _compute(self, arrays, mask): - """ - Call our second filter on its inputs, masking out any inputs rejected - by our first filter. - """ - first_result, then_inputs = arrays[0], arrays[1:] - return self._then._compute( - then_inputs, - mask & first_result, - ) diff --git a/zipline/modelling/graph.py b/zipline/modelling/graph.py index 5ed0150b..5ef4eaa3 100644 --- a/zipline/modelling/graph.py +++ b/zipline/modelling/graph.py @@ -6,6 +6,7 @@ from networkx import ( topological_sort, ) from six import itervalues, iteritems +from zipline.utils.memoize import lazyval class CyclicDependency(Exception): @@ -14,7 +15,7 @@ class CyclicDependency(Exception): class TermGraph(DiGraph): """ - Graph represention of FFC Term dependencies + Graph represention of FFC Term dependencies. Each node in the graph has an `extra_rows` attribute, indicating how many, if any, extra rows we should compute for the node. Extra rows are most @@ -25,13 +26,13 @@ class TermGraph(DiGraph): Parameters ---------- terms : dict - A dict mapping names to terms. + A dict mapping names to final output terms. Attributes ---------- outputs + offset extra_rows - max_extra_rows Methods ------- @@ -47,12 +48,106 @@ class TermGraph(DiGraph): assert not parents self._outputs = terms - self._extra_rows = { + self._ordered = topological_sort(self) + + @lazyval + def offset(self): + """ + For all pairs (term, input) such that `input` is an input to `term`, + compute a mapping: + + (term, input) -> offset(term, input) + + where `offset(term, input)` is defined as + + Max number of extra rows needed by any term depending on `input` + minus + Number of extra rows needed by `term`. + + Example + ------- + + Factor A needs 5 extra rows of USEquityPricing.close, and Factor B + needs 3 extra rows of the same. Factor A also requires 5 extra rows of + USEquityPricing.high, which no other Factor users. + + We load 5 extra rows of both `price` and `high` to ensure we can + service Factor A, and the following offsets get computed: + + self.offset[Factor A, USEquityPricing.close] == 0 + self.offset[Factor A, USEquityPricing.high] == 0 + self.offset[Factor B, USEquityPricing.close] == 2 + self.offset[Factor B, USEquityPricing.high] raises KeyError. + + Notes + ----- + `offset(term, input) >= 0` for all valid pairs, since `input` must be + an input to `term` if the pair appears in the mapping. + + This value is useful because we load enough rows of each input to serve + all possible dependencies. However, for any given dependency, we only + want to compute using the actual number of required extra rows for that + dependency. We can do so by truncating off the first `offset` rows of + the loaded data for `input`. + + See Also + -------- + zipline.modelling.graph.TermGraph.offset + zipline.modelling.engine.SimpleFFCEngine._inputs_for_term + zipline.modelling.engine.SimpleFFCEngine._mask_for_term + """ + out = {} + for term in self: + extra_input_rows = term.extra_input_rows + for input_ in term.inputs: + out[term, input_] = self.extra_rows[input_] - extra_input_rows + mask = term.mask + if term.mask is not None: + out[term, mask] = self.extra_rows[mask] - extra_input_rows + return out + + @lazyval + def extra_rows(self): + """ + A dict mapping `term` -> `# of extra rows to load/compute of `term`. + + This is always the maximum number of extra **input** rows required by + any Filter/Factor for which `term` is an input. + + Notes + ---- + This value depends on the other terms in the graph that require `term` + **as an input**. This is not to be confused with + `term.extra_input_rows`, which is how many extra rows of `term`'s + inputs we need to load, and which is determined entirely by `Term` + itself. + + Example + ------- + Our graph contains the following terms: + + A = SimpleMovingAverage([USEquityPricing.high], window_length=5) + B = SimpleMovingAverage([USEquityPricing.high], window_length=10) + C = SimpleMovingAverage([USEquityPricing.low], window_length=8) + + To compute N rows of A, we need N + 4 extra rows of `high`. + To compute N rows of B, we need N + 9 extra rows of `high`. + To compute N rows of C, we need N + 7 extra rows of `low`. + + We store the following extra_row requirements: + + self.extra_rows[high] = 9 # Ensures that we can service B. + self.extra_rows[low] = 7 + + See Also + -------- + zipline.modelling.graph.TermGraph.offset + zipline.modelling.term.Term.extra_input_rows + """ + return { term: attrs['extra_rows'] for term, attrs in iteritems(self.node) } - self._max_extra_rows = max(itervalues(self._extra_rows)) - self._ordered = topological_sort(self) @property def outputs(self): @@ -61,20 +156,6 @@ class TermGraph(DiGraph): """ return self._outputs - @property - def extra_rows(self): - """ - Dict mapping term -> number of extra rows to compute for term. - """ - return self._extra_rows - - @property - def max_extra_rows(self): - """ - Maximum number of extra rows required to compute any term in the graph. - """ - return self._max_extra_rows - def ordered(self): """ Return a topologically-sorted iterator over the terms in `self`. @@ -92,25 +173,35 @@ class TermGraph(DiGraph): raise CyclicDependency(term) parents.add(term) - try: - existing = self.node[term] - except KeyError: - # `term` is not yet in the graph: add it with the specified number - # of extra rows. - self.add_node(term, extra_rows=extra_rows) - else: - # `term` is already in the graph because we've been traversed by - # another parent. Ensure that we have enough extra rows to satisfy - # all of our parents. - existing['extra_rows'] = max(extra_rows, existing['extra_rows']) + # Idempotent if term is already in the graph. + self.add_node(term) - extra_rows_for_subterms = extra_rows + term.extra_input_rows - for subterm in term.inputs: + # Make sure we're going to compute at least `extra_rows` of `term`. + self._ensure_extra_rows(term, extra_rows) + + # Number of extra rows we need to compute for this term's dependencies. + dependency_extra_rows = extra_rows + term.extra_input_rows + + # Recursively add dependencies. + for dependency in term.inputs: self._add_to_graph( - subterm, + dependency, parents, - extra_rows=extra_rows_for_subterms + extra_rows=dependency_extra_rows, ) - self.add_edge(subterm, term) + self.add_edge(dependency, term) + + # Add term's mask, which is really just a specially-enumerated input. + mask = term.mask + if mask is not None: + self._add_to_graph(mask, parents, extra_rows=dependency_extra_rows) + self.add_edge(mask, term) parents.remove(term) + + def _ensure_extra_rows(self, term, N): + """ + Ensure that we're going to compute at least N extra rows of `term`. + """ + attrs = self.node[term] + attrs['extra_rows'] = max(N, attrs.get('extra_rows', 0)) diff --git a/zipline/modelling/term.py b/zipline/modelling/term.py index e7b18798..9204f98d 100644 --- a/zipline/modelling/term.py +++ b/zipline/modelling/term.py @@ -1,15 +1,12 @@ """ Base class for Filters, Factors and Classifiers """ -from numpy import ( - empty, - float64, - full, - nan, -) from weakref import WeakValueDictionary +from numpy import bool_, full, nan + from zipline.errors import ( + DTypeNotSpecified, InputTermNotAtomic, TermInputsNotSpecified, WindowLengthNotPositive, @@ -18,25 +15,48 @@ from zipline.errors import ( from zipline.utils.memoize import lazyval -NotSpecified = (object(),) +@object.__new__ # bind a single instance to the name 'NotSpecified' +class NotSpecified(object): + """ + Singleton sentinel value used for Term defaults. + """ + __slots__ = ('__weakref__',) + + def __new__(cls): + raise TypeError("Can't construct new instances of NotSpecified") + + def __repr__(self): + return type(self).__name__ + + def __reduce__(self): + return type(self).__name__ + + def __deepcopy__(self, _memo): + return self + + def __copy__(self): + return self class Term(object): """ Base class for terms in an FFC API compute graph. """ + # These are NotSpecified because a subclass is required to provide them. inputs = NotSpecified window_length = NotSpecified - domain = None - dtype = float64 + dtype = NotSpecified + mask = NotSpecified + domain = NotSpecified _term_cache = WeakValueDictionary() def __new__(cls, - inputs=None, - window_length=None, - domain=None, - dtype=None, + inputs=NotSpecified, + mask=NotSpecified, + window_length=NotSpecified, + domain=NotSpecified, + dtype=NotSpecified, *args, **kwargs): """ @@ -49,22 +69,35 @@ class Term(object): Caching previously-constructed Terms is **sane** because terms and their inputs are both conceptually immutable. """ - if inputs is None: - inputs = tuple(cls.inputs) - else: + # Class-level attributes can be used to provide defaults for Term + # subclasses. + + if inputs is NotSpecified: + inputs = cls.inputs + # Having inputs = NotSpecified is an error, but we handle it later + # in self._validate rather than here. + if inputs is not NotSpecified: + # Allow users to specify lists as class-level defaults, but + # normalize to a tuple so that inputs is hashable. inputs = tuple(inputs) - if window_length is None: + if mask is NotSpecified: + mask = cls.mask + if mask is NotSpecified: + mask = AssetExists() + + if window_length is NotSpecified: window_length = cls.window_length - if domain is None: + if domain is NotSpecified: domain = cls.domain - if dtype is None: + if dtype is NotSpecified: dtype = cls.dtype identity = cls.static_identity( inputs=inputs, + mask=mask, window_length=window_length, domain=domain, dtype=dtype, @@ -77,6 +110,7 @@ class Term(object): new_instance = cls._term_cache[identity] = \ super(Term, cls).__new__(cls)._init( inputs=inputs, + mask=mask, window_length=window_length, domain=domain, dtype=dtype, @@ -100,8 +134,9 @@ class Term(object): """ pass - def _init(self, inputs, window_length, domain, dtype): + def _init(self, inputs, mask, window_length, domain, dtype): self.inputs = inputs + self.mask = mask self.window_length = window_length self.domain = domain self.dtype = dtype @@ -110,7 +145,7 @@ class Term(object): return self @classmethod - def static_identity(cls, inputs, window_length, domain, dtype): + def static_identity(cls, inputs, mask, window_length, domain, dtype): """ Return the identity of the Term that would be constructed from the given arguments. @@ -122,7 +157,7 @@ class Term(object): This is a classmethod so that it can be called from Term.__new__ to determine whether to produce a new instance. """ - return (cls, inputs, window_length, domain, dtype) + return (cls, inputs, mask, window_length, domain, dtype) def _validate(self): """ @@ -133,6 +168,11 @@ class Term(object): raise TermInputsNotSpecified(termname=type(self).__name__) if self.window_length is NotSpecified: raise WindowLengthNotSpecified(termname=type(self).__name__) + if self.dtype is NotSpecified: + raise DTypeNotSpecified(termname=type(self).__name__) + if self.mask is NotSpecified and not self.atomic: + # This isn't user error, this is a bug in our code. + raise AssertionError("{term} has no mask".format(term=self)) if self.window_length: for child in self.inputs: @@ -172,7 +212,7 @@ class Term(object): """ return max(0, self.window_length - 1) - def _compute(self, inputs, mask): + def _compute(self, inputs, dates, assets, mask): """ Subclasses should implement this to perform actual computation. @@ -188,6 +228,7 @@ class Term(object): type=type(self).__name__, inputs=self.inputs, window_length=self.window_length, + mask=self.mask, ) @@ -209,9 +250,9 @@ class SingleInputMixin(object): class RequiredWindowLengthMixin(object): def _validate(self): - if self.windowed or self.window_length is NotSpecified: - return super(RequiredWindowLengthMixin, self)._validate() - raise WindowLengthNotPositive(window_length=self.window_length) + if not self.windowed: + raise WindowLengthNotPositive(window_length=self.window_length) + return super(RequiredWindowLengthMixin, self)._validate() class CustomTermMixin(object): @@ -230,14 +271,13 @@ class CustomTermMixin(object): """ raise NotImplementedError() - def _compute(self, windows, mask): + def _compute(self, windows, dates, assets, mask): """ Call the user's `compute` function on each window with a pre-built output array. """ # TODO: Make mask available to user's `compute`. compute = self.compute - dates, assets = mask.index, mask.columns out = full(mask.shape, nan, dtype=self.dtype) with self.ctx: # TODO: Consider pre-filtering columns that are all-nan at each @@ -249,5 +289,32 @@ class CustomTermMixin(object): out[idx], *(next(w) for w in windows) ) - out[~mask.values] = nan + out[~mask] = nan return out + + +class AssetExists(Term): + """ + Pseudo-filter describing whether or not an asset existed on a given day. + This is the default mask for all terms that haven't been passed a mask + explicitly. + + This is morally a Filter, in the sense that it produces a boolean value for + every asset on every date. We don't subclass Filter, however, because + `AssetExists` is computed directly by the FFCEngine. + + See Also + -------- + zipline.assets.AssetFinder.lifetimes + """ + inputs = () + dtype = bool_ + window_length = 0 + mask = None + + def _compute(self, *args, **kwargs): + # TODO: Consider moving the bulk of the logic from + # SimpleFFCEngine._compute_root_mask here. + raise NotImplementedError( + "Direct computation of AssetExists is not supported!" + ) diff --git a/zipline/modelling/visualize.py b/zipline/modelling/visualize.py index 7839888f..d5c9e079 100644 --- a/zipline/modelling/visualize.py +++ b/zipline/modelling/visualize.py @@ -5,14 +5,14 @@ from functools import partial from contextlib import contextmanager from logbook import Logger, StderrHandler from networkx import topological_sort -from six import iteritems, itervalues +from six import iteritems import subprocess from zipline.data.dataset import BoundColumn from zipline.modelling import Filter, Factor, Classifier, Term +from zipline.modelling.term import AssetExists from zipline.modelling.graph import TermGraph - logger = Logger('Visualize') @@ -70,7 +70,7 @@ def roots(g): return set(n for n, d in iteritems(g.in_degree()) if d == 0) -def write_graph(g, filename, formats=('svg',)): +def write_graph(g, filename, formats=('svg',), include_asset_exists=False): """ Write the dependency graph of `terms` as a dot graph. @@ -95,6 +95,8 @@ def write_graph(g, filename, formats=('svg',)): # Write inputs cluster. with cluster(f, 'Input', **cluster_attrs): for term in in_nodes: + if term is AssetExists() and not include_asset_exists: + continue add_term_node(f, term) # Write intermediate results. @@ -105,6 +107,8 @@ def write_graph(g, filename, formats=('svg',)): # Write edges for source, dest in g.edges(): + if source is AssetExists() and not include_asset_exists: + continue add_edge(f, id(source), id(dest)) outs = [] @@ -116,7 +120,7 @@ def write_graph(g, filename, formats=('svg',)): return outs -def show_graph(g): +def show_graph(g, include_asset_exists=False): """ Display a TermGraph interactively with IPython """ @@ -124,7 +128,12 @@ def show_graph(g): from IPython.display import SVG except ImportError: raise Exception("IPython is not installed. Can't show term graph.") - result = write_graph(g, 'temp', ('svg',))[0] + result = write_graph( + g, + 'temp', + ('svg',), + include_asset_exists=include_asset_exists, + )[0] return SVG(filename=result) diff --git a/zipline/protocol.py b/zipline/protocol.py index 26aa7832..c8f80b10 100644 --- a/zipline/protocol.py +++ b/zipline/protocol.py @@ -504,7 +504,14 @@ class BarData(object): if today > self._factor_matrix_expires: self._factor_matrix, self._factor_matrix_expires = \ algo.compute_factor_matrix(today) - return self._factor_matrix.loc[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/pandas_utils.py b/zipline/utils/pandas_utils.py new file mode 100644 index 00000000..e488ce84 --- /dev/null +++ b/zipline/utils/pandas_utils.py @@ -0,0 +1,12 @@ +""" +Utilities for working with pandas objects. +""" + + +def explode(df): + """ + Take a DataFrame and return a triple of + + (df.index, df.columns, df.values) + """ + return df.index, df.columns, df.values