/tests/pipeline/test_consensus_estimates.py

"""
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, {})