wassname/catalyst
1
0
Fork 0
mirror of https://github.com/wassname/catalyst.git synced 2026-07-15 11:22:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1063 from quantopian/dividends-in-pipeline

Browse Source 
Dividends in pipeline
This commit is contained in:
Maya Tydykov
2016-03-29 14:07:41 -04:00
parent 8c059e6da6 e8185a1512
commit 9b27f3439d
17 changed files with 1229 additions and 431 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/pipeline/base.py
+1 -122
View File
@@ -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)
tests/pipeline/test_buyback_auth.py
+59 -114
View File
@@ -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, {})
tests/pipeline/test_dividends.py
+436
View File
@@ -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, {})
tests/pipeline/test_earnings.py
+110 -121
View File
@@ -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, {})
zipline/pipeline/common.py
+12
View File
@@ -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'
zipline/pipeline/data/dividends.py
+25
View File
@@ -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)
zipline/pipeline/factors/events.py
+49
View File
@@ -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]
zipline/pipeline/loaders/blaze/buyback_auth.py
+2 -36
View File
@@ -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
zipline/pipeline/loaders/blaze/dividends.py
+187
View File
@@ -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
zipline/pipeline/loaders/blaze/earnings.py
+1 -15
View File
@@ -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
zipline/pipeline/loaders/blaze/events.py
+5 -1
View File
@@ -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):
zipline/pipeline/loaders/dividends.py
+116
View File
@@ -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
)
zipline/pipeline/loaders/events.py
+23 -3
View File
@@ -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):
zipline/pipeline/loaders/utils.py
+77 -9
View File
@@ -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]),
)
zipline/testing/__init__.py
-1
View File
@@ -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,
zipline/testing/core.py
-8
View File
@@ -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.
zipline/testing/fixtures.py
+126 -1
View File
@@ -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)