From 6c6a33c73b7545f96e6217e9fdccf1fbad90cb22 Mon Sep 17 00:00:00 2001 From: Maya Tydykov Date: Fri, 5 Aug 2016 11:53:29 -0400 Subject: [PATCH] ENH: add loader for estimates --- tests/pipeline/test_quarters_estimates.py | 14 ++ zipline/pipeline/common.py | 2 + zipline/pipeline/loaders/blaze/estimates.py | 146 ++++++++++++ zipline/pipeline/loaders/blaze/events.py | 32 +-- zipline/pipeline/loaders/events.py | 12 +- zipline/pipeline/loaders/quarter_estimates.py | 221 ++++++++++++++++++ zipline/pipeline/loaders/utils.py | 32 +++ 7 files changed, 420 insertions(+), 39 deletions(-) create mode 100644 tests/pipeline/test_quarters_estimates.py create mode 100644 zipline/pipeline/loaders/blaze/estimates.py create mode 100644 zipline/pipeline/loaders/quarter_estimates.py diff --git a/tests/pipeline/test_quarters_estimates.py b/tests/pipeline/test_quarters_estimates.py new file mode 100644 index 00000000..201011ea --- /dev/null +++ b/tests/pipeline/test_quarters_estimates.py @@ -0,0 +1,14 @@ +def test_shift_quarters_forward(): + quarters = list(range(1, 5)) + shifts = list(range(5)) + expected = [(x, i) for ] + expected = ((0, 1), (0, 2), (0, 3), (0, 4), (1, 1), + (0, 2), (0, 3), (0, 4), (1, 1), (1, 2)) + for quarter in quarters: + for shift in shifts: + yrs_to_shift, new_qtr = EstimizeLoader.calc_forward_shift(quarter, + shift) + if quarter + shift <= 4: + assert yrs_to_shift == 0 + assert new_qtr == quarter + shift + else: diff --git a/zipline/pipeline/common.py b/zipline/pipeline/common.py index e64b1dc9..7b48e271 100644 --- a/zipline/pipeline/common.py +++ b/zipline/pipeline/common.py @@ -6,6 +6,8 @@ ANNOUNCEMENT_FIELD_NAME = 'announcement_date' CASH_FIELD_NAME = 'cash' DAYS_SINCE_PREV = 'days_since_prev' DAYS_TO_NEXT = 'days_to_next' +FISCAL_QUARTER_FIELD_NAME = 'fiscal_quarter' +FISCAL_YEAR_FIELD_NAME = 'fiscal_year' NEXT_ANNOUNCEMENT = 'next_announcement' PREVIOUS_AMOUNT = 'previous_amount' PREVIOUS_ANNOUNCEMENT = 'previous_announcement' diff --git a/zipline/pipeline/loaders/blaze/estimates.py b/zipline/pipeline/loaders/blaze/estimates.py new file mode 100644 index 00000000..6a89fa7b --- /dev/null +++ b/zipline/pipeline/loaders/blaze/estimates.py @@ -0,0 +1,146 @@ +from datashape import istabular + +from .core import ( + bind_expression_to_resources, + ffill_query_in_range, +) +from zipline.pipeline.loaders.base import PipelineLoader +from zipline.pipeline.loaders.events import ( + EventsLoader, + required_event_fields, +) +from zipline.pipeline.common import ( + SID_FIELD_NAME, + TS_FIELD_NAME, +) +from zipline.pipeline.loaders.quarter_estimates import \ + NextQuartersEstimatesLoader, PreviousQuartersEstimatesLoader +from zipline.pipeline.loaders.utils import ( + check_data_query_args, + normalize_data_query_bounds, + normalize_timestamp_to_query_time, + load_raw_data) +from zipline.utils.input_validation import ensure_timezone, optionally +from zipline.utils.preprocess import preprocess + + +class BlazeEstimatesLoader(PipelineLoader): + """An abstract pipeline loader for the estimates datasets that loads + data from a blaze expression. + + Parameters + ---------- + expr : Expr + The expression representing the data to load. + resources : dict, optional + Mapping from the loadable terms of ``expr`` to actual data resources. + odo_kwargs : dict, optional + Extra keyword arguments to pass to odo when executing the expression. + data_query_time : time, optional + The time to use for the data query cutoff. + data_query_tz : tzinfo or str + The timezeone to use for the data query cutoff. + dataset : DataSet + The DataSet object for which this loader loads data. + + Notes + ----- + The expression should have a tabular dshape of:: + + Dim * {{ + {SID_FIELD_NAME}: int64, + {TS_FIELD_NAME}: datetime, + }} + + And other dataset-specific fields, where each row of the table is a + record including the sid to identify the company, the timestamp where we + learned about the announcement, and the date when the earnings will be z + announced. + + If the '{TS_FIELD_NAME}' field is not included it is assumed that we + start the backtest with knowledge of all announcements. + """ + + @preprocess(data_query_tz=optionally(ensure_timezone)) + def __init__(self, + expr, + columns, + resources=None, + odo_kwargs=None, + data_query_time=None, + data_query_tz=None, + loader=None): + + dshape = expr.dshape + if not istabular(dshape): + raise ValueError( + 'expression dshape must be tabular, got: %s' % dshape, + ) + + required_cols = list( + required_event_fields(columns) + ) + self._expr = bind_expression_to_resources( + expr[required_cols], + resources, + ) + self._columns = columns + self._odo_kwargs = odo_kwargs if odo_kwargs is not None else {} + check_data_query_args(data_query_time, data_query_tz) + self._data_query_time = data_query_time + self._data_query_tz = data_query_tz + self.loader = loader + + def load_adjusted_array(self, columns, dates, assets, mask): + raw = load_raw_data(assets, dates, self._data_query_time, + self._data_query_tz, self._exp, self._odo_kwargs) + + return self.loader( + events=raw, + next_value_columns=self._columns, + ).load_adjusted_array( + columns, + dates, + assets, + mask, + ) + + +class BlazeNextEstimatesLoader(BlazeEstimatesLoader): + loader = NextQuartersEstimatesLoader + + def __init__(self, + expr, + columns, + resources=None, + odo_kwargs=None, + data_query_time=None, + data_query_tz=None, + loader=None): + super(BlazeNextEstimatesLoader).__init__(expr, + columns, + resources, + odo_kwargs, + data_query_time, + data_query_tz, + loader) + + +class BlazePreviousEstimatesLoader(BlazeEstimatesLoader): + loader = PreviousQuartersEstimatesLoader + + def __init__(self, + expr, + columns, + resources=None, + odo_kwargs=None, + data_query_time=None, + data_query_tz=None, + loader=None): + super(BlazeNextEstimatesLoader).__init__(expr, + columns, + resources, + odo_kwargs, + data_query_time, + data_query_tz, + loader) diff --git a/zipline/pipeline/loaders/blaze/events.py b/zipline/pipeline/loaders/blaze/events.py index c71646cb..4165166b 100644 --- a/zipline/pipeline/loaders/blaze/events.py +++ b/zipline/pipeline/loaders/blaze/events.py @@ -17,7 +17,7 @@ from zipline.pipeline.loaders.utils import ( check_data_query_args, normalize_data_query_bounds, normalize_timestamp_to_query_time, -) + load_raw_data) from zipline.utils.input_validation import ensure_timezone, optionally from zipline.utils.preprocess import preprocess @@ -90,34 +90,8 @@ class BlazeEventsLoader(PipelineLoader): self._data_query_tz = data_query_tz def load_adjusted_array(self, columns, dates, assets, mask): - data_query_time = self._data_query_time - data_query_tz = self._data_query_tz - lower_dt, upper_dt = normalize_data_query_bounds( - dates[0], - dates[-1], - data_query_time, - data_query_tz, - ) - - raw = ffill_query_in_range( - self._expr, - lower_dt, - upper_dt, - self._odo_kwargs, - ) - sids = raw.loc[:, SID_FIELD_NAME] - raw.drop( - sids[~sids.isin(assets)].index, - inplace=True - ) - if data_query_time is not None: - normalize_timestamp_to_query_time( - raw, - data_query_time, - data_query_tz, - inplace=True, - ts_field=TS_FIELD_NAME, - ) + raw = load_raw_data(assets, dates, self._data_query_time, + self._data_query_tz, self._expr, self._odo_kwargs) return EventsLoader( events=raw, diff --git a/zipline/pipeline/loaders/events.py b/zipline/pipeline/loaders/events.py index 645143d8..1c49779c 100644 --- a/zipline/pipeline/loaders/events.py +++ b/zipline/pipeline/loaders/events.py @@ -41,16 +41,8 @@ def validate_column_specs(events, next_value_columns, previous_value_columns): serve the BoundColumns described by ``next_value_columns`` and ``previous_value_columns``. """ - required = { - TS_FIELD_NAME, - SID_FIELD_NAME, - EVENT_DATE_FIELD_NAME, - }.union( - # We also expect any of the field names that our loadable columns - # are mapped to. - viewvalues(next_value_columns), - viewvalues(previous_value_columns), - ) + required = required_event_fields(next_value_columns, + previous_value_columns) received = set(events.columns) missing = required - received if missing: diff --git a/zipline/pipeline/loaders/quarter_estimates.py b/zipline/pipeline/loaders/quarter_estimates.py new file mode 100644 index 00000000..495bc075 --- /dev/null +++ b/zipline/pipeline/loaders/quarter_estimates.py @@ -0,0 +1,221 @@ +from itertools import groupby +import numpy as np +import pandas as pd +from six import viewvalues +from zipline.pipeline.common import AD_FIELD_NAME, SID_FIELD_NAME, \ + EVENT_DATE_FIELD_NAME, FISCAL_QUARTER_FIELD_NAME, FISCAL_YEAR_FIELD_NAME +from zipline.pipeline.loaders.base import PipelineLoader +from zipline.pipeline.loaders.frame import DataFrameLoader + + +def required_event_fields(columns): + """ + Compute the set of resource columns required to serve + ``next_value_columns`` and ``previous_value_columns``. + """ + # These metadata columns are used to align event indexers. + return { + AD_FIELD_NAME, + SID_FIELD_NAME, + EVENT_DATE_FIELD_NAME, + FISCAL_QUARTER_FIELD_NAME, + FISCAL_YEAR_FIELD_NAME + }.union( + # We also expect any of the field names that our loadable columns + # are mapped to. + viewvalues(columns), + ) + + +def validate_column_specs(events, columns): + """ + Verify that the columns of ``events`` can be used by an EventsLoader to + serve the BoundColumns described by ``next_value_columns`` and + ``previous_value_columns``. + """ + required = required_event_fields(columns) + received = set(events.columns) + missing = required - received + if missing: + raise ValueError( + "EventsLoader missing required columns {missing}.\n" + "Got Columns: {received}\n" + "Expected Columns: {required}".format( + missing=sorted(missing), + received=sorted(received), + required=sorted(required), + ) + ) + + +def calc_forward_shift(qtr, num_shifts): + yrs_to_shift, new_qtr = divmod(qtr + num_shifts, 4) + if yrs_to_shift == 1 and new_qtr == 0: + yrs_to_shift = 0 + new_qtr = 4 + return yrs_to_shift, new_qtr + + +def calc_backward_shift(qtr, num_shifts): + yrs_to_shift, new_qtr = divmod(abs(num_shifts - qtr), 4) + if yrs_to_shift == 0 and new_qtr == 0: + yrs_to_shift = 1 + new_qtr = 4 + yrs_to_shift = -yrs_to_shift + return yrs_to_shift, new_qtr + + +class QuarterEstimatesLoader(PipelineLoader): + def __init__(self, + events, + columns): + validate_column_specs( + events, + columns + ) + + self.events = events[ + events[EVENT_DATE_FIELD_NAME].notnull() and + events[FISCAL_QUARTER_FIELD_NAME].notnull() and + events[FISCAL_YEAR_FIELD_NAME].notnull() + ] + + self.columns = columns + + def load_quarters(self, next_releases, num_quarters, dates_sids, gb): + pass + + def load_adjusted_array(self, columns, dates, assets, mask): + groups = groupby(lambda x: x.dataset.num_quarters, columns) + out = {} + date_values = pd.DataFrame(dates, columns=['dates']) + date_values['key'] = 1 + self.events['key'] = 1 + merged = pd.merge(date_values, self.events, on='key') + asset_df = pd.DataFrame(assets, columns=['sid']) + asset_df['key'] = 1 + dates_sids = pd.merge(date_values, asset_df, on='key') + for num_quarters in groups: + columns = groups[num_quarters] + # First, group by sid, fiscal year, and fiscal quarter and only + # keep the last estimate made. + final_releases_per_qtr = merged[merged.asof_date <= + merged.dates].sort( + ['dates', 'asof_date'] + ).groupby( + ['dates', 'sid', 'fiscal_year', 'fiscal_quarter'] + ).last() + gb = final_releases_per_qtr.reset_index().groupby(['dates', 'sid']) + # Split the date-sid combinations into ones with a next release + # and ones without + eligible_next_releases = pd.concat([group[1] for group in gb if ( + group[1][EVENT_DATE_FIELD_NAME] >= group[1]['dates'] + ).any()]) + + eligible_next_releases.sort(EVENT_DATE_FIELD_NAME) + # For each sid, get the next release/year/quarter that we care + # about. + next_releases = eligible_next_releases.groupby( + ['dates', 'sid'] + ).min() + next_releases = next_releases.rename( + columns={'fiscal_year': 'next_fiscal_year', + 'fiscal_quarter': 'next_fiscal_quarter'} + ) + + result = self.load_quarters(next_releases, + num_quarters, + dates_sids) + + for c in columns: + column_name = self.columns[c.name] + # Need to pass a DataFrame that has dates as the index and + # all sids as columns with column values being the value in + # 'result' for column c + loader = DataFrameLoader( + c, + result.pivot(index='dates', + columns='sid', + values=column_name), + adjustments=None + ) + out[c] = loader.load_adjusted_array([c], dates, assets, mask)[c] + return out + + +class NextQuartersEstimatesLoader(QuarterEstimatesLoader): + def __init__(self, + events, + columns): + super(NextQuartersEstimatesLoader).__init__(events, columns) + + def load_quarters(self, next_releases, num_quarters, dates_sids, gb): + # `next_qtr` is already the next quarter over, + # so we should offest `num_shifts` by 1. + next_releases['fiscal_quarter'] = next_releases.apply( + lambda x: calc_forward_shift(x['next_fiscal_quarter'], + num_quarters - 1)[1], + axis=1 + ) + next_releases['fiscal_year'] = next_releases.apply( + lambda x: + x['next_fiscal_year'] + + calc_forward_shift(x['next_fiscal_quarter'], + num_quarters - 1)[0], + axis=1 + ) + # Merge to get the rows we care about for each date + result = dates_sids.merge(next_releases.reset_index(), + on=(['dates', 'sid']), + how='left') + return result + + +class PreviousQuartersEstimatesLoader(QuarterEstimatesLoader): + def __init__(self, + events, + columns): + super(PreviousQuartersEstimatesLoader).__init__(events, columns) + + def load_quarters(self, next_releases, num_quarters, dates_sids, gb): + next_releases['fiscal_quarter'] = next_releases.apply( + lambda x: calc_backward_shift(x['next_fiscal_quarter'], + num_quarters)[1], + axis=1 + ) + next_releases['fiscal_year'] = next_releases.apply( + lambda x: + x['next_fiscal_year'] + + calc_backward_shift(x['next_fiscal_quarter'], + num_quarters)[0], + axis=1 + ) + only_previous_releases = pd.concat([group[1] for group in gb if ( + group[1][EVENT_DATE_FIELD_NAME] < group[1]['dates'] + ).all()]) + only_previous_releases.sort(EVENT_DATE_FIELD_NAME) + # For each sid, get the latest release we knew about prior to + # each simulation date. + previous_releases = only_previous_releases.groupby(['dates', + 'sid']).max() + previous_releases = previous_releases.rename(columns={ + 'fiscal_year': 'previous_fiscal_year', + 'fiscal_quarter': 'previous_fiscal_quarter' + }) + previous_releases['fiscal_quarter'] = previous_releases.apply( + lambda x: calc_backward_shift(x['previous_fiscal_quarter'], + num_quarters)[1], + axis=1 + ) + previous_releases['fiscal_year'] = previous_releases.apply( + lambda x: + x['previous_fiscal_year'] + + calc_backward_shift(x['previous_fiscal_quarter'], + num_quarters)[0], + axis=1 + ) + all_releases = pd.concat([next_releases, previous_releases]) + # Merge to get the rows we care about for each date + result = dates_sids.merge(all_releases.reset_index(), + on=(['dates', 'sid']), how='left') + return result diff --git a/zipline/pipeline/loaders/utils.py b/zipline/pipeline/loaders/utils.py index 77a9f447..2f388810 100644 --- a/zipline/pipeline/loaders/utils.py +++ b/zipline/pipeline/loaders/utils.py @@ -2,6 +2,8 @@ import datetime import numpy as np import pandas as pd +from zipline.pipeline.common import TS_FIELD_NAME, SID_FIELD_NAME +from zipline.pipeline.loaders.blaze.core import ffill_query_in_range from zipline.utils.pandas_utils import mask_between_time @@ -272,3 +274,33 @@ def check_data_query_args(data_query_time, data_query_tz): data_query_tz, ), ) + + +def load_raw_data(assets, dates, data_query_time, data_query_tz, expr, + odo_kwargs): + lower_dt, upper_dt = normalize_data_query_bounds( + dates[0], + dates[-1], + data_query_time, + data_query_tz, + ) + raw = ffill_query_in_range( + expr, + lower_dt, + upper_dt, + odo_kwargs, + ) + sids = raw.loc[:, SID_FIELD_NAME] + raw.drop( + sids[~sids.isin(assets)].index, + inplace=True + ) + if data_query_time is not None: + normalize_timestamp_to_query_time( + raw, + data_query_time, + data_query_tz, + inplace=True, + ts_field=TS_FIELD_NAME, + ) + return raw