wassname/catalyst
1
0
Fork 0
mirror of https://github.com/wassname/catalyst.git synced 2026-07-07 15:27:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

MAINT: accept more data structures, verify, and select loader dynamically.

Browse Source 
MAINT: add fields based on changes to events loader.

MAINT: modify based on expectations of events loader.

MAINT: modify args.

TST: clean up and clarify df access.

TST: fix bugs in test that didn't properly split datasets' data.

MAINT: fix merge error.
This commit is contained in:
Maya Tydykov
2016-02-16 17:57:01 -05:00
parent a877fcfdb6
commit 7100e60474
10 changed files with 134 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/pipeline/test_buyback_auth.py
+45 -36
View File
@@ -32,7 +32,7 @@ from zipline.pipeline.loaders.blaze import (
TS_FIELD_NAME,
CASH_FIELD_NAME
)
from zipline.utils.numpy_utils import make_datetime64D, np_NaT
from zipline.utils.numpy_utils import make_datetime64D, NaTD
from zipline.utils.test_utils import (
gen_calendars,
make_simple_equity_info,
@@ -217,7 +217,7 @@ class BuybackAuthLoaderCommonTest:
# Set NaTs to 0 temporarily because busday_count doesn't support NaT.
# We fill these entries with NaNs later.
whereNaT = raw_announce_dates == np_NaT
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
@@ -260,8 +260,6 @@ class CashBuybackAuthLoaderTestCase(TestCase, BuybackAuthLoaderCommonTest):
"""
Test for cash buyback authorizations dataset.
"""
buyback_authorizations = {sid: df.drop(SHARE_COUNT_FIELD_NAME, 1)
for sid, df in iteritems(buyback_authorizations)}
pipeline_columns = {
'previous_buyback_cash':
CashBuybackAuthorizations.previous_value.latest,
@@ -278,7 +276,9 @@ class CashBuybackAuthLoaderTestCase(TestCase, BuybackAuthLoaderCommonTest):
tmp_asset_finder(equities=equity_info),
)
cls.cols = {}
cls.buyback_authorizations = buyback_authorizations
cls.buyback_authorizations = {sid: df.drop(SHARE_COUNT_FIELD_NAME, 1)
for sid, df in
iteritems(buyback_authorizations)}
cls.loader_type = CashBuybackAuthorizationsLoader
@classmethod
@@ -325,8 +325,6 @@ class ShareBuybackAuthLoaderTestCase(BuybackAuthLoaderCommonTest, TestCase):
"""
Test for share buyback authorizations dataset.
"""
buyback_authorizations = {sid: df.drop(CASH_FIELD_NAME, 1)
for sid, df in iteritems(buyback_authorizations)}
pipeline_columns = {
'previous_buyback_share_count':
ShareBuybackAuthorizations.previous_share_count.latest,
@@ -343,7 +341,9 @@ class ShareBuybackAuthLoaderTestCase(BuybackAuthLoaderCommonTest, TestCase):
tmp_asset_finder(equities=equity_info),
)
cls.cols = {}
cls.buyback_authorizations = buyback_authorizations
cls.buyback_authorizations = {sid: df.drop(CASH_FIELD_NAME, 1)
for sid, df in
iteritems(buyback_authorizations)}
cls.loader_type = ShareBuybackAuthorizationsLoader
@classmethod
@@ -386,23 +386,6 @@ class ShareBuybackAuthLoaderTestCase(BuybackAuthLoaderCommonTest, TestCase):
self._test_compute_buyback_auth(dates)
def mapping_to_df(mapping):
return (bz.Data(pd.concat(
pd.DataFrame({
BUYBACK_ANNOUNCEMENT_FIELD_NAME:
frame[BUYBACK_ANNOUNCEMENT_FIELD_NAME],
SHARE_COUNT_FIELD_NAME:
frame[SHARE_COUNT_FIELD_NAME],
CASH_FIELD_NAME:
frame[CASH_FIELD_NAME],
TS_FIELD_NAME:
frame[TS_FIELD_NAME],
SID_FIELD_NAME: sid,
})
for sid, frame in iteritems(mapping)
).reset_index(drop=True)),)
class BlazeCashBuybackAuthLoaderTestCase(CashBuybackAuthLoaderTestCase):
""" Test case for loading via blaze.
"""
@@ -416,7 +399,18 @@ class BlazeCashBuybackAuthLoaderTestCase(CashBuybackAuthLoaderTestCase):
BlazeCashBuybackAuthLoaderTestCase,
self,
).loader_args(dates)
return mapping_to_df(mapping)
return (bz.Data(pd.concat(
pd.DataFrame({
BUYBACK_ANNOUNCEMENT_FIELD_NAME:
frame[BUYBACK_ANNOUNCEMENT_FIELD_NAME],
CASH_FIELD_NAME:
frame[CASH_FIELD_NAME],
TS_FIELD_NAME:
frame[TS_FIELD_NAME],
SID_FIELD_NAME: sid,
})
for sid, frame in iteritems(mapping)
).reset_index(drop=True)),)
class BlazeShareBuybackAuthLoaderTestCase(ShareBuybackAuthLoaderTestCase):
@@ -432,7 +426,18 @@ class BlazeShareBuybackAuthLoaderTestCase(ShareBuybackAuthLoaderTestCase):
BlazeShareBuybackAuthLoaderTestCase,
self,
).loader_args(dates)
return mapping_to_df(mapping)
return (bz.Data(pd.concat(
pd.DataFrame({
BUYBACK_ANNOUNCEMENT_FIELD_NAME:
frame[BUYBACK_ANNOUNCEMENT_FIELD_NAME],
SHARE_COUNT_FIELD_NAME:
frame[SHARE_COUNT_FIELD_NAME],
TS_FIELD_NAME:
frame[TS_FIELD_NAME],
SID_FIELD_NAME: sid,
})
for sid, frame in iteritems(mapping)
).reset_index(drop=True)),)
class BlazeShareBuybackAuthLoaderNotInteractiveTestCase(
@@ -458,20 +463,24 @@ class BlazeCashBuybackAuthLoaderNotInteractiveTestCase(
).loader_args(dates)
return swap_resources_into_scope(bound_expr, {})
dtx = pd.date_range('2014-01-01', '2014-01-10')
class BuybackAuthLoaderInferTimestampTestCase(TestCase):
@parameterized.expand([[CashBuybackAuthorizationsLoader],
[ShareBuybackAuthorizationsLoader]])
def test_infer_timestamp(self, loader):
dtx = pd.date_range('2014-01-01', '2014-01-10')
# 'fields' needs to match expected fields for the given loader to
# satisfy column check in constructor.
@parameterized.expand([[CashBuybackAuthorizationsLoader,
{BUYBACK_ANNOUNCEMENT_FIELD_NAME: dtx,
CASH_FIELD_NAME: [0] * 10}],
[ShareBuybackAuthorizationsLoader,
{BUYBACK_ANNOUNCEMENT_FIELD_NAME: dtx,
SHARE_COUNT_FIELD_NAME: [0] * 10}]])
def test_infer_timestamp(self, loader, fields):
events_by_sid = {
# No timestamp column - should index by first given date
0: pd.DataFrame({BUYBACK_ANNOUNCEMENT_FIELD_NAME: dtx}),
0: pd.DataFrame(fields),
# timestamp column exists - should index by it
1: pd.DataFrame(
{BUYBACK_ANNOUNCEMENT_FIELD_NAME: dtx,
TS_FIELD_NAME: dtx}
)
1: pd.DataFrame(dict(fields, **{TS_FIELD_NAME: dtx}))
}
loader = loader(
dtx,
tests/pipeline/test_earnings.py
+4 -4
View File
@@ -386,16 +386,16 @@ class EarningsCalendarLoaderInferTimestampTestCase(TestCase):
announcement_dates.keys(),
)
assert_series_equal(
pd.Series(loader.events_by_sid[0][ANNOUNCEMENT_FIELD_NAME]),
loader.events_by_sid[0].loc[:, ANNOUNCEMENT_FIELD_NAME],
pd.Series(index=[dtx[0]] * 10,
data=dtx,
name=ANNOUNCEMENT_FIELD_NAME),
)
assert_series_equal(
pd.Series(loader.events_by_sid[1][ANNOUNCEMENT_FIELD_NAME]),
pd.Series(index=announcement_dates[1][TS_FIELD_NAME],
loader.events_by_sid[1].loc[:, ANNOUNCEMENT_FIELD_NAME],
pd.Series(index=announcement_dates[1].loc[:, TS_FIELD_NAME],
data=np.array(
announcement_dates[1][ANNOUNCEMENT_FIELD_NAME]
announcement_dates[1].loc[:, ANNOUNCEMENT_FIELD_NAME]
),
name=ANNOUNCEMENT_FIELD_NAME)
)
tests/pipeline/test_events.py
+1
View File
@@ -0,0 +1 @@
__author__ = 'mtydykov'
zipline/pipeline/factors/events.py
+1 -1
View File
@@ -42,7 +42,7 @@ class BusinessDaysSincePreviousEvents(Factor):
announce_dates = arrays[0].astype(datetime64D_dtype)
# Set masked values to NaT.
announce_dates[~mask] = np_NaT
announce_dates[~mask] = NaTD
# Convert row labels into a column vector for broadcasted comparison.
reference_dates = dates.values.astype(datetime64D_dtype)[:, newaxis]
zipline/pipeline/loaders/blaze/buyback_auth.py
+6 -6
View File
@@ -32,7 +32,7 @@ class BlazeCashBuybackAuthorizationsLoader(BlazeEventsCalendarLoader):
The timezeone to use for the data query cutoff.
dataset: DataSet
The DataSet object for which this loader loads data.
loader: EventsLoader
concrete_loader: EventsLoader
The reference loader to use for this dataset.
Notes
@@ -74,7 +74,7 @@ class BlazeCashBuybackAuthorizationsLoader(BlazeEventsCalendarLoader):
data_query_time=None,
data_query_tz=None,
dataset=CashBuybackAuthorizations,
loader=CashBuybackAuthorizationsLoader,
concrete_loader=CashBuybackAuthorizationsLoader,
**kwargs):
super(
BlazeCashBuybackAuthorizationsLoader, self
@@ -84,7 +84,7 @@ class BlazeCashBuybackAuthorizationsLoader(BlazeEventsCalendarLoader):
data_query_time=data_query_time,
data_query_tz=data_query_tz,
dataset=dataset,
loader=loader,
concrete_loader=concrete_loader,
**kwargs)
@@ -106,7 +106,7 @@ class BlazeShareBuybackAuthorizationsLoader(BlazeEventsCalendarLoader):
The timezeone to use for the data query cutoff.
dataset: DataSet
The DataSet object for which this loader loads data.
loader: EventsLoader
concrete_loader: EventsLoader
The reference loader to use for this dataset.
Notes
@@ -148,7 +148,7 @@ class BlazeShareBuybackAuthorizationsLoader(BlazeEventsCalendarLoader):
data_query_time=None,
data_query_tz=None,
dataset=ShareBuybackAuthorizations,
loader=ShareBuybackAuthorizationsLoader,
concrete_loader=ShareBuybackAuthorizationsLoader,
**kwargs):
super(
BlazeShareBuybackAuthorizationsLoader, self
@@ -158,5 +158,5 @@ class BlazeShareBuybackAuthorizationsLoader(BlazeEventsCalendarLoader):
data_query_time=data_query_time,
data_query_tz=data_query_tz,
dataset=dataset,
loader=loader,
concrete_loader=concrete_loader,
**kwargs)
zipline/pipeline/loaders/blaze/earnings.py
+5 -4
View File
@@ -26,7 +26,7 @@ class BlazeEarningsCalendarLoader(BlazeEventsCalendarLoader):
The timezeone to use for the data query cutoff.
dataset: DataSet
The DataSet object for which this loader loads data.
loader: EventsLoader
concrete_loader: EventsLoader
The reference loader to use for this dataset.
Notes
@@ -66,10 +66,11 @@ class BlazeEarningsCalendarLoader(BlazeEventsCalendarLoader):
data_query_time=None,
data_query_tz=None,
dataset=EarningsCalendar,
loader=EarningsCalendarLoader,
concrete_loader=EarningsCalendarLoader,
**kwargs):
super(
BlazeEarningsCalendarLoader, self
).__init__(expr, dataset=dataset, loader=loader, resources=resources,
odo_kwargs=odo_kwargs, data_query_time=data_query_time,
).__init__(expr, dataset=dataset, concrete_loader=concrete_loader,
resources=resources, odo_kwargs=odo_kwargs,
data_query_time=data_query_time,
data_query_tz=data_query_tz, **kwargs)
zipline/pipeline/loaders/buyback_auth.py
+16 -28
View File
@@ -24,31 +24,25 @@ class CashBuybackAuthorizationsLoader(EventsLoader):
event date, cash value)]
"""
expected_cols = frozenset([BUYBACK_ANNOUNCEMENT_FIELD_NAME,
CASH_FIELD_NAME])
def __init__(self,
all_dates,
events_by_sid,
infer_timestamps=False,
dataset=CashBuybackAuthorizations):
dataset=CashBuybackAuthorizations,
expected_cols=expected_cols):
super(CashBuybackAuthorizationsLoader, self).__init__(
all_dates,
events_by_sid,
infer_timestamps=infer_timestamps,
dataset=dataset
dataset=dataset,
expected_cols=expected_cols,
)
def get_loader(self, column):
"""dispatch to the loader for ``column``.
"""
if column is self.dataset.previous_value:
return self.previous_buyback_value_loader
elif column is self.dataset.previous_announcement_date:
return self.previous_event_date_loader
else:
raise ValueError("Don't know how to load column '%s'." % column)
@lazyval
def previous_buyback_value_loader(self):
def previous_value_loader(self):
return self._previous_event_value_loader(
self.dataset.previous_value,
BUYBACK_ANNOUNCEMENT_FIELD_NAME,
@@ -56,7 +50,7 @@ class CashBuybackAuthorizationsLoader(EventsLoader):
)
@lazyval
def previous_event_date_loader(self):
def previous_announcement_date_loader(self):
return self._previous_event_date_loader(
self.dataset.previous_announcement_date,
BUYBACK_ANNOUNCEMENT_FIELD_NAME,
@@ -75,31 +69,25 @@ class ShareBuybackAuthorizationsLoader(EventsLoader):
event date, share value)]
"""
expected_cols = frozenset([BUYBACK_ANNOUNCEMENT_FIELD_NAME,
SHARE_COUNT_FIELD_NAME])
def __init__(self,
all_dates,
events_by_sid,
infer_timestamps=False,
dataset=ShareBuybackAuthorizations):
dataset=ShareBuybackAuthorizations,
expected_cols=expected_cols):
super(ShareBuybackAuthorizationsLoader, self).__init__(
all_dates,
events_by_sid,
infer_timestamps=infer_timestamps,
dataset=dataset
dataset=dataset,
expected_cols=expected_cols,
)
def get_loader(self, column):
"""dispatch to the loader for ``column``.
"""
if column is self.dataset.previous_share_count:
return self.previous_buyback_share_count_loader
elif column is self.dataset.previous_announcement_date:
return self.previous_event_date_loader
else:
raise ValueError("Don't know how to load column '%s'." % column)
@lazyval
def previous_buyback_share_count_loader(self):
def previous_share_count_loader(self):
return self._previous_event_value_loader(
self.dataset.previous_share_count,
BUYBACK_ANNOUNCEMENT_FIELD_NAME,
@@ -107,7 +95,7 @@ class ShareBuybackAuthorizationsLoader(EventsLoader):
)
@lazyval
def previous_event_date_loader(self):
def previous_announcement_date_loader(self):
return self._previous_event_date_loader(
self.dataset.previous_announcement_date,
BUYBACK_ANNOUNCEMENT_FIELD_NAME,
zipline/pipeline/loaders/earnings.py
+9 -15
View File
@@ -10,22 +10,16 @@ ANNOUNCEMENT_FIELD_NAME = "announcement_date"
class EarningsCalendarLoader(EventsLoader):
def __init__(self, all_dates, events_by_sid, infer_timestamps=False,
dataset=EarningsCalendar):
super(EarningsCalendarLoader, self).__init__(all_dates,
events_by_sid,
infer_timestamps,
dataset=dataset)
expected_cols = frozenset([ANNOUNCEMENT_FIELD_NAME])
def get_loader(self, column):
"""Dispatch to the loader for ``column``.
"""
if column is self.dataset.next_announcement:
return self.next_announcement_loader
elif column is self.dataset.previous_announcement:
return self.previous_announcement_loader
else:
raise ValueError("Don't know how to load column '%s'." % column)
def __init__(self, all_dates, events_by_sid,
infer_timestamps=False,
dataset=EarningsCalendar,
expected_cols=expected_cols):
super(EarningsCalendarLoader, self).__init__(
all_dates, events_by_sid, infer_timestamps, dataset=dataset,
expected_cols=expected_cols
)
@lazyval
def next_announcement_loader(self):
zipline/pipeline/loaders/events.py
+45 -13
View File
@@ -10,6 +10,7 @@ from .frame import DataFrameLoader
from .utils import next_date_frame, previous_date_frame, previous_value
TS_FIELD_NAME = "timestamp"
SID_FIELD_NAME = "sid"
class EventsLoader(PipelineLoader):
@@ -32,18 +33,21 @@ class EventsLoader(PipelineLoader):
If the DataFrames do not contain a "timestamp" column, we assume we
knew about the event on all prior dates. This mode is only supported
if ``infer_timestamp`` is explicitly passed as a truthy value.
infer_timestamps : bool, optional
Whether to allow omitting the "timestamp" column.
dataset : DataSet
The DataSet object for which this loader loads data.
expected_cols : frozenset
Set of expected columns for the dataset, without timestamp.
"""
def __init__(self,
all_dates,
events_by_sid,
infer_timestamps=False,
dataset=None):
dataset=None,
expected_cols=frozenset()):
self.all_dates = all_dates
# Do not modify the original in place, since it may be used for other
# purposes.
self.events_by_sid = (
@@ -52,24 +56,52 @@ class EventsLoader(PipelineLoader):
dates = self.all_dates.values
for k, v in iteritems(events_by_sid):
if "timestamp" not in v.columns:
# First, must convert to DataFrame.
if isinstance(v, pd.Series):
# If Series was passed, DateTime index is assumed.
self.events_by_sid[k] = pd.DataFrame(v)
elif isinstance(v, pd.DatetimeIndex):
if not infer_timestamps:
raise ValueError(
"Got DataFrame without a 'timestamp' column for "
"sid %d.\n"
"Got DatetimeIndex for sid %d.\n"
"Pass `infer_timestamps=True` to use the first date in"
" `all_dates` as implicit timestamp."
" `all_dates` as implicit timestamp."% k
)
self.events_by_sid[k] = v = v.copy()
self.events_by_sid[k] = pd.DataFrame(v)
v.index = [dates[0]] * len(v)
# Already a DataFrame
elif isinstance(v, pd.DataFrame):
if TS_FIELD_NAME not in v.columns:
if not infer_timestamps:
raise ValueError(
"Got DataFrame without a '%s' column for sid %d.\n"
"Pass `infer_timestamps=True` to use the first "
"date in `all_dates` as implicit timestamp."%
(TS_FIELD_NAME, k)
)
self.events_by_sid[k] = v = v.copy()
v.index = [dates[0]] * len(v)
else:
self.events_by_sid[k] = v.set_index(TS_FIELD_NAME)
else:
self.events_by_sid[k] = v.set_index("timestamp")
raise ValueError("Data for sid %s must be in DataFrame, "
"Series, or DatetimeIndex."% k)
# Once data is in a DF, make sure columns are correct.
cols_except_ts = (set(v.columns.values) -
{TS_FIELD_NAME} -
{SID_FIELD_NAME})
# Check that all columns other than timestamp are as expected.
if cols_except_ts != expected_cols:
raise ValueError(
"Expected columns %s for sid %s but got columns %s." %
(expected_cols, k, v.columns.values)
)
self.dataset = dataset
@abstractmethod
def get_loader(self):
raise NotImplementedError("Must implement 'get_loader'.")
def get_loader(self, column):
if column in self.dataset.columns:
return getattr(self, "%s_loader" % column.name)
raise ValueError("Don't know how to load column '%s'." % column)
def load_adjusted_array(self, columns, dates, assets, mask):
return merge(
zipline/pipeline/loaders/utils.py
+2 -2
View File
@@ -6,7 +6,7 @@ import pandas as pd
from six import iteritems
from six.moves import zip
from zipline.utils.numpy_utils import NaTns
from zipline.utils.numpy_utils import NaTns, NaTD
def next_date_frame(dates, events_by_sid):
@@ -83,7 +83,7 @@ def previous_date_frame(date_index, events_by_sid):
next_date_frame
"""
sids = list(events_by_sid)
out = np.full((len(date_index), len(sids)), np_NaT, dtype='datetime64[ns]')
out = np.full((len(date_index), len(sids)), NaTD, dtype='datetime64[ns]')
d_n = date_index[-1].asm8
for col_idx, sid in enumerate(sids):
# events_by_sid[sid] is Series mapping knowledge_date to actual