diff --git a/.travis.yml b/.travis.yml index 58bcb62e..5a7d2414 100644 --- a/.travis.yml +++ b/.travis.yml @@ -19,8 +19,13 @@ install: - conda create -n testenv --yes pip python=$TRAVIS_PYTHON_VERSION - source activate testenv - conda install --yes -c https://conda.binstar.org/Quantopian numpy=$NUMPY_VERSION pandas=$PANDAS_VERSION scipy matplotlib Cython patsy statsmodels tornado pyparsing xlrd mock pytz requests six dateutil ta-lib logbook + - grep bottleneck== etc/requirements.txt | xargs pip install - grep cyordereddict== etc/requirements.txt | xargs pip install - grep contextlib2== etc/requirements.txt | xargs pip install + - grep click== etc/requirements.txt | xargs pip install + - grep networkx== etc/requirements.txt | xargs pip install + - grep numexpr== etc/requirements.txt | xargs pip install + - grep bcolz== etc/requirements.txt | xargs pip install - grep pyflakes== etc/requirements_dev.txt | xargs pip install - grep pep8== etc/requirements_dev.txt | xargs pip install - grep mccabe== etc/requirements_dev.txt | xargs pip install @@ -28,6 +33,7 @@ install: - grep nose== etc/requirements_dev.txt | xargs pip install --upgrade --force-reinstall - grep nose-parameterized== etc/requirements_dev.txt | xargs pip install - grep nose-ignore-docstring== etc/requirements_dev.txt | xargs pip install + - grep testfixtures== etc/requirements_dev.txt | xargs pip install - pip install coveralls - pip install nose-timer - python setup.py build_ext --inplace diff --git a/etc/requirements.txt b/etc/requirements.txt index e274c0f5..a6e74816 100644 --- a/etc/requirements.txt +++ b/etc/requirements.txt @@ -33,3 +33,15 @@ cyordereddict==0.2.2 bottleneck==1.0.0 contextlib2==0.4.0 + +# Graph algorithms used by zipline.modelling +networkx==1.9.1 + +# NumericalExpression modelling terms. +numexpr==2.4.3 + +# On disk storage format for modelling data. +bcolz==0.9.0 + +# Command line interface helper +click==4.0.0 diff --git a/etc/requirements_dev.txt b/etc/requirements_dev.txt index 178b7be6..cf2fa5b0 100644 --- a/etc/requirements_dev.txt +++ b/etc/requirements_dev.txt @@ -19,6 +19,9 @@ pbr==1.3.0 mock==1.3.0 +# Temp Directories for testing +testfixtures==4.1.2 + # Linting flake8==2.4.1 diff --git a/setup.py b/setup.py index 7738b678..32ccf686 100644 --- a/setup.py +++ b/setup.py @@ -24,6 +24,21 @@ ext_modules = [ ['zipline/assets/_assets.pyx'], include_dirs=[np.get_include()], ), + Extension( + 'zipline.lib.adjusted_array', + ['zipline/lib/adjusted_array.pyx'], + include_dirs=[np.get_include()], + ), + Extension( + 'zipline.lib.adjustment', + ['zipline/lib/adjustment.pyx'], + include_dirs=[np.get_include()], + ), + Extension( + 'zipline.data.ffc.loaders._us_equity_pricing', + ['zipline/data/ffc/loaders/_us_equity_pricing.pyx'], + include_dirs=[np.get_include()], + ), ] setup( diff --git a/tests/modelling/__init__.py b/tests/modelling/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/modelling/test_adjusted_array.py b/tests/modelling/test_adjusted_array.py new file mode 100644 index 00000000..4f94f27d --- /dev/null +++ b/tests/modelling/test_adjusted_array.py @@ -0,0 +1,325 @@ +""" +Tests for chunked adjustments. +""" +from unittest import TestCase + +from nose_parameterized import parameterized +from numpy import ( + arange, + array, + full, +) +from numpy.testing import assert_array_equal +from six.moves import zip_longest + +from zipline.lib.adjustment import ( + Float64Multiply, + Float64Overwrite, +) +from zipline.lib.adjusted_array import ( + adjusted_array, + NOMASK, +) +from zipline.errors import ( + WindowLengthNotPositive, + WindowLengthTooLong, +) + + +def num_windows_of_length_M_on_buffers_of_length_N(M, N): + """ + For a window of length M rolling over a buffer of length N, + there are (N - M) + 1 legal windows. + + Example: + If my array has N=4 rows, and I want windows of length M=2, there are + 3 legal windows: data[0:2], data[1:3], and data[2:4]. + """ + return N - M + 1 + + +def valid_window_lengths(underlying_buffer_length): + """ + An iterator of all legal window lengths on a buffer of a given length. + + Returns values from 1 to underlying_buffer_length. + """ + return iter(range(1, underlying_buffer_length + 1)) + + +def _gen_unadjusted_cases(dtype): + + nrows = 6 + ncols = 3 + data = arange(nrows * ncols, dtype=dtype).reshape(nrows, ncols) + + for windowlen in valid_window_lengths(nrows): + + num_legal_windows = num_windows_of_length_M_on_buffers_of_length_N( + windowlen, nrows + ) + + yield ( + "length_%d" % windowlen, + data, + windowlen, + {}, + [ + data[offset:offset + windowlen] + for offset in range(num_legal_windows) + ], + ) + + +def _gen_multiplicative_adjustment_cases(dtype): + """ + Generate expected moving windows on a buffer with adjustments. + + We proceed by constructing, at each row, the view of the array we expect in + in all windows anchored on or after that row. + + In general, if we have an adjustment to be applied once we process the row + at index N, should see that adjustment applied to the underlying buffer for + any window containing the row at index N. + + We then build all legal windows over these buffers. + """ + adjustment_type = { + float: Float64Multiply, + }[dtype] + + nrows, ncols = 6, 3 + adjustments = {} + buffer_as_of = [None] * 6 + baseline = buffer_as_of[0] = full((nrows, ncols), 1, dtype=dtype) + + # Note that row indices are inclusive! + adjustments[1] = [ + adjustment_type(0, 0, 0, dtype(2)), + ] + buffer_as_of[1] = array([[2, 1, 1], + [1, 1, 1], + [1, 1, 1], + [1, 1, 1], + [1, 1, 1], + [1, 1, 1]], dtype=dtype) + + # No adjustment at index 2. + buffer_as_of[2] = buffer_as_of[1] + + adjustments[3] = [ + adjustment_type(1, 2, 1, dtype(3)), + adjustment_type(0, 1, 0, dtype(4)), + ] + buffer_as_of[3] = array([[8, 1, 1], + [4, 3, 1], + [1, 3, 1], + [1, 1, 1], + [1, 1, 1], + [1, 1, 1]], dtype=dtype) + + adjustments[4] = [ + adjustment_type(0, 3, 2, dtype(5)) + ] + buffer_as_of[4] = array([[8, 1, 5], + [4, 3, 5], + [1, 3, 5], + [1, 1, 5], + [1, 1, 1], + [1, 1, 1]], dtype=dtype) + + adjustments[5] = [ + adjustment_type(0, 4, 1, dtype(6)), + adjustment_type(2, 2, 2, dtype(7)), + ] + buffer_as_of[5] = array([[8, 6, 5], + [4, 18, 5], + [1, 18, 35], + [1, 6, 5], + [1, 6, 1], + [1, 1, 1]], dtype=dtype) + + return _gen_expectations(baseline, adjustments, buffer_as_of, nrows) + + +def _gen_overwrite_adjustment_cases(dtype): + """ + Generate test cases for overwrite adjustments. + + The algorithm used here is the same as the one used above for + multiplicative adjustments. The only difference is the semantics of how + the adjustments are expected to modify the arrays. + """ + + adjustment_type = { + float: Float64Overwrite, + }[dtype] + + nrows, ncols = 6, 3 + adjustments = {} + buffer_as_of = [None] * 6 + baseline = buffer_as_of[0] = full((nrows, ncols), 2, dtype=dtype) + + # Note that row indices are inclusive! + adjustments[1] = [ + adjustment_type(0, 0, 0, dtype(1)), + ] + buffer_as_of[1] = array([[1, 2, 2], + [2, 2, 2], + [2, 2, 2], + [2, 2, 2], + [2, 2, 2], + [2, 2, 2]], dtype=dtype) + + # No adjustment at index 2. + buffer_as_of[2] = buffer_as_of[1] + + adjustments[3] = [ + adjustment_type(1, 2, 1, dtype(3)), + adjustment_type(0, 1, 0, dtype(4)), + ] + buffer_as_of[3] = array([[4, 2, 2], + [4, 3, 2], + [2, 3, 2], + [2, 2, 2], + [2, 2, 2], + [2, 2, 2]], dtype=dtype) + + adjustments[4] = [ + adjustment_type(0, 3, 2, dtype(5)) + ] + buffer_as_of[4] = array([[4, 2, 5], + [4, 3, 5], + [2, 3, 5], + [2, 2, 5], + [2, 2, 2], + [2, 2, 2]], dtype=dtype) + + adjustments[5] = [ + adjustment_type(0, 4, 1, dtype(6)), + adjustment_type(2, 2, 2, dtype(7)), + ] + buffer_as_of[5] = array([[4, 6, 5], + [4, 6, 5], + [2, 6, 7], + [2, 6, 5], + [2, 6, 2], + [2, 2, 2]], dtype=dtype) + + return _gen_expectations( + baseline, + adjustments, + buffer_as_of, + nrows, + ) + + +def _gen_expectations(baseline, adjustments, buffer_as_of, nrows): + + for windowlen in valid_window_lengths(nrows): + + num_legal_windows = num_windows_of_length_M_on_buffers_of_length_N( + windowlen, nrows + ) + + yield ( + "length_%d" % windowlen, + baseline, + windowlen, + adjustments, + [ + # This is a nasty expression... + # + # Reading from right to left: we want a slice of length + # 'windowlen', starting at 'offset', from the buffer on which + # we've applied all adjustments corresponding to the last row + # of the data, which will be (offset + windowlen - 1). + buffer_as_of[offset + windowlen - 1][offset:offset + windowlen] + for offset in range(num_legal_windows) + ], + ) + + +class AdjustedArrayTestCase(TestCase): + + @parameterized.expand(_gen_unadjusted_cases(float)) + def test_no_adjustments(self, + name, + data, + lookback, + adjustments, + expected): + array = adjusted_array( + data, + NOMASK, + adjustments, + ) + for _ in range(2): # Iterate 2x ensure adjusted_arrays are re-usable. + window_iter = array.traverse(lookback) + for yielded, expected_yield in zip_longest(window_iter, expected): + assert_array_equal(yielded, expected_yield) + + @parameterized.expand(_gen_multiplicative_adjustment_cases(float)) + def test_multiplicative_adjustments(self, + name, + data, + lookback, + adjustments, + expected): + array = adjusted_array( + data, + NOMASK, + adjustments, + ) + for _ in range(2): # Iterate 2x ensure adjusted_arrays are re-usable. + window_iter = array.traverse(lookback) + for yielded, expected_yield in zip_longest(window_iter, expected): + assert_array_equal(yielded, expected_yield) + + @parameterized.expand(_gen_overwrite_adjustment_cases(float)) + def test_overwrite_adjustment_cases(self, + name, + data, + lookback, + adjustments, + expected): + array = adjusted_array( + data, + NOMASK, + adjustments, + ) + for _ in range(2): # Iterate 2x ensure adjusted_arrays are re-usable. + window_iter = array.traverse(lookback) + for yielded, expected_yield in zip_longest(window_iter, expected): + assert_array_equal(yielded, expected_yield) + + def test_invalid_lookback(self): + + data = arange(30, dtype=float).reshape(6, 5) + adj_array = adjusted_array(data, NOMASK, {}) + + with self.assertRaises(WindowLengthTooLong): + adj_array.traverse(7) + + with self.assertRaises(WindowLengthNotPositive): + adj_array.traverse(0) + + with self.assertRaises(WindowLengthNotPositive): + adj_array.traverse(-1) + + def test_array_views_arent_writable(self): + + data = arange(30, dtype=float).reshape(6, 5) + adj_array = adjusted_array(data, NOMASK, {}) + + for frame in adj_array.traverse(3): + with self.assertRaises(ValueError): + frame[0, 0] = 5.0 + + def test_bad_input(self): + msg = "Mask shape \(2, 3\) != data shape \(5, 5\)" + data = arange(25).reshape(5, 5) + bad_mask = array([[0, 1, 1], [0, 0, 1]], dtype=bool) + + with self.assertRaisesRegexp(ValueError, msg): + adjusted_array(data, bad_mask, {}) diff --git a/tests/modelling/test_engine.py b/tests/modelling/test_engine.py new file mode 100644 index 00000000..41c1752d --- /dev/null +++ b/tests/modelling/test_engine.py @@ -0,0 +1,422 @@ +""" +Tests for SimpleFFCEngine +""" +from __future__ import division +from unittest import TestCase + +from numpy import ( + full, + isnan, + nan, +) +from numpy.testing import assert_array_equal +from pandas import ( + DataFrame, + date_range, + Int64Index, + rolling_mean, + Timestamp, +) +from pandas.util.testing import assert_frame_equal +from testfixtures import TempDirectory + +from zipline.assets import AssetFinder +from zipline.data.equities import USEquityPricing +from zipline.data.ffc.synthetic import ( + ConstantLoader, + MultiColumnLoader, + NullAdjustmentReader, + SyntheticDailyBarWriter, +) +from zipline.data.ffc.frame import ( + DataFrameFFCLoader, + MULTIPLY, +) +from zipline.data.ffc.loaders.us_equity_pricing import ( + BcolzDailyBarReader, + USEquityPricingLoader, +) +from zipline.finance.trading import TradingEnvironment +from zipline.modelling.engine import SimpleFFCEngine +from zipline.modelling.factor import TestingFactor +from zipline.modelling.factor.technical import ( + MaxDrawdown, + SimpleMovingAverage, +) +from zipline.utils.lazyval import lazyval +from zipline.utils.test_utils import ( + make_rotating_asset_info, + make_simple_asset_info, + product_upper_triangle, +) + + +class RollingSumDifference(TestingFactor): + window_length = 3 + inputs = [USEquityPricing.open, USEquityPricing.close] + + def from_windows(self, open, close): + return (open - close).sum(axis=0) + + +class ConstantInputTestCase(TestCase): + + def setUp(self): + self.constants = { + # Every day, assume every stock starts at 2, goes down to 1, + # goes up to 4, and finishes at 3. + USEquityPricing.low: 1, + USEquityPricing.open: 2, + USEquityPricing.close: 3, + USEquityPricing.high: 4, + } + self.assets = [1, 2, 3] + self.dates = date_range('2014-01-01', '2014-02-01', freq='D', tz='UTC') + self.loader = ConstantLoader( + constants=self.constants, + dates=self.dates, + assets=self.assets, + ) + + self.asset_info = make_simple_asset_info( + self.assets, + start_date=self.dates[0], + end_date=self.dates[-1], + ) + self.asset_finder = AssetFinder(self.asset_info) + + def test_single_factor(self): + loader = self.loader + engine = SimpleFFCEngine(loader, self.dates, self.asset_finder) + result_shape = (num_dates, num_assets) = (5, len(self.assets)) + dates = self.dates[10:10 + num_dates] + + factor = RollingSumDifference() + + result = engine.factor_matrix({'f': factor}, dates[0], dates[-1]) + self.assertEqual(set(result.columns), {'f'}) + + assert_array_equal( + result['f'].unstack().values, + full(result_shape, -factor.window_length), + ) + + def test_multiple_rolling_factors(self): + + loader = self.loader + engine = SimpleFFCEngine(loader, self.dates, self.asset_finder) + shape = num_dates, num_assets = (5, len(self.assets)) + dates = self.dates[10:10 + num_dates] + + short_factor = RollingSumDifference(window_length=3) + long_factor = RollingSumDifference(window_length=5) + high_factor = RollingSumDifference( + window_length=3, + inputs=[USEquityPricing.open, USEquityPricing.high], + ) + + results = engine.factor_matrix( + {'short': short_factor, 'long': long_factor, 'high': high_factor}, + dates[0], + dates[-1], + ) + self.assertEqual(set(results.columns), {'short', 'high', 'long'}) + + # row-wise sum over an array whose values are all (1 - 2) + assert_array_equal( + results['short'].unstack().values, + full(shape, -short_factor.window_length), + ) + assert_array_equal( + results['long'].unstack().values, + full(shape, -long_factor.window_length), + ) + # row-wise sum over an array whose values are all (1 - 3) + assert_array_equal( + results['high'].unstack().values, + full(shape, -2 * high_factor.window_length), + ) + + def test_numeric_factor(self): + constants = self.constants + loader = self.loader + engine = SimpleFFCEngine(loader, self.dates, self.asset_finder) + num_dates = 5 + dates = self.dates[10:10 + num_dates] + high, low = USEquityPricing.high, USEquityPricing.low + open, close = USEquityPricing.open, USEquityPricing.close + + high_minus_low = RollingSumDifference(inputs=[high, low]) + open_minus_close = RollingSumDifference(inputs=[open, close]) + avg = (high_minus_low + open_minus_close) / 2 + + results = engine.factor_matrix( + { + 'high_low': high_minus_low, + 'open_close': open_minus_close, + 'avg': avg, + }, + dates[0], + dates[-1], + ) + + high_low_result = results['high_low'].unstack() + expected_high_low = 3.0 * (constants[high] - constants[low]) + assert_frame_equal( + high_low_result, + 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, + ) + ) + + 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, + ) + ) + + +class FrameInputTestCase(TestCase): + + def setUp(self): + env = TradingEnvironment.instance() + day = env.trading_day + + self.assets = Int64Index([1, 2, 3]) + self.dates = date_range( + '2015-01-01', + '2015-01-31', + freq=day, + tz='UTC', + ) + + asset_info = make_simple_asset_info( + self.assets, + start_date=self.dates[0], + end_date=self.dates[-1], + ) + self.asset_finder = AssetFinder(asset_info) + + @lazyval + def base_mask(self): + return self.make_frame(True) + + def make_frame(self, data): + return DataFrame(data, columns=self.assets, index=self.dates) + + def test_compute_with_adjustments(self): + dates, assets = self.dates, self.assets + low, high = USEquityPricing.low, USEquityPricing.high + apply_idxs = [3, 10, 16] + + def apply_date(idx, offset=0): + return dates[apply_idxs[idx] + offset] + + adjustments = DataFrame.from_records( + [ + dict( + kind=MULTIPLY, + sid=assets[1], + value=2.0, + start_date=None, + end_date=apply_date(0, offset=-1), + apply_date=apply_date(0), + ), + dict( + kind=MULTIPLY, + sid=assets[1], + value=3.0, + start_date=None, + end_date=apply_date(1, offset=-1), + apply_date=apply_date(1), + ), + dict( + kind=MULTIPLY, + sid=assets[1], + value=5.0, + start_date=None, + end_date=apply_date(2, offset=-1), + apply_date=apply_date(2), + ), + ] + ) + low_base = DataFrame(self.make_frame(30.0)) + low_loader = DataFrameFFCLoader(low, low_base.copy(), adjustments=None) + + # Pre-apply inverse of adjustments to the baseline. + high_base = DataFrame(self.make_frame(30.0)) + high_base.iloc[:apply_idxs[0], 1] /= 2.0 + high_base.iloc[:apply_idxs[1], 1] /= 3.0 + high_base.iloc[:apply_idxs[2], 1] /= 5.0 + + high_loader = DataFrameFFCLoader(high, high_base, adjustments) + loader = MultiColumnLoader({low: low_loader, high: high_loader}) + + engine = SimpleFFCEngine(loader, self.dates, self.asset_finder) + + for window_length in range(1, 4): + low_mavg = SimpleMovingAverage( + inputs=[USEquityPricing.low], + window_length=window_length, + ) + high_mavg = SimpleMovingAverage( + inputs=[USEquityPricing.high], + window_length=window_length, + ) + bounds = product_upper_triangle(range(window_length, len(dates))) + for start, stop in bounds: + results = engine.factor_matrix( + {'low': low_mavg, 'high': high_mavg}, + dates[start], + dates[stop], + ) + self.assertEqual(set(results.columns), {'low', 'high'}) + iloc_bounds = slice(start, stop + 1) # +1 to include end date + + low_results = results.unstack()['low'] + assert_frame_equal(low_results, low_base.iloc[iloc_bounds]) + + high_results = results.unstack()['high'] + assert_frame_equal(high_results, high_base.iloc[iloc_bounds]) + + +class SyntheticBcolzTestCase(TestCase): + + @classmethod + def setUpClass(cls): + cls.first_asset_start = Timestamp('2015-04-01', tz='UTC') + cls.env = TradingEnvironment.instance() + cls.trading_day = cls.env.trading_day + cls.asset_info = make_rotating_asset_info( + num_assets=6, + first_start=cls.first_asset_start, + frequency=cls.trading_day, + periods_between_starts=4, + asset_lifetime=8, + ) + 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.finder = AssetFinder(cls.asset_info) + + 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, + ) + + cls.ffc_loader = USEquityPricingLoader( + BcolzDailyBarReader(table), + NullAdjustmentReader(), + ) + + @classmethod + def tearDownClass(cls): + cls.temp_dir.cleanup() + + 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 + SMA = SimpleMovingAverage( + inputs=(USEquityPricing.close,), + window_length=window_length, + ) + + results = engine.factor_matrix( + {'sma': SMA}, + dates[window_length], + dates[-1], + ) + raw_closes = self.writer.expected_values_2d(dates, assets, 'close') + expected_sma_result = rolling_mean( + raw_closes, + 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, + ), + ) + + def test_drawdown(self): + # The monotonically-increasing data produced by SyntheticDailyBarWriter + # exercises two pathological cases for MaxDrawdown. The actual + # computed results are pretty much useless (everything is either NaN) + # or zero, but verifying we correctly handle those corner cases is + # valuable. + engine = SimpleFFCEngine( + self.ffc_loader, + self.env.trading_days, + self.finder, + ) + dates, assets = self.all_dates, self.all_assets + window_length = 5 + drawdown = MaxDrawdown( + inputs=(USEquityPricing.close,), + window_length=window_length, + ) + + results = engine.factor_matrix( + {'drawdown': drawdown}, + dates[window_length], + dates[-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, + ), + ) diff --git a/tests/modelling/test_factor.py b/tests/modelling/test_factor.py new file mode 100644 index 00000000..5286d5b2 --- /dev/null +++ b/tests/modelling/test_factor.py @@ -0,0 +1,89 @@ +""" +Tests for Factor terms. +""" +from unittest import TestCase + +from numpy import ( + array, +) +from numpy.testing import assert_array_equal +from pandas import ( + DataFrame, + date_range, + Int64Index, +) +from six import iteritems + +from zipline.errors import UnknownRankMethod +from zipline.modelling.factor import TestingFactor + + +class F(TestingFactor): + inputs = () + window_length = 0 + + +class FactorTestCase(TestCase): + + def setUp(self): + self.f = F() + self.dates = date_range('2014-01-01', periods=5, freq='D') + self.assets = Int64Index(range(5)) + self.mask = DataFrame(True, index=self.dates, columns=self.assets) + + def tearDown(self): + pass + + def test_bad_input(self): + + with self.assertRaises(UnknownRankMethod): + self.f.rank("not a real rank method") + + def test_rank(self): + + # Generated with: + # 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]]) + expected_ranks = { + 'ordinal': 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.]]), + 'average': array([[1.5, 3., 4., 5., 1.5], + [2.5, 4., 5., 1., 2.5], + [3.5, 5., 1., 2., 3.5], + [4.5, 1., 2., 3., 4.5], + [1.5, 3., 4., 5., 1.5]]), + 'min': array([[1., 3., 4., 5., 1.], + [2., 4., 5., 1., 2.], + [3., 5., 1., 2., 3.], + [4., 1., 2., 3., 4.], + [1., 3., 4., 5., 1.]]), + 'max': array([[2., 3., 4., 5., 2.], + [3., 4., 5., 1., 3.], + [4., 5., 1., 2., 4.], + [5., 1., 2., 3., 5.], + [2., 3., 4., 5., 2.]]), + 'dense': array([[1., 2., 3., 4., 1.], + [2., 3., 4., 1., 2.], + [3., 4., 1., 2., 3.], + [4., 1., 2., 3., 4.], + [1., 2., 3., 4., 1.]]), + } + + # Test with the default, which should be 'ordinal'. + default_result = self.f.rank().compute_from_arrays([data], self.mask) + assert_array_equal(default_result, expected_ranks['ordinal']) + + # Test with each method passed explicitly. + for method, expected_result in iteritems(expected_ranks): + result = self.f.rank(method=method).compute_from_arrays( + [data], + self.mask, + ) + assert_array_equal(result, expected_ranks[method]) diff --git a/tests/modelling/test_filter.py b/tests/modelling/test_filter.py new file mode 100644 index 00000000..b415364b --- /dev/null +++ b/tests/modelling/test_filter.py @@ -0,0 +1,216 @@ +""" +Tests for filter terms. +""" +from unittest import TestCase + +from numpy import ( + arange, + array, + eye, + float64, + nan, + nanpercentile, + ones_like, + putmask, +) +from numpy.testing import assert_array_equal + +from pandas import ( + DataFrame, + date_range, + Int64Index, +) + +from zipline.errors import BadPercentileBounds +from zipline.modelling.factor import TestingFactor + + +class SomeFactor(TestingFactor): + inputs = () + window_length = 0 + + +class FilterTestCase(TestCase): + + def setUp(self): + self.f = SomeFactor() + self.dates = date_range('2014-01-01', periods=5, freq='D') + self.assets = Int64Index(range(5)) + self.mask = DataFrame(True, index=self.dates, columns=self.assets) + + def tearDown(self): + pass + + def maskframe(self, array): + return DataFrame( + array, + index=date_range('2014-01-01', periods=array.shape[0], freq='D'), + columns=arange(array.shape[1]), + ) + + def test_bad_input(self): + f = self.f + + bad_percentiles = [ + (-.1, 10), + (10, 100.1), + (20, 10), + (50, 50), + ] + for min_, max_ in bad_percentiles: + with self.assertRaises(BadPercentileBounds): + f.percentile_between(min_, max_) + + def test_rank_percentile_nice_partitions(self): + # Test case with nicely-defined partitions. + eye5 = eye(5, dtype=float64) + eye6 = eye(6, dtype=float64) + nanmask = array([[0, 0, 0, 0, 0, 1], + [1, 0, 0, 0, 0, 0], + [0, 1, 0, 0, 0, 0], + [0, 0, 1, 0, 0, 0], + [0, 0, 0, 1, 0, 0], + [0, 0, 0, 0, 1, 0]], dtype=bool) + nandata = eye6.copy() + putmask(nandata, nanmask, nan) + + for quintile in range(5): + factor = self.f.percentile_between( + quintile * 20.0, + (quintile + 1) * 20.0, + ) + # Test w/o any NaNs + result = factor.compute_from_arrays( + [eye5], + self.maskframe(ones_like(eye5, dtype=bool)), + ) + # Test with NaNs in the data. + nandata_result = factor.compute_from_arrays( + [nandata], + self.maskframe(ones_like(nandata, dtype=bool)), + ) + # Test with Falses in the mask. + nanmask_result = factor.compute_from_arrays( + [eye6], + self.maskframe(~nanmask), + ) + + assert_array_equal(nandata_result, nanmask_result) + + if quintile < 4: + # There are 4 0s and one 1 in each row, so the first 4 + # quintiles should be all the locations with zeros in the input + # array. + assert_array_equal(result, ~eye5.astype(bool)) + # Should reject all the ones, plus the nans. + assert_array_equal( + nandata_result, + ~(nanmask | eye6.astype(bool)) + ) + + else: + # The last quintile should contain all the 1s. + assert_array_equal(result, eye(5, dtype=bool)) + # Should accept all the 1s. + assert_array_equal(nandata_result, eye(6, dtype=bool)) + + def test_rank_percentile_nasty_partitions(self): + # Test case with nasty partitions: divide up 5 assets into quartiles. + data = arange(25, dtype=float).reshape(5, 5) % 4 + nandata = data.copy() + nandata[eye(5, dtype=bool)] = nan + for quartile in range(4): + lower_bound = quartile * 25.0 + upper_bound = (quartile + 1) * 25.0 + factor = self.f.percentile_between(lower_bound, upper_bound) + + # There isn't a nice definition of correct behavior here, so for + # now we guarantee the behavior of numpy.nanpercentile. + + result = factor.compute_from_arrays([data], self.mask) + min_value = nanpercentile(data, lower_bound, axis=1, keepdims=True) + max_value = nanpercentile(data, upper_bound, axis=1, keepdims=True) + assert_array_equal( + result, + (min_value <= data) & (data <= max_value), + ) + + nanresult = factor.compute_from_arrays([nandata], self.mask) + min_value = nanpercentile( + nandata, + lower_bound, + axis=1, + keepdims=True, + ) + max_value = nanpercentile( + nandata, + upper_bound, + axis=1, + keepdims=True, + ) + assert_array_equal( + nanresult, + (min_value <= nandata) & (nandata <= max_value), + ) + + def test_sequenced_filter(self): + first = SomeFactor() < 1 + first_input = eye(5) + first_result = first.compute_from_arrays([first_input], self.mask) + assert_array_equal(first_result, ~eye(5, dtype=bool)) + + # Second should pick out the fourth column. + second = SomeFactor().eq(3.0) + second_input = arange(25, dtype=float).reshape(5, 5) % 5 + + sequenced = first.then(second) + + result = sequenced.compute_from_arrays( + [first_result, second_input], + self.mask, + ) + expected_result = (first_result & (second_input == 3.0)) + assert_array_equal(result, expected_result) + + def test_sequenced_filter_order_dependent(self): + f = SomeFactor() < 1 + f_input = eye(5) + f_result = f.compute_from_arrays([f_input], self.mask) + assert_array_equal(f_result, ~eye(5, dtype=bool)) + + g = SomeFactor().percentile_between(80, 100) + g_input = arange(25, dtype=float).reshape(5, 5) % 5 + g_result = g.compute_from_arrays([g_input], self.mask) + assert_array_equal(g_result, g_input == 4) + + result = f.then(g).compute_from_arrays( + [f_result, g_input], + self.mask, + ) + # Input data is strictly increasing, so the result should be the top + # value not filtered by first. + expected_result = array( + [[0, 0, 0, 0, 1], + [0, 0, 0, 0, 1], + [0, 0, 0, 0, 1], + [0, 0, 0, 0, 1], + [0, 0, 0, 1, 0]], + dtype=bool, + ) + assert_array_equal(result, expected_result) + + result = g.then(f).compute_from_arrays( + [g_result, f_input], + self.mask, + ) + + # Percentile calculated first, then diagonal is removed. + expected_result = array( + [[0, 0, 0, 0, 1], + [0, 0, 0, 0, 1], + [0, 0, 0, 0, 1], + [0, 0, 0, 0, 1], + [0, 0, 0, 0, 0]], + dtype=bool, + ) + assert_array_equal(result, expected_result) diff --git a/tests/modelling/test_frameload.py b/tests/modelling/test_frameload.py new file mode 100644 index 00000000..c99101f3 --- /dev/null +++ b/tests/modelling/test_frameload.py @@ -0,0 +1,218 @@ +""" +Tests for zipline.data.ffc.frame.DataFrameFFCLoader +""" +from unittest import TestCase + +from mock import patch +from numpy import arange +from numpy.testing import assert_array_equal +from pandas import ( + DataFrame, + DatetimeIndex, + Int64Index, +) + +from zipline.lib.adjustment import ( + Float64Add, + Float64Multiply, + Float64Overwrite, +) +from zipline.data.equities import USEquityPricing +from zipline.data.ffc.frame import ( + ADD, + DataFrameFFCLoader, + MULTIPLY, + OVERWRITE, +) +from zipline.utils.tradingcalendar import trading_day + + +class DataFrameFFCLoaderTestCase(TestCase): + + def setUp(self): + self.nsids = 5 + self.ndates = 20 + + self.sids = Int64Index(range(self.nsids)) + self.dates = DatetimeIndex( + start='2014-01-02', + freq=trading_day, + periods=self.ndates, + ) + + self.mask = DataFrame( + True, + index=self.dates, + columns=self.sids, + dtype=bool, + ) + + def tearDown(self): + pass + + def test_bad_input(self): + data = arange(100).reshape(self.ndates, self.nsids) + baseline = DataFrame(data, index=self.dates, columns=self.sids) + loader = DataFrameFFCLoader( + USEquityPricing.close, + baseline, + ) + + with self.assertRaises(ValueError): + # Wrong column. + loader.load_adjusted_array([USEquityPricing.open], self.mask) + + with self.assertRaises(ValueError): + # Too many columns. + loader.load_adjusted_array( + [USEquityPricing.open, USEquityPricing.close], + self.mask + ) + + def test_baseline(self): + data = arange(100).reshape(self.ndates, self.nsids) + baseline = DataFrame(data, index=self.dates, columns=self.sids) + loader = DataFrameFFCLoader( + USEquityPricing.close, + baseline, + ) + + dates_slice = slice(None, 10, None) + sids_slice = slice(1, 3, None) + adj_array = loader.load_adjusted_array( + [USEquityPricing.close], + self.mask.iloc[dates_slice, sids_slice] + ) + + for idx, window in enumerate(adj_array.traverse(window_length=3)): + expected = baseline.values[dates_slice, sids_slice][idx:idx + 3] + assert_array_equal(window, expected) + + def test_adjustments(self): + data = arange(100).reshape(self.ndates, self.nsids) + baseline = DataFrame(data, index=self.dates, columns=self.sids) + + # Use the dates from index 10 on and sids 1-3. + dates_slice = slice(10, None, None) + sids_slice = slice(1, 4, None) + + # Adjustments that should actually affect the output. + relevant_adjustments = [ + { + 'sid': 1, + 'start_date': None, + 'end_date': self.dates[15], + 'apply_date': self.dates[16], + 'value': 0.5, + 'kind': MULTIPLY, + }, + { + 'sid': 2, + 'start_date': self.dates[5], + 'end_date': self.dates[15], + 'apply_date': self.dates[16], + 'value': 1.0, + 'kind': ADD, + }, + { + 'sid': 2, + 'start_date': self.dates[15], + 'end_date': self.dates[16], + 'apply_date': self.dates[17], + 'value': 1.0, + 'kind': ADD, + }, + { + 'sid': 3, + 'start_date': self.dates[16], + 'end_date': self.dates[17], + 'apply_date': self.dates[18], + 'value': 99.0, + 'kind': OVERWRITE, + }, + ] + + # These adjustments shouldn't affect the output. + irrelevant_adjustments = [ + { # Sid Not Requested + 'sid': 0, + 'start_date': self.dates[16], + 'end_date': self.dates[17], + 'apply_date': self.dates[18], + 'value': -9999.0, + 'kind': OVERWRITE, + }, + { # Sid Unknown + 'sid': 9999, + 'start_date': self.dates[16], + 'end_date': self.dates[17], + 'apply_date': self.dates[18], + 'value': -9999.0, + 'kind': OVERWRITE, + }, + { # Date Not Requested + 'sid': 2, + 'start_date': self.dates[1], + 'end_date': self.dates[2], + 'apply_date': self.dates[3], + 'value': -9999.0, + 'kind': OVERWRITE, + }, + { # Date Before Known Data + 'sid': 2, + 'start_date': self.dates[0] - (2 * trading_day), + 'end_date': self.dates[0] - trading_day, + 'apply_date': self.dates[0] - trading_day, + 'value': -9999.0, + 'kind': OVERWRITE, + }, + { # Date After Known Data + 'sid': 2, + 'start_date': self.dates[-1] + trading_day, + 'end_date': self.dates[-1] + (2 * trading_day), + 'apply_date': self.dates[-1] + (3 * trading_day), + 'value': -9999.0, + 'kind': OVERWRITE, + }, + ] + + adjustments = DataFrame(relevant_adjustments + irrelevant_adjustments) + loader = DataFrameFFCLoader( + USEquityPricing.close, + baseline, + adjustments=adjustments, + ) + + expected_baseline = baseline.iloc[dates_slice, sids_slice] + + formatted_adjustments = loader.format_adjustments( + self.dates[dates_slice], + self.sids[sids_slice], + ) + expected_formatted_adjustments = { + 6: [ + Float64Multiply(first_row=0, last_row=5, col=0, value=0.5), + Float64Add(first_row=0, last_row=5, col=1, value=1.0), + ], + 7: [ + Float64Add(first_row=5, last_row=6, col=1, value=1.0), + ], + 8: [ + Float64Overwrite(first_row=6, last_row=7, col=2, value=99.0) + ], + } + self.assertEqual(formatted_adjustments, expected_formatted_adjustments) + + mask = self.mask.iloc[dates_slice, sids_slice] + with patch('zipline.data.ffc.frame.adjusted_array') as m: + loader.load_adjusted_array( + columns=[USEquityPricing.close], + mask=mask, + ) + + self.assertEqual(m.call_count, 1) + + args, kwargs = m.call_args + assert_array_equal(kwargs['data'], expected_baseline.values) + assert_array_equal(kwargs['mask'], mask.values) + self.assertEqual(kwargs['adjustments'], expected_formatted_adjustments) diff --git a/tests/modelling/test_modelling_algo.py b/tests/modelling/test_modelling_algo.py new file mode 100644 index 00000000..4932d161 --- /dev/null +++ b/tests/modelling/test_modelling_algo.py @@ -0,0 +1,215 @@ +""" +Tests for Algorithms running the full FFC stack. +""" +from unittest import TestCase +from os.path import ( + dirname, + join, + realpath, +) + +from numpy import ( + array, + full_like, + nan, +) +from numpy.testing import assert_almost_equal +from pandas import ( + concat, + DataFrame, + DatetimeIndex, + Panel, + read_csv, + Series, + Timestamp, +) +from six import iteritems +from testfixtures import TempDirectory + +from zipline.algorithm import TradingAlgorithm +from zipline.api import ( + # add_filter, + add_factor, + get_datetime, +) +from zipline.assets import AssetFinder +# from zipline.data.equities import USEquityPricing +from zipline.data.ffc.loaders.us_equity_pricing import ( + BcolzDailyBarReader, + DailyBarWriterFromCSVs, + SQLiteAdjustmentReader, + SQLiteAdjustmentWriter, + USEquityPricingLoader, +) +# from zipline.modelling.factor import CustomFactor +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 + + +TEST_RESOURCE_PATH = join( + dirname(dirname(realpath(__file__))), # zipline_repo/tests + 'resources', + 'modelling_inputs', +) + + +def rolling_vwap(df, length): + "Simple rolling vwap implementation for testing" + closes = df['close'].values + volumes = df['volume'].values + product = closes * volumes + out = full_like(closes, nan) + for upper_bound in range(length, len(closes) + 1): + bounds = slice(upper_bound - length, upper_bound) + out[upper_bound - 1] = product[bounds].sum() / volumes[bounds].sum() + + return Series(out, index=df.index) + + +class FFCAlgorithmTestCase(TestCase): + + @classmethod + def setUpClass(cls): + cls.AAPL = 1 + cls.MSFT = 2 + cls.BRK_A = 3 + cls.assets = [cls.AAPL, cls.MSFT, cls.BRK_A] + asset_info = make_simple_asset_info( + cls.assets, + Timestamp('2014'), + Timestamp('2015'), + ['AAPL', 'MSFT', 'BRK_A'], + ) + cls.asset_finder = AssetFinder(asset_info) + cls.tempdir = tempdir = TempDirectory() + tempdir.create() + try: + cls.raw_data, cls.bar_reader = cls.create_bar_reader(tempdir) + cls.adj_reader = cls.create_adjustment_reader(tempdir) + cls.ffc_loader = USEquityPricingLoader( + cls.bar_reader, cls.adj_reader + ) + except: + cls.tempdir.cleanup() + raise + + cls.dates = cls.raw_data[cls.AAPL].index.tz_localize('UTC') + + @classmethod + def create_bar_reader(cls, tempdir): + resources = { + cls.AAPL: join(TEST_RESOURCE_PATH, 'AAPL.csv'), + cls.MSFT: join(TEST_RESOURCE_PATH, 'MSFT.csv'), + cls.BRK_A: join(TEST_RESOURCE_PATH, 'BRK-A.csv'), + } + raw_data = { + asset: read_csv(path, parse_dates=['day']).set_index('day') + for asset, path in iteritems(resources) + } + # Add 'price' column as an alias because all kinds of stuff in zipline + # depends on it being present. :/ + for frame in raw_data.values(): + frame['price'] = frame['close'] + + writer = DailyBarWriterFromCSVs(resources) + data_path = tempdir.getpath('testdata.bcolz') + table = writer.write(data_path, trading_days, cls.assets) + return raw_data, BcolzDailyBarReader(table) + + @classmethod + def create_adjustment_reader(cls, tempdir): + dbpath = tempdir.getpath('adjustments.sqlite') + writer = SQLiteAdjustmentWriter(dbpath) + splits = DataFrame.from_records([ + { + 'effective_date': str_to_seconds('2014-06-09'), + 'ratio': (1 / 7.0), + 'sid': cls.AAPL, + } + ]) + mergers = dividends = DataFrame( + { + # Hackery to make the dtypes correct on an empty frame. + 'effective_date': array([], dtype=int), + 'ratio': array([], dtype=float), + 'sid': array([], dtype=int), + }, + index=DatetimeIndex([], tz='UTC'), + columns=['effective_date', 'ratio', 'sid'], + ) + writer.write(splits, mergers, dividends) + return SQLiteAdjustmentReader(dbpath) + + @classmethod + def tearDownClass(cls): + cls.tempdir.cleanup() + + def make_source(self): + return Panel(self.raw_data).tz_localize('UTC', axis=1) + + 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 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 + + window_lengths = [1, 2, 5, 10] + # length -> asset -> expected vwap + vwaps = {length: {} for length in window_lengths} + vwap_keys = {} + 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) + vwaps[length][asset] = concat( + [ + raw[:split_loc], + adj[split_loc:] + ] + ) + + def initialize(context): + context.vwaps = [] + for length, key in iteritems(vwap_keys): + context.vwaps.append(VWAP(window_length=length)) + add_factor(context.vwaps[-1], name=key) + + def handle_data(context, data): + today = get_datetime() + factors = data.factors + for length, key in iteritems(vwap_keys): + 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) + + algo = TradingAlgorithm( + initialize=initialize, + handle_data=handle_data, + data_frequency='daily', + ffc_loader=self.ffc_loader, + asset_finder=self.asset_finder, + start=self.dates[max(window_lengths)], + end=self.dates[-1], + ) + + algo.run( + source=self.make_source(), + # Yes, I really do want to use the start and end dates I passed to + # TradingAlgorithm. + overwrite_sim_params=False, + ) diff --git a/tests/modelling/test_numerical_expression.py b/tests/modelling/test_numerical_expression.py new file mode 100644 index 00000000..1afd2d20 --- /dev/null +++ b/tests/modelling/test_numerical_expression.py @@ -0,0 +1,409 @@ +from operator import ( + and_, + ge, + gt, + le, + lt, + methodcaller, + ne, + or_, +) +from unittest import TestCase + +import numpy +from numpy import ( + arange, + eye, + full, + isnan, + zeros, +) +from numpy.testing import assert_array_equal +from pandas import ( + DataFrame, + date_range, + Int64Index, +) + +from zipline.modelling.expression import ( + NumericalExpression, + NUMEXPR_MATH_FUNCS, +) +from zipline.modelling.factor import TestingFactor + + +class F(TestingFactor): + inputs = () + window_length = 0 + + +class G(TestingFactor): + inputs = () + window_length = 0 + + +class H(TestingFactor): + inputs = () + window_length = 0 + + +class NumericalExpressionTestCase(TestCase): + + def setUp(self): + self.dates = date_range('2014-01-01', periods=5, freq='D') + self.assets = Int64Index(range(5)) + self.f = F() + self.g = G() + self.h = H() + self.fake_raw_data = { + self.f: full((5, 5), 3), + self.g: full((5, 5), 2), + self.h: full((5, 5), 1), + } + self.mask = DataFrame(True, index=self.dates, columns=self.assets) + + def check_output(self, expr, expected): + result = expr.compute_from_arrays( + [self.fake_raw_data[input_] for input_ in expr.inputs], + self.mask, + ) + assert_array_equal(result, full((5, 5), expected)) + + def check_constant_output(self, expr, expected): + self.assertFalse(isnan(expected)) + return self.check_output(expr, full((5, 5), expected)) + + def test_validate_good(self): + f = self.f + g = self.g + + NumericalExpression("x_0", (f,)) + NumericalExpression("x_0 ", (f,)) + NumericalExpression("x_0 + x_0", (f,)) + NumericalExpression("x_0 + 2", (f,)) + NumericalExpression("2 * x_0", (f,)) + NumericalExpression("x_0 + x_1", (f, g)) + NumericalExpression("x_0 + x_1 + x_0", (f, g)) + NumericalExpression("x_0 + 1 + x_1", (f, g)) + + def test_validate_bad(self): + f, g, h = F(), G(), H() + + # Too few inputs. + with self.assertRaises(ValueError): + NumericalExpression("x_0", ()) + with self.assertRaises(ValueError): + NumericalExpression("x_0 + x_1", (f,)) + + # Too many inputs. + with self.assertRaises(ValueError): + NumericalExpression("x_0", (f, g)) + with self.assertRaises(ValueError): + NumericalExpression("x_0 + x_1", (f, g, h)) + + # Invalid variable name. + with self.assertRaises(ValueError): + NumericalExpression("x_0x_1", (f,)) + with self.assertRaises(ValueError): + NumericalExpression("x_0x_1", (f, g)) + + # Variable index must start at 0. + with self.assertRaises(ValueError): + NumericalExpression("x_1", (f,)) + + # Scalar operands must be numeric. + with self.assertRaises(TypeError): + "2" + f + with self.assertRaises(TypeError): + f + "2" + with self.assertRaises(TypeError): + f > "2" + + # Boolean binary operators must be between filters. + with self.assertRaises(TypeError): + f + (f > 2) + with self.assertRaises(TypeError): + (f > f) > f + + def test_negate(self): + f, g = self.f, self.g + + self.check_constant_output(-f, -3.0) + self.check_constant_output(--f, 3.0) + self.check_constant_output(---f, -3.0) + + self.check_constant_output(-(f + f), -6.0) + self.check_constant_output(-f + -f, -6.0) + self.check_constant_output(-(-f + -f), 6.0) + + self.check_constant_output(f + -g, 1.0) + self.check_constant_output(f - -g, 5.0) + + self.check_constant_output(-(f + g) + (f + g), 0.0) + self.check_constant_output((f + g) + -(f + g), 0.0) + self.check_constant_output(-(f + g) + -(f + g), -10.0) + + def test_add(self): + f, g = self.f, self.g + + self.check_constant_output(f + g, 5.0) + + self.check_constant_output((1 + f) + g, 6.0) + self.check_constant_output(1 + (f + g), 6.0) + self.check_constant_output((f + 1) + g, 6.0) + self.check_constant_output(f + (1 + g), 6.0) + self.check_constant_output((f + g) + 1, 6.0) + self.check_constant_output(f + (g + 1), 6.0) + + self.check_constant_output((f + f) + f, 9.0) + self.check_constant_output(f + (f + f), 9.0) + + self.check_constant_output((f + g) + f, 8.0) + self.check_constant_output(f + (g + f), 8.0) + + self.check_constant_output((f + g) + (f + g), 10.0) + self.check_constant_output((f + g) + (g + f), 10.0) + self.check_constant_output((g + f) + (f + g), 10.0) + self.check_constant_output((g + f) + (g + f), 10.0) + + def test_subtract(self): + f, g = self.f, self.g + + self.check_constant_output(f - g, 1.0) # 3 - 2 + + self.check_constant_output((1 - f) - g, -4.) # (1 - 3) - 2 + self.check_constant_output(1 - (f - g), 0.0) # 1 - (3 - 2) + self.check_constant_output((f - 1) - g, 0.0) # (3 - 1) - 2 + self.check_constant_output(f - (1 - g), 4.0) # 3 - (1 - 2) + self.check_constant_output((f - g) - 1, 0.0) # (3 - 2) - 1 + self.check_constant_output(f - (g - 1), 2.0) # 3 - (2 - 1) + + self.check_constant_output((f - f) - f, -3.) # (3 - 3) - 3 + self.check_constant_output(f - (f - f), 3.0) # 3 - (3 - 3) + + self.check_constant_output((f - g) - f, -2.) # (3 - 2) - 3 + self.check_constant_output(f - (g - f), 4.0) # 3 - (2 - 3) + + self.check_constant_output((f - g) - (f - g), 0.0) # (3 - 2) - (3 - 2) + self.check_constant_output((f - g) - (g - f), 2.0) # (3 - 2) - (2 - 3) + self.check_constant_output((g - f) - (f - g), -2.) # (2 - 3) - (3 - 2) + self.check_constant_output((g - f) - (g - f), 0.0) # (2 - 3) - (2 - 3) + + def test_multiply(self): + f, g = self.f, self.g + + self.check_constant_output(f * g, 6.0) + + self.check_constant_output((2 * f) * g, 12.0) + self.check_constant_output(2 * (f * g), 12.0) + self.check_constant_output((f * 2) * g, 12.0) + self.check_constant_output(f * (2 * g), 12.0) + self.check_constant_output((f * g) * 2, 12.0) + self.check_constant_output(f * (g * 2), 12.0) + + self.check_constant_output((f * f) * f, 27.0) + self.check_constant_output(f * (f * f), 27.0) + + self.check_constant_output((f * g) * f, 18.0) + self.check_constant_output(f * (g * f), 18.0) + + self.check_constant_output((f * g) * (f * g), 36.0) + self.check_constant_output((f * g) * (g * f), 36.0) + self.check_constant_output((g * f) * (f * g), 36.0) + self.check_constant_output((g * f) * (g * f), 36.0) + + self.check_constant_output(f * f * f * 0 * f * f, 0.0) + + def test_divide(self): + f, g = self.f, self.g + + self.check_constant_output(f / g, 3.0 / 2.0) + + self.check_constant_output( + (2 / f) / g, + (2 / 3.0) / 2.0 + ) + self.check_constant_output( + 2 / (f / g), + 2 / (3.0 / 2.0), + ) + self.check_constant_output( + (f / 2) / g, + (3.0 / 2) / 2.0, + ) + self.check_constant_output( + f / (2 / g), + 3.0 / (2 / 2.0), + ) + self.check_constant_output( + (f / g) / 2, + (3.0 / 2.0) / 2, + ) + self.check_constant_output( + f / (g / 2), + 3.0 / (2.0 / 2), + ) + self.check_constant_output( + (f / f) / f, + (3.0 / 3.0) / 3.0 + ) + self.check_constant_output( + f / (f / f), + 3.0 / (3.0 / 3.0), + ) + self.check_constant_output( + (f / g) / f, + (3.0 / 2.0) / 3.0, + ) + self.check_constant_output( + f / (g / f), + 3.0 / (2.0 / 3.0), + ) + + self.check_constant_output( + (f / g) / (f / g), + (3.0 / 2.0) / (3.0 / 2.0), + ) + self.check_constant_output( + (f / g) / (g / f), + (3.0 / 2.0) / (2.0 / 3.0), + ) + self.check_constant_output( + (g / f) / (f / g), + (2.0 / 3.0) / (3.0 / 2.0), + ) + self.check_constant_output( + (g / f) / (g / f), + (2.0 / 3.0) / (2.0 / 3.0), + ) + + def test_pow(self): + f, g = self.f, self.g + + self.check_constant_output(f ** g, 3.0 ** 2) + self.check_constant_output(2 ** f, 2.0 ** 3) + self.check_constant_output(f ** 2, 3.0 ** 2) + + self.check_constant_output((f + g) ** 2, (3.0 + 2.0) ** 2) + self.check_constant_output(2 ** (f + g), 2 ** (3.0 + 2.0)) + + self.check_constant_output(f ** (f ** g), 3.0 ** (3.0 ** 2.0)) + self.check_constant_output((f ** f) ** g, (3.0 ** 3.0) ** 2.0) + + self.check_constant_output((f ** g) ** (f ** g), 9.0 ** 9.0) + self.check_constant_output((f ** g) ** (g ** f), 9.0 ** 8.0) + self.check_constant_output((g ** f) ** (f ** g), 8.0 ** 9.0) + self.check_constant_output((g ** f) ** (g ** f), 8.0 ** 8.0) + + def test_mod(self): + f, g = self.f, self.g + + self.check_constant_output(f % g, 3.0 % 2.0) + self.check_constant_output(f % 2.0, 3.0 % 2.0) + self.check_constant_output(g % f, 2.0 % 3.0) + + self.check_constant_output((f + g) % 2, (3.0 + 2.0) % 2) + self.check_constant_output(2 % (f + g), 2 % (3.0 + 2.0)) + + self.check_constant_output(f % (f % g), 3.0 % (3.0 % 2.0)) + self.check_constant_output((f % f) % g, (3.0 % 3.0) % 2.0) + + self.check_constant_output((f + g) % (f * g), 5.0 % 6.0) + + def test_math_functions(self): + f, g = self.f, self.g + + fake_raw_data = self.fake_raw_data + alt_fake_raw_data = { + self.f: full((5, 5), .5), + self.g: full((5, 5), -.5), + } + + for funcname in NUMEXPR_MATH_FUNCS: + method = methodcaller(funcname) + func = getattr(numpy, funcname) + + # These methods have domains in [0, 1], so we need alternate inputs + # that are in the domain. + if funcname in ('arcsin', 'arccos', 'arctanh'): + self.fake_raw_data = alt_fake_raw_data + else: + self.fake_raw_data = fake_raw_data + + f_val = self.fake_raw_data[f][0, 0] + g_val = self.fake_raw_data[g][0, 0] + + self.check_constant_output(method(f), func(f_val)) + self.check_constant_output(method(g), func(g_val)) + + self.check_constant_output(method(f) + 1, func(f_val) + 1) + self.check_constant_output(1 + method(f), 1 + func(f_val)) + + self.check_constant_output(method(f + .25), func(f_val + .25)) + self.check_constant_output(method(.25 + f), func(.25 + f_val)) + + self.check_constant_output( + method(f) + method(g), + func(f_val) + func(g_val), + ) + self.check_constant_output( + method(f + g), + func(f_val + g_val), + ) + + def test_comparisons(self): + f, g, h = self.f, self.g, self.h + self.fake_raw_data = { + f: arange(25).reshape(5, 5), + g: arange(25).reshape(5, 5) - eye(5), + h: full((5, 5), 5), + } + f_data = self.fake_raw_data[f] + g_data = self.fake_raw_data[g] + + cases = [ + # Sanity Check with hand-computed values. + (f, g, eye(5), zeros((5, 5))), + (f, 10, f_data, 10), + (10, f, 10, f_data), + (f, f, f_data, f_data), + (f + 1, f, f_data + 1, f_data), + (1 + f, f, 1 + f_data, f_data), + (f, g, f_data, g_data), + (f + 1, g, f_data + 1, g_data), + (f, g + 1, f_data, g_data + 1), + (f + 1, g + 1, f_data + 1, g_data + 1), + ((f + g) / 2, f ** 2, (f_data + g_data) / 2, f_data ** 2), + ] + for op in (gt, ge, lt, le, ne): + for expr_lhs, expr_rhs, expected_lhs, expected_rhs in cases: + self.check_output( + op(expr_lhs, expr_rhs), + op(expected_lhs, expected_rhs), + ) + + def test_boolean_binops(self): + f, g, h = self.f, self.g, self.h + self.fake_raw_data = { + f: arange(25).reshape(5, 5), + g: arange(25).reshape(5, 5) - eye(5), + h: full((5, 5), 5), + } + + # Should be True on the diagonal. + eye_filter = f > g + # Should be True in the first row only. + first_row_filter = f < h + + eye_mask = eye(5, dtype=bool) + first_row_mask = zeros((5, 5), dtype=bool) + first_row_mask[0] = 1 + + self.check_output(eye_filter, eye_mask) + self.check_output(first_row_filter, first_row_mask) + + for op in (and_, or_): # NumExpr doesn't support xor. + self.check_output( + op(eye_filter, first_row_filter), + op(eye_mask, first_row_mask), + ) diff --git a/tests/modelling/test_term.py b/tests/modelling/test_term.py new file mode 100644 index 00000000..41a2105e --- /dev/null +++ b/tests/modelling/test_term.py @@ -0,0 +1,246 @@ +""" +Tests for Term. +""" +from itertools import product +from unittest import TestCase + +from networkx import topological_sort +from numpy import ( + float32, + uint32, + uint8, +) + +from zipline.data.dataset import ( + Column, + DataSet, +) +from zipline.errors import InputTermNotAtomic +from zipline.modelling.engine import build_dependency_graph +from zipline.modelling.factor import Factor +from zipline.modelling.expression import NUMEXPR_MATH_FUNCS + + +class SomeDataSet(DataSet): + + foo = Column(float32) + bar = Column(uint32) + buzz = Column(uint8) + + +class SomeFactor(Factor): + window_length = 5 + inputs = [SomeDataSet.foo, SomeDataSet.bar] + + +class NoLookbackFactor(Factor): + window_length = 0 + + +class SomeOtherFactor(Factor): + window_length = 5 + inputs = [SomeDataSet.bar, SomeDataSet.buzz] + + +SomeFactorAlias = SomeFactor + + +def gen_equivalent_factors(): + """ + Return an iterator of SomeFactor instances that should all be the same + object. + """ + yield SomeFactor() + yield SomeFactor(inputs=None) + 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( + [SomeDataSet.foo, SomeDataSet.bar], + window_length=SomeFactor.window_length, + ) + yield SomeFactorAlias() + + +class DependencyResolutionTestCase(TestCase): + + def setup(self): + pass + + def teardown(self): + pass + + def test_single_factor(self): + """ + Test dependency resolution for a single factor. + """ + + build_dependency_graph([SomeFactor()]) + + def check_output(graph): + + resolution_order = topological_sort(graph) + + self.assertEqual(len(resolution_order), 3) + self.assertEqual( + set([resolution_order[0], resolution_order[1]]), + set([SomeDataSet.foo, SomeDataSet.bar]), + ) + self.assertEqual(resolution_order[-1], SomeFactor()) + self.assertEqual(graph.node[SomeDataSet.foo]['extra_rows'], 4) + self.assertEqual(graph.node[SomeDataSet.bar]['extra_rows'], 4) + + for foobar in gen_equivalent_factors(): + check_output(build_dependency_graph([foobar])) + + def test_single_factor_instance_args(self): + """ + Test dependency resolution for a single factor with arguments passed to + the constructor. + """ + graph = build_dependency_graph( + [SomeFactor([SomeDataSet.bar, SomeDataSet.buzz], window_length=5)] + ) + resolution_order = topological_sort(graph) + + self.assertEqual(len(resolution_order), 3) + self.assertEqual( + set([resolution_order[0], resolution_order[1]]), + set([SomeDataSet.bar, SomeDataSet.buzz]), + ) + self.assertEqual( + resolution_order[-1], + SomeFactor([SomeDataSet.bar, SomeDataSet.buzz], window_length=5), + ) + self.assertEqual(graph.node[SomeDataSet.bar]['extra_rows'], 4) + self.assertEqual(graph.node[SomeDataSet.buzz]['extra_rows'], 4) + + def test_reuse_atomic_terms(self): + """ + Test that raw inputs only show up in the dependency graph once. + """ + f1 = SomeFactor([SomeDataSet.foo, SomeDataSet.bar]) + f2 = SomeOtherFactor([SomeDataSet.bar, SomeDataSet.buzz]) + + graph = build_dependency_graph([f1, f2]) + resolution_order = topological_sort(graph) + + # bar should only appear once. + self.assertEqual(len(resolution_order), 5) + indices = { + term: resolution_order.index(term) + for term in resolution_order + } + + # Verify that f1's dependencies will be computed before f1. + self.assertLess(indices[SomeDataSet.foo], indices[f1]) + self.assertLess(indices[SomeDataSet.bar], indices[f1]) + + # Verify that f2's dependencies will be computed before f2. + self.assertLess(indices[SomeDataSet.bar], indices[f2]) + self.assertLess(indices[SomeDataSet.buzz], indices[f2]) + + def test_disallow_recursive_lookback(self): + + with self.assertRaises(InputTermNotAtomic): + SomeFactor(inputs=[SomeFactor(), SomeDataSet.foo]) + + +class ObjectIdentityTestCase(TestCase): + + def assertSameObject(self, *objs): + first = objs[0] + for obj in objs: + self.assertIs(first, obj) + + def test_instance_caching(self): + + self.assertSameObject(*gen_equivalent_factors()) + self.assertIs( + SomeFactor(window_length=SomeFactor.window_length + 1), + SomeFactor(window_length=SomeFactor.window_length + 1), + ) + + self.assertIs( + SomeFactor(dtype=int), + SomeFactor(dtype=int), + ) + + self.assertIs( + SomeFactor(inputs=[SomeFactor.inputs[1], SomeFactor.inputs[0]]), + SomeFactor(inputs=[SomeFactor.inputs[1], SomeFactor.inputs[0]]), + ) + + def test_instance_non_caching(self): + + f = SomeFactor() + + # Different window_length. + self.assertIsNot( + f, + SomeFactor(window_length=SomeFactor.window_length + 1), + ) + + # Different dtype + self.assertIsNot( + f, + SomeFactor(dtype=int) + ) + + # Reordering inputs changes semantics. + self.assertIsNot( + f, + SomeFactor(inputs=[SomeFactor.inputs[1], SomeFactor.inputs[0]]), + ) + + def test_instance_non_caching_redefine_class(self): + + orig_foobar_instance = SomeFactorAlias() + + class SomeFactor(Factor): + window_length = 5 + inputs = [SomeDataSet.foo, SomeDataSet.bar] + + self.assertIsNot(orig_foobar_instance, SomeFactor()) + + def test_instance_caching_binops(self): + f = SomeFactor() + g = SomeOtherFactor() + for lhs, rhs in product([f, g], [f, g]): + self.assertIs((lhs + rhs), (lhs + rhs)) + self.assertIs((lhs - rhs), (lhs - rhs)) + self.assertIs((lhs * rhs), (lhs * rhs)) + self.assertIs((lhs / rhs), (lhs / rhs)) + self.assertIs((lhs ** rhs), (lhs ** rhs)) + + self.assertIs((1 + rhs), (1 + rhs)) + self.assertIs((rhs + 1), (rhs + 1)) + + self.assertIs((1 - rhs), (1 - rhs)) + self.assertIs((rhs - 1), (rhs - 1)) + + self.assertIs((2 * rhs), (2 * rhs)) + self.assertIs((rhs * 2), (rhs * 2)) + + self.assertIs((2 / rhs), (2 / rhs)) + self.assertIs((rhs / 2), (rhs / 2)) + + self.assertIs((2 ** rhs), (2 ** rhs)) + self.assertIs((rhs ** 2), (rhs ** 2)) + + self.assertIs((f + g) + (f + g), (f + g) + (f + g)) + + def test_instance_caching_unary_ops(self): + f = SomeFactor() + self.assertIs(-f, -f) + self.assertIs(--f, --f) + self.assertIs(---f, ---f) + + def test_instance_caching_math_funcs(self): + f = SomeFactor() + for funcname in NUMEXPR_MATH_FUNCS: + method = getattr(f, funcname) + self.assertIs(method(), method()) diff --git a/tests/modelling/test_us_equity_pricing_loader.py b/tests/modelling/test_us_equity_pricing_loader.py new file mode 100644 index 00000000..0aa4f929 --- /dev/null +++ b/tests/modelling/test_us_equity_pricing_loader.py @@ -0,0 +1,643 @@ +# +# 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. +""" +Tests for zipline.data.ffc.loaders.us_equity_pricing +""" +from unittest import TestCase + +from nose_parameterized import parameterized +from numpy import ( + arange, + datetime64, + uint32, +) +from numpy.testing import ( + assert_allclose, + assert_array_equal, +) +from pandas import ( + concat, + DataFrame, + DatetimeIndex, + Timestamp, +) +from pandas.util.testing import assert_index_equal +from testfixtures import TempDirectory + +from zipline.lib.adjustment import Float64Multiply +from zipline.data.equities import USEquityPricing +from zipline.data.ffc.synthetic import ( + NullAdjustmentReader, + SyntheticDailyBarWriter, +) +from zipline.data.ffc.loaders.us_equity_pricing import ( + BcolzDailyBarReader, + SQLiteAdjustmentReader, + SQLiteAdjustmentWriter, + USEquityPricingLoader, +) +from zipline.errors import WindowLengthTooLong +from zipline.finance.trading import TradingEnvironment +from zipline.utils.test_utils import ( + seconds_to_timestamp, + str_to_seconds, +) + +# Test calendar ranges over the month of June 2015 +# June 2015 +# Mo Tu We Th Fr Sa Su +# 1 2 3 4 5 6 7 +# 8 9 10 11 12 13 14 +# 15 16 17 18 19 20 21 +# 22 23 24 25 26 27 28 +# 29 30 +TEST_CALENDAR_START = Timestamp('2015-06-01', tz='UTC') +TEST_CALENDAR_STOP = Timestamp('2015-06-30', tz='UTC') + +TEST_QUERY_START = Timestamp('2015-06-10', tz='UTC') +TEST_QUERY_STOP = Timestamp('2015-06-19', tz='UTC') + +# One asset for each of the cases enumerated in load_raw_arrays_from_bcolz. +EQUITY_INFO = DataFrame( + [ + # 1) The equity's trades start and end before query. + {'start_date': '2015-06-01', 'end_date': '2015-06-05'}, + # 2) The equity's trades start and end after query. + {'start_date': '2015-06-22', 'end_date': '2015-06-30'}, + # 3) The equity's data covers all dates in range. + {'start_date': '2015-06-02', 'end_date': '2015-06-30'}, + # 4) The equity's trades start before the query start, but stop + # before the query end. + {'start_date': '2015-06-01', 'end_date': '2015-06-15'}, + # 5) The equity's trades start and end during the query. + {'start_date': '2015-06-12', 'end_date': '2015-06-18'}, + # 6) The equity's trades start during the query, but extend through + # the whole query. + {'start_date': '2015-06-15', 'end_date': '2015-06-25'}, + ], + index=arange(1, 7), + columns=['start_date', 'end_date'], +).astype(datetime64) + +TEST_QUERY_ASSETS = EQUITY_INFO.index + + +class BcolzDailyBarTestCase(TestCase): + + def setUp(self): + all_trading_days = TradingEnvironment.instance().trading_days + self.trading_days = all_trading_days[ + all_trading_days.get_loc(TEST_CALENDAR_START): + all_trading_days.get_loc(TEST_CALENDAR_STOP) + 1 + ] + + self.asset_info = EQUITY_INFO + self.writer = SyntheticDailyBarWriter( + self.asset_info, + self.trading_days, + ) + + self.dir_ = TempDirectory() + self.dir_.create() + self.dest = self.dir_.getpath('daily_equity_pricing.bcolz') + + def tearDown(self): + self.dir_.cleanup() + + @property + def assets(self): + return self.asset_info.index + + def trading_days_between(self, start, end): + return self.trading_days[self.trading_days.slice_indexer(start, end)] + + def asset_start(self, asset_id): + return self.writer.asset_start(asset_id) + + def asset_end(self, asset_id): + return self.writer.asset_end(asset_id) + + def dates_for_asset(self, asset_id): + start, end = self.asset_start(asset_id), self.asset_end(asset_id) + return self.trading_days_between(start, end) + + def test_write_ohlcv_content(self): + result = self.writer.write(self.dest, self.trading_days, self.assets) + for column in SyntheticDailyBarWriter.OHLCV: + idx = 0 + data = result[column][:] + multiplier = 1 if column == 'volume' else 1000 + for asset_id in self.assets: + for date in self.dates_for_asset(asset_id): + self.assertEqual( + SyntheticDailyBarWriter.expected_value( + asset_id, + date, + column + ) * multiplier, + data[idx], + ) + idx += 1 + self.assertEqual(idx, len(data)) + + def test_write_day_and_id(self): + result = self.writer.write(self.dest, self.trading_days, self.assets) + idx = 0 + ids = result['id'] + days = result['day'] + for asset_id in self.assets: + for date in self.dates_for_asset(asset_id): + self.assertEqual(ids[idx], asset_id) + self.assertEqual(date, seconds_to_timestamp(days[idx])) + idx += 1 + + def test_write_attrs(self): + result = self.writer.write(self.dest, self.trading_days, self.assets) + expected_first_row = { + '1': 0, + '2': 5, # Asset 1 has 5 trading days. + '3': 12, # Asset 2 has 7 trading days. + '4': 33, # Asset 3 has 21 trading days. + '5': 44, # Asset 4 has 11 trading days. + '6': 49, # Asset 5 has 5 trading days. + } + expected_last_row = { + '1': 4, + '2': 11, + '3': 32, + '4': 43, + '5': 48, + '6': 57, # Asset 6 has 9 trading days. + } + expected_calendar_offset = { + '1': 0, # Starts on 6-01, 1st trading day of month. + '2': 15, # Starts on 6-22, 16th trading day of month. + '3': 1, # Starts on 6-02, 2nd trading day of month. + '4': 0, # Starts on 6-01, 1st trading day of month. + '5': 9, # Starts on 6-12, 10th trading day of month. + '6': 10, # Starts on 6-15, 11th trading day of month. + } + self.assertEqual(result.attrs['first_row'], expected_first_row) + self.assertEqual(result.attrs['last_row'], expected_last_row) + self.assertEqual( + result.attrs['calendar_offset'], + expected_calendar_offset, + ) + assert_index_equal( + self.trading_days, + DatetimeIndex(result.attrs['calendar'], tz='UTC'), + ) + + 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) + 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, + self.writer.expected_values_2d( + dates, + assets, + column.name, + ) + ) + + @parameterized.expand([ + ([USEquityPricing.open],), + ([USEquityPricing.close, USEquityPricing.volume],), + ([USEquityPricing.volume, USEquityPricing.high, USEquityPricing.low],), + (USEquityPricing.columns,), + ]) + def test_read(self, columns): + self._check_read_results( + columns, + self.assets, + TEST_QUERY_START, + TEST_QUERY_STOP, + ) + + def test_start_on_asset_start(self): + """ + Test loading with queries that starts on the first day of each asset's + lifetime. + """ + columns = [USEquityPricing.high, USEquityPricing.volume] + for asset in self.assets: + self._check_read_results( + columns, + self.assets, + start_date=self.asset_start(asset), + end_date=self.trading_days[-1], + ) + + def test_start_on_asset_end(self): + """ + Test loading with queries that start on the last day of each asset's + lifetime. + """ + columns = [USEquityPricing.close, USEquityPricing.volume] + for asset in self.assets: + self._check_read_results( + columns, + self.assets, + start_date=self.asset_end(asset), + end_date=self.trading_days[-1], + ) + + def test_end_on_asset_start(self): + """ + Test loading with queries that end on the first day of each asset's + lifetime. + """ + columns = [USEquityPricing.close, USEquityPricing.volume] + for asset in self.assets: + self._check_read_results( + columns, + self.assets, + start_date=self.trading_days[0], + end_date=self.asset_start(asset), + ) + + def test_end_on_asset_end(self): + """ + Test loading with queries that end on the last day of each asset's + lifetime. + """ + columns = [USEquityPricing.close, USEquityPricing.volume] + for asset in self.assets: + self._check_read_results( + columns, + self.assets, + start_date=self.trading_days[0], + end_date=self.asset_end(asset), + ) + + +# ADJUSTMENTS use the following scheme to indicate information about the value +# upon inspection. +# +# 1s place is the equity +# +# 0.1s place is the action type, with: +# +# splits, 1 +# mergers, 2 +# dividends, 3 +# +# 0.001s is the date +SPLITS = DataFrame( + [ + # Before query range, should be excluded. + {'effective_date': str_to_seconds('2015-06-03'), + 'ratio': 1.103, + 'sid': 1}, + # First day of query range, should be excluded. + {'effective_date': str_to_seconds('2015-06-10'), + 'ratio': 3.110, + 'sid': 3}, + # Third day of query range, should have last_row of 2 + {'effective_date': str_to_seconds('2015-06-12'), + 'ratio': 3.112, + 'sid': 3}, + # After query range, should be excluded. + {'effective_date': str_to_seconds('2015-06-21'), + 'ratio': 6.121, + 'sid': 6}, + # Another action in query range, should have last_row of 1 + {'effective_date': str_to_seconds('2015-06-11'), + 'ratio': 3.111, + 'sid': 3}, + # Last day of range. Should have last_row of 7 + {'effective_date': str_to_seconds('2015-06-19'), + 'ratio': 3.119, + 'sid': 3}, + ], + columns=['effective_date', 'ratio', 'sid'], +) + + +MERGERS = DataFrame( + [ + # Before query range, should be excluded. + {'effective_date': str_to_seconds('2015-06-03'), + 'ratio': 1.203, + 'sid': 1}, + # First day of query range, should be excluded. + {'effective_date': str_to_seconds('2015-06-10'), + 'ratio': 3.210, + 'sid': 3}, + # Third day of query range, should have last_row of 2 + {'effective_date': str_to_seconds('2015-06-12'), + 'ratio': 3.212, + 'sid': 3}, + # After query range, should be excluded. + {'effective_date': str_to_seconds('2015-06-25'), + 'ratio': 6.225, + 'sid': 6}, + # Another action in query range, should have last_row of 2 + {'effective_date': str_to_seconds('2015-06-12'), + 'ratio': 4.212, + 'sid': 4}, + # Last day of range. Should have last_row of 7 + {'effective_date': str_to_seconds('2015-06-19'), + 'ratio': 3.219, + 'sid': 3}, + ], + columns=['effective_date', 'ratio', 'sid'], +) + + +DIVIDENDS = DataFrame( + [ + # Before query range, should be excluded. + {'effective_date': str_to_seconds('2015-06-01'), + 'ratio': 1.301, + 'sid': 1}, + # First day of query range, should be excluded. + {'effective_date': str_to_seconds('2015-06-10'), + 'ratio': 3.310, + 'sid': 3}, + # Third day of query range, should have last_row of 2 + {'effective_date': str_to_seconds('2015-06-12'), + 'ratio': 3.312, + 'sid': 3}, + # After query range, should be excluded. + {'effective_date': str_to_seconds('2015-06-25'), + 'ratio': 6.325, + 'sid': 6}, + # Another action in query range, should have last_row of 3 + {'effective_date': str_to_seconds('2015-06-15'), + 'ratio': 3.315, + 'sid': 3}, + # Last day of range. Should have last_row of 7 + {'effective_date': str_to_seconds('2015-06-19'), + 'ratio': 3.319, + 'sid': 3}, + ], + columns=['effective_date', 'ratio', 'sid'], +) + + +class USEquityPricingLoaderTestCase(TestCase): + + @classmethod + def setUpClass(cls): + cls.test_data_dir = TempDirectory() + cls.db_path = cls.test_data_dir.getpath('adjustments.db') + writer = SQLiteAdjustmentWriter(cls.db_path) + writer.write(SPLITS, MERGERS, DIVIDENDS) + + cls.assets = TEST_QUERY_ASSETS + all_days = TradingEnvironment.instance().trading_days + cls.calendar_days = all_days[ + all_days.slice_indexer(TEST_CALENDAR_START, TEST_CALENDAR_STOP) + ] + + cls.asset_info = EQUITY_INFO + cls.bcolz_writer = SyntheticDailyBarWriter( + cls.asset_info, + cls.calendar_days, + ) + cls.bcolz_path = cls.test_data_dir.getpath('equity_pricing.bcolz') + cls.bcolz_writer.write(cls.bcolz_path, cls.calendar_days, cls.assets) + + @classmethod + def tearDownClass(cls): + cls.test_data_dir.cleanup() + + def test_input_sanity(self): + # Ensure that the input data doesn't contain adjustments during periods + # where the corresponding asset didn't exist. + for table in SPLITS, MERGERS, DIVIDENDS: + for eff_date_secs, _, sid in table.itertuples(index=False): + eff_date = Timestamp(eff_date_secs, unit='s') + asset_start, asset_end = EQUITY_INFO.ix[ + sid, ['start_date', 'end_date'] + ] + 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 expected_adjustments(self, start_date, end_date): + price_adjustments = {} + volume_adjustments = {} + query_days = self.calendar_days_between(start_date, end_date) + start_loc = query_days.get_loc(start_date) + + for table in SPLITS, MERGERS, DIVIDENDS: + 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): + continue + + eff_date_loc = query_days.get_loc(eff_date) + delta = eff_date_loc - start_loc + + # Pricing adjusments 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, + col=sid - 1, + value=ratio, + ) + ) + # Volume is *inversely* affected by *splits only*. + if table is SPLITS: + volume_adjustments.setdefault(delta, []).append( + Float64Multiply( + first_row=0, + last_row=delta - 1, + col=sid - 1, + value=1.0 / ratio, + ) + ) + return price_adjustments, volume_adjustments + + def test_load_adjustments_from_sqlite(self): + reader = SQLiteAdjustmentReader(self.db_path) + columns = [USEquityPricing.close, USEquityPricing.volume] + query_days = self.calendar_days_between( + TEST_QUERY_START, + TEST_QUERY_STOP + ) + + adjustments = reader.load_adjustments( + columns, + query_days, + self.assets, + ) + + close_adjustments = adjustments[0] + volume_adjustments = adjustments[1] + + expected_close_adjustments, expected_volume_adjustments = \ + self.expected_adjustments(TEST_QUERY_START, TEST_QUERY_STOP) + self.assertEqual(close_adjustments, expected_close_adjustments) + self.assertEqual(volume_adjustments, expected_volume_adjustments) + + def test_read_no_adjustments(self): + adjustment_reader = NullAdjustmentReader() + columns = [USEquityPricing.close, USEquityPricing.volume] + query_days = self.calendar_days_between( + TEST_QUERY_START, + TEST_QUERY_STOP + ) + + adjustments = adjustment_reader.load_adjustments( + columns, + query_days, + self.assets, + ) + self.assertEqual(adjustments, [{}, {}]) + + baseline_reader = BcolzDailyBarReader(self.bcolz_path) + pricing_loader = USEquityPricingLoader( + baseline_reader, + adjustment_reader, + ) + + closes, volumes = pricing_loader.load_adjusted_array( + columns, + DataFrame(True, index=query_days, columns=self.assets), + ) + + expected_baseline_closes = self.bcolz_writer.expected_values_2d( + query_days, + self.assets, + 'close', + ) + expected_baseline_volumes = self.bcolz_writer.expected_values_2d( + query_days, + self.assets, + 'volume', + ) + + # AdjustedArrays should yield the same data as the expected baseline. + for windowlen in range(1, len(query_days) + 1): + for offset, window in enumerate(closes.traverse(windowlen)): + assert_array_equal( + expected_baseline_closes[offset:offset + windowlen], + window, + ) + + for offset, window in enumerate(volumes.traverse(windowlen)): + assert_array_equal( + expected_baseline_volumes[offset:offset + windowlen], + window, + ) + + # Verify that we checked up to the longest possible window. + with self.assertRaises(WindowLengthTooLong): + closes.traverse(windowlen + 1) + with self.assertRaises(WindowLengthTooLong): + volumes.traverse(windowlen + 1) + + def apply_adjustments(self, dates, assets, baseline_values, adjustments): + min_date, max_date = dates[[0, -1]] + values = baseline_values.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: + 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 + + def test_read_with_adjustments(self): + columns = [USEquityPricing.high, USEquityPricing.volume] + query_days = self.calendar_days_between( + TEST_QUERY_START, + TEST_QUERY_STOP + ) + + baseline_reader = BcolzDailyBarReader(self.bcolz_path) + adjustment_reader = SQLiteAdjustmentReader(self.db_path) + pricing_loader = USEquityPricingLoader( + baseline_reader, + adjustment_reader, + ) + + closes, volumes = pricing_loader.load_adjusted_array( + columns, + DataFrame(True, index=query_days, columns=arange(1, 7)), + ) + + expected_baseline_highs = self.bcolz_writer.expected_values_2d( + query_days, + self.assets, + 'high', + ) + expected_baseline_volumes = self.bcolz_writer.expected_values_2d( + query_days, + self.assets, + 'volume', + ) + + # 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)): + baseline = expected_baseline_highs[offset:offset + windowlen] + baseline_dates = query_days[offset:offset + windowlen] + expected_adjusted_highs = self.apply_adjustments( + baseline_dates, + self.assets, + baseline, + # Apply all adjustments. + concat([SPLITS, MERGERS, DIVIDENDS], ignore_index=True), + ) + assert_allclose(expected_adjusted_highs, window) + + for offset, window in enumerate(volumes.traverse(windowlen)): + baseline = expected_baseline_volumes[offset:offset + windowlen] + baseline_dates = query_days[offset:offset + windowlen] + # 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, + baseline, + adjustments, + ) + # FIXME: Make AdjustedArray properly support integral types. + assert_array_equal( + expected_adjusted_volumes, + window.astype(uint32), + ) + + # Verify that we checked up to the longest possible window. + with self.assertRaises(WindowLengthTooLong): + closes.traverse(windowlen + 1) + with self.assertRaises(WindowLengthTooLong): + volumes.traverse(windowlen + 1) diff --git a/tests/resources/modelling_inputs/AAPL.csv b/tests/resources/modelling_inputs/AAPL.csv new file mode 100644 index 00000000..dad5c845 --- /dev/null +++ b/tests/resources/modelling_inputs/AAPL.csv @@ -0,0 +1,128 @@ +day,open,high,low,close,volume +2014-03-03,523.4200440000001,530.649956,522.8099900000001,527.76001,59695300 +2014-03-04,530.999977,532.6400150000001,527.769997,531.240036,64785000 +2014-03-05,530.919975,534.750023,529.1299740000001,532.360008,50015700 +2014-03-06,532.790031,534.4400019999999,528.100044,530.7499849999999,46372200 +2014-03-07,531.0900190000001,531.9799730000001,526.050011,530.4399639999999,55182400 +2014-03-10,528.3600230000001,533.330017,528.339996,530.919975,44646000 +2014-03-11,535.450012,538.740021,532.590027,536.090027,69806100 +2014-03-12,534.509964,537.350029,532.0,536.6099849999999,49831600 +2014-03-13,537.4399639999999,539.660042,529.160042,530.649956,64435700 +2014-03-14,528.789993,530.8900150000001,523.000008,524.68998,59299800 +2014-03-17,527.699982,529.969994,525.850006,526.740013,49886200 +2014-03-18,525.899994,531.9699860000001,525.200005,531.40004,52411800 +2014-03-19,532.259979,536.23999,528.9999849999999,531.26001,56189000 +2014-03-20,529.889992,532.669975,527.34996,528.700005,52099600 +2014-03-21,531.929985,533.75,526.330017,532.870033,93511600 +2014-03-24,538.41996,540.500008,535.0599900000001,539.1899639999999,88925200 +2014-03-25,541.499977,545.750008,539.590027,544.98999,70573300 +2014-03-26,546.5200120000001,549.0000150000001,538.8600230000001,539.779991,74942000 +2014-03-27,540.019997,541.499977,535.1199650000001,537.4599910000001,55507900 +2014-03-28,538.320038,538.940025,534.2499849999999,536.8600309999999,50141000 +2014-03-31,539.2300190000001,540.809975,535.930023,536.739975,42167300 +2014-04-01,537.760025,541.8700259999999,536.769989,541.649994,50190000 +2014-04-02,542.379997,543.479996,540.260002,542.549988,45105200 +2014-04-03,541.3900150000001,542.5,537.6399690000001,538.7900089999999,40586000 +2014-04-04,539.809952,540.000023,530.579994,531.820023,68812800 +2014-04-07,528.019989,530.900002,521.8899690000001,523.4700320000001,72462600 +2014-04-08,525.1899639999999,526.1200259999999,518.699989,523.4399639999999,60972100 +2014-04-09,522.639999,530.490005,522.0200120000001,530.320015,51542400 +2014-04-10,530.680023,532.240005,523.1699980000001,523.4800190000001,59913000 +2014-04-11,519.000023,522.830017,517.139954,519.6100230000001,67929400 +2014-04-14,521.899956,522.160042,517.209969,521.679977,51418500 +2014-04-15,520.2700120000001,521.6399769999999,511.32999400000006,517.9600519999999,66622500 +2014-04-16,518.049988,521.090004,514.139992,519.01001,53691400 +2014-04-17,519.999992,527.76001,519.200027,524.940025,71083600 +2014-04-21,525.339981,532.1399769999999,523.96003,531.170021,45637200 +2014-04-22,528.3100360000001,531.8299559999999,526.500008,531.699966,50640800 +2014-04-23,529.060013,531.129967,524.450027,524.750008,98735000 +2014-04-24,568.210014,570.0000150000001,560.730003,567.770004,189977900 +2014-04-25,564.529984,571.990021,563.959984,571.93998,97568800 +2014-04-28,572.799973,595.749977,572.550034,594.0900190000001,167371400 +2014-04-29,593.739998,595.979996,589.509995,592.329979,84344400 +2014-04-30,592.639999,599.430008,589.799988,590.089981,114160200 +2014-05-01,591.9999849999999,594.799995,586.359962,591.4799730000001,61012000 +2014-05-02,592.3400190000001,594.199982,589.7100519999999,592.580025,47878600 +2014-05-05,590.1399690000001,600.999977,589.999992,600.96003,71766800 +2014-05-06,601.799995,604.410042,594.409973,594.409973,93641100 +2014-05-07,595.249992,597.289986,587.7300339999999,592.329979,70716100 +2014-05-08,588.249992,594.409973,586.400017,587.990013,57574300 +2014-05-09,584.5399480000001,586.25,580.330025,585.53997,72899400 +2014-05-12,587.489975,593.659996,587.40004,592.830017,53302200 +2014-05-13,591.9999849999999,594.540016,590.699982,593.760025,39934300 +2014-05-14,592.430008,597.400002,591.740005,593.869987,41601000 +2014-05-15,594.699966,596.599983,588.0399480000001,588.819992,57711500 +2014-05-16,588.6299740000001,597.529991,585.399994,597.509964,69064100 +2014-05-19,597.849998,607.3300019999999,597.33004,604.5900190000001,79438800 +2014-05-20,604.509964,606.399994,600.730011,604.709969,58709000 +2014-05-21,603.8300019999999,606.700027,602.059975,606.309952,49214900 +2014-05-22,606.599998,609.850006,604.100021,607.2700269999999,50190000 +2014-05-23,607.25,614.730011,606.470009,614.129997,58052400 +2014-05-27,615.879997,625.8599849999999,615.630005,625.629967,87216500 +2014-05-28,626.019989,629.8299559999999,623.779991,624.01001,78870400 +2014-05-29,627.850044,636.87001,627.769989,635.37999,94118500 +2014-05-30,637.979996,644.1700440000001,628.900002,633.0000150000001,141005200 +2014-06-02,633.959984,634.830017,622.5000150000001,628.649956,92337700 +2014-06-03,628.4599910000001,638.740013,628.25,637.539986,73177300 +2014-06-04,637.4400099999999,647.8899690000001,636.110046,644.819992,83870500 +2014-06-05,646.20005,649.370003,642.610008,647.349983,75951400 +2014-06-06,649.900002,651.259979,644.469971,645.570023,87484600 +2014-06-09,92.699997,93.879997,91.75,93.699997,75415000 +2014-06-10,94.730003,95.050003,93.57,94.25,62777000 +2014-06-11,94.129997,94.760002,93.470001,93.860001,45681000 +2014-06-12,94.040001,94.120003,91.900002,92.290001,54749000 +2014-06-13,92.199997,92.440002,90.879997,91.279999,54525000 +2014-06-16,91.510002,92.75,91.449997,92.199997,35561000 +2014-06-17,92.309998,92.699997,91.800003,92.08000200000001,29726000 +2014-06-18,92.269997,92.290001,91.349998,92.18,33514000 +2014-06-19,92.290001,92.300003,91.339996,91.860001,35528000 +2014-06-20,91.849998,92.550003,90.900002,90.910004,100898000 +2014-06-23,91.32,91.620003,90.599998,90.83000200000001,43694000 +2014-06-24,90.75,91.739998,90.190002,90.279999,39036000 +2014-06-25,90.209999,90.699997,89.650002,90.360001,36869000 +2014-06-26,90.370003,91.050003,89.800003,90.900002,32629000 +2014-06-27,90.82,92.0,90.769997,91.980003,64029000 +2014-06-30,92.099998,93.730003,92.089996,92.93,49589000 +2014-07-01,93.519997,94.07,93.129997,93.519997,38223000 +2014-07-02,93.870003,94.059998,93.089996,93.480003,28465000 +2014-07-03,93.66999799999999,94.099998,93.199997,94.029999,22891800 +2014-07-07,94.139999,95.989998,94.099998,95.970001,56468000 +2014-07-08,96.269997,96.800003,93.91999799999999,95.349998,65222000 +2014-07-09,95.440002,95.949997,94.760002,95.389999,36436000 +2014-07-10,93.760002,95.550003,93.519997,95.040001,39686000 +2014-07-11,95.360001,95.889999,94.860001,95.220001,34018000 +2014-07-14,95.860001,96.889999,95.650002,96.449997,42810000 +2014-07-15,96.800003,96.849998,95.029999,95.32,45477900 +2014-07-16,96.970001,97.099998,94.739998,94.779999,53502000 +2014-07-17,95.029999,95.279999,92.57,93.089996,57298000 +2014-07-18,93.620003,94.739998,93.019997,94.43,49988000 +2014-07-21,94.989998,95.0,93.720001,93.940002,39079000 +2014-07-22,94.68,94.889999,94.120003,94.720001,55197000 +2014-07-23,95.41999799999999,97.879997,95.16999799999999,97.190002,92918000 +2014-07-24,97.040001,97.32,96.41999799999999,97.029999,45729000 +2014-07-25,96.849998,97.839996,96.639999,97.66999799999999,43469000 +2014-07-28,97.82,99.239998,97.550003,99.019997,55318000 +2014-07-29,99.33000200000001,99.440002,98.25,98.379997,43143000 +2014-07-30,98.440002,98.699997,97.66999799999999,98.150002,33010000 +2014-07-31,97.160004,97.449997,95.33000200000001,95.599998,56843000 +2014-08-01,94.900002,96.620003,94.809998,96.129997,48511000 +2014-08-04,96.370003,96.58000200000001,95.16999799999999,95.589996,39958000 +2014-08-05,95.360001,95.68,94.360001,95.120003,55933000 +2014-08-06,94.75,95.480003,94.709999,94.959999,38558000 +2014-08-07,94.93,95.949997,94.099998,94.480003,46711000 +2014-08-08,94.260002,94.82,93.279999,94.739998,41865000 +2014-08-11,95.269997,96.08000200000001,94.839996,95.989998,36585000 +2014-08-12,96.040001,96.879997,95.610001,95.970001,33795000 +2014-08-13,96.150002,97.239998,96.040001,97.239998,31916000 +2014-08-14,97.33000200000001,97.57,96.800003,97.5,28116000 +2014-08-15,97.900002,98.190002,96.860001,97.980003,48951000 +2014-08-18,98.489998,99.370003,97.980003,99.160004,47572000 +2014-08-19,99.410004,100.68,99.32,100.529999,69399000 +2014-08-20,100.440002,101.089996,99.949997,100.57,52699000 +2014-08-21,100.57,100.940002,100.110001,100.58000200000001,33478000 +2014-08-22,100.290001,101.470001,100.190002,101.32,44184000 +2014-08-25,101.790001,102.16999799999999,101.279999,101.540001,40270000 +2014-08-26,101.41999799999999,101.5,100.860001,100.889999,33152000 +2014-08-27,101.019997,102.57,100.699997,102.129997,52369000 +2014-08-28,101.589996,102.779999,101.559998,102.25,68460000 +2014-08-29,102.860001,102.900002,102.199997,102.5,44595000 diff --git a/tests/resources/modelling_inputs/BRK-A.csv b/tests/resources/modelling_inputs/BRK-A.csv new file mode 100644 index 00000000..1f61240d --- /dev/null +++ b/tests/resources/modelling_inputs/BRK-A.csv @@ -0,0 +1,128 @@ +day,open,high,low,close,volume +2014-03-03,174100.0,174997.0,172759.0,174500.0,800 +2014-03-04,175651.0,177989.0,175651.0,177989.0,700 +2014-03-05,177489.0,178750.0,177389.0,178655.0,400 +2014-03-06,179449.0,182214.0,179000.0,182175.0,600 +2014-03-07,183199.0,184355.0,182345.0,183772.0,400 +2014-03-10,183999.0,186400.0,183601.0,186400.0,400 +2014-03-11,186925.0,187490.0,185910.0,187101.0,600 +2014-03-12,186498.0,187832.0,186005.0,187750.0,300 +2014-03-13,188150.0,188852.0,185254.0,185750.0,700 +2014-03-14,185825.0,186507.0,183418.0,183860.0,600 +2014-03-17,184350.0,185790.0,184350.0,185050.0,400 +2014-03-18,185400.0,185400.0,183860.0,184860.0,200 +2014-03-19,184860.0,185489.0,182764.0,183860.0,200 +2014-03-20,183999.0,186742.0,183630.0,186540.0,300 +2014-03-21,187925.0,188598.0,187213.0,187850.0,400 +2014-03-24,188014.0,188380.0,186000.0,186520.0,200 +2014-03-25,187647.0,187647.0,185831.0,186587.0,200 +2014-03-26,187190.0,187472.0,184512.0,184540.0,300 +2014-03-27,184950.0,185735.0,184090.0,185200.0,300 +2014-03-28,185800.0,186433.0,183000.0,185149.0,200 +2014-03-31,186005.0,187400.0,186005.0,187350.0,600 +2014-04-01,187500.0,188000.0,186609.0,187213.0,300 +2014-04-02,186700.0,187280.0,186289.0,186759.0,700 +2014-04-03,186608.0,187024.0,185915.0,186287.0,300 +2014-04-04,186715.0,187925.0,185740.0,185753.0,200 +2014-04-07,185351.0,186133.0,184613.0,184700.0,300 +2014-04-08,184700.0,185735.0,183845.0,184640.0,900 +2014-04-09,185400.0,186121.0,184800.0,185897.0,300 +2014-04-10,186300.0,186984.0,183401.0,183402.0,400 +2014-04-11,183000.0,183735.0,182245.0,182759.0,300 +2014-04-14,183500.0,183700.0,181785.0,183212.0,400 +2014-04-15,183900.0,185870.0,183215.0,185640.0,400 +2014-04-16,186201.0,188984.0,186201.0,188900.0,400 +2014-04-17,188880.0,191506.0,188527.0,190639.0,400 +2014-04-21,191475.0,191475.0,189400.0,189482.0,300 +2014-04-22,189499.0,191031.0,188887.0,190720.0,600 +2014-04-23,191475.0,191880.0,190100.0,190800.0,100 +2014-04-24,191500.0,191833.0,189965.0,190500.0,200 +2014-04-25,190480.0,191315.0,189918.0,190686.0,200 +2014-04-28,192399.0,192399.0,189400.0,191400.0,200 +2014-04-29,191530.0,193098.0,191530.0,192545.0,300 +2014-04-30,192799.0,193789.0,192500.0,193275.0,200 +2014-05-01,193295.0,193800.0,192220.0,193482.0,200 +2014-05-02,194020.0,194530.0,191798.0,192255.0,200 +2014-05-05,189800.0,190979.0,189198.0,189790.0,400 +2014-05-06,189102.0,189346.0,186250.0,187375.0,400 +2014-05-07,188724.0,194670.0,188602.0,191550.0,500 +2014-05-08,191375.0,191375.0,189555.0,190100.0,100 +2014-05-09,189500.0,191035.0,188954.0,190883.0,100 +2014-05-12,192300.0,192320.0,190432.0,191115.0,200 +2014-05-13,191500.0,191826.0,190716.0,191826.0,100 +2014-05-14,191200.0,191900.0,191200.0,191420.0,200 +2014-05-15,190860.0,190860.0,189028.0,189371.0,200 +2014-05-16,190400.0,190400.0,188660.0,190210.0,200 +2014-05-19,190000.0,190746.0,189283.0,190400.0,200 +2014-05-20,190550.0,190550.0,188665.0,189201.0,200 +2014-05-21,189499.0,190366.0,189452.0,189975.0,200 +2014-05-22,189668.0,190737.0,189668.0,190535.0,200 +2014-05-23,190751.0,191000.0,189686.0,190205.0,100 +2014-05-27,191140.0,191649.0,190949.0,191300.0,200 +2014-05-28,191280.0,191800.0,191000.0,191356.0,200 +2014-05-29,191065.0,192300.0,190640.0,192300.0,200 +2014-05-30,191650.0,193000.0,191650.0,192000.0,300 +2014-06-02,192300.0,192522.0,191145.0,191748.0,100 +2014-06-03,191000.0,191632.0,190000.0,190417.0,200 +2014-06-04,190725.0,191357.0,190323.0,190766.0,200 +2014-06-05,190620.0,192309.0,190620.0,192100.0,300 +2014-06-06,192380.0,192908.0,192140.0,192895.0,300 +2014-06-09,192800.0,192878.0,191442.0,191917.0,300 +2014-06-10,191575.0,192420.0,191400.0,192306.0,200 +2014-06-11,191938.0,192658.0,191796.0,192357.0,100 +2014-06-12,192623.0,192720.0,190425.0,190950.0,200 +2014-06-13,190950.0,191266.0,188575.0,189520.0,1200 +2014-06-16,189000.0,189490.0,188721.0,189250.0,300 +2014-06-17,189110.0,189545.0,188493.0,188990.0,300 +2014-06-18,188990.0,190960.0,188220.0,190675.0,500 +2014-06-19,191000.0,191042.0,189899.0,190491.0,300 +2014-06-20,191721.0,191860.0,189789.0,190500.0,400 +2014-06-23,190215.0,190461.0,189477.0,189900.0,300 +2014-06-24,189889.0,191000.0,189624.0,190171.0,300 +2014-06-25,190000.0,190983.0,189733.0,190660.0,200 +2014-06-26,190849.0,190849.0,189232.0,190576.0,300 +2014-06-27,190400.0,191059.0,190050.0,190559.0,100 +2014-06-30,190200.0,190604.0,189799.0,189900.0,100 +2014-07-01,190100.0,191405.0,190100.0,190500.0,200 +2014-07-02,190800.0,191508.0,190333.0,191499.0,200 +2014-07-03,191940.0,193649.0,190900.0,193600.0,200 +2014-07-07,193333.0,193597.0,192333.0,193000.0,200 +2014-07-08,192560.0,193555.0,192295.0,192400.0,300 +2014-07-09,192617.0,193249.0,192502.0,193040.0,200 +2014-07-10,191911.0,193015.0,191799.0,192500.0,200 +2014-07-11,192220.0,193006.0,191945.0,192900.0,100 +2014-07-14,193503.0,193860.0,193030.0,193380.0,200 +2014-07-15,193480.0,193750.0,192000.0,192776.0,200 +2014-07-16,193145.0,193145.0,191849.0,192072.0,100 +2014-07-17,191920.0,192101.0,189696.0,189811.0,200 +2014-07-18,190860.0,192631.0,190555.0,192487.0,200 +2014-07-21,191860.0,193058.0,191554.0,193030.0,100 +2014-07-22,193296.0,193450.0,192339.0,192640.0,200 +2014-07-23,192577.0,192827.0,192140.0,192205.0,100 +2014-07-24,192200.0,192985.0,192107.0,192441.0,100 +2014-07-25,192436.0,192436.0,191190.0,191224.0,200 +2014-07-28,191120.0,192682.0,190000.0,192580.0,200 +2014-07-29,192310.0,192948.0,191083.0,191158.0,200 +2014-07-30,191200.0,192549.0,190483.0,191716.0,200 +2014-07-31,190807.0,191553.0,188124.0,188124.0,400 +2014-08-01,187152.0,190368.0,185005.0,189279.0,400 +2014-08-04,192126.0,195005.0,190942.0,194305.0,800 +2014-08-05,194500.0,194500.0,191358.0,192499.0,700 +2014-08-06,192189.0,194460.0,191890.0,193700.0,200 +2014-08-07,194670.0,194800.0,193079.0,194001.0,600 +2014-08-08,193850.0,196391.0,193570.0,196253.0,400 +2014-08-11,197699.0,199072.0,196500.0,198000.0,200 +2014-08-12,197895.0,199799.0,197895.0,199562.0,200 +2014-08-13,199960.0,199999.0,198800.0,199609.0,400 +2014-08-14,199600.0,203081.0,199600.0,202850.0,600 +2014-08-15,203350.0,203355.0,199721.0,201227.0,300 +2014-08-18,202400.0,203000.0,202100.0,202419.0,300 +2014-08-19,202805.0,202986.0,201941.0,202600.0,300 +2014-08-20,202364.0,202600.0,201891.0,202388.0,400 +2014-08-21,202799.0,205169.0,202679.0,205159.0,400 +2014-08-22,204600.0,204800.0,203455.0,203532.0,300 +2014-08-25,204600.0,204999.0,204233.0,204579.0,300 +2014-08-26,204880.0,205097.0,204358.0,204577.0,200 +2014-08-27,204577.0,205445.0,204000.0,204741.0,400 +2014-08-28,204300.0,204640.0,203895.0,204040.0,200 +2014-08-29,204050.0,205880.0,204050.0,205880.0,400 diff --git a/tests/resources/modelling_inputs/MSFT.csv b/tests/resources/modelling_inputs/MSFT.csv new file mode 100644 index 00000000..18fc0ea0 --- /dev/null +++ b/tests/resources/modelling_inputs/MSFT.csv @@ -0,0 +1,128 @@ +day,open,high,low,close,volume +2014-03-03,37.919998,38.130001,37.490002000000004,37.779999,29717500 +2014-03-04,38.200001,38.48,38.07,38.41,26802400 +2014-03-05,38.25,38.27,37.93,38.110001000000004,20520100 +2014-03-06,38.139998999999996,38.240002000000004,37.889998999999996,38.150002,23582200 +2014-03-07,38.279999,38.360001000000004,37.689999,37.900002,26591600 +2014-03-10,37.990002000000004,38.009997999999996,37.720001,37.82,19006600 +2014-03-11,37.869999,38.23,37.720001,38.02,25216400 +2014-03-12,37.799999,38.43,37.790001000000004,38.27,30494100 +2014-03-13,38.419998,38.450001,37.639998999999996,37.889998999999996,32169700 +2014-03-14,37.650002,38.139998999999996,37.509997999999996,37.700001,27195600 +2014-03-17,37.900002,38.41,37.790001000000004,38.049999,20479600 +2014-03-18,38.259997999999996,39.900002,38.220001,39.549999,64063900 +2014-03-19,39.470001,39.549999,38.91,39.27,35597200 +2014-03-20,39.25,40.650002,39.240002000000004,40.330002,59269800 +2014-03-21,40.720001,40.939999,40.009997999999996,40.16,80721800 +2014-03-24,40.34,40.639998999999996,39.860001000000004,40.5,46098400 +2014-03-25,40.66,40.990002000000004,39.959998999999996,40.34,43193100 +2014-03-26,40.48,40.709998999999996,39.599998,39.790001000000004,41977500 +2014-03-27,39.740002000000004,39.970001,39.34,39.360001000000004,35369200 +2014-03-28,39.790001000000004,40.639998999999996,39.68,40.299999,43472700 +2014-03-31,40.43,41.5,40.400002,40.990002000000004,46886300 +2014-04-01,41.150002,41.59,41.07,41.419998,32605000 +2014-04-02,41.439999,41.66,41.169998,41.349998,28666700 +2014-04-03,41.290001000000004,41.290001000000004,40.709998999999996,41.009997999999996,30139600 +2014-04-04,41.25,41.389998999999996,39.639998999999996,39.869999,51409600 +2014-04-07,39.959998999999996,40.27,39.740002000000004,39.799999,37559600 +2014-04-08,39.75,39.93,39.200001,39.82,35918600 +2014-04-09,39.93,40.549999,39.880001,40.470001,27398700 +2014-04-10,40.439999,40.689999,39.09,39.360001000000004,45960800 +2014-04-11,39.0,39.790001000000004,39.0,39.209998999999996,34330200 +2014-04-14,39.110001000000004,39.41,38.900002,39.18,32006600 +2014-04-15,39.34,39.959998999999996,39.049999,39.75,33968700 +2014-04-16,40.060001,40.419998,39.91,40.400002,30615800 +2014-04-17,40.009997999999996,40.200001,39.509997999999996,40.009997999999996,36689400 +2014-04-21,40.130001,40.150002,39.790001000000004,39.939999,22221200 +2014-04-22,39.959998999999996,40.139998999999996,39.830002,39.990002000000004,27056700 +2014-04-23,39.990002000000004,39.990002000000004,39.470001,39.689999,24602800 +2014-04-24,39.740002000000004,39.970001,39.299999,39.860001000000004,42381600 +2014-04-25,40.290001000000004,40.68,39.75,39.91,56876800 +2014-04-28,40.139998999999996,41.290001000000004,40.09,40.869999,50610200 +2014-04-29,41.099998,41.189999,40.389998999999996,40.509997999999996,29636200 +2014-04-30,40.400002,40.5,40.169998,40.400002,35458700 +2014-05-01,40.240002000000004,40.360001000000004,39.950001,40.0,28787400 +2014-05-02,40.310001,40.34,39.66,39.689999,43416600 +2014-05-05,39.52,39.639998999999996,39.299999,39.43,22460900 +2014-05-06,39.290001000000004,39.349998,38.950001,39.060001,27112400 +2014-05-07,39.220001,39.509997999999996,38.509997999999996,39.419998,41744500 +2014-05-08,39.34,39.900002,38.970001,39.639998999999996,32120400 +2014-05-09,39.540001000000004,39.849998,39.369999,39.540001000000004,29647600 +2014-05-12,39.740002000000004,40.02,39.650002,39.970001,22782600 +2014-05-13,39.919998,40.5,39.849998,40.419998,27004800 +2014-05-14,40.299999,40.450001,40.049999,40.240002000000004,18818700 +2014-05-15,40.09,40.400002,39.509997999999996,39.599998,37793200 +2014-05-16,39.669998,39.84,39.27,39.830002,29867100 +2014-05-19,39.610001000000004,39.82,39.459998999999996,39.75,24537400 +2014-05-20,39.68,39.939999,39.459998999999996,39.68,21320900 +2014-05-21,39.799999,40.349998,39.740002000000004,40.349998,22398700 +2014-05-22,40.290001000000004,40.349998,39.849998,40.099998,20201800 +2014-05-23,40.369999,40.369999,40.0,40.119999,18020000 +2014-05-27,40.259997999999996,40.259997999999996,39.810001,40.189999,26160600 +2014-05-28,40.139998999999996,40.189999,39.82,40.009997999999996,25711500 +2014-05-29,40.150002,40.349998,39.91,40.34,19888200 +2014-05-30,40.450001,40.970001,40.25,40.939999,34567600 +2014-06-02,40.950001,41.09,40.68,40.790001000000004,18504300 +2014-06-03,40.599998,40.68,40.25,40.290001000000004,18068900 +2014-06-04,40.209998999999996,40.369999,39.860001000000004,40.32,23209000 +2014-06-05,40.59,41.25,40.400002,41.209998999999996,31865200 +2014-06-06,41.48,41.66,41.240002000000004,41.48,24060500 +2014-06-09,41.389998999999996,41.48,41.02,41.27,15019200 +2014-06-10,41.029999,41.16,40.860001000000004,41.110001000000004,15117700 +2014-06-11,40.93,41.07,40.77,40.860001000000004,18040000 +2014-06-12,40.810001,40.880001,40.290001000000004,40.580002,29818900 +2014-06-13,41.099998,41.57,40.860001000000004,41.23,26310000 +2014-06-16,41.040001000000004,41.610001000000004,41.040001000000004,41.5,24205300 +2014-06-17,41.290001000000004,41.91,40.34,41.68,22518600 +2014-06-18,41.610001000000004,41.740002000000004,41.18,41.650002,27097000 +2014-06-19,41.57,41.77,41.330002,41.509997999999996,19828200 +2014-06-20,41.450001,41.830002,41.380001,41.68,47764900 +2014-06-23,41.73,42.0,41.689999,41.990002000000004,18743900 +2014-06-24,41.830002,41.939999,41.560001,41.75,26509100 +2014-06-25,41.700001,42.049999,41.459998999999996,42.029999,20049100 +2014-06-26,41.93,41.939999,41.43,41.720001,23604400 +2014-06-27,41.610001000000004,42.290001000000004,41.509997999999996,42.25,74640000 +2014-06-30,42.169998,42.209998999999996,41.700001,41.700001,30805500 +2014-07-01,41.860001000000004,42.150002,41.689999,41.869999,26917000 +2014-07-02,41.73,41.900002,41.529999,41.900002,20208100 +2014-07-03,41.91,41.990002000000004,41.560001,41.799999,15969300 +2014-07-07,41.75,42.119999,41.709998999999996,41.990002000000004,21952400 +2014-07-08,41.869999,42.0,41.610001000000004,41.779999,31218200 +2014-07-09,41.98,41.990002000000004,41.529999,41.669998,18445900 +2014-07-10,41.369999,42.0,41.049999,41.689999,21854700 +2014-07-11,41.700001,42.09,41.48,42.09,24083000 +2014-07-14,42.220001,42.450001,42.040001000000004,42.139998999999996,21881100 +2014-07-15,42.330002,42.470001,42.029999,42.450001,28748700 +2014-07-16,42.509997999999996,44.310001,42.48,44.080002,63318000 +2014-07-17,45.450001,45.709998999999996,44.25,44.529999,82180300 +2014-07-18,44.650002,44.84,44.25,44.689999,43407500 +2014-07-21,44.560001,45.16,44.220001,44.84,37604400 +2014-07-22,45.0,45.150002,44.59,44.830002,43095800 +2014-07-23,45.450001,45.450001,44.619999,44.869999,52362900 +2014-07-24,44.93,45.0,44.32,44.400002,30725300 +2014-07-25,44.299999,44.66,44.299999,44.5,26737700 +2014-07-28,44.360001000000004,44.509997999999996,43.93,43.970001,29684200 +2014-07-29,43.91,44.09,43.639998999999996,43.889998999999996,27763100 +2014-07-30,44.07,44.099998,43.290001000000004,43.580002,31921400 +2014-07-31,43.380001,43.689999,43.080002,43.16,31537500 +2014-08-01,43.209998999999996,43.25,42.599998,42.860001000000004,31170300 +2014-08-04,42.970001,43.470001,42.810001,43.369999,34277400 +2014-08-05,43.310001,43.459998999999996,42.830002,43.080002,26266400 +2014-08-06,42.740002000000004,43.169998,42.209998999999996,42.740002000000004,24634000 +2014-08-07,42.84,43.450001,42.650002,43.23,30314900 +2014-08-08,43.23,43.32,42.91,43.200001,28942700 +2014-08-11,43.259997999999996,43.450001,43.02,43.200001,20351600 +2014-08-12,43.040001000000004,43.59,43.0,43.52,21431100 +2014-08-13,43.68,44.18,43.52,44.080002,22889500 +2014-08-14,44.080002,44.419998,44.009997999999996,44.27,19313200 +2014-08-15,44.580002,44.900002,44.400002,44.790001000000004,41611300 +2014-08-18,44.939999,45.110001000000004,44.68,45.110001000000004,26891100 +2014-08-19,44.970001,45.34,44.830002,45.330002,28139500 +2014-08-20,45.34,45.400002,44.900002,44.950001,24770500 +2014-08-21,44.84,45.25,44.830002,45.220001,22285500 +2014-08-22,45.349998,45.470001,45.07,45.150002,18294500 +2014-08-25,45.400002,45.439999,45.040001000000004,45.169998,16910000 +2014-08-26,45.310001,45.400002,44.939999,45.009997999999996,14873100 +2014-08-27,44.900002,45.0,44.759997999999996,44.869999,21287900 +2014-08-28,44.75,44.98,44.610001000000004,44.880001,17657600 +2014-08-29,45.09,45.439999,44.860001000000004,45.43,21607600 diff --git a/tests/resources/modelling_inputs/generate.py b/tests/resources/modelling_inputs/generate.py new file mode 100644 index 00000000..6dfb563c --- /dev/null +++ b/tests/resources/modelling_inputs/generate.py @@ -0,0 +1,36 @@ +""" +Quick and dirty script to generate test case inputs. +""" +from __future__ import print_function +from os.path import ( + dirname, + join, +) +from pandas.io.data import get_data_yahoo + +here = join(dirname(__file__)) + + +def main(): + symbols = ['AAPL', 'MSFT', 'BRK-A'] + # Specifically chosen to include the AAPL split on June 9, 2014. + for symbol in symbols: + data = get_data_yahoo(symbol, start='2014-03-01', end='2014-09-01') + data.rename( + columns={ + 'Open': 'open', + 'High': 'high', + 'Low': 'low', + 'Close': 'close', + 'Volume': 'volume', + }, + inplace=True, + ) + del data['Adj Close'] + + dest = join(here, symbol + '.csv') + print("Writing %s -> %s" % (symbol, dest)) + data.to_csv(dest, index_label='day') + +if __name__ == '__main__': + main() diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py index f3fa3177..a18895a8 100644 --- a/tests/test_algorithm.py +++ b/tests/test_algorithm.py @@ -23,9 +23,10 @@ from unittest import TestCase import numpy as np import pandas as pd +from zipline.api_support import ZiplineAPI from zipline.assets import AssetFinder +from zipline.utils.control_flow import nullctx from zipline.utils.test_utils import ( - nullctx, setup_logger, teardown_logger ) @@ -150,6 +151,26 @@ class TestMiscellaneousAPI(TestCase): def tearDown(self): teardown_logger(self) + def test_zipline_api_resolves_dynamically(self): + # Make a dummy algo. + algo = TradingAlgorithm( + initialize=lambda context: None, + handle_data=lambda context, data: None, + sim_params=self.sim_params, + ) + + # Verify that api methods get resolved dynamically by patching them out + # and then calling them + for method in algo.all_api_methods(): + name = method.__name__ + sentinel = object() + + def fake_method(*args, **kwargs): + return sentinel + setattr(algo, name, fake_method) + with ZiplineAPI(algo): + self.assertIs(sentinel, getattr(zipline.api, name)()) + def test_get_environment(self): expected_env = { 'arena': 'backtest', diff --git a/tests/test_assets.py b/tests/test_assets.py index dafa7e7a..f70ac69a 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -24,10 +24,13 @@ from datetime import datetime, timedelta import pickle import uuid import warnings + import pandas as pd from pandas.tseries.tools import normalize_date +from pandas.util.testing import assert_frame_equal from nose_parameterized import parameterized +from numpy import full from zipline.assets import Asset, Equity, Future, AssetFinder from zipline.assets.futures import FutureChain @@ -37,6 +40,11 @@ from zipline.errors import ( SidAssignmentError, RootSymbolNotFound, ) +from zipline.finance.trading import with_environment +from zipline.utils.test_utils import ( + all_subindices, + make_rotating_asset_info, +) def build_lookup_generic_cases(): @@ -608,6 +616,49 @@ class AssetFinderTestCase(TestCase): post_map = finder.map_identifier_index_to_sids(pre_map, dt) self.assertListEqual([201, 2, 200, 1], post_map) + @with_environment() + def test_compute_lifetimes(self, env=None): + num_assets = 4 + trading_day = env.trading_day + first_start = pd.Timestamp('2015-04-01', tz='UTC') + + frame = make_rotating_asset_info( + num_assets=num_assets, + first_start=first_start, + frequency=env.trading_day, + periods_between_starts=3, + asset_lifetime=5 + ) + finder = AssetFinder(frame) + + all_dates = pd.date_range( + start=first_start, + end=frame.end_date.max(), + freq=trading_day, + ) + + for dates in all_subindices(all_dates): + expected_mask = full( + shape=(len(dates), num_assets), + fill_value=False, + dtype=bool, + ) + + for i, date in enumerate(dates): + it = frame[['start_date', 'end_date']].itertuples() + for j, start, end in it: + if start <= date <= end: + expected_mask[i, j] = True + + # Filter out columns with all-empty columns. + expected_result = pd.DataFrame( + data=expected_mask, + index=dates, + columns=frame.sid.values, + ) + actual_result = finder.lifetimes(dates) + assert_frame_equal(actual_result, expected_result) + class TestFutureChain(TestCase): metadata = { diff --git a/tests/test_doctests.py b/tests/test_doctests.py new file mode 100644 index 00000000..2bc8b35b --- /dev/null +++ b/tests/test_doctests.py @@ -0,0 +1,54 @@ +from __future__ import print_function +import sys +import doctest +from unittest import TestCase + +from zipline.lib import adjustment +from zipline.modelling import ( + engine, + expression, +) +from zipline.utils import ( + lazyval, + test_utils, +) + + +class DoctestTestCase(TestCase): + + @classmethod + def setUpClass(cls): + import pdb + # Workaround for the issue addressed by this (unmerged) PR to pdbpp: + # https://bitbucket.org/antocuni/pdb/pull-request/40/fix-ensure_file_can_write_unicode/diff # noqa + if '_pdbpp_path_hack' in pdb.__file__: + cls._skip = True + else: + cls._skip = False + + def _check_docs(self, module): + if self._skip: + # Printing this directly to __stdout__ so that it doesn't get + # captured by nose. + print("Warning: Skipping doctests for %s because " + "pdbpp is installed." % module.__name__, file=sys.__stdout__) + return + try: + doctest.testmod(module, verbose=True, raise_on_error=True) + except doctest.UnexpectedException as e: + raise e.exc_info[1] + + def test_adjustment_docs(self): + self._check_docs(adjustment) + + def test_expression_docs(self): + self._check_docs(expression) + + def test_engine_docs(self): + self._check_docs(engine) + + def test_lazyval_docs(self): + self._check_docs(lazyval) + + def test_test_utils_docs(self): + self._check_docs(test_utils) diff --git a/zipline/algorithm.py b/zipline/algorithm.py index 43aea515..34d6c787 100644 --- a/zipline/algorithm.py +++ b/zipline/algorithm.py @@ -31,16 +31,16 @@ from six import ( from operator import attrgetter from zipline.errors import ( + AddTermPostInit, OrderDuringInitialize, OverrideCommissionPostInit, OverrideSlippagePostInit, - RegisterTradingControlPostInit, RegisterAccountControlPostInit, + RegisterTradingControlPostInit, UnsupportedCommissionModel, UnsupportedOrderParameters, UnsupportedSlippageModel, ) - from zipline.finance.trading import TradingEnvironment from zipline.finance.blotter import Blotter from zipline.finance.commission import PerShare, PerTrade, PerDollar @@ -68,8 +68,16 @@ from zipline.assets import Asset, Future from zipline.assets.futures import FutureChain from zipline.gens.composites import date_sorted_sources from zipline.gens.tradesimulation import AlgorithmSimulator +from zipline.modelling.engine import ( + NoOpFFCEngine, + SimpleFFCEngine, +) from zipline.sources import DataFrameSource, DataPanelSource -from zipline.utils.api_support import ZiplineAPI, api_method +from zipline.utils.api_support import ( + api_method, + require_not_initialized, + ZiplineAPI, +) import zipline.utils.events from zipline.utils.events import ( EventManager, @@ -203,6 +211,21 @@ class TradingAlgorithm(object): # Pull in the environment's new AssetFinder for quick reference self.asset_finder = self.trading_environment.asset_finder + ffc_loader = kwargs.get('ffc_loader', None) + if ffc_loader is not None: + self.engine = SimpleFFCEngine( + ffc_loader, + self.trading_environment.trading_days, + self.asset_finder, + ) + else: + self.engine = NoOpFFCEngine() + + # Maps from name to Term + self._filters = {} + self._factors = {} + self._classifiers = {} + self.blotter = kwargs.pop('blotter', None) if not self.blotter: self.blotter = Blotter() @@ -1223,6 +1246,49 @@ class TradingAlgorithm(object): """ self.register_trading_control(LongOnly()) + ########### + # FFC API # + ########### + @api_method + @require_not_initialized(AddTermPostInit()) + def add_factor(self, factor, name): + if name in self._factors: + raise ValueError("Name %r is already a factor!" % name) + self._factors[name] = factor + + @api_method + @require_not_initialized(AddTermPostInit()) + def add_filter(self, filter): + name = "anon_filter_%d" % len(self._filters) + self._filters[name] = filter + + # Note: add_classifier is not yet implemented since you can't do anything + # useful with classifiers yet. + + def _all_terms(self): + # Merge all three dicts. + return dict( + chain.from_iterable( + iteritems(terms) + for terms in (self._filters, self._factors, self._classifiers) + ) + ) + + def compute_factor_matrix(self, start_date): + """ + Compute a factor matrix starting at start_date. + """ + days = self.trading_environment.trading_days + start_date_loc = days.get_loc(start_date) + sim_end = self.sim_params.period_end + end_loc = min(start_date_loc + 252, days.get_loc(sim_end)) + end_date = days[end_loc] + return self.engine.factor_matrix( + self._all_terms(), + start_date, + end_date, + ), end_date + def current_universe(self): return self._current_universe diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index 0ee4ac13..c847daf6 100644 --- a/zipline/assets/assets.py +++ b/zipline/assets/assets.py @@ -1,4 +1,3 @@ -# # Copyright 2015 Quantopian, Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -149,6 +148,9 @@ class AssetFinder(object): self._asset_type_cache = {} + # Populated on first call to `lifetimes`. + self._asset_lifetimes = None + def create_db_tables(self): c = self.conn.cursor() @@ -898,6 +900,70 @@ class AssetFinder(object): self._insert_metadata(identifier, **metadata_dict) self.conn.commit() + def _compute_asset_lifetimes(self): + """ + Compute and cache a recarry of asset lifetimes. + + FUTURE OPTIMIZATION: We're looping over a big array, which means this + probably should be in C/Cython. + """ + with self.conn as transaction: + results = transaction.execute( + 'SELECT sid, start_date, end_date from equities' + ).fetchall() + + lifetimes = np.recarray( + shape=(len(results),), + dtype=[('sid', 'i8'), ('start', 'i8'), ('end', 'i8')], + ) + + # TODO: This is **WAY** slower than it could be because we have to + # check for None everywhere. If we represented "no start date" as + # 0, and "no end date" as MAX_INT in our metadata, this would be + # significantly faster. + NO_START = 0 + NO_END = np.iinfo(int).max + for idx, (sid, start, end) in enumerate(results): + lifetimes[idx] = ( + sid, + start if start is not None else NO_START, + end if end is not None else NO_END, + ) + return lifetimes + + def lifetimes(self, dates): + """ + Compute a DataFrame representing asset lifetimes for the specified date + range. + + Parameters + ---------- + dates : pd.DatetimeIndex + The dates for which to compute lifetimes. + + Returns + ------- + lifetimes : pd.DataFrame + A frame of dtype bool with `dates` as index and an Int64Index of + assets as columns. The value at `lifetimes.loc[date, asset]` will + be True iff `asset` existed on `data`. + + See Also + -------- + numpy.putmask + """ + # This is a less than ideal place to do this, because if someone adds + # assets to the finder after we've touched lifetimes we won't have + # those new assets available. Mutability is not my favorite + # programming feature. + if self._asset_lifetimes is None: + self._asset_lifetimes = self._compute_asset_lifetimes() + lifetimes = self._asset_lifetimes + + raw_dates = dates.asi8[:, None] + mask = (lifetimes.start <= raw_dates) & (raw_dates <= lifetimes.end) + return pd.DataFrame(mask, index=dates, columns=lifetimes.sid) + class AssetConvertible(with_metaclass(ABCMeta)): """ @@ -908,6 +974,7 @@ class AssetConvertible(with_metaclass(ABCMeta)): """ pass + AssetConvertible.register(Integral) AssetConvertible.register(Asset) # Use six.string_types for Python2/3 compatibility diff --git a/zipline/data/dataset.py b/zipline/data/dataset.py new file mode 100644 index 00000000..f711b083 --- /dev/null +++ b/zipline/data/dataset.py @@ -0,0 +1,103 @@ +""" +dataset.py +""" +from six import ( + iteritems, + with_metaclass, +) + +from zipline.modelling.term import Term + + +class Column(object): + """ + An abstract column of data, not yet associated with a dataset. + """ + + def __init__(self, dtype): + self.dtype = dtype + + def bind(self, dataset, name): + """ + Bind a column to a concrete dataset. + """ + return BoundColumn(dtype=self.dtype, dataset=dataset, name=name) + + +class BoundColumn(Term): + """ + A Column of data that's been concretely bound to a particular dataset. + """ + + def __new__(cls, dtype, dataset, name): + return super(BoundColumn, cls).__new__( + cls, + inputs=(), + window_length=0, + domain=dataset.domain, + dtype=dtype, + dataset=dataset, + name=name, + ) + + def _init(self, dataset, name, *args, **kwargs): + self._dataset = dataset + self._name = name + return super(BoundColumn, self)._init(*args, **kwargs) + + @classmethod + def static_identity(cls, dataset, name, *args, **kwargs): + return ( + super(BoundColumn, cls).static_identity(*args, **kwargs), + dataset, + name, + ) + + @property + def dataset(self): + return self._dataset + + @property + def name(self): + return self._name + + @property + def qualname(self): + """ + Fully qualified of this column. + """ + return '.'.join([self.dataset.__name__, self.name]) + + def __repr__(self): + return "{qualname}::{dtype}".format( + qualname=self.qualname, + dtype=self.dtype.__name__, + ) + + +class DataSetMeta(type): + """ + Metaclass for DataSets + + Supplies name and dataset information to Column attributes. + """ + + def __new__(mcls, name, bases, dict_): + newtype = type.__new__(mcls, name, bases, dict_) + _columns = [] + for maybe_colname, maybe_column in iteritems(dict_): + if isinstance(maybe_column, Column): + bound_column = maybe_column.bind(newtype, maybe_colname) + setattr(newtype, maybe_colname, bound_column) + _columns.append(bound_column) + + newtype._columns = _columns + return newtype + + @property + def columns(self): + return self._columns + + +class DataSet(with_metaclass(DataSetMeta)): + domain = None diff --git a/zipline/data/equities.py b/zipline/data/equities.py new file mode 100644 index 00000000..945a35f1 --- /dev/null +++ b/zipline/data/equities.py @@ -0,0 +1,18 @@ +from numpy import ( + float64, + uint32, +) + +from zipline.data.dataset import ( + Column, + DataSet, +) + + +class USEquityPricing(DataSet): + + open = Column(float64) + high = Column(float64) + low = Column(float64) + close = Column(float64) + volume = Column(uint32) diff --git a/zipline/data/ffc/__init__.py b/zipline/data/ffc/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/zipline/data/ffc/base.py b/zipline/data/ffc/base.py new file mode 100644 index 00000000..0554d7cd --- /dev/null +++ b/zipline/data/ffc/base.py @@ -0,0 +1,21 @@ +""" +Base class for FFC data loaders. +""" +from abc import ( + ABCMeta, + abstractmethod, +) + + +from six import with_metaclass + + +class FFCLoader(with_metaclass(ABCMeta)): + """ + ABC for classes that can load data for use with zipline.modelling pipeline. + + TODO: DOCUMENT THIS MORE! + """ + @abstractmethod + def load_adjusted_array(self, columns, mask): + pass diff --git a/zipline/data/ffc/frame.py b/zipline/data/ffc/frame.py new file mode 100644 index 00000000..7f78a5e6 --- /dev/null +++ b/zipline/data/ffc/frame.py @@ -0,0 +1,185 @@ +""" +FFC Loader accepting a DataFrame as input. +""" +from numpy import ( + ix_, + zeros, +) +from pandas import ( + DataFrame, + DatetimeIndex, + Index, + Int64Index, +) +from zipline.lib.adjusted_array import adjusted_array +from zipline.lib.adjustment import ( + Float64Add, + Float64Multiply, + Float64Overwrite, +) +from zipline.data.ffc.base import FFCLoader + + +ADD, MULTIPLY, OVERWRITE = range(3) +ADJUSTMENT_CONSTRUCTORS = { + ADD: Float64Add.from_assets_and_dates, + MULTIPLY: Float64Multiply.from_assets_and_dates, + OVERWRITE: Float64Overwrite.from_assets_and_dates, +} +ADJUSTMENT_COLUMNS = Index([ + 'sid', + 'value', + 'kind', + 'start_date', + 'end_date', + 'apply_date', +]) + + +class DataFrameFFCLoader(FFCLoader): + """ + An FFCLoader that reads its input from DataFrames. + + Mostly useful for testing, but can also be used for real work if your data + fits in memory. + + Parameters + ---------- + column : zipline.data.dataset.BoundColumn + The column whose data is loadable by this loader. + + baseline : pandas.DataFrame + A DataFrame with index of type DatetimeIndex and columns of type + Int64Index. + + adjustments : pandas.DataFrame, default=None + A DataFrame with the following columns: + sid : int + value : any + kind : int (zipline.data.ffc.frame.ADJUSTMENT_TYPES) + start_date : datetime64 (can be NaT) + end_date : datetime64 (must be set) + apply_date : datetime64 (must be set) + + The default of None is interpreted as "no adjustments to the baseline". + """ + + def __init__(self, column, baseline, adjustments=None): + self.column = column + self.baseline = baseline.values + self.dates = baseline.index + self.assets = baseline.columns + + if adjustments is None: + adjustments = DataFrame( + index=DatetimeIndex([]), + columns=ADJUSTMENT_COLUMNS, + ) + else: + # Ensure that columns are in the correct order. + adjustments = adjustments.reindex_axis(ADJUSTMENT_COLUMNS, axis=1) + adjustments.sort(['apply_date', 'sid'], inplace=True) + + self.adjustments = adjustments + self.adjustment_apply_dates = DatetimeIndex(adjustments.apply_date) + self.adjustment_end_dates = DatetimeIndex(adjustments.end_date) + self.adjustment_sids = Int64Index(adjustments.sid) + + def format_adjustments(self, dates, assets): + """ + Build a dict of Adjustment objects in the format expected by + adjusted_array. + + Returns a dict of the form: + { + # Integer index into `dates` for the date on which we should + # apply the list of adjustments. + 1 : [ + Float64Multiply(first_row=2, last_row=4, col=3, value=0.5), + Float64Overwrite(first_row=3, last_row=5, col=1, value=2.0), + ... + ], + ... + } + """ + min_date, max_date = dates[[0, -1]] + # TODO: Consider porting this to Cython. + if len(self.adjustments) == 0: + return {} + + # Mask for adjustments whose apply_dates are in the requested window of + # dates. + date_bounds = self.adjustment_apply_dates.slice_indexer( + min_date, + max_date, + ) + dates_filter = zeros(len(self.adjustments), dtype='bool') + dates_filter[date_bounds] = True + # Ignore adjustments whose apply_date is in range, but whose end_date + # is out of range. + dates_filter &= (self.adjustment_end_dates >= min_date) + + # Mask for adjustments whose sids are in the requested assets. + sids_filter = self.adjustment_sids.isin(assets.values) + + adjustments_to_use = self.adjustments.loc[ + dates_filter & sids_filter + ].set_index('apply_date') + + # For each apply_date on which we have an adjustment, compute + # the integer index of that adjustment's apply_date in `dates`. + # Then build a list of Adjustment objects for that apply_date. + # This logic relies on the sorting applied on the previous line. + out = {} + previous_apply_date = object() + for row in adjustments_to_use.itertuples(): + # This expansion depends on the ordering of the DataFrame columns, + # defined above. + apply_date, sid, value, kind, start_date, end_date = row + if apply_date != previous_apply_date: + # Get the next apply date if no exact match. + row_loc = dates.get_loc(apply_date, method='bfill') + current_date_adjustments = out[row_loc] = [] + previous_apply_date = apply_date + + # Look up the approprate Adjustment constructor based on the value + # of `kind`. + current_date_adjustments.append( + ADJUSTMENT_CONSTRUCTORS[kind]( + dates, + assets, + start_date, + end_date, + sid, + value, + ), + ) + return out + + def load_adjusted_array(self, columns, mask): + """ + Load data from our stored baseline. + """ + if len(columns) != 1: + raise ValueError( + "Can't load multiple columns with DataFrameLoader" + ) + elif columns[0] != self.column: + raise ValueError("Can't load unknown column %s" % columns[0]) + + dates, assets, mask_values = mask.index, mask.columns, mask.values + + date_indexer = self.dates.get_indexer(dates) + assets_indexer = self.assets.get_indexer(assets) + + # Boolean arrays with True on matched entries + good_dates = (date_indexer != -1) + good_assets = (assets_indexer != -1) + + return adjusted_array( + # Pull out requested columns/rows from our baseline data. + data=self.baseline[ix_(date_indexer, assets_indexer)], + # Mask out requested columns/rows that didnt match. + mask=(good_assets & good_dates[:, None]) & mask_values, + adjustments=self.format_adjustments(dates, assets), + ) diff --git a/zipline/data/ffc/loaders/__init__.py b/zipline/data/ffc/loaders/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/zipline/data/ffc/loaders/_us_equity_pricing.pyx b/zipline/data/ffc/loaders/_us_equity_pricing.pyx new file mode 100644 index 00000000..d71cdf39 --- /dev/null +++ b/zipline/data/ffc/loaders/_us_equity_pricing.pyx @@ -0,0 +1,459 @@ +# +# 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. +from cpython cimport ( + PyDict_Contains, + 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 + +from zipline.lib.adjustment import Float64Multiply + +_SID_QUERY_TEMPLATE = """ +SELECT DISTINCT sid FROM {0} +WHERE effective_date >= ? 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 + int last_row + 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: + date_loc = dates.get_loc( + Timestamp(eff_date, unit='s', tz='UTC'), + # 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) + 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, last_row, 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: + date_loc = dates.get_loc( + Timestamp(eff_date, unit='s', tz='UTC'), + # 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) + 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: + date_loc = dates.get_loc( + Timestamp(eff_date, unit='s', tz='UTC'), + # 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) + 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 + + +@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 new file mode 100644 index 00000000..90b6e17c --- /dev/null +++ b/zipline/data/ffc/loaders/us_equity_pricing.py @@ -0,0 +1,638 @@ +# 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. +from abc import ( + ABCMeta, + abstractmethod, +) +from contextlib import contextmanager +from errno import ENOENT +from os import remove +from os.path import exists + +from bcolz import ( + carray, + ctable, +) +from click import progressbar +from numpy import ( + array, + array_equal, + float64, + floating, + full, + iinfo, + integer, + issubdtype, + uint32, +) +from pandas import ( + DatetimeIndex, + read_csv, + Timestamp, +) +from six import ( + iteritems, + string_types, + with_metaclass, +) +import sqlite3 + + +from zipline.data.ffc.base import FFCLoader +from zipline.data.ffc.loaders._us_equity_pricing import ( + _compute_row_slices, + _read_bcolz_data, + load_adjustments_from_sqlite, +) +from zipline.lib.adjusted_array import ( + adjusted_array, +) +from zipline.errors import NoFurtherDataError + +OHLC = frozenset(['open', 'high', 'low', 'close']) +US_EQUITY_PRICING_BCOLZ_COLUMNS = [ + 'open', 'high', 'low', 'close', 'volume', 'day', 'id' +] +DAILY_US_EQUITY_PRICING_DEFAULT_FILENAME = 'daily_us_equity_pricing.bcolz' +SQLITE_ADJUSTMENT_COLUMNS = frozenset(['effective_date', 'ratio', 'sid']) +SQLITE_ADJUSTMENT_COLUMN_DTYPES = { + 'effective_date': integer, + 'ratio': floating, + 'sid': integer, +} +SQLITE_ADJUSTMENT_TABLENAMES = frozenset(['splits', 'dividends', 'mergers']) + +UINT32_MAX = iinfo(uint32).max + + +@contextmanager +def passthrough(obj): + yield obj + + +class BcolzDailyBarWriter(with_metaclass(ABCMeta)): + """ + Class capable of writing daily OHLCV data to disk in a format that can be + read efficiently by BcolzDailyOHLCVReader. + + See Also + -------- + BcolzDailyBarReader : Consumer of the data written by this class. + """ + + @abstractmethod + def gen_tables(self, assets): + """ + Return an iterator of pairs of (asset_id, bcolz.ctable). + """ + raise NotImplementedError() + + @abstractmethod + def to_uint32(self, array, colname): + """ + Convert raw column values produced by gen_tables into uint32 values. + + Parameters + ---------- + array : np.array + An array of raw values. + colname : str, {'open', 'high', 'low', 'close', 'volume', 'day'} + The name of the column being loaded. + + For output being read by the default BcolzOHLCVReader, data should be + stored in the following manner: + + - Pricing columns (Open, High, Low, Close) should be stored as 1000 * + as-traded dollar value. + - Volume should be the as-traded volume. + - Dates should be stored as seconds since midnight UTC, Jan 1, 1970. + """ + raise NotImplementedError() + + def write(self, filename, calendar, assets, show_progress=False): + """ + Parameters + ---------- + filename : str + The location at which we should write our output. + calendar : pandas.DatetimeIndex + Calendar to use to compute asset calendar offsets. + assets : pandas.Int64Index + The assets for which to write data. + show_progress : bool + Whether or not to show a progress bar while writing. + + Returns + ------- + table : bcolz.ctable + The newly-written table. + """ + _iterator = self.gen_tables(assets) + if show_progress: + pbar = progressbar( + _iterator, + length=len(assets), + item_show_func=lambda i: i if i is None else str(i[0]), + label="Merging asset files:", + ) + with pbar as pbar_iterator: + return self._write_internal(filename, calendar, pbar_iterator) + return self._write_internal(filename, calendar, _iterator) + + def _write_internal(self, filename, calendar, iterator): + """ + Internal implementation of write. + + `iterator` should be an iterator yielding pairs of (asset, ctable). + """ + total_rows = 0 + first_row = {} + last_row = {} + calendar_offset = {} + + # Maps column name -> output carray. + columns = { + k: carray(array([], dtype=uint32)) + for k in US_EQUITY_PRICING_BCOLZ_COLUMNS + } + + for asset_id, table in iterator: + nrows = len(table) + for column_name in columns: + if column_name == 'id': + # We know what the content of this column is, so don't + # bother reading it. + columns['id'].append(full((nrows,), asset_id)) + continue + columns[column_name].append( + self.to_uint32(table[column_name][:], column_name) + ) + + # Bcolz doesn't support ints as keys in `attrs`, so convert + # assets to strings for use as attr keys. + asset_key = str(asset_id) + + # Calculate the index into the array of the first and last row + # for this asset. This allows us to efficiently load single + # assets when querying the data back out of the table. + first_row[asset_key] = total_rows + last_row[asset_key] = total_rows + nrows - 1 + total_rows += nrows + + # Calculate the number of trading days between the first date + # in the stored data and the first date of **this** asset. This + # offset used for output alignment by the reader. + + # HACK: Index with a list so that we get back an array we can pass + # to self.to_uint32. We could try to extract this in the loop + # above, but that makes the logic a lot messier. + asset_first_day = self.to_uint32(table['day'][[0]], 'day')[0] + calendar_offset[asset_key] = calendar.get_loc( + Timestamp(asset_first_day, unit='s', tz='UTC'), + ) + + # This writes the table to disk. + full_table = ctable( + columns=[ + columns[colname] + for colname in US_EQUITY_PRICING_BCOLZ_COLUMNS + ], + names=US_EQUITY_PRICING_BCOLZ_COLUMNS, + rootdir=filename, + mode='w', + ) + full_table.attrs['first_row'] = first_row + full_table.attrs['last_row'] = last_row + full_table.attrs['calendar_offset'] = calendar_offset + full_table.attrs['calendar'] = calendar.asi8.tolist() + return full_table + + +class DailyBarWriterFromCSVs(BcolzDailyBarWriter): + """ + BcolzDailyBarWriter constructed from a map from csvs to assets. + + Parameters + ---------- + asset_map : dict + A map from asset_id -> path to csv with data for that asset. + + CSVs should have the following columns: + day : datetime64 + open : float64 + high : float64 + low : float64 + close : float64 + volume : int64 + """ + _csv_dtypes = { + 'open': float64, + 'high': float64, + 'low': float64, + 'close': float64, + 'volume': float64, + } + + def __init__(self, asset_map): + self._asset_map = asset_map + + def gen_tables(self, assets): + """ + Read CSVs as DataFrames from our asset map. + """ + dtypes = self._csv_dtypes + for asset in assets: + path = self._asset_map.get(asset) + if path is None: + raise KeyError("No path supplied for asset %s" % asset) + data = read_csv(path, parse_dates=['day'], dtype=dtypes) + yield asset, ctable.fromdataframe(data) + + def to_uint32(self, array, colname): + arrmax = array.max() + if colname in OHLC: + self.check_uint_safe(arrmax * 1000, colname) + return (array * 1000).astype(uint32) + elif colname == 'volume': + self.check_uint_safe(arrmax, colname) + return array.astype(uint32) + elif colname == 'day': + nanos_per_second = (1000 * 1000 * 1000) + self.check_uint_safe(arrmax.view(int) / nanos_per_second, colname) + return (array.view(int) / nanos_per_second).astype(uint32) + + @staticmethod + def check_uint_safe(value, colname): + if value >= UINT32_MAX: + raise ValueError( + "Value %s from column '%s' is too large" % (value, colname) + ) + + +class BcolzDailyBarReader(object): + """ + Reader for raw pricing data written by BcolzDailyOHLCVWriter. + + A Bcolz CTable is comprised of Columns and Attributes. + + Columns + ------- + The table with which this loader interacts contains the following columns: + + ['open', 'high', 'low', 'close', 'volume', 'day', 'id']. + + The data in these columns is interpreted as follows: + + - Price columns ('open', 'high', 'low', 'close') are interpreted as 1000 * + as-traded dollar value. + - Volume is interpreted as as-traded volume. + - Day is interpreted as seconds since midnight UTC, Jan 1, 1970. + - Id is the asset id of the row. + + The data in each column is grouped by asset and then sorted by day within + each asset block. + + The table is built to represent a long time range of data, e.g. ten years + of equity data, so the lengths of each asset block is not equal to each + other. The blocks are clipped to the known start and end date of each asset + to cut down on the number of empty values that would need to be included to + make a regular/cubic dataset. + + When read across the open, high, low, close, and volume with the same + index should represent the same asset and day. + + Attributes + ---------- + The table with which this loader interacts contains the following + attributes: + + first_row : dict + Map from asset_id -> index of first row in the dataset with that id. + last_row : dict + Map from asset_id -> index of last row in the dataset with that id. + calendar_offset : dict + Map from asset_id -> calendar index of first row. + calendar : list[int64] + Calendar used to compute offsets, in asi8 format (ns since EPOCH). + + We use first_row and last_row together to quickly find ranges of rows to + load when reading an asset's data into memory. + + We use calendar_offset and calendar to orient loaded blocks within a + range of queried dates. + """ + def __init__(self, table): + if isinstance(table, string_types): + table = ctable(rootdir=table, mode='r') + + self._table = table + self._calendar = DatetimeIndex(table.attrs['calendar'], tz='UTC') + self._first_rows = { + int(asset_id): start_index + for asset_id, start_index in iteritems(table.attrs['first_row']) + } + self._last_rows = { + int(asset_id): end_index + for asset_id, end_index in iteritems(table.attrs['last_row']) + } + self._calendar_offsets = { + int(id_): offset + 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): + """ + Compute the raw row indices to load for each asset on a query for the + given dates. + + Parameters + ---------- + dates : pandas.DatetimeIndex + Dates of the query on which we want to compute row indices. + assets : pandas.Int64Index + Assets for which we want to compute row indices + + Returns + ------- + A 3-tuple of (first_rows, last_rows, offsets): + first_rows : np.array[intp] + Array with length == len(assets) containing the index of the first + row to load for each asset in `assets`. + last_rows : np.array[intp] + Array with length == len(assets) containing the index of the last + row to load for each asset in `assets`. + offset : np.array[intp] + Array with length == (len(asset) containing the index in a buffer + of length `dates` corresponding to the first row of each asset. + + The value of offset[i] will be 0 if asset[i] existed at the start + 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, + assets, + ) + + def load_raw_arrays(self, columns, dates, assets): + first_rows, last_rows, offsets = self._compute_slices(dates, assets) + return _read_bcolz_data( + self._table, + (len(dates), len(assets)), + [column.name for column in columns], + first_rows, + last_rows, + offsets, + ) + + +class SQLiteAdjustmentWriter(object): + """ + Writer for data to be read by SQLiteAdjustmentWriter + + Parameters + ---------- + conn_or_path : str or sqlite3.Connection + A handle to the target sqlite database. + overwrite : bool, optional, default=False + If True and conn_or_path is a string, remove any existing files at the + given path before connecting. + + See Also + -------- + SQLiteAdjustmentReader + """ + + def __init__(self, conn_or_path, overwrite=False): + if isinstance(conn_or_path, sqlite3.Connection): + self.conn = conn_or_path + elif isinstance(conn_or_path, str): + if overwrite and exists(conn_or_path): + try: + remove(conn_or_path) + except OSError as e: + if e.errno != ENOENT: + raise + self.conn = sqlite3.connect(conn_or_path) + else: + raise TypeError("Unknown connection type %s" % type(conn_or_path)) + + def write_frame(self, tablename, frame): + if frozenset(frame.columns) != SQLITE_ADJUSTMENT_COLUMNS: + raise ValueError( + "Unexpected frame columns:\n" + "Expected Columns: %s\n" + "Received Columns: %s" % ( + SQLITE_ADJUSTMENT_COLUMNS, + frame.columns.tolist(), + ) + ) + elif tablename not in SQLITE_ADJUSTMENT_TABLENAMES: + raise ValueError( + "Adjustment table %s not in %s" % ( + tablename, SQLITE_ADJUSTMENT_TABLENAMES + ) + ) + + expected_dtypes = SQLITE_ADJUSTMENT_COLUMN_DTYPES + actual_dtypes = frame.dtypes + for colname, expected in iteritems(expected_dtypes): + actual = actual_dtypes[colname] + if not issubdtype(actual, expected): + raise TypeError( + "Expected data of type {expected} for column '{colname}', " + "but got {actual}.".format( + expected=expected, + colname=colname, + actual=actual, + ) + ) + return frame.to_sql(tablename, self.conn) + + def write(self, splits, mergers, dividends): + """ + Writes data to a SQLite file to be read by SQLiteAdjustmentReader. + + Parameters + ---------- + splits : pandas.DataFrame + Dataframe containing split data. + mergers : pandas.DataFrame + DataFrame containing merger data. + dividends : pandas.DataFrame + DataFrame containing dividend data. + + Notes + ----- + DataFrame input (`splits`, `mergers`, and `dividends`) should all have + the following columns: + + effective_date : int + The date, represented as seconds since Unix epoch, on which the + adjustment should be applied. + ratio : float + A value to apply to all data earlier than the effective date. + sid : int + The asset id associated with this adjustment. + + The ratio column is interpreted as follows: + - For all adjustment types, multiply price fields ('open', 'high', + 'low', and 'close') by the ratio. + - For **splits only**, **divide** volume by the adjustment ratio. + + Dividend ratios should be calculated as + 1.0 - (dividend_value / "close on day prior to dividend ex_date"). + + Returns + ------- + None + + See Also + -------- + SQLiteAdjustmentReader : Consumer for the data written by this class + """ + self.write_frame('splits', splits) + self.write_frame('mergers', mergers) + self.write_frame('dividends', dividends) + self.conn.execute( + "CREATE INDEX splits_sids " + "ON splits(sid)" + ) + self.conn.execute( + "CREATE INDEX splits_effective_date " + "ON splits(effective_date)" + ) + self.conn.execute( + "CREATE INDEX mergers_sids " + "ON mergers(sid)" + ) + self.conn.execute( + "CREATE INDEX mergers_effective_date " + "ON mergers(effective_date)" + ) + self.conn.execute( + "CREATE INDEX dividends_sid " + "ON dividends(sid)" + ) + self.conn.execute( + "CREATE INDEX dividends_effective_date " + "ON dividends(effective_date)" + ) + + def close(self): + self.conn.close() + + +class SQLiteAdjustmentReader(object): + """ + Loads adjustments based on corporate actions from a SQLite database. + + Expects data written in the format output by `SQLiteAdjustmentWriter`. + + Parameters + ---------- + conn : str or sqlite3.Connection + Connection from which to load data. + """ + + def __init__(self, conn): + if isinstance(conn, str): + conn = sqlite3.connect(conn) + self.conn = conn + + def load_adjustments(self, columns, dates, assets): + return load_adjustments_from_sqlite( + self.conn, + [column.name for column in columns], + dates, + assets, + ) + + +class USEquityPricingLoader(FFCLoader): + """ + FFCLoader for US Equity Pricing + + Delegates loading of baselines and adjustments. + """ + + def __init__(self, raw_price_loader, adjustments_loader): + self.raw_price_loader = raw_price_loader + self.adjustments_loader = adjustments_loader + + def load_adjusted_array(self, columns, mask): + dates, assets = mask.index, mask.columns + raw_arrays = self.raw_price_loader.load_raw_arrays( + columns, + dates, + assets, + ) + adjustments = self.adjustments_loader.load_adjustments( + columns, + dates, + assets, + ) + + return [ + adjusted_array(raw_array, mask.values, col_adjustments) + for raw_array, col_adjustments in zip(raw_arrays, adjustments) + ] diff --git a/zipline/data/ffc/synthetic.py b/zipline/data/ffc/synthetic.py new file mode 100644 index 00000000..19e8cb33 --- /dev/null +++ b/zipline/data/ffc/synthetic.py @@ -0,0 +1,267 @@ +""" +Synthetic data loaders for testing. +""" +from bcolz import ctable + +from numpy import ( + arange, + array, + float64, + full, + iinfo, + uint32, +) +from pandas import ( + DataFrame, + Timestamp, +) +from sqlite3 import connect as sqlite3_connect + +from six import iteritems + +from zipline.data.ffc.base import FFCLoader +from zipline.data.ffc.frame import DataFrameFFCLoader +from zipline.data.ffc.loaders.us_equity_pricing import ( + BcolzDailyBarWriter, + SQLiteAdjustmentReader, + SQLiteAdjustmentWriter, + US_EQUITY_PRICING_BCOLZ_COLUMNS, +) + + +UINT_32_MAX = iinfo(uint32).max + + +def nanos_to_seconds(nanos): + return nanos / (1000 * 1000 * 1000) + + +class MultiColumnLoader(FFCLoader): + """ + FFCLoader that can delegate to sub-loaders. + + Parameters + ---------- + loaders : dict + Dictionary mapping columns -> loader + """ + def __init__(self, loaders): + self._loaders = loaders + + def load_adjusted_array(self, columns, mask): + """ + Load by delegating to sub-loaders. + """ + out = [] + for column in columns: + try: + loader = self._loaders[column] + except KeyError: + raise ValueError("Couldn't find loader for %s" % column) + out.append(loader.load_adjusted_array([column], mask)) + return out + + +class ConstantLoader(MultiColumnLoader): + """ + Synthetic FFCLoader that returns a constant value for each column. + + Parameters + ---------- + constants : dict + Map from column to value(s) to use for that column. + Values can be anything that can be passed as the first positional + argument to a DataFrame of the same shape as `mask`. + mask : pandas.DataFrame + Mask indicating when assets existed. + Indices of this frame are used to align input queries. + + Notes + ----- + Adjustments are unsupported with ConstantLoader. + """ + def __init__(self, constants, dates, assets): + loaders = {} + for column, const in iteritems(constants): + frame = DataFrame( + const, + index=dates, + columns=assets, + dtype=column.dtype, + ) + loaders[column] = DataFrameFFCLoader( + column=column, + baseline=frame, + adjustments=None, + ) + + super(ConstantLoader, self).__init__(loaders) + + +class SyntheticDailyBarWriter(BcolzDailyBarWriter): + """ + Bcolz writer that creates synthetic data based on asset lifetime metadata. + + For a given asset/date/column combination, we generate a corresponding raw + value using the following formula for OHLCV columns: + + data(asset, date, column) = (100,000 * asset_id) + + (10,000 * column_num) + + (date - Jan 1 2000).days # ~6000 for 2015 + where: + column_num('open') = 0 + column_num('high') = 1 + column_num('low') = 2 + column_num('close') = 3 + column_num('volume') = 4 + + We use days since Jan 1, 2000 to guarantee that there are no collisions + while also the produced values smaller than UINT32_MAX / 1000. + + For 'day' and 'id', we use the standard format expected by the base class. + + Parameters + ---------- + asset_info : DataFrame + DataFrame with asset_id as index and 'start_date'/'end_date' columns. + calendar : DatetimeIndex + Calendar to use for constructing asset lifetimes. + """ + OHLCV = ('open', 'high', 'low', 'close', 'volume') + OHLC = ('open', 'high', 'low', 'close') + PSEUDO_EPOCH = Timestamp('2000-01-01', tz='UTC') + + def __init__(self, asset_info, calendar): + super(SyntheticDailyBarWriter, self).__init__() + assert ( + # Using .value here to avoid having to care about UTC-aware dates. + self.PSEUDO_EPOCH.value < + calendar.min().value <= + asset_info['start_date'].min().value + ) + assert (asset_info['start_date'] < asset_info['end_date']).all() + self._asset_info = asset_info + self._calendar = calendar + + def _raw_data_for_asset(self, asset_id): + """ + Generate 'raw' data that encodes information about the asset. + + See class docstring for a description of the data format. + """ + # Get the dates for which this asset existed according to our asset + # info. + dates = self._calendar[ + self._calendar.slice_indexer( + self.asset_start(asset_id), self.asset_end(asset_id) + ) + ] + + data = full( + (len(dates), len(US_EQUITY_PRICING_BCOLZ_COLUMNS)), + asset_id * (100 * 1000), + dtype=uint32, + ) + + # Add 10,000 * column-index to OHLCV columns + data[:, :5] += arange(5) * (10 * 1000) + + # Add days since Jan 1 2001 for OHLCV columns. + data[:, :5] += (dates - self.PSEUDO_EPOCH).days[:, None] + + frame = DataFrame( + data, + index=dates, + columns=US_EQUITY_PRICING_BCOLZ_COLUMNS, + ) + + frame['day'] = nanos_to_seconds(dates.asi8) + frame['id'] = asset_id + + return ctable.fromdataframe(frame) + + def asset_start(self, asset): + ret = self._asset_info.loc[asset]['start_date'] + if ret.tz is None: + ret = ret.tz_localize('UTC') + assert ret.tzname() == 'UTC', "Unexpected non-UTC timestamp" + return ret + + def asset_end(self, asset): + ret = self._asset_info.loc[asset]['end_date'] + if ret.tz is None: + ret = ret.tz_localize('UTC') + assert ret.tzname() == 'UTC', "Unexpected non-UTC timestamp" + return ret + + @classmethod + def expected_value(cls, asset_id, date, colname): + """ + Check that the raw value for an asset/date/column triple is as + expected. + + Used by tests to verify data written by a writer. + """ + from_asset = asset_id * 100 * 1000 + from_colname = cls.OHLCV.index(colname) * (10 * 1000) + from_date = (date - cls.PSEUDO_EPOCH).days + return from_asset + from_colname + from_date + + def expected_values_2d(self, dates, assets, colname): + """ + Return an 2D array containing cls.expected_value(asset_id, date, + colname) for each date/asset pair in the inputs. + + Values before/after an assets lifetime are filled with 0 for volume and + NaN for price columns. + """ + if colname == 'volume': + dtype = uint32 + missing = 0 + else: + dtype = float64 + missing = float('nan') + + data = full((len(dates), len(assets)), missing, dtype=dtype) + for j, asset in enumerate(assets): + start, end = self.asset_start(asset), self.asset_end(asset) + for i, date in enumerate(dates): + # No value expected for dates outside the asset's start/end + # date. + if not (start <= date <= end): + continue + data[i, j] = self.expected_value(asset, date, colname) + return data + + # BEGIN SUPERCLASS INTERFACE + def gen_tables(self, assets): + for asset in assets: + yield asset, self._raw_data_for_asset(asset) + + def to_uint32(self, array, colname): + if colname in {'open', 'high', 'low', 'close'}: + # Data is stored as 1000 * raw value. + assert array.max() < (UINT_32_MAX / 1000), "Test data overflow!" + return array * 1000 + else: + assert colname in ('volume', 'day'), "Unknown column: %s" % colname + return array + # END SUPERCLASS INTERFACE + + +class NullAdjustmentReader(SQLiteAdjustmentReader): + """ + A SQLiteAdjustmentReader that stores no adjustments and uses in-memory + SQLite. + """ + + def __init__(self): + conn = sqlite3_connect(':memory:') + writer = SQLiteAdjustmentWriter(conn) + empty = DataFrame({ + 'sid': array([], dtype=uint32), + 'effective_date': array([], dtype=uint32), + 'ratio': array([], dtype=float), + }) + writer.write(splits=empty, mergers=empty, dividends=empty) + super(NullAdjustmentReader, self).__init__(conn) diff --git a/zipline/data/loader.py b/zipline/data/loader.py index abc4caed..8a001f4e 100644 --- a/zipline/data/loader.py +++ b/zipline/data/loader.py @@ -16,7 +16,6 @@ import importlib import os -from os.path import expanduser from collections import OrderedDict from datetime import timedelta @@ -30,25 +29,16 @@ from six import iteritems from . import benchmarks from . benchmarks import get_benchmark_returns +from .paths import ( + cache_root, + data_root, +) from zipline.utils.tradingcalendar import trading_day as trading_day_nyse from zipline.utils.tradingcalendar import trading_days as trading_days_nyse logger = logbook.Logger('Loader') -# TODO: Make this path customizable. -DATA_PATH = os.path.join( - expanduser("~"), - '.zipline', - 'data' -) - -CACHE_PATH = os.path.join( - expanduser("~"), - '.zipline', - 'cache' -) - # Mapping from index symbol to appropriate bond data INDEX_MAPPING = { '^GSPC': @@ -66,18 +56,20 @@ def get_data_filepath(name): Creates containing directory, if needed. """ + dr = data_root() - if not os.path.exists(DATA_PATH): - os.makedirs(DATA_PATH) + if not os.path.exists(dr): + os.makedirs(dr) - return os.path.join(DATA_PATH, name) + return os.path.join(dr, name) def get_cache_filepath(name): - if not os.path.exists(CACHE_PATH): - os.makedirs(CACHE_PATH) + cr = cache_root() + if not os.path.exists(cr): + os.makedirs(cr) - return os.path.join(CACHE_PATH, name) + return os.path.join(cr, name) def dump_treasury_curves(module='treasuries', filename='treasury_curves.csv'): diff --git a/zipline/data/paths.py b/zipline/data/paths.py new file mode 100644 index 00000000..59b0ab98 --- /dev/null +++ b/zipline/data/paths.py @@ -0,0 +1,91 @@ +""" +Canonical path locations for zipline data. + +Paths are rooted at $ZIPLINE_ROOT if that environment variable is set. +Otherwise default to expanduser(~/.zipline) +""" +import os +from os.path import ( + expanduser, + join, +) + + +def zipline_root(environ=None): + """ + Get the root directory for all zipline-managed files. + + For testing purposes, this accepts a dictionary to interpret as the os + environment. + + Parameters + ---------- + environ : dict, optional + A dict to interpret as the os environment. + + Returns + ------- + root : string + Path to the zipline root dir. + """ + if environ is None: + environ = os.environ.copy() + + root = environ.get('ZIPLINE_ROOT', None) + if root is None: + root = expanduser('~/.zipline') + + return root + + +def zipline_root_path(path, environ=None): + """ + Get a path relative to the zipline root. + + Parameters + ---------- + path : str + The requested path. + environ : dict, optional + An environment dict to forward to zipline_root. + + Returns + ------- + newpath : str + The requested path joined with the zipline root. + """ + return join(zipline_root(environ=environ), path) + + +def data_root(environ=None): + """ + The root directory for zipline data files. + + Parameters + ---------- + environ : dict, optional + An environment dict to forward to zipline_root. + + Returns + ------- + data_root : str + The zipline data root. + """ + return zipline_root_path('data', environ=environ) + + +def cache_root(environ=None): + """ + The root directory for zipline cache files. + + Parameters + ---------- + environ : dict, optional + An environment dict to forward to zipline_root. + + Returns + ------- + cache_root : str + The zipline cache root. + """ + return zipline_root_path('cache', environ=environ) diff --git a/zipline/errors.py b/zipline/errors.py index 54128814..c1a04831 100644 --- a/zipline/errors.py +++ b/zipline/errors.py @@ -303,3 +303,104 @@ Only one simulation date given. Please specify both the 'start' and 'end' for the simulation, or neither. If neither is given, the start and end of the DataSource will be used. Given start = '{start}', end = '{end}' """.strip() + + +class WindowLengthTooLong(ZiplineError): + """ + Raised when a trailing window is instantiated with a lookback greater than + the length of the underlying array. + """ + msg = ( + "Can't construct a rolling window of length " + "{window_length} on an array of length {nrows}." + ).strip() + + +class WindowLengthNotPositive(ZiplineError): + """ + Raised when a trailing window would be instantiated with a length less than + 1. + """ + msg = ( + "Expected a window_length greater than 0, got {window_length}." + ).strip() + + +class InputTermNotAtomic(ZiplineError): + """ + Raised when a non-atomic term is specified as an input to an FFC term with + a lookback window. + """ + msg = ( + "Can't compute {parent} with non-atomic input {child}." + ) + + +class TermInputsNotSpecified(ZiplineError): + """ + Raised if a user attempts to construct a term without specifying inputs and + that term does not have class-level default inputs. + """ + msg = "{termname} requires inputs, but no inputs list was passed." + + +class WindowLengthNotSpecified(ZiplineError): + """ + Raised if a user attempts to construct a term without specifying inputs and + that term does not have class-level default inputs. + """ + msg = ( + "{termname} requires a window_length, but no window_length was passed." + ) + + +class BadPercentileBounds(ZiplineError): + """ + Raised by API functions accepting percentile bounds when the passed bounds + are invalid. + """ + msg = ( + "Percentile bounds must fall between 0.0 and 100.0, and min must be " + "less than max." + "\nInputs were min={min_percentile}, max={max_percentile}." + ) + + +class UnknownRankMethod(ZiplineError): + """ + Raised during construction of a Rank factor when supplied a bad Rank + method. + """ + msg = ( + "Unknown ranking method: '{method}'. " + "`method` must be one of {choices}" + ) + + +class AddTermPostInit(ZiplineError): + """ + Raised when a user tries to call add_{filter,factor,classifier} + outside of initialize. + """ + msg = ( + "Attempted to add a new filter, factor, or classifier " + "outside of initialize.\n" + "New FFC terms may only be added during initialize." + ) + + +class UnsupportedDataType(ZiplineError): + """ + Raised by FFC CustomFactors with unsupported dtypes. + """ + msg = "CustomFactors with dtype {dtype} are not supported." + + +class NoFurtherDataError(ZiplineError): + """ + Raised by calendar operations that would ask for dates beyond the extent of + our known data. + """ + # This accepts an arbitrary message string because it's used in more places + # that can be usefully templated. + msg = '{msg}' diff --git a/zipline/finance/trading.py b/zipline/finance/trading.py index fe822c4e..0df434a5 100644 --- a/zipline/finance/trading.py +++ b/zipline/finance/trading.py @@ -24,7 +24,10 @@ import numpy as np from zipline.data.loader import load_market_data from zipline.utils import tradingcalendar from zipline.assets import AssetFinder -from zipline.errors import UpdateAssetFinderTypeError +from zipline.errors import ( + NoFurtherDataError, + UpdateAssetFinderTypeError, +) log = logbook.Logger('Trading') @@ -69,13 +72,6 @@ log = logbook.Logger('Trading') environment = None -class NoFurtherDataError(Exception): - """ - Thrown when next trading is attempted at the end of available data. - """ - pass - - class TradingEnvironment(object): @classmethod @@ -258,7 +254,9 @@ class TradingEnvironment(object): idx = self.get_index(date) + n if idx < 0 or idx >= len(self.trading_days): - raise NoFurtherDataError('Cannot add %d days to %s' % (n, date)) + raise NoFurtherDataError( + msg='Cannot add %d days to %s' % (n, date) + ) return self.trading_days[idx] @@ -299,8 +297,9 @@ class TradingEnvironment(object): if next_open is None: raise NoFurtherDataError( - "Attempt to backtest beyond available history. \ -Last successful date: %s" % self.last_trading_day) + msg=("Attempt to backtest beyond available history. " + "Last known date: %s" % self.last_trading_day) + ) return self.get_open_and_close(next_open) @@ -313,8 +312,9 @@ Last successful date: %s" % self.last_trading_day) if previous is None: raise NoFurtherDataError( - "Attempt to backtest beyond available history. " - "First successful date: %s" % self.first_trading_day) + msg=("Attempt to backtest beyond available history. " + "First known date: %s" % self.first_trading_day) + ) return self.get_open_and_close(previous) def next_market_minute(self, start): diff --git a/zipline/lib/__init__.py b/zipline/lib/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/zipline/lib/adjusted_array.pyx b/zipline/lib/adjusted_array.pyx new file mode 100644 index 00000000..f7c7c0a4 --- /dev/null +++ b/zipline/lib/adjusted_array.pyx @@ -0,0 +1,225 @@ +""" +Class capable of yielding adjusted chunks of an ndarray. +""" +from cpython cimport ( + Py_EQ, + PyObject_RichCompare, +) +from numpy import ( + asarray, + bool_, + float64, + full, + uint8, +) +from numpy cimport ( + float64_t, + ndarray, + uint8_t, +) + +from zipline.errors import ( + WindowLengthNotPositive, + WindowLengthTooLong, +) + +cdef extern from "math.h" nogil: + float NAN + + +NOMASK = None + + +def ensure_ndarray(ndarray_or_adjusted_array): + """ + Return the input as a numpy ndarray. + + This is a no-op if the input is already an ndarray. If the input is an + adjusted_array, this extracts a read-only view of its internal data buffer. + + Parameters + ---------- + ndarray_or_adjusted_array : numpy.ndarray | zipline.data.adjusted_array + + Returns + ------- + out : The input, converted to an ndarray. + """ + if isinstance(ndarray_or_adjusted_array, ndarray): + return ndarray_or_adjusted_array + elif isinstance(ndarray_or_adjusted_array, AdjustedArray): + return ndarray_or_adjusted_array.data + else: + raise TypeError( + "Can't convert %s to ndarray" % + type(ndarray_or_adjusted_array).__name__ + ) + + +cpdef adjusted_array(ndarray data, ndarray mask, dict adjustments): + """ + Factory function for producing adjusted arrays on inputs of different + dtypes. + + If mask is None, the array is assumed to contain all valid data points. + Otherwise mask should be an array of uint8 of the same shape + as data, containing 0s for valid values and 1s for invalid values. + """ + if data.dtype != float64: + data = data.astype(float64) + if mask is not NOMASK: + if mask.dtype == bool_: + # Cython isn't smart enough to make this coercion even though the + # arrays are bools internally. + mask = mask.view(uint8) + + return Float64AdjustedArray(data, mask, adjustments) + + +cdef _check_window_length(object data, int window_length): + + if window_length < 1: + raise WindowLengthNotPositive(window_length=window_length) + + if window_length > data.shape[0]: + raise WindowLengthTooLong( + nrows=data.shape[0], + window_length=window_length, + ) + + +cdef class AdjustedArray: + + property data: + def __get__(self): + out = asarray(self._data, dtype=self.dtype) + out.setflags(write=False) + return out + + +cdef class Float64AdjustedArray(AdjustedArray): + """ + Adjusted array of float64. + """ + cdef: + readonly float64_t[:, :] _data + dict adjustments + + def __cinit__(self, + float64_t[:, :] data not None, + uint8_t[:, :] mask, # None is equivalent to all 0s. + dict adjustments): + cdef Py_ssize_t row, col + + if mask is not NOMASK: + if not PyObject_RichCompare(mask.shape, data.shape, Py_EQ): + raise ValueError( + "Mask shape %s != data shape %s" % ( + (mask.shape[0], mask.shape[1]), + (data.shape[0], data.shape[1]), + ) + ) + # Fill in NaNs for the mask. + for row in range(mask.shape[0]): + for col in range(mask.shape[1]): + if not mask[row, col]: + data[row, col] = NAN + + self._data = data + self.adjustments = adjustments + + property dtype: + def __get__(self): + return float64 + + cpdef traverse(self, Py_ssize_t window_length, Py_ssize_t offset=0): + return _Float64AdjustedArrayWindow( + self._data.copy(), + self.adjustments, + window_length, + offset, + ) + + +cdef class _Float64AdjustedArrayWindow: + """ + An iterator representing a moving view over an AdjustedArray. + + This object stores a copy of the data from the AdjustedArray over which + it's iterating. At each step in the iteration, it mutates its copy to + allow us to show different data when looking back over the array. + + The arrays yielded by this iterator are always views over the underlying + data. + """ + + cdef float64_t[:, :] data + cdef readonly Py_ssize_t window_length + cdef Py_ssize_t anchor, max_anchor, next_adj + cdef dict adjustments + cdef list adjustment_indices + + def __cinit__(self, + float64_t[:, :] data, + dict adjustments, + Py_ssize_t window_length, + Py_ssize_t offset): + + _check_window_length(data, window_length) + + self.data = data + self.window_length = window_length + + # anchor is the index of the row **after** the row from which we're + # looking back. + self.anchor = window_length + offset + self.max_anchor = data.shape[0] + + self.adjustments = adjustments + self.adjustment_indices = sorted(adjustments, reverse=True) + + if len(self.adjustment_indices) > 0: + self.next_adj = self.adjustment_indices.pop() + else: + self.next_adj = self.max_anchor + + def __iter__(self): + return self + + def __next__(self): + cdef: + ndarray[float64_t, ndim=2] out + object adjustment + Py_ssize_t start, anchor + + anchor = self.anchor + if anchor > self.max_anchor: + raise StopIteration() + + # Apply any adjustments that occured before our current anchor. + # Equivalently, apply any adjustments known **on or before** the date + # for which we're calculating a window. + while self.next_adj < anchor: + + for adjustment in self.adjustments[self.next_adj]: + adjustment.mutate(self.data) + + if len(self.adjustment_indices) > 0: + self.next_adj = self.adjustment_indices.pop() + else: + self.next_adj = self.max_anchor + + start = anchor - self.window_length + out = asarray(self.data[start:self.anchor]) + out.setflags(write=False) + + self.anchor += 1 + return out + + def __repr__(self): + return "%s(window_length=%d, anchor=%d, max_anchor=%d)" % ( + type(self).__name__, + self.window_length, + self.anchor, + self.max_anchor, + ) diff --git a/zipline/lib/adjustment.pyx b/zipline/lib/adjustment.pyx new file mode 100644 index 00000000..776bd6a1 --- /dev/null +++ b/zipline/lib/adjustment.pyx @@ -0,0 +1,228 @@ +from cpython cimport Py_EQ + +from pandas import isnull +from numpy cimport float64_t, uint8_t +# Purely for readability. There aren't C-level declarations for these types. +ctypedef object Int64Index_t +ctypedef object DatetimeIndex_t +ctypedef object Timestamp_t + + +cpdef tuple get_adjustment_locs(DatetimeIndex_t dates_index, + Int64Index_t assets_index, + Timestamp_t start_date, + Timestamp_t end_date, + int asset_id): + """ + Compute indices suitable for passing to an Adjustment constructor. + + If the specified dates aren't in dates_index, we return the index of the + first date **BEFORE** the supplied date. + + Example: + + >>> from pandas import date_range, Int64Index, Timestamp + >>> dates = date_range('2014-01-01', '2014-01-07') + >>> assets = Int64Index(range(10)) + >>> get_adjustment_locs( + ... dates, + ... assets, + ... Timestamp('2014-01-03'), + ... Timestamp('2014-01-05'), + ... 3, + ... ) + (2, 4, 3) + """ + cdef int start_date_loc + + # None or NaT signifies "All values before the end_date". + if isnull(start_date): + start_date_loc = 0 + else: + # Location of earliest date on or after start_date. + 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'), + assets_index.get_loc(asset_id), # Must be exact match. + ) + + +cpdef _from_assets_and_dates(cls, + DatetimeIndex_t dates_index, + Int64Index_t assets_index, + Timestamp_t start_date, + Timestamp_t end_date, + int asset_id, + object value): + """ + Helper for constructing an Adjustment instance from coordinates in + assets/dates indices. + + Example + ------- + + >>> from pandas import date_range, Int64Index, Timestamp + >>> dates = date_range('2014-01-01', '2014-01-07') + >>> assets = Int64Index(range(10)) + >>> Float64Multiply.from_assets_and_dates( + ... dates, + ... assets, + ... Timestamp('2014-01-03'), + ... Timestamp('2014-01-05'), + ... 3, + ... 0.5, + ... ) + Float64Multiply(first_row=2, last_row=4, col=3, value=0.500000) + """ + cdef: + Py_ssize_t first_row, last_row, col + + first_row, last_row, col = get_adjustment_locs( + dates_index, + assets_index, + start_date, + end_date, + asset_id, + ) + return cls(first_row, last_row, col, value) + + +cdef class Float64Adjustment: + """ + Base class for adjustments that operate on Float64 buffers. + """ + cdef: + readonly Py_ssize_t col, first_row, last_row + readonly float64_t value + + def __cinit__(self, + Py_ssize_t first_row, + Py_ssize_t last_row, + Py_ssize_t col, + object value): + assert 0 <= first_row <= last_row + + self.first_row = first_row + self.last_row = last_row + self.col = col + self.value = float(value) + + from_assets_and_dates = classmethod(_from_assets_and_dates) + + def __repr__(self): + return "%s(first_row=%d, last_row=%d, col=%d, value=%f)" % ( + type(self).__name__, + self.first_row, + self.last_row, + self.col, + self.value, + ) + + def __richcmp__(self, object other, int op): + """ + Rich comparison method. Only Equality is defined. + """ + if op != Py_EQ or type(self) != type(other): + return NotImplemented + + return ( + (self.first_row, self.last_row, self.col, self.value) == \ + (other.first_row, other.last_row, other.col, other.value) + ) + +cdef class Float64Multiply(Float64Adjustment): + """ + An adjustment that multiplies by a scalar. + + Example + ------- + + >>> import numpy as np + >>> arr = np.arange(9, dtype=float).reshape(3, 3) + >>> arr + array([[ 0., 1., 2.], + [ 3., 4., 5.], + [ 6., 7., 8.]]) + + >>> adj = Float64Multiply(first_row=1, last_row=2, col=1, value=4.0) + >>> adj.mutate(arr) + >>> arr + array([[ 0., 1., 2.], + [ 3., 16., 5.], + [ 6., 28., 8.]]) + """ + + cpdef mutate(self, float64_t[:, :] data): + cdef Py_ssize_t row, col + col = self.col + + # last_row + 1 because last_row should also be affected. + for row in range(self.first_row, self.last_row + 1): + data[row, col] *= self.value + + +cdef class Float64Overwrite(Float64Adjustment): + """ + An adjustment that overwrites with a scalar. + + Example + ------- + + >>> import numpy as np + >>> arr = np.arange(9, dtype=float).reshape(3, 3) + >>> arr + array([[ 0., 1., 2.], + [ 3., 4., 5.], + [ 6., 7., 8.]]) + + >>> adj = Float64Overwrite(first_row=1, last_row=2, col=1, value=0.0) + >>> adj.mutate(arr) + >>> arr + array([[ 0., 1., 2.], + [ 3., 0., 5.], + [ 6., 0., 8.]]) + """ + + cpdef mutate(self, float64_t[:, :] data): + cdef Py_ssize_t row, col + col = self.col + + # last_row + 1 because last_row should also be affected. + for row in range(self.first_row, self.last_row + 1): + data[row, col] = self.value + + +cdef class Float64Add(Float64Adjustment): + """ + An adjustment that adds a scalar. + + Example + ------- + + >>> import numpy as np + >>> arr = np.arange(9, dtype=float).reshape(3, 3) + >>> arr + array([[ 0., 1., 2.], + [ 3., 4., 5.], + [ 6., 7., 8.]]) + + >>> adj = Float64Add(first_row=1, last_row=2, col=1, value=1.0) + >>> adj.mutate(arr) + >>> arr + array([[ 0., 1., 2.], + [ 3., 5., 5.], + [ 6., 8., 8.]]) + """ + + cpdef mutate(self, float64_t[:, :] data): + cdef Py_ssize_t row, col + col = self.col + + # last_row + 1 because last_row should also be affected. + for row in range(self.first_row, self.last_row + 1): + data[row, col] += self.value diff --git a/zipline/modelling/__init__.py b/zipline/modelling/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/zipline/modelling/classifier.py b/zipline/modelling/classifier.py new file mode 100644 index 00000000..09b58c24 --- /dev/null +++ b/zipline/modelling/classifier.py @@ -0,0 +1,9 @@ +""" +classifier.py +""" + +from zipline.modelling.term import Term + + +class Classifier(Term): + pass diff --git a/zipline/modelling/engine.py b/zipline/modelling/engine.py new file mode 100644 index 00000000..5814f67a --- /dev/null +++ b/zipline/modelling/engine.py @@ -0,0 +1,471 @@ +""" +Compute Engine for FFC API +""" +from abc import ( + ABCMeta, + abstractmethod, +) +from operator import and_ +from six import ( + iteritems, + with_metaclass, +) +from six.moves import ( + reduce, + zip, + zip_longest, +) + +from networkx import ( + DiGraph, + get_node_attributes, + topological_sort, +) +from numpy import ( + add, + empty_like, +) +from pandas import ( + DataFrame, + MultiIndex, +) + +from zipline.lib.adjusted_array import ensure_ndarray +from zipline.errors import NoFurtherDataError +from zipline.modelling.factor import Factor +from zipline.modelling.filter import Filter + + +# TODO: Move this somewhere else. +class CyclicDependency(Exception): + pass + + +def build_dependency_graph(terms): + """ + Build a dependency graph containing the given terms and their dependencies. + + Parameters + ---------- + terms : iterable + An iterable of zipline.modelling.term.Term. + + Returns + ------- + dependencies : networkx.DiGraph + A directed graph representing the dependencies of the desired inputs. + + 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 often needed when a term is an input to a rolling window + computation. For example, if we compute a 30 day moving average of + price from day X to day Y, we need to load price data for the range + from day (X - 29) to day Y. + """ + dependencies = DiGraph() + parents = set() + for term in terms: + _add_to_graph( + term, + dependencies, + parents, + extra_rows=0, + ) + # No parents should be left between top-level terms. + assert not parents + return dependencies + + +def _add_to_graph(term, + dependencies, + parents, + extra_rows): + """ + Add the term and all its inputs to dependencies. + """ + # If we've seen this node already as a parent of the current traversal, + # it means we have an unsatisifiable dependency. This should only be + # possible if the term's inputs are mutated after construction. + if term in parents: + raise CyclicDependency(term) + parents.add(term) + + try: + existing = dependencies.node[term] + except KeyError: + # We're not yet in the graph: add the term with the specified number of + # extra rows. + dependencies.add_node(term, extra_rows=extra_rows) + else: + # We're 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']) + + for subterm in term.inputs: + _add_to_graph( + subterm, + dependencies, + parents, + extra_rows=extra_rows + term.extra_input_rows, + ) + dependencies.add_edge(subterm, term) + + parents.remove(term) + + +class FFCEngine(with_metaclass(ABCMeta)): + + @abstractmethod + def factor_matrix(self, terms, start_date, end_date): + """ + Compute values for `terms` between `start_date` and `end_date`. + + Returns a DataFrame with a MultiIndex of (date, asset) pairs on the + index. On each date, we return a row for each asset that passed all + instances of `Filter` in `terms, and the columns of the returned frame + will be the keys in `terms` whose values are instances of `Factor`. + + 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. + + Returns + ------- + matrix : pd.DataFrame + A matrix of factors + """ + raise NotImplementedError("factor_matrix") + + +class NoOpFFCEngine(FFCEngine): + """ + FFCEngine that doesn't do anything. + """ + + def factor_matrix(self, terms, start, end): + return DataFrame(index=[], columns=sorted(terms.keys())) + + +class SimpleFFCEngine(object): + """ + FFC Engine class that computes each term independently. + + Parameters + ---------- + loader : FFCLoader + A loader to use to retrieve raw data for atomic terms. + calendar : DatetimeIndex + Array of dates to consider as trading days when computing a range + between a fixed start and end. + asset_finder : zipline.assets.AssetFinder + An AssetFinder instance. We depend on the AssetFinder to determine + which assets are in the top-level universe at any point in time. + """ + __slots__ = [ + '_loader', + '_calendar', + '_finder', + '__weakref__', + ] + + def __init__(self, loader, calendar, asset_finder): + self._loader = loader + self._calendar = calendar + self._finder = asset_finder + + def factor_matrix(self, terms, start_date, end_date): + """ + Compute a factor matrix. + + Parameters + ---------- + terms : dict[str -> zipline.modelling.term.Term] + Dict mapping term names to instances. The supplied names are used + as column names in our output frame. + start_date : pd.Timestamp + Start date of the computed matrix. + end_date : pd.Timestamp + End date of the computed matrix. + + The algorithm implemented here can be broken down into the following + stages: + + 0. Build a dependency graph of all terms in `terms`. Topologically + sort the graph to determine an order in which we can compute the terms. + + 1. Ask our AssetFinder for a "lifetimes matrix", which should contain, + for each date between start_date and end_date, a boolean value for each + known asset indicating whether the asset existed on that date. + + 2. Compute each term in the dependency order determined in (0), caching + the results in a a dictionary to that they can be fed into future + terms. + + 3. For each date, determine the number of assets passing **all** + filters. The sum, N, of all these values is the total number of rows in + our output frame, so we pre-allocate an output array of length N for + each factor in `terms`. + + 4. Fill in the arrays allocated in (3) by copying computed values from + our output cache into the corresponding rows. + + 5. Stick the values computed in (4) into a DataFrame and return it. + + Step 0 is performed in `build_dependency_graph`. + Step 1 is performed in `self.build_lifetimes_matrix`. + Step 2 is performed in `self.compute_chunk`. + Steps 3, 4, and 5 are performed in self._format_factor_matrix. + + See Also + -------- + FFCEngine.factor_matrix + """ + graph = build_dependency_graph(terms.values()) + ordered_terms = topological_sort(graph) + extra_row_counts = get_node_attributes(graph, 'extra_rows') + max_extra_rows = max(extra_row_counts.values()) + + lifetimes = self.build_lifetimes_matrix( + start_date, + end_date, + max_extra_rows, + ) + lifetimes_between_dates = lifetimes[max_extra_rows:] + + dates = lifetimes_between_dates.index.values + assets = lifetimes_between_dates.columns.values + + raw_outputs = self.compute_chunk( + ordered_terms, + extra_row_counts, + lifetimes, + ) + + # We only need filters and factors to compute the final output matrix. + raw_filters = [lifetimes_between_dates.values] + raw_factors = [] + factor_names = [] + for name, term in iteritems(terms): + extra = extra_row_counts[term] + if isinstance(term, Factor): + factor_names.append(name) + raw_factors.append(raw_outputs[term][extra:]) + + elif isinstance(term, Filter): + raw_filters.append(raw_outputs[term][extra:]) + + return self._format_factor_matrix( + dates, + assets, + raw_filters, + raw_factors, + factor_names, + ) + + def build_lifetimes_matrix(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. + + Parameters + ---------- + start_date : pd.Timestamp + Base start date for the matrix. + end_date : pd.Timestamp + End date for the matrix. + extra_rows : int + Number of rows prior to `start_date` to include. + Extra rows are needed by terms like moving averages that require a + trailing window of data to compute. + + Returns + ------- + lifetimes : pd.DataFrame + Frame of dtype `bool` containing dates from `extra_rows` days + before `start_date`, continuing through to `end_date`. The + returned frame contains as columns all assets in our AssetFinder + that existed for at least one day between `start_date` and + `end_date`. + """ + calendar = self._calendar + finder = self._finder + start_idx, end_idx = self._calendar.slice_locs(start_date, end_date) + if start_idx < extra_rows: + raise NoFurtherDataError( + msg="Insufficient data to compute FFC Matrix: " + "start date was %s, " + "earliest known date was %s, " + "and %d extra rows were requested." % ( + start_date, calendar[0], extra_rows, + ), + ) + + # Build lifetimes matrix reaching back as far start_date plus + # max_extra_rows. + lifetimes = finder.lifetimes( + calendar[start_idx - extra_rows:end_idx] + ) + assert lifetimes.index[extra_rows] == start_date + assert lifetimes.index[-1] == end_date + + # Filter out columns that didn't exist between the requested start and + # end dates. + existed = lifetimes.iloc[extra_rows:].any() + return lifetimes.loc[:, existed] + + def _inputs_for_term(self, term, workspace, extra_row_counts): + """ + 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. + """ + term_extra_rows = term.extra_input_rows + if term.windowed: + return [ + workspace[input_].traverse( + term.window_length, + offset=extra_row_counts[input_] - term_extra_rows + ) + for input_ in term.inputs + ] + else: + return [ + ensure_ndarray( + workspace[input_][ + extra_row_counts[input_] - term_extra_rows: + ], + ) + for input_ in term.inputs + ] + + def compute_chunk(self, ordered_terms, extra_row_counts, base_mask): + """ + Compute the FFC terms in the graph based on the assets and dates + defined by base_mask. + + Returns a dictionary mapping terms to computed arrays. + """ + loader = self._loader + max_extra_rows = max(extra_row_counts.values()) + workspace = {term: None for term in ordered_terms} + + for term in ordered_terms: + base_mask_for_term = base_mask.iloc[ + max_extra_rows - extra_row_counts[term]: + ] + 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, + ) + for loaded_term, adj_array in zip_longest(to_load, loaded): + workspace[loaded_term] = adj_array + else: + if term.windowed: + compute = term.compute_from_windows + else: + compute = term.compute_from_arrays + workspace[term] = compute( + self._inputs_for_term(term, workspace, extra_row_counts), + base_mask_for_term, + ) + return workspace + + def _format_factor_matrix(self, + dates, + assets, + filter_data, + factor_data, + factor_names): + """ + Convert raw computed filters/factors into a DataFrame for public APIs. + + Parameters + ---------- + dates : np.array[datetime64] + Index for raw data in filter_data/factor_data. + assets : np.array[int64] + Column labels for raw data in filter_data/factor_data. + filter_data : list[ndarray[bool]] + Raw filters data. + factor_data : list[ndarray] + Raw factor data. + factor_names : list[str] + Names of factors to use as keys. + + Returns + ------- + factor_matrix : pd.DataFrame + A DataFrame with the following indices: + + index : two-tiered MultiIndex of (date, asset). For each date, we + return a row for each asset that passed all filters on that + date. + columns : keys from `factor_data` + + Each date/asset/factor triple contains the computed value of the given + factor on the given date for the given asset. + """ + # FUTURE OPTIMIZATION: Cythonize all of this. + + # Boolean mask of values that passed all filters. + unioned = reduce(and_, filter_data) + + # Parallel arrays of (x,y) coords for all date/asset pairs that passed + # all filters. Each entry here will correspond to a row in our output + # frame. + nonzero_xs, nonzero_ys = unioned.nonzero() + + raw_dates_index = empty_like(nonzero_xs, dtype='datetime64[ns]') + raw_assets_index = empty_like(nonzero_xs, dtype=int) + factor_outputs = [ + empty_like(nonzero_xs, dtype=factor.dtype) + for factor in factor_data + ] + + # This is tricky. + + # unioned.sum(axis=1) gives us an array of the same size as `dates` + # containing, for each date, the number of assets that passed our + # filters on that date. + + # Running this through add.accumulate gives us an array containing, for + # each date, the running total of the number of assets that passed our + # filters on or before that date. + + # This means that (bounds[i - 1], bounds[i]) gives us the slice bounds + # of rows in our output DataFrame corresponding to each date. + dt_start = 0 + bounds = add.accumulate(unioned.sum(axis=1)) + for dt_idx, dt_end in enumerate(bounds): + + bounds = slice(dt_start, dt_end) + column_indices = nonzero_ys[bounds] + + raw_dates_index[bounds] = dates[dt_idx] + raw_assets_index[bounds] = assets[column_indices] + for computed, output in zip(factor_data, factor_outputs): + output[bounds] = computed[dt_idx, column_indices] + + # Upper bound of current row becomes lower bound for next row. + dt_start = dt_end + + return DataFrame( + dict(zip(factor_names, factor_outputs)), + index=MultiIndex.from_arrays( + [raw_dates_index, raw_assets_index], + ) + ).tz_localize('UTC', level=0) diff --git a/zipline/modelling/expression.py b/zipline/modelling/expression.py new file mode 100644 index 00000000..7d2078a6 --- /dev/null +++ b/zipline/modelling/expression.py @@ -0,0 +1,307 @@ +""" +NumericalExpression term. +""" +from itertools import chain +import re + +import numexpr +from numexpr.necompiler import getExprNames +from numpy import ( + empty, + find_common_type, +) +from six import integer_types + +from zipline.modelling.term import Term + +_VARIABLE_NAME_RE = re.compile("^(x_)([0-9]+)$") + +# Map from op symbol to equivalent Python magic method name. +_ops_to_methods = { + '+': '__add__', + '-': '__sub__', + '*': '__mul__', + '/': '__div__', + '%': '__mod__', + '**': '__pow__', + '&': '__and__', + '|': '__or__', + '^': '__xor__', + '<': '__lt__', + '<=': '__le__', + '==': '__eq__', + '!=': '__ne__', + '>=': '__ge__', + '>': '__gt__', +} +# Map from op symbol to equivalent Python magic method name after flipping +# arguments. +_ops_to_commuted_methods = { + '+': '__radd__', + '-': '__rsub__', + '*': '__rmul__', + '/': '__rdiv__', + '%': '__rmod__', + '**': '__rpow__', + '&': '__rand__', + '|': '__ror__', + '^': '__rxor__', + '<': '__gt__', + '<=': '__ge__', + '==': '__eq__', + '!=': '__ne__', + '>=': '__le__', + '>': '__lt__', +} +UNARY_OPS = {'-'} +MATH_BINOPS = {'+', '-', '*', '/', '**', '%'} +FILTER_BINOPS = {'&', '|'} # NumExpr doesn't support xor. +COMPARISONS = {'<', '<=', '!=', '>=', '>', '=='} + +NUMERIC_TYPES = (float,) + integer_types +NUMEXPR_MATH_FUNCS = { + 'sin', + 'cos', + 'tan', + 'arcsin', + 'arccos', + 'arctan', + 'sinh', + 'cosh', + 'tanh', + 'arcsinh', + 'arccosh', + 'arctanh', + 'log', + 'log10', + 'log1p', + 'exp', + 'expm1', + 'sqrt', + 'abs', +} + + +def _ensure_element(tup, elem): + """ + Create a tuple containing all elements of tup, plus elem. + + Returns the new tuple and the index of elem in the new tuple. + """ + try: + return tup, tup.index(elem) + except ValueError: + return tuple(chain(tup, (elem,))), len(tup) + + +class BadBinaryOperator(TypeError): + """ + Called when a bad binary operation is encountered. + + Parameters + ---------- + op : str + The attempted operation + left : zipline.computable.Term + The left hand side of the operation. + right : zipline.computable.Term + The right hand side of the operation. + """ + def __init__(self, op, left, right): + super(BadBinaryOperator, self).__init__( + "Can't compute {left} {op} {right}".format( + op=op, + left=type(left).__name__, + right=type(right).__name__, + ) + ) + + +def method_name_for_op(op, commute=False): + """ + Get the name of the Python magic method corresponding to `op`. + + Parameters + ---------- + op : str {'+','-','*', '/','**','&','|','^','<','<=','==','!=','>=','>'} + The requested operation. + commute : bool + Whether to return the name of an equivalent method after flipping args. + + Returns + ------- + method_name : str + The name of the Python magic method corresponding to `op`. + If `commute` is True, returns the name of a method equivalent to `op` + with inputs flipped. + + Examples + -------- + >>> method_name_for_op('+') + '__add__' + >>> method_name_for_op('+', commute=True) + '__radd__' + >>> method_name_for_op('>') + '__gt__' + >>> method_name_for_op('>', commute=True) + '__lt__' + """ + if commute: + return _ops_to_commuted_methods[op] + return _ops_to_methods[op] + + +def is_comparison(op): + return op in COMPARISONS + + +class NumericalExpression(Term): + """ + Term binding to a numexpr expression. + + Parameters + ---------- + expr : string + A string suitable for passing to numexpr. All variables in 'expr' + should be of the form "x_i", where i is the index of the corresponding + factor input in 'binds'. + binds : tuple + A tuple of factors to use as inputs. + """ + window_length = 0 + + def __new__(cls, expr, binds): + + # If our class doesn't have an explicit dtype set, infer one from the + # inputs. + + # FIXME: This doesn't take into account dtypes of constants, so it will + # break if we have something like + # 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: + dtype = cls.dtype + else: + dtype = find_common_type( + [factor.dtype for factor in binds], + [], + ) + return super(NumericalExpression, cls).__new__( + cls, + inputs=binds, + expr=expr, + dtype=dtype, + ) + + def _init(self, expr, *args, **kwargs): + self._expr = expr + return super(NumericalExpression, self)._init(*args, **kwargs) + + @classmethod + def static_identity(cls, expr, *args, **kwargs): + return ( + super(NumericalExpression, cls).static_identity(*args, **kwargs), + expr, + ) + + def _validate(self): + """ + Ensure that our expression string has variables of the form x_0, x_1, + ... x_(N - 1), where N is the length of our inputs. + """ + variable_names, _unused = getExprNames(self._expr, {}) + expr_indices = [] + for name in variable_names: + match = _VARIABLE_NAME_RE.match(name) + if not match: + raise ValueError("%r is not a valid variable name" % name) + expr_indices.append(int(match.group(2))) + + expr_indices.sort() + expected_indices = list(range(len(self.inputs))) + if expr_indices != expected_indices: + raise ValueError( + "Expected %s for variable indices, but got %s" % ( + expected_indices, expr_indices, + ) + ) + return super(NumericalExpression, self)._validate() + + def compute_from_arrays(self, arrays, mask): + """ + Compute our stored expression string with numexpr. + """ + out = empty(mask.shape, dtype=self.dtype) + # This writes directly into our output buffer. + numexpr.evaluate( + self._expr, + local_dict={ + "x_%d" % idx: array + for idx, array in enumerate(arrays) + }, + global_dict={}, + out=out, + ) + return out + + def _rebind_variables(self, new_inputs): + """ + Return self._expr with all variables rebound to the indices implied by + new_inputs. + """ + expr = self._expr + for idx, input_ in enumerate(self.inputs): + old_varname = "x_%d" % idx + # Temporarily rebind to x_temp_N so that we don't overwrite the + # same value multiple times. + temp_new_varname = "x_temp_%d" % new_inputs.index(input_) + expr = expr.replace(old_varname, temp_new_varname) + # Clear out the temp variables now that we've finished iteration. + return expr.replace("_temp_", "_") + + def _merge_expressions(self, other): + """ + Merge the inputs of two NumericalExpressions into a single input tuple, + rewriting their respective string expressions to make input names + resolve correctly. + + Returns a tuple of (new_self_expr, new_other_expr, new_inputs) + """ + new_inputs = tuple(set(self.inputs).union(other.inputs)) + new_self_expr = self._rebind_variables(new_inputs) + new_other_expr = other._rebind_variables(new_inputs) + return new_self_expr, new_other_expr, new_inputs + + def build_binary_op(self, op, other): + """ + Compute new expression strings and a new inputs tuple for combining + self and other with a binary operator. + """ + if isinstance(other, NumericalExpression): + self_expr, other_expr, new_inputs = self._merge_expressions(other) + elif isinstance(other, Term): + self_expr = self._expr + new_inputs, other_idx = _ensure_element(self.inputs, other) + other_expr = "x_%d" % other_idx + elif isinstance(other, NUMERIC_TYPES): + self_expr = self._expr + other_expr = str(other) + new_inputs = self.inputs + else: + raise BadBinaryOperator(op, other) + return self_expr, other_expr, new_inputs + + @property + def bindings(self): + return { + "x_%d" % i: input_ + for i, input_ in enumerate(self.inputs) + } + + def __repr__(self): + return "{typename}(expr='{expr}', bindings={bindings})".format( + typename=type(self).__name__, + expr=self._expr, + bindings=self.bindings, + ) diff --git a/zipline/modelling/factor/__init__.py b/zipline/modelling/factor/__init__.py new file mode 100644 index 00000000..4d55df50 --- /dev/null +++ b/zipline/modelling/factor/__init__.py @@ -0,0 +1,11 @@ +from .factor import ( + Factor, + TestingFactor, + CustomFactor, +) + +__all__ = [ + 'Factor', + 'TestingFactor', + 'CustomFactor', +] diff --git a/zipline/modelling/factor/factor.py b/zipline/modelling/factor/factor.py new file mode 100644 index 00000000..a6679f8f --- /dev/null +++ b/zipline/modelling/factor/factor.py @@ -0,0 +1,413 @@ +""" +factor.py +""" +from operator import attrgetter +from numpy import ( + apply_along_axis, + float64, + nan, +) +from scipy.stats import rankdata + +from zipline.errors import ( + UnknownRankMethod, + UnsupportedDataType, +) +from zipline.modelling.term import ( + CustomTermMixin, + RequiredWindowLengthMixin, + SingleInputMixin, + Term, + TestingTermMixin, +) +from zipline.modelling.expression import ( + BadBinaryOperator, + COMPARISONS, + is_comparison, + MATH_BINOPS, + method_name_for_op, + NUMERIC_TYPES, + NumericalExpression, + NUMEXPR_MATH_FUNCS, + UNARY_OPS, +) +from zipline.modelling.filter import ( + NumExprFilter, + PercentileFilter, +) +from zipline.utils.control_flow import nullctx + + +_RANK_METHODS = frozenset(['average', 'min', 'max', 'dense', 'ordinal']) + + +def binop_return_type(op): + if is_comparison(op): + return NumExprFilter + else: + return NumExprFactor + + +def binary_operator(op): + """ + Factory function for making binary operator methods on a Factor subclass. + + Returns a function, "binary_operator" suitable for implementing functions + like __add__. + """ + # When combining a Factor with a NumericalExpression, we use this + # attrgetter instance to defer to the commuted implementation of the + # NumericalExpression operator. + commuted_method_getter = attrgetter(method_name_for_op(op, commute=True)) + + def binary_operator(self, other): + # This can't be hoisted up a scope because the types returned by + # binop_return_type aren't defined when the top-level function is + # invoked in the class body of Factor. + return_type = binop_return_type(op) + if isinstance(self, NumExprFactor): + self_expr, other_expr, new_inputs = self.build_binary_op( + op, other, + ) + return return_type( + "({left}) {op} ({right})".format( + left=self_expr, + op=op, + right=other_expr, + ), + new_inputs, + ) + elif isinstance(other, NumExprFactor): + # NumericalExpression overrides ops to correctly handle merging of + # inputs. Look up and call the appropriate reflected operator with + # ourself as the input. + return commuted_method_getter(other)(self) + elif isinstance(other, Factor): + if self is other: + return return_type( + "x_0 {op} x_0".format(op=op), + (self,), + ) + return return_type( + "x_0 {op} x_1".format(op=op), + (self, other), + ) + elif isinstance(other, NUMERIC_TYPES): + return return_type( + "x_0 {op} ({constant})".format(op=op, constant=other), + binds=(self,), + ) + raise BadBinaryOperator(op, self, other) + + return binary_operator + + +def reflected_binary_operator(op): + """ + Factory function for making binary operator methods on a Factor. + + Returns a function, "reflected_binary_operator" suitable for implementing + functions like __radd__. + """ + assert not is_comparison(op) + + def reflected_binary_operator(self, other): + + if isinstance(self, NumericalExpression): + self_expr, other_expr, new_inputs = self.build_binary_op( + op, other + ) + return NumExprFactor( + "({left}) {op} ({right})".format( + left=other_expr, + right=self_expr, + op=op, + ), + new_inputs, + ) + + # Only have to handle the numeric case because in all other valid cases + # the corresponding left-binding method will be called. + elif isinstance(other, NUMERIC_TYPES): + return NumExprFactor( + "{constant} {op} x_0".format(op=op, constant=other), + binds=(self,), + ) + raise BadBinaryOperator(op, other, self) + return reflected_binary_operator + + +def unary_operator(op): + """ + Factory function for making unary operator methods for Factors. + """ + # Only negate is currently supported for all our possible input types. + valid_ops = {'-'} + if op not in valid_ops: + raise ValueError("Invalid unary operator %s." % op) + + def unary_operator(self): + # This can't be hoisted up a scope because the types returned by + # unary_op_return_type aren't defined when the top-level function is + # invoked. + if isinstance(self, NumericalExpression): + return NumExprFactor( + "{op}({expr})".format(op=op, expr=self._expr), + self.inputs, + ) + else: + return NumExprFactor("{op}x_0".format(op=op), (self,)) + return unary_operator + + +def function_application(func): + """ + Factory function for producing function application methods for Factor + subclasses. + """ + if func not in NUMEXPR_MATH_FUNCS: + raise ValueError("Unsupported mathematical function '%s'" % func) + + def mathfunc(self): + if isinstance(self, NumericalExpression): + return NumExprFactor( + "{func}({expr})".format(func=func, expr=self._expr), + self.inputs, + ) + else: + return NumExprFactor("{func}(x_0)".format(func=func), (self,)) + return mathfunc + + +class Factor(Term): + """ + A transformation yielding a timeseries of scalar values associated with an + Asset. + """ + # Dynamically add functions for creating NumExprFactor/NumExprFilter + # instances. + clsdict = locals() + clsdict.update( + { + method_name_for_op(op): binary_operator(op) + # Don't override __eq__ because it breaks comparisons on tuples of + # Factors. + for op in MATH_BINOPS.union(COMPARISONS - {'=='}) + } + ) + clsdict.update( + { + method_name_for_op(op, commute=True): reflected_binary_operator(op) + for op in MATH_BINOPS + } + ) + clsdict.update( + { + '__neg__': unary_operator(op) + for op in UNARY_OPS + } + ) + clsdict.update( + { + funcname: function_application(funcname) + for funcname in NUMEXPR_MATH_FUNCS + } + ) + + __truediv__ = clsdict['__div__'] + __rtruediv__ = clsdict['__rdiv__'] + + eq = binary_operator('==') + + def rank(self, method='ordinal'): + """ + Construct a new Factor representing the sorted rank of each column + within each row. + + Returns + ------- + ranks : zipline.modelling.factor.Rank + A new factor that will compute the sorted indices of the data + produced by `self`. + method : str, {'ordinal', 'min', 'max', 'dense', 'average'} + The method used to assign ranks to tied elements. Default is + 'ordinal'. See `scipy.stats.rankdata` for a full description of + the semantics for each ranking method. + + The default is 'ordinal'. + + Notes + ----- + The default value for `method` is different from the default for + `scipy.stats.rankdata`. See that function's documentation for a full + description of the valid inputs to `method`. + + Missing or non-existent data on a given day will cause an asset to be + given a rank of NaN for that day. + + See Also + -------- + scipy.stats.rankdata : Underlying ranking algorithm. + zipline.modelling.factor.Rank : Class implementing core functionality. + """ + return Rank(self, method=method) + + def percentile_between(self, min_percentile, max_percentile): + """ + Construct a new Filter representing entries from the output of this + Factor that fall within the percentile range defined by min_percentile + and max_percentile. + + Parameters + ---------- + min_percentile : float [0.0, 100.0] + max_percentile : float [0.0, 100.0] + + Returns + ------- + out : zipline.modelling.filter.PercentileFilter + A new filter that will compute the specified percentile-range mask. + + See Also + -------- + zipline.modelling.filter.PercentileFilter + """ + return PercentileFilter( + self, + min_percentile=min_percentile, + max_percentile=max_percentile, + ) + + +class NumExprFactor(NumericalExpression, Factor): + """ + Factor computed from a numexpr expression. + + Parameters + ---------- + expr : string + A string suitable for passing to numexpr. All variables in 'expr' + should be of the form "x_i", where i is the index of the corresponding + factor input in 'binds'. + binds : tuple + A tuple of factors to use as inputs. + + Notes + ----- + NumExprFactors are constructed by numerical operators like `+` and `-`. + Users should rarely need to construct a NumExprFactor directly. + """ + pass + + +class Rank(SingleInputMixin, Factor): + """ + A Factor representing the row-wise rank data of another Factor. + + Parameters + ---------- + factor : zipline.modelling.factor.Factor + The factor on which to compute ranks. + method : str, {'average', 'min', 'max', 'dense', 'ordinal'} + The method used to assign ranks to tied elements. See + `scipy.stats.rankdata` for a full description of the semantics for each + ranking method. + + See Also + -------- + scipy.stats.rankdata : Underlying ranking algorithm. + zipline.factor.Factor.rank : Method-style interface to same functionality. + + Notes + ----- + Most users should call Factor.rank rather than directly construct an + instance of this class. + """ + dtype = float64 + window_length = 0 + domain = None + + def __new__(cls, factor, method): + return super(Rank, cls).__new__( + cls, + inputs=(factor,), + method=method, + ) + + def _init(self, method, *args, **kwargs): + self._method = method + return super(Rank, self)._init(*args, **kwargs) + + @classmethod + def static_identity(cls, method, *args, **kwargs): + return ( + super(Rank, cls).static_identity(*args, **kwargs), + method, + ) + + def _validate(self): + """ + Verify that the stored rank method is valid. + """ + if self._method not in _RANK_METHODS: + raise UnknownRankMethod( + method=self._method, + choices=set(_RANK_METHODS), + ) + return super(Rank, self)._validate() + + def compute_from_arrays(self, arrays, mask): + """ + For each row in the input, compute a like-shaped array of per-row + ranks. + """ + # FUTURE OPTIMIZATION: + # Write a less general `apply_to_rows` method in + # Cython that doesn't do all the extra work that apply_over_axis does. + + # FUTURE OPTIMIZATION: + # Look at bottleneck.nanrankdata, which is ~30% faster than numpy here, + # and does what we want with NaNs, but doesn't support `method`. + 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 + return result + + def __repr__(self): + return "{type}({input_}, method='{method}')".format( + type=type(self).__name__, + input_=self.inputs[0], + method=self._method, + ) + + +class CustomFactor(RequiredWindowLengthMixin, CustomTermMixin, Factor): + """ + Base class for user-defined Factors operating on windows of raw data. + + TODO: This is basically the most important class to document in the whole + FFC API... + + We currently only support CustomFactors of type float64. + """ + dtype = float64 + ctx = nullctx() + + def _validate(self): + if self.dtype != float64: + raise UnsupportedDataType(self.dtype) + return super(CustomFactor, self)._validate() + + +class TestingFactor(TestingTermMixin, Factor): + """ + Base class for testing engines that asserts all inputs are correctly + shaped. + """ + pass diff --git a/zipline/modelling/factor/technical.py b/zipline/modelling/factor/technical.py new file mode 100644 index 00000000..93f656a0 --- /dev/null +++ b/zipline/modelling/factor/technical.py @@ -0,0 +1,91 @@ +""" +Technical Analysis Factors +-------------------------- +""" +from bottleneck import ( + nanargmax, + nanmax, + nanmean, + nansum, +) +from numpy import ( + clip, + diff, + fmax, + inf, + isnan, + NINF, +) +from numexpr import evaluate + +from zipline.data.equities import USEquityPricing +from zipline.modelling.term import SingleInputMixin +from zipline.utils.control_flow import ignore_nanwarnings +from .factor import CustomFactor + + +class RSI(CustomFactor, SingleInputMixin): + """ + Factor computing rolling relative-strength index on a DataSet. + + Default Input: USEquityPricing.close + Default Window Length: 14 + """ + window_length = 14 + inputs = (USEquityPricing.close,) + + def compute(self, today, assets, out, closes): + diffs = diff(closes) + ups = nanmean(clip(diffs, 0, inf), axis=0) + downs = nanmean(clip(diffs, -inf, 0), axis=0) + return evaluate( + "100 - (100 / (1 + (ups / downs)))", + locals_dict={'ups': ups, 'downs': downs}, + globals_dict={}, + out=out, + ) + + +class SimpleMovingAverage(CustomFactor, SingleInputMixin): + """ + Factor computing moving averages on a DataSet. + """ + # numpy's nan functions throw warnings when passed an array containing only + # nans, but they still returns the desired value (nan), so we ignore the + # warning. + ctx = ignore_nanwarnings() + + def compute(self, today, assets, out, data): + out[:] = nanmean(data, axis=0) + + +class WeightedAverageValue(CustomFactor): + """ + Helper for VWAP-like computations. + """ + def compute(self, today, assets, out, base, weight): + out[:] = nansum(base * weight, axis=0) / nansum(weight, axis=0) + + +class VWAP(WeightedAverageValue): + """ + Volume-weighted average price + """ + inputs = (USEquityPricing.close, USEquityPricing.volume) + + +class MaxDrawdown(CustomFactor, SingleInputMixin): + """ + Max Drawdown over a window + """ + ctx = ignore_nanwarnings() + + def compute(self, today, assets, out, data): + drawdowns = fmax.accumulate(data, axis=0) - data + drawdowns[isnan(drawdowns)] = NINF + drawdown_ends = nanargmax(drawdowns, axis=0) + + # TODO: Accelerate this loop in Cython or Numba. + for i, end in enumerate(drawdown_ends): + peak = nanmax(data[:end + 1, i]) + out[i] = (peak - data[end, i]) / data[end, i] diff --git a/zipline/modelling/filter.py b/zipline/modelling/filter.py new file mode 100644 index 00000000..d6ede154 --- /dev/null +++ b/zipline/modelling/filter.py @@ -0,0 +1,282 @@ +""" +filter.py +""" +from numpy import ( + bool_, + float64, + nan, + nanpercentile, +) +from itertools import chain +from operator import attrgetter + +from zipline.errors import ( + BadPercentileBounds, +) +from zipline.modelling.term import ( + SingleInputMixin, + Term, + TestingTermMixin, +) +from zipline.modelling.expression import ( + BadBinaryOperator, + FILTER_BINOPS, + method_name_for_op, + NumericalExpression, +) + + +def concat_tuples(*tuples): + """ + Concatenate a sequence of tuples into one tuple. + """ + return tuple(chain(*tuples)) + + +def binary_operator(op): + """ + Factory function for making binary operator methods on a Filter subclass. + + Returns a function "binary_operator" suitable for implementing functions + like __and__ or __or__. + """ + # When combining a Filter with a NumericalExpression, we use this + # attrgetter instance to defer to the commuted interpretation of the + # NumericalExpression operator. + commuted_method_getter = attrgetter(method_name_for_op(op, commute=True)) + + def binary_operator(self, other): + if isinstance(self, NumericalExpression): + self_expr, other_expr, new_inputs = self.build_binary_op( + op, other, + ) + return NumExprFilter( + "({left}) {op} ({right})".format( + left=self_expr, + op=op, + right=other_expr, + ), + new_inputs, + ) + elif isinstance(other, NumericalExpression): + # NumericalExpression overrides numerical ops to correctly handle + # merging of inputs. Look up and call the appropriate + # right-binding operator with ourself as the input. + return commuted_method_getter(other)(self) + elif isinstance(other, Filter): + if self is other: + return NumExprFilter( + "x_0 {op} x_0".format(op=op), + (self,), + ) + return NumExprFilter( + "x_0 {op} x_1".format(op=op), + (self, other), + ) + elif isinstance(other, int): # Note that this is true for bool as well + return NumExprFilter( + "x_0 {op} ({constant})".format(op=op, constant=int(other)), + binds=(self,), + ) + raise BadBinaryOperator(op, self, other) + return binary_operator + + +class Filter(Term): + """ + A boolean predicate on a universe of Assets. + """ + domain = None + dtype = bool_ + + clsdict = locals() + clsdict.update( + { + method_name_for_op(op): binary_operator(op) + for op in FILTER_BINOPS + } + ) + + 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_from_arrays(self, arrays, mask): + """ + Compute our result with numexpr, then apply `mask`. + """ + numexpr_result = super(NumExprFilter, self).compute_from_arrays( + arrays, + mask, + ) + return numexpr_result & mask + + +class PercentileFilter(SingleInputMixin, Filter): + """ + A Filter representing assets falling between percentile bounds of a Factor. + + Parameters + ---------- + factor : zipline.modelling.factor.Factor + The factor over which to compute percentile bounds. + min_percentile : float [0.0, 1.0] + The minimum percentile rank of an asset that will pass the filter. + max_percentile : float [0.0, 1.0] + The maxiumum percentile rank of an asset that will pass the filter. + """ + window_length = 0 + + def __new__(cls, factor, min_percentile, max_percentile): + return super(PercentileFilter, cls).__new__( + cls, + inputs=(factor,), + min_percentile=min_percentile, + max_percentile=max_percentile, + ) + + def _init(self, min_percentile, max_percentile, *args, **kwargs): + self._min_percentile = min_percentile + self._max_percentile = max_percentile + return super(PercentileFilter, self)._init(*args, **kwargs) + + @classmethod + def static_identity(cls, min_percentile, max_percentile, *args, **kwargs): + return ( + super(PercentileFilter, cls).static_identity(*args, **kwargs), + min_percentile, + max_percentile, + ) + + def _validate(self): + """ + Ensure that our percentile bounds are well-formed. + """ + if not 0.0 <= self._min_percentile < self._max_percentile <= 100.0: + raise BadPercentileBounds( + min_percentile=self._min_percentile, + max_percentile=self._max_percentile, + ) + return super(PercentileFilter, self)._validate() + + def compute_from_arrays(self, arrays, 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 + + # FIXME: np.nanpercentile **should** support computing multiple bounds + # at once, but there's a bug in the logic for multiple bounds in numpy + # 1.9.2. It will be fixed in 1.10. + # c.f. https://github.com/numpy/numpy/pull/5981 + lower_bounds = nanpercentile( + data, + self._min_percentile, + axis=1, + keepdims=True, + ) + upper_bounds = nanpercentile( + data, + self._max_percentile, + axis=1, + 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_from_arrays(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_from_arrays( + then_inputs, + mask & first_result, + ) + + +class TestingFilter(TestingTermMixin, Filter): + """ + Base class for testing engines that asserts all inputs are correctly + shaped. + """ + pass diff --git a/zipline/modelling/term.py b/zipline/modelling/term.py new file mode 100644 index 00000000..c680075b --- /dev/null +++ b/zipline/modelling/term.py @@ -0,0 +1,294 @@ +""" +Base class for Filters, Factors and Classifiers +""" +from numpy import ( + empty, + float64, + full, + nan, +) +from weakref import WeakValueDictionary + +from zipline.errors import ( + InputTermNotAtomic, + TermInputsNotSpecified, + WindowLengthNotPositive, + WindowLengthNotSpecified, +) +from zipline.utils.lazyval import lazyval + + +NotSpecified = (object(),) + + +class Term(object): + """ + Base class for terms in an FFC API compute graph. + """ + inputs = NotSpecified + window_length = NotSpecified + domain = None + dtype = float64 + + _term_cache = WeakValueDictionary() + + def __new__(cls, + inputs=None, + window_length=None, + domain=None, + dtype=None, + *args, + **kwargs): + """ + Memoized constructor for Terms. + + Caching previously-constructed Terms is useful because it allows us to + only compute equivalent sub-expressions once when traversing an FFC + dependency graph. + + Caching previously-constructed Terms is **sane** because terms and + their inputs are both conceptually immutable. + """ + if inputs is None: + inputs = tuple(cls.inputs) + else: + inputs = tuple(inputs) + + if window_length is None: + window_length = cls.window_length + + if domain is None: + domain = cls.domain + + if dtype is None: + dtype = cls.dtype + + identity = cls.static_identity( + inputs=inputs, + window_length=window_length, + domain=domain, + dtype=dtype, + *args, **kwargs + ) + + try: + return cls._term_cache[identity] + except KeyError: + new_instance = cls._term_cache[identity] = \ + super(Term, cls).__new__(cls)._init( + inputs=inputs, + window_length=window_length, + domain=domain, + dtype=dtype, + *args, **kwargs + ) + return new_instance + + def __init__(self, *args, **kwargs): + """ + Noop constructor to play nicely with our caching __new__. Subclasses + should implement _init instead of this method. + + When a class' __new__ returns an instance of that class, Python will + automatically call __init__ on the object, even if a new object wasn't + actually constructed. Because we memoize instances, we often return an + object that was already initialized from __new__, in which case we + don't want to call __init__ again. + + Subclasses that need to initialize new instances should override _init, + which is guaranteed to be called only once. + """ + pass + + def _init(self, inputs, window_length, domain, dtype): + self.inputs = inputs + self.window_length = window_length + self.domain = domain + self.dtype = dtype + + self._validate() + return self + + @classmethod + def static_identity(cls, inputs, window_length, domain, dtype): + """ + Return the identity of the Term that would be constructed from the + given arguments. + + Identities that compare equal will cause us to return a cached instance + rather than constructing a new one. We do this primarily because it + makes dependency resolution easier. + + 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) + + def _validate(self): + """ + Assert that this term is well-formed. This should be called exactly + once, at the end of Term._init(). + """ + if self.inputs is NotSpecified: + raise TermInputsNotSpecified(termname=type(self).__name__) + if self.window_length is NotSpecified: + raise WindowLengthNotSpecified(termname=type(self).__name__) + + if self.window_length: + for child in self.inputs: + if not child.atomic: + raise InputTermNotAtomic(parent=self, child=child) + + @lazyval + def atomic(self): + """ + Whether or not this term has dependencies. + + If term.atomic is truthy, it should have dataset and dtype attributes. + """ + return len(self.inputs) == 0 + + @lazyval + def windowed(self): + """ + Whether or not this term represents a trailing window computation. + + If term.windowed is truthy, its compute_from_windows method will be + called with instances of AdjustedArray as inputs. + + If term.windowed is falsey, its compute_from_baseline will be called + with instances of np.ndarray as inputs. + """ + return ( + self.window_length is not NotSpecified + and self.window_length > 0 + ) + + @lazyval + def extra_input_rows(self): + """ + The number of extra rows needed for each of our inputs to compute this + term. + """ + return max(0, self.window_length - 1) + + def compute_from_windows(self, windows, mask): + """ + Subclasses should implement this for computations requiring moving + windows of continually-adjusting data. + """ + raise NotImplementedError() + + def compute_from_arrays(self, arrays, mask): + """ + Subclasses should implement this for computations that can be expressed + directly as array computations. + """ + raise NotImplementedError() + + def __repr__(self): + return ( + "{type}({inputs}, window_length={window_length})" + ).format( + type=type(self).__name__, + inputs=self.inputs, + window_length=self.window_length, + ) + + +# TODO: Move mixins to a separate file? +class SingleInputMixin(object): + + def _validate(self): + num_inputs = len(self.inputs) + if num_inputs != 1: + raise ValueError( + "{typename} expects only one input, " + "but received {num_inputs} instead.".format( + typename=type(self).__name__, + num_inputs=num_inputs + ) + ) + return super(SingleInputMixin, self)._validate() + + +class RequiredWindowLengthMixin(object): + def _validate(self): + if self.windowed: + return super(RequiredWindowLengthMixin, self)._validate() + if self.window_length is NotSpecified: + raise WindowLengthNotSpecified() + raise WindowLengthNotPositive(window_length=self.window_length) + + +class CustomTermMixin(object): + """ + Mixin for user-defined rolling-window Terms. + + Implements `compute_from_windows` in terms of a user-defined `compute` + function, which is mapped over the input windows. + + Used by CustomFactor, CustomFilter, CustomClassifier, etc. + """ + + def compute(self, today, assets, out, *arrays): + """ + Override this method with a function that writes a value into `out`. + """ + raise NotImplementedError() + + def compute_from_windows(self, windows, 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 + # time-step? + for idx, date in enumerate(dates): + compute( + date, + assets, + out[idx], + *(next(w) for w in windows) + ) + out[~mask.values] = nan + return out + + +class TestingTermMixin(object): + """ + Mixin for Term subclasses testing engines that asserts all inputs are + correctly shaped. + + Used by TestingTerm, TestingFilter, TestingClassifier, etc. + """ + def compute_from_windows(self, windows, mask): + assert self.window_length > 0 + dates, assets = mask.index, mask.columns + outbuf = empty(mask.shape, dtype=self.dtype) + for idx, _ in enumerate(dates): + result = self.from_windows(*(next(w) for w in windows)) + assert result.shape == (len(assets),) + outbuf[idx] = result + + for window in windows: + try: + next(window) + except StopIteration: + pass + else: + raise AssertionError("window %s was not exhausted" % window) + return outbuf + + def compute_from_arrays(self, arrays, mask): + assert self.window_length == 0 + outbuf = empty(mask.shape, dtype=self.dtype) + for array in arrays: + assert array.shape == outbuf.shape + outbuf[:] = self.from_arrays(*arrays) + return outbuf diff --git a/zipline/protocol.py b/zipline/protocol.py index f5d3f999..83dada7b 100644 --- a/zipline/protocol.py +++ b/zipline/protocol.py @@ -17,6 +17,7 @@ from copy import copy from six import iteritems, iterkeys import pandas as pd +from pandas.tseries.tools import normalize_date import numpy as np from . utils.protocol_utils import Enum @@ -494,6 +495,17 @@ class BarData(object): def __init__(self, data=None): self._data = data or {} self._contains_override = None + self._factor_matrix = None + self._factor_matrix_expires = pd.Timestamp(0, tz='UTC') + + @property + def factors(self): + algo = get_algo_instance() + today = normalize_date(algo.get_datetime()) + if today > self._factor_matrix_expires: + self._factor_matrix, self._factor_matrix_expires = \ + algo.compute_factor_matrix(today) + return self._factor_matrix.loc[today] def __contains__(self, name): if self._contains_override: diff --git a/zipline/utils/api_support.py b/zipline/utils/api_support.py index bdf9cf87..0def3881 100644 --- a/zipline/utils/api_support.py +++ b/zipline/utils/api_support.py @@ -54,3 +54,25 @@ def api_method(f): zipline.api.__all__.append(f.__name__) f.is_api_method = True return f + + +def require_not_initialized(exception): + """ + Decorator for API methods that should only be called during or before + TradingAlgorithm.initialize. `exception` will be raised if the method is + called after initialize. + + Usage + ----- + @required_not_initialized(SomeException, "Don't do that!") + def method(self): + # Do stuff that should only be allowed during initialize. + """ + def decorator(method): + @wraps(method) + def wrapped_method(self, *args, **kwargs): + if self.initialized: + raise exception + return method(self, *args, **kwargs) + return wrapped_method + return decorator diff --git a/zipline/utils/control_flow.py b/zipline/utils/control_flow.py new file mode 100644 index 00000000..f45f3c0e --- /dev/null +++ b/zipline/utils/control_flow.py @@ -0,0 +1,56 @@ +""" +Control flow utilities. +""" +from warnings import ( + catch_warnings, + filterwarnings, +) + + +class nullctx(object): + """ + Null context manager. Useful for conditionally adding a contextmanager in + a single line, e.g.: + + with SomeContextManager() if some_expr else nullctx(): + do_stuff() + """ + def __enter__(self): + return self + + def __exit__(*args): + return False + + +class WarningContext(object): + """ + Re-entrant contextmanager for contextually managing warnings. + """ + def __init__(self, *warning_specs): + self._warning_specs = warning_specs + self._catchers = [] + + def __enter__(self): + catcher = catch_warnings() + catcher.__enter__() + self._catchers.append(catcher) + for args, kwargs in self._warning_specs: + filterwarnings(*args, **kwargs) + return catcher + + def __exit__(self, *exc_info): + catcher = self._catchers.pop() + return catcher.__exit__(*exc_info) + + +def ignore_nanwarnings(): + """ + Helper for building a WarningContext that ignores warnings from numpy's + nanfunctions. + """ + return WarningContext( + ( + ('ignore',), + {'category': RuntimeWarning, 'module': 'numpy.lib.nanfunctions'}, + ) + ) diff --git a/zipline/utils/lazyval.py b/zipline/utils/lazyval.py new file mode 100644 index 00000000..8db2f2e3 --- /dev/null +++ b/zipline/utils/lazyval.py @@ -0,0 +1,46 @@ +""" +An immutable, lazily loaded value descriptor. +""" + + +from weakref import WeakKeyDictionary + + +class lazyval(object): + """ + Decorator that marks that an attribute should not be computed until + needed, and that the value should be memoized. + + Example + ------- + + >>> from zipline.utils.lazyval import lazyval + >>> class C(object): + ... def __init__(self): + ... self.count = 0 + ... @lazyval + ... def val(self): + ... self.count += 1 + ... return "val" + ... + >>> c = C() + >>> c.count + 0 + >>> c.val, c.count + ('val', 1) + >>> c.val, c.count + ('val', 1) + """ + def __init__(self, get): + self._get = get + self._cache = WeakKeyDictionary() + + def __get__(self, instance, owner): + if instance is None: + return self + + try: + return self._cache[instance] + except KeyError: + self._cache[instance] = val = self._get(instance) + return val diff --git a/zipline/utils/math_utils.py b/zipline/utils/math_utils.py index e1b19bee..be2b66b4 100644 --- a/zipline/utils/math_utils.py +++ b/zipline/utils/math_utils.py @@ -26,9 +26,17 @@ try: nanmean = bn.nanmean nanstd = bn.nanstd nansum = bn.nansum + nanmax = bn.nanmax + nanmin = bn.nanmin + nanargmax = bn.nanargmax + nanargmin = bn.nanargmin except ImportError: # slower numpy import numpy as np nanmean = np.nanmean nanstd = np.nanstd nansum = np.nansum + nanmax = np.nanmax + nanmin = np.nanmin + nanargmax = np.nanargmax + nanargmin = np.nanargmin diff --git a/zipline/utils/test_utils.py b/zipline/utils/test_utils.py index fb217edf..96dc7d7b 100644 --- a/zipline/utils/test_utils.py +++ b/zipline/utils/test_utils.py @@ -1,21 +1,48 @@ from contextlib import contextmanager +from itertools import ( + product, +) from logbook import FileHandler from mock import patch +import operator from zipline.finance.blotter import ORDER_STATUS from zipline.utils import security_list -from six import itervalues +from six import ( + itervalues, +) +from six.moves import filter import os import pandas as pd import shutil import tempfile +EPOCH = pd.Timestamp(0, tz='UTC') + + +def seconds_to_timestamp(seconds): + return pd.Timestamp(seconds, unit='s', tz='UTC') + def to_utc(time_str): + """Convert a string in US/Eastern time to UTC""" return pd.Timestamp(time_str, tz='US/Eastern').tz_convert('UTC') +def str_to_seconds(s): + """ + Convert a pandas-intelligible string to (integer) seconds since UTC. + + >>> from pandas import Timestamp + >>> (Timestamp('2014-01-01') - Timestamp(0)).total_seconds() + 1388534400.0 + >>> str_to_seconds('2014-01-01') + 1388534400 + """ + return int((pd.Timestamp(s, tz='UTC') - EPOCH).total_seconds()) + + def setup_logger(test, path='test.log'): test.log_handler = FileHandler(path) test.log_handler.push_application() @@ -111,18 +138,6 @@ class ExceptionSource(object): 5 / 0 -@contextmanager -def nullctx(): - """ - Null context manager. Useful for conditionally adding a contextmanager in - a single line, e.g.: - - with SomeContextManager() if some_expr else nullctx(): - do_stuff() - """ - yield - - @contextmanager def security_list_copy(): old_dir = security_list.SECURITY_LISTS_DIR @@ -159,3 +174,137 @@ def add_security_data(adds, deletes): for sym in adds: f.write(sym) f.write('\n') + + +def all_pairs_matching_predicate(values, pred): + """ + Return an iterator of all pairs, (v0, v1) from values such that + + `pred(v0, v1) == True` + + Parameters + ---------- + values : iterable + pred : function + + Returns + ------- + pairs_iterator : generator + Generator yielding pairs matching `pred`. + + Examples + -------- + >>> from zipline.utils.test_utils import all_pairs_matching_predicate + >>> from operator import eq, lt + >>> list(all_pairs_matching_predicate(range(5), eq)) + [(0, 0), (1, 1), (2, 2), (3, 3), (4, 4)] + >>> list(all_pairs_matching_predicate("abcd", lt)) + [('a', 'b'), ('a', 'c'), ('a', 'd'), ('b', 'c'), ('b', 'd'), ('c', 'd')] + """ + return filter(lambda pair: pred(*pair), product(values, repeat=2)) + + +def product_upper_triangle(values, include_diagonal=False): + """ + Return an iterator over pairs, (v0, v1), drawn from values. + + If `include_diagonal` is True, returns all pairs such that v0 <= v1. + If `include_diagonal` is False, returns all pairs such that v0 < v1. + """ + return all_pairs_matching_predicate( + values, + operator.le if include_diagonal else operator.lt, + ) + + +def all_subindices(index): + """ + Return all valid sub-indices of a pandas Index. + """ + return ( + index[start:stop] + for start, stop in product_upper_triangle(range(len(index) + 1)) + ) + + +def make_rotating_asset_info(num_assets, + first_start, + frequency, + periods_between_starts, + asset_lifetime): + """ + Create a DataFrame representing lifetimes of assets that are constantly + rotating in and out of existence. + + Parameters + ---------- + num_assets : int + How many assets to create. + first_start : pd.Timestamp + The start date for the first asset. + frequency : str or pd.tseries.offsets.Offset (e.g. trading_day) + Frequency used to interpret next two arguments. + periods_between_starts : int + Create a new asset every `frequency` * `periods_between_new` + asset_lifetime : int + Each asset exists for `frequency` * `asset_lifetime` days. + + Returns + ------- + info : pd.DataFrame + DataFrame representing newly-created assets. + """ + return pd.DataFrame( + { + 'sid': range(num_assets), + 'symbol': [chr(ord('A') + i) for i in range(num_assets)], + 'asset_type': ['equity'] * num_assets, + # Start a new asset every `periods_between_starts` days. + 'start_date': pd.date_range( + first_start, + freq=(periods_between_starts * frequency), + periods=num_assets, + ), + # Each asset lasts for `asset_lifetime` days. + 'end_date': pd.date_range( + first_start + (asset_lifetime * frequency), + freq=(periods_between_starts * frequency), + periods=num_assets, + ), + 'exchange': 'TEST', + } + ) + + +def make_simple_asset_info(assets, start_date, end_date, symbols=None): + """ + Create a DataFrame representing assets that exist for the full duration + between `start_date` and `end_date`. + + Parameters + ---------- + assets : array-like + start_date : pd.DatetimeIndex + end_date : pd.DatetimeIndex + symbols : list, optional + Symbols to use for the assets. + If not provided, symbols are generated from upper-case letters. + + Returns + ------- + info : pd.DataFrame + DataFrame representing newly-created assets. + """ + num_assets = len(assets) + if symbols is None: + symbols = [chr(ord('A') + i) for i in range(num_assets)] + return pd.DataFrame( + { + 'sid': assets, + 'symbol': symbols, + 'asset_type': ['equity'] * num_assets, + 'start_date': [start_date] * num_assets, + 'end_date': [end_date] * num_assets, + 'exchange': 'TEST', + } + )