ENH: add loader for estimates

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:

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'

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)

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,

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:

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

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