mirror of
https://github.com/wassname/catalyst.git
synced 2026-08-01 12:20:21 +08:00
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:
@@ -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
|
||||
@@ -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, {})
|
||||
@@ -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, {})
|
||||
@@ -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, {})
|
||||
@@ -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
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user