diff --git a/tests/pipeline/base.py b/tests/pipeline/base.py index e8261ccb..1f89e56f 100644 --- a/tests/pipeline/base.py +++ b/tests/pipeline/base.py @@ -1,35 +1,25 @@ """ Base class for Pipeline API unittests. """ -import abc from functools import wraps from unittest import TestCase -from nose_parameterized import parameterized import numpy as np from numpy import arange, prod -import pandas as pd from pandas import date_range, Int64Index, DataFrame -from pandas.util.testing import assert_series_equal from six import iteritems -from zipline.pipeline import Pipeline, TermGraph +from zipline.pipeline import TermGraph from zipline.pipeline.engine import SimplePipelineEngine from zipline.pipeline.term import AssetExists from zipline.testing import ( check_arrays, ExplodingObject, - gen_calendars, make_simple_equity_info, - num_days_in_range, tmp_asset_finder, ) from zipline.utils.functional import dzip_exact -from zipline.utils.numpy_utils import ( - NaTD, - make_datetime64D -) from zipline.utils.pandas_utils import explode from zipline.utils.tradingcalendar import trading_day @@ -176,114 +166,3 @@ class BasePipelineTestCase(TestCase): @with_default_shape def ones_mask(self, shape): return np.ones(shape, dtype=bool) - - -class EventLoaderCommonMixin(object): - @abc.abstractproperty - def get_sids(cls): - raise NotImplementedError('get_sids') - - @classmethod - def get_equity_info(cls): - return make_simple_equity_info( - cls.get_sids(), - start_date=pd.Timestamp('2013-01-01', tz='UTC'), - end_date=pd.Timestamp('2015-01-01', tz='UTC'), - ) - - def zip_with_floats(self, dates, flts): - return pd.Series(flts, index=dates).astype('float') - - def num_days_between(self, dates, start_date, end_date): - return num_days_in_range(dates, start_date, end_date) - - def zip_with_dates(self, index_dates, dts): - return pd.Series(pd.to_datetime(dts), index=index_dates) - - def loader_args(self, dates): - """Construct the base object to pass to the loader. - - Parameters - ---------- - dates : pd.DatetimeIndex - The dates we can serve. - - Returns - ------- - args : tuple[any] - The arguments to forward to the loader positionally. - """ - return dates, self.dataset - - def setup_engine(self, dates): - """ - Make a Pipeline Enigne object based on the given dates. - """ - loader = self.loader_type(*self.loader_args(dates)) - return SimplePipelineEngine(lambda _: loader, dates, self.finder) - - @staticmethod - def _compute_busday_offsets(announcement_dates): - """ - Compute expected business day offsets from a DataFrame of announcement - dates. - """ - # Column-vector of dates on which factor `compute` will be called. - raw_call_dates = announcement_dates.index.values.astype( - 'datetime64[D]' - )[:, None] - - # 2D array of dates containining expected nexg announcement. - raw_announce_dates = ( - announcement_dates.values.astype('datetime64[D]') - ) - - # Set NaTs to 0 temporarily because busday_count doesn't support NaT. - # We fill these entries with NaNs later. - whereNaT = raw_announce_dates == NaTD - raw_announce_dates[whereNaT] = make_datetime64D(0) - - # The abs call here makes it so that we can use this function to - # compute offsets for both next and previous earnings (previous - # earnings offsets come back negative). - expected = abs(np.busday_count( - raw_call_dates, - raw_announce_dates - ).astype(float)) - - expected[whereNaT] = np.nan - return pd.DataFrame( - data=expected, - columns=announcement_dates.columns, - index=announcement_dates.index, - ) - - @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', - ], utc=True), - )) - def test_compute(self, dates): - engine = self.setup_engine(dates) - self.setup(dates) - - pipe = Pipeline( - columns=self.pipeline_columns - ) - - result = engine.run_pipeline( - pipe, - start_date=dates[0], - end_date=dates[-1], - ) - - for sid in self.get_sids(): - for col_name in self.cols.keys(): - assert_series_equal(result[col_name].xs(sid, level=1), - self.cols[col_name][sid], - check_names=False) diff --git a/tests/pipeline/test_buyback_auth.py b/tests/pipeline/test_buyback_auth.py index 35fe592f..e86abda5 100644 --- a/tests/pipeline/test_buyback_auth.py +++ b/tests/pipeline/test_buyback_auth.py @@ -1,16 +1,10 @@ """ Tests for the reference loader for Buyback Authorizations. """ -from functools import partial -from unittest import TestCase - import blaze as bz from blaze.compute.core import swap_resources_into_scope -from contextlib2 import ExitStack -import itertools import pandas as pd from six import iteritems -from .base import EventLoaderCommonMixin from zipline.pipeline.common import( BUYBACK_ANNOUNCEMENT_FIELD_NAME, @@ -39,7 +33,14 @@ from zipline.pipeline.loaders.blaze import ( BlazeCashBuybackAuthorizationsLoader, BlazeShareBuybackAuthorizationsLoader, ) -from zipline.testing import tmp_asset_finder +from zipline.pipeline.loaders.utils import ( + get_values_for_date_ranges, + zip_with_floats, + zip_with_dates +) +from zipline.testing.fixtures import ( + WithPipelineEventDataLoader, ZiplineTestCase +) date_intervals = [[None, '2014-01-04'], ['2014-01-05', '2014-01-09'], ['2014-01-10', None]] @@ -62,48 +63,20 @@ buyback_authorizations_cases = [ ] -def get_values_for_date_ranges(zip_with_floats_dates, - num_days_between_dates, - vals_for_date_intervals): - # Fill in given values for given date ranges. - return zip_with_floats_dates( - list( - itertools.chain(*[ - [val] * num_days_between_dates(*date_intervals[i]) - for i, val in enumerate(vals_for_date_intervals) - ]) - ) - ) - - -def get_expected_previous_values(zip_with_floats_dates, - num_days_between_dates, +def get_expected_previous_values(zip_date_index_with_vals, dates, vals_for_date_intervals): return pd.DataFrame({ - 0: get_values_for_date_ranges(zip_with_floats_dates, - num_days_between_dates, - vals_for_date_intervals), - 1: zip_with_floats_dates(['NaN'] * len(dates)), + 0: get_values_for_date_ranges(zip_date_index_with_vals, + vals_for_date_intervals, + date_intervals, + dates), + 1: zip_date_index_with_vals(dates, ['NaN'] * len(dates)), }, index=dates) -def get_expected_previous_dates(zip_with_dates_for_dates, - num_days_between_for_dates, - dates): - return pd.DataFrame({ - 0: zip_with_dates_for_dates( - ['NaT'] * num_days_between_for_dates(None, '2014-01-04') + - ['2014-01-04'] * num_days_between_for_dates('2014-01-05', - '2014-01-09') + - ['2014-01-09'] * num_days_between_for_dates('2014-01-10', - None), - ), - 1: zip_with_dates_for_dates(['NaT'] * len(dates)) - }) - - -class CashBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonMixin): +class CashBuybackAuthLoaderTestCase(WithPipelineEventDataLoader, + ZiplineTestCase): """ Test for cash buyback authorizations dataset. """ @@ -121,43 +94,33 @@ class CashBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonMixin): return range(2) @classmethod - def setUpClass(cls): - cls._cleanup_stack = stack = ExitStack() - cls.finder = stack.enter_context( - tmp_asset_finder(equities=cls.get_equity_info()), - ) - cls.cols = {} - cls.dataset = {sid: - frame.drop(SHARE_COUNT_FIELD_NAME, axis=1) - for sid, frame - in enumerate(buyback_authorizations_cases)} - cls.loader_type = CashBuybackAuthorizationsLoader + def get_dataset(cls): + return {sid: + frame.drop(SHARE_COUNT_FIELD_NAME, axis=1) + for sid, frame + in enumerate(buyback_authorizations_cases)} - @classmethod - def tearDownClass(cls): - cls._cleanup_stack.close() + loader_type = CashBuybackAuthorizationsLoader def setup(self, dates): - zip_with_floats_dates = partial(self.zip_with_floats, dates) - num_days_between_dates = partial(self.num_days_between, dates) - num_days_between_for_dates = partial(self.num_days_between, dates) - zip_with_dates_for_dates = partial(self.zip_with_dates, dates) + cols = {} _expected_previous_cash = get_expected_previous_values( - zip_with_floats_dates, num_days_between_dates, dates, + zip_with_floats, dates, ['NaN', 10, 20] ) - self.cols[ + cols[ PREVIOUS_BUYBACK_ANNOUNCEMENT - ] = get_expected_previous_dates(zip_with_dates_for_dates, - num_days_between_for_dates, - dates) - self.cols[PREVIOUS_BUYBACK_CASH] = _expected_previous_cash - self.cols[DAYS_SINCE_PREV] = self._compute_busday_offsets( - self.cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] + ] = get_expected_previous_values(zip_with_dates, dates, + ['NaT', '2014-01-04', '2014-01-09']) + cols[PREVIOUS_BUYBACK_CASH] = _expected_previous_cash + cols[DAYS_SINCE_PREV] = self._compute_busday_offsets( + cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] ) + return cols -class ShareBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonMixin): +class ShareBuybackAuthLoaderTestCase(WithPipelineEventDataLoader, + ZiplineTestCase): """ Test for share buyback authorizations dataset. """ @@ -175,56 +138,41 @@ class ShareBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonMixin): return range(2) @classmethod - def setUpClass(cls): - cls._cleanup_stack = stack = ExitStack() - cls.finder = stack.enter_context( - tmp_asset_finder(equities=cls.get_equity_info()), - ) - cls.cols = {} - cls.dataset = {sid: - frame.drop(CASH_FIELD_NAME, axis=1) - for sid, frame - in enumerate(buyback_authorizations_cases)} - cls.loader_type = ShareBuybackAuthorizationsLoader + def get_dataset(cls): + return {sid: + frame.drop(CASH_FIELD_NAME, axis=1) + for sid, frame + in enumerate(buyback_authorizations_cases)} - @classmethod - def tearDownClass(cls): - cls._cleanup_stack.close() + loader_type = ShareBuybackAuthorizationsLoader def setup(self, dates): - zip_with_floats_dates = partial(self.zip_with_floats, dates) - num_days_between_dates = partial(self.num_days_between, dates) - num_days_between_for_dates = partial(self.num_days_between, dates) - zip_with_dates_for_dates = partial(self.zip_with_dates, dates) - - self.cols[ + cols = {} + cols[ PREVIOUS_BUYBACK_SHARE_COUNT - ] = get_expected_previous_values(zip_with_floats_dates, - num_days_between_dates, dates, + ] = get_expected_previous_values(zip_with_floats, + dates, ['NaN', 1, 15]) - self.cols[ + cols[ PREVIOUS_BUYBACK_ANNOUNCEMENT - ] = get_expected_previous_dates(zip_with_dates_for_dates, - num_days_between_for_dates, - dates) - self.cols[DAYS_SINCE_PREV] = self._compute_busday_offsets( - self.cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] + ] = get_expected_previous_values(zip_with_dates, dates, + ['NaT', '2014-01-04', '2014-01-09']) + cols[DAYS_SINCE_PREV] = self._compute_busday_offsets( + cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] ) + return cols class BlazeCashBuybackAuthLoaderTestCase(CashBuybackAuthLoaderTestCase): """ Test case for loading via blaze. """ - @classmethod - def setUpClass(cls): - super(BlazeCashBuybackAuthLoaderTestCase, cls).setUpClass() - cls.loader_type = BlazeCashBuybackAuthorizationsLoader + loader_type = BlazeCashBuybackAuthorizationsLoader - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): _, mapping = super( BlazeCashBuybackAuthLoaderTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return (bz.data(pd.concat( pd.DataFrame({ BUYBACK_ANNOUNCEMENT_FIELD_NAME: @@ -242,16 +190,13 @@ class BlazeCashBuybackAuthLoaderTestCase(CashBuybackAuthLoaderTestCase): class BlazeShareBuybackAuthLoaderTestCase(ShareBuybackAuthLoaderTestCase): """ Test case for loading via blaze. """ - @classmethod - def setUpClass(cls): - super(BlazeShareBuybackAuthLoaderTestCase, cls).setUpClass() - cls.loader_type = BlazeShareBuybackAuthorizationsLoader + loader_type = BlazeShareBuybackAuthorizationsLoader - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): _, mapping = super( BlazeShareBuybackAuthLoaderTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return (bz.data(pd.concat( pd.DataFrame({ BUYBACK_ANNOUNCEMENT_FIELD_NAME: @@ -270,11 +215,11 @@ class BlazeShareBuybackAuthLoaderNotInteractiveTestCase( BlazeShareBuybackAuthLoaderTestCase): """Test case for passing a non-interactive symbol and a dict of resources. """ - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): (bound_expr,) = super( BlazeShareBuybackAuthLoaderNotInteractiveTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return swap_resources_into_scope(bound_expr, {}) @@ -282,9 +227,9 @@ class BlazeCashBuybackAuthLoaderNotInteractiveTestCase( BlazeCashBuybackAuthLoaderTestCase): """Test case for passing a non-interactive symbol and a dict of resources. """ - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): (bound_expr,) = super( BlazeCashBuybackAuthLoaderNotInteractiveTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return swap_resources_into_scope(bound_expr, {}) diff --git a/tests/pipeline/test_dividends.py b/tests/pipeline/test_dividends.py new file mode 100644 index 00000000..431078b3 --- /dev/null +++ b/tests/pipeline/test_dividends.py @@ -0,0 +1,436 @@ +""" +Tests for the reference loader for Dividends datasets. +""" +import blaze as bz +from blaze.compute.core import swap_resources_into_scope +import pandas as pd +from six import iteritems + +from zipline.pipeline.common import ( + ANNOUNCEMENT_FIELD_NAME, + DAYS_SINCE_PREV_DIVIDEND_ANNOUNCEMENT, + DAYS_SINCE_PREV_EX_DATE, + DAYS_TO_NEXT_EX_DATE, + NEXT_AMOUNT, + NEXT_EX_DATE, + NEXT_PAY_DATE, + PREVIOUS_ANNOUNCEMENT, + PREVIOUS_EX_DATE, + PREVIOUS_PAY_DATE, + PREVIOUS_AMOUNT, + SID_FIELD_NAME, + TS_FIELD_NAME, + CASH_AMOUNT_FIELD_NAME, + EX_DATE_FIELD_NAME, + PAY_DATE_FIELD_NAME +) +from zipline.pipeline.data.dividends import ( + DividendsByAnnouncementDate, + DividendsByExDate, + DividendsByPayDate +) +from zipline.pipeline.factors.events import ( + BusinessDaysSinceDividendAnnouncement, + BusinessDaysSincePreviousExDate, + BusinessDaysUntilNextExDate +) +from zipline.pipeline.loaders.blaze.dividends import ( + BlazeDividendsByAnnouncementDateLoader, + BlazeDividendsByPayDateLoader, + BlazeDividendsByExDateLoader +) +from zipline.pipeline.loaders.dividends import ( + DividendsByAnnouncementDateLoader, + DividendsByExDateLoader, + DividendsByPayDateLoader +) +from zipline.pipeline.loaders.utils import ( + get_values_for_date_ranges, + zip_with_dates, + zip_with_floats +) +from zipline.testing.fixtures import ( + WithPipelineEventDataLoader, + ZiplineTestCase +) + +dividends_cases = [ + # K1--K2--A1--A2. + pd.DataFrame({ + CASH_AMOUNT_FIELD_NAME: [1, 15], + EX_DATE_FIELD_NAME: pd.to_datetime(['2014-01-15', '2014-01-20']), + PAY_DATE_FIELD_NAME: pd.to_datetime(['2014-01-15', '2014-01-20']), + TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-10']), + ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-04', '2014-01-09']) + }), + # K1--K2--A2--A1. + pd.DataFrame({ + CASH_AMOUNT_FIELD_NAME: [7, 13], + EX_DATE_FIELD_NAME: pd.to_datetime(['2014-01-20', '2014-01-15']), + PAY_DATE_FIELD_NAME: pd.to_datetime(['2014-01-20', '2014-01-15']), + TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-10']), + ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-04', '2014-01-09']) + }), + # K1--A1--K2--A2. + pd.DataFrame({ + CASH_AMOUNT_FIELD_NAME: [3, 1], + EX_DATE_FIELD_NAME: pd.to_datetime(['2014-01-10', '2014-01-20']), + PAY_DATE_FIELD_NAME: pd.to_datetime(['2014-01-10', '2014-01-20']), + TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-15']), + ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-04', '2014-01-14']) + }), + # K1 == K2. + pd.DataFrame({ + CASH_AMOUNT_FIELD_NAME: [6, 23], + EX_DATE_FIELD_NAME: pd.to_datetime(['2014-01-10', '2014-01-15']), + PAY_DATE_FIELD_NAME: pd.to_datetime(['2014-01-10', '2014-01-15']), + TS_FIELD_NAME: pd.to_datetime(['2014-01-05'] * 2), + ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-04', '2014-01-04']) + }), + pd.DataFrame( + columns=[CASH_AMOUNT_FIELD_NAME, + EX_DATE_FIELD_NAME, + PAY_DATE_FIELD_NAME, + TS_FIELD_NAME, + ANNOUNCEMENT_FIELD_NAME], + dtype='datetime64[ns]' + ), +] + +prev_date_intervals = [ + [ + [None, '2014-01-14'], ['2014-01-15', '2014-01-19'], + ['2014-01-20', None] + ], + [ + [None, '2014-01-14'], ['2014-01-15', '2014-01-19'], + ['2014-01-20', None] + ], + [ + [None, '2014-01-09'], ['2014-01-10', '2014-01-19'], + ['2014-01-20', None] + ], + [ + [None, '2014-01-09'], ['2014-01-10', '2014-01-14'], + ['2014-01-15', None] + ] +] + +next_date_intervals = [ + [ + [None, '2014-01-04'], ['2014-01-05', '2014-01-15'], + ['2014-01-16', '2014-01-20'], ['2014-01-21', None] + ], + [ + [None, '2014-01-04'], ['2014-01-05', '2014-01-09'], + ['2014-01-10', '2014-01-15'], ['2014-01-16', '2014-01-20'], + ['2014-01-21', None] + ], + [ + [None, '2014-01-04'], ['2014-01-05', '2014-01-10'], + ['2014-01-11', '2014-01-14'], ['2014-01-15', '2014-01-20'], + ['2014-01-21', None] + ], + [ + [None, '2014-01-04'], ['2014-01-05', '2014-01-10'], + ['2014-01-11', '2014-01-15'], ['2014-01-16', None] + ] +] + +next_ex_and_pay_dates = [['NaT', '2014-01-15', '2014-01-20', 'NaT'], + ['NaT', '2014-01-20', '2014-01-15', '2014-01-20', + 'NaT'], + ['NaT', '2014-01-10', 'NaT', '2014-01-20', 'NaT'], + ['NaT', '2014-01-10', '2014-01-15', 'NaT']] + +prev_ex_and_pay_dates = [['NaT', '2014-01-15', '2014-01-20'], + ['NaT', '2014-01-15', '2014-01-20'], + ['NaT', '2014-01-10', '2014-01-20'], + ['NaT', '2014-01-10', '2014-01-15']] + +prev_amounts = [['NaN', 1, 15], + ['NaN', 13, 7], + ['NaN', 3, 1], + ['NaN', 6, 23]] + +next_amounts = [['NaN', 1, 15, 'NaN'], + ['NaN', 7, 13, 7, 'NaN'], + ['NaN', 3, 'NaN', 1, 'NaN'], + ['NaN', 6, 23, 'NaN']] + + +def get_vals_for_dates(zip_date_index_with_vals, + vals, + date_invervals, + dates): + return pd.DataFrame({ + 0: get_values_for_date_ranges(zip_date_index_with_vals, + vals[0], + date_invervals[0], + dates), + 1: get_values_for_date_ranges(zip_date_index_with_vals, + vals[1], + date_invervals[1], + dates), + 2: get_values_for_date_ranges(zip_date_index_with_vals, + vals[2], + date_invervals[2], + dates), + # Assume the latest of 2 cash values is used if we find out about 2 + # announcements that happened on the same day for the same sid. + 3: get_values_for_date_ranges(zip_date_index_with_vals, + vals[3], + date_invervals[3], + dates), + 4: zip_date_index_with_vals(dates, ['NaN'] * len(dates)), + }, index=dates) + + +class DividendsByAnnouncementDateTestCase(WithPipelineEventDataLoader, + ZiplineTestCase): + """ + Tests for loading the dividends by announcement date data. + """ + pipeline_columns = { + PREVIOUS_ANNOUNCEMENT: + DividendsByAnnouncementDate.previous_announcement_date.latest, + PREVIOUS_AMOUNT: DividendsByAnnouncementDate.previous_amount.latest, + DAYS_SINCE_PREV_DIVIDEND_ANNOUNCEMENT: + BusinessDaysSinceDividendAnnouncement(), + } + + @classmethod + def get_dataset(cls): + return {sid: + frame.drop([EX_DATE_FIELD_NAME, + PAY_DATE_FIELD_NAME], axis=1) + for sid, frame + in enumerate(dividends_cases)} + + loader_type = DividendsByAnnouncementDateLoader + + def setup(self, dates): + date_intervals = [ + [ + [None, '2014-01-04'], ['2014-01-05', '2014-01-09'], + ['2014-01-10', None] + ], + [ + [None, '2014-01-04'], ['2014-01-05', '2014-01-09'], + ['2014-01-10', None] + ], + [ + [None, '2014-01-04'], ['2014-01-05', '2014-01-14'], + ['2014-01-15', None] + ], + [ + [None, '2014-01-04'], ['2014-01-05', None] + ] + ] + announcement_dates = [['NaT', '2014-01-04', '2014-01-09'], + ['NaT', '2014-01-04', '2014-01-09'], + ['NaT', '2014-01-04', '2014-01-14'], + ['NaT', '2014-01-04']] + amounts = [['NaN', 1, 15], ['NaN', 7, 13], ['NaN', 3, 1], ['NaN', 23]] + cols = {} + cols[PREVIOUS_ANNOUNCEMENT] = get_vals_for_dates( + zip_with_dates, announcement_dates, date_intervals, dates + ) + + cols[PREVIOUS_AMOUNT] = get_vals_for_dates( + zip_with_floats, amounts, date_intervals, dates + ) + + cols[ + DAYS_SINCE_PREV_DIVIDEND_ANNOUNCEMENT + ] = self._compute_busday_offsets(cols[PREVIOUS_ANNOUNCEMENT]) + return cols + + +class BlazeDividendsByAnnouncementDateTestCase( + DividendsByAnnouncementDateTestCase +): + loader_type = BlazeDividendsByAnnouncementDateLoader + + def pipeline_event_loader_args(self, dates): + _, mapping = super( + BlazeDividendsByAnnouncementDateTestCase, + self, + ).pipeline_event_loader_args(dates) + return (bz.Data(pd.concat( + pd.DataFrame({ + ANNOUNCEMENT_FIELD_NAME: df[ANNOUNCEMENT_FIELD_NAME], + TS_FIELD_NAME: df[TS_FIELD_NAME], + SID_FIELD_NAME: sid, + CASH_AMOUNT_FIELD_NAME: df[CASH_AMOUNT_FIELD_NAME] + }) + for sid, df in iteritems(mapping) + ).reset_index(drop=True)),) + + +class BlazeDividendsByAnnouncementDateNotInteractiveTestCase( + BlazeDividendsByAnnouncementDateTestCase): + """Test case for passing a non-interactive symbol and a dict of resources. + """ + + def pipeline_event_loader_args(self, dates): + (bound_expr,) = super( + BlazeDividendsByAnnouncementDateNotInteractiveTestCase, + self, + ).pipeline_event_loader_args(dates) + return swap_resources_into_scope(bound_expr, {}) + + +class DividendsByExDateTestCase(WithPipelineEventDataLoader, ZiplineTestCase): + """ + Tests for loading the dividends by ex date data. + """ + pipeline_columns = { + NEXT_EX_DATE: DividendsByExDate.next_date.latest, + PREVIOUS_EX_DATE: DividendsByExDate.previous_date.latest, + NEXT_AMOUNT: DividendsByExDate.next_amount.latest, + PREVIOUS_AMOUNT: DividendsByExDate.previous_amount.latest, + DAYS_TO_NEXT_EX_DATE: BusinessDaysUntilNextExDate(), + DAYS_SINCE_PREV_EX_DATE: BusinessDaysSincePreviousExDate() + } + + @classmethod + def get_dataset(cls): + return {sid: + frame.drop([ANNOUNCEMENT_FIELD_NAME, + PAY_DATE_FIELD_NAME], axis=1) + for sid, frame + in enumerate(dividends_cases)} + + loader_type = DividendsByExDateLoader + + def setup(self, dates): + cols = {} + cols[NEXT_EX_DATE] = get_vals_for_dates( + zip_with_dates, next_ex_and_pay_dates, next_date_intervals, dates, + ) + + cols[PREVIOUS_EX_DATE] = get_vals_for_dates( + zip_with_dates, prev_ex_and_pay_dates, prev_date_intervals, dates + ) + + cols[NEXT_AMOUNT] = get_vals_for_dates( + zip_with_floats, next_amounts, next_date_intervals, dates + ) + + cols[PREVIOUS_AMOUNT] = get_vals_for_dates( + zip_with_floats, prev_amounts, prev_date_intervals, dates + ) + + cols[DAYS_TO_NEXT_EX_DATE] = self._compute_busday_offsets( + cols[NEXT_EX_DATE] + ) + + cols[DAYS_SINCE_PREV_EX_DATE] = self._compute_busday_offsets( + cols[PREVIOUS_EX_DATE] + ) + return cols + + +class BlazeDividendsByExDateLoaderTestCase(DividendsByExDateTestCase): + loader_type = BlazeDividendsByExDateLoader + + def pipeline_event_loader_args(self, dates): + _, mapping = super( + BlazeDividendsByExDateLoaderTestCase, + self, + ).pipeline_event_loader_args(dates) + return (bz.Data(pd.concat( + pd.DataFrame({ + EX_DATE_FIELD_NAME: df[EX_DATE_FIELD_NAME], + TS_FIELD_NAME: df[TS_FIELD_NAME], + SID_FIELD_NAME: sid, + CASH_AMOUNT_FIELD_NAME: df[CASH_AMOUNT_FIELD_NAME] + }) + for sid, df in iteritems(mapping) + ).reset_index(drop=True)),) + + +class BlazeDividendsByExDateLoaderNotInteractiveTestCase( + BlazeDividendsByExDateLoaderTestCase): + """Test case for passing a non-interactive symbol and a dict of resources. + """ + + def pipeline_event_loader_args(self, dates): + (bound_expr,) = super( + BlazeDividendsByExDateLoaderNotInteractiveTestCase, + self, + ).pipeline_event_loader_args(dates) + return swap_resources_into_scope(bound_expr, {}) + + +class DividendsByPayDateTestCase(WithPipelineEventDataLoader, ZiplineTestCase): + """ + Tests for loading the dividends by pay date data. + """ + pipeline_columns = { + NEXT_PAY_DATE: DividendsByPayDate.next_date.latest, + PREVIOUS_PAY_DATE: DividendsByPayDate.previous_date.latest, + NEXT_AMOUNT: DividendsByPayDate.next_amount.latest, + PREVIOUS_AMOUNT: DividendsByPayDate.previous_amount.latest, + } + + @classmethod + def get_dataset(cls): + return {sid: + frame.drop([ANNOUNCEMENT_FIELD_NAME, + EX_DATE_FIELD_NAME], axis=1) + for sid, frame + in enumerate(dividends_cases)} + + loader_type = DividendsByPayDateLoader + + def setup(self, dates): + cols = {} + cols[NEXT_PAY_DATE] = get_vals_for_dates( + zip_with_dates, next_ex_and_pay_dates, next_date_intervals, dates + ) + cols[PREVIOUS_PAY_DATE] = get_vals_for_dates( + zip_with_dates, prev_ex_and_pay_dates, prev_date_intervals, dates + ) + + cols[NEXT_AMOUNT] = get_vals_for_dates( + zip_with_floats, next_amounts, next_date_intervals, dates + ) + + cols[PREVIOUS_AMOUNT] = get_vals_for_dates( + zip_with_floats, prev_amounts, prev_date_intervals, dates + ) + return cols + + +class BlazeDividendsByPayDateLoaderTestCase(DividendsByPayDateTestCase): + loader_type = BlazeDividendsByPayDateLoader + + def pipeline_event_loader_args(self, dates): + _, mapping = super( + BlazeDividendsByPayDateLoaderTestCase, + self, + ).pipeline_event_loader_args(dates) + return (bz.Data(pd.concat( + pd.DataFrame({ + PAY_DATE_FIELD_NAME: df[PAY_DATE_FIELD_NAME], + TS_FIELD_NAME: df[TS_FIELD_NAME], + SID_FIELD_NAME: sid, + CASH_AMOUNT_FIELD_NAME: df[CASH_AMOUNT_FIELD_NAME] + }) + for sid, df in iteritems(mapping) + ).reset_index(drop=True)),) + + +class BlazeDividendsByPayDateLoaderNotInteractiveTestCase( + BlazeDividendsByPayDateLoaderTestCase): + """Test case for passing a non-interactive symbol and a dict of resources. + """ + + def pipeline_event_loader_args(self, dates): + (bound_expr,) = super( + BlazeDividendsByPayDateLoaderNotInteractiveTestCase, + self, + ).pipeline_event_loader_args(dates) + return swap_resources_into_scope(bound_expr, {}) diff --git a/tests/pipeline/test_earnings.py b/tests/pipeline/test_earnings.py index 18ec67ca..43f502c2 100644 --- a/tests/pipeline/test_earnings.py +++ b/tests/pipeline/test_earnings.py @@ -1,15 +1,10 @@ """ Tests for the reference loader for EarningsCalendar. """ -from functools import partial -from unittest import TestCase - import blaze as bz from blaze.compute.core import swap_resources_into_scope -from contextlib2 import ExitStack import pandas as pd from six import iteritems -from .base import EventLoaderCommonMixin from zipline.pipeline.common import ( ANNOUNCEMENT_FIELD_NAME, @@ -26,11 +21,16 @@ from zipline.pipeline.factors.events import ( BusinessDaysUntilNextEarnings, ) from zipline.pipeline.loaders.earnings import EarningsCalendarLoader -from zipline.pipeline.loaders.blaze import ( - BlazeEarningsCalendarLoader, +from zipline.pipeline.loaders.blaze import BlazeEarningsCalendarLoader +from zipline.pipeline.loaders.utils import ( + get_values_for_date_ranges, + zip_with_dates ) -from zipline.testing import tmp_asset_finder +from zipline.testing.fixtures import ( + WithPipelineEventDataLoader, + ZiplineTestCase +) earnings_cases = [ # K1--K2--A1--A2. @@ -60,8 +60,61 @@ earnings_cases = [ ), ] +next_date_intervals = [ + [[None, '2014-01-04'], + ['2014-01-05', '2014-01-15'], + ['2014-01-16', '2014-01-20'], + ['2014-01-21', None]], + [[None, '2014-01-04'], + ['2014-01-05', '2014-01-09'], + ['2014-01-10', '2014-01-15'], + ['2014-01-16', '2014-01-20'], + ['2014-01-21', None]], + [[None, '2014-01-04'], + ['2014-01-05', '2014-01-10'], + ['2014-01-11', '2014-01-14'], + ['2014-01-15', '2014-01-20'], + ['2014-01-21', None]], + [[None, '2014-01-04'], + ['2014-01-05', '2014-01-10'], + ['2014-01-11', '2014-01-15'], + ['2014-01-16', None]] +] -class EarningsCalendarLoaderTestCase(TestCase, EventLoaderCommonMixin): +next_dates = [ + ['NaT', '2014-01-15', '2014-01-20', 'NaT'], + ['NaT', '2014-01-20', '2014-01-15', '2014-01-20', 'NaT'], + ['NaT', '2014-01-10', 'NaT', '2014-01-20', 'NaT'], + ['NaT', '2014-01-10', '2014-01-15', 'NaT'], + ['NaT'] +] + +prev_date_intervals = [ + [[None, '2014-01-14'], + ['2014-01-15', '2014-01-19'], + ['2014-01-20', None]], + [[None, '2014-01-14'], + ['2014-01-15', '2014-01-19'], + ['2014-01-20', None]], + [[None, '2014-01-09'], + ['2014-01-10', '2014-01-19'], + ['2014-01-20', None]], + [[None, '2014-01-09'], + ['2014-01-10', '2014-01-14'], + ['2014-01-15', None]] +] + +prev_dates = [ + ['NaT', '2014-01-15', '2014-01-20'], + ['NaT', '2014-01-15', '2014-01-20'], + ['NaT', '2014-01-10', '2014-01-20'], + ['NaT', '2014-01-10', '2014-01-15'], + ['NaT'] +] + + +class EarningsCalendarLoaderTestCase(WithPipelineEventDataLoader, + ZiplineTestCase): """ Tests for loading the earnings announcement data. """ @@ -73,111 +126,53 @@ class EarningsCalendarLoaderTestCase(TestCase, EventLoaderCommonMixin): } @classmethod - def get_sids(cls): - return range(5) + def get_dataset(cls): + return {sid: df for sid, df in enumerate(earnings_cases)} - @classmethod - def setUpClass(cls): - cls._cleanup_stack = stack = ExitStack() - cls.cols = {} - cls.dataset = {sid: df for sid, df in enumerate(earnings_cases)} - cls.finder = stack.enter_context( - tmp_asset_finder(equities=cls.get_equity_info()), - ) - - cls.loader_type = EarningsCalendarLoader + loader_type = EarningsCalendarLoader def get_expected_next_event_dates(self, dates): - num_days_between_for_dates = partial(self.num_days_between, dates) - zip_with_dates_for_dates = partial(self.zip_with_dates, dates) return pd.DataFrame({ - 0: zip_with_dates_for_dates( - ['NaT'] * - num_days_between_for_dates(None, '2014-01-04') + - ['2014-01-15'] * - num_days_between_for_dates('2014-01-05', '2014-01-15') + - ['2014-01-20'] * - num_days_between_for_dates('2014-01-16', '2014-01-20') + - ['NaT'] * - num_days_between_for_dates('2014-01-21', None) - ), - 1: zip_with_dates_for_dates( - ['NaT'] * - num_days_between_for_dates(None, '2014-01-04') + - ['2014-01-20'] * - num_days_between_for_dates('2014-01-05', '2014-01-09') + - ['2014-01-15'] * - num_days_between_for_dates('2014-01-10', '2014-01-15') + - ['2014-01-20'] * - num_days_between_for_dates('2014-01-16', '2014-01-20') + - ['NaT'] * - num_days_between_for_dates('2014-01-21', None) - ), - 2: zip_with_dates_for_dates( - ['NaT'] * - num_days_between_for_dates(None, '2014-01-04') + - ['2014-01-10'] * - num_days_between_for_dates('2014-01-05', '2014-01-10') + - ['NaT'] * - num_days_between_for_dates('2014-01-11', '2014-01-14') + - ['2014-01-20'] * - num_days_between_for_dates('2014-01-15', '2014-01-20') + - ['NaT'] * - num_days_between_for_dates('2014-01-21', None) - ), - 3: zip_with_dates_for_dates( - ['NaT'] * - num_days_between_for_dates(None, '2014-01-04') + - ['2014-01-10'] * - num_days_between_for_dates('2014-01-05', '2014-01-10') + - ['2014-01-15'] * - num_days_between_for_dates('2014-01-11', '2014-01-15') + - ['NaT'] * - num_days_between_for_dates('2014-01-16', None) - ), - 4: zip_with_dates_for_dates(['NaT'] * - len(dates)), + 0: get_values_for_date_ranges(zip_with_dates, + next_dates[0], + next_date_intervals[0], + dates), + 1: get_values_for_date_ranges(zip_with_dates, + next_dates[1], + next_date_intervals[1], + dates), + 2: get_values_for_date_ranges(zip_with_dates, + next_dates[2], + next_date_intervals[2], + dates), + 3: get_values_for_date_ranges(zip_with_dates, + next_dates[3], + next_date_intervals[3], + dates), + 4: zip_with_dates(dates, ['NaT'] * len(dates)), }, index=dates) def get_expected_previous_event_dates(self, dates): - num_days_between_for_dates = partial(self.num_days_between, dates) - zip_with_dates_for_dates = partial(self.zip_with_dates, dates) return pd.DataFrame({ - 0: zip_with_dates_for_dates( - ['NaT'] * num_days_between_for_dates(None, '2014-01-14') + - ['2014-01-15'] * num_days_between_for_dates('2014-01-15', - '2014-01-19') + - ['2014-01-20'] * num_days_between_for_dates('2014-01-20', - None), - ), - 1: zip_with_dates_for_dates( - ['NaT'] * num_days_between_for_dates(None, '2014-01-14') + - ['2014-01-15'] * num_days_between_for_dates('2014-01-15', - '2014-01-19') + - ['2014-01-20'] * num_days_between_for_dates('2014-01-20', - None), - ), - 2: zip_with_dates_for_dates( - ['NaT'] * num_days_between_for_dates(None, '2014-01-09') + - ['2014-01-10'] * num_days_between_for_dates('2014-01-10', - '2014-01-19') + - ['2014-01-20'] * num_days_between_for_dates('2014-01-20', - None), - ), - 3: zip_with_dates_for_dates( - ['NaT'] * num_days_between_for_dates(None, '2014-01-09') + - ['2014-01-10'] * num_days_between_for_dates('2014-01-10', - '2014-01-14') + - ['2014-01-15'] * num_days_between_for_dates('2014-01-15', - None), - ), - 4: zip_with_dates_for_dates(['NaT'] * len(dates)), + 0: get_values_for_date_ranges(zip_with_dates, + prev_dates[0], + prev_date_intervals[0], + dates), + 1: get_values_for_date_ranges(zip_with_dates, + prev_dates[1], + prev_date_intervals[1], + dates), + 2: get_values_for_date_ranges(zip_with_dates, + prev_dates[2], + prev_date_intervals[2], + dates), + 3: get_values_for_date_ranges(zip_with_dates, + prev_dates[3], + prev_date_intervals[3], + dates), + 4: zip_with_dates(dates, ['NaT'] * len(dates)), }, index=dates) - @classmethod - def tearDownClass(cls): - cls._cleanup_stack.close() - def setup(self, dates): _expected_next_announce = self.get_expected_next_event_dates(dates) @@ -191,23 +186,22 @@ class EarningsCalendarLoaderTestCase(TestCase, EventLoaderCommonMixin): _expected_previous_busday_offsets = self._compute_busday_offsets( _expected_previous_announce ) - self.cols[PREVIOUS_ANNOUNCEMENT] = _expected_previous_announce - self.cols[NEXT_ANNOUNCEMENT] = _expected_next_announce - self.cols[DAYS_TO_NEXT] = _expected_next_busday_offsets - self.cols[DAYS_SINCE_PREV] = _expected_previous_busday_offsets + cols = {} + cols[PREVIOUS_ANNOUNCEMENT] = _expected_previous_announce + cols[NEXT_ANNOUNCEMENT] = _expected_next_announce + cols[DAYS_TO_NEXT] = _expected_next_busday_offsets + cols[DAYS_SINCE_PREV] = _expected_previous_busday_offsets + return cols class BlazeEarningsCalendarLoaderTestCase(EarningsCalendarLoaderTestCase): - @classmethod - def setUpClass(cls): - super(BlazeEarningsCalendarLoaderTestCase, cls).setUpClass() - cls.loader_type = BlazeEarningsCalendarLoader + loader_type = BlazeEarningsCalendarLoader - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): _, mapping = super( BlazeEarningsCalendarLoaderTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return (bz.data(pd.concat( pd.DataFrame({ ANNOUNCEMENT_FIELD_NAME: df[ANNOUNCEMENT_FIELD_NAME], @@ -222,15 +216,10 @@ class BlazeEarningsCalendarLoaderNotInteractiveTestCase( BlazeEarningsCalendarLoaderTestCase): """Test case for passing a non-interactive symbol and a dict of resources. """ - @classmethod - def setUpClass(cls): - super(BlazeEarningsCalendarLoaderNotInteractiveTestCase, - cls).setUpClass() - cls.loader_type = BlazeEarningsCalendarLoader - def loader_args(self, dates): + def pipeline_event_loader_args(self, dates): (bound_expr,) = super( BlazeEarningsCalendarLoaderNotInteractiveTestCase, self, - ).loader_args(dates) + ).pipeline_event_loader_args(dates) return swap_resources_into_scope(bound_expr, {}) diff --git a/zipline/pipeline/common.py b/zipline/pipeline/common.py index de225409..aa71d3a9 100644 --- a/zipline/pipeline/common.py +++ b/zipline/pipeline/common.py @@ -4,14 +4,26 @@ Common constants for Pipeline. AD_FIELD_NAME = 'asof_date' ANNOUNCEMENT_FIELD_NAME = 'announcement_date' CASH_FIELD_NAME = 'cash' +CASH_AMOUNT_FIELD_NAME = 'cash_amount' BUYBACK_ANNOUNCEMENT_FIELD_NAME = 'buyback_date' DAYS_SINCE_PREV = 'days_since_prev' +DAYS_SINCE_PREV_DIVIDEND_ANNOUNCEMENT = 'days_since_prev_dividend_announcement' +DAYS_SINCE_PREV_EX_DATE = 'days_since_prev_ex_date' DAYS_TO_NEXT = 'days_to_next' +DAYS_TO_NEXT_EX_DATE = 'days_to_next_ex_date' +EX_DATE_FIELD_NAME = 'ex_date' +NEXT_AMOUNT = 'next_amount' NEXT_ANNOUNCEMENT = 'next_announcement' +NEXT_EX_DATE = 'next_ex_date' +NEXT_PAY_DATE = 'next_pay_date' +PAY_DATE_FIELD_NAME = 'pay_date' +PREVIOUS_AMOUNT = 'previous_amount' PREVIOUS_ANNOUNCEMENT = 'previous_announcement' PREVIOUS_BUYBACK_ANNOUNCEMENT = 'previous_buyback_announcement' PREVIOUS_BUYBACK_CASH = 'previous_buyback_cash' PREVIOUS_BUYBACK_SHARE_COUNT = 'previous_buyback_share_count' +PREVIOUS_EX_DATE = 'previous_ex_date' +PREVIOUS_PAY_DATE = 'previous_pay_date' SHARE_COUNT_FIELD_NAME = 'share_count' SID_FIELD_NAME = 'sid' TS_FIELD_NAME = 'timestamp' diff --git a/zipline/pipeline/data/dividends.py b/zipline/pipeline/data/dividends.py new file mode 100644 index 00000000..176aac43 --- /dev/null +++ b/zipline/pipeline/data/dividends.py @@ -0,0 +1,25 @@ +""" +Dataset representing dates of upcoming dividends. +""" +from zipline.utils.numpy_utils import datetime64ns_dtype, float64_dtype + +from .dataset import Column, DataSet + + +class DividendsByExDate(DataSet): + next_date = Column(datetime64ns_dtype) + previous_date = Column(datetime64ns_dtype) + next_amount = Column(float64_dtype) + previous_amount = Column(float64_dtype) + + +class DividendsByPayDate(DataSet): + next_date = Column(datetime64ns_dtype) + previous_date = Column(datetime64ns_dtype) + next_amount = Column(float64_dtype) + previous_amount = Column(float64_dtype) + + +class DividendsByAnnouncementDate(DataSet): + previous_announcement_date = Column(datetime64ns_dtype) + previous_amount = Column(float64_dtype) diff --git a/zipline/pipeline/factors/events.py b/zipline/pipeline/factors/events.py index a64fe3e2..cdaa5573 100644 --- a/zipline/pipeline/factors/events.py +++ b/zipline/pipeline/factors/events.py @@ -7,6 +7,10 @@ from zipline.pipeline.data.buyback_auth import ( CashBuybackAuthorizations, ShareBuybackAuthorizations ) +from zipline.pipeline.data.dividends import ( + DividendsByAnnouncementDate, + DividendsByExDate +) from zipline.pipeline.data.earnings import EarningsCalendar from zipline.utils.numpy_utils import ( NaTD, @@ -156,3 +160,48 @@ class BusinessDaysSinceShareBuybackAuth( zipline.pipeline.factors.BusinessDaysSinceShareBuybackAuth """ inputs = [ShareBuybackAuthorizations.announcement_date] + + +class BusinessDaysSinceDividendAnnouncement( + BusinessDaysSincePreviousEvents +): + """ + Factor returning the number of **business days** (not trading days!) since + the most recent dividend announcement for each asset. + + + See Also + -------- + zipline.pipeline.factors.BusinessDaysSinceDividendAnnouncement + """ + inputs = [DividendsByAnnouncementDate.previous_announcement_date] + + +class BusinessDaysUntilNextExDate( + BusinessDaysUntilNextEvents +): + """ + Factor returning the number of **business days** (not trading days!) until + the next ex date for each asset. + + + See Also + -------- + zipline.pipeline.factors.BusinessDaysSinceDividendAnnouncement + """ + inputs = [DividendsByExDate.next_date] + + +class BusinessDaysSincePreviousExDate( + BusinessDaysSincePreviousEvents +): + """ + Factor returning the number of **business days** (not trading days!) since + the most recent ex date for each asset. + + + See Also + -------- + zipline.pipeline.factors.BusinessDaysSinceDividendAnnouncement + """ + inputs = [DividendsByExDate.previous_date] diff --git a/zipline/pipeline/loaders/blaze/buyback_auth.py b/zipline/pipeline/loaders/blaze/buyback_auth.py index 44b35c3e..9a1acfc9 100644 --- a/zipline/pipeline/loaders/blaze/buyback_auth.py +++ b/zipline/pipeline/loaders/blaze/buyback_auth.py @@ -68,24 +68,7 @@ class BlazeCashBuybackAuthorizationsLoader(BlazeEventsLoader): }) concrete_loader = CashBuybackAuthorizationsLoader - - def __init__(self, - expr, - resources=None, - odo_kwargs=None, - data_query_time=None, - data_query_tz=None, - dataset=CashBuybackAuthorizations, - **kwargs): - super( - BlazeCashBuybackAuthorizationsLoader, self - ).__init__(expr, - resources=resources, - odo_kwargs=odo_kwargs, - data_query_time=data_query_time, - data_query_tz=data_query_tz, - dataset=dataset, - **kwargs) + default_dataset = CashBuybackAuthorizations class BlazeShareBuybackAuthorizationsLoader(BlazeEventsLoader): @@ -140,21 +123,4 @@ class BlazeShareBuybackAuthorizationsLoader(BlazeEventsLoader): }) concrete_loader = ShareBuybackAuthorizationsLoader - - def __init__(self, - expr, - resources=None, - odo_kwargs=None, - data_query_time=None, - data_query_tz=None, - dataset=ShareBuybackAuthorizations, - **kwargs): - super( - BlazeShareBuybackAuthorizationsLoader, self - ).__init__(expr, - resources=resources, - odo_kwargs=odo_kwargs, - data_query_time=data_query_time, - data_query_tz=data_query_tz, - dataset=dataset, - **kwargs) + default_dataset = ShareBuybackAuthorizations diff --git a/zipline/pipeline/loaders/blaze/dividends.py b/zipline/pipeline/loaders/blaze/dividends.py new file mode 100644 index 00000000..7bba9e05 --- /dev/null +++ b/zipline/pipeline/loaders/blaze/dividends.py @@ -0,0 +1,187 @@ +from zipline.pipeline.common import ( + ANNOUNCEMENT_FIELD_NAME, + CASH_AMOUNT_FIELD_NAME, + EX_DATE_FIELD_NAME, + PAY_DATE_FIELD_NAME, + SID_FIELD_NAME, + TS_FIELD_NAME, +) +from zipline.pipeline.data.dividends import ( + DividendsByExDate, + DividendsByAnnouncementDate, + DividendsByPayDate +) +from zipline.pipeline.loaders.dividends import ( + DividendsByAnnouncementDateLoader, + DividendsByPayDateLoader, + DividendsByExDateLoader +) +from .events import BlazeEventsLoader + + +class BlazeDividendsByAnnouncementDateLoader(BlazeEventsLoader): + """A pipeline loader for the ``DividendsByAnnouncementDate`` dataset that + loads data from a blaze expression. + + Parameters + ---------- + expr : Expr + The expression representing the data to load. + resources : dict, optional + Mapping from the atomic 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, + {CASH_AMOUNT_FIELD_NAME}: ?datetime, + {ANNOUNCEMENT_FIELD_NAME}: ?datetime, + }} + + Where each row of the table is a record including the sid to identify the + company, the timestamp where we learned about the announcement, the + date when the dividends will be announced, and the cash amount. + + If the '{TS_FIELD_NAME}' field is not included it is assumed that we + start the backtest with knowledge of all announcements. + """ + + __doc__ = __doc__.format( + TS_FIELD_NAME=TS_FIELD_NAME, + SID_FIELD_NAME=SID_FIELD_NAME, + CASH_AMOUNT_FIELD_NAME=CASH_AMOUNT_FIELD_NAME, + ANNOUNCEMENT_FIELD_NAME=ANNOUNCEMENT_FIELD_NAME + ) + + _expected_fields = frozenset({ + TS_FIELD_NAME, + SID_FIELD_NAME, + CASH_AMOUNT_FIELD_NAME, + ANNOUNCEMENT_FIELD_NAME + }) + + concrete_loader = DividendsByAnnouncementDateLoader + default_dataset = DividendsByAnnouncementDate + + +class BlazeDividendsByExDateLoader(BlazeEventsLoader): + """A pipeline loader for the ``DividendsByExDate`` dataset that loads + data from a blaze expression. + + Parameters + ---------- + expr : Expr + The expression representing the data to load. + resources : dict, optional + Mapping from the atomic 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, + {EX_DATE_FIELD_NAME}: ?datetime, + {CASH_AMOUNT_FIELD_NAME}: ?datetime, + }} + + Where each row of the table is a record including the sid to identify the + company, the timestamp where we learned about the ex date, the + ex date, and the associated cash amount. + + If the '{TS_FIELD_NAME}' field is not included it is assumed that we + start the backtest with knowledge of all announcements. + """ + + __doc__ = __doc__.format( + TS_FIELD_NAME=TS_FIELD_NAME, + SID_FIELD_NAME=SID_FIELD_NAME, + EX_DATE_FIELD_NAME=EX_DATE_FIELD_NAME, + CASH_AMOUNT_FIELD_NAME=CASH_AMOUNT_FIELD_NAME, + ) + + _expected_fields = frozenset({ + TS_FIELD_NAME, + SID_FIELD_NAME, + EX_DATE_FIELD_NAME, + CASH_AMOUNT_FIELD_NAME, + }) + + concrete_loader = DividendsByExDateLoader + default_dataset = DividendsByExDate + + +class BlazeDividendsByPayDateLoader(BlazeEventsLoader): + """A pipeline loader for the ``DividendsByPayDate`` dataset that loads + data from a blaze expression. + + Parameters + ---------- + expr : Expr + The expression representing the data to load. + resources : dict, optional + Mapping from the atomic 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, + {PAY_DATE_FIELD_NAME}: ?datetime, + {CASH_AMOUNT_FIELD_NAME}: ?datetime, + }} + + Where each row of the table is a record including the sid to identify the + company, the timestamp where we learned about the pay date, the pay date, + and the associated cash amount. + + If the '{TS_FIELD_NAME}' field is not included it is assumed that we + start the backtest with knowledge of all announcements. + """ + + __doc__ = __doc__.format( + TS_FIELD_NAME=TS_FIELD_NAME, + SID_FIELD_NAME=SID_FIELD_NAME, + PAY_DATE_FIELD_NAME=PAY_DATE_FIELD_NAME, + CASH_AMOUNT_FIELD_NAME=CASH_AMOUNT_FIELD_NAME, + ) + + _expected_fields = frozenset({ + TS_FIELD_NAME, + SID_FIELD_NAME, + PAY_DATE_FIELD_NAME, + CASH_AMOUNT_FIELD_NAME, + }) + + concrete_loader = DividendsByPayDateLoader + default_dataset = DividendsByPayDate diff --git a/zipline/pipeline/loaders/blaze/earnings.py b/zipline/pipeline/loaders/blaze/earnings.py index 17e76c63..9d7e9b5c 100644 --- a/zipline/pipeline/loaders/blaze/earnings.py +++ b/zipline/pipeline/loaders/blaze/earnings.py @@ -58,18 +58,4 @@ class BlazeEarningsCalendarLoader(BlazeEventsLoader): }) concrete_loader = EarningsCalendarLoader - - def __init__(self, - expr, - resources=None, - odo_kwargs=None, - data_query_time=None, - data_query_tz=None, - dataset=EarningsCalendar, - **kwargs): - super( - BlazeEarningsCalendarLoader, self - ).__init__(expr, dataset=dataset, - resources=resources, odo_kwargs=odo_kwargs, - data_query_time=data_query_time, - data_query_tz=data_query_tz, **kwargs) + default_dataset = EarningsCalendar diff --git a/zipline/pipeline/loaders/blaze/events.py b/zipline/pipeline/loaders/blaze/events.py index 326c97ad..6fae418f 100644 --- a/zipline/pipeline/loaders/blaze/events.py +++ b/zipline/pipeline/loaders/blaze/events.py @@ -56,6 +56,7 @@ class BlazeEventsLoader(PipelineLoader): If the '{TS_FIELD_NAME}' field is not included it is assumed that we start the backtest with knowledge of all announcements. """ + default_dataset = None @preprocess(data_query_tz=optionally(ensure_timezone)) def __init__(self, @@ -64,7 +65,10 @@ class BlazeEventsLoader(PipelineLoader): odo_kwargs=None, data_query_time=None, data_query_tz=None, - dataset=None): + dataset=default_dataset): + if dataset is None: + dataset = self.default_dataset + dshape = expr.dshape if not istabular(dshape): diff --git a/zipline/pipeline/loaders/dividends.py b/zipline/pipeline/loaders/dividends.py new file mode 100644 index 00000000..16a4585e --- /dev/null +++ b/zipline/pipeline/loaders/dividends.py @@ -0,0 +1,116 @@ +from zipline.pipeline.common import ( + EX_DATE_FIELD_NAME, + PAY_DATE_FIELD_NAME, + CASH_AMOUNT_FIELD_NAME, + ANNOUNCEMENT_FIELD_NAME +) +from zipline.pipeline.loaders.events import EventsLoader +from zipline.pipeline.data.dividends import ( + DividendsByExDate, + DividendsByAnnouncementDate, + DividendsByPayDate +) +from zipline.utils.memoize import lazyval + + +class DividendsByAnnouncementDateLoader(EventsLoader): + expected_cols = frozenset([ANNOUNCEMENT_FIELD_NAME, + CASH_AMOUNT_FIELD_NAME]) + + def __init__(self, all_dates, events_by_sid, + infer_timestamps=False, + dataset=DividendsByAnnouncementDate): + super(DividendsByAnnouncementDateLoader, self).__init__( + all_dates, events_by_sid, infer_timestamps, dataset=dataset, + ) + + @lazyval + def previous_announcement_date_loader(self): + return self._previous_event_date_loader( + self.dataset.previous_announcement_date, + ANNOUNCEMENT_FIELD_NAME + ) + + @lazyval + def previous_amount_loader(self): + return self._previous_event_value_loader( + self.dataset.previous_amount, + ANNOUNCEMENT_FIELD_NAME, + CASH_AMOUNT_FIELD_NAME + ) + + +class DividendsByPayDateLoader(EventsLoader): + expected_cols = frozenset([PAY_DATE_FIELD_NAME, + CASH_AMOUNT_FIELD_NAME]) + + def __init__(self, all_dates, events_by_sid, + infer_timestamps=False, + dataset=DividendsByPayDate): + super(DividendsByPayDateLoader, self).__init__( + all_dates, events_by_sid, infer_timestamps, dataset=dataset, + ) + + @lazyval + def next_date_loader(self): + return self._next_event_date_loader(self.dataset.next_date, + PAY_DATE_FIELD_NAME) + + @lazyval + def previous_date_loader(self): + return self._previous_event_date_loader( + self.dataset.previous_date, + PAY_DATE_FIELD_NAME + ) + + @lazyval + def next_amount_loader(self): + return self._next_event_value_loader(self.dataset.next_amount, + PAY_DATE_FIELD_NAME, + CASH_AMOUNT_FIELD_NAME) + + @lazyval + def previous_amount_loader(self): + return self._previous_event_value_loader( + self.dataset.previous_amount, + PAY_DATE_FIELD_NAME, + CASH_AMOUNT_FIELD_NAME + ) + + +class DividendsByExDateLoader(EventsLoader): + expected_cols = frozenset([EX_DATE_FIELD_NAME, + CASH_AMOUNT_FIELD_NAME]) + + def __init__(self, all_dates, events_by_sid, + infer_timestamps=False, + dataset=DividendsByExDate): + super(DividendsByExDateLoader, self).__init__( + all_dates, events_by_sid, infer_timestamps, dataset=dataset, + ) + + @lazyval + def next_date_loader(self): + return self._next_event_date_loader(self.dataset.next_date, + EX_DATE_FIELD_NAME) + + @lazyval + def previous_date_loader(self): + return self._previous_event_date_loader( + self.dataset.previous_date, + EX_DATE_FIELD_NAME + ) + + @lazyval + def next_amount_loader(self): + return self._next_event_value_loader(self.dataset.next_amount, + EX_DATE_FIELD_NAME, + CASH_AMOUNT_FIELD_NAME) + + @lazyval + def previous_amount_loader(self): + return self._previous_event_value_loader( + self.dataset.previous_amount, + EX_DATE_FIELD_NAME, + CASH_AMOUNT_FIELD_NAME + ) diff --git a/zipline/pipeline/loaders/events.py b/zipline/pipeline/loaders/events.py index 27cebd50..e88d7966 100644 --- a/zipline/pipeline/loaders/events.py +++ b/zipline/pipeline/loaders/events.py @@ -5,7 +5,7 @@ from toolz import merge from .base import PipelineLoader from .frame import DataFrameLoader -from .utils import previous_event_frame, next_date_frame +from .utils import previous_event_frame, next_event_frame from zipline.pipeline.common import TS_FIELD_NAME from zipline.utils.numpy_utils import NaTD @@ -167,14 +167,34 @@ class EventsLoader(PipelineLoader): def _next_event_date_loader(self, next_date_field, event_date_field_name): return DataFrameLoader( next_date_field, - next_date_frame( - self.all_dates, + next_event_frame( self.events_by_sid, + self.all_dates, + next_date_field.missing_value, + next_date_field.dtype, + event_date_field_name, event_date_field_name ), adjustments=None, ) + def _next_event_value_loader(self, + next_value_field, + event_date_field_name, + value_field_name): + return DataFrameLoader( + next_value_field, + next_event_frame( + self.events_by_sid, + self.all_dates, + next_value_field.missing_value, + next_value_field.dtype, + event_date_field_name, + value_field_name + ), + adjustments=None, + ) + def _previous_event_date_loader(self, prev_date_field, event_date_field_name): diff --git a/zipline/pipeline/loaders/utils.py b/zipline/pipeline/loaders/utils.py index 8b2cc9a2..3b79253f 100644 --- a/zipline/pipeline/loaders/utils.py +++ b/zipline/pipeline/loaders/utils.py @@ -8,9 +8,15 @@ from six.moves import zip from zipline.utils.numpy_utils import NaTns -def next_date_frame(dates, events_by_sid, event_date_field_name): +def next_event_frame(events_by_sid, + dates, + missing_value, + field_dtype, + event_date_field_name, + return_field_name): """ - Make a DataFrame representing the simulated next known date for an event. + Make a DataFrame representing the simulated next known dates or values + for an event. Parameters ---------- @@ -36,28 +42,36 @@ def next_date_frame(dates, events_by_sid, event_date_field_name): -------- previous_date_frame """ - cols = { + date_cols = { equity: np.full_like(dates, NaTns) for equity in events_by_sid } + value_cols = { + equity: np.full(len(dates), missing_value, dtype=field_dtype) + for equity in events_by_sid + } + raw_dates = dates.values for equity, df in iteritems(events_by_sid): event_dates = df[event_date_field_name] - data = cols[equity] + values = df[return_field_name] + data = date_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: + iter_date_vals = zip(event_dates.index.values, event_dates.values, + values) + for knowledge_date, event_date, value in iter_date_vals: date_mask = ( (knowledge_date <= raw_dates) & (raw_dates <= event_date) ) value_mask = (event_date <= data) | (data == NaTns) - data[date_mask & value_mask] = event_date - - return pd.DataFrame(index=dates, data=cols) + data_indices = np.where(date_mask & value_mask) + data[data_indices] = event_date + value_cols[equity][data_indices] = value + return pd.DataFrame(index=dates, data=value_cols) def previous_event_frame(events_by_sid, @@ -260,3 +274,57 @@ def check_data_query_args(data_query_time, data_query_tz): data_query_tz, ), ) + + +def zip_with_floats(dates, flts): + return pd.Series(flts, index=dates, dtype='float') + + +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 zip_with_dates(index_dates, dts): + return pd.Series(pd.to_datetime(dts), index=index_dates) + + +def get_values_for_date_ranges(zip_date_index_with_vals, + vals_for_date_intervals, + date_intervals, + date_index): + """ + Returns a Series of values indexed by date based on values for the given + date intervals. + + Parameters + ---------- + zip_date_index_with_vals : callable + A function that takes in a list of dates and a list of values and + returns a pd.Series with the values indexed by the dates. + vals_for_date_intervals : list + A list of values for each date interval in `date_intervals`. + date_intervals : list + A list of pairs of dates, where each pair represents a date interval + that corresponds to the value at the same index in + `vals_for_date_intervals`. + date_index : DatetimeIndex + The DatetimeIndex containing all dates for which values were requested. + + Returns + ------- + date_index_with_vals : pd.Series + A Series indexed by the given DatetimeIndex and with values assigned + to dates based on the given date intervals. + """ + # Fill in given values for given date ranges. + return zip_date_index_with_vals( + date_index, + np.repeat(vals_for_date_intervals, + [num_days_in_range(date_index, *date_interval) + for date_interval in + date_intervals]), + ) diff --git a/zipline/testing/__init__.py b/zipline/testing/__init__.py index 34bf9db8..cedfa74d 100644 --- a/zipline/testing/__init__.py +++ b/zipline/testing/__init__.py @@ -19,7 +19,6 @@ from .core import ( # noqa make_simple_equity_info, make_test_handler, make_trade_panel_for_asset_info, - num_days_in_range, parameter_space, permute_rows, powerset, diff --git a/zipline/testing/core.py b/zipline/testing/core.py index 3d96c4ec..1766b2d4 100644 --- a/zipline/testing/core.py +++ b/zipline/testing/core.py @@ -786,14 +786,6 @@ def to_series(knowledge_dates, 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. diff --git a/zipline/testing/fixtures.py b/zipline/testing/fixtures.py index c8b8f04c..0e792dce 100644 --- a/zipline/testing/fixtures.py +++ b/zipline/testing/fixtures.py @@ -2,13 +2,19 @@ from unittest import TestCase from contextlib2 import ExitStack from logbook import NullHandler +from nose_parameterized import parameterized +import numpy as np import pandas as pd +from pandas.util.testing import assert_series_equal from six import with_metaclass -from .core import tmp_asset_finder +from .core import tmp_asset_finder, make_simple_equity_info, gen_calendars from ..finance.trading import TradingEnvironment from ..utils import tradingcalendar, factory from ..utils.final import FinalMeta, final +from zipline.pipeline import Pipeline, SimplePipelineEngine +from zipline.utils.numpy_utils import make_datetime64D +from zipline.utils.numpy_utils import NaTD class ZiplineTestCase(with_metaclass(FinalMeta, TestCase)): @@ -292,3 +298,122 @@ class WithNYSETradingDays(object): start_loc = end_loc - cls.TRADING_DAY_COUNT cls.trading_days = all_days[start_loc:end_loc + 1] + + +class WithPipelineEventDataLoader(WithAssetFinder): + """ + ZiplineTestCase mixin providing common test methods/behaviors for event + data loaders. + + `get_sids` must return the sids being tested. + `get_dataset` must return {sid -> pd.DataFrame} + `loader_type` must return the loader class to use for loading the dataset + `make_asset_finder` returns a default asset finder which can be overridden. + """ + @classmethod + def get_sids(cls): + return range(0, 5) + + @classmethod + def get_dataset(cls): + return {sid: pd.DataFrame() for sid in cls.get_sids()} + + @classmethod + def loader_type(self): + return None + + @classmethod + def make_equities_info(cls): + return make_simple_equity_info( + cls.get_sids(), + start_date=pd.Timestamp('2013-01-01', tz='UTC'), + end_date=pd.Timestamp('2015-01-01', tz='UTC'), + ) + + def pipeline_event_loader_args(self, dates): + """Construct the base object to pass to the loader. + + Parameters + ---------- + dates : pd.DatetimeIndex + The dates we can serve. + + Returns + ------- + args : tuple[any] + The arguments to forward to the loader positionally. + """ + return dates, self.get_dataset() + + def pipeline_event_setup_engine(self, dates): + """ + Make a Pipeline Enigne object based on the given dates. + """ + loader = self.loader_type(*self.pipeline_event_loader_args(dates)) + return SimplePipelineEngine(lambda _: loader, dates, self.asset_finder) + + @staticmethod + def _compute_busday_offsets(announcement_dates): + """ + Compute expected business day offsets from a DataFrame of announcement + dates. + """ + # Column-vector of dates on which factor `compute` will be called. + raw_call_dates = announcement_dates.index.values.astype( + 'datetime64[D]' + )[:, None] + + # 2D array of dates containining expected nexg announcement. + raw_announce_dates = ( + announcement_dates.values.astype('datetime64[D]') + ) + + # Set NaTs to 0 temporarily because busday_count doesn't support NaT. + # We fill these entries with NaNs later. + whereNaT = raw_announce_dates == NaTD + raw_announce_dates[whereNaT] = make_datetime64D(0) + + # The abs call here makes it so that we can use this function to + # compute offsets for both next and previous earnings (previous + # earnings offsets come back negative). + expected = abs(np.busday_count( + raw_call_dates, + raw_announce_dates + ).astype(float)) + + expected[whereNaT] = np.nan + return pd.DataFrame( + data=expected, + columns=announcement_dates.columns, + index=announcement_dates.index, + ) + + @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', + ], utc=True), + )) + def test_compute(self, dates): + engine = self.pipeline_event_setup_engine(dates) + cols = self.setup(dates) + + pipe = Pipeline( + columns=self.pipeline_columns + ) + + result = engine.run_pipeline( + pipe, + start_date=dates[0], + end_date=dates[-1], + ) + + for sid in self.get_sids(): + for col_name in cols.keys(): + assert_series_equal(result[col_name].xs(sid, level=1), + cols[col_name][sid], + check_names=False)