mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-13 17:42:42 +08:00
Maya Tydykov
346 lines
9.2 KiB
Python
346 lines
9.2 KiB
Python
"""
|
|
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 (
|
|
ACTUAL_VALUE_FIELD_NAME,
|
|
COUNT_FIELD_NAME,
|
|
FISCAL_QUARTER_FIELD_NAME,
|
|
FISCAL_YEAR_FIELD_NAME,
|
|
HIGH_FIELD_NAME,
|
|
LOW_FIELD_NAME,
|
|
MEAN_FIELD_NAME,
|
|
NEXT_COUNT,
|
|
NEXT_FISCAL_QUARTER,
|
|
NEXT_FISCAL_YEAR,
|
|
NEXT_HIGH,
|
|
NEXT_LOW,
|
|
NEXT_RELEASE_DATE,
|
|
NEXT_STANDARD_DEVIATION,
|
|
PREVIOUS_ACTUAL_VALUE,
|
|
PREVIOUS_COUNT,
|
|
PREVIOUS_FISCAL_QUARTER,
|
|
PREVIOUS_FISCAL_YEAR,
|
|
PREVIOUS_HIGH,
|
|
PREVIOUS_LOW,
|
|
PREVIOUS_MEAN, NEXT_MEAN,
|
|
PREVIOUS_RELEASE_DATE,
|
|
PREVIOUS_STANDARD_DEVIATION,
|
|
RELEASE_DATE_FIELD_NAME,
|
|
STANDARD_DEVIATION_FIELD_NAME,
|
|
SID_FIELD_NAME)
|
|
from zipline.pipeline.data import ConsensusEstimates
|
|
from zipline.pipeline.loaders.consensus_estimates import (
|
|
ConsensusEstimatesLoader
|
|
)
|
|
from zipline.pipeline.loaders.blaze import BlazeConsensusEstimatesLoader
|
|
from zipline.pipeline.loaders.utils import (
|
|
zip_with_floats
|
|
)
|
|
from zipline.testing.fixtures import (
|
|
ZiplineTestCase,
|
|
WithNextAndPreviousEventDataLoader
|
|
)
|
|
|
|
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),
|
|
NEXT_RELEASE_DATE: self.get_expected_next_event_dates(dates)
|
|
}
|
|
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
|
|
)
|
|
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, {})
|