MAINT: Rework event datasets.

- Refactored EventsLoader and BlazeEventsLoader to not require a
  subclass per dataset.  Instead, you now pass a map from columns to
  event fields directly to the EventsLoader constructor.

- Removed a large number of Quantopian-specific datasets and associated
  tests.

- Rewrote the core logic of EventsLoader and BlazeEventsLoader to share
  index calculations across multiple requested columns.

- Fixed a bug where event fields were incorrectly forward-filled when
  null values were present in an event.
This commit is contained in:
Scott Sanderson
2016-06-10 19:22:27 -04:00
parent 5a6b870cd5
commit bc302beec9
33 changed files with 780 additions and 3556 deletions
-115
View File
@@ -1,115 +0,0 @@
"""
Tests for the reference loader for 13d filings.
"""
import pandas as pd
from zipline.pipeline.common import TS_FIELD_NAME
from zipline.pipeline.data import _13DFilings
from zipline.pipeline.factors.events import BusinessDaysSince13DFilingsDate
from zipline.pipeline.loaders._13d_filings import (
_13DFilingsLoader,
DISCLOSURE_DATE,
NUM_SHARES,
PERCENT_SHARES,
)
from zipline.pipeline.loaders.utils import (
zip_with_floats,
zip_with_dates
)
from zipline.testing.fixtures import WithPipelineEventDataLoader
from zipline.testing.fixtures import ZiplineTestCase
DAYS_SINCE_PREV_DISCLOSURE = 'days_since_prev_disclosure'
PREVIOUS_DISCLOSURE_DATE = 'previous_disclosure_date'
PREVIOUS_NUM_SHARES = 'previous_number_shares'
PREVIOUS_PERCENT_SHARES = 'previous_percentage'
date_intervals = [
[['2014-01-01', '2014-01-04'],
['2014-01-05', '2014-01-09'],
['2014-01-10', '2014-01-31']]
]
empty_df = pd.DataFrame(
columns=[NUM_SHARES,
PERCENT_SHARES,
DISCLOSURE_DATE,
TS_FIELD_NAME],
)
empty_df[NUM_SHARES] = empty_df[NUM_SHARES].astype('float')
empty_df[PERCENT_SHARES] = empty_df[PERCENT_SHARES].astype('float')
empty_df[TS_FIELD_NAME] = empty_df[TS_FIELD_NAME].astype('datetime64[ns]')
empty_df[DISCLOSURE_DATE] = empty_df[DISCLOSURE_DATE].astype('datetime64[ns]')
_13d_filings_cases = [
pd.DataFrame({
NUM_SHARES: [1, 15],
PERCENT_SHARES: [10, 20],
TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-10']),
DISCLOSURE_DATE: pd.to_datetime(['2014-01-04', '2014-01-09'])
}),
empty_df
]
class _13DFilingsLoaderTestCase(WithPipelineEventDataLoader,
ZiplineTestCase):
"""
Test for _13_filings dataset.
"""
pipeline_columns = {
PREVIOUS_NUM_SHARES:
_13DFilings.number_shares.latest,
PREVIOUS_PERCENT_SHARES:
_13DFilings.percent_shares.latest,
PREVIOUS_DISCLOSURE_DATE:
_13DFilings.disclosure_date.latest,
DAYS_SINCE_PREV_DISCLOSURE:
BusinessDaysSince13DFilingsDate(),
}
@classmethod
def get_sids(cls):
return range(2)
@classmethod
def get_dataset(cls):
return {sid: frame
for sid, frame
in enumerate(_13d_filings_cases)}
loader_type = _13DFilingsLoader
def setup(self, dates):
cols = {
PREVIOUS_DISCLOSURE_DATE: self.get_sids_to_frames(
zip_with_dates,
[['NaT', '2014-01-04', '2014-01-09']],
date_intervals,
dates,
'datetime64[ns]',
'NaN'
),
PREVIOUS_NUM_SHARES: self.get_sids_to_frames(
zip_with_floats,
[['NaN', 1, 15]],
date_intervals,
dates,
'float',
'NaN'
),
PREVIOUS_PERCENT_SHARES: self.get_sids_to_frames(
zip_with_floats,
[['NaN', 10, 20]],
date_intervals,
dates,
'float',
'NaN'
)
}
cols[DAYS_SINCE_PREV_DISCLOSURE] = self._compute_busday_offsets(
cols[PREVIOUS_DISCLOSURE_DATE]
)
return cols
-172
View File
@@ -1,172 +0,0 @@
"""
Tests for the reference loader for Buyback Authorizations.
"""
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(
DAYS_SINCE_PREV,
SID_FIELD_NAME,
TS_FIELD_NAME,
)
from zipline.pipeline.data import BuybackAuthorizations
from zipline.pipeline.factors.events import BusinessDaysSinceBuybackAuth
from zipline.pipeline.loaders.buyback_auth import (
BUYBACK_AMOUNT_FIELD_NAME,
BUYBACK_ANNOUNCEMENT_FIELD_NAME,
BUYBACK_TYPE_FIELD_NAME,
BUYBACK_UNIT_FIELD_NAME,
BuybackAuthorizationsLoader,
)
from zipline.pipeline.loaders.blaze import BlazeBuybackAuthorizationsLoader
from zipline.pipeline.loaders.utils import (
zip_with_dates,
zip_with_floats,
zip_with_strs
)
from zipline.testing.fixtures import (
WithPipelineEventDataLoader, ZiplineTestCase
)
PREVIOUS_BUYBACK_AMOUNT = 'previous_value'
PREVIOUS_BUYBACK_ANNOUNCEMENT = 'previous_buyback_announcement'
PREVIOUS_BUYBACK_CASH = 'previous_buyback_cash'
PREVIOUS_BUYBACK_SHARE_COUNT = 'previous_buyback_share_count'
PREVIOUS_BUYBACK_TYPE = 'previous_buyback_type'
PREVIOUS_BUYBACK_UNIT = 'previous_buyback_unit'
date_intervals = [
[['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-09'],
['2014-01-10', '2014-01-31']]
]
buyback_authorizations_cases = [
pd.DataFrame({
BUYBACK_AMOUNT_FIELD_NAME: [1, 15],
BUYBACK_UNIT_FIELD_NAME: ["$M", "Mshares"],
BUYBACK_TYPE_FIELD_NAME: ["New", "Additional"],
TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-10']),
BUYBACK_ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-04',
'2014-01-09'])
}),
pd.DataFrame(
columns=[BUYBACK_AMOUNT_FIELD_NAME,
BUYBACK_UNIT_FIELD_NAME,
BUYBACK_TYPE_FIELD_NAME,
BUYBACK_ANNOUNCEMENT_FIELD_NAME,
TS_FIELD_NAME],
dtype='datetime64[ns]'
),
]
class BuybackAuthLoaderTestCase(WithPipelineEventDataLoader, ZiplineTestCase):
"""
Test for cash buyback authorizations dataset.
"""
pipeline_columns = {
PREVIOUS_BUYBACK_AMOUNT:
BuybackAuthorizations.previous_amount.latest,
PREVIOUS_BUYBACK_ANNOUNCEMENT:
BuybackAuthorizations.previous_date.latest,
PREVIOUS_BUYBACK_UNIT:
BuybackAuthorizations.previous_unit.latest,
PREVIOUS_BUYBACK_TYPE:
BuybackAuthorizations.previous_type.latest,
DAYS_SINCE_PREV:
BusinessDaysSinceBuybackAuth(),
}
@classmethod
def get_sids(cls):
return range(2)
@classmethod
def get_dataset(cls):
return {sid: frame
for sid, frame
in enumerate(buyback_authorizations_cases)}
loader_type = BuybackAuthorizationsLoader
def setup(self, dates):
cols = {
PREVIOUS_BUYBACK_AMOUNT: self.get_sids_to_frames(zip_with_floats,
[['NaN', 1, 15]],
date_intervals,
dates,
'float',
'NaN'),
PREVIOUS_BUYBACK_ANNOUNCEMENT: self.get_sids_to_frames(
zip_with_dates,
[['NaT', '2014-01-04', '2014-01-09']],
date_intervals,
dates,
'datetime64[ns]',
'NaN'
),
PREVIOUS_BUYBACK_UNIT: self.get_sids_to_frames(
zip_with_strs,
[[None, "$M", "Mshares"]],
date_intervals,
dates,
'category',
None
),
PREVIOUS_BUYBACK_TYPE: self.get_sids_to_frames(
zip_with_strs,
[[None, "New", "Additional"]],
date_intervals,
dates,
'category',
None
)
}
cols[DAYS_SINCE_PREV] = self._compute_busday_offsets(
cols[PREVIOUS_BUYBACK_ANNOUNCEMENT]
)
return cols
class BlazeBuybackAuthLoaderTestCase(BuybackAuthLoaderTestCase):
""" Test case for loading via blaze.
"""
loader_type = BlazeBuybackAuthorizationsLoader
def pipeline_event_loader_args(self, dates):
_, mapping = super(
BlazeBuybackAuthLoaderTestCase,
self,
).pipeline_event_loader_args(dates)
return (bz.data(pd.concat(
pd.DataFrame({
BUYBACK_ANNOUNCEMENT_FIELD_NAME:
frame[BUYBACK_ANNOUNCEMENT_FIELD_NAME],
BUYBACK_AMOUNT_FIELD_NAME:
frame[BUYBACK_AMOUNT_FIELD_NAME],
BUYBACK_UNIT_FIELD_NAME:
frame[BUYBACK_UNIT_FIELD_NAME],
BUYBACK_TYPE_FIELD_NAME:
frame[BUYBACK_TYPE_FIELD_NAME],
TS_FIELD_NAME:
frame[TS_FIELD_NAME],
SID_FIELD_NAME: sid,
})
for sid, frame in iteritems(mapping)
).reset_index(drop=True)),)
class BlazeBuybackAuthLoaderNotInteractiveTestCase(
BlazeBuybackAuthLoaderTestCase
):
"""Test case for passing a non-interactive symbol and a dict of resources.
"""
def pipeline_event_loader_args(self, dates):
(bound_expr,) = super(
BlazeBuybackAuthLoaderNotInteractiveTestCase,
self,
).pipeline_event_loader_args(dates)
return swap_resources_into_scope(bound_expr, {})
-353
View File
@@ -1,353 +0,0 @@
"""
Tests for the reference loader for ConsensusEstimates.
"""
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 SID_FIELD_NAME
from zipline.pipeline.data import ConsensusEstimates
from zipline.pipeline.loaders.consensus_estimates import (
ACTUAL_VALUE_FIELD_NAME,
ConsensusEstimatesLoader,
COUNT_FIELD_NAME,
FISCAL_QUARTER_FIELD_NAME,
FISCAL_YEAR_FIELD_NAME,
HIGH_FIELD_NAME,
LOW_FIELD_NAME,
MEAN_FIELD_NAME,
RELEASE_DATE_FIELD_NAME,
STANDARD_DEVIATION_FIELD_NAME,
)
from zipline.pipeline.loaders.blaze import BlazeConsensusEstimatesLoader
from zipline.pipeline.loaders.utils import (
zip_with_floats
)
from zipline.testing.fixtures import (
ZiplineTestCase,
WithNextAndPreviousEventDataLoader
)
NEXT_COUNT = 'next_count'
NEXT_FISCAL_QUARTER = 'next_fiscal_quarter'
NEXT_FISCAL_YEAR = 'next_fiscal_year'
NEXT_HIGH = 'next_high'
NEXT_LOW = 'next_low'
NEXT_MEAN = 'next_mean'
NEXT_RELEASE_DATE = 'next_release_date'
NEXT_STANDARD_DEVIATION = 'next_standard_deviation'
PREVIOUS_ACTUAL_VALUE = 'previous_actual_value'
PREVIOUS_COUNT = 'previous_count'
PREVIOUS_FISCAL_QUARTER = 'previous_fiscal_quarter'
PREVIOUS_FISCAL_YEAR = 'previous_fiscal_year'
PREVIOUS_HIGH = 'previous_high'
PREVIOUS_LOW = 'previous_low'
PREVIOUS_MEAN = 'previous_mean'
PREVIOUS_RELEASE_DATE = 'previous_release_date'
PREVIOUS_STANDARD_DEVIATION = 'previous_standard_deviation'
consensus_estimates_cases = [
# K1--K2--A1--A2.
pd.DataFrame({
ACTUAL_VALUE_FIELD_NAME: (100, 200),
STANDARD_DEVIATION_FIELD_NAME: (.5, .6),
COUNT_FIELD_NAME: (1, 2),
FISCAL_QUARTER_FIELD_NAME: (1, 1),
HIGH_FIELD_NAME: (.6, .7),
MEAN_FIELD_NAME: (.1, .2),
FISCAL_YEAR_FIELD_NAME: (2014, 2014),
LOW_FIELD_NAME: (.05, .06),
}),
# K1--K2--A2--A1.
pd.DataFrame({
ACTUAL_VALUE_FIELD_NAME: (200, 300),
STANDARD_DEVIATION_FIELD_NAME: (.6, .7),
COUNT_FIELD_NAME: (2, 3),
FISCAL_QUARTER_FIELD_NAME: (1, 1),
HIGH_FIELD_NAME: (.7, .8),
MEAN_FIELD_NAME: (.2, .3),
FISCAL_YEAR_FIELD_NAME: (2014, 2014),
LOW_FIELD_NAME: (.06, .07),
}),
# K1--A1--K2--A2.
pd.DataFrame({
ACTUAL_VALUE_FIELD_NAME: (300, 400),
STANDARD_DEVIATION_FIELD_NAME: (.7, .8),
COUNT_FIELD_NAME: (3, 4),
FISCAL_QUARTER_FIELD_NAME: (1, 1),
HIGH_FIELD_NAME: (.8, .9),
MEAN_FIELD_NAME: (.3, .4),
FISCAL_YEAR_FIELD_NAME: (2014, 2014),
LOW_FIELD_NAME: (.07, .08),
}),
# K1 == K2.
pd.DataFrame({
ACTUAL_VALUE_FIELD_NAME: (400, 500),
STANDARD_DEVIATION_FIELD_NAME: (.8, .9),
COUNT_FIELD_NAME: (4, 5),
FISCAL_QUARTER_FIELD_NAME: (1, 1),
HIGH_FIELD_NAME: (.9, 1.0),
MEAN_FIELD_NAME: (.4, .5),
FISCAL_YEAR_FIELD_NAME: (2014, 2014),
LOW_FIELD_NAME: (.08, .09),
}),
pd.DataFrame(
columns=[ACTUAL_VALUE_FIELD_NAME,
STANDARD_DEVIATION_FIELD_NAME,
COUNT_FIELD_NAME,
FISCAL_QUARTER_FIELD_NAME,
HIGH_FIELD_NAME,
MEAN_FIELD_NAME,
FISCAL_YEAR_FIELD_NAME,
LOW_FIELD_NAME],
dtype='datetime64[ns]'
),
]
prev_actual_value = [
['NaN', 100, 200],
['NaN', 300, 200],
['NaN', 300, 400],
['NaN', 400, 500],
['NaN']
]
next_standard_deviation = [
['NaN', .5, .6, 'NaN'],
['NaN', .6, .7, .6, 'NaN'],
['NaN', .7, 'NaN', .8, 'NaN'],
['NaN', .8, .9, 'NaN'],
['NaN']
]
prev_standard_deviation = [
['NaN', .5, .6],
['NaN', .7, .6],
['NaN', .7, .8],
['NaN', .8, .9],
['NaN']
]
next_count = [
['NaN', 1, 2, 'NaN'],
['NaN', 2, 3, 2, 'NaN'],
['NaN', 3, 'NaN', 4, 'NaN'],
['NaN', 4, 5, 'NaN'],
['NaN']
]
prev_count = [
['NaN', 1, 2],
['NaN', 3, 2],
['NaN', 3, 4],
['NaN', 4, 5],
['NaN']
]
next_fiscal_quarter = [
['NaN', 1, 1, 'NaN'],
['NaN', 1, 1, 1, 'NaN'],
['NaN', 1, 'NaN', 1, 'NaN'],
['NaN', 1, 1, 'NaN'],
['NaN']
]
prev_fiscal_quarter = [
['NaN', 1, 1],
['NaN', 1, 1],
['NaN', 1, 1],
['NaN', 1, 1],
['NaN']
]
next_high = [
['NaN', .6, .7, 'NaN'],
['NaN', .7, .8, .7, 'NaN'],
['NaN', .8, 'NaN', .9, 'NaN'],
['NaN', .9, 1.0, 'NaN'],
['NaN']
]
prev_high = [
['NaN', .6, .7],
['NaN', .8, .7],
['NaN', .8, .9],
['NaN', .9, 1.0],
['NaN']
]
next_mean = [
['NaN', .1, .2, 'NaN'],
['NaN', .2, .3, .2, 'NaN'],
['NaN', .3, 'NaN', .4, 'NaN'],
['NaN', .4, .5, 'NaN'],
['NaN']
]
prev_mean = [
['NaN', .1, .2],
['NaN', .3, .2],
['NaN', .3, .4],
['NaN', .4, .5],
['NaN']
]
next_fiscal_year = [
['NaN', 2014, 2014, 'NaN'],
['NaN', 2014, 2014, 2014, 'NaN'],
['NaN', 2014, 'NaN', 2014, 'NaN'],
['NaN', 2014, 2014, 'NaN'],
['NaN']
]
prev_fiscal_year = [
['NaN', 2014, 2014],
['NaN', 2014, 2014],
['NaN', 2014, 2014],
['NaN', 2014, 2014],
['NaN']
]
next_low = [
['NaN', .05, .06, 'NaN'],
['NaN', .06, .07, .06, 'NaN'],
['NaN', .07, 'NaN', .08, 'NaN'],
['NaN', .08, .09, 'NaN'],
['NaN']
]
prev_low = [
['NaN', .05, .06],
['NaN', .07, .06],
['NaN', .07, .08],
['NaN', .08, .09],
['NaN']
]
field_name_to_expected_col = {
PREVIOUS_ACTUAL_VALUE: prev_actual_value,
PREVIOUS_STANDARD_DEVIATION: prev_standard_deviation,
NEXT_STANDARD_DEVIATION: next_standard_deviation,
PREVIOUS_COUNT: prev_count,
NEXT_COUNT: next_count,
PREVIOUS_FISCAL_QUARTER: prev_fiscal_quarter,
NEXT_FISCAL_QUARTER: next_fiscal_quarter,
PREVIOUS_HIGH: prev_high,
NEXT_HIGH: next_high,
PREVIOUS_MEAN: prev_mean,
NEXT_MEAN: next_mean,
PREVIOUS_FISCAL_YEAR: prev_fiscal_year,
NEXT_FISCAL_YEAR: next_fiscal_year,
PREVIOUS_LOW: prev_low,
NEXT_LOW: next_low
}
class ConsensusEstimatesLoaderTestCase(WithNextAndPreviousEventDataLoader,
ZiplineTestCase):
"""
Tests for loading the consensus estimates data.
"""
pipeline_columns = {
PREVIOUS_ACTUAL_VALUE:
ConsensusEstimates.previous_actual_value.latest,
NEXT_RELEASE_DATE:
ConsensusEstimates.next_release_date.latest,
PREVIOUS_RELEASE_DATE:
ConsensusEstimates.previous_release_date.latest,
PREVIOUS_STANDARD_DEVIATION:
ConsensusEstimates.previous_standard_deviation.latest,
NEXT_STANDARD_DEVIATION:
ConsensusEstimates.next_standard_deviation.latest,
PREVIOUS_COUNT:
ConsensusEstimates.previous_count.latest,
NEXT_COUNT:
ConsensusEstimates.next_count.latest,
PREVIOUS_FISCAL_QUARTER:
ConsensusEstimates.previous_fiscal_quarter.latest,
NEXT_FISCAL_QUARTER:
ConsensusEstimates.next_fiscal_quarter.latest,
PREVIOUS_HIGH:
ConsensusEstimates.previous_high.latest,
NEXT_HIGH:
ConsensusEstimates.next_high.latest,
PREVIOUS_MEAN:
ConsensusEstimates.previous_mean.latest,
NEXT_MEAN:
ConsensusEstimates.next_mean.latest,
PREVIOUS_FISCAL_YEAR:
ConsensusEstimates.previous_fiscal_year.latest,
NEXT_FISCAL_YEAR:
ConsensusEstimates.next_fiscal_year.latest,
PREVIOUS_LOW:
ConsensusEstimates.previous_low.latest,
NEXT_LOW:
ConsensusEstimates.next_low.latest
}
@classmethod
def get_dataset(cls):
return {sid:
pd.concat([
cls.base_cases[sid].rename(columns={
'other_date': RELEASE_DATE_FIELD_NAME
}),
df
], axis=1)
for sid, df in enumerate(consensus_estimates_cases)}
loader_type = ConsensusEstimatesLoader
def setup(self, dates):
cols = {
PREVIOUS_RELEASE_DATE:
self.get_expected_previous_event_dates(
dates, 'datetime64[ns]', 'NaN'
),
NEXT_RELEASE_DATE: self.get_expected_next_event_dates(
dates, 'datetime64[ns]', 'NaN'
)
}
for field_name in field_name_to_expected_col:
cols[field_name] = self.get_sids_to_frames(
zip_with_floats, field_name_to_expected_col[field_name],
self.prev_date_intervals
if field_name.startswith("previous")
else self.next_date_intervals,
dates,
'float',
'NaN'
)
return cols
class BlazeConsensusEstimatesLoaderTestCase(ConsensusEstimatesLoaderTestCase):
loader_type = BlazeConsensusEstimatesLoader
def pipeline_event_loader_args(self, dates):
_, mapping = super(
BlazeConsensusEstimatesLoaderTestCase,
self,
).pipeline_event_loader_args(dates)
frames = []
for sid, df in iteritems(mapping):
frame = df.copy()
frame[SID_FIELD_NAME] = sid
frames.append(frame)
return bz.data(pd.concat(frames).reset_index(drop=True)),
class BlazeConsensusEstimatesLoaderNotInteractiveTestCase(
BlazeConsensusEstimatesLoaderTestCase
):
"""Test case for passing a non-interactive symbol and a dict of resources.
"""
def pipeline_event_loader_args(self, dates):
(bound_expr,) = super(
BlazeConsensusEstimatesLoaderNotInteractiveTestCase,
self,
).pipeline_event_loader_args(dates)
return swap_resources_into_scope(bound_expr, {})
-517
View File
@@ -1,517 +0,0 @@
"""
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,
PREVIOUS_AMOUNT,
PREVIOUS_ANNOUNCEMENT,
SID_FIELD_NAME,
TS_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 (
CASH_AMOUNT_FIELD_NAME,
CURRENCY_FIELD_NAME,
DIVIDEND_TYPE_FIELD_NAME,
DividendsByAnnouncementDateLoader,
DividendsByExDateLoader,
DividendsByPayDateLoader,
EX_DATE_FIELD_NAME,
PAY_DATE_FIELD_NAME,
)
from zipline.pipeline.loaders.utils import (
zip_with_dates,
zip_with_floats,
zip_with_strs,
)
from zipline.testing.fixtures import (
WithPipelineEventDataLoader,
ZiplineTestCase
)
DAYS_SINCE_PREV_DIVIDEND_ANNOUNCEMENT = 'days_since_prev_dividend_announcement'
DAYS_SINCE_PREV_EX_DATE = 'days_since_prev_ex_date'
DAYS_TO_NEXT_EX_DATE = 'days_to_next_ex_date'
NEXT_AMOUNT = 'next_amount'
NEXT_CURRENCY_TYPE = 'next_currency_type'
NEXT_DIVIDEND_TYPE = 'next_dividend_type'
NEXT_EX_DATE = 'next_ex_date'
NEXT_PAY_DATE = 'next_pay_date'
PREVIOUS_CURRENCY_TYPE = 'previous_currency_type'
PREVIOUS_DIVIDEND_TYPE = 'previous_dividend_type'
PREVIOUS_EX_DATE = 'previous_ex_date'
PREVIOUS_PAY_DATE = 'previous_pay_date'
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']),
CURRENCY_FIELD_NAME: ["$", "EUR"],
DIVIDEND_TYPE_FIELD_NAME: ["Stock", "Mixed"]
}),
# 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']),
CURRENCY_FIELD_NAME: ["EUR", "$"],
DIVIDEND_TYPE_FIELD_NAME: ["Mixed", "Stock"]
}),
# 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']),
CURRENCY_FIELD_NAME: ["$", "EUR"],
DIVIDEND_TYPE_FIELD_NAME: ["Stock", "Mixed"]
}),
# 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']),
CURRENCY_FIELD_NAME: ["$", "EUR"],
DIVIDEND_TYPE_FIELD_NAME: ["Stock", "Mixed"]
}),
pd.DataFrame(
columns=[CASH_AMOUNT_FIELD_NAME,
EX_DATE_FIELD_NAME,
PAY_DATE_FIELD_NAME,
TS_FIELD_NAME,
ANNOUNCEMENT_FIELD_NAME,
CURRENCY_FIELD_NAME,
DIVIDEND_TYPE_FIELD_NAME],
dtype='datetime64[ns]'
),
]
prev_date_intervals = [
[
['2014-01-01', '2014-01-14'], ['2014-01-15', '2014-01-19'],
['2014-01-20', '2014-01-31']
],
[
['2014-01-01', '2014-01-14'], ['2014-01-15', '2014-01-19'],
['2014-01-20', '2014-01-31']
],
[
['2014-01-01', '2014-01-09'], ['2014-01-10', '2014-01-19'],
['2014-01-20', '2014-01-31']
],
[
['2014-01-01', '2014-01-09'], ['2014-01-10', '2014-01-14'],
['2014-01-15', '2014-01-31']
]
]
next_date_intervals = [
[
['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-15'],
['2014-01-16', '2014-01-20'], ['2014-01-21', '2014-01-31']
],
[
['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-09'],
['2014-01-10', '2014-01-15'], ['2014-01-16', '2014-01-20'],
['2014-01-21', '2014-01-31']
],
[
['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-10'],
['2014-01-11', '2014-01-14'], ['2014-01-15', '2014-01-20'],
['2014-01-21', '2014-01-31']
],
[
['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-10'],
['2014-01-11', '2014-01-15'], ['2014-01-16', '2014-01-31']
]
]
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']]
prev_currency_types = [[None, "$", "EUR"],
[None, "$", "EUR"],
[None, "$", "EUR"],
[None, "$", "EUR"]]
next_currency_types = [[None, "$", "EUR", None],
[None, "EUR", "$", "EUR", None],
[None, "$", None, "EUR", None],
[None, "$", "EUR", None]]
prev_dividend_types = [[None, "Stock", "Mixed"],
[None, "Stock", "Mixed"],
[None, "Stock", "Mixed"],
[None, "Stock", "Mixed"]]
next_dividend_types = [[None, "Stock", "Mixed", None],
[None, "Mixed", "Stock", "Mixed", None],
[None, "Stock", None, "Mixed", None],
[None, "Stock", "Mixed", None]]
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(),
PREVIOUS_CURRENCY_TYPE:
DividendsByAnnouncementDate.previous_currency.latest,
PREVIOUS_DIVIDEND_TYPE:
DividendsByAnnouncementDate.previous_type.latest,
}
@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 = [
[
['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-09'],
['2014-01-10', '2014-01-31']
],
[
['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-09'],
['2014-01-10', '2014-01-31']
],
[
['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-14'],
['2014-01-15', '2014-01-31']
],
[
['2014-01-01', '2014-01-04'], ['2014-01-05', '2014-01-31']
]
]
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]]
currency_types = [[None, "$", "EUR"], [None, "EUR", "$"],
[None, "$", "EUR"], [None, "EUR"]]
dividend_types = [[None, "Stock", "Mixed"], [None, "Mixed", "Stock"],
[None, "Stock", "Mixed"], [None, "Mixed"]]
cols = {
PREVIOUS_ANNOUNCEMENT: self.get_sids_to_frames(
zip_with_dates, announcement_dates, date_intervals, dates,
'datetime64[ns]', 'NaN'
),
PREVIOUS_AMOUNT: self.get_sids_to_frames(
zip_with_floats, amounts, date_intervals, dates, 'float', 'NaN'
),
PREVIOUS_CURRENCY_TYPE: self.get_sids_to_frames(
zip_with_strs, currency_types, date_intervals, dates,
'category', None
),
PREVIOUS_DIVIDEND_TYPE: self.get_sids_to_frames(
zip_with_strs, dividend_types, date_intervals, dates,
'category', None
),
}
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],
CURRENCY_FIELD_NAME: df[CURRENCY_FIELD_NAME],
DIVIDEND_TYPE_FIELD_NAME: df[DIVIDEND_TYPE_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,
PREVIOUS_CURRENCY_TYPE: DividendsByExDate.previous_currency.latest,
NEXT_CURRENCY_TYPE: DividendsByExDate.next_currency.latest,
PREVIOUS_DIVIDEND_TYPE: DividendsByExDate.previous_type.latest,
NEXT_DIVIDEND_TYPE: DividendsByExDate.next_type.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 = {
NEXT_EX_DATE: self.get_sids_to_frames(
zip_with_dates, next_ex_and_pay_dates, next_date_intervals,
dates,
'datetime64[ns]', 'NaN'
),
PREVIOUS_EX_DATE: self.get_sids_to_frames(
zip_with_dates, prev_ex_and_pay_dates, prev_date_intervals,
dates,
'datetime64[ns]', 'NaN'
),
NEXT_AMOUNT: self.get_sids_to_frames(
zip_with_floats, next_amounts, next_date_intervals, dates,
'float', 'NaN'
),
PREVIOUS_AMOUNT: self.get_sids_to_frames(
zip_with_floats, prev_amounts, prev_date_intervals, dates,
'float', 'NaN'
),
PREVIOUS_CURRENCY_TYPE: self.get_sids_to_frames(
zip_with_strs, prev_currency_types, prev_date_intervals, dates,
'category', None
),
NEXT_CURRENCY_TYPE: self.get_sids_to_frames(
zip_with_strs, next_currency_types, next_date_intervals, dates,
'category', None
),
PREVIOUS_DIVIDEND_TYPE: self.get_sids_to_frames(
zip_with_strs, prev_dividend_types, prev_date_intervals, dates,
'category', None
),
NEXT_DIVIDEND_TYPE: self.get_sids_to_frames(
zip_with_strs, next_dividend_types, next_date_intervals, dates,
'category', None
),
}
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],
CURRENCY_FIELD_NAME: df[CURRENCY_FIELD_NAME],
DIVIDEND_TYPE_FIELD_NAME: df[DIVIDEND_TYPE_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,
PREVIOUS_CURRENCY_TYPE: DividendsByPayDate.previous_currency.latest,
NEXT_CURRENCY_TYPE: DividendsByPayDate.next_currency.latest,
PREVIOUS_DIVIDEND_TYPE: DividendsByPayDate.previous_type.latest,
NEXT_DIVIDEND_TYPE: DividendsByPayDate.next_type.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):
return {
NEXT_PAY_DATE: self.get_sids_to_frames(
zip_with_dates, next_ex_and_pay_dates, next_date_intervals,
dates,
'datetime64[ns]', 'NaN'
),
PREVIOUS_PAY_DATE: self.get_sids_to_frames(
zip_with_dates, prev_ex_and_pay_dates, prev_date_intervals,
dates,
'datetime64[ns]', 'NaN'
),
NEXT_AMOUNT: self.get_sids_to_frames(
zip_with_floats, next_amounts, next_date_intervals, dates,
'float', 'NaN'
),
PREVIOUS_AMOUNT: self.get_sids_to_frames(
zip_with_floats, prev_amounts, prev_date_intervals, dates,
'float', 'NaN'
),
PREVIOUS_CURRENCY_TYPE: self.get_sids_to_frames(
zip_with_strs, prev_currency_types, prev_date_intervals, dates,
'category', None
),
NEXT_CURRENCY_TYPE: self.get_sids_to_frames(
zip_with_strs, next_currency_types, next_date_intervals, dates,
'category', None
),
PREVIOUS_DIVIDEND_TYPE: self.get_sids_to_frames(
zip_with_strs, prev_dividend_types, prev_date_intervals, dates,
'category', None
),
NEXT_DIVIDEND_TYPE: self.get_sids_to_frames(
zip_with_strs, next_dividend_types, next_date_intervals, dates,
'category', None
),
}
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],
CURRENCY_FIELD_NAME: df[CURRENCY_FIELD_NAME],
DIVIDEND_TYPE_FIELD_NAME: df[DIVIDEND_TYPE_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, {})
-98
View File
@@ -1,98 +0,0 @@
"""
Tests for the reference loader for EarningsCalendar.
"""
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,
DAYS_TO_NEXT,
NEXT_ANNOUNCEMENT,
PREVIOUS_ANNOUNCEMENT,
SID_FIELD_NAME,
TS_FIELD_NAME
)
from zipline.pipeline.data import EarningsCalendar
from zipline.pipeline.factors.events import (
BusinessDaysSincePreviousEarnings,
BusinessDaysUntilNextEarnings,
)
from zipline.pipeline.loaders.earnings import EarningsCalendarLoader
from zipline.pipeline.loaders.blaze import BlazeEarningsCalendarLoader
from zipline.testing.fixtures import (
ZiplineTestCase,
WithNextAndPreviousEventDataLoader
)
class EarningsCalendarLoaderTestCase(WithNextAndPreviousEventDataLoader,
ZiplineTestCase):
"""
Tests for loading the earnings announcement data.
"""
pipeline_columns = {
NEXT_ANNOUNCEMENT: EarningsCalendar.next_announcement.latest,
PREVIOUS_ANNOUNCEMENT: EarningsCalendar.previous_announcement.latest,
DAYS_SINCE_PREV: BusinessDaysSincePreviousEarnings(),
DAYS_TO_NEXT: BusinessDaysUntilNextEarnings(),
}
@classmethod
def get_dataset(cls):
return {sid: df.rename(
columns={'other_date': ANNOUNCEMENT_FIELD_NAME}
) for sid, df in enumerate(cls.base_cases)}
loader_type = EarningsCalendarLoader
def setup(self, dates):
cols = {
PREVIOUS_ANNOUNCEMENT: self.get_expected_previous_event_dates(
dates,
'datetime64[ns]', 'NaN'
),
NEXT_ANNOUNCEMENT: self.get_expected_next_event_dates(
dates, 'datetime64[ns]', 'NaN'
),
}
cols[DAYS_TO_NEXT] = self._compute_busday_offsets(
cols[NEXT_ANNOUNCEMENT]
)
cols[DAYS_SINCE_PREV] = self._compute_busday_offsets(
cols[PREVIOUS_ANNOUNCEMENT]
)
return cols
class BlazeEarningsCalendarLoaderTestCase(EarningsCalendarLoaderTestCase):
loader_type = BlazeEarningsCalendarLoader
def pipeline_event_loader_args(self, dates):
_, mapping = super(
BlazeEarningsCalendarLoaderTestCase,
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,
})
for sid, df in iteritems(mapping)
).reset_index(drop=True)),)
class BlazeEarningsCalendarLoaderNotInteractiveTestCase(
BlazeEarningsCalendarLoaderTestCase):
"""Test case for passing a non-interactive symbol and a dict of resources.
"""
def pipeline_event_loader_args(self, dates):
(bound_expr,) = super(
BlazeEarningsCalendarLoaderNotInteractiveTestCase,
self,
).pipeline_event_loader_args(dates)
return swap_resources_into_scope(bound_expr, {})
+425 -266
View File
@@ -1,299 +1,458 @@
"""
Tests for setting up an EventsLoader and a BlazeEventsLoader.
"""
import re
from unittest import TestCase
from itertools import product
import blaze as bz
from nose_parameterized import parameterized
import numpy as np
from numpy.testing import assert_array_equal
import pandas as pd
from pandas.util.testing import assert_series_equal
from zipline.pipeline import Pipeline, SimplePipelineEngine
from zipline.pipeline.common import (
ANNOUNCEMENT_FIELD_NAME,
EVENT_DATE_FIELD_NAME,
TS_FIELD_NAME,
SID_FIELD_NAME,
TS_FIELD_NAME
)
from zipline.pipeline.data import DataSet, Column
from zipline.pipeline.loaders.events import EventsLoader
from zipline.pipeline.loaders.blaze.events import BlazeEventsLoader
from zipline.pipeline.loaders.events import (
DF_NO_TS_NOT_INFER_TS_ERROR,
DTINDEX_NOT_INFER_TS_ERROR,
EventsLoader,
SERIES_NO_DTINDEX_ERROR,
WRONG_COLS_ERROR,
WRONG_MANY_COL_DATA_FORMAT_ERROR,
WRONG_SINGLE_COL_DATA_FORMAT_ERROR
from zipline.pipeline.loaders.utils import (
previous_event_indexer,
next_event_indexer,
)
from zipline.testing import ZiplineTestCase
from zipline.testing.fixtures import (
WithAssetFinder,
WithNYSETradingDays,
)
from zipline.testing.predicates import assert_equal
from zipline.utils.numpy_utils import (
categorical_dtype,
datetime64ns_dtype,
float64_dtype,
int64_dtype,
)
from zipline.utils.memoize import lazyval
from zipline.utils.numpy_utils import datetime64ns_dtype
OTHER_FIELD = "other_field"
ABSTRACT_CONCRETE_LOADER_ERROR = 'abstract methods concrete_loader'
ABSTRACT_EXPECTED_COLS_ERROR = 'abstract methods event_date_col, expected_cols'
class EventDataSet(DataSet):
previous_announcement = Column(datetime64ns_dtype)
next_announcement = Column(datetime64ns_dtype)
previous_event_date = Column(dtype=datetime64ns_dtype)
next_event_date = Column(dtype=datetime64ns_dtype)
class EventDataSetLoader(EventsLoader):
expected_cols = frozenset([ANNOUNCEMENT_FIELD_NAME])
previous_float = Column(dtype=float64_dtype)
next_float = Column(dtype=float64_dtype)
event_date_col = ANNOUNCEMENT_FIELD_NAME
previous_datetime = Column(dtype=datetime64ns_dtype)
next_datetime = Column(dtype=datetime64ns_dtype)
def __init__(self,
all_dates,
events_by_sid,
infer_timestamps=False,
dataset=EventDataSet):
super(EventDataSetLoader, self).__init__(
all_dates,
events_by_sid,
infer_timestamps=infer_timestamps,
dataset=dataset,
)
previous_int = Column(dtype=int64_dtype, missing_value=-1)
next_int = Column(dtype=int64_dtype, missing_value=-1)
@lazyval
def previous_announcement_loader(self):
return self._previous_event_date_loader(
self.dataset.previous_announcement,
)
previous_string = Column(dtype=categorical_dtype, missing_value=None)
next_string = Column(dtype=categorical_dtype, missing_value=None)
@lazyval
def next_announcement_loader(self):
return self._next_event_date_loader(
self.dataset.next_announcement,
)
# Test case just for catching an error when multiple columns are in the wrong
# data format, so no loader defined.
class EventDataSetLoaderMultipleExpectedColsNoColumnLoaders(EventsLoader):
expected_cols = frozenset([ANNOUNCEMENT_FIELD_NAME, OTHER_FIELD])
event_date_col = ANNOUNCEMENT_FIELD_NAME
class EventDataSetLoaderNoExpectedCols(EventsLoader):
def __init__(self,
all_dates,
events_by_sid,
infer_timestamps=False,
dataset=EventDataSet):
super(EventDataSetLoaderNoExpectedCols, self).__init__(
all_dates,
events_by_sid,
infer_timestamps=infer_timestamps,
dataset=dataset,
)
dtx = pd.date_range('2014-01-01', '2014-01-10')
class EventLoaderTestCase(TestCase):
def test_null_in_event_date_col(self):
# Tests that if there is a null date in the event date column, it is
# filtered out and does not break on loading the adjusted array.
dates_with_null = pd.Series(dtx)
dates_with_null[2] = pd.NaT
events_by_sid = {0: pd.DataFrame({ANNOUNCEMENT_FIELD_NAME:
dates_with_null,
TS_FIELD_NAME: dtx})}
loader = EventDataSetLoader(
dtx,
events_by_sid,
)
prev_result = loader.load_adjusted_array({
EventDataSet.previous_announcement
}, dtx, [0], [True])[EventDataSet.previous_announcement].data[:, 0]
next_result = loader.load_adjusted_array({
EventDataSet.next_announcement
}, dtx, [0], [True])[EventDataSet.next_announcement].data[:, 0]
expected_prev = dates_with_null[:]
expected_prev[2] = dtx[1]
assert_array_equal(prev_result, expected_prev)
expected_next = dates_with_null[:]
expected_next[2] = np.datetime64('NaT')
assert_array_equal(next_result, expected_next)
def assert_loader_error(self, events_by_sid, error, msg,
infer_timestamps, loader):
with self.assertRaisesRegexp(error, re.escape(msg)):
loader(
dtx, events_by_sid, infer_timestamps=infer_timestamps,
)
def test_no_expected_cols_defined(self):
events_by_sid = {0: pd.DataFrame({ANNOUNCEMENT_FIELD_NAME: dtx})}
self.assert_loader_error(events_by_sid, TypeError,
ABSTRACT_EXPECTED_COLS_ERROR,
True, EventDataSetLoaderNoExpectedCols)
def test_wrong_cols(self):
wrong_col_name = 'some_other_col'
# Test wrong cols (cols != expected)
events_by_sid = {0: pd.DataFrame({wrong_col_name: dtx})}
self.assert_loader_error(
events_by_sid, ValueError, WRONG_COLS_ERROR.format(
expected_columns=list(EventDataSetLoader.expected_cols),
sid=0,
resulting_columns=[wrong_col_name],
),
True,
EventDataSetLoader
)
@parameterized.expand([
# DataFrame without timestamp column and infer_timestamps = True
[pd.DataFrame({ANNOUNCEMENT_FIELD_NAME: dtx}), True],
# DataFrame with timestamp column
[pd.DataFrame({ANNOUNCEMENT_FIELD_NAME: dtx,
TS_FIELD_NAME: dtx}), False],
# DatetimeIndex with infer_timestamps = True
[pd.DatetimeIndex(dtx), True],
# Series with DatetimeIndex as index and infer_timestamps = False
[pd.Series(dtx, index=dtx), False]
])
def test_conversion_to_df(self, df, infer_timestamps):
events_by_sid = {0: df}
loader = EventDataSetLoader(
dtx,
events_by_sid,
infer_timestamps=infer_timestamps,
)
self.assertEqual(
loader.events_by_sid.keys(),
events_by_sid.keys(),
)
if infer_timestamps:
expected = pd.Series(index=[dtx[0]] * 10, data=dtx,
name=ANNOUNCEMENT_FIELD_NAME)
else:
expected = pd.Series(index=dtx, data=dtx,
name=ANNOUNCEMENT_FIELD_NAME)
expected.index.name = TS_FIELD_NAME
# Check that index by first given date has been added
assert_series_equal(
loader.events_by_sid[0][ANNOUNCEMENT_FIELD_NAME],
expected,
)
@parameterized.expand(
[
# DataFrame without timestamp column and infer_timestamps = True
[
pd.DataFrame({ANNOUNCEMENT_FIELD_NAME: dtx}),
False,
DF_NO_TS_NOT_INFER_TS_ERROR.format(
timestamp_column_name=TS_FIELD_NAME,
sid=0
),
EventDataSetLoader
],
# DatetimeIndex with infer_timestamps = False
[
pd.DatetimeIndex(dtx, name=ANNOUNCEMENT_FIELD_NAME),
False,
DTINDEX_NOT_INFER_TS_ERROR.format(sid=0),
EventDataSetLoader
],
# Series with DatetimeIndex as index and infer_timestamps = False
[
pd.Series(dtx, name=ANNOUNCEMENT_FIELD_NAME),
False,
SERIES_NO_DTINDEX_ERROR.format(sid=0),
EventDataSetLoader
],
# Below, 2 cases repeated for infer_timestamps = True and False.
# Shouldn't make a difference in the outcome.
# We expected 1 column but got a data structure other than a
# DataFrame, Series, or DatetimeIndex
[
[dtx],
True,
WRONG_SINGLE_COL_DATA_FORMAT_ERROR.format(sid=0),
EventDataSetLoader
],
# We expected multiple columns but got a data structure other
# than a DataFrame
[
[dtx, dtx],
True,
WRONG_MANY_COL_DATA_FORMAT_ERROR.format(sid=0),
EventDataSetLoaderMultipleExpectedColsNoColumnLoaders
],
[
[dtx],
False,
WRONG_SINGLE_COL_DATA_FORMAT_ERROR.format(sid=0),
EventDataSetLoader
],
# We expected multiple columns but got a data structure other
# than a DataFrame
[
[dtx, dtx],
False,
WRONG_MANY_COL_DATA_FORMAT_ERROR.format(sid=0),
EventDataSetLoaderMultipleExpectedColsNoColumnLoaders
]
]
previous_string_custom_missing = Column(
dtype=categorical_dtype,
missing_value=u"<<NULL>>",
)
next_string_custom_missing = Column(
dtype=categorical_dtype,
missing_value=u"<<NULL>>",
)
def test_bad_conversion_to_df(self, df, infer_timestamps, msg, loader):
events_by_sid = {0: df}
self.assert_loader_error(events_by_sid, ValueError, msg,
infer_timestamps, loader)
class BlazeEventDataSetLoaderNoConcreteLoader(BlazeEventsLoader):
def __init__(self,
expr,
dataset=EventDataSet,
**kwargs):
super(
BlazeEventDataSetLoaderNoConcreteLoader, self
).__init__(expr,
dataset=dataset,
**kwargs)
critical_dates = pd.to_datetime([
'2014-01-05',
'2014-01-10',
'2014-01-15',
'2014-01-20',
])
class BlazeEventLoaderTestCase(TestCase):
# Blaze loader: need to test failure if no concrete loader
def test_no_concrete_loader_defined(self):
with self.assertRaisesRegexp(
TypeError, re.escape(ABSTRACT_CONCRETE_LOADER_ERROR)
):
BlazeEventDataSetLoaderNoConcreteLoader(
bz.data(
pd.DataFrame({ANNOUNCEMENT_FIELD_NAME: dtx,
SID_FIELD_NAME: 0})
def make_events_for_sid(sid, event_dates, event_timestamps):
num_events = len(event_dates)
return pd.DataFrame({
'sid': np.full(num_events, sid, dtype=np.int64),
'timestamp': event_timestamps,
'event_date': event_dates,
'float': np.arange(num_events, dtype=np.float64) + sid,
'int': np.arange(num_events) + sid,
'datetime': pd.date_range('1990-01-01', periods=num_events).shift(sid),
'string': ['-'.join([str(sid), str(i)]) for i in range(num_events)],
})
def make_null_event_date_events(all_sids, timestamp):
"""
Make an event with a null event_date for all sids.
Used to test that EventsLoaders filter out null events.
"""
return pd.DataFrame({
'sid': all_sids,
'timestamp': timestamp,
'event_date': pd.Timestamp('NaT'),
'float': -9999.0,
'int': -9999,
'datetime': pd.Timestamp('1980'),
'string': 'should be ignored',
})
def make_events(add_nulls):
"""
Every event has at least three pieces of data associated with it:
1. sid : The ID of the asset associated with the event.
2. event_date : The date on which an event occurred.
3. timestamp : The date on which we learned about the event.
This can be before the occurence_date in the case of an
announcement about an upcoming event.
Events for two different sids shouldn't interact in any way, so the
interesting cases are determined by the possible interleavings of
event_date and timestamp for a single sid.
Fix two events with dates e1, e2 and timestamps t1 and t2.
Without loss of generality, assume that e1 < e2. (If two events have the
same occurrence date, the behavior of next/previous event is undefined).
The remaining possible sequences of events are given by taking all possible
4-tuples of four ascending dates. For each possible interleaving, we
generate a set of fake events with those dates and assign them to a new
sid.
"""
def gen_date_interleavings():
for e1, e2, t1, t2 in product(*[critical_dates] * 4):
if e1 < e2:
yield (e1, e2, t1, t2)
event_frames = []
for sid, (e1, e2, t1, t2) in enumerate(gen_date_interleavings()):
event_frames.append(make_events_for_sid(sid, [e1, e2], [t1, t2]))
if add_nulls:
for date in critical_dates:
event_frames.append(
make_null_event_date_events(
np.arange(sid + 1),
timestamp=date,
)
)
return pd.concat(event_frames, ignore_index=True)
class BlazeEventDataSetLoader(BlazeEventsLoader):
concrete_loader = EventDataSetLoader
_expected_fields = frozenset({ANNOUNCEMENT_FIELD_NAME,
TS_FIELD_NAME,
SID_FIELD_NAME})
def __init__(self,
expr,
dataset=EventDataSet,
**kwargs):
super(
BlazeEventDataSetLoader, self
).__init__(expr,
dataset=dataset,
**kwargs)
class EventIndexerTestCase(ZiplineTestCase):
@classmethod
def init_class_fixtures(cls):
super(EventIndexerTestCase, cls).init_class_fixtures()
cls.events = make_events(add_nulls=False).sort('event_date')
cls.events.reset_index(inplace=True)
def test_previous_event_indexer(self):
events = self.events
event_sids = events['sid'].values
event_dates = events['event_date'].values
event_timestamps = events['timestamp'].values
all_dates = pd.date_range('2014', '2014-01-31')
all_sids = np.unique(event_sids)
indexer = previous_event_indexer(
all_dates,
all_sids,
event_dates,
event_timestamps,
event_sids,
)
# Compute expected results without knowledge of null events.
for i, sid in enumerate(all_sids):
self.check_previous_event_indexer(
events,
all_dates,
sid,
indexer[:, i],
)
def check_previous_event_indexer(self,
events,
all_dates,
sid,
indexer):
relevant_events = events[events.sid == sid]
self.assertEqual(len(relevant_events), 2)
ix1, ix2 = relevant_events.index
# An event becomes a possible value once we're past both its event_date
# and its timestamp.
event1_first_eligible = max(
relevant_events.loc[ix1, ['event_date', 'timestamp']],
)
event2_first_eligible = max(
relevant_events.loc[ix2, ['event_date', 'timestamp']],
)
for date, computed_index in zip(all_dates, indexer):
if date >= event2_first_eligible:
# If we've seen event 2, it should win even if we've seen event
# 1, because events are sorted by event_date.
self.assertEqual(computed_index, ix2)
elif date >= event1_first_eligible:
# If we've seen event 1 but not event 2, event 1 should win.
self.assertEqual(computed_index, ix1)
else:
# If we haven't seen either event, then we should have -1 as
# sentinel.
self.assertEqual(computed_index, -1)
def test_next_event_indexer(self):
events = self.events
event_sids = events['sid'].values
event_dates = events['event_date'].values
event_timestamps = events['timestamp'].values
all_dates = pd.date_range('2014', '2014-01-31')
all_sids = np.unique(event_sids)
indexer = next_event_indexer(
all_dates,
all_sids,
event_dates,
event_timestamps,
event_sids,
)
# Compute expected results without knowledge of null events.
for i, sid in enumerate(all_sids):
self.check_next_event_indexer(
events,
all_dates,
sid,
indexer[:, i],
)
def check_next_event_indexer(self,
events,
all_dates,
sid,
indexer):
relevant_events = events[events.sid == sid]
self.assertEqual(len(relevant_events), 2)
ix1, ix2 = relevant_events.index
e1, e2 = relevant_events['event_date']
t1, t2 = relevant_events['timestamp']
for date, computed_index in zip(all_dates, indexer):
# An event is eligible to be the next event if it's between the
# timestamp and the event_date, inclusive.
if t1 <= date <= e1:
# If e1 is eligible, it should be chosen even if e2 is
# eligible, since it's earlier.
self.assertEqual(computed_index, ix1)
elif t2 <= date <= e2:
# e2 is eligible and e1 is not, so e2 should be chosen.
self.assertEqual(computed_index, ix2)
else:
# Neither event is eligible. Return -1 as a sentinel.
self.assertEqual(computed_index, -1)
class EventsLoaderTestCase(WithAssetFinder,
WithNYSETradingDays,
ZiplineTestCase):
START_DATE = pd.Timestamp('2014-01-01')
END_DATE = pd.Timestamp('2014-01-30')
@classmethod
def init_class_fixtures(cls):
# This is a rare case where we actually want to do work **before** we
# call init_class_fixtures. We choose our sids for WithAssetFinder
# based on the events generated by make_event_data.
cls.raw_events = make_events(add_nulls=True)
cls.raw_events_no_nulls = cls.raw_events[
cls.raw_events['event_date'].notnull()
]
cls.next_value_columns = {
EventDataSet.next_datetime: 'datetime',
EventDataSet.next_event_date: 'event_date',
EventDataSet.next_float: 'float',
EventDataSet.next_int: 'int',
EventDataSet.next_string: 'string',
EventDataSet.next_string_custom_missing: 'string'
}
cls.previous_value_columns = {
EventDataSet.previous_datetime: 'datetime',
EventDataSet.previous_event_date: 'event_date',
EventDataSet.previous_float: 'float',
EventDataSet.previous_int: 'int',
EventDataSet.previous_string: 'string',
EventDataSet.previous_string_custom_missing: 'string'
}
cls.loader = cls.make_loader(
events=cls.raw_events,
next_value_columns=cls.next_value_columns,
previous_value_columns=cls.previous_value_columns,
)
cls.ASSET_FINDER_EQUITY_SIDS = list(cls.raw_events['sid'].unique())
cls.ASSET_FINDER_EQUITY_SYMBOLS = [
's' + str(n) for n in cls.ASSET_FINDER_EQUITY_SIDS
]
super(EventsLoaderTestCase, cls).init_class_fixtures()
@classmethod
def make_loader(cls, events, next_value_columns, previous_value_columns):
# This method exists to be overridden by BlazeEventsLoaderTestCase
return EventsLoader(events, next_value_columns, previous_value_columns)
def test_load_with_trading_calendar(self):
engine = SimplePipelineEngine(
lambda x: self.loader,
self.trading_days,
self.asset_finder,
)
results = engine.run_pipeline(
Pipeline({c.name: c.latest for c in EventDataSet.columns}),
start_date=self.trading_days[0],
end_date=self.trading_days[-1],
)
for c in EventDataSet.columns:
if c in self.next_value_columns:
self.check_next_value_results(c, results[c.name].unstack())
elif c in self.previous_value_columns:
self.check_previous_value_results(c, results[c.name].unstack())
else:
raise AssertionError("Unexpected column %s." % c)
def assert_result_contains_all_sids(self, results):
assert_equal(
list(map(int, results.columns)),
self.ASSET_FINDER_EQUITY_SIDS,
)
def check_previous_value_results(self, column, results):
"""
Check previous value results for a single column.
"""
# Verify that we got a result for every sid.
self.assert_result_contains_all_sids(results)
events = self.raw_events_no_nulls
# Remove timezone info from trading days, since the outputs
# from pandas won't be tz_localized.
dates = self.trading_days.tz_localize(None)
for asset, asset_result in results.iterkv():
relevant_events = events[events.sid == asset.sid]
self.assertEqual(len(relevant_events), 2)
v1, v2 = relevant_events[self.previous_value_columns[column]]
event1_first_eligible = max(
# .ix doesn't work here because the frame index contains
# integers, so 0 is still interpreted as a key.
relevant_events.iloc[0].loc[['event_date', 'timestamp']],
)
event2_first_eligible = max(
relevant_events.iloc[1].loc[['event_date', 'timestamp']]
)
for date, computed_value in zip(dates, asset_result):
if date >= event2_first_eligible:
# If we've seen event 2, it should win even if we've seen
# event 1, because events are sorted by event_date.
self.assertEqual(computed_value, v2)
elif date >= event1_first_eligible:
# If we've seen event 1 but not event 2, event 1 should
# win.
self.assertEqual(computed_value, v1)
else:
# If we haven't seen either event, then we should have
# column.missing_value.
assert_equal(
computed_value,
column.missing_value,
# Coerce from Timestamp to datetime64.
allow_datetime_coercions=True,
)
def check_next_value_results(self, column, results):
"""
Check results for a single column.
"""
self.assert_result_contains_all_sids(results)
events = self.raw_events_no_nulls
# Remove timezone info from trading days, since the outputs
# from pandas won't be tz_localized.
dates = self.trading_days.tz_localize(None)
for asset, asset_result in results.iterkv():
relevant_events = events[events.sid == asset.sid]
self.assertEqual(len(relevant_events), 2)
v1, v2 = relevant_events[self.next_value_columns[column]]
e1, e2 = relevant_events['event_date']
t1, t2 = relevant_events['timestamp']
for date, computed_value in zip(dates, asset_result):
if t1 <= date <= e1:
# If we've seen event 2, it should win even if we've seen
# event 1, because events are sorted by event_date.
self.assertEqual(computed_value, v1)
elif t2 <= date <= e2:
# If we've seen event 1 but not event 2, event 1 should
# win.
self.assertEqual(computed_value, v2)
else:
# If we haven't seen either event, then we should have
# column.missing_value.
assert_equal(
computed_value,
column.missing_value,
# Coerce from Timestamp to datetime64.
allow_datetime_coercions=True,
)
def test_wrong_cols(self):
# Test wrong cols (cols != expected)
events = pd.DataFrame({
'c': [5],
SID_FIELD_NAME: [1],
TS_FIELD_NAME: [pd.Timestamp('2014')],
EVENT_DATE_FIELD_NAME: [pd.Timestamp('2014')],
})
EventsLoader(events, {EventDataSet.next_float: 'c'}, {})
EventsLoader(events, {}, {EventDataSet.previous_float: 'c'})
with self.assertRaises(ValueError) as e:
EventsLoader(events, {EventDataSet.next_float: 'd'}, {})
msg = str(e.exception)
expected = (
"EventsLoader missing required columns ['d'].\n"
"Got Columns: ['c', 'event_date', 'sid', 'timestamp']\n"
"Expected Columns: ['d', 'event_date', 'sid', 'timestamp']"
)
self.assertEqual(msg, expected)
class BlazeEventsLoaderTestCase(EventsLoaderTestCase):
"""
Run the same tests as EventsLoaderTestCase, but using a BlazeEventsLoader.
"""
@classmethod
def make_loader(cls, events, next_value_columns, previous_value_columns):
return BlazeEventsLoader(
bz.data(events),
next_value_columns,
previous_value_columns,
)