mirror of
https://github.com/wassname/catalyst.git
synced 2026-09-10 11:50:32 +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:
Scott Sanderson
@@ -23,7 +23,6 @@ from .core import ( # noqa
|
||||
empty_asset_finder,
|
||||
empty_assets_db,
|
||||
empty_trading_env,
|
||||
gen_calendars,
|
||||
make_test_handler,
|
||||
make_trade_data_for_asset_info,
|
||||
parameter_space,
|
||||
|
||||
+2
-314
@@ -1,14 +1,10 @@
|
||||
from abc import ABCMeta, abstractproperty
|
||||
import sqlite3
|
||||
from unittest import TestCase
|
||||
|
||||
from contextlib2 import ExitStack
|
||||
from logbook import NullHandler, Logger
|
||||
from nose_parameterized import parameterized
|
||||
from pandas.util.testing import assert_series_equal
|
||||
from six import with_metaclass
|
||||
from toolz import flip
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import responses
|
||||
|
||||
@@ -36,17 +32,9 @@ from ..finance.trading import TradingEnvironment
|
||||
from ..utils import factory
|
||||
from ..utils.classproperty import classproperty
|
||||
from ..utils.final import FinalMeta, final
|
||||
from ..utils.metautils import with_metaclasses
|
||||
from .core import tmp_asset_finder, make_simple_equity_info, gen_calendars
|
||||
from zipline.pipeline import Pipeline, SimplePipelineEngine
|
||||
from .core import tmp_asset_finder, make_simple_equity_info
|
||||
from zipline.pipeline import SimplePipelineEngine
|
||||
from zipline.pipeline.loaders.testing import make_seeded_random_loader
|
||||
from zipline.utils.numpy_utils import make_datetime64D
|
||||
from zipline.utils.numpy_utils import NaTD
|
||||
from zipline.pipeline.common import TS_FIELD_NAME
|
||||
from zipline.pipeline.loaders.utils import (
|
||||
get_values_for_date_ranges,
|
||||
zip_with_dates
|
||||
)
|
||||
from zipline.utils.calendars import (
|
||||
get_calendar,
|
||||
ExchangeTradingSchedule,
|
||||
@@ -908,182 +896,6 @@ class WithAdjustmentReader(WithBcolzDailyBarReader):
|
||||
cls.adjustment_reader = SQLiteAdjustmentReader(conn)
|
||||
|
||||
|
||||
class WithPipelineEventDataLoader(
|
||||
with_metaclasses((type(ZiplineTestCase), ABCMeta), WithAssetFinder)):
|
||||
"""
|
||||
ZiplineTestCase mixin providing common test methods/behaviors for event
|
||||
data loaders.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
loader_type : PipelineLoader
|
||||
The type of loader to use. This must be overridden by subclasses.
|
||||
|
||||
Methods
|
||||
-------
|
||||
get_sids() -> iterable[int]
|
||||
Class method which returns the sids that need to be available to the
|
||||
tests.
|
||||
get_dataset() -> dict[int -> pd.DataFrmae]
|
||||
Class method which returns a mapping from sid to data for that sid.
|
||||
By default this is empty for every sid.
|
||||
pipeline_event_loader_args(dates: pd.DatetimeIndex) -> tuple[any]
|
||||
The arguments to pass to the ``loader_type`` to construct the pipeline
|
||||
loader for this test.
|
||||
"""
|
||||
@classmethod
|
||||
def get_sids(cls):
|
||||
return range(0, 5)
|
||||
|
||||
@classmethod
|
||||
def get_dataset(cls):
|
||||
return {sid: pd.DataFrame() for sid in cls.get_sids()}
|
||||
|
||||
@abstractproperty
|
||||
def loader_type(self):
|
||||
raise NotImplementedError('loader_type')
|
||||
|
||||
@classmethod
|
||||
def make_equity_info(cls):
|
||||
return make_simple_equity_info(
|
||||
cls.get_sids(),
|
||||
start_date=pd.Timestamp('2013-01-01', tz='UTC'),
|
||||
end_date=pd.Timestamp('2015-01-01', tz='UTC'),
|
||||
)
|
||||
|
||||
def pipeline_event_loader_args(self, dates):
|
||||
"""Construct the base object to pass to the loader.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
dates : pd.DatetimeIndex
|
||||
The dates we can serve.
|
||||
|
||||
Returns
|
||||
-------
|
||||
args : tuple[any]
|
||||
The arguments to forward to the loader positionally.
|
||||
"""
|
||||
return dates, self.get_dataset()
|
||||
|
||||
def pipeline_event_setup_engine(self, dates):
|
||||
"""
|
||||
Make a Pipeline Enigne object based on the given dates.
|
||||
"""
|
||||
loader = self.loader_type(*self.pipeline_event_loader_args(dates))
|
||||
return SimplePipelineEngine(lambda _: loader, dates, self.asset_finder)
|
||||
|
||||
def get_sids_to_frames(self,
|
||||
zip_date_index_with_vals,
|
||||
vals,
|
||||
date_intervals,
|
||||
dates,
|
||||
dtype_name,
|
||||
missing_dtype):
|
||||
"""
|
||||
Construct a DataFrame that maps sid to the expected values for the
|
||||
given dates.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
zip_date_index_with_vals: callable
|
||||
A function that returns a series of `vals` repeated based on the
|
||||
number of days in the date interval for each val, indexed by the
|
||||
dates in `dates`.
|
||||
vals: iterable
|
||||
An iterable with values that correspond to each interval in
|
||||
`date_intervals`.
|
||||
date_intervals: list
|
||||
A list of date intervals for each sid that correspond to values in
|
||||
`vals`.
|
||||
dates: DatetimeIndex
|
||||
The dates which will serve as the index for each Series for each
|
||||
sid in the DataFrame.
|
||||
dtype_name: str
|
||||
The name of the dtype of the values in `vals`.
|
||||
missing_dtype: str
|
||||
The name of the value that should be used as the missing value
|
||||
for the dtype of `vals` - e.g., 'NaN' for floats.
|
||||
"""
|
||||
frame = pd.DataFrame({sid: get_values_for_date_ranges(
|
||||
zip_date_index_with_vals,
|
||||
vals[sid],
|
||||
pd.DatetimeIndex(list(zip(*date_intervals[sid]))[0]),
|
||||
pd.DatetimeIndex(list(zip(*date_intervals[sid]))[1]),
|
||||
dates
|
||||
).astype(dtype_name) for sid in self.get_sids()[:-1]})
|
||||
frame[self.get_sids()[-1]] = zip_date_index_with_vals(
|
||||
dates, [missing_dtype] * len(dates)
|
||||
).astype(dtype_name)
|
||||
return frame
|
||||
|
||||
@staticmethod
|
||||
def _compute_busday_offsets(announcement_dates):
|
||||
"""
|
||||
Compute expected business day offsets from a DataFrame of announcement
|
||||
dates.
|
||||
"""
|
||||
# Column-vector of dates on which factor `compute` will be called.
|
||||
raw_call_dates = announcement_dates.index.values.astype(
|
||||
'datetime64[D]'
|
||||
)[:, None]
|
||||
|
||||
# 2D array of dates containining expected nexg announcement.
|
||||
raw_announce_dates = (
|
||||
announcement_dates.values.astype('datetime64[D]')
|
||||
)
|
||||
|
||||
# Set NaTs to 0 temporarily because busday_count doesn't support NaT.
|
||||
# We fill these entries with NaNs later.
|
||||
whereNaT = raw_announce_dates == NaTD
|
||||
raw_announce_dates[whereNaT] = make_datetime64D(0)
|
||||
|
||||
# The abs call here makes it so that we can use this function to
|
||||
# compute offsets for both next and previous earnings (previous
|
||||
# earnings offsets come back negative).
|
||||
expected = abs(np.busday_count(
|
||||
raw_call_dates,
|
||||
raw_announce_dates
|
||||
).astype(float))
|
||||
|
||||
expected[whereNaT] = np.nan
|
||||
return pd.DataFrame(
|
||||
data=expected,
|
||||
columns=announcement_dates.columns,
|
||||
index=announcement_dates.index,
|
||||
)
|
||||
|
||||
@parameterized.expand(gen_calendars(
|
||||
'2014-01-01',
|
||||
'2014-01-31',
|
||||
critical_dates=pd.to_datetime([
|
||||
'2014-01-05',
|
||||
'2014-01-10',
|
||||
'2014-01-15',
|
||||
'2014-01-20',
|
||||
], utc=True),
|
||||
))
|
||||
def test_compute(self, dates):
|
||||
engine = self.pipeline_event_setup_engine(dates)
|
||||
cols = self.setup(dates)
|
||||
|
||||
pipe = Pipeline(
|
||||
columns=self.pipeline_columns
|
||||
)
|
||||
|
||||
result = engine.run_pipeline(
|
||||
pipe,
|
||||
start_date=dates[0],
|
||||
end_date=dates[-1],
|
||||
)
|
||||
|
||||
for sid in self.get_sids():
|
||||
for col_name in cols.keys():
|
||||
assert_series_equal(result[col_name].unstack(1)[sid],
|
||||
cols[col_name][sid],
|
||||
check_names=False)
|
||||
|
||||
|
||||
class WithSeededRandomPipelineEngine(WithNYSETradingDays, WithAssetFinder):
|
||||
"""
|
||||
ZiplineTestCase mixin providing class-level fixtures for running pipelines
|
||||
@@ -1239,127 +1051,3 @@ class WithResponses(object):
|
||||
self.responses = self.enter_instance_context(
|
||||
responses.RequestsMock(),
|
||||
)
|
||||
|
||||
|
||||
class WithNextAndPreviousEventDataLoader(WithPipelineEventDataLoader):
|
||||
"""
|
||||
ZiplineTestCase mixin extending common functionality for event data
|
||||
loader tests that have both next and previous events.
|
||||
|
||||
`base_cases` should be used as the template to test cases that combine
|
||||
knowledge date (timestamp) and some 'other_date' in various ways.
|
||||
`next_date_intervals` gives the date intervals for the next event based
|
||||
on the dates given in `base_cases`.
|
||||
`next_dates` gives the next date from `other_date` which is known about at
|
||||
each interval.
|
||||
`prev_date_intervals` gives the date intervals for each sid for the
|
||||
previous event based on the dates given in `base_cases`.
|
||||
`prev_dates` gives the previous date from `other_date` which is known
|
||||
about at each interval.
|
||||
`get_expected_previous_event_dates` is a convenience function that fills
|
||||
a DataFrame with the previously known dates for each sid for the given
|
||||
dates.
|
||||
`get_expected_next_event_dates` is a convenience function that fills
|
||||
a DataFrame with the next known dates for each sid for the given
|
||||
dates.
|
||||
"""
|
||||
base_cases = [
|
||||
# K1--K2--A1--A2.
|
||||
pd.DataFrame({
|
||||
TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-10']),
|
||||
'other_date': pd.to_datetime(['2014-01-15', '2014-01-20']),
|
||||
}),
|
||||
# K1--K2--A2--A1.
|
||||
pd.DataFrame({
|
||||
TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-10']),
|
||||
'other_date': pd.to_datetime(['2014-01-20', '2014-01-15']),
|
||||
}),
|
||||
# K1--A1--K2--A2.
|
||||
pd.DataFrame({
|
||||
TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-15']),
|
||||
'other_date': pd.to_datetime(['2014-01-10', '2014-01-20']),
|
||||
}),
|
||||
# K1 == K2.
|
||||
pd.DataFrame({
|
||||
TS_FIELD_NAME: pd.to_datetime(['2014-01-05'] * 2),
|
||||
'other_date': pd.to_datetime(['2014-01-10', '2014-01-15']),
|
||||
}),
|
||||
pd.DataFrame(
|
||||
columns=['other_date',
|
||||
TS_FIELD_NAME],
|
||||
dtype='datetime64[ns]'
|
||||
),
|
||||
]
|
||||
|
||||
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_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'],
|
||||
['NaT']
|
||||
]
|
||||
|
||||
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']]
|
||||
]
|
||||
|
||||
prev_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'],
|
||||
['NaT']
|
||||
]
|
||||
|
||||
def get_expected_previous_event_dates(self, dates, dtype_name,
|
||||
missing_dtype):
|
||||
return self.get_sids_to_frames(
|
||||
zip_with_dates,
|
||||
self.prev_dates,
|
||||
self.prev_date_intervals,
|
||||
dates,
|
||||
dtype_name,
|
||||
missing_dtype
|
||||
)
|
||||
|
||||
def get_expected_next_event_dates(self, dates, dtype_name, missing_dtype):
|
||||
return self.get_sids_to_frames(
|
||||
zip_with_dates,
|
||||
self.next_dates,
|
||||
self.next_date_intervals,
|
||||
dates,
|
||||
dtype_name,
|
||||
missing_dtype
|
||||
)
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import datetime
|
||||
from functools import partial
|
||||
import inspect
|
||||
|
||||
@@ -339,6 +340,46 @@ def assert_adjustment_equal(result, expected, path=(), **kwargs):
|
||||
)
|
||||
|
||||
|
||||
@assert_equal.register(
|
||||
(datetime.datetime, np.datetime64),
|
||||
(datetime.datetime, np.datetime64),
|
||||
)
|
||||
def assert_timestamp_and_datetime_equal(result,
|
||||
expected,
|
||||
path=(),
|
||||
msg='',
|
||||
allow_datetime_coercions=False,
|
||||
compare_nat_equal=True,
|
||||
**kwargs):
|
||||
"""
|
||||
Branch for comparing python datetime (which includes pandas Timestamp) and
|
||||
np.datetime64 as equal.
|
||||
|
||||
Returns raises unless ``allow_datetime_coercions`` is passed as True.
|
||||
"""
|
||||
assert allow_datetime_coercions or type(result) == type(expected), (
|
||||
"%sdatetime types (%s, %s) don't match and "
|
||||
"allow_datetime_coercions was not set.\n%s" % (
|
||||
_fmt_msg(msg),
|
||||
type(result),
|
||||
type(expected),
|
||||
_fmt_path(path),
|
||||
)
|
||||
)
|
||||
|
||||
result = pd.Timestamp(result)
|
||||
expected = pd.Timestamp(result)
|
||||
if compare_nat_equal and pd.isnull(result) and pd.isnull(expected):
|
||||
return
|
||||
|
||||
assert_equal.dispatch(object, object)(
|
||||
result,
|
||||
expected,
|
||||
path=path,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
|
||||
try:
|
||||
# pull the dshape cases in
|
||||
from datashape.util.testing import assert_dshape_equal
|
||||
|
||||
Reference in New Issue
Block a user