From 7100e60474c13f754d676dec8e108f7275f020ed Mon Sep 17 00:00:00 2001 From: Maya Tydykov Date: Tue, 16 Feb 2016 17:57:01 -0500 Subject: [PATCH] MAINT: accept more data structures, verify, and select loader dynamically. 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. --- tests/pipeline/test_buyback_auth.py | 81 ++++++++++--------- tests/pipeline/test_earnings.py | 8 +- tests/pipeline/test_events.py | 1 + zipline/pipeline/factors/events.py | 2 +- .../pipeline/loaders/blaze/buyback_auth.py | 12 +-- zipline/pipeline/loaders/blaze/earnings.py | 9 ++- zipline/pipeline/loaders/buyback_auth.py | 44 ++++------ zipline/pipeline/loaders/earnings.py | 24 +++--- zipline/pipeline/loaders/events.py | 58 ++++++++++--- zipline/pipeline/loaders/utils.py | 4 +- 10 files changed, 134 insertions(+), 109 deletions(-) create mode 100644 tests/pipeline/test_events.py diff --git a/tests/pipeline/test_buyback_auth.py b/tests/pipeline/test_buyback_auth.py index 6b7218c0..7bc9d6c8 100644 --- a/tests/pipeline/test_buyback_auth.py +++ b/tests/pipeline/test_buyback_auth.py @@ -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, diff --git a/tests/pipeline/test_earnings.py b/tests/pipeline/test_earnings.py index aecc2c88..308f9977 100644 --- a/tests/pipeline/test_earnings.py +++ b/tests/pipeline/test_earnings.py @@ -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) ) diff --git a/tests/pipeline/test_events.py b/tests/pipeline/test_events.py new file mode 100644 index 00000000..8710fd29 --- /dev/null +++ b/tests/pipeline/test_events.py @@ -0,0 +1 @@ +__author__ = 'mtydykov' diff --git a/zipline/pipeline/factors/events.py b/zipline/pipeline/factors/events.py index 127aae29..cee82d2d 100644 --- a/zipline/pipeline/factors/events.py +++ b/zipline/pipeline/factors/events.py @@ -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] diff --git a/zipline/pipeline/loaders/blaze/buyback_auth.py b/zipline/pipeline/loaders/blaze/buyback_auth.py index e92c3b9c..c5505079 100644 --- a/zipline/pipeline/loaders/blaze/buyback_auth.py +++ b/zipline/pipeline/loaders/blaze/buyback_auth.py @@ -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) diff --git a/zipline/pipeline/loaders/blaze/earnings.py b/zipline/pipeline/loaders/blaze/earnings.py index a08c2fd8..7483f951 100644 --- a/zipline/pipeline/loaders/blaze/earnings.py +++ b/zipline/pipeline/loaders/blaze/earnings.py @@ -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) diff --git a/zipline/pipeline/loaders/buyback_auth.py b/zipline/pipeline/loaders/buyback_auth.py index 038a89c0..538919da 100644 --- a/zipline/pipeline/loaders/buyback_auth.py +++ b/zipline/pipeline/loaders/buyback_auth.py @@ -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, diff --git a/zipline/pipeline/loaders/earnings.py b/zipline/pipeline/loaders/earnings.py index e2030430..33c421bb 100644 --- a/zipline/pipeline/loaders/earnings.py +++ b/zipline/pipeline/loaders/earnings.py @@ -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): diff --git a/zipline/pipeline/loaders/events.py b/zipline/pipeline/loaders/events.py index 9920a729..9165581c 100644 --- a/zipline/pipeline/loaders/events.py +++ b/zipline/pipeline/loaders/events.py @@ -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( diff --git a/zipline/pipeline/loaders/utils.py b/zipline/pipeline/loaders/utils.py index ff4ee899..5be70420 100644 --- a/zipline/pipeline/loaders/utils.py +++ b/zipline/pipeline/loaders/utils.py @@ -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