From acf345e1d389e4fff3710e7d44d1d06a6241c3f8 Mon Sep 17 00:00:00 2001 From: Andrew Daniels Date: Tue, 11 Apr 2017 16:35:04 -0400 Subject: [PATCH] PERF: Optimize session close lookups in resample bar reader (#1749) Optimize session close lookups in MinuteResampleSessionBarReader: - Adds `session_closes_in_range` method (along with `session_opens_in_range`) to TradingCalendar to allow vectorized retrieval of all values in a range of sessions. - Improves code path for resampling a single session's worth of data (as is the case when calling `get_value`), since we don't actually need to look up the close minute. --- tests/calendars/test_trading_calendar.py | 17 +++++++ zipline/data/resample.py | 52 ++++++++++++++------- zipline/utils/calendars/trading_calendar.py | 12 +++++ 3 files changed, 63 insertions(+), 18 deletions(-) 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