diff --git a/tests/history_cases.py b/tests/history_cases.py new file mode 100644 index 00000000..2ba6e0a2 --- /dev/null +++ b/tests/history_cases.py @@ -0,0 +1,578 @@ +""" +Test case definitions for history tests. +""" + +import pandas as pd +import numpy as np + +from zipline.finance.trading import TradingEnvironment +from zipline.history.history import HistorySpec +from zipline.protocol import BarData + + +def to_utc(time_str): + return pd.Timestamp(time_str, tz='US/Eastern').tz_convert('UTC') + + +def mixed_frequency_expected_index(count, frequency): + """ + Helper for enumerating expected indices for test_mixed_frequency. + """ + env = TradingEnvironment.instance() + minute = MIXED_FREQUENCY_MINUTES[count] + + if frequency == '1d': + return [env.previous_open_and_close(minute)[1], minute] + elif frequency == '1m': + return [env.previous_market_minute(minute), minute] + + +def mixed_frequency_expected_data(count, frequency): + """ + Helper for enumerating expected data test_mixed_frequency. + """ + if frequency == '1d': + if count < 390: + return [np.nan, count] + else: + return [389, count] + elif frequency == '1m': + if count == 0: + return [np.nan, count] + else: + return [count - 1, count] + + +MIXED_FREQUENCY_MINUTES = TradingEnvironment.instance().market_minute_window( + to_utc('2013-06-28 9:31AM'), 780, +) +DAILY_OPEN_CLOSE_SPECS = [ + HistorySpec(3, '1d', 'open_price', False), + HistorySpec(3, '1d', 'close_price', False), +] +ILLIQUID_PRICES_SPECS = [ + HistorySpec(3, '1m', 'price', False), + HistorySpec(5, '1m', 'price', True), +] +MIXED_FREQUENCY_SPECS = [ + HistorySpec(2, '1m', 'price', False), + HistorySpec(2, '1d', 'price', False), +] +MIXED_FIELDS_SPECS = [ + HistorySpec(3, '1m', 'price', True), + HistorySpec(3, '1m', 'open_price', True), + HistorySpec(3, '1m', 'close_price', True), + HistorySpec(3, '1m', 'high', True), + HistorySpec(3, '1m', 'low', True), + HistorySpec(3, '1m', 'volume', True), +] + + +HISTORY_CONTAINER_TEST_CASES = { + # June 2013 + # Su Mo Tu We Th Fr Sa + # 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 daily open close': { + # A list of HistorySpec objects. + 'specs': DAILY_OPEN_CLOSE_SPECS, + + # Sids for the test. + 'sids': [1], + + # Start date for test. + 'dt': to_utc('2013-06-21 9:31AM'), + + # Sequence of updates to the container + 'updates': [ + BarData( + { + 1: { + 'open_price': 10, + 'close_price': 11, + 'dt': to_utc('2013-06-21 10:00AM'), + }, + }, + ), + BarData( + { + 1: { + 'open_price': 12, + 'close_price': 13, + 'dt': to_utc('2013-06-21 3:30PM'), + }, + }, + ), + BarData( + { + 1: { + 'open_price': 14, + 'close_price': 15, + # Wait a full market day before the next bar. + # We should end up with nans for Monday the 24th. + 'dt': to_utc('2013-06-25 9:31AM'), + }, + }, + ), + ], + + # Dictionary mapping spec_key -> list of expected outputs + 'expected': { + # open + DAILY_OPEN_CLOSE_SPECS[0].key_str: [ + pd.DataFrame( + data={ + 1: [np.nan, np.nan, 10] + }, + index=[ + to_utc('2013-06-19 4:00PM'), + to_utc('2013-06-20 4:00PM'), + to_utc('2013-06-21 10:00AM'), + ], + ), + + pd.DataFrame( + data={ + 1: [np.nan, np.nan, 10] + }, + index=[ + to_utc('2013-06-19 4:00PM'), + to_utc('2013-06-20 4:00PM'), + to_utc('2013-06-21 3:30PM'), + ], + ), + + pd.DataFrame( + data={ + 1: [10, np.nan, 14] + }, + index=[ + to_utc('2013-06-21 4:00PM'), + to_utc('2013-06-24 4:00PM'), + to_utc('2013-06-25 9:31AM'), + ], + ), + ], + # close + DAILY_OPEN_CLOSE_SPECS[1].key_str: [ + pd.DataFrame( + data={ + 1: [np.nan, np.nan, 11] + }, + index=[ + to_utc('2013-06-19 4:00PM'), + to_utc('2013-06-20 4:00PM'), + to_utc('2013-06-21 10:00AM'), + ], + ), + + pd.DataFrame( + data={ + 1: [np.nan, np.nan, 13] + }, + index=[ + to_utc('2013-06-19 4:00PM'), + to_utc('2013-06-20 4:00PM'), + to_utc('2013-06-21 3:30PM'), + ], + ), + + pd.DataFrame( + data={ + 1: [13, np.nan, 15] + }, + index=[ + to_utc('2013-06-21 4:00PM'), + to_utc('2013-06-24 4:00PM'), + to_utc('2013-06-25 9:31AM'), + ], + ), + ], + }, + }, + + 'test illiquid prices': { + + # A list of HistorySpec objects. + 'specs': ILLIQUID_PRICES_SPECS, + + # Sids for the test. + 'sids': [1], + + # Start date for test. + 'dt': to_utc('2013-06-28 9:31AM'), + + # Sequence of updates to the container + 'updates': [ + BarData( + { + 1: { + 'price': 10, + 'dt': to_utc('2013-06-28 9:31AM'), + }, + }, + ), + BarData( + { + 1: { + 'price': 11, + 'dt': to_utc('2013-06-28 9:32AM'), + }, + }, + ), + BarData( + { + 1: { + 'price': 12, + 'dt': to_utc('2013-06-28 9:33AM'), + }, + }, + ), + BarData( + { + 1: { + 'price': 13, + # Note: Skipping 9:34 to simulate illiquid bar/missing + # data. + 'dt': to_utc('2013-06-28 9:35AM'), + }, + }, + ), + ], + + # Dictionary mapping spec_key -> list of expected outputs + 'expected': { + ILLIQUID_PRICES_SPECS[0].key_str: [ + pd.DataFrame( + data={ + 1: [np.nan, np.nan, 10], + }, + index=[ + to_utc('2013-06-27 3:59PM'), + to_utc('2013-06-27 4:00PM'), + to_utc('2013-06-28 9:31AM'), + ], + ), + + pd.DataFrame( + data={ + 1: [np.nan, 10, 11], + }, + index=[ + to_utc('2013-06-27 4:00PM'), + to_utc('2013-06-28 9:31AM'), + to_utc('2013-06-28 9:32AM'), + ], + ), + + pd.DataFrame( + data={ + 1: [10, 11, 12], + }, + index=[ + to_utc('2013-06-28 9:31AM'), + to_utc('2013-06-28 9:32AM'), + to_utc('2013-06-28 9:33AM'), + ], + ), + + # Since there's no update for 9:34, this is called at 9:35. + pd.DataFrame( + data={ + 1: [12, np.nan, 13], + }, + index=[ + to_utc('2013-06-28 9:33AM'), + to_utc('2013-06-28 9:34AM'), + to_utc('2013-06-28 9:35AM'), + ], + ), + ], + + ILLIQUID_PRICES_SPECS[1].key_str: [ + pd.DataFrame( + data={ + 1: [np.nan, np.nan, np.nan, np.nan, 10], + }, + index=[ + to_utc('2013-06-27 3:57PM'), + to_utc('2013-06-27 3:58PM'), + to_utc('2013-06-27 3:59PM'), + to_utc('2013-06-27 4:00PM'), + to_utc('2013-06-28 9:31AM'), + ], + ), + + pd.DataFrame( + data={ + 1: [np.nan, np.nan, np.nan, 10, 11], + }, + index=[ + to_utc('2013-06-27 3:58PM'), + to_utc('2013-06-27 3:59PM'), + to_utc('2013-06-27 4:00PM'), + to_utc('2013-06-28 9:31AM'), + to_utc('2013-06-28 9:32AM'), + ], + ), + + pd.DataFrame( + data={ + 1: [np.nan, np.nan, 10, 11, 12], + }, + index=[ + to_utc('2013-06-27 3:59PM'), + to_utc('2013-06-27 4:00PM'), + to_utc('2013-06-28 9:31AM'), + to_utc('2013-06-28 9:32AM'), + to_utc('2013-06-28 9:33AM'), + ], + ), + + # Since there's no update for 9:34, this is called at 9:35. + # The 12 value from 9:33 should be forward-filled. + pd.DataFrame( + data={ + 1: [10, 11, 12, 12, 13], + }, + index=[ + to_utc('2013-06-28 9:31AM'), + to_utc('2013-06-28 9:32AM'), + to_utc('2013-06-28 9:33AM'), + to_utc('2013-06-28 9:34AM'), + to_utc('2013-06-28 9:35AM'), + ], + ), + ], + }, + }, + + 'test mixed frequencies': { + + # A list of HistorySpec objects. + 'specs': MIXED_FREQUENCY_SPECS, + + # Sids for the test. + 'sids': [1], + + # Start date for test. + 'dt': to_utc('2013-06-28 9:31AM'), + + # Sequence of updates to the container + 'updates': [ + BarData( + { + 1: { + 'price': count, + 'dt': dt, + } + } + ) + for count, dt in enumerate(MIXED_FREQUENCY_MINUTES) + ], + + # Dictionary mapping spec_key -> list of expected outputs. + 'expected': { + MIXED_FREQUENCY_SPECS[0].key_str: [ + pd.DataFrame( + data={ + 1: mixed_frequency_expected_data(count, '1m'), + }, + index=mixed_frequency_expected_index(count, '1m'), + ) + for count in range(len(MIXED_FREQUENCY_MINUTES)) + ], + + MIXED_FREQUENCY_SPECS[1].key_str: [ + pd.DataFrame( + data={ + 1: mixed_frequency_expected_data(count, '1d'), + }, + index=mixed_frequency_expected_index(count, '1d'), + ) + for count in range(len(MIXED_FREQUENCY_MINUTES)) + ] + }, + }, + + 'test multiple fields and sids': { + + # A list of HistorySpec objects. + 'specs': MIXED_FIELDS_SPECS, + + # Sids for the test. + 'sids': [1, 10], + + # Start date for test. + 'dt': to_utc('2013-06-28 9:31AM'), + + # Sequence of updates to the container + 'updates': [ + BarData( + { + 1: { + 'dt': dt, + 'price': count, + 'open_price': count, + 'close_price': count, + 'high': count, + 'low': count, + 'volume': count, + }, + 10: { + 'dt': dt, + 'price': count * 10, + 'open_price': count * 10, + 'close_price': count * 10, + 'high': count * 10, + 'low': count * 10, + 'volume': count * 10, + }, + }, + ) + for count, dt in enumerate([ + to_utc('2013-06-28 9:31AM'), + to_utc('2013-06-28 9:32AM'), + to_utc('2013-06-28 9:33AM'), + # NOTE: No update for 9:34 + to_utc('2013-06-28 9:35AM'), + ]) + ], + + # Dictionary mapping spec_key -> list of expected outputs + 'expected': dict( + + # Build a dict from a list of tuples. Doing it this way because + # there are two distinct cases we want to test: forward-fillable + # fields and non-forward-fillable fields. + [ + ( + # Non forward-fill fields + key, + [ + pd.DataFrame( + data={ + 1: [np.nan, np.nan, 0], + 10: [np.nan, np.nan, 0], + }, + index=[ + to_utc('2013-06-27 3:59PM'), + to_utc('2013-06-27 4:00PM'), + to_utc('2013-06-28 9:31AM'), + ], + + # Missing volume data should manifest as 0's rather + # than nans. + ).fillna(0 if 'volume' in key else np.nan), + pd.DataFrame( + data={ + 1: [np.nan, 0, 1], + 10: [np.nan, 0, 10], + }, + index=[ + to_utc('2013-06-27 4:00PM'), + to_utc('2013-06-28 9:31AM'), + to_utc('2013-06-28 9:32AM'), + ], + ).fillna(0 if 'volume' in key else np.nan), + + pd.DataFrame( + data={ + 1: [0, 1, 2], + 10: [0, 10, 20], + }, + index=[ + to_utc('2013-06-28 9:31AM'), + to_utc('2013-06-28 9:32AM'), + to_utc('2013-06-28 9:33AM'), + ], + + # Note: Calling fillna() here even though there are + # no NaNs because this makes it less likely + # for us to introduce a stupid bug by + # copy/pasting in the future. + ).fillna(0 if 'volume' in key else np.nan), + pd.DataFrame( + data={ + 1: [2, np.nan, 3], + 10: [20, np.nan, 30], + }, + index=[ + to_utc('2013-06-28 9:33AM'), + to_utc('2013-06-28 9:34AM'), + to_utc('2013-06-28 9:35AM'), + ], + ).fillna(0 if 'volume' in key else np.nan), + ], + ) + for key in [spec.key_str for spec in MIXED_FIELDS_SPECS + if spec.field not in HistorySpec.FORWARD_FILLABLE] + ] + + + # Concatenate the expected results for non-ffillable with + # expected result for ffillable. + [ + ( + # Forward-fillable fields + key, + [ + pd.DataFrame( + data={ + 1: [np.nan, np.nan, 0], + 10: [np.nan, np.nan, 0], + }, + index=[ + to_utc('2013-06-27 3:59PM'), + to_utc('2013-06-27 4:00PM'), + to_utc('2013-06-28 9:31AM'), + ], + ), + + pd.DataFrame( + data={ + 1: [np.nan, 0, 1], + 10: [np.nan, 0, 10], + }, + index=[ + to_utc('2013-06-27 4:00PM'), + to_utc('2013-06-28 9:31AM'), + to_utc('2013-06-28 9:32AM'), + ], + ), + + pd.DataFrame( + data={ + 1: [0, 1, 2], + 10: [0, 10, 20], + }, + index=[ + to_utc('2013-06-28 9:31AM'), + to_utc('2013-06-28 9:32AM'), + to_utc('2013-06-28 9:33AM'), + ], + ), + + pd.DataFrame( + data={ + 1: [2, 2, 3], + 10: [20, 20, 30], + }, + index=[ + to_utc('2013-06-28 9:33AM'), + to_utc('2013-06-28 9:34AM'), + to_utc('2013-06-28 9:35AM'), + ], + ), + ], + ) + for key in [spec.key_str for spec in MIXED_FIELDS_SPECS + if spec.field in HistorySpec.FORWARD_FILLABLE] + ] + ), + }, +} diff --git a/tests/test_history.py b/tests/test_history.py index 3ca7e10b..dc8fc7e1 100644 --- a/tests/test_history.py +++ b/tests/test_history.py @@ -24,10 +24,14 @@ from zipline.history.history_container import HistoryContainer from zipline.protocol import BarData import zipline.utils.factory as factory from zipline import TradingAlgorithm -from zipline.finance.trading import SimulationParameters +from zipline.finance.trading import SimulationParameters, TradingEnvironment from zipline.sources import RandomWalkSource +from .history_cases import ( + HISTORY_CONTAINER_TEST_CASES, +) + # Cases are over the July 4th holiday, to ensure use of trading calendar. # March 2013 @@ -73,7 +77,7 @@ from zipline.sources import RandomWalkSource # Times to be converted via: # pd.Timestamp('2013-07-05 9:31', tz='US/Eastern').tz_convert('UTC')}, -MINUTE_CASES_RAW = { +INDEX_TEST_CASES_RAW = { 'week of daily data': { 'input': {'bar_count': 5, 'frequency': '1d', @@ -86,6 +90,18 @@ MINUTE_CASES_RAW = { '2013-07-05 9:31AM', ] }, + 'five minutes on july 5th open': { + 'input': {'bar_count': 5, + 'frequency': '1m', + 'algo_dt': '2013-07-05 9:31AM'}, + 'expected': [ + '2013-07-03 12:57PM', + '2013-07-03 12:58PM', + '2013-07-03 12:59PM', + '2013-07-03 1:00PM', + '2013-07-05 9:31AM', + ] + }, } @@ -104,28 +120,31 @@ def convert_cases(cases): in case['expected']]) return cases -MINUTE_CASES = convert_cases(MINUTE_CASES_RAW) +INDEX_TEST_CASES = convert_cases(INDEX_TEST_CASES_RAW) -def index_at_dt(case_input): +def get_index_at_dt(case_input): history_spec = history.HistorySpec( case_input['bar_count'], case_input['frequency'], None, False ) - return history.index_at_dt(history_spec, - case_input['algo_dt']) + return history.index_at_dt(history_spec, case_input['algo_dt']) class TestHistoryIndex(TestCase): + @classmethod + def setUpClass(cls): + cls.environment = TradingEnvironment.instance() + @parameterized.expand( [(name, case['input'], case['expected']) - for name, case in MINUTE_CASES.items()] + for name, case in INDEX_TEST_CASES.items()] ) def test_index_at_dt(self, name, case_input, expected): - history_index = index_at_dt(case_input) + history_index = get_index_at_dt(case_input) history_series = pd.Series(index=history_index) expected_series = pd.Series(index=expected) @@ -135,9 +154,64 @@ class TestHistoryIndex(TestCase): class TestHistoryContainer(TestCase): + @classmethod + def setUpClass(cls): + cls.env = TradingEnvironment.instance() + + def bar_data_dt(self, bar_data, require_unique=True): + """ + Get a dt to associate with the given BarData object. + + If require_unique == True, throw an error if multiple unique dt's are + encountered. Otherwise, return the earliest dt encountered. + """ + dts = {sid_data['dt'] for sid_data in bar_data.values()} + if require_unique and len(dts) > 1: + self.fail("Multiple unique dts ({0}) in {1}".format(dts, bar_data)) + + return sorted(dts)[0] + + @parameterized.expand( + [(name, + case['specs'], + case['sids'], + case['dt'], + case['updates'], + case['expected']) + for name, case in HISTORY_CONTAINER_TEST_CASES.items()] + ) + def test_history_container(self, + name, + specs, + sids, + dt, + updates, + expected): + + for spec in specs: + # Sanity check on test input. + self.assertEqual(len(expected[spec.key_str]), len(updates)) + + container = HistoryContainer( + {spec.key_str: spec for spec in specs}, sids, dt + ) + + for update_count, update in enumerate(updates): + + bar_dt = self.bar_data_dt(update) + container.update(update, bar_dt) + + for spec in specs: + pd.util.testing.assert_frame_equal( + container.get_history(spec, bar_dt), + expected[spec.key_str][update_count], + check_dtype=False, + check_column_type=True, + check_index_type=True, + check_frame_type=True, + ) + def test_container_nans_and_daily_roll(self): - # set up trading environment - factory.create_simulation_parameters(num_days=4) spec = history.HistorySpec( bar_count=3, diff --git a/zipline/history/history.py b/zipline/history/history.py index 2adbd103..e5b147b3 100644 --- a/zipline/history/history.py +++ b/zipline/history/history.py @@ -38,7 +38,7 @@ class Frequency(object): the elapsed minutes of the (incomplete) current day - "1m" provides data at minute frequency """ - SUPPORTED_FREQUENCIES = frozenset({'1d'}) + SUPPORTED_FREQUENCIES = frozenset({'1d', '1m'}) MAX_MINUTES = {'m': 1, 'd': 390} def __init__(self, freq_str): diff --git a/zipline/protocol.py b/zipline/protocol.py index e3983de8..7d78119d 100644 --- a/zipline/protocol.py +++ b/zipline/protocol.py @@ -157,8 +157,8 @@ class BarData(object): usage of what this replaced as a dictionary subclass. """ - def __init__(self): - self._data = {} + def __init__(self, data=None): + self._data = data or {} self._contains_override = None def __contains__(self, name): @@ -217,3 +217,6 @@ class BarData(object): def __len__(self): return len(self.keys()) + + def __repr__(self): + return '{0}({1})'.format(self.__class__.__name__, self._data)