mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-13 17:42:42 +08:00
Maya Tydykov
1102 lines
41 KiB
Python
1102 lines
41 KiB
Python
import blaze as bz
|
|
import itertools
|
|
from nose.tools import assert_true
|
|
from nose_parameterized import parameterized
|
|
import numpy as np
|
|
from numpy.testing import assert_array_equal
|
|
import pandas as pd
|
|
from toolz import merge
|
|
|
|
from zipline.pipeline import SimplePipelineEngine, Pipeline, CustomFactor
|
|
from zipline.pipeline.common import (
|
|
EVENT_DATE_FIELD_NAME,
|
|
FISCAL_QUARTER_FIELD_NAME,
|
|
FISCAL_YEAR_FIELD_NAME,
|
|
SID_FIELD_NAME,
|
|
TS_FIELD_NAME,
|
|
)
|
|
from zipline.pipeline.data import DataSet
|
|
from zipline.pipeline.data import Column
|
|
from zipline.pipeline.loaders.blaze.estimates import (
|
|
BlazeNextEstimatesLoader,
|
|
BlazePreviousEstimatesLoader
|
|
)
|
|
from zipline.pipeline.loaders.earnings_estimates import (
|
|
INVALID_NUM_QTRS_MESSAGE,
|
|
NextEarningsEstimatesLoader,
|
|
normalize_quarters,
|
|
PreviousEarningsEstimatesLoader,
|
|
split_normalized_quarters,
|
|
)
|
|
from zipline.testing.fixtures import (
|
|
WithAssetFinder,
|
|
WithTradingSessions,
|
|
ZiplineTestCase,
|
|
)
|
|
from zipline.testing.predicates import assert_equal, assert_raises_regex
|
|
from zipline.utils.numpy_utils import datetime64ns_dtype
|
|
from zipline.utils.numpy_utils import float64_dtype
|
|
|
|
|
|
class Estimates(DataSet):
|
|
event_date = Column(dtype=datetime64ns_dtype)
|
|
fiscal_quarter = Column(dtype=float64_dtype)
|
|
fiscal_year = Column(dtype=float64_dtype)
|
|
estimate = Column(dtype=float64_dtype)
|
|
|
|
|
|
def QuartersEstimates(announcements_out):
|
|
class QtrEstimates(Estimates):
|
|
num_announcements = announcements_out
|
|
name = Estimates
|
|
return QtrEstimates
|
|
|
|
|
|
def QuartersEstimatesNoNumQuartersAttr(num_qtr):
|
|
class QtrEstimates(Estimates):
|
|
name = Estimates
|
|
return QtrEstimates
|
|
|
|
|
|
class WithEstimates(WithTradingSessions, WithAssetFinder):
|
|
"""
|
|
ZiplineTestCase mixin providing cls.loader and cls.events as class
|
|
level fixtures.
|
|
|
|
|
|
Methods
|
|
-------
|
|
make_loader(events, columns) -> PipelineLoader
|
|
Method which returns the loader to be used throughout tests.
|
|
|
|
events : pd.DataFrame
|
|
The raw events to be used as input to the pipeline loader.
|
|
columns : dict[str -> str]
|
|
The dictionary mapping the names of BoundColumns to the
|
|
associated column name in the events DataFrame.
|
|
"""
|
|
|
|
# Short window defined in order for test to run faster.
|
|
START_DATE = pd.Timestamp('2014-12-28')
|
|
END_DATE = pd.Timestamp('2015-02-04')
|
|
|
|
@classmethod
|
|
def make_loader(cls, events, columns):
|
|
raise NotImplementedError('make_loader')
|
|
|
|
@classmethod
|
|
def make_events(cls):
|
|
raise NotImplementedError('make_events')
|
|
|
|
@classmethod
|
|
def get_sids(cls):
|
|
return cls.events[SID_FIELD_NAME].unique()
|
|
|
|
@classmethod
|
|
def init_class_fixtures(cls):
|
|
cls.events = cls.make_events()
|
|
cls.ASSET_FINDER_EQUITY_SIDS = cls.get_sids()
|
|
cls.columns = {
|
|
Estimates.event_date: 'event_date',
|
|
Estimates.fiscal_quarter: 'fiscal_quarter',
|
|
Estimates.fiscal_year: 'fiscal_year',
|
|
Estimates.estimate: 'estimate'
|
|
}
|
|
cls.loader = cls.make_loader(cls.events, {column.name: val for
|
|
column, val in
|
|
cls.columns.items()})
|
|
cls.ASSET_FINDER_EQUITY_SYMBOLS = [
|
|
's' + str(n) for n in cls.ASSET_FINDER_EQUITY_SIDS
|
|
]
|
|
super(WithEstimates, cls).init_class_fixtures()
|
|
|
|
|
|
class WithWrongLoaderDefinition(WithEstimates):
|
|
"""
|
|
ZiplineTestCase mixin providing cls.events as a class level fixture and
|
|
defining a test for all inheritors to use.
|
|
|
|
Attributes
|
|
----------
|
|
events : pd.DataFrame
|
|
A simple DataFrame with columns needed for estimates and a single sid
|
|
and no other data.
|
|
|
|
Tests
|
|
------
|
|
test_wrong_num_announcements_passed()
|
|
Tests that loading with an incorrect quarter number raises an error.
|
|
test_no_num_announcements_attr()
|
|
Tests that the loader throws an AssertionError if the dataset being
|
|
loaded has no `num_announcements` attribute.
|
|
"""
|
|
|
|
@classmethod
|
|
def make_events(cls):
|
|
return pd.DataFrame({SID_FIELD_NAME: 0},
|
|
columns=[SID_FIELD_NAME,
|
|
TS_FIELD_NAME,
|
|
EVENT_DATE_FIELD_NAME,
|
|
FISCAL_QUARTER_FIELD_NAME,
|
|
FISCAL_YEAR_FIELD_NAME,
|
|
'estimate'],
|
|
index=[0])
|
|
|
|
def test_wrong_num_announcements_passed(self):
|
|
bad_dataset1 = QuartersEstimates(-1)
|
|
bad_dataset2 = QuartersEstimates(-2)
|
|
good_dataset = QuartersEstimates(1)
|
|
engine = SimplePipelineEngine(
|
|
lambda x: self.loader,
|
|
self.trading_days,
|
|
self.asset_finder,
|
|
)
|
|
columns = {c.name + str(dataset.num_announcements): c.latest
|
|
for dataset in (bad_dataset1,
|
|
bad_dataset2,
|
|
good_dataset)
|
|
for c in dataset.columns}
|
|
p = Pipeline(columns)
|
|
|
|
with self.assertRaises(ValueError) as e:
|
|
engine.run_pipeline(
|
|
p,
|
|
start_date=self.trading_days[0],
|
|
end_date=self.trading_days[-1],
|
|
)
|
|
assert_raises_regex(e, INVALID_NUM_QTRS_MESSAGE % "-1,-2")
|
|
|
|
def test_no_num_announcements_attr(self):
|
|
dataset = QuartersEstimatesNoNumQuartersAttr(1)
|
|
engine = SimplePipelineEngine(
|
|
lambda x: self.loader,
|
|
self.trading_days,
|
|
self.asset_finder,
|
|
)
|
|
p = Pipeline({c.name: c.latest for c in dataset.columns})
|
|
|
|
with self.assertRaises(AttributeError):
|
|
engine.run_pipeline(
|
|
p,
|
|
start_date=self.trading_days[0],
|
|
end_date=self.trading_days[-1],
|
|
)
|
|
|
|
|
|
class PreviousWithWrongNumQuarters(WithWrongLoaderDefinition,
|
|
ZiplineTestCase):
|
|
"""
|
|
Tests that previous quarter loader correctly breaks if an incorrect
|
|
number of quarters is passed.
|
|
"""
|
|
@classmethod
|
|
def make_loader(cls, events, columns):
|
|
return PreviousEarningsEstimatesLoader(events, columns)
|
|
|
|
|
|
class NextWithWrongNumQuarters(WithWrongLoaderDefinition,
|
|
ZiplineTestCase):
|
|
"""
|
|
Tests that next quarter loader correctly breaks if an incorrect
|
|
number of quarters is passed.
|
|
"""
|
|
@classmethod
|
|
def make_loader(cls, events, columns):
|
|
return NextEarningsEstimatesLoader(events, columns)
|
|
|
|
|
|
class WithEstimatesTimeZero(WithEstimates):
|
|
"""
|
|
ZiplineTestCase mixin providing cls.events as a class level fixture and
|
|
defining a test for all inheritors to use.
|
|
|
|
Attributes
|
|
----------
|
|
cls.events : pd.DataFrame
|
|
Generated dynamically in order to test inter-leavings of estimates and
|
|
event dates for multiple quarters to make sure that we select the
|
|
right immediate 'next' or 'previous' quarter relative to each date -
|
|
i.e., the right 'time zero' on the timeline. We care about selecting
|
|
the right 'time zero' because we use that to calculate which quarter's
|
|
data needs to be returned for each day.
|
|
|
|
Methods
|
|
-------
|
|
get_expected_estimate(q1_knowledge,
|
|
q2_knowledge,
|
|
comparable_date) -> pd.DataFrame
|
|
Retrieves the expected estimate given the latest knowledge about each
|
|
quarter and the date on which the estimate is being requested. If
|
|
there is no expected estimate, returns an empty DataFrame.
|
|
|
|
Tests
|
|
------
|
|
test_estimates()
|
|
Tests that we get the right 'time zero' value on each day for each
|
|
sid and for each column.
|
|
"""
|
|
# Shorter date range for performance
|
|
END_DATE = pd.Timestamp('2015-01-28')
|
|
|
|
q1_knowledge_dates = [pd.Timestamp('2015-01-01'),
|
|
pd.Timestamp('2015-01-04'),
|
|
pd.Timestamp('2015-01-07'),
|
|
pd.Timestamp('2015-01-11')]
|
|
q2_knowledge_dates = [pd.Timestamp('2015-01-14'),
|
|
pd.Timestamp('2015-01-17'),
|
|
pd.Timestamp('2015-01-20'),
|
|
pd.Timestamp('2015-01-23')]
|
|
# We want to model the possibility of an estimate predicting a release date
|
|
# that doesn't match the actual release. This could be done by dynamically
|
|
# generating more combinations with different release dates, but that
|
|
# significantly increases the amount of time it takes to run the tests.
|
|
# These hard-coded cases are sufficient to know that we can update our
|
|
# beliefs when we get new information.
|
|
q1_release_dates = [pd.Timestamp('2015-01-13'),
|
|
pd.Timestamp('2015-01-14')] # One day late
|
|
q2_release_dates = [pd.Timestamp('2015-01-25'), # One day early
|
|
pd.Timestamp('2015-01-26')]
|
|
|
|
@classmethod
|
|
def make_events(cls):
|
|
"""
|
|
In order to determine which estimate we care about for a particular
|
|
sid, we need to look at all estimates that we have for that sid and
|
|
their associated event dates.
|
|
|
|
We define q1 < q2, and thus event1 < event2 since event1 occurs
|
|
during q1 and event2 occurs during q2 and we assume that there can
|
|
only be 1 event per quarter. We assume that there can be multiple
|
|
estimates per quarter leading up to the event. We assume that estimates
|
|
will not surpass the relevant event date. We will look at 2 estimates
|
|
for an event before the event occurs, since that is the simplest
|
|
scenario that covers the interesting edge cases:
|
|
- estimate values changing
|
|
- a release date changing
|
|
- estimates for different quarters interleaving
|
|
|
|
Thus, we generate all possible inter-leavings of 2 estimates per
|
|
quarter-event where estimate1 < estimate2 and all estimates are < the
|
|
relevant event and assign each of these inter-leavings to a
|
|
different sid.
|
|
"""
|
|
|
|
sid_estimates = []
|
|
sid_releases = []
|
|
# We want all permutations of 2 knowledge dates per quarter.
|
|
it = enumerate(
|
|
itertools.permutations(cls.q1_knowledge_dates +
|
|
cls.q2_knowledge_dates,
|
|
4)
|
|
)
|
|
for sid, (q1e1, q1e2, q2e1, q2e2) in it:
|
|
# We're assuming that estimates must come before the relevant
|
|
# release.
|
|
if (q1e1 < q1e2 and
|
|
q2e1 < q2e2 and
|
|
# All estimates are < Q2's event, so just constrain Q1
|
|
# estimates.
|
|
q1e1 < cls.q1_release_dates[0] and
|
|
q1e2 < cls.q1_release_dates[0]):
|
|
sid_estimates.append(cls.create_estimates_df(q1e1,
|
|
q1e2,
|
|
q2e1,
|
|
q2e2,
|
|
sid))
|
|
sid_releases.append(cls.create_releases_df(sid))
|
|
return pd.concat(sid_estimates +
|
|
sid_releases).reset_index(drop=True)
|
|
|
|
@classmethod
|
|
def get_sids(cls):
|
|
sids = cls.events[SID_FIELD_NAME].unique()
|
|
# Tack on an extra sid to make sure that sids with no data are
|
|
# included but have all-null columns.
|
|
return list(sids) + [max(sids) + 1]
|
|
|
|
@classmethod
|
|
def create_releases_df(cls, sid):
|
|
# Final release dates never change. The quarters have very tight date
|
|
# ranges in order to reduce the number of dates we need to iterate
|
|
# through when testing.
|
|
return pd.DataFrame({
|
|
TS_FIELD_NAME: [pd.Timestamp('2015-01-13'),
|
|
pd.Timestamp('2015-01-26')],
|
|
EVENT_DATE_FIELD_NAME: [pd.Timestamp('2015-01-13'),
|
|
pd.Timestamp('2015-01-26')],
|
|
'estimate': [0.5, 0.8],
|
|
FISCAL_QUARTER_FIELD_NAME: [1.0, 2.0],
|
|
FISCAL_YEAR_FIELD_NAME: [2015.0, 2015.0],
|
|
SID_FIELD_NAME: sid
|
|
})
|
|
|
|
@classmethod
|
|
def create_estimates_df(cls,
|
|
q1e1,
|
|
q1e2,
|
|
q2e1,
|
|
q2e2,
|
|
sid):
|
|
return pd.DataFrame({
|
|
EVENT_DATE_FIELD_NAME: cls.q1_release_dates + cls.q2_release_dates,
|
|
'estimate': [.1, .2, .3, .4],
|
|
FISCAL_QUARTER_FIELD_NAME: [1.0, 1.0, 2.0, 2.0],
|
|
FISCAL_YEAR_FIELD_NAME: [2015.0, 2015.0, 2015.0, 2015.0],
|
|
TS_FIELD_NAME: [q1e1, q1e2, q2e1, q2e2],
|
|
SID_FIELD_NAME: sid,
|
|
})
|
|
|
|
@classmethod
|
|
def init_class_fixtures(cls):
|
|
super(WithEstimatesTimeZero, cls).init_class_fixtures()
|
|
|
|
def get_expected_estimate(self,
|
|
q1_knowledge,
|
|
q2_knowledge,
|
|
comparable_date):
|
|
return pd.DataFrame()
|
|
|
|
def test_estimates(self):
|
|
dataset = QuartersEstimates(1)
|
|
engine = SimplePipelineEngine(
|
|
lambda x: self.loader,
|
|
self.trading_days,
|
|
self.asset_finder,
|
|
)
|
|
results = engine.run_pipeline(
|
|
Pipeline({c.name: c.latest for c in dataset.columns}),
|
|
start_date=self.trading_days[1],
|
|
end_date=self.trading_days[-2],
|
|
)
|
|
for sid in self.ASSET_FINDER_EQUITY_SIDS:
|
|
sid_estimates = results.xs(sid, level=1)
|
|
# Separate assertion for all-null DataFrame to avoid setting
|
|
# column dtypes on `all_expected`.
|
|
if sid == max(self.ASSET_FINDER_EQUITY_SIDS):
|
|
assert_true(sid_estimates.isnull().all().all())
|
|
else:
|
|
ts_sorted_estimates = self.events[
|
|
self.events[SID_FIELD_NAME] == sid
|
|
].sort(TS_FIELD_NAME)
|
|
q1_knowledge = ts_sorted_estimates[
|
|
ts_sorted_estimates[FISCAL_QUARTER_FIELD_NAME] == 1
|
|
]
|
|
q2_knowledge = ts_sorted_estimates[
|
|
ts_sorted_estimates[FISCAL_QUARTER_FIELD_NAME] == 2
|
|
]
|
|
all_expected = pd.concat(
|
|
[self.get_expected_estimate(
|
|
q1_knowledge[q1_knowledge[TS_FIELD_NAME] <=
|
|
date.tz_localize(None)],
|
|
q2_knowledge[q2_knowledge[TS_FIELD_NAME] <=
|
|
date.tz_localize(None)],
|
|
date.tz_localize(None),
|
|
).set_index([[date]]) for date in sid_estimates.index],
|
|
axis=0)
|
|
assert_equal(all_expected[sid_estimates.columns],
|
|
sid_estimates)
|
|
|
|
|
|
class NextEstimate(WithEstimatesTimeZero, ZiplineTestCase):
|
|
@classmethod
|
|
def make_loader(cls, events, columns):
|
|
return NextEarningsEstimatesLoader(events, columns)
|
|
|
|
def get_expected_estimate(self,
|
|
q1_knowledge,
|
|
q2_knowledge,
|
|
comparable_date):
|
|
# If our latest knowledge of q1 is that the release is
|
|
# happening on this simulation date or later, then that's
|
|
# the estimate we want to use.
|
|
if (not q1_knowledge.empty and
|
|
q1_knowledge[EVENT_DATE_FIELD_NAME].iloc[-1] >=
|
|
comparable_date):
|
|
return q1_knowledge.iloc[-1:]
|
|
# If q1 has already happened or we don't know about it
|
|
# yet and our latest knowledge indicates that q2 hasn't
|
|
# happened yet, then that's the estimate we want to use.
|
|
elif (not q2_knowledge.empty and
|
|
q2_knowledge[EVENT_DATE_FIELD_NAME].iloc[-1] >=
|
|
comparable_date):
|
|
return q2_knowledge.iloc[-1:]
|
|
return pd.DataFrame(columns=q1_knowledge.columns,
|
|
index=[comparable_date])
|
|
|
|
|
|
class BlazeNextEstimateLoaderTestCase(NextEstimate):
|
|
"""
|
|
Run the same tests as EventsLoaderTestCase, but using a BlazeEventsLoader.
|
|
"""
|
|
|
|
@classmethod
|
|
def make_loader(cls, events, columns):
|
|
return BlazeNextEstimatesLoader(
|
|
bz.data(events),
|
|
columns,
|
|
)
|
|
|
|
|
|
class PreviousEstimate(WithEstimatesTimeZero, ZiplineTestCase):
|
|
@classmethod
|
|
def make_loader(cls, events, columns):
|
|
return PreviousEarningsEstimatesLoader(events, columns)
|
|
|
|
def get_expected_estimate(self,
|
|
q1_knowledge,
|
|
q2_knowledge,
|
|
comparable_date):
|
|
|
|
# The expected estimate will be for q2 if the last thing
|
|
# we've seen is that the release date already happened.
|
|
# Otherwise, it'll be for q1, as long as the release date
|
|
# for q1 has already happened.
|
|
if (not q2_knowledge.empty and
|
|
q2_knowledge[EVENT_DATE_FIELD_NAME].iloc[-1] <=
|
|
comparable_date):
|
|
return q2_knowledge.iloc[-1:]
|
|
elif (not q1_knowledge.empty and
|
|
q1_knowledge[EVENT_DATE_FIELD_NAME].iloc[-1] <=
|
|
comparable_date):
|
|
return q1_knowledge.iloc[-1:]
|
|
return pd.DataFrame(columns=q1_knowledge.columns,
|
|
index=[comparable_date])
|
|
|
|
|
|
class BlazePreviousEstimateLoaderTestCase(PreviousEstimate):
|
|
"""
|
|
Run the same tests as EventsLoaderTestCase, but using a BlazeEventsLoader.
|
|
"""
|
|
|
|
@classmethod
|
|
def make_loader(cls, events, columns):
|
|
return BlazePreviousEstimatesLoader(
|
|
bz.data(events),
|
|
columns,
|
|
)
|
|
|
|
|
|
class WithEstimateMultipleQuarters(WithEstimates):
|
|
"""
|
|
ZiplineTestCase mixin providing cls.events, cls.make_expected_out as
|
|
class-level fixtures and self.test_multiple_qtrs_requested as a test.
|
|
|
|
Attributes
|
|
----------
|
|
events : pd.DataFrame
|
|
Simple DataFrame with estimates for 2 quarters for a single sid.
|
|
|
|
Methods
|
|
-------
|
|
make_expected_out() --> pd.DataFrame
|
|
Returns the DataFrame that is expected as a result of running a
|
|
Pipeline where estimates are requested for multiple quarters out.
|
|
fill_expected_out(expected)
|
|
Fills the expected DataFrame with data.
|
|
|
|
Tests
|
|
------
|
|
test_multiple_qtrs_requested()
|
|
Runs a Pipeline that calculate which estimates for multiple quarters
|
|
out and checks that the returned columns contain data for the correct
|
|
number of quarters out.
|
|
"""
|
|
|
|
@classmethod
|
|
def make_events(cls):
|
|
return pd.DataFrame({
|
|
SID_FIELD_NAME: [0] * 2,
|
|
TS_FIELD_NAME: [pd.Timestamp('2015-01-01'),
|
|
pd.Timestamp('2015-01-06')],
|
|
EVENT_DATE_FIELD_NAME: [pd.Timestamp('2015-01-10'),
|
|
pd.Timestamp('2015-01-20')],
|
|
'estimate': [1., 2.],
|
|
FISCAL_QUARTER_FIELD_NAME: [1, 2],
|
|
FISCAL_YEAR_FIELD_NAME: [2015, 2015]
|
|
})
|
|
|
|
@classmethod
|
|
def init_class_fixtures(cls):
|
|
super(WithEstimateMultipleQuarters, cls).init_class_fixtures()
|
|
cls.expected_out = cls.make_expected_out()
|
|
|
|
@classmethod
|
|
def make_expected_out(cls):
|
|
expected = pd.DataFrame(columns=[cls.columns[col] + '1'
|
|
for col in cls.columns] +
|
|
[cls.columns[col] + '2'
|
|
for col in cls.columns],
|
|
index=cls.trading_days)
|
|
|
|
for (col, raw_name), suffix in itertools.product(
|
|
cls.columns.items(), ('1', '2')
|
|
):
|
|
expected_name = raw_name + suffix
|
|
if col.dtype == datetime64ns_dtype:
|
|
expected[expected_name] = pd.to_datetime(
|
|
expected[expected_name]
|
|
)
|
|
else:
|
|
expected[expected_name] = expected[
|
|
expected_name
|
|
].astype(col.dtype)
|
|
cls.fill_expected_out(expected)
|
|
return expected.reindex(cls.trading_days)
|
|
|
|
def test_multiple_qtrs_requested(self):
|
|
dataset1 = QuartersEstimates(1)
|
|
dataset2 = QuartersEstimates(2)
|
|
engine = SimplePipelineEngine(
|
|
lambda x: self.loader,
|
|
self.trading_days,
|
|
self.asset_finder,
|
|
)
|
|
|
|
results = engine.run_pipeline(
|
|
Pipeline(
|
|
merge([{c.name + '1': c.latest for c in dataset1.columns},
|
|
{c.name + '2': c.latest for c in dataset2.columns}])
|
|
),
|
|
start_date=self.trading_days[0],
|
|
end_date=self.trading_days[-1],
|
|
)
|
|
q1_columns = [col.name + '1' for col in self.columns]
|
|
q2_columns = [col.name + '2' for col in self.columns]
|
|
|
|
# We now expect a column for 1 quarter out and a column for 2
|
|
# quarters out for each of the dataset columns.
|
|
assert_equal(sorted(np.array(q1_columns + q2_columns)),
|
|
sorted(results.columns.values))
|
|
assert_equal(self.expected_out.sort(axis=1),
|
|
results.xs(0, level=1).sort(axis=1))
|
|
|
|
|
|
class NextEstimateMultipleQuarters(
|
|
WithEstimateMultipleQuarters, ZiplineTestCase
|
|
):
|
|
@classmethod
|
|
def make_loader(cls, events, columns):
|
|
return NextEarningsEstimatesLoader(events, columns)
|
|
|
|
@classmethod
|
|
def fill_expected_out(cls, expected):
|
|
# Fill columns for 1 Q out
|
|
for raw_name in cls.columns.values():
|
|
expected.loc[
|
|
pd.Timestamp('2015-01-01'):pd.Timestamp('2015-01-11'),
|
|
raw_name + '1'
|
|
] = cls.events[raw_name].iloc[0]
|
|
expected.loc[
|
|
pd.Timestamp('2015-01-11'):pd.Timestamp('2015-01-20'),
|
|
raw_name + '1'
|
|
] = cls.events[raw_name].iloc[1]
|
|
|
|
# Fill columns for 2 Q out
|
|
# We only have an estimate and event date for 2 quarters out before
|
|
# Q1's event happens; after Q1's event, we know 1 Q out but not 2 Qs
|
|
# out.
|
|
for col_name in ['estimate', 'event_date']:
|
|
expected.loc[
|
|
pd.Timestamp('2015-01-06'):pd.Timestamp('2015-01-10'),
|
|
col_name + '2'
|
|
] = cls.events[col_name].iloc[1]
|
|
# But we know what FQ and FY we'd need in both Q1 and Q2
|
|
# because we know which FQ is next and can calculate from there
|
|
expected.loc[
|
|
pd.Timestamp('2015-01-01'):pd.Timestamp('2015-01-09'),
|
|
FISCAL_QUARTER_FIELD_NAME + '2'
|
|
] = 2
|
|
expected.loc[
|
|
pd.Timestamp('2015-01-12'):pd.Timestamp('2015-01-20'),
|
|
FISCAL_QUARTER_FIELD_NAME + '2'
|
|
] = 3
|
|
expected.loc[
|
|
pd.Timestamp('2015-01-01'):pd.Timestamp('2015-01-20'),
|
|
FISCAL_YEAR_FIELD_NAME + '2'
|
|
] = 2015
|
|
|
|
return expected
|
|
|
|
|
|
class PreviousEstimateMultipleQuarters(
|
|
WithEstimateMultipleQuarters,
|
|
ZiplineTestCase
|
|
):
|
|
|
|
@classmethod
|
|
def make_loader(cls, events, columns):
|
|
return PreviousEarningsEstimatesLoader(events, columns)
|
|
|
|
@classmethod
|
|
def fill_expected_out(cls, expected):
|
|
# Fill columns for 1 Q out
|
|
for raw_name in cls.columns.values():
|
|
expected[raw_name + '1'].loc[
|
|
pd.Timestamp('2015-01-12'):pd.Timestamp('2015-01-19')
|
|
] = cls.events[raw_name].iloc[0]
|
|
expected[raw_name + '1'].loc[
|
|
pd.Timestamp('2015-01-20'):
|
|
] = cls.events[raw_name].iloc[1]
|
|
|
|
# Fill columns for 2 Q out
|
|
for col_name in ['estimate', 'event_date']:
|
|
expected[col_name + '2'].loc[
|
|
pd.Timestamp('2015-01-20'):
|
|
] = cls.events[col_name].iloc[0]
|
|
expected[
|
|
FISCAL_QUARTER_FIELD_NAME + '2'
|
|
].loc[pd.Timestamp('2015-01-12'):pd.Timestamp('2015-01-20')] = 4
|
|
expected[
|
|
FISCAL_YEAR_FIELD_NAME + '2'
|
|
].loc[pd.Timestamp('2015-01-12'):pd.Timestamp('2015-01-20')] = 2014
|
|
expected[
|
|
FISCAL_QUARTER_FIELD_NAME + '2'
|
|
].loc[pd.Timestamp('2015-01-20'):] = 1
|
|
expected[
|
|
FISCAL_YEAR_FIELD_NAME + '2'
|
|
].loc[pd.Timestamp('2015-01-20'):] = 2015
|
|
return expected
|
|
|
|
|
|
class WithVaryingNumEstimates(WithEstimates):
|
|
"""
|
|
ZiplineTestCase mixin providing fixtures and a test to ensure that we
|
|
have the correct overwrites when the event date changes. We want to make
|
|
sure that if we have a quarter with an event date that gets pushed back,
|
|
we don't start overwriting for the next quarter early. Likewise,
|
|
if we have a quarter with an event date that gets pushed forward, we want
|
|
to make sure that we start applying adjustments at the appropriate, earlier
|
|
date, rather than the later date.
|
|
|
|
Methods
|
|
-------
|
|
assert_compute()
|
|
Defines how to determine that results computed for the `SomeFactor`
|
|
factor are correct.
|
|
|
|
Tests
|
|
-----
|
|
test_windows_with_varying_num_estimates()
|
|
Tests that we create the correct overwrites from 2015-01-13 to
|
|
2015-01-14 regardless of how event dates were updated for each
|
|
quarter for each sid.
|
|
"""
|
|
|
|
@classmethod
|
|
def make_events(cls):
|
|
return pd.DataFrame({
|
|
SID_FIELD_NAME: [0] * 3 + [1] * 3,
|
|
TS_FIELD_NAME: [pd.Timestamp('2015-01-09'),
|
|
pd.Timestamp('2015-01-12'),
|
|
pd.Timestamp('2015-01-13')] * 2,
|
|
EVENT_DATE_FIELD_NAME: [pd.Timestamp('2015-01-12'),
|
|
pd.Timestamp('2015-01-13'),
|
|
pd.Timestamp('2015-01-20'),
|
|
pd.Timestamp('2015-01-13'),
|
|
pd.Timestamp('2015-01-12'),
|
|
pd.Timestamp('2015-01-20')],
|
|
'estimate': [11., 12., 21.] * 2,
|
|
FISCAL_QUARTER_FIELD_NAME: [1, 1, 2] * 2,
|
|
FISCAL_YEAR_FIELD_NAME: [2015] * 6
|
|
})
|
|
|
|
@classmethod
|
|
def assert_compute(cls, estimate, today):
|
|
raise NotImplementedError('assert_compute')
|
|
|
|
def test_windows_with_varying_num_estimates(self):
|
|
dataset = QuartersEstimates(1)
|
|
assert_compute = self.assert_compute
|
|
|
|
class SomeFactor(CustomFactor):
|
|
inputs = [dataset.estimate]
|
|
window_length = 3
|
|
|
|
def compute(self, today, assets, out, estimate):
|
|
assert_compute(estimate, today)
|
|
|
|
engine = SimplePipelineEngine(
|
|
lambda x: self.loader,
|
|
self.trading_days,
|
|
self.asset_finder,
|
|
)
|
|
engine.run_pipeline(
|
|
Pipeline({'est': SomeFactor()}),
|
|
start_date=pd.Timestamp('2015-01-13', tz='utc'),
|
|
# last event date we have
|
|
end_date=pd.Timestamp('2015-01-14', tz='utc'),
|
|
)
|
|
|
|
|
|
class PreviousVaryingNumEstimates(
|
|
WithVaryingNumEstimates,
|
|
ZiplineTestCase
|
|
):
|
|
def assert_compute(self, estimate, today):
|
|
if today == pd.Timestamp('2015-01-13', tz='utc'):
|
|
assert_array_equal(estimate[:, 0],
|
|
np.array([np.NaN, np.NaN, 12]))
|
|
assert_array_equal(estimate[:, 1],
|
|
np.array([np.NaN, 12, 12]))
|
|
else:
|
|
assert_array_equal(estimate[:, 0],
|
|
np.array([np.NaN, 12, 12]))
|
|
assert_array_equal(estimate[:, 1],
|
|
np.array([12, 12, 12]))
|
|
|
|
@classmethod
|
|
def make_loader(cls, events, columns):
|
|
return PreviousEarningsEstimatesLoader(events, columns)
|
|
|
|
|
|
class NextVaryingNumEstimates(
|
|
WithVaryingNumEstimates,
|
|
ZiplineTestCase
|
|
):
|
|
|
|
def assert_compute(self, estimate, today):
|
|
if today == pd.Timestamp('2015-01-13', tz='utc'):
|
|
assert_array_equal(estimate[:, 0],
|
|
np.array([11, 12, 12]))
|
|
assert_array_equal(estimate[:, 1],
|
|
np.array([np.NaN, np.NaN, 21]))
|
|
else:
|
|
assert_array_equal(estimate[:, 0],
|
|
np.array([np.NaN, 21, 21]))
|
|
assert_array_equal(estimate[:, 1],
|
|
np.array([np.NaN, 21, 21]))
|
|
|
|
@classmethod
|
|
def make_loader(cls, events, columns):
|
|
return NextEarningsEstimatesLoader(events, columns)
|
|
|
|
|
|
class WithEstimateWindows(WithEstimates):
|
|
"""
|
|
ZiplineTestCase mixin providing fixures and a test to test running a
|
|
Pipeline with an estimates loader over differently-sized windows.
|
|
|
|
Attributes
|
|
----------
|
|
events : pd.DataFrame
|
|
DataFrame with estimates for 2 quarters for 2 sids.
|
|
window_test_start_date : pd.Timestamp
|
|
The date from which the window should start.
|
|
timelines : dict[int -> pd.DataFrame]
|
|
A dictionary mapping to the number of quarters out to
|
|
snapshots of how the data should look on each date in the date range.
|
|
|
|
Methods
|
|
-------
|
|
make_expected_timelines() -> dict[int -> pd.DataFrame]
|
|
Creates a dictionary of expected data. See `timelines`, above.
|
|
|
|
Tests
|
|
-----
|
|
test_estimate_windows_at_quarter_boundaries()
|
|
Tests that we overwrite values with the correct quarter's estimate at
|
|
the correct dates when we have a factor that asks for a window of data.
|
|
"""
|
|
|
|
window_test_start_date = pd.Timestamp('2015-01-05')
|
|
critical_dates = [pd.Timestamp('2015-01-09', tz='utc'),
|
|
pd.Timestamp('2015-01-12', tz='utc'),
|
|
pd.Timestamp('2015-01-15', tz='utc'),
|
|
pd.Timestamp('2015-01-20', tz='utc')]
|
|
# window length, starting date, num quarters out, timeline. Parameterizes
|
|
# over number of quarters out.
|
|
window_test_cases = list(itertools.product(critical_dates, (1, 2)))
|
|
|
|
@classmethod
|
|
def make_events(cls):
|
|
sid_0_timeline = pd.DataFrame({
|
|
TS_FIELD_NAME: [pd.Timestamp('2015-01-05'),
|
|
pd.Timestamp('2015-01-07'),
|
|
pd.Timestamp('2015-01-05'),
|
|
pd.Timestamp('2015-01-17')],
|
|
EVENT_DATE_FIELD_NAME:
|
|
[pd.Timestamp('2015-01-10'),
|
|
pd.Timestamp('2015-01-10'),
|
|
pd.Timestamp('2015-01-20'),
|
|
pd.Timestamp('2015-01-20')],
|
|
'estimate': [100., 101.] + [200., 201.],
|
|
FISCAL_QUARTER_FIELD_NAME: [1] * 2 + [2] * 2,
|
|
FISCAL_YEAR_FIELD_NAME: 2015,
|
|
SID_FIELD_NAME: 0,
|
|
})
|
|
|
|
sid_1_timeline = pd.DataFrame({
|
|
TS_FIELD_NAME: [pd.Timestamp('2015-01-09'),
|
|
pd.Timestamp('2015-01-12'),
|
|
pd.Timestamp('2015-01-09'),
|
|
pd.Timestamp('2015-01-15')],
|
|
EVENT_DATE_FIELD_NAME:
|
|
[pd.Timestamp('2015-01-12'), pd.Timestamp('2015-01-12'),
|
|
pd.Timestamp('2015-01-15'), pd.Timestamp('2015-01-15')],
|
|
'estimate': [110., 111.] + [310., 311.],
|
|
FISCAL_QUARTER_FIELD_NAME: [1] * 2 + [3] * 2,
|
|
FISCAL_YEAR_FIELD_NAME: 2015,
|
|
SID_FIELD_NAME: 1
|
|
})
|
|
|
|
# Extra sid to make sure we have correct overwrites when sid quarter
|
|
# boundaries collide.
|
|
sid_3_timeline = pd.DataFrame({
|
|
TS_FIELD_NAME: [pd.Timestamp('2015-01-05'),
|
|
pd.Timestamp('2015-01-07'),
|
|
pd.Timestamp('2015-01-05'),
|
|
pd.Timestamp('2015-01-17')],
|
|
EVENT_DATE_FIELD_NAME:
|
|
[pd.Timestamp('2015-01-10'),
|
|
pd.Timestamp('2015-01-10'),
|
|
pd.Timestamp('2015-01-20'),
|
|
pd.Timestamp('2015-01-20')],
|
|
'estimate': [120., 121.] + [220., 221.],
|
|
FISCAL_QUARTER_FIELD_NAME: [1] * 2 + [2] * 2,
|
|
FISCAL_YEAR_FIELD_NAME: 2015,
|
|
SID_FIELD_NAME: 2
|
|
})
|
|
return pd.concat([sid_0_timeline, sid_1_timeline, sid_3_timeline])
|
|
|
|
@classmethod
|
|
def make_expected_timelines(cls):
|
|
return {}
|
|
|
|
@classmethod
|
|
def init_class_fixtures(cls):
|
|
super(WithEstimateWindows, cls).init_class_fixtures()
|
|
cls.timelines = cls.make_expected_timelines()
|
|
|
|
@classmethod
|
|
def create_expected_df(cls, tuples, end_date):
|
|
"""
|
|
Given a list of tuples of new data we get for each sid on each critical
|
|
date (when information changes), create a DataFrame that fills that
|
|
data through a date range ending at `end_date`.
|
|
"""
|
|
df = pd.DataFrame(tuples,
|
|
columns=[SID_FIELD_NAME,
|
|
'estimate',
|
|
'knowledge_date'])
|
|
df = df.pivot_table(columns=SID_FIELD_NAME,
|
|
values='estimate',
|
|
index='knowledge_date')
|
|
df = df.reindex(
|
|
pd.date_range(cls.window_test_start_date, end_date)
|
|
)
|
|
# Index name is lost during reindex.
|
|
df.index = df.index.rename('knowledge_date')
|
|
df['at_date'] = end_date.tz_localize('utc')
|
|
df = df.set_index(['at_date', df.index.tz_localize('utc')]).ffill()
|
|
return df
|
|
|
|
@parameterized.expand(window_test_cases)
|
|
def test_estimate_windows_at_quarter_boundaries(self,
|
|
start_idx,
|
|
num_announcements_out):
|
|
dataset = QuartersEstimates(num_announcements_out)
|
|
trading_days = self.trading_days
|
|
timelines = self.timelines
|
|
# The window length should be from the starting index back to the first
|
|
# date on which we got data. The goal is to ensure that as we
|
|
# progress through the timeline, all data we got, starting from that
|
|
# first date, is correctly overwritten.
|
|
window_len = (
|
|
self.trading_days.get_loc(start_idx) -
|
|
self.trading_days.get_loc(self.window_test_start_date) + 1
|
|
)
|
|
|
|
class SomeFactor(CustomFactor):
|
|
inputs = [dataset.estimate]
|
|
window_length = window_len
|
|
|
|
def compute(self, today, assets, out, estimate):
|
|
today_idx = trading_days.get_loc(today)
|
|
today_timeline = timelines[
|
|
num_announcements_out
|
|
].loc[today].reindex(
|
|
trading_days[:today_idx + 1]
|
|
).values
|
|
timeline_start_idx = (len(today_timeline) - window_len)
|
|
assert_equal(estimate,
|
|
today_timeline[timeline_start_idx:])
|
|
engine = SimplePipelineEngine(
|
|
lambda x: self.loader,
|
|
self.trading_days,
|
|
self.asset_finder,
|
|
)
|
|
engine.run_pipeline(
|
|
Pipeline({'est': SomeFactor()}),
|
|
start_date=start_idx,
|
|
# last event date we have
|
|
end_date=pd.Timestamp('2015-01-20', tz='utc'),
|
|
)
|
|
|
|
|
|
class PreviousEstimateWindows(WithEstimateWindows, ZiplineTestCase):
|
|
@classmethod
|
|
def make_loader(cls, events, columns):
|
|
return PreviousEarningsEstimatesLoader(events, columns)
|
|
|
|
@classmethod
|
|
def make_expected_timelines(cls):
|
|
oneq_previous = pd.concat([
|
|
cls.create_expected_df(
|
|
[(0, np.NaN, cls.window_test_start_date),
|
|
(1, np.NaN, cls.window_test_start_date),
|
|
(2, np.NaN, cls.window_test_start_date)],
|
|
pd.Timestamp('2015-01-09')
|
|
),
|
|
cls.create_expected_df(
|
|
[(0, 101, pd.Timestamp('2015-01-10')),
|
|
(1, 111, pd.Timestamp('2015-01-12')),
|
|
(2, 121, pd.Timestamp('2015-01-10'))],
|
|
pd.Timestamp('2015-01-12')
|
|
),
|
|
cls.create_expected_df(
|
|
[(0, 101, pd.Timestamp('2015-01-10')),
|
|
(1, 111, pd.Timestamp('2015-01-12')),
|
|
(2, 121, pd.Timestamp('2015-01-10'))],
|
|
pd.Timestamp('2015-01-13')
|
|
),
|
|
cls.create_expected_df(
|
|
[(0, 101, pd.Timestamp('2015-01-10')),
|
|
(1, 111, pd.Timestamp('2015-01-12')),
|
|
(2, 121, pd.Timestamp('2015-01-10'))],
|
|
pd.Timestamp('2015-01-14')
|
|
),
|
|
cls.create_expected_df(
|
|
[(0, 101, pd.Timestamp('2015-01-10')),
|
|
(1, 311, pd.Timestamp('2015-01-15')),
|
|
(2, 121, pd.Timestamp('2015-01-10'))],
|
|
pd.Timestamp('2015-01-15')
|
|
),
|
|
cls.create_expected_df(
|
|
[(0, 101, pd.Timestamp('2015-01-10')),
|
|
(1, 311, pd.Timestamp('2015-01-15')),
|
|
(2, 121, pd.Timestamp('2015-01-10'))],
|
|
pd.Timestamp('2015-01-16')
|
|
),
|
|
cls.create_expected_df(
|
|
[(0, 201, pd.Timestamp('2015-01-17')),
|
|
(1, 311, pd.Timestamp('2015-01-15')),
|
|
(2, 221, pd.Timestamp('2015-01-17'))],
|
|
pd.Timestamp('2015-01-20')
|
|
),
|
|
])
|
|
|
|
twoq_previous = pd.concat(
|
|
[cls.create_expected_df(
|
|
[(0, np.NaN, cls.window_test_start_date),
|
|
(1, np.NaN, cls.window_test_start_date),
|
|
(2, np.NaN, cls.window_test_start_date)],
|
|
end_date
|
|
) for end_date in pd.date_range('2015-01-09', '2015-01-19')] +
|
|
[cls.create_expected_df(
|
|
[(0, 101, pd.Timestamp('2015-01-20')),
|
|
(1, np.NaN, cls.window_test_start_date),
|
|
(2, 121, pd.Timestamp('2015-01-20'))],
|
|
pd.Timestamp('2015-01-20')
|
|
)]
|
|
)
|
|
return {
|
|
1: oneq_previous,
|
|
2: twoq_previous
|
|
}
|
|
|
|
|
|
class NextEstimateWindows(WithEstimateWindows, ZiplineTestCase):
|
|
@classmethod
|
|
def make_loader(cls, events, columns):
|
|
return NextEarningsEstimatesLoader(events, columns)
|
|
|
|
@classmethod
|
|
def make_expected_timelines(cls):
|
|
oneq_next = pd.concat([
|
|
cls.create_expected_df(
|
|
[(0, 100, cls.window_test_start_date),
|
|
(0, 101, pd.Timestamp('2015-01-07')),
|
|
(1, 110, pd.Timestamp('2015-01-09')),
|
|
(2, 120, cls.window_test_start_date),
|
|
(2, 121, pd.Timestamp('2015-01-07'))],
|
|
pd.Timestamp('2015-01-09')
|
|
),
|
|
cls.create_expected_df(
|
|
[(0, 200, cls.window_test_start_date),
|
|
(1, 110, pd.Timestamp('2015-01-09')),
|
|
(1, 111, pd.Timestamp('2015-01-12')),
|
|
(2, 220, cls.window_test_start_date)],
|
|
pd.Timestamp('2015-01-12')
|
|
),
|
|
cls.create_expected_df(
|
|
[(0, 200, cls.window_test_start_date),
|
|
(1, 310, pd.Timestamp('2015-01-09')),
|
|
(2, 220, cls.window_test_start_date)],
|
|
pd.Timestamp('2015-01-13')
|
|
),
|
|
cls.create_expected_df(
|
|
[(0, 200, cls.window_test_start_date),
|
|
(1, 310, pd.Timestamp('2015-01-09')),
|
|
(2, 220, cls.window_test_start_date)],
|
|
pd.Timestamp('2015-01-14')
|
|
),
|
|
cls.create_expected_df(
|
|
[(0, 200, cls.window_test_start_date),
|
|
(1, 310, pd.Timestamp('2015-01-09')),
|
|
(1, 311, pd.Timestamp('2015-01-15')),
|
|
(2, 220, cls.window_test_start_date)],
|
|
pd.Timestamp('2015-01-15')
|
|
),
|
|
cls.create_expected_df(
|
|
[(0, 200, cls.window_test_start_date),
|
|
(1, np.NaN, cls.window_test_start_date),
|
|
(2, 220, cls.window_test_start_date)],
|
|
pd.Timestamp('2015-01-16')
|
|
),
|
|
cls.create_expected_df(
|
|
[(0, 200, cls.window_test_start_date),
|
|
(0, 201, pd.Timestamp('2015-01-17')),
|
|
(1, np.NaN, cls.window_test_start_date),
|
|
(2, 220, cls.window_test_start_date),
|
|
(2, 221, pd.Timestamp('2015-01-17'))],
|
|
pd.Timestamp('2015-01-20')
|
|
),
|
|
])
|
|
|
|
twoq_next = pd.concat(
|
|
[cls.create_expected_df(
|
|
[(0, 200, pd.Timestamp(cls.window_test_start_date)),
|
|
(1, np.NaN, pd.Timestamp(cls.window_test_start_date)),
|
|
(2, 220, pd.Timestamp(cls.window_test_start_date))],
|
|
pd.Timestamp('2015-01-09')
|
|
)] +
|
|
[cls.create_expected_df(
|
|
[(0, np.NaN, pd.Timestamp(cls.window_test_start_date)),
|
|
(1, np.NaN, pd.Timestamp(cls.window_test_start_date)),
|
|
(2, np.NaN, pd.Timestamp(cls.window_test_start_date))],
|
|
end_date
|
|
) for end_date in pd.date_range('2015-01-12', '2015-01-20')]
|
|
)
|
|
|
|
return {
|
|
1: oneq_next,
|
|
2: twoq_next
|
|
}
|
|
|
|
|
|
class QuarterShiftTestCase(ZiplineTestCase):
|
|
"""
|
|
This tests, in isolation, quarter calculation logic for shifting quarters
|
|
backwards/forwards from a starting point.
|
|
"""
|
|
def test_quarter_normalization(self):
|
|
input_yrs = pd.Series(range(2011, 2015), dtype=np.int64)
|
|
input_qtrs = pd.Series(range(1, 5), dtype=np.int64)
|
|
result_years, result_quarters = split_normalized_quarters(
|
|
normalize_quarters(input_yrs, input_qtrs)
|
|
)
|
|
# Can't use assert_series_equal here with check_names=False
|
|
# because that still fails due to name differences.
|
|
assert_equal(input_yrs, result_years)
|
|
assert_equal(input_qtrs, result_quarters)
|