diff --git a/tests/calendars/test_trading_calendar.py b/tests/calendars/test_trading_calendar.py index 761c82a1..0ede45c3 100644 --- a/tests/calendars/test_trading_calendar.py +++ b/tests/calendars/test_trading_calendar.py @@ -34,6 +34,7 @@ from zipline.errors import ( InvalidCalendarName, ) +from zipline.testing.predicates import assert_equal from zipline.utils.calendars import( register_calendar, deregister_calendar, @@ -683,6 +684,22 @@ class ExchangeCalendarTestBase(object): self.assertEqual(open_answer, found_open) self.assertEqual(close_answer, found_close) + def test_session_opens_in_range(self): + found_opens = self.calendar.session_opens_in_range( + self.answers.index[0], + self.answers.index[-1], + ) + + assert_equal(found_opens, self.answers['market_open']) + + def test_session_closes_in_range(self): + found_closes = self.calendar.session_closes_in_range( + self.answers.index[0], + self.answers.index[-1], + ) + + assert_equal(found_closes, self.answers['market_close']) + def test_daylight_savings(self): # 2004 daylight savings switches: # Sunday 2004-04-04 and Sunday 2004-10-31 diff --git a/zipline/data/resample.py b/zipline/data/resample.py index b5632493..4676a18f 100644 --- a/zipline/data/resample.py +++ b/zipline/data/resample.py @@ -514,28 +514,49 @@ class MinuteResampleSessionBarReader(SessionBarReader): self._calendar = calendar self._minute_bar_reader = minute_bar_reader - def _get_resampled(self, columns, start_dt, end_dt, assets): + def _get_resampled(self, columns, start_session, end_session, assets): + range_open = self._calendar.session_open(start_session) + range_close = self._calendar.session_close(end_session) + minute_data = self._minute_bar_reader.load_raw_arrays( - columns, start_dt, end_dt, assets) - dts = self._calendar.minutes_in_range(start_dt, end_dt).values - sessions = self._calendar.sessions_in_range(start_dt, end_dt) - m_closes = np.zeros(len(sessions), dtype=np.dtype('datetime64[ns]')) - for i, s in enumerate(sessions): - close = self._calendar.session_close(s) - m_closes[i] = close.value - m_locs = np.searchsorted(dts, m_closes) + columns, + range_open, + range_close, + assets, + ) + + # Get the index of the close minute for each session in the range. + # If the range contains only one session, the only close in the range + # is the last minute in the data. Otherwise, we need to get all the + # session closes and find their indices in the range of minutes. + if start_session == end_session: + close_ilocs = np.array([len(minute_data[0]) - 1], dtype=np.int64) + else: + minutes = self._calendar.minutes_in_range( + range_open, + range_close, + ) + session_closes = self._calendar.session_closes_in_range( + start_session, + end_session, + ) + close_ilocs = minutes.searchsorted(session_closes.values) + results = [] - shape = (len(sessions), len(assets)) + shape = (len(close_ilocs), len(assets)) + for col in columns: if col != 'volume': out = np.full(shape, np.nan) else: out = np.zeros(shape, dtype=np.uint32) results.append(out) + for i in range(len(assets)): for j, column in enumerate(columns): data = minute_data[j][:, i] - minute_to_session(column, m_locs, data, results[j][:, i]) + minute_to_session(column, close_ilocs, data, results[j][:, i]) + return results @property @@ -543,19 +564,14 @@ class MinuteResampleSessionBarReader(SessionBarReader): return self._calendar def load_raw_arrays(self, columns, start_dt, end_dt, sids): - range_open, _ = self._calendar.open_and_close_for_session( - start_dt) - _, range_close = self._calendar.open_and_close_for_session( - end_dt) - return self._get_resampled(columns, range_open, range_close, sids) + return self._get_resampled(columns, start_dt, end_dt, sids) def get_value(self, sid, session, colname): # WARNING: This will need caching or other optimization if used in a # tight loop. # This was developed to complete interface, but has not been tuned # for real world use. - start, end = self._calendar.open_and_close_for_session(session) - return self._get_resampled([colname], start, end, [sid])[0][0][0] + return self._get_resampled([colname], session, session, [sid])[0][0][0] @lazyval def sessions(self): diff --git a/zipline/utils/calendars/trading_calendar.py b/zipline/utils/calendars/trading_calendar.py index 6e63c124..9a09c566 100644 --- a/zipline/utils/calendars/trading_calendar.py +++ b/zipline/utils/calendars/trading_calendar.py @@ -660,6 +660,18 @@ class TradingCalendar(with_metaclass(ABCMeta)): 'market_close' ].tz_localize('UTC') + def session_opens_in_range(self, start_session_label, end_session_label): + return self.schedule.loc[ + start_session_label:end_session_label, + 'market_open', + ].dt.tz_localize('UTC') + + def session_closes_in_range(self, start_session_label, end_session_label): + return self.schedule.loc[ + start_session_label:end_session_label, + 'market_close', + ].dt.tz_localize('UTC') + @property def all_sessions(self): return self.schedule.index