diff --git a/tests/pipeline/test_earnings.py b/tests/pipeline/test_earnings.py index 871fb155..8c52986f 100644 --- a/tests/pipeline/test_earnings.py +++ b/tests/pipeline/test_earnings.py @@ -27,57 +27,15 @@ from zipline.pipeline.loaders.blaze import ( TS_FIELD_NAME, ) from zipline.utils.numpy_utils import make_datetime64D, np_NaT -from zipline.utils.tradingcalendar import trading_days from zipline.utils.test_utils import ( make_simple_equity_info, - powerset, tmp_asset_finder, + gen_calendars, + to_series, + num_days_in_range, ) -def _to_series(knowledge_dates, earning_dates): - """ - Helper for converting a dict of strings to a Series of datetimes. - - This is just for making the test cases more readable. - """ - return pd.Series( - index=pd.to_datetime(knowledge_dates), - data=pd.to_datetime(earning_dates), - ) - - -def num_days_in_range(dates, start, end): - """ - Return the number of days in `dates` between start and end, inclusive. - """ - start_idx, stop_idx = dates.slice_locs(start, end) - return stop_idx - start_idx - - -def gen_calendars(): - """ - Generate calendars to use as inputs to test_compute_latest. - """ - start, stop = '2014-01-01', '2014-01-31' - all_dates = pd.date_range(start, stop, tz='utc') - - # These dates are the points where announcements or knowledge dates happen. - # Test every combination of them being absent. - critical_dates = pd.to_datetime([ - '2014-01-05', - '2014-01-10', - '2014-01-15', - '2014-01-20', - ]) - for to_drop in map(list, powerset(critical_dates)): - # Have to yield tuples. - yield (all_dates.drop(to_drop),) - - # Also test with the trading calendar. - yield (trading_days[trading_days.slice_indexer(start, stop)],) - - class EarningsCalendarLoaderTestCase(TestCase): """ Tests for loading the earnings announcement data. @@ -99,22 +57,22 @@ class EarningsCalendarLoaderTestCase(TestCase): cls.earnings_dates = { # K1--K2--E1--E2. - A: _to_series( + A: to_series( knowledge_dates=['2014-01-05', '2014-01-10'], earning_dates=['2014-01-15', '2014-01-20'], ), # K1--K2--E2--E1. - B: _to_series( + B: to_series( knowledge_dates=['2014-01-05', '2014-01-10'], earning_dates=['2014-01-20', '2014-01-15'] ), # K1--E1--K2--E2. - C: _to_series( + C: to_series( knowledge_dates=['2014-01-05', '2014-01-15'], earning_dates=['2014-01-10', '2014-01-20'] ), # K1 == K2. - D: _to_series( + D: to_series( knowledge_dates=['2014-01-05'] * 2, earning_dates=['2014-01-10', '2014-01-15'], ), @@ -294,7 +252,16 @@ class EarningsCalendarLoaderTestCase(TestCase): index=announcement_dates.index, ) - @parameterized.expand(gen_calendars()) + @parameterized.expand(gen_calendars( + '2014-01-01', + '2014-01-31', + critical_dates=pd.to_datetime([ + '2014-01-05', + '2014-01-10', + '2014-01-15', + '2014-01-20', + ]), + )) def test_compute_earnings(self, dates): ( diff --git a/zipline/pipeline/loaders/blaze/core.py b/zipline/pipeline/loaders/blaze/core.py index 8db633a6..556eb3f8 100644 --- a/zipline/pipeline/loaders/blaze/core.py +++ b/zipline/pipeline/loaders/blaze/core.py @@ -154,7 +154,7 @@ from toolz import ( memoize, ) import toolz.curried.operator as op -from six import with_metaclass, PY2, itervalues +from six import with_metaclass, PY2, itervalues, iteritems from zipline.pipeline.data.dataset import DataSet, Column @@ -903,3 +903,86 @@ class BlazeLoader(dict): ) global_loader = BlazeLoader.global_instance() + + +def bind_expression_to_resources(expr, resources): + """ + Bind a Blaze expression to resources. + + Parameters + ---------- + expr : bz.Expr + The expression to which we want to bind resources. + resources : dict[bz.Symbol -> any] + Mapping from the atomic terms of ``expr`` to actual data resources. + + Returns + ------- + bound_expr : bz.Expr + ``expr`` with bound resources. + """ + # bind the resources into the expression + if resources is None: + resources = {} + + # _subs stands for substitute. It's not actually private, blaze just + # prefixes symbol-manipulation methods with underscores to prevent + # collisions with data column names. + return expr._subs({ + k: bz.Data(v, dshape=k.dshape) for k, v in iteritems(resources) + }) + + +def ffill_query_in_range(expr, + lower, + upper, + odo_kwargs=None, + ts_field=TS_FIELD_NAME, + sid_field=SID_FIELD_NAME): + """Query a blaze expression in a given time range properly forward filling + from values that fall before the lower date. + + Parameters + ---------- + expr : Expr + Bound blaze expression. + lower : datetime + The lower date to query for. + upper : datetime + The upper date to query for. + odo_kwargs : dict, optional + The extra keyword arguments to pass to ``odo``. + ts_field : str, optional + The name of the timestamp field in the given blaze expression. + sid_field : str, optional + The name of the sid field in the given blaze expression. + + Returns + ------- + raw : pd.DataFrame + A strict dataframe for the data in the given date range. This may + start before the requested start date if a value is needed to ffill. + """ + odo_kwargs = odo_kwargs or {} + filtered = expr[expr[ts_field] <= lower] + computed_lower = odo( + bz.by( + filtered[sid_field], + timestamp=filtered[ts_field].max(), + ).timestamp.min(), + pd.Timestamp, + **odo_kwargs + ) + if pd.isnull(computed_lower): + # If there is no lower date, just query for data in the date + # range. It must all be null anyways. + computed_lower = lower + + return odo( + expr[ + (expr[ts_field] >= computed_lower) & + (expr[ts_field] <= upper) + ], + pd.DataFrame, + **odo_kwargs + ) diff --git a/zipline/pipeline/loaders/blaze/earnings.py b/zipline/pipeline/loaders/blaze/earnings.py index 6b5c1695..2c4f8dfb 100644 --- a/zipline/pipeline/loaders/blaze/earnings.py +++ b/zipline/pipeline/loaders/blaze/earnings.py @@ -1,11 +1,13 @@ -import blaze as bz from datashape import istabular -from odo import odo import pandas as pd -from six import iteritems from toolz import valmap -from .core import TS_FIELD_NAME, SID_FIELD_NAME +from .core import ( + TS_FIELD_NAME, + SID_FIELD_NAME, + bind_expression_to_resources, + ffill_query_in_range, +) from zipline.pipeline.data import EarningsCalendar from zipline.pipeline.loaders.base import PipelineLoader from zipline.pipeline.loaders.earnings import EarningsCalendarLoader @@ -14,34 +16,6 @@ from zipline.pipeline.loaders.earnings import EarningsCalendarLoader ANNOUNCEMENT_FIELD_NAME = 'announcement_date' -def bind_expression_to_resources(expr, resources): - """ - Bind a Blaze expression to resources. - - Parameters - ---------- - expr : bz.Expr - The expression to which we want to bind resources. - resources : dict[bz.Symbol -> any] - Mapping from the atomic terms of ``expr`` to actual data resources. - - Returns - ------- - bound_expr : bz.Expr - ``expr`` with bound resources. - """ - # bind the resources into the expression - if resources is None: - resources = {} - - # _subs stands for substitute. It's not actually private, blaze just - # prefixes symbol-manipulation methods with underscores to prevent - # collisions with data column names. - return expr._subs({ - k: bz.Data(v, dshape=k.dshape) for k, v in iteritems(resources) - }) - - class BlazeEarningsCalendarLoader(PipelineLoader): """A pipeline loader for the ``EarningsCalendar`` dataset that loads data from a blaze expression. @@ -61,8 +35,8 @@ class BlazeEarningsCalendarLoader(PipelineLoader): Dim * {{ {SID_FIELD_NAME}: int64, - {TS_FIELD_NAME}: datetime64, - {ANNOUNCEMENT_FIELD_NAME}: datetime64, + {TS_FIELD_NAME}: datetime, + {ANNOUNCEMENT_FIELD_NAME}: ?datetime, }} Where each row of the table is a record including the sid to identify the @@ -87,7 +61,6 @@ class BlazeEarningsCalendarLoader(PipelineLoader): def __init__(self, expr, resources=None, - compute_kwargs=None, odo_kwargs=None, dataset=EarningsCalendar): dshape = expr.dshape @@ -106,33 +79,15 @@ class BlazeEarningsCalendarLoader(PipelineLoader): self._dataset = dataset def load_adjusted_array(self, columns, dates, assets, mask): - expr = self._expr - filtered = expr[expr[TS_FIELD_NAME] <= dates[0]] - lower = odo( - bz.by( - filtered[SID_FIELD_NAME], - timestamp=filtered[TS_FIELD_NAME].max(), - ).timestamp.min(), - pd.Timestamp, - **self._odo_kwargs + raw = ffill_query_in_range( + self._expr, + dates[0], + dates[-1], + self._odo_kwargs, ) - if pd.isnull(lower): - # If there is no lower date, just query for data in the date - # range. It must all be null anyways. - lower = dates[0] - - raw = odo( - expr[ - (expr[TS_FIELD_NAME] >= lower) & - (expr[TS_FIELD_NAME] <= dates[-1]) - ], - pd.DataFrame, - **self._odo_kwargs - ) - sids = raw.loc[:, SID_FIELD_NAME] raw.drop( - sids[~(sids.isin(assets) | sids.notnull())].index, + sids[~sids.isin(assets)].index, inplace=True ) diff --git a/zipline/pipeline/loaders/earnings.py b/zipline/pipeline/loaders/earnings.py index b5792a1d..9912faff 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, ), @@ -108,85 +106,3 @@ class EarningsCalendarLoader(PipelineLoader): ) for column in columns ) - - -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..43dedfce --- /dev/null +++ b/zipline/pipeline/loaders/utils.py @@ -0,0 +1,95 @@ +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, events_by_sid): + """ + Make a DataFrame representing the simulated next known date for an event. + + Parameters + ---------- + dates : pd.DatetimeIndex. + The index of the returned DataFrame. + events_by_sid : dict[int -> pd.Series] + Dict mapping sids to a series of dates. Each k:v pair of the series + represents the date we learned of the event mapping to the date the + event will occur. + Returns + ------- + next_events: pd.DataFrame + A DataFrame where each column is a security from `events_by_sid` where + the values are the dates of the next known event with the knowledge we + had on the date of the index. Entries falling after the last date will + have `NaT` as the result in the output. + + + See Also + -------- + previous_date_frame + """ + cols = { + equity: np.full_like(dates, np_NaT) for equity in events_by_sid + } + raw_dates = dates.values + for equity, event_dates in iteritems(events_by_sid): + data = cols[equity] + if not event_dates.index.is_monotonic_increasing: + event_dates = event_dates.sort_index() + + # Iterate over the raw Series values, since we're comparing against + # numpy arrays anyway. + iterkv = zip(event_dates.index.values, event_dates.values) + for knowledge_date, event_date in iterkv: + date_mask = ( + (knowledge_date <= raw_dates) & + (raw_dates <= event_date) + ) + value_mask = (event_date <= data) | (data == np_NaT) + data[date_mask & value_mask] = event_date + + return pd.DataFrame(index=dates, data=cols) + + +def previous_date_frame(date_index, events_by_sid): + """ + Make a DataFrame representing simulated next earnings date_index. + + Parameters + ---------- + date_index : DatetimeIndex. + The index of the returned DataFrame. + events_by_sid : dict[int -> DatetimeIndex] + Dict mapping sids to a series of dates. Each k:v pair of the series + represents the date we learned of the event mapping to the date the + event will occur. + + Returns + ------- + previous_events: pd.DataFrame + A DataFrame where each column is a security from `events_by_sid` where + the values are the dates of the previous event that occured on the date + of the index. Entries falling before the first date will have `NaT` as + the result in the output. + + See Also + -------- + next_date_frame + """ + sids = list(events_by_sid) + out = np.full((len(date_index), len(sids)), np_NaT, dtype='datetime64[ns]') + dn = date_index[-1].asm8 + for col_idx, sid in enumerate(sids): + # events_by_sid[sid] is Series mapping knowledge_date to actual + # event_date. We don't care about the knowledge date for + # computing previous earnings. + values = events_by_sid[sid].values + values = values[values <= dn] + out[date_index.searchsorted(values), col_idx] = values + + frame = pd.DataFrame(out, index=date_index, columns=sids) + frame.ffill(inplace=True) + return frame diff --git a/zipline/utils/memoize.py b/zipline/utils/memoize.py index aeedb01c..70401582 100644 --- a/zipline/utils/memoize.py +++ b/zipline/utils/memoize.py @@ -7,9 +7,8 @@ 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. + """Decorator that marks that an attribute of an instance should not be + computed until needed, and that the value should be memoized. Example ------- @@ -57,6 +56,35 @@ class lazyval(object): del self._cache[instance] +class classlazyval(lazyval): + """ Decorator that marks that an attribute of a class should not be + computed until needed, and that the value should be memoized. + + Example + ------- + + >>> from zipline.utils.memoize import classlazyval + >>> class C(object): + ... count = 0 + ... @classlazyval + ... def val(cls): + ... cls.count += 1 + ... return "val" + ... + >>> C.count + 0 + >>> C.val, C.count + ('val', 1) + >>> C.val, C.count + ('val', 1) + """ + # We don't reassign the name on the class to implement the caching because + # then we would need to use a metaclass to track the name of the + # descriptor. + def __get__(self, instance, owner): + return super(classlazyval, self).__get__(owner, owner) + + def remember_last(f): """ Decorator that remembers the last computed value of a function and doesn't diff --git a/zipline/utils/test_utils.py b/zipline/utils/test_utils.py index 31846868..cbcd3fda 100644 --- a/zipline/utils/test_utils.py +++ b/zipline/utils/test_utils.py @@ -26,6 +26,7 @@ from zipline.assets.asset_writer import AssetDBWriterFromDataFrame from zipline.assets.futures import CME_CODE_TO_MONTH from zipline.finance.order import ORDER_STATUS from zipline.utils import security_list +from zipline.utils.tradingcalendar import trading_days EPOCH = pd.Timestamp(0, tz='UTC') @@ -630,3 +631,36 @@ def powerset(values): Return the power set (i.e., the set of all subsets) of entries in `values`. """ return concat(combinations(values, i) for i in range(len(values) + 1)) + + +def to_series(knowledge_dates, earning_dates): + """ + Helper for converting a dict of strings to a Series of datetimes. + + This is just for making the test cases more readable. + """ + return pd.Series( + index=pd.to_datetime(knowledge_dates), + data=pd.to_datetime(earning_dates), + ) + + +def num_days_in_range(dates, start, end): + """ + Return the number of days in `dates` between start and end, inclusive. + """ + start_idx, stop_idx = dates.slice_locs(start, end) + return stop_idx - start_idx + + +def gen_calendars(start, stop, critical_dates): + """ + Generate calendars to use as inputs. + """ + all_dates = pd.date_range(start, stop, tz='utc') + for to_drop in map(list, powerset(critical_dates)): + # Have to yield tuples. + yield (all_dates.drop(to_drop),) + + # Also test with the trading calendar. + yield (trading_days[trading_days.slice_indexer(start, stop)],)