diff --git a/zipline/pipeline/loaders/earnings.py b/zipline/pipeline/loaders/earnings.py index b5792a1d..5972345d 100644 --- a/zipline/pipeline/loaders/earnings.py +++ b/zipline/pipeline/loaders/earnings.py @@ -3,16 +3,14 @@ Reference implementation for EarningsCalendar loaders. """ from itertools import repeat -from numpy import full_like, full import pandas as pd from six import iteritems -from six.moves import zip from toolz import merge from .base import PipelineLoader from .frame import DataFrameLoader +from .utils import next_date_frame, previous_date_frame from ..data.earnings import EarningsCalendar -from zipline.utils.numpy_utils import np_NaT from zipline.utils.memoize import lazyval @@ -83,7 +81,7 @@ class EarningsCalendarLoader(PipelineLoader): def next_announcement_loader(self): return DataFrameLoader( self.dataset.next_announcement, - next_earnings_date_frame( + next_date_frame( self.all_dates, self.announcement_dates, ), @@ -94,7 +92,7 @@ class EarningsCalendarLoader(PipelineLoader): def previous_announcement_loader(self): return DataFrameLoader( self.dataset.previous_announcement, - previous_earnings_date_frame( + previous_date_frame( self.all_dates, self.announcement_dates, ), @@ -110,83 +108,3 @@ class EarningsCalendarLoader(PipelineLoader): ) -def next_earnings_date_frame(dates, announcement_dates): - """ - Make a DataFrame representing simulated next earnings dates. - - Parameters - ---------- - dates : pd.DatetimeIndex. - The index of the returned DataFrame. - announcement_dates : dict[int -> pd.Series] - Dict mapping sids to an index of dates on which earnings were announced - for that sid. - - Returns - ------- - next_earnings: pd.DataFrame - A DataFrame representing, for each (label, date) pair, the first entry - in `earnings_calendars[label]` on or after `date`. Entries falling - after the last date in a calendar will have `np_NaT` as the result in - the output. - - See Also - -------- - previous_earnings_date_frame - """ - cols = {equity: full_like(dates, np_NaT) for equity in announcement_dates} - raw_dates = dates.values - for equity, earnings_dates in iteritems(announcement_dates): - data = cols[equity] - if not earnings_dates.index.is_monotonic_increasing: - earnings_dates = earnings_dates.sort_index() - - # Iterate over the raw Series values, since we're comparing against - # numpy arrays anyway. - iterkv = zip(earnings_dates.index.values, earnings_dates.values) - for timestamp, announce_date in iterkv: - date_mask = (timestamp <= raw_dates) & (raw_dates <= announce_date) - value_mask = (announce_date <= data) | (data == np_NaT) - data[date_mask & value_mask] = announce_date - - return pd.DataFrame(index=dates, data=cols) - - -def previous_earnings_date_frame(dates, announcement_dates): - """ - Make a DataFrame representing simulated next earnings dates. - - Parameters - ---------- - dates : DatetimeIndex. - The index of the returned DataFrame. - announcement_dates : dict[int -> DatetimeIndex] - Dict mapping sids to an index of dates on which earnings were announced - for that sid. - - Returns - ------- - prev_earnings: pd.DataFrame - A DataFrame representing, for (label, date) pair, the first entry in - `announcement_dates[label]` strictly before `date`. Entries falling - before the first date in a calendar will have `NaT` as the result in - the output. - - See Also - -------- - next_earnings_date_frame - """ - sids = list(announcement_dates) - out = full((len(dates), len(sids)), np_NaT, dtype='datetime64[ns]') - dn = dates[-1].asm8 - for col_idx, sid in enumerate(sids): - # announcement_dates[sid] is Series mapping knowledge_date to actual - # announcement date. We don't care about the knowledge date for - # computing previous earnings. - values = announcement_dates[sid].values - values = values[values <= dn] - out[dates.searchsorted(values), col_idx] = values - - frame = pd.DataFrame(out, index=dates, columns=sids) - frame.ffill(inplace=True) - return frame diff --git a/zipline/pipeline/loaders/utils.py b/zipline/pipeline/loaders/utils.py new file mode 100644 index 00000000..cea08789 --- /dev/null +++ b/zipline/pipeline/loaders/utils.py @@ -0,0 +1,90 @@ +import numpy as np +import pandas as pd +from six import iteritems +from six.moves import zip + +from zipline.utils.numpy_utils import np_NaT + + +def next_date_frame(dates, announcement_dates): + """ + Make a DataFrame representing simulated next earnings dates. + + Parameters + ---------- + dates : pd.DatetimeIndex. + The index of the returned DataFrame. + announcement_dates : dict[int -> pd.Series] + Dict mapping sids to an index of dates on which earnings were announced + for that sid. + + Returns + ------- + next_earnings: pd.DataFrame + A DataFrame representing, for each (label, date) pair, the first entry + in `earnings_calendars[label]` on or after `date`. Entries falling + after the last date in a calendar will have `np_NaT` as the result in + the output. + + See Also + -------- + previous_earnings_date_frame + """ + cols = { + equity: np.full_like(dates, np_NaT) for equity in announcement_dates + } + raw_dates = dates.values + for equity, earnings_dates in iteritems(announcement_dates): + data = cols[equity] + if not earnings_dates.index.is_monotonic_increasing: + earnings_dates = earnings_dates.sort_index() + + # Iterate over the raw Series values, since we're comparing against + # numpy arrays anyway. + iterkv = zip(earnings_dates.index.values, earnings_dates.values) + for timestamp, announce_date in iterkv: + date_mask = (timestamp <= raw_dates) & (raw_dates <= announce_date) + value_mask = (announce_date <= data) | (data == np_NaT) + data[date_mask & value_mask] = announce_date + + return pd.DataFrame(index=dates, data=cols) + + +def previous_date_frame(dates, announcement_dates): + """ + Make a DataFrame representing simulated next earnings dates. + + Parameters + ---------- + dates : DatetimeIndex. + The index of the returned DataFrame. + announcement_dates : dict[int -> DatetimeIndex] + Dict mapping sids to an index of dates on which earnings were announced + for that sid. + + Returns + ------- + prev_earnings: pd.DataFrame + A DataFrame representing, for (label, date) pair, the first entry in + `announcement_dates[label]` strictly before `date`. Entries falling + before the first date in a calendar will have `NaT` as the result in + the output. + + See Also + -------- + next_earnings_date_frame + """ + sids = list(announcement_dates) + out = np.full((len(dates), len(sids)), np_NaT, dtype='datetime64[ns]') + dn = dates[-1].asm8 + for col_idx, sid in enumerate(sids): + # announcement_dates[sid] is Series mapping knowledge_date to actual + # announcement date. We don't care about the knowledge date for + # computing previous earnings. + values = announcement_dates[sid].values + values = values[values <= dn] + out[dates.searchsorted(values), col_idx] = values + + frame = pd.DataFrame(out, index=dates, columns=sids) + frame.ffill(inplace=True) + return frame